PyTorch version of Learning Mesh-Based Simulation with Graph Networks (ICLR 2021)

Source repository: https://github.com/deepmind/deepmind-research/tree/master/meshgraphnets

Video site: sites.google.com/view/meshgraphnets

Paper: arxiv.org/abs/2010.03409

Overview

The code in this repository is the PyTorch version of Learning Mesh-Based Simulation with Graph Networks. Currently, the code of cloth simulation can be run on both windows and linux. Development environment is PyCharm 2021.1.3. Package dependencies are defined in requirements.txt, please install all dependencies before runnning. The result of this version of MGN network is not 100% conform to the original code, please tune it according to your need. Other GNN networks code will be added to this repository in the future to explore the GNN performance for physical simulation.

New Features compared to original MeshGraphNets

The novel ripple model, inspired by water ripple, is based on deepmind's meshgraphnets and utilizes ripples to enhance the information propagation. Further new features such as different aggregation methods, attention and stochastic message passing are also already added.

Setup

Install dependencies:

pip install -r requirements.txt

Download a dataset:

mkdir -p ${DATA}
bash meshgraphnets/download_dataset.sh flag_simple ${DATA}

Go to the dataset directory and generate .idx file(needed by package tfrecord for reading .tfrecord file in PyTorch):

python -m tfrecord.tools.tfrecord2idx <file>.tfrecord <file>.id

Configure running goals by setting the variables and flags variables at the beginning of run_model.py, which includes running mode (training/evaluation), model, epochs, saving path and etc.

Running the model

Run the code after configurating the code:

python run_model.py

Plot a trajectory (rollout.pkl path can be defined inside plot_cloth.py):

python plot_cloth.py

Datasets

Datasets can be downloaded using the script download_dataset.sh. They contain a metadata file describing the available fields and their shape, and tfrecord datasets for train, valid and test splits. Dataset names match the naming in the paper. The following datasets are available:

airfoil
cylinder_flow
deforming_plate
flag_minimal
flag_simple
flag_dynamic
sphere_simple
sphere_dynamic

flag_minimal is a truncated version of flag_simple, and is only used for integration tests.