Metadata-Version: 2.4
Name: feax
Version: 0.1.0
Summary: A GPU-accelerated finite element analysis framework with JAX.
Author-email: Naruki-Ichihara <ichihara.naruki@nihon-u.ac.jp>
Maintainer-email: Naruki-Ichihara <ichihara.naruki@nihon-u.ac.jp>
License: GPL-3.0-or-later
Project-URL: Homepage, https://github.com/Naruki-Ichihara/feax
Project-URL: Documentation, https://github.com/Naruki-Ichihara/feax
Project-URL: Repository, https://github.com/Naruki-Ichihara/feax.git
Project-URL: Issues, https://github.com/Naruki-Ichihara/feax/issues
Keywords: JAX,CUDA,Simulation,PINN,finite-element,FEA,GPU
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: scipy
Requires-Dist: matplotlib
Requires-Dist: meshio
Requires-Dist: fenics-basix
Requires-Dist: gmsh
Requires-Dist: equinox
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx; extra == "docs"
Requires-Dist: sphinx-rtd-theme; extra == "docs"
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]; extra == "cuda12"
Requires-Dist: nvidia-cudss-cu12; extra == "cuda12"
Requires-Dist: nvidia-cublas-cu12; extra == "cuda12"
Requires-Dist: nvidia-cuda-runtime-cu12; extra == "cuda12"
Requires-Dist: nvidia-cudnn-cu12; extra == "cuda12"
Requires-Dist: nvidia-cuda-nvcc-cu12; extra == "cuda12"
Requires-Dist: nanobind; extra == "cuda12"
Requires-Dist: scikit-build-core>=0.5; extra == "cuda12"
Requires-Dist: ninja; extra == "cuda12"
Dynamic: license-file


# FEAX 

[![License](https://img.shields.io/badge/license-GPL%20v3-blue.svg)](LICENSE)
[![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
[![JAX](https://img.shields.io/badge/JAX-0.7%2B-green.svg)](https://github.com/google/jax) <img src="assets/logo.svg" alt="logo" width=150 align="right"></img>

**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

<div align="center">
<img src="assets/feax_jax_transformations.svg" alt="JAX Transformations" width="800">
</div>

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

```bash
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](LICENSE) for the full license text.

## Acknowledgments

FEAX builds upon the excellent work of:
- [JAX](https://github.com/google/jax) for automatic differentiation and compilation
- [JAX-FEM](https://github.com/tianjuxue/jax_fem) for inspiration and reference implementations

---
<div align="center">
<img src="assets/logo.svg" alt="logo" width=150></img>
</div>

