OneScience/LagrangianMGN

LagrangianMGN

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

Use Case	Description
Multi-material interaction simulation	Model interactions among fluids, particles, rigid bodies, and deformable materials
Long-horizon physical prediction	Use autoregressive rollouts to predict system evolution over hundreds or thousands of steps
Rapid local validation	Use synthetic data to validate data loading, model training, inference, and result visualization

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

modelscope download --model OneScience/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:

modelscope download --dataset OneScience/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

Platform	OneScience Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

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.