ketiswp/meshgraphnet-hf

MeshGraphNet — cylinder_flow

MeshGraphNet (DeepMind, ICLR 2021) for flow around a cylinder. Given the velocity at every mesh node, it predicts the velocity change for the next time step.

Usage

from transformers import AutoModel

model = AutoModel.from_pretrained("ketiswp/meshgraphnet-hf", trust_remote_code=True)

inputs = model.prepare_inputs(cells, mesh_pos, node_type, velocity)  # one cylinder_flow time step
delta = model(**inputs).delta_velocity                              # [N, 2]
trajectory = model.rollout(**inputs, num_steps=599)                 # [600, N, 2]

Training

Trained on DeepMind cylinder_flow with the original settings: one frame per step (frames 1..598, shuffle buffer 10000), velocity noise std 0.02, Adam with lr 1e-4 * 0.1^(step/5e6) + 1e-6, 10M steps, first 1000 steps only collect normalizer statistics, Sonnet weight init, float32.

split_first_linear holds the first Linear of each layer per input (W [h_s; h_r; e] = W_s h_s + W_r h_r + W_e e), so sender/receiver projections are computed once per node instead of building the per-edge concatenation. Same weights and outputs as without it.

Results

CylinderFlow, in the format of Table 1 of the paper.

	# nodes (avg.)	# steps	t_model ms/step	t_full ms/step	t_GT ms/step	RMSE 1-step ×10⁻³	RMSE rollout-50 ×10⁻³	RMSE rollout-all ×10⁻³	GPU
paper	1885	600	21	23	820	2.34 ± 0.12	6.3 ± 0.7	40.88 ± 7.2	NVIDIA Tesla V100
this model	1885	600	4.0	4.4	–	2.35 ± 0.09	6.67 ± 0.30	55.83 ± 3.97	NVIDIA RTX 6000 Ada Generation

This model: 100 test trajectories; ± is the standard error over trajectories. Timings on one GPU, batch 1, float32; t_full includes rebuilding the graph from the mesh every step.

License

Apache-2.0, as the original DeepMind and medav/meshgraphnets-torch code.