Installation

This guide provides detailed instructions for installing FEAX and its dependencies.

JAX

FEAX requires JAX (0.11 or later, required by asdex) as its core dependency. JAX provides automatic differentiation and JIT compilation capabilities. FEAX does not install JAX for you — no extra pulls it in — so install the variant matching your hardware first.

CPU-only Installation

For CPU-only usage (Linux/macOS/Windows):

pip install -U jax

This is sufficient for learning and small-scale problems.

NVIDIA GPU (CUDA 12)

For NVIDIA GPU acceleration with CUDA 12:

pip install -U "jax[cuda12]"

Requirements: NVIDIA driver ≥525

NVIDIA GPU (CUDA 13)

For CUDA 13:

pip install -U "jax[cuda13]"

For more details, see the official JAX installation guide.

PyPI

Install FEAX

Once JAX is installed, install FEAX using pip:

pip install feax

This will automatically install all required dependencies including:

numpy, scipy - Numerical computing
pandas - Data handling
fenics-basix - Finite element basis functions
tetgen - Tetrahedralization behind sphere_mesh / cylinder_mesh
pyvista - Mesh I/O and visualization
scikit-image - Marching cubes for density / SDF surface extraction
matplotlib - Plotting
equinox, optax - Pytree utilities and optimizers
asdex - Sparse automatic differentiation (feax.asd)

Gmsh is not a required dependency any more. Core meshing (box_mesh, sphere_mesh, cylinder_mesh, structured grids) works without it; only the feax.gene adaptive remeshing and SDF / density → volume-mesh extraction need it, via the gmsh extra below.

Install from Source

To get the latest development version:

pip install git+https://github.com/Naruki-Ichihara/feax.git@main

For development with editable install:

git clone https://github.com/Naruki-Ichihara/feax.git
cd feax
pip install -e .

Optional Extras

FEAX provides optional dependency groups via pyproject.toml:

Extra	Contents	Usage
`.[cuda12]`	cuDSS + cuBLAS + cuDNN for CUDA 12	GPU acceleration (CUDA 12, without JAX)
`.[cuda13]`	cuDSS + cuBLAS + cuDNN for CUDA 13	GPU acceleration (CUDA 13, without JAX)
`.[sksparse]`	`scikit-sparse`	CPU host-side direct solvers (`cholmod`, `umfpack`)
`.[amg]`	`amjax` + `pyamg`	Algebraic multigrid preconditioner (`fe.amg` / `fe.AMGOptions`)
`.[matfree]`	`matfree`	Matrix-free eigensolver utilities (linear buckling)
`.[gmsh]`	`gmsh`	`feax.gene` adaptive remeshing and SDF / density → volume-mesh extraction
`.[cad]`	`cadjoint` (installed from GitHub, not PyPI)	Differentiable CAD front end used as a preprocessor; see `examples/experimental/sdf_shape_optimization.py`
`.[mesh]`	(empty)	Kept so `pip install feax[mesh]` keeps working; TetGen is a core dependency now
`.[dev]`	pytest, black, ruff, mypy	Development and testing

scikit-sparse (Optional)

The cholmod and umfpack direct solvers wrap scikit-sparse, which compiles against the SuiteSparse C library. Without it, DirectSolverOptions() (auto) falls back to JAX's spsolve on CPU or cudss on GPU — no functionality is lost for typical use.

To enable cholmod / umfpack, install the system library first, then the extra:

# Debian / Ubuntu / Google Colab
apt-get install -y libsuitesparse-dev
pip install feax[sksparse]

For GPU use outside Docker, install JAX first, then the CUDA extra:

pip install -U "jax[cuda13]"
pip install "feax[cuda13]"

cuDSS Direct Solver

When using the cuDSS direct solver, the spineax package is also required. Install it after feax and with --no-deps, so it does not pull in (and reinstall) jax[cuda13] / NumPy over the versions your environment already has:

pip install --no-deps lineax spineax

lineax is spineax's only pure-Python dependency not already provided by feax[cuda13] (which supplies equinox, jaxtyping, and the cuDSS / cuSPARSE CUDA libraries). Installing spineax without --no-deps reinstalls jax[cuda13] and upgrades NumPy, which breaks any environment that must keep its existing JAX (Docker's NVCR base, Colab).

Note: The official PyPI release ships an x86_64-only wheel. On other architectures (e.g. aarch64 / GB10), build from source instead (still --no-deps):
pip install --no-deps lineax
pip install --no-build-isolation --no-deps git+https://github.com/johnviljoen/spineax.git
The Docker image does this automatically (official wheel on x86_64, source build on other arches).

Docker

FEAX provides a Dockerfile based on NVIDIA's JAX image (nvcr.io/nvidia/jax:26.08-py3, JAX 0.11 / CUDA 13), which includes JAX pre-compiled for CUDA 13. JAX is not reinstalled during the build: the image's jax / jaxlib versions are pinned through PIP_CONSTRAINT=/etc/pip-constraints-jax.txt, which also stays active inside the container.

Note: Inside the container, a pip install whose dependencies require a different JAX fails with ResolutionImpossible instead of replacing the base JAX (a PyPI jax/jaxlib over the NVCR build breaks the CUDA plugin with PJRT_FFI_UserData_Add_Args size errors). Install such packages with --no-deps.

Build Arguments

Argument	Default	Description
`INSTALL_DOCS`	`false`	Install Node.js 20 + pydoc-markdown + Docusaurus dependencies

Build Docker Image

Standard build:

git clone https://github.com/Naruki-Ichihara/feax.git
cd feax
docker build -t feax:latest .

Build with Docusaurus support (for docs development):

docker build --build-arg INSTALL_DOCS=true -t feax:latest .

Run Container

docker run --gpus all -it feax:latest

Docker Compose (Recommended)

For development with persistent volumes and GPU support:

cd feax
docker-compose up -d
docker exec -it feax bash

The docker-compose configuration includes:

GPU support (deploy.resources.reservations.devices)
Volume mounting for development (./:/workspace)
WSL2 display support (for GUI applications like gmsh)
Shared memory configuration (shm_size: 4gb)

Docs Development

To develop the documentation site locally, Node.js 20+ and pydoc-markdown are required.

Install Dependencies

# Node.js 20
curl -fsSL https://deb.nodesource.com/setup_20.x | bash -
apt-get install -y nodejs

# pydoc-markdown (API doc generator)
pip install pydoc-markdown

# Docusaurus + npm packages
cd docs
npm install

Start Dev Server

Use the provided script to generate API docs and start the server in one step:

./docs/dev.sh

Or manually:

cd docs
npm run api:generate   # Generate API reference from Python source
npm run start          # Start dev server at http://localhost:3000

Build for Production

cd docs
npm run build

Colab

You can use FEAX in Google Colab notebooks with free GPU/TPU access.

Basic Setup

In a Colab notebook cell, install system dependencies via the bundled helper script. It mirrors the Dockerfile: installs gmsh + Python bindings, swig, NLopt, BLAS/LAPACK, gmsh runtime libraries, builds SuiteSparse 7.x from source (with SUITESPARSE_USE_CUDA=OFF to avoid an nvcc compute_52 failure on Colab GPU), and installs the nlopt Python bindings:

!curl -fsSL https://raw.githubusercontent.com/Naruki-Ichihara/feax/main/scripts/colab_setup.sh | bash
!pip install git+https://github.com/Naruki-Ichihara/feax.git

The setup step takes ~5–8 minutes on Colab (one-time, dominated by the SuiteSparse build).

Note: scikit-sparse is not required — the default DirectSolverOptions() auto-selects cudss on GPU or spsolve on CPU. Once colab_setup.sh has run, you can opt in to cholmod / umfpack with:
!SUITESPARSE_INCLUDE_DIR=/usr/local/include/suitesparse \
 SUITESPARSE_LIBRARY_DIR=/usr/local/lib \
 pip install "feax[sksparse] @ git+https://github.com/Naruki-Ichihara/feax.git"
The SUITESPARSE_* env vars point scikit-sparse's setup.py at the source-built SuiteSparse under /usr/local — its default search only covers /usr/include/suitesparse (the apt path).

With GPU Support

Colab provides CUDA-enabled GPUs by default. Run the setup script, then install the GPU extras:

!curl -fsSL https://raw.githubusercontent.com/Naruki-Ichihara/feax/main/scripts/colab_setup.sh | bash
!SUITESPARSE_INCLUDE_DIR=/usr/local/include/suitesparse \
 SUITESPARSE_LIBRARY_DIR=/usr/local/lib \
 pip install "feax[cuda13,sksparse] @ git+https://github.com/Naruki-Ichihara/feax.git"
!pip install --no-deps lineax spineax

# Verify GPU is available
import jax
print(jax.devices())  # Should show GPU

--no-deps is required. A plain pip install spineax pulls spineax's declared jax[cuda13] dependency, which reinstalls JAX and upgrades NumPy to 2.x — that forces a "You must restart the runtime" and clashes with Colab's NumPy-pinned packages (e.g. numba wants numpy<2.1). With --no-deps, only spineax and its one extra pure-Python dep lineax are installed; equinox / jaxtyping / the cuDSS + cuSPARSE libraries come from feax[cuda13]. Colab is x86_64, so the prebuilt wheel installs directly (nothing is compiled).

Verification

After installation, verify your setup:

import jax
import feax
import jax.numpy as np

print(f"JAX version: {jax.__version__}")
print(f"FEAX version: {feax.__version__}")
print(f"JAX devices: {jax.devices()}")
print(f"JAX 64-bit enabled: {jax.config.jax_enable_x64}")

Expected output:

JAX version: 0.4.x
FEAX version: 0.6.1
JAX devices: [CpuDevice(id=0)] or [GpuDevice(id=0)]
JAX 64-bit enabled: True

Next Steps

Try the Basic Tutorial to get started
Check out example scripts
Read the API documentation

JAX​

CPU-only Installation​

NVIDIA GPU (CUDA 12)​

NVIDIA GPU (CUDA 13)​

PyPI​

Install FEAX​

Install from Source​

Optional Extras​

scikit-sparse (Optional)​

cuDSS Direct Solver​

Docker​

Build Arguments​

Build Docker Image​

Run Container​

Docker Compose (Recommended)​

Docs Development​

Install Dependencies​

Start Dev Server​

Build for Production​

Colab​

Basic Setup​

With GPU Support​

Verification​

Next Steps​

JAX