CARDIAX is a GPU-accelerated differentiable finite element package based on JAX, forked from JAX-FEM. CARDIAX focuses on solving cardiac mechanics problems in a single, unified framework from Laplace-Dirichlet-Rule-Based method for fiber generation to biventricular inverse models. This package is actively managed by the Willerson Center for Cardiovascular Modeling and Simulation (WCCMS) and is constantly adapting to accommodate the suite of problems we are intereseted in solving. We are only focused on GPU development.
Before installing cardiax be sure to install jax at JAX Install. Verify that the GPU is seen by running the following to see if CUDA devices are found
import jax
print(jax.devices())Once the jax installation is working, the easiest option is to build all the dependencies through a conda environment using environment.yaml which also installs JAX with CUDA. Another option is to build a virtual environment with venv and pip for the dependencies. We have a SC_requirements.txt for installing on the supercomputer systems (TACC) where only venv is allowed. This can be modified to run on your own cluster if there are strict versioning requirements. These files set up the dependecies, but to install cardiax, you must go inside the directory ../CARDIAX and run
pip install -e .
Warning: JAX just upgraded to 0.9, so we haven't created the latest environment yet.
In the documentation, there are examples that walk through how to use the code. These are under tutorials, but the files are markdown format to explain functionality. The corresponding *.py files live in the CARDIAX/tutorials directory. More examples are also given in the demos but without the lengthy markdown explanations.
While JAX supports CPU, CARDIAX is not being tested on CPU environments. We created this codebase to fully leverage GPUs, but the functionality should remain consistent. Also, multi-GPU functionality is also not available. The problems we are currently solving can fit on the memory of a single GPU, so we will not develop this parallelization until needed.
The finite element limitations are quite vast compared to some other codes. Not to say these can't be addressed, but they would take some work. This repo is consistently being developed, but here are some problems we don't support:
- Mixed FE fields with off-diagonal block entries
- Fluid-structure interaction
- Manipulation of test function
This project is licensed under the GNU General Public License v3 - see the LICENSE for details.
If you're using this project, you can cite this work here. Please also cite JAX-FEM's work as well! Here is their repo.
We'll add a list of others papers built upon this framework below: