LagrangianMGN

<p align="center"> <strong> <span style="font-size: 30px;">LagrangianMGN</span> </strong> </p>

Model Overview

LagrangianMGN is a graph-network model developed by researchers at DeepMind for simulating complex physical systems. It rapidly predicts the dynamics of particle-based systems, including fluids, rigid bodies, and deformable materials.

Paper: Learning to simulate complex physics with graph networks https://arxiv.org/abs/2002.09405

Model Description

LagrangianMGN uses a message-passing graph-network architecture trained on the Lagrangian particle simulation dataset to perform long-horizon dynamical simulations of complex systems involving fluids, rigid bodies, and deformable materials.

Use Cases

Usage

1. OneCode

Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:

Launch OneCode for one-click AI4S programming

2. Manual Setup

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.

Download the Model Package

hf download OneScience-Group/LagrangianMGN --local-dir ./LagrangianMGN
cd LagrangianMGN

Set Up the Runtime Environment

DCU Environment

# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Installation with uv is also supported
pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Activate Conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# Installation with uv is also supported
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Training Data

The OneScience community provides the DeepMind Lagrangian dataset for training. Download it with the command below and verify that the data path in conf/config.yaml is configured correctly:

hf download --repo-type dataset OneScience-Group/lagrangian --local-dir ./data/

Training

Single GPU:

python scripts/train.py

Multiple GPUs:

torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py

Model Weights

This repository will provide weights trained on the DeepMind Lagrangian dataset in the weights/ directory. The weights will be uploaded soon.

Inference

python scripts/inference.py

The inference script loads a checkpoint from resume_dir, whose default value is weight/checkpoints.

Evaluation and Visualization

python scripts/result.py

Official OneScience Resources

Citations and License

• Original LagrangianMGN paper: Learning to simulate complex physics with graph networks.
• This repository retains the relevant source and attribution notices. Follow all applicable license requirements when using, modifying, or distributing its contents.