Nvidia PhysicsNemo Notes

Notes on NVIDIA PhysicsNeMo, read straight from the main branch as of 2 August 2026. Everything below is grounded in the source and docstrings of examples/weather and torch-harmonics rather than in the documentation prose. The organising idea: the sphere refuses a regular grid, and every model in the package is a different way of paying that debt.

Section 1

Geometry: what the sphere costs you

Before any architecture, one structural fact determines the shape of every design in this package. It is worth stating precisely, because it explains why two apparently unrelated grids — the cubed sphere and HEALPix — carry exactly the same number of defects.

1.1  Discretization

Tile the sphere with quadrilaterals. Euler's formula gives V − E + F = 2, and every face having four edges gives 4F = 2E, hence V = F + 2. Summing the degree deficit over all vertices:

Σ_v (4 − deg v) = 4V − 2E = 4(F+2) − 4F = 8

No quadrilateral mesh on a sphere can have every vertex of degree four. The total deficit is always exactly 8. A cubed sphere pays it as eight three-valent corners. HEALPix, whose twelve base pixels form a rhombic-dodecahedron topology (12 faces, 24 edges, 14 vertices — 8 of degree 3 and 6 of degree 4), pays exactly the same eight. An equirectangular latitude–longitude grid dumps all eight onto the two poles, buying regular indexing at the price of unbounded anisotropy. The triangular analogue is Σ(6 − deg v) = 12, which is why every icosahedral refinement carries exactly twelve pentagonal vertices, forever.

The five discretizations that appear in PhysicsNeMo, and where each one lives in the code:

Cross-check: is HEALPix consistent with GFDL X-SHiELD?

No — they are different tilings, and the mismatch is worth understanding before anyone tries to train a HEALPix model on X-SHiELD output.

SHiELD is powered by FV3, the Finite-Volume Cubed-Sphere dynamical core. Its storm-resolving configuration, X-SHiELD, runs a 3.25 km global domain on a C3072 cubed sphere with 79 hybrid-sigma-pressure layers — a figure visible even in the DYAMOND archive path, 20200120.00Z.C3072.L79x2. So X-SHiELD's geometry belongs to the same family as PhysicsNeMo's dlwp example, not its dlwp_healpix one.

What the two grids share is exactly the topology established above: both are quadrilateral tilings, and both carry precisely eight three-valent corners. What they do not share:

• Equal area. HEALPix is equal-area by construction. FV3's gnomonic cubed sphere is not — cell areas vary between face centre and corner by roughly a factor of two for the equidistant gnomonic variant, and about 1.3 to 1.5 for the equiangular one.
• Cell orientation. This is the decisive one for machine learning. Karlbauer et al. (2024) moved DLWP from the cubed sphere to HEALPix precisely because every HEALPix cell shares a consistent east–west orientation, which permits a single location-invariant convolution kernel across the whole globe. On the cubed sphere the polar faces are oriented differently from the equatorial ones, so one shared kernel does not transport patterns correctly across them.
• Iso-latitude rings. HEALPix has them, which is why it was designed for CMB spherical-harmonic analysis in the first place; the cubed sphere does not.
• Local refinement. FV3 supports stretched and nested grids — C-SHiELD, T-SHiELD and Tele-SHiELD all use nests down to 1.4 km. HEALPix has no equivalent; its refinement is globally uniform.

The practical consequence: X-SHiELD output must be conservatively remapped before a HEALPix model can consume it. That is exactly the class of operation TempestRemap exists for, and PhysicsNeMo already ships it in the DLWP pipeline (there it runs lat–lon to cubed sphere; the same tool handles cubed sphere to HEALPix).

One informative data point on what practitioners actually chose. GFDL's own collaboration with Ai2 Climate Modeling to emulate X-SHiELD did not go to HEALPix. ACE2 coarse-grains to a roughly 1° lat–lon grid and uses a Spherical Fourier Neural Operator; HiRO-ACE trains on a decade of X-SHiELD for emulation and downscaling on the same footing. The lesson generalises to the next subsection: if the operator is properly spherical, you no longer need a clever grid to fix the poles.

1.2  Operators, and which discretization each one requires

A grid is only a set of sample points. What actually determines geometric fidelity is how the operator is defined. The dividing line is whether it lives on tensor indices or on the continuous sphere, discretized by quadrature. The second kind gets two properties for free: resolution-agnosticism (change the grid without changing the architecture) and approximate SO(3) equivariance, because the Haar measure is rotation-invariant.

This pairing is what the rest of the package is built on — not every operator can run on every grid:

Note what this implies about the package layout. Spherical attention, SFNO, DISCO and the SHT are not in PhysicsNeMo at all. They live in torch-harmonics — attention/attention.py, disco/convolution.py, spectral_convolution.py and sht.py. PhysicsNeMo has no models/sfno directory; SFNO enters its model registry through an entry point once makani is installed, which is how the unified_recipe example trains it. The three attention modules that are in PhysicsNeMo are not spherical operators: healda/attention_layers.py is ragged local cross-attention with Triton grouped-query kernels, flare_attention.py is low-rank attention routing for Transolver, and physics_attention.py belongs to the same family.

Equivariance

Atmospheric dynamics has no preferred longitude origin. An operator satisfying K[R·f] = R·K[f] has that symmetry written into the architecture rather than learned from data, which is a direct source of sample efficiency and rollout stability.

The full comparison

Does the geometry actually buy you a later blow-up?

This is testable, and it has been tested independently. Lehmann et al. (2026) roll nine models out for two years — 2,920 autoregressive six-hourly steps from a 1 January 2021 initialisation — and define blow-up as the onset of unbounded exponential growth in the spatial extremes: the first 30-day window in which the global minimum or maximum grows log-linearly with R² > 0.9. Metrics and notebooks: lehmannfa/ai-weather-stability.

Days until blow-up; higher is better, and > 730 means it never blew up inside the test window. The AFNO model fails at 8 to 9 days. SFNO survives the whole two years on all nine variables the study tracks. That is at least a ninety-fold gap, and it is a lower bound.

But the causal story is not "spherical geometry, therefore stable." Look at the same table again: Pangu is a 3-D Swin transformer on a plain equirectangular grid with no spherical machinery whatsoever, and it is exactly as stable as SFNO. Sphericity is evidently sufficient but not necessary. Lehmann et al. locate the mechanism elsewhere — in the treatment of small spatio-temporal scales, where unstable models amplify high-frequency energy while stable ones behave as denoisers when noise is injected into their inputs. Their ablations on window size, shifting, patch size, normalisation, vertical levels, static fields and time embedding all leave stability unchanged. Earlier work in the same direction: McCabe et al. (2023) traced autoregressive error growth partly to the Double Fourier Sphere trick creating artificial discontinuities and unphysical cross-polar communication, and reported an 800% forecast-horizon improvement on ERA5 from controlling instability-inducing operations.

Two caveats on scope. The blow-up benchmark covers no spherical-attention model, so Attention on the Sphere (which validates on shallow water, spherical segmentation and depth estimation) has no equivalent year-scale evidence yet. And FourCastNet 3, which does report 60-day rollouts with no large-scale blow-up and invariant spectral slopes, gets there with spherical convolutions, local and global, rather than attention.

Which is best for the sphere

Section 2

Global medium-range emulators (autoregressive, mostly deterministic)

Six examples in examples/weather map an atmospheric state to its successor and iterate. They differ far less in "CNN versus transformer" than in which discretization from §1.1 they commit to.

fcn_afno — FourCastNet v1

A vision transformer whose self-attention is replaced by AFNO: token mixing happens in the Fourier domain with a block-diagonal MLP and soft thresholding, dropping the cost from O(N²) to O(N log N). Trained on a 20-channel ERA5 subset at 0.25°, step 6 h, autoregressive. A week-long forecast takes under two seconds. Paper: Pathak et al. 2022. Model code: models/afno/afno.py.

graphcast

A re-implementation of GraphCast. Encoder, processor, decoder over an icosahedral multi-mesh: the lat-lon grid is encoded onto a mesh built by repeated subdivision, a deep message-passing network propagates information across edges spanning several scales at once, and the decoder maps back to the grid. 73 ERA5 channels, step 6 h, with up to 12 steps of autoregressive fine-tuning. A synthetic-dataset switch (synthetic_dataset=true) lets you smoke-test without ERA5. Code: graph_cast_net.py, mesh construction in icosahedral_mesh.py, loss in graphcast_loss.py.

pangu_weather

A re-implementation of Pangu-Weather. A 3-D Earth-Specific Transformer treats pressure levels as a genuine third dimension for windowed attention and adds a position-dependent earth-specific bias per window. Inputs come in three blocks: 4 surface channels (t2m, 10u, 10v, msl), 5 upper-air variables on 13 pressure levels (65 channels), and 3 static fields (land-sea mask, soil type, topography). Requires NVIDIA Apex for the optimizer. Code: models/pangu.

dlwp — cubed sphere

A CNN mapping u(t) to u(t+6h) on a cubed sphere, 6 faces of 64×64, remapped from ERA5 with TempestRemap. U-Net with skip connections, 7 ERA5 variables, 18 input and 14 output channels in the shipped config. At 1.4° it produces 320-member six-week ensembles in minutes. Papers: Weyn et al. 2020 and Weyn et al. 2021.

dlwp_healpix — HEALPix, and a coupled ocean

The same idea on the HEALPix mesh, with a recurrent U-Net built from ConvNeXt blocks and ConvGRU cells (HEALPixRecUNet.py). The distinguishing feature is atmosphere-ocean coupling: a DLOM (Deep Learning Ocean Model) forecasts sea-surface temperature on a slower 48-hour step and is trained jointly with the 6-hour atmospheric model, configured through config_hpx32_coupled_dlwp and config_hpx32_coupled_dlom. Coupling logic in datapipes/healpix/couplers.py. Paper: Karlbauer et al. 2024; original reference implementation: CognitiveModeling/dlwp-hpx.

Data plumbing

dataset_download pulls ERA5 through the Copernicus CDS API into Zarr (a resumable checkpoint format) and then into HDF5 split across train/, test/, out_of_sample/ and stats/. Fixed at 0.25°; dt must divide 24; roughly 15 GB per variable per year. It does not handle invariant fields, so land-sea mask and orography must be sourced separately.

unified_recipe is a single training scaffold: pull from ARCO-ERA5, curate and regrid to Zarr, then swap configs to train AFNO, GraphCast or SFNO. SFNO requires makani, which registers the model through PhysicsNeMo's entry-point registry. This is the right starting point for architecture comparisons on identical data.

Section 3

Regional downscaling and generative high resolution (two-stage regression plus diffusion)

Both examples here share one design: a deterministic regression U-Net predicts the conditional mean, and a diffusion model learns only the residual. That split matters. A pure diffusion model must learn the large-scale mean and the small-scale texture at once, which trains poorly; absorbing the mean into an L2 regression leaves the diffusion model a zero-mean, much lower-variance target, cutting both sample count and instability. The price is two training runs, and the fact that regression bias propagates into the diffusion stage. Both use EDM preconditioning over SongUNet backbones (models/diffusion_unets).

corrdiff — km-scale downscaling

CorrDiff maps coarse ERA5 or GEFS onto a high-resolution regional target. Three dataset configurations, and the differences between them are easy to get wrong:

• Taiwan / CWA (cwb.yaml): 12 input channels to a 448×448 target at 4× downscaling. The four output channels are maximum radar reflectivity, 10u, 10v and t2m — confirmed by the normalisation constants in datasets/cwb.py. There is no precipitation channel in this configuration.
• GEFS-HRRR / CONUS (gefs_hrrr.yaml): this is where precipitation lives. Outputs are u10m, v10m, t2m, precip plus four precipitation-type probabilities (snow, ice, freezing rain, rain). Hence the lt_aware_ce_regression variant — cross-entropy for the categorical phase, and a lead-time embedding over 9 forecast steps because the input is a GEFS forecast rather than a reanalysis.
• HRRR-Mini: a deliberately small network and dataset (about 10 A100-hours) for learning the codebase; outputs only 10u and 10v. Available on NGC.

A patched_diffusion variant tiles a large domain and runs multi-diffusion with 100 learnable positional channels, which is what makes training on a 1057×1796 domain tractable.

stormcast — convection-allowing emulation

StormCast emulates NOAA's 3 km HRRR autoregressively. The regression model predicts the next HRRR step from the previous HRRR state plus an ERA5 background; the diffusion model corrects it. The crucial contrast with CorrDiff: StormCast advances in time (1-hour autoregressive rollout), CorrDiff does single-frame spatial downscaling. Domain 512×640, roughly 100 variables (u/v/t/q on many model levels plus surface fields and reflectivity). The same recipe also trains Stormscope (arXiv:2601.17268), which drops the regression stage for a diffusion-only DiT and nowcasts clouds and precipitation from GOES radiances and MRMS radar. Config: config/model/stormscope.yaml.

Section 4

Diagnostics, post-processing and ensembling

diagnostic

A diagnostic model does not advance time; it infers an extra variable from the simultaneous atmospheric state. The shipped example is an AFNO with 27 input channels — the 20 ERA5 variables from Pathak et al. 2022, plus cosine zenith angle, geopotential, land-sea mask and four lat/lon encoding channels — producing one output, tp06 (6-hour accumulated precipitation). This is FourCastNet's precipitation module as a standalone, and the template for bolting a diagnostic head onto any backbone. Config: diagnostic_precip.yaml.

temporal_interpolation

AI forecast models step natively at 6 hours; this one interpolates to 1 hour. ModAFNO takes both endpoint states plus auxiliary fields such as orography, and an embedding network converts the conditioning variable into shift and scale parameters for each AFNO block — FiLM-style modulation. It is deterministic, trained with a latitude-weighted L2 loss, and approximates the conditional expectation of the intermediate state given both endpoints and the offset. The released model used 73 variables, about 100 TB of data and 64 H100s. A pretrained wrapper exists in Earth2Studio. Code: models/afno/modafno.py.

mixture_of_experts — MoWE

Mixture of Weather Experts. Forecast fields from Aurora, Pangu-Weather and FourCastNet 3 are stacked as a multi-channel image; a Diffusion Transformer acts as the gating network, conditioned on lead time (sampled between 6 hours and 2 days during training), and emits per-expert weights and a bias so that the blended forecast is a weighted sum of the experts. A probabilistic variant additionally consumes a noise vector and trains with a CRPS loss. All expert forecasts must be generated offline first with Earth2Studio — multiple TB. DiT code: models/dit.

Section 5

Data assimilation

regen — score-based DA

The trick is that no conditional model is ever trained. An unconditional diffusion model is fit to HRRR over a 128×128 Oklahoma window with channels 10u, 10v, tp. At inference, score-based data assimilation modifies only the sampling loop: at each denoising step the current estimate is passed through an observation operator, compared against real station data, and the resulting posterior score — the prior score plus the likelihood score — steers the next update. Observations come from NOAA's Integrated Surface Database. Assimilating 40 stations lowers RMSE at held-out stations by about 10% relative to HRRR itself, the first demonstration of this at kilometre scale. Paper: Manshausen et al. 2024.

healda — stateless DA on HEALPix

Marked under active construction in the repository. The goal is stateless assimilation: no background-forecast cycle, just conventional and satellite observations (ATMS, MHS, conventional) in, one global ERA5-compatible analysis out, on a HEALPix level-6 padded XY grid. The example directory holds only a README, normalisation statistics and requirements; the model and datapipe live in experimental/models/healda and experimental/datapipes/healda. Architecturally it is a video DiT with temporal self-attention plus ragged pixel-to-observation cross-attention, backed by Triton grouped-query-attention kernels (pixel_attention.py) and adaLN conditioning. Most of the README is a guide to plugging in your own observation loader and transform.

Section 6

Hydrology

flood_modeling/hydrographnet

A physics-informed graph neural network surrogate for the shallow-water equations. Encoder, processor, decoder message passing on an unstructured mesh autoregressively predicts increments of water depth and volume with residual connections. Three design choices carry the result: a mass-conservation loss built from volume-continuity inequalities (no automatic differentiation required), the pushforward trick to suppress autoregressive error accumulation, and Kolmogorov-Arnold Networks replacing MLPs for interpretability. Node features mix static attributes (elevation, slope, roughness) with dynamic history, plus global forcings from the inflow hydrograph and precipitation. Validated on the White River near Muncie, Indiana: 4,787 nodes, ground truth from HEC-RAS, dataset auto-downloaded from Zenodo. Paper: Taghizadeh et al. 2025.

Quick reference

One line each

Sources: NVIDIA/physicsnemo and NVIDIA/torch-harmonics main branch source and docstrings, read 2 August 2026; GFDL SHiELD configuration tables for the X-SHiELD cross-check; rollout blow-up times from Lehmann et al. 2026 (Table 1) and its companion repository lehmannfa/ai-weather-stability.