APACHE-2.0 License
After graduation, the original host (UCLA's Google Drive) for the dataset is unavailable. I took this opportunity to significantly improve the repository as follows:
misc
folder, which contains the full documentation.I would like to thank Zijie Huang for inviting me to help integrate this framework into Nvidia's Modulus in July 2024. Many of the code revisions in this version were developed during that collaboration project. However, due to time constraints, not all functions, such as multi-batch training, were implemented. This provided the motivation for me to complete the full revision.
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.
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.
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.
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:
This process ensures that the connectivity is preserved at any depth of coarser level.
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.
I now use Poetry for environment management. Setting up is simple:
conda env create -f environment.yml
conda activate bsms-gnn
poetry install --no-root
All parameters are now handled by Hydra configuration. Please refer to the configs folder to understand the parameters better.
# Run a training, defaulting to the airfoil case
# Set board=true to upload info to a WandB board
python src/train.py board=true
# Run a rollout using a specific checkpoint
# Assign the datasets, model to be cylinder_flow
python src/rollout.py restore_dir=./bsms-res/ckpts/train/cylinder_flow/20240806-140921/ restore_step=0 datasets=cylinder_flow model=cylinder_flow
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}
}