Home

CARDIAX-NNFE  —  developed and maintained by

GPU-accelerated Neural Network Finite Element for parameterized cardiac mechanics — a scientific machine learning framework for learning parameter-to-solution maps defined by PDE residuals, built natively on JAX with CARDIAX as the finite element backend.

Get Started View on GitHub

---

What is NNFE?

The Neural Network Finite Element (NNFE) method learns the parameter-to-solution mapping of a PDE system directly from its residual. Rather than solving a new boundary value problem for each parameter configuration, NNFE trains a network to approximate the solution — enabling fast, accurate evaluations across a continuous parameterized space.

---

Key Features

• GPU-Native

  ---

  Built entirely on JAX. JIT compilation, automatic differentiation, and vmap-based batching are first-class citizens throughout the codebase.

• Residual-Based Training

  ---

  No labeled solution data required. Training is driven entirely by the finite element residual, computed via CARDIAX.

• Configuration-Driven

  ---

  Full experiments are reproducible from YAML configuration files. Each run spawns a unique directory with saved configs, model weights, and plots.

• Modular Architecture

  ---

  Networks, optimizers, samplers, and FE problems are independently configurable. Mix and match components or extend with custom implementations.

---

Stack

Component	Role
JAX	Automatic differentiation, JIT compilation, GPU acceleration
CARDIAX	Finite element residual computation
Equinox	Neural network construction and JAX-compatible PyTree handling
Optax	Optimization and learning rate scheduling
Lineax	GPU-accelerated linear operators

---

Installation

Install JAX with CUDA support first — see the JAX installation guide. Then build the full environment:

conda env create -f environment.yaml
conda activate cardiax

Install CARDIAX from GitHub, then install NNFE in editable mode:

pip install -e .

Verify GPU visibility before proceeding:

import jax
print(jax.devices())  # Should show CUDA devices

---

Citation

If you use CARDIAX-NNFE in your research, please cite CITATION.cff.

Full citation details will be updated upon paper acceptance in SoftwareX.