Online version: https://www.sciencedirect.com/science/article/pii/S0045782523004759
This repository contains Python scripts to perform the Physics-Informed Graph Neural Network (PI-GNN) emulation experiments as described in the above paper. Please cite this paper if you use the code:
@article{daltonPIGNN2023,
title = {Physics-informed graph neural network emulation of soft-tissue mechanics},
journal = {Computer Methods in Applied Mechanics and Engineering},
volume = {417},
pages = {116351},
year = {2023},
issn = {0045-7825},
author = {David Dalton and Dirk Husmeier and Hao Gao},
keywords = {Soft-tissue mechanics, Graph neural networks, Physics-informed machine learning},
}
Experiments are performed in Python, using the JAX, Flax and Optax libraries.
Assuming conda is installed, a virtual environment with the required packages can be set up using the below shell commands. Firstly, create and activate a new environment.
conda create --name pignnEmulationEnv
conda activate pignnEmulationEnv
OPTIONAL - install Paraview, which is used to create 3D visualisations of the results.
conda install -c conda-forge paraview
Finally, use pip to install remaining packages.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install flax
pip install tensorflow
pip install absl-py
pip install -U scikit-learn
The specific versions of these packages we used were jax 0.4.16
(jaxlib 0.4.16+cuda11.cudnn86
), flax 0.7.4
, tensorflow 2.13.0
, absl-py 2.0.0
and scikit-learn 1.3.1
. Other required packages, such as optax, are installed automatically with the above commands.
Unfortunately I have found it difficult sometimes to install paraview on Linux, depending on the machine (it generally works fine on Windows). Or sometimes it will install correctly, but then a warning appears when I try to load the paraview.simple
package. For additional installation instructions, see here and here. If you cannot get paraview installed correctly, the rest of the code relating to training and evaluating the PIGNN emulator will still work.
Various datasets from the manuscript are included in the data
subdirectory. Physics-informed training and subsequent evaluation on the test data can then be run by calling main.py
from the shell. For example, to train for 1000 epochs on the Liver model and then predict on the test data, run:
python -m main --mode="train_and_evaluate" --data_path="Liver" --n_epochs=1000
Pre-trained emulation parameters for the TwistingBeam, Liver and LeftVentricle models are stored in emulationResults/trainedParameters
. Instructions on how to make predictions using these parameters are given in PRETRAINED_EMULATORS.md
.
Tensorboard can be used to monitor training, for example for the Liver data by running:
tensorboard --logdir=emulationResults/Liver
and then following the instructions printed to the console.
data/LeftVentricle
assumes a fixed blood pressure of 8 mmHg is applied to the inner surface of the LV. This is slightly different to the model considered in the paper, where the blood pressure varied in the range [4,10] mmHg. The data handling code had been written assumed that pressure would be fixed. During the revision process, a reviewer asked us to consider emulation of the LeftVentricle model where the pressure value varies. To account for varying pressure values, we made use of an adjustment which breaks generalisability of the code, so we do not make it public. The results in either case are almost identical.main.py
Main script to call PI-GNN emulation
models.py
Implements PrimalGraphEmulator GNN emulation architecture
utils_*
Various script containing utility functions for data handling, potential energy evaluation, model training, etc
PRETRAINED_EMULATORS.md
Details how predictions can be made using pre-trained emulation parameters
data
Stores data on each model (such as geometry, constitutive law, etc) and test simulation data generated using FEniCS used for evaluating of emulation performance. See the data/README.md
file for further details
emulationResults
Stores the trained emulator parameters and predictions.