/BSMS-GNN

Primary LanguagePythonApache License 2.0Apache-2.0

Bi-stride Multi-Scale GNN

This repository contains the code implementations for Efficient Learning of Mesh-Based Physical Simulation with BSMS-GNN (ICML 2023). The paper is also available on Arxiv.

Motivations

We focus on developing a multi-scale graph neural network for physics-based simulation. Previous works have certain limitations when it comes to building multi-scale connectivity.

  • GraphUNet (Gao et al., 2019) has additional scoring modules to select the most informative nodes for constructing coarser levels. They adopt a power-of-2 adjacency enhancement to prevent loss of connectivity. However, this enhancement does not guarantee connectivity preservation.
  • MS-GNN-Grid (Lino et al., 2021) uses background helper grids to build the coarser levels. However, this approach can blur boundaries that are spatially close but not necessarily geodesically close.
  • MultiScale MeshGraphNets (Fortunato et al., 2022) uses manually drawn coarser meshes for the same domain, but this requires a significant amount of additional labor.

We aimed to find a solution that would be consistent across any input graphs, without introducing blurring effects on cross-boundary edges, while preserving correct connectivity and minimizing additional labor.

Method

Bi-stride pooling

We drew inspiration from bipartite graphs, where nodes can be split into two groups, and the minimum geometric distance between the two groups is exactly one hop away. This property allows a simple power-of-2 adjacency enhancement to preserve connectivity. We extend this idea to a general mesh:

  1. Select an initial node.
  2. Perform Breadth-First-Search on a general mesh, marking the geodesic distance to the initial node.
  3. Pool nodes at every other level and apply the power-of-2 adjacency enhancement.

This process ensures that the connectivity is preserved at any depth of coarser level.

Pipeline

  1. Before training, we employ bi-stride pooling as a pre-processing step to determine the multi-level graph for the input mesh.
  2. Based on the multi-scale connectivities, we then determine non-parametric transition modules.
  3. These advantages eliminate the need for additional overhead such as scoring modules or matrix enhancement during training.

Results

Our dataset includes the following: 1) cylinder flow, 2) compressible flow around an airfoil, 3) elastic plate, and 4) inflating elastic surface. The multi-scale structure of these datasets, achieved through bi-stride pooling, is shown below:

The method performs well on all datasets, demonstrating significant improvements in training and inference time as well as RAM consumption.

The absence of cross-boundary edges helps avoid artificial blurring effects.

Bi-stride pooling consistently works on unseen geometry, leading to higher accuracy.

Overall, we achieve the lowest inference error compared to previous methods in the most contact-rich test case.

Requirements

  • Pytorch
  • PyG
  • Numpy
  • h5py
  • TensorBoard
  • SciPy
  • scikit-learn
  • sparse-dot-mkl

Download datasets and pretrained models

We host the datasets and pretrained models on this link. Please maintain the file structure shown below to run the script by default.

this project
│   ...    
│
└───data
│   └───cylinder
|       └───outputs_test
|       └───outputs_train
|       └───outputs_valid
│       │   meta.json
│   └───...
└───res
│   └───cylinder
|       └───ours
|       |   └───ckpts
|       |       |   *.pt
│   └───...

If you store the data and result folders somewhere else, you can modify the data_dir and dump_dir in the config files accordingly.

How to use

# ./run_BSMS.sh $case_name ./configs/$case_name $mode $restart_epoch

# case_name: [cylinder, airfoil, plate, font]
# ./configs/$case_name: stores the corresponding config files of a case
# mode: [0:train, 1:local test, 2: global rollout]
# restart_epoch: -1 (or leave blank) to train from the start; 0, 1... to reload the stored ckpts of a certain frame

# e.g. train font from scratch
./run_BSMS.sh font ./configs/font 0 -1
# e.g. local test RMSE of cylinder at epoch 19
./run_BSMS.sh cylinder ./configs/cylinder 1 19
# e.g. global rollout RMSE of airfoil at epoch 39
./run_BSMS.sh airfoil ./configs/airfoil 2 39

Citation

If you find this method useful, please cite it using the following format:

@inproceedings{cao2023efficient,
  title     = {Efficient Learning of Mesh-Based Physical Simulation with Bi-Stride Multi-Scale Graph Neural Network},
  author    = {Cao, Yadi and Chai, Menglei and Li, Minchen and Jiang, Chenfanfu},
  booktitle = {International Conference on Machine Learning},
  year      = {2023},
  url       = {https://openreview.net/forum?id=2Mbo7IEtZW}
}

ICML Poster

poster