Add AeroJEPA model + SuperWing tutorial recipe (experimental)

[fgiral000]

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.

[mnabian]

Please use physicsnemo's Mlp.

[mnabian]

Please use physicsnemo's Mlp.

[mnabian]

Please use physicsnemo's Mlp.

[mnabian]

Please use physicsnemo's Mlp: Mlp(dim, dim, num_heads, act_layer=nn.GELU)

[mnabian]

Please use physicsnemo's Mlp: Mlp(3, dim, dim, act_layer=nn.GELU)

[mnabian]

Please use physicsnemo's Mlp: Mlp(dim, dim, num_heads, act_layer=nn.GELU)

[mnabian]

Please use physicsnemo's Mlp: Mlp(3, dim, dim, act_layer=nn.GELU)

[mnabian]

Drop SineLayer, use physicsnemo.nn.SirenLayer; keep SirenHead as a thin local composition but build it from the library's SirenLayer instead of the local one.

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
LayerNorm(int(token_dim))

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.norm = LayerNorm(int(dim))

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.norm = LayerNorm(self.dim)

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.norm_q = LayerNorm(self.dim)

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.norm_kv = LayerNorm(self.dim)

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.out_norm = LayerNorm(self.hidden_dim)

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.out_norm = LayerNorm(self.point_feature_dim)

[mnabian]

Use PhysicsNeMo's TE-aware LayerNorm instead of torch.nn.LayerNorm to pick up Transformer Engine's high-performance LayerNorm (notably the faster backward pass).

from physicsnemo.nn.module.layer_norm import LayerNorm
...
self.out_norm = LayerNorm(int(token_dim))

[mnabian]

cluster_size=None for the neighborhood-cluster strategies silently defaults to max_tokens (e.g. 512 neighbors/token) — almost certainly unintended, and inconsistent with the data_prototype_cluster path which defaults to 16. Recommend failing fast: require cluster_size for random_cluster/fps_cluster/voxel_fps_cluster, and drop the max_tokens fallback.

suggestion: In __init__, after the strategy/prototype validation:

if (
    self.strategy in {"random_cluster", "fps_cluster", "voxel_fps_cluster"}
    and self.cluster_size is None
):
    raise ValueError(
        f"cluster_size must be provided for tokenizer_strategy='{self.strategy}'."
    )

Then in tokenize_with_clusters, replace the fallback (lines 423–425):

cluster_size = (
    max_tokens if self.cluster_size is None else int(self.cluster_size)
)

with:

cluster_size = int(self.cluster_size)

[mnabian]

data_prototype_cluster with zero-point inputs calls KNN on empty tensor before early-return guard.

[mnabian]

encode_geometry has a non-trivial branch when include_geometry_global_in_decoder_cond=True (lines 245-252): it pulls the context global token and concatenates it onto cond_global, widening the decoder conditioning. Every test builds the trunk with this flag False, so the concat branch and the resulting wider cond_global are never exercised. Add a model built with the flag True and assert cond_global length grows by the context-global token dim.

Update: Based on the newer previous comment, you can ally this comment directly to trunk._build_cond_global_single.

[mnabian]

AeroJEPA.forward raises ValueError when context_pos/context_feat are not rank 2 (line 474) or when their point counts disagree (line 479), but no test triggers either guard. Add two tests asserting pytest.raises(ValueError, match="rank 2") for a rank-3 context_pos and match="agree on the point count" for mismatched N between pos and feat.

[mnabian]

Hot path: QueryTokenDecoder._decode_chunk, run for every query chunk over the full surface/volume field (millions of query points at inference via decode_field_chunked). Each LocalTokenCrossAttentionBlock.forward independently calls chunked_knn_indices(query_coords=query_pos, key_coords=token_coords, ...) (decoder.py:458, attention_blocks.py), but within a chunk query_pos and token_coords are identical across all cross_attention_layers blocks and the cross-blocks use no dilation, so the same k-NN graph is recomputed cross_attention_layers times per chunk. Compute the neighbor index once per chunk and pass it into the blocks (e.g. an optional precomputed-idx argument), reusing it across layers. Same redundancy pattern applies to PointTransformer.encode_single (point.py:606) when the dilation schedule is constant.

[mnabian]

The point_weights argument is documented and implemented in all four losses (mse_loss 106-111/122-124, relative_l2_loss 230-235, relative_mse_loss 414-432, plus the hybrid), but no test ever passes point_weights. None of the per-point weighting / weighted-denominator branches are exercised. Add at least one test asserting that uniform point_weights=ones matches the unweighted result and that a zeroed weight on a row drops that row's contribution.

[mnabian]

test_mse_mask_drops_invalid_rows (test/experimental/models/aerojepa/losses/test_reconstruction.py:77-84) only ever passes an all-True mask (torch.ones(8, dtype=torch.bool)) and asserts the result equals the no-mask path. With an all-True mask, err * mask_t is a no-op and denom_weights.sum() equals the full row count, so the masked branch (reconstruction.py:113-120, including the denom_weights.sum().clamp_min(1.0) denominator) is only covered vacuously. The test's own docstring claims "masked-out rows contribute neither to the numerator nor the denominator," yet no row is ever masked out — a regression that stopped excluding masked rows would still pass. Add a test with a partial mask: set some rows to mask=False, fill those target rows with large values, and assert the loss equals mse_loss computed on only the kept rows.

[mnabian]

payload = torch.load(ckpt_path, map_location=device, weights_only=False) deserializes the checkpoint with full pickle support. Loading a .pt from an untrusted source would execute arbitrary code via a malicious pickle __reduce__. Impact is limited here: the loaded payload is consumed only as tensor state_dicts (payload["model"], payload["ema_shadow"]) plus a scalar ema_decay, so weights_only=True would work and is the safer default.

[mnabian]

decode_field_chunked is only tested with precision="fp32" (test_aerojepa.py:179), so the autocast-enabled path (lines 393-401, where enabled=True for fp16/bf16) is never run. On a CUDA device this is the path that actually matters for the memory-bounded use case. Add a parametrized test over precision in {"fp16","bf16"} asserting the returned tensor is finite, CPU, and shape (Nq, C).

[mnabian]

build_target_token_coords raises ValueError when the target encoder lacks the _tokenize_single path (lines 449-452), but only the happy path is tested (test_aerojepa.py:137). The defensive branch that protects callers swapping in a non-transformer target encoder is untested. A small test passing a stub trunk whose target_encoder has no encoder._tokenize_single and asserting the ValueError would cover it.

[mnabian]

_assign_points is a k=1 nearest-center query — replace the manual cdist+min+chunk loop with nn.functional.knn(centers, points, k=1), which chunks internally and uses the cuML/scipy backends. The returned distance is unused by _run_chunked_kmeans (assign, _ = …), so the dist and chunk_size args can be dropped too.

from physicsnemo.nn.functional import knn

def _assign_points(points: torch.Tensor, centers: torch.Tensor) -> torch.Tensor:
    """Assign each point to its nearest center (k=1 kNN)."""
    idx, _ = knn(centers, points, k=1)   # neighbors of `points` among `centers`
    return idx[:, 0]

and the call site in _run_chunked_kmeans:

assign = _assign_points(points, centers)

(chunk_size is no longer threaded into _assign_points.)

[mnabian]

The AeroJEPA CHANGELOG is far longer and more granular than any other entry. Consider condensing to a single concise bullet (optionally split model vs. recipe) matching the house style.

[mnabian]

The kNN distance search is recomputed in every block although coordinates are static within each stack. Compute the sorted neighbor list once (at k = neighbor_k × max(dilation)), then have each block apply its own dilation stride - a pure indexing op. This removes N−1 redundant distance/sort passes per stack without changing behavior.

[mnabian]

same static-coords kNN recomputation here

[mnabian]

same static-coords kNN recomputation here

[mnabian]

same static-coords kNN recomputation here

[mnabian]

same static-coords kNN recomputation here

[mnabian]

AeroJEPA subclasses physicsnemo.core.Module, so its __init__ args must be JSON-serializable. The only exception is args that are themselves physicsnemo.Module instances. But trunk and predictor here are plain torch.nn.Module:

model = AeroJEPA(trunk=trunk, predictor=predictor)
model.save("m.mdlus")  # raises during json.dumps

Suggested fix: You can take JSON-serializable config in __init__ and build trunk/predictor internally — the FNO/DoMINO pattern. Alternatively, convert the full submodule tree to physicsnemo.Module.

[mnabian]

The blocks are actually always constructed with conditioning_dim=self.hidden_dim and run AdaLN on the zero cond_embed — so they apply a learned constant modulation, not nothing. Either gate conditioning_dim on cond_dim>0 to truly disable, or fix the docstring.

[mnabian]

This is a reimplementation of the trunk._build_cond_global_single. Call the trunk helper instead.

[mnabian]

query_sdf=None not validated against decoder.use_sdf. Add an early guard.