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](https://arxiv.org/abs/2411.15178) # 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 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 ```bash 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** ```bash # 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** ```bash # 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. ```bash 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`. ```bash python scripts/train.py ``` Multi-device training: ```bash torchrun --standalone --nproc_per_node=