siddik-lanl/spus-pde-unet-36m-v2

SPUS: Small PDE U-Net Solver (36M) — Version 2

A lightweight residual U-Net foundation model for solving partial differential equations, released by Los Alamos National Laboratory.

This 36M-parameter convolutional U-Net was pretrained from scratch using the four Compressible Euler datasets used for SPUS v1, together with two additional Navier–Stokes datasets: NS-Sines and NS-Gaussians. It uses the same general architecture as v1, but its weights were trained independently from a fresh initialization.

• —Original model: https://huggingface.co/siddik-lanl/spus-pde-unet-36m
• —Code: https://github.com/lanl/SPUS-Small-PDE-U-net-Solver
• —Paper: https://arxiv.org/abs/2510.01370

Usage

Download model.py with the model weights, then run:

import torch
from huggingface_hub import PyTorchModelHubMixin
from model import Unet2D

class SPUSUnet2D(Unet2D, PyTorchModelHubMixin):
    pass

model = SPUSUnet2D.from_pretrained(
    "siddik-lanl/spus-pde-unet-36m-v2"
)
model.eval()

x = torch.randn(1, 1, 5, 128, 128)

with torch.no_grad():
    prediction = model(x)

Released under the BSD-3-Clause License.