Hyper-Graph Network Simulator (Hyper-GNS)

This repository accompanies the mini-project "Simulating Physics using Hypergraph Networks" which was part of the course "Representation Learning on Graphs and Networks" (L45) at the University of Cambridge in Lent term 2023. The main goal was to investigate whether hypergraphs can improve the prediction accuracy of physics simulations by capturing higher-order interactions between particles. A sneak peek at our experimental results can be found below:

	Ground truth	GNS	Hyper-GNS	Hybrid model
WaterDrops
WaterRamps
Sand
MultiMaterial

Graph Network Simulator (GNS)

Krishna Kumar, The University of Texas at Austin.

Joseph Vantassel, Texas Advanced Computing Center, UT Austin.

Graph Network-based Simulator (GNS) is a framework for developing generalizable, efficient, and accurate machine learning (ML)-based surrogate models 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.

Run GNS

Training

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

Resume training

To resume training specify model_file and train_state_file:

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

Rollout

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"

Render

python3 -m gns.render_rollout --output_mode="gif" --rollout_dir="<path-containing-rollout-file>" --rollout_name="rollout_0"

The renderer also writes .vtu files to visualize in ParaView.

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

Datasets

The data loader provided with this PyTorch implementation utilizes the more general .npz format. The .npz format includes a list of tuples of arbitrary length where each tuple is for a different training trajectory and is of the form (position, particle_type). 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, ...):

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

We provide the following datasets:

• WaterDropSample (smallest dataset)
• Sand
• SandRamps

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

Building 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 based on https://arxiv.org/abs/2002.09405 and https://github.com/deepmind/deepmind-research/tree/master/learning_to_simulate.

Acknowledgement

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

Citation

Repo

Kumar, K., & Vantassel, J. (2022). Graph Network Simulator: v1.0.1 (Version v1.0.1) [Computer software]. https://doi.org/10.5281/zenodo.6658322

Dataset

Vantassel, Joseph; Kumar, Krishna (2022) “Graph Network Simulator Datasets.” DesignSafe-CI. https://doi.org/10.17603/ds2-0phb-dg64 v1