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

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},
}