Add AeroJEPA model + SuperWing tutorial recipe (experimental) by fgiral000 · Pull Request #1690 · NVIDIA/physicsnemo

fgiral000 · 2026-06-01T18:07:19Z

PhysicsNeMo Pull Request

Description

Adds the AeroJEPA model and a SuperWing tutorial recipe under
physicsnemo.experimental and examples/cfd/external_aerodynamics/.
AeroJEPA is a Joint-Embedding Predictive Architecture for 3D
aerodynamic surrogate modeling: instead of mapping geometry directly to
a flow field, it predicts a latent representation of the flow from a
latent representation of the geometry and operating conditions, and
reconstructs the field through a continuous implicit decoder when
needed (Giral et al., arXiv:2605.05586).

What this PR delivers:

Model at physicsnemo.experimental.models.aerojepa.
AeroJEPA composes a context encoder, a target encoder, a query-token
field decoder (collectively AeroJEPATrunk), and a JEPA predictor
head (PrototypeTokenJEPAHead) into a single
physicsnemo.core.module.Module. The training path takes context
positions/features, independent target encoder surface/volume inputs,
and operating conditions; the predictor predicts target tokens, and
the decoder evaluates the field at user-supplied query points.
predict is a no-grad inference wrapper; decode_field_chunked
supports memory-bounded evaluation over very large query sets.
Concrete encoders (ContextTransformer, TargetTransformer,
PointTransformer), the QueryTokenDecoder, and the encoder ABCs
are all exposed as composable components.
Building blocks at
physicsnemo.experimental.models.aerojepa.layers. TokenSet and
EncoderOutput token dataclasses, a deterministic
FourierPositionalEncoding, ResidualMLP, the
LocalPointTransformerBlock / LocalTokenCrossAttentionBlock
attention blocks (with optional AdaLN / AdaLN-Zero conditioning), the
PointCloudTokenizer (seven center-selection strategies with k-NN
cluster pooling), token batching / mask / k-NN helpers, and prototype
anchor build / load utilities. TokenSet and EncoderOutput are
re-exported from the model package for convenience.
Losses at physicsnemo.experimental.models.aerojepa.losses.
SIGReg and TokenLatentSIGReg (a sketch isotropic-Gaussian
regularizer for latent-token distributions, with a padding-aware
wrapper), the flatten_valid_token_features /
reshape_token_features_for_sigreg masking helpers, and the
reconstruction loss family (MSELoss / RelativeL2Loss /
RelativeMSELoss / RelativeL2MSELoss, each with functional and
nn.Module forms, optional per-channel weights stored as a
persistent buffer, optional per-point weights, and an optional
validity mask).
Tutorial recipe at
examples/cfd/external_aerodynamics/aerojepa. End-to-end Hydra-driven
workflow on the public SuperWing dataset (Yang et al.,
arXiv:2512.14397): dataset download via the Hugging Face Hub
(yunplus/SuperWing), automatic split-by-geometry manifest and
per-channel normalization stats, JEPA training (reconstruction +
latent + SIGReg with linear warmups; AdamW +
warmup-cosine; optional EMA), checkpointed inference with chunked
decoding, three-panel GT | Pred | |Error| field plots for the three
surface channels (Cp, Cf_tau, Cf_z), per-channel relative-L2 /
RMSE / MAE metrics on the test split, and a pressure-only CL/CD
post-processor that integrates the surface field and emits a per-case
CSV plus a parity scatter.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.
The CHANGELOG.md is up to date with these changes.
An issue is linked to this pull request.
If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Tests

193 unit tests under test/experimental/models/aerojepa/
(constructor + attribute checks, non-regression shape checks on the
encoders, decoder, predictor, trunk, top-level model, layers, and
losses). pytest test/experimental/models/aerojepa/ -q passes
locally on CPU (~20 s).
Full SuperWing end-to-end smoke-tested on a single GPU:
train.py -> inference.py -> superwing_metrics -> superwing_forces.
Training losses decrease monotonically; inference produces field
plots, per-case field-error metrics, and a force-coefficient parity
scatter.

Dependencies

No new core dependencies. The example recipe adds optional
example-side dependencies in
examples/cfd/external_aerodynamics/aerojepa/requirements.txt
(Hugging Face Hub for the dataset download, plotting and
post-processing utilities). Pre-commit hooks, ruff, interrogate,
markdownlint, and the SPDX license check pass on every file in the
PR.

copy-pr-bot · 2026-06-01T18:07:23Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-06-01T18:15:45Z

Greptile Summary

This PR introduces the AeroJEPA model — a Joint-Embedding Predictive Architecture for 3D aerodynamic surrogate modeling — along with all its building blocks (encoders, decoder, predictor, layers, losses) under physicsnemo.experimental, plus a full Hydra-driven SuperWing tutorial recipe and 193 unit tests.

Model core (aerojepa.py, trunk.py, predictor.py, decoder.py): The context/target encoder–predictor–decoder pipeline is well-structured; batched and single-sample forward paths handle edge cases correctly.
Training recipe (train.py, runtime.py): The validation path in _run_epoch builds full autograd graphs without a torch.no_grad() guard, and get_autocast_context exposes fp16 without a paired GradScaler.
masked_mean: Returns (B, F) for rank-3 no-mask input but documents (B, 1, F).

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/aerojepa/aerojepa.py	Top-level AeroJEPA Module composing trunk + predictor; forward/predict/decode_field_chunked paths look correct; build_target_token_coords uses a private _tokenize_single method (noqa-suppressed).
physicsnemo/experimental/models/aerojepa/trunk.py	AeroJEPATrunk wiring encoder/decoder; encode_context, decode_queries, forward_single/forward_batch all look correct.
physicsnemo/experimental/models/aerojepa/decoder.py	QueryTokenDecoder with chunked cross-attention, SIREN options, wall-velocity gate, and batched forward; logic appears sound.
physicsnemo/experimental/models/aerojepa/predictor.py	PrototypeTokenJEPAHead with interleaved self/cross attention; batch handling and conditioning logic look correct.
physicsnemo/experimental/models/aerojepa/layers/token_utils.py	Batch flattening, k-NN, and TokenSet utilities; masked_mean has a docstring/implementation shape inconsistency for rank-3 no-mask input (returns (B,F) not (B,1,F) as documented).
examples/cfd/external_aerodynamics/aerojepa/train.py	Hydra training entry point; validation forward pass in _run_epoch builds unnecessary autograd graphs because there is no torch.no_grad() guard when is_train=False, wasting GPU memory.
examples/cfd/external_aerodynamics/aerojepa/src/training/runtime.py	get_autocast_context enables fp16 autocast without a paired GradScaler; safe with the default bf16 config but could silently corrupt training if users set precision: fp16.

Comments Outside Diff (1)

physicsnemo/experimental/models/aerojepa/layers/token_utils.py, line 1350-1356 (link)

masked_mean return shape mismatch between mask=None and mask≠None paths for rank-3 input

The docstring states the function returns (B, 1, F) for rank-3 input, but the mask is None branch uses features.mean(dim=1) (no keepdim) and actually returns (B, F). The masked branch correctly uses keepdim=True and returns (B, 1, F). This inconsistency could cause silent shape mismatches if a caller passes rank-3 features without a mask and expects the documented (B, 1, F) layout.

_{Reviews (1): Last reviewed commit: "changelog: move SuperWing recipe bullet ..." | Re-trigger Greptile}

peterdsharpe · 2026-06-01T18:18:52Z

Hi @fgiral000, thanks for the PR! To keep PR size reviewable, would it be possible to:

a) split this PR up into two separate PRs, one of which adds the model ("PR 1"), and a later follow-on that adds the example ("PR 2").

b) In PR 1, please re-use shared PhysicsNeMo tooling where possible. (E.g., _gpu_knn.py should re-use existing KNN implementations in physicsnemo.nn.functional; conditioning MLPs should use FullyConnected, many losses duplicate existing code)

c) All functions should use jaxtyping annotations for tensor shapes. Please use Literal types for enumerations rather than str, etc.

d) In PR 2, if possible, please add AeroJEPA as an example within ./examples/external_aerodynamics/unified_external_aero_recipe/, rather than as a standalone aerojepa folder.

peterdsharpe · 2026-06-01T18:20:33Z

Actually, it might be worth splitting out a third PR as well for addition of the SuperWing dataset utils.

The training loop stacked every sample's loss and ran a single backward on the batch mean, keeping all sample autograd graphs alive until then, so peak memory scaled with batch size. Backward each sample's loss divided by the batch size inside the loop instead: mathematically identical to the batch-mean backward but only one sample graph is alive at a time. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

review) Validation runs on the EMA shadow weights, but the checkpoint saved the live training weights (ema.restore had already run), so the selected best checkpoint did not reproduce the reported val metric. Persist ema.shadow as the model state when EMA is enabled; it shadows the full state_dict, so it is a complete model state. The live optimizer/EMA state is still saved for resuming. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

…DIA#1690 review) Move the per-call TokenSet import to the module top, and document that the flat-to-grid reshape assumes the eval_full_grid_query query_pos is emitted in row-major (H, W) order (which the SuperWing dataset guarantees). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

Fix the epochs comment to match the actual 200-epoch schedule, note that relative_mse_mode is ignored by the active relative_l2_mse blend, and drop the unused mask loss block. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

The normalization stats saved per-channel target_min/target_max, but nothing consumes them (the dataset z-scores targets with mean/std). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

…A#1690 review) The decoder out_dim comment named the channels Cp, Cf_x, Cf_y while the target encoder and dataset use Cp, Cf_tau, Cf_z. Same three quantities; align the decoder comment to the canonical names. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

Note that the recipe currently trains on a single GPU and that multi-GPU / multi-node distributed training is planned as a follow-up. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

compute_latent_loss detached its target, a leftover from older research code that slipped in while porting the recipe to the PhysicsNeMo implementation. AeroJEPA's anti-collapse mechanism is the SIGReg regularizer, not EMA and stop-gradient, so the latent alignment term is meant to flow gradients into both the predictor and the target encoder. With the stray detach the target encoder got no signal from the alignment loss (reconstruction flows through the predictor output), leaving it shaped only by SIGReg. Drop the detach so the predictor and target encoder are trained jointly to agree in latent space, with SIGReg preventing the trivial constant-collapse solution. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

Restore the original two-phase schedule as an opt-in config: phase 1 trains the context + target encoders and the predictor in latent space (latent + SIGReg, decoder frozen, reconstruction off); phase 2 freezes the encoders and predictor and trains only the decoder to reconstruct the field from the frozen latents. Freezing is done via requires_grad on the single optimizer, and the decoder forward is skipped when the reconstruction weight is zero. Also add an optional SIGReg regularizer on the context latents (a separate module from the target-latent SIGReg), defaulting to weight 0 so it is off unless explicitly enabled. Disabled by default the training is a single phase with all loss terms active and every parameter trainable, matching the prior behaviour. Tests cover the target-encoder gradient flow, phase-1/phase-2 freezing, phase weight gating, phase resolution, and the context-SIGReg wiring. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

train.py had no way to load a checkpoint at startup. Add two optional, mutually exclusive modes: - training.resume: restore model + optimizer + LR scheduler + EMA + epoch and best-val counters and continue the schedule from where a run stopped. - training.init_from_checkpoint: start a new run from pretrained weights, loading the model weights only (fresh optimizer, epoch 0). With strict=false a checkpoint holding only a subset of modules loads those and leaves the rest at initialization; combined with two-phase training this trains a decoder on top of frozen pretrained encoders + predictor. Checkpoints now also record best_val_loss so a resumed run keeps its model selection. Both modes are off by default. Tests cover the weights-only init, the full resume, the fresh-start default, and non-strict subset loading. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

For rank-3 input with no mask, masked_mean reduced with mean(dim=1) and no keepdim, returning (B, F) while the docstring and the masked path return (B, 1, F). Add keepdim=True so both paths agree with the documented shape. No current caller relied on the (B, F) form. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

…1690 review) get_autocast_context can enable fp16 autocast, but the training loop stepped the optimizer with no GradScaler, so fp16 gradients could overflow to inf/nan and silently corrupt the weights. Add a GradScaler enabled only for fp16 on CUDA; for bf16 / fp32 / CPU it is disabled and every scaler call is a transparent no-op, leaving the default bf16 path unchanged. Gradients are unscaled before clipping so the clip norm is computed on true gradients. A new test runs one _run_epoch train step end-to-end through the scaler path. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>

loliverhennigh

LGTM

mnabian

LGTM now. Thanks for addressing all the comments!

mnabian · 2026-07-15T21:41:46Z

/ok to test b0566e8

mnabian · 2026-07-15T22:56:39Z

/ok to test afca308

fgiral000 requested review from ktangsali and megnvidia as code owners June 1, 2026 18:07

greptile-apps Bot reviewed Jun 1, 2026

View reviewed changes

Comment thread examples/cfd/external_aerodynamics/aerojepa/src/training/runtime.py

Comment thread examples/cfd/external_aerodynamics/aerojepa/train.py

mnabian self-requested a review June 1, 2026 18:40

megnvidia approved these changes Jun 1, 2026

View reviewed changes

Comment thread CHANGELOG.md Outdated

Comment thread CHANGELOG.md Outdated

Comment thread CHANGELOG.md Outdated

Comment thread CHANGELOG.md Outdated

mnabian reviewed Jun 4, 2026

View reviewed changes

Comment thread physicsnemo/experimental/models/aerojepa/layers/_gpu_knn.py Outdated

Comment thread physicsnemo/experimental/models/aerojepa/layers/token_utils.py Outdated

Comment thread physicsnemo/experimental/models/aerojepa/layers/positional_encoding.py Outdated

mnabian reviewed Jun 4, 2026

View reviewed changes

Comment thread physicsnemo/experimental/nn/attention_blocks.py Outdated