FEA in JAX

fea-in-jax is a Finite Element Analysis (FEA) library written in JAX. It leverages JAX's composable function transformations—JIT compilation, automatic differentiation, and vectorization—to provide a high-performance solver capable of running on GPUs and TPUs.

Features

• GPU Acceleration: Native support for hardware acceleration via JAX.
• Differentiability: Differentiate through the physics simulation for gradient-based optimization and machine learning integration.
• Batched Computation: Designed to efficiently handle large batches of elements and quadrature points.

Project Structure

• src/fe_jax: Core library source code, including element definitions, quadrature rules, and solver implementations.
• tests: extensive test suite that also serves as a catalogue of usage examples.
• docs: Documentation and theoretical background.

Getting Started

Prerequisites

• Python 3.10+
• CUDA Toolkit (optional, strictly for GPU acceleration)

Installation

1. Clone the repository:

   git clone <repository_url>
   cd fea-in-jax

2. Set up a virtual environment (recommended):

   python -m venv .venv
   source .venv/bin/activate  # On Windows: .venv\Scripts\activate

3. Install dependencies:

   pip install -r requirements.txt

   Note: jax installation instructions vary depending on your hardware (CPU, GPU, TPU). Please refer to the JAX installation guide if the default pip install does not match your system configuration.

   For development, install the package in editable mode with the test dependency:

   pip install -e ".[dev]"

4. (Optional) Install pyamgx: To enable GPU-accelerated algebraic multigrid preconditioners:

   1. Install NVIDIA's AMGX.
   2. Install pyamgx.

Running Tests

To verify the installation and run the test suite:

pytest tests

Usage

The tests directory contains numerous examples demonstrating how to define meshes, apply boundary conditions, and solve boundary value problems.

• Basic Linear Elasticity: See tests/test_simple_fea_solve.py for a straightforward example.
• Complex Scenarios: See tests/test_fea_solve.py.

Theory and Implementation

For detailed information on the nonlinear solver derivation, handling of Dirichlet boundary conditions, and internal variable definitions, please refer to docs/theory.md.

Resources

• JAX Interoperability: External Callbacks
• Scientific Computing in JAX:
  • Wrapping Scipy KD trees
  • HPC Lecture Notes
• Performance Optimization:
  • JAX GPU Performance Tips
  • JAX AOT Compilation
  • Multi-Process/Distributed Support

Profiling Performance

• JAX Profiling Docs
• NVIDIA JAX Toolbox
• NSys-JAX Wrapper
• JAX Device Memory Profiling
• jax-smi (GPU Memory Tracking)

To profile time and memory for JIT-compiled sections:

1. Collect trace: jax.profiler.start_trace("<directory>/prof")
2. Visualize: using TensorBoard or xprof.

Public Release Information

Distribution Statement A. Approved for public release: distribution is unlimited. Case #: AFRL-2025-4644