feax 0.1.0

A GPU-accelerated finite element analysis framework with JAX.

FEAX

FEAX (Finite Element Analysis with JAX) is a compact, high-performance finite element analysis engine built on JAX. It provides an API for solving partial differential equations on XLA.

What is FEAX?

FEAX combines automatic differentiation with finite element methods. It's designed for:

• Differentiable Physics: Compute gradients through entire FE simulations for optimization, inverse problems, and machine learning
• High Performance: JIT compilation and vectorization through JAX for maximum computational efficiency

JAX Transformations in FEAX

FEAX leverages JAX's powerful transformation system to enable:

• Automatic Differentiation: Compute exact gradients through finite element solvers
• JIT Compilation: Compile to optimized machine code for maximum performance
• Vectorization: Efficiently process multiple scenarios in parallel with vmap
• Parallelization: Scale across multiple devices with pmap

Installation

pip install feax[cuda12]
pip install --no-build-isolation git+https://github.com/Naruki-Ichihara/spineax.git

Note: spineax requires CUDA toolkit and cuDSS to be available at build time. The --no-build-isolation flag is required so that spineax can find JAX and NVIDIA libraries installed by feax[cuda12].

feax.flat

Flat (Feax Lattice) is a utility for asymptotic homogenization of lattice unit cell.

feax.gene

Gene (Generative design in FEAX) is a comprehensive toolkit for topology optimization and generative design. It provides efficient, JAX-native implementations of common topology optimization components.

Key Features

• Response Functions: Compliance and volume fraction calculations optimized for topology optimization
• Filtering Methods:
  • PDE-based Helmholtz filter for smooth, physically-motivated designs
  • Distance-based density filter for efficient spatial smoothing
  • Sensitivity filter for mesh-independent gradient smoothing
• Constrained Optimization: MDMM (Modified Differential Multiplier Method) for handling equality and inequality constraints with automatic differentiation
• Pure JAX Implementation: Fully differentiable and compatible with optax optimizers

License

FEAX is licensed under the GNU General Public License v3.0. See LICENSE for the full license text.

Acknowledgments

FEAX builds upon the excellent work of:

• JAX for automatic differentiation and compilation
• JAX-FEM for inspiration and reference implementations

---