AMG

Model Overview

AMG is a multi-graph neural operator for solving partial differential equations on arbitrary geometries. It was proposed by researchers from the Hong Kong University of Science and Technology (Guangzhou), Beihang University, and other institutions. The model uses multiscale graphs, a physics graph, and the dynamic graph-attention mechanism in GraphFormer to learn PDE solution operators on regularly or irregularly discretized domains. This package is trained for the ShapeNet-Car/CarCFD experiment. Given a three-dimensional vehicle geometry and mesh-node features, it predicts surface pressure and the surrounding airflow velocity field for rapid aerodynamic evaluation and drag-coefficient estimation.

This repository is an independent reproduction of the AMG ShapeNet-Car experiment, implemented through the OneScience workflow from the paper description and official configuration.

Paper: Harnessing Scale and Physics: A Multi-Graph Neural Operator Framework for PDEs on Arbitrary Geometries

Model Description

AMG uses an encoder, a multi-graph processor, and a decoder. It takes discrete mesh-point coordinates and node features as input. A local graph captures high-frequency local variations, a global graph models long-range spatial dependencies, and a physics graph aggregates and propagates latent physical properties. The GraphFormer module performs cross-node interaction with dynamic graph attention and produces the PDE solution at the original discrete points. For the ShapeNet-Car/CarCFD task, the model predicts three velocity components and pressure on a three-dimensional unstructured mesh containing 32,186 points.

Use Cases

Use Case Description
Three-dimensional automotive external-flow prediction Predict the three-dimensional velocity field around complex vehicle geometries and pressure on the vehicle surface
Automotive aerodynamic evaluation Use predicted velocity and pressure fields to support drag-coefficient estimation and aerodynamic assessment
PDE learning on unstructured meshes Build PDE surrogate models for three-dimensional unstructured meshes and irregular geometric domains
CFD surrogate modeling Replace part of a costly numerical-simulation workflow to accelerate vehicle-shape design and batch evaluation of design candidates

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 from Hugging Face

pip install -U huggingface_hub
hf download OneScience-Group/AMG-ShapeNet-Car --local-dir ./AMG-ShapeNet-Car
cd AMG-ShapeNet-Car

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 ShapeNetCar training data on Hugging Face. Download it with the command below, then update data.root in config/config.yaml so that it points to the downloaded preprocessed_data directory.

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

The dataset contains 889 three-dimensional automotive external-flow samples organized into nine shards, param0 through param8. It includes 690 training samples, 99 validation samples, and 100 test samples. Each sample contains 32,186 unstructured mesh points and the following files:

  • x.npy: node input features with shape (32186, 7); AMG uses the four features at indices [3, 4, 5, 6]
  • y.npy: ground-truth CFD solution with shape (32186, 4), containing three velocity components followed by pressure
  • pos.npy: three-dimensional mesh-point coordinates with shape (32186, 3)
  • surf.npy: vehicle-surface node mask with shape (32186,)
  • edge_index.npy: unstructured-mesh edge connectivity with shape (2, E)

Training

Single-device training:

  • Training parameters are read from the data, model, and training sections of config/config.yaml.
  • The default configuration follows the paper and trains for 500 epochs. Add --smoke for a minimal end-to-end test.
  • Before training, set data.root to the ShapeNetCar dataset's preprocessed_data directory.
  • For the two-device, eight-hour profile, use --profile fast_8h_2dcu.
python scripts/train.py

Multi-device training:

torchrun --standalone --nproc_per_node=<num_GPUs> scripts/train.py

Two-device, eight-hour training profile:

torchrun --standalone --nproc_per_node=2 scripts/train.py --profile fast_8h_2dcu

Training reports the training loss, validation pressure relative L2 error, and validation velocity relative L2 error. The best checkpoint is selected using the mean validation relative L2 error and saved to:

./weight/best_model.pt

Model Weights

The weight/ directory contains best_model.pt, the best AMG checkpoint obtained from this ShapeNet-Car/CarCFD reproduction. It can be loaded directly for inference.

Inference

python scripts/inference.py

Predictions, per-sample metrics, and aggregate metrics are saved under results/.

Evaluation and Visualization

python scripts/result.py

The evaluation report, paper-metric comparison plots, training curves, and flow-field visualizations are saved under results/.

Official OneScience Resources

Citations and License

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train OneScience-Group/AMG-ShapeNet-Car

Paper for OneScience-Group/AMG-ShapeNet-Car