Skip to content

Models ​

physicsnemo.models holds 25 architecture families. You do not need to know them all; you need to know which shape of data you have, because that is what picks the family.

Pick by data shape ​

Your data isUseWhat meshio++ gives you
A regular grid, same resolution in and outFNO, AFNO, UNetthe grid data path and the fno/afno families in TrainSpec (2-D through Grid.Squeeze); grid, voxelize and write_vti produce the lattice
A coarse grid in, a fine grid outSRResNet (srrn)the grid data path, a paired Target in the manifest, and the srresnet family in TrainSpec
An unstructured mesh with connectivityMeshGraphNet familygraph_sample — nodes, edges, features, targets
An unordered point cloudTransolver, FIGConvNet, FLAREfeature_matrix gives the node table, select_points the token budget
A CAD surface plus a volume (external aero)DoMINOextract_surface plus compute_sdf as a nodal field
A time series of statesRNN, mesh_reducedTimeSeries and the target_offset pairing
Particles with trajectoriesMeshGraphNet, VFGNmesh edges only; proximity graphs are a gap
A distribution, not a single answerdiffusion U-Nets, DiTnothing yet
The globeGraphCast, Pangu, FengWu, DLWPnothing; no mesh counterpart at all

The same choice as a chart — start from the shape of one sample and follow the arrows:

A decision chart from the shape of one sample to an architecture family

Figure 1: What one sample looks like decides the family. Everything else is detail.

The families ​

Neural operators on grids ​

FamilyClassesNotes
fnoFNOThe reference neural operator. dimension=1..4; the 4-D case is a spatio-temporal block operator
afnoAFNO, ModAFNOFourier operator with an adaptive token mixer. ModAFNO takes a timestep input, so simulated time conditions it
dpotDPOTNetDenoising pre-trained operator transformer, 2-D and 3-D
unetUNetA plain 3-D convolutional U-Net
pix2pixPix2PixImage-to-image translation
srrnSRResNetSuper-resolution residual network — coarse in, fine out

Graphs and meshes ​

FamilyClassesNotes
meshgraphnetMeshGraphNet, BiStrideMeshGraphNet, HybridMeshGraphNet, MeshGraphKANEncode–process–decode on an edge graph. The variants add a multiscale hierarchy, proximity "world" edges, and KAN layers
mesh_reducedMesh_Reduced, Sequence_ModelReduce a mesh, then learn the dynamics in the reduced space with temporal attention
graphcastGraphCastNetThe weather architecture, on an icosahedral grid
vfgnVFGNLearnedSimulator and its encode/process/decode partsVirtual Foundry GraphNet, for sintering and additive manufacturing

Point clouds and transformers ​

FamilyClassesNotes
transolverTransolverAttention over learned "physics slices" of a point cloud
geotransolverGeoTransolverGeometry-aware successor
figconvnetFIGConvUNetFactorized implicit grids; per-point fields plus a scalar (drag-style) head
flareFLAREExperimental attention successor (also under experimental.models)
dominoDoMINOExternal aerodynamics: CAD surface plus volume, with pretrained checkpoints

Generative ​

FamilyClassesNotes
diffusion_unetsSongUNet, SongUNetPosEmbd, DhariwalUNet, CorrDiffRegressionUNet, UNet, StormCastUNetThe denoiser backbones. 2-D image oriented
ditDiTDiffusion transformer — the denoiser to reach for on non-image data
topodiffTopoDiffDiffusion for topology optimization

A volumetric 3-D denoiser exists too, under physicsnemo.experimental.models.diffusion_unets.DiffusionUNet3D — see Companion packages.

Sequences and weather ​

rnn (One2ManyRNN, Seq2SeqRNN), dlwp (DLWP), dlwp_healpix (HEALPixUNet, HEALPixRecUNet), pangu (Pangu), fengwu (Fengwu), swinvrnn (SwinRNN).

The plain one ​

mlp (FullyConnected) — a multilayer perceptron. It is the right first model far more often than it looks: if your input is a handful of case parameters and your output is a field, you want this, not a neural operator.

Worth knowing about, on nobody's critical path ​

FamilyClassWhat it would bring
dpotDPOTNeta PDE foundation model (AFNO mixing, pretrained across equation families) to fine-tune on your own grids
topodiffTopoDiffgenerative topology optimization with constraint channels, on compliance data
pix2pixPix2Pix, Pix2PixUneta plain convolutional image-to-image translator; fits any grid pipeline mechanically
experimental.xdeeponetDeepONetbranch (parameters) plus trunk (coordinates) operator learning — parameters in, field at the mesh nodes out, without a POD basis; the deeponet family in TrainSpec, over each entry's Metadata
experimental.globeGLOBEboundary-driven elliptic problems, from named boundary meshes
experimental.aerojepaAeroJEPAself-supervised pretraining on geometry alone, before any labels exist
experimental.strata, experimental.healdaStrata, VideoHealDAweather emulation on the sphere and HEALPix data assimilation; the assimilation idea matters for digital twins, the API is calendar-shaped
pangu, fengwu, swinvrnn, dlwp_healpixas namedglobal weather architectures with no mesh counterpart

Building blocks ​

physicsnemo.nn holds the layers the models are made of — around 150 of them. Two are worth knowing by name:

  • ConcreteDropout — dropout with a learned rate, which makes MC-dropout uncertainty estimates meaningfully calibrated instead of arbitrary;
  • physicsnemo.nn.functional.derivatives — differential operators used by the physics-informed path.

In meshio++ ​

Five families are wired end to end, one per data shape in the chart: MeshGraphNet, through TrainSpec's Model.Name: "meshgraphnet", graph_sample for the tensors and predict for the write-back; SRResNet, through Model.Name: "srresnet" over the grid data path's coarse/fine pairs; FNO and AFNO, through "fno"/"afno" over the same grid path with the coarse grid paired with itself, 2-D through the thin-axis squeeze (AFNO is 2-D only and patches a fixed image, so the squeeze is required and the sample shape must divide the patch); and the experimental DeepONet, through "deeponet", whose branch reads each entry's Metadata parameters and whose trunk is the mesh's own points — parameters in, field out, on a fixed geometry. See the neural-operator families and parameters in, field out. The graph family came first for a reason — a mesh with connectivity is the shape meshio++ natively holds, and PyG batches ragged graphs of different sizes natively, which is what makes a dataset of differently-sized meshes trainable with no padding convention of its own.

Everything else on the chart above is reachable by hand: feature_matrix gives you the node table any point-cloud model wants, select_points reduces it to a token budget, to_torch/to_dlpack move it to the device without a file round trip, and the model is then ordinary PyTorch. The rest of the dataset half has since shipped too — proximity_graph (v10.31.0) builds the radius/kNN graphs a particle method needs, and make_window/iter_windows/rollout (v10.33.0) give a sequence model its history and its autoregressive evaluation loop.

Next: Data and datapipes — feeding them.

Released under the MIT License.