Graph Network Simulator (GNS) and MeshNet

Krishna Kumar, The University of Texas at Austin. Joseph Vantassel, Texas Advanced Computing Center, UT Austin. Yongjin Choi, The University of Texas at Austin.

Graph Network-based Simulator (GNS) is a generalizable, efficient, and accurate machine learning (ML)-based surrogate simulator for particulate and fluid systems using Graph Neural Networks (GNNs). GNS code is a viable surrogate for numerical methods such as Material Point Method, Smooth Particle Hydrodynamics and Computational Fluid dynamics. GNS exploits distributed data parallelism to achieve fast multi-GPU training. The GNS code can handle complex boundary conditions and multi-material interactions.

MeshNet is a scalable surrogate simulator for any mesh-based models like Finite Element Analysis (FEA), Computational Fluid Dynammics (CFD), and Finite Difference Methods (FDM).

Run GNS/MeshNet

Training GNS/MeshNet on simulation data

# For particulate domain,
python3 -m gns.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>" -ntraining_steps=100
# For mesh-based domain,
python3 -m meshnet.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>" -ntraining_steps=100

Resume training

To resume training specify model_file and train_state_file:

# For particulate domain,
python3 -m gns.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>"  --model_file="model.pt" --train_state_file="train_state.pt" -ntraining_steps=100
# For mesh-based domain,
python3 -m meshnet.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>"  --model_file="model.pt" --train_state_file="train_state.pt" -ntraining_steps=100

Rollout prediction

# For particulate domain,
python3 -m gns.train --mode="rollout" --data_path="<input-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>" --model_file="model.pt" --train_state_file="train_state.pt"
# For mesh-based domain,
python3 -m meshnet.train --mode="rollout" --data_path="<input-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>" --model_file="model.pt" --train_state_file="train_state.pt"

Render

# For particulate domain,
python3 -m gns.render_rollout --output_mode="gif" --rollout_dir="<path-containing-rollout-file>" --rollout_name="<name-of-rollout-file>"
# For mesh-based domain,
python3 -m gns.render --rollout_dir="<path-containing-rollout-file>" --rollout_name="<name-of-rollout-file>"

In particulate domain, the renderer also writes .vtu files to visualize in ParaView.

GNS prediction of Sand rollout after training for 2 million steps.

In mesh-based domain, the renderer writes .gif animation.

Meshnet GNS prediction of cylinder flow after training for 1 million steps.

Datasets

Particulate domain:

We use the numpy .npz format for storing positional data for GNS training. The .npz format includes a list of tuples of arbitrary length where each tuple corresponds to a differenet training trajectory and is of the form (position, particle_type). The data loader provides INPUT_SEQUENCE_LENGTH positions, set equal to six by default, to provide the GNS with the last INPUT_SEQUENCE_LENGTH minus one positions as input to predict the position at the next time step. The position is a 3-D tensor of shape (n_time_steps, n_particles, n_dimensions) and particle_type is a 1-D tensor of shape (n_particles).

The dataset contains:

Metadata file with dataset information (sequence length, dimensionality, box bounds, default connectivity radius, statistics for normalization, ...):

{
  "bounds": [[0.1, 0.9], [0.1, 0.9]], 
  "sequence_length": 320, 
  "default_connectivity_radius": 0.015, 
  "dim": 2, 
  "dt": 0.0025, 
  "vel_mean": [5.123277536458455e-06, -0.0009965205918140803], 
  "vel_std": [0.0021978993231675805, 0.0026653552458701774], 
  "acc_mean": [5.237611158734309e-07, 2.3633027988858656e-07], 
  "acc_std": [0.0002582944917306106, 0.00029554531667679154]
}

npz containing data for all trajectories (particle types, positions, global context, ...):

Training datasets for Sand, SandRamps, and WaterDropSample are available on DesignSafe Data Depot [@vantassel2022gnsdata].

We provide the following datasets:

WaterDropSample (smallest dataset)
Sand
SandRamps

Download the dataset DesignSafe DataDepot. If you are using this dataset please cite Vantassel and Kumar., 2022

Mesh-based domain:

We also use the numpy .npz format for storing data for training meshnet GNS.

The dataset contains:

npz containing python dictionary describing mesh data and relevant dynamics at mesh nodes for all trajectories. The dictionary includes {pos: (ntimestep, nnodes, ndims), node_type: (ntimestep, nnodes, ntypes), velocity: (ntimestep, nnodes, ndims), pressure: (ntimestep, nnodes, 1), cells: (ntimestep, ncells, 3)}

The dataset is shared on DesignSafe DataDepot. If you are using this dataset please cite Kumar and Choi., 2023

Installation

GNS uses pytorch geometric and CUDA. These packages have specific requirements, please see [PyG installation]((https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) for details.

CPU-only installation on Linux

conda install -y pytorch torchvision torchaudio cpuonly -c pytorch
conda install -y pyg -c pyg
conda install -y pytorch-cluster -c pyg
conda install -y absl-py -c anaconda 
conda install -y numpy dm-tree matplotlib-base pyevtk -c conda-forge

You can use the WaterDropletSample dataset to check if your gns code is working correctly.

To test the code you can run:

pytest test/

To test on the small waterdroplet sample:

git clone https://github.com/geoelements/gns-sample

TMP_DIR="./gns-sample"
DATASET_NAME="WaterDropSample"

mkdir -p ${TMP_DIR}/${DATASET_NAME}/models/
mkdir -p ${TMP_DIR}/${DATASET_NAME}/rollout/

DATA_PATH="${TMP_DIR}/${DATASET_NAME}/dataset/"
MODEL_PATH="${TMP_DIR}/${DATASET_NAME}/models/"
ROLLOUT_PATH="${TMP_DIR}/${DATASET_NAME}/rollout/"

python -m gns.train --data_path=${DATA_PATH} --model_path=${MODEL_PATH} --ntraining_steps=10

Building GNS environment on TACC (LS6 and Frontera)

to setup a virtualenv

sh ./build_venv.sh

check tests run sucessfully.
start your environment

source start_venv.sh

Inspiration

PyTorch version of Graph Network Simulator and Mesh Graph Network Simulator are based on:

Acknowledgement

This code is based upon work supported by the National Science Foundation under Grant OAC-2103937.

Citation

Repo

Kumar, K., & Vantassel, J. (2023). GNS: A generalizable Graph Neural Network-based simulator for particulate and fluid modeling. Journal of Open Source Software, 8(88), 5025. https://doi.org/10.21105/joss.05025

Dataset

Vantassel, Joseph; Kumar, Krishna (2022) “Graph Network Simulator Datasets.” DesignSafe-CI. https://doi.org/10.17603/ds2-0phb-dg64 v1
Kumar, K., Y. Choi. (2023) "Cylinder flow with graph neural network-based simulator." DesignSafe-CI. https://doi.org/10.17603/ds2-fzg7-1719

summ2020 / GNS