High-performance Jax/CUDA-based library for differentiable materials modelling, designed to integrate physics-based models into modern machine learning workflows.
Enables gradient-based optimisation of ML surrogate models for material simulations, by combining efficient computation with scalable ML infrastructure which can run on CPUs and GPUs.
Materials with fine microstructure, such as carbon fibre composites, are expensive to simulate with classical PDE methods. Upscaling methods require a large number of simulations to infer distributions of material parameters. In addition, it is often desirable to provide
- Sensitivity of output to input parameters
- Support for running on CPU and GPU hardware
Machine learning surrogate models can reduce runtime but require:
- Fast code for data generation and inference
- Differentiable solvers to use methods like Physics Enhanced Deep Surrogates
- Differentiable materials models implemented in JAX
- Automatic differentiation for optimisation and ML training
- GPU acceleration of differentiable solvers via JAX
- Bespoke CUDA solvers for data generation and inference
- Modular design for extending models and components
- Compatible with ML pipelines and optimisation frameworks
The following figures compares the performance of the (an-) isotropic JAX and CUDA solvers when applied to an isotropic material.
Clone and run
pip install jaxmaterials
for the JAX-Python library. See detailed instructions below for CUDA support.
This repository contains the following code:
A highly efficient CUDA accelerated solver of the linear elasticity equation in isotropic materials based on the Lippmann Schwinger method by [Moulinec and Suquet, 1998. Computer Methods in Applied Mechanics and Engineering, 157(1-2), pp.69-94].
A Jax implementation of the same method, which allows back-propagation through the solver for later use in a ML setting. Solvers for both isotropic and anisotropic materials have been implemented.
In addition to the plain Lippmann Schwinger solver, the code also supports Anderson acceleration as described in [Wicht, Schneider and Boehlke, T., 2021. International Journal for Numerical Methods in Engineering, 122(9), pp.2287-2311]. Since any Jax code is inherently differentiable, the solver can be used as a building block in a machine learning framework (see below).
Both solvers use the same discretisation as the AMITEX solver, which is described in [Gelebart 2020. Comptes Rendus. Mecanique, 348(8-9), pp.693-704]. For mathematical details see the ./doc subdirectory.
Code for sampling the Lame parameters which can be used as an input to the solvers. The generated samples contain cross-ply layers of fibres as shown in the following figure:
A ML toolchain which uses the above code to train machine learning models that predict the efficient elasticity tensor for a given pair of Lame-parameter fields will be added later.
The CUDA solver requires a working cuda installation, including the NVidia CUDA Toolkit which contains the NVidia CuFFT library. A working C++ compiler and CMake is required to compile and install the solver. To run the automated tests, the GoogleTest is required.
See pyproject.toml for a list of required Python packages.
The following instructions should work on Linux machines, but will need to be adapted on Windows and Mac.
- Clone this repository
- Change to the
cudasubdirectory - Configure in the
builddirectory with
cmake -B build -DCMAKE_INSTALL_PREFIX=<INSTALL_DIR>
where <INSTALL_DIR> is the directory where the solver library should be installed. If the -DCMAKE_INSTALL_PREFIX=<INSTALL_DIR> is omitted, the default (usually /usr/lib/) is used, and you will likely need root access to install the library in this location.
- Build the solver with
cmake --build build
- (Optionally), if the google test framework is installed, test the library by running
./build/bin/test
- Install the library by running
cmake --install build
- Add the install directory to
LD_LIBRARY_PATHto ensure that it can be loaded from Python
export LD_LIBRARY_PATH=<INSTALL_DIR>:${LD_LIBRARY_PATH}
In the main directory of the repository run
python -m pip install .
Optionally, add --editable flag for an editable install.
Run the tests suite with
python -m pytest