wjenkinson/sph-particle-gnn

SPH Particle Dynamics GNN

A physics-informed Graph Neural Network that predicts per-particle displacements and velocities for 2D SPH (Smoothed Particle Hydrodynamics) fluid simulations.

Model Description

Property	Value
Architecture	Edge-conditioned GNN with residual message passing
Node features	[is_fluid, is_solid, is_wall, is_piston, vx, vy] (6 channels)
Edge features	[Δx, Δy, Δz, ‖r‖] (4 channels) — relative positions + distance
Output heads	Displacement (dx, dy, dz) + Velocity (vx, vy)
Particle types	1 = fluid, 2 = solid, 3 = wall, 4 = piston
Neighbour radius	0.003 (with periodic BC in x)
Domain	x ∈ [−0.1, 0.1], y ∈ [−0.08, 0.13]

Training

• Loss: grid-based SPH density field MSE + kinetic energy density MSE (fluid) + per-node MSE (non-fluid boundaries)
• Multistep rollout loss: backpropagates through K autoregressive steps to mitigate error accumulation
• Grid splatting: sparse local-patch cubic spline kernel on a 70×67 cell grid

Intended Use

Autoregressive rollout of 2D particle-based fluid dynamics (e.g. extrusion, dam-break, jet impact).

Not intended for: 3D simulations, turbulent flows, or real-time safety-critical applications.

How to Use

import torch
from model import SimpleGnnPredictor, load_checkpoint

model, hparams = load_checkpoint("model.pt")

# Build a torch_geometric Data object for your particle state
# (see inference.py for a complete example)
displacement, velocity = model(data)
next_pos = current_pos + displacement

See inference.py for a self-contained single-step and multi-step rollout example.

Checkpoint Format

{
    "model_state_dict": OrderedDict,   # model weights
    "model_hparams": {
        "in_channels": 6,
        "hidden_channels": int,        # e.g. 64 or 128
        "edge_channels": 4,
        "num_layers": int,             # e.g. 2 or 3
    },
}

Limitations

• Autoregressive rollout accumulates error over long horizons
• Trained on a specific 2D extrusion geometry; generalisation to other setups requires fine-tuning
• CPU-only training/inference (no GPU optimisation yet)

Citation

If you use this model, please cite the repository:

@software{sph_gnn_2025,
  title  = {SPH Particle Dynamics GNN},
  url    = {https://huggingface.co/YOUR_USERNAME/sph-particle-gnn},
  year   = {2025},
}