Relative L2 loss - #1746
Merged
Merged
Conversation
mnabian
marked this pull request as ready for review
June 23, 2026 18:08
Contributor
Greptile SummaryThis PR adds
Important Files Changed
Reviews (1): Last reviewed commit: "add relative losses" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: Mohammad Amin Nabian <m.a.nabiyan@gmail.com>
Collaborator
Author
|
/ok to test 741b1aa |
6 tasks
fgiral000
added a commit
to fgiral000/physicsnemo
that referenced
this pull request
Jul 3, 2026
…IA#1690 review) The reconstruction loss wrapper now lives in the recipe at src/losses/reconstruction.py and builds on the core physicsnemo.metrics losses. Import the channel-weighted / hybrid loss modules from that local wrapper in src/losses/builders.py, and stop re-exporting the reconstruction family from the experimental losses package (which no longer ships it). Relocate the reconstruction-loss tests next to the recipe code, with a recipe-level conftest that puts the recipe root on sys.path and provides the device fixture. Note: the wrapper depends on physicsnemo.metrics.general.relative_error and the weights-enabled mse from PR NVIDIA#1746, which is not merged yet; the recipe imports will resolve once NVIDIA#1746 lands. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es>
4 tasks
This was referenced Jul 9, 2026
coreyjadams
approved these changes
Jul 15, 2026
coreyjadams
left a comment
coreyjadams
left a comment
Collaborator
There was a problem hiding this comment.
Approving to unblock; I left a few small comments.
|
|
||
|
|
||
| def mse(pred: Tensor, target: Tensor, dim: int = None) -> Union[Tensor, float]: | ||
| def mse( |
Collaborator
There was a problem hiding this comment.
PyTorch has https://docs.pytorch.org/docs/2.13/generated/torch.nn.functional.mse_loss.html, which is quite similar. Can we document here in the docstring any relevant differences users should be aware of?
Collaborator
Author
There was a problem hiding this comment.
Added some explanation in the docstrings.
Comment on lines
+100
to
+101
| eps: float = 1e-8, | ||
| weights: Tensor | None = None, |
Collaborator
There was a problem hiding this comment.
It's a minor nit, but can we put these in the same order in relative_error and mse? [pred, target, dim, weights, eps]` makes sense but I care more about consistency than detailed order, honestly.
| def relative_mse( | ||
| pred: Tensor, | ||
| target: Tensor, | ||
| *, |
Collaborator
There was a problem hiding this comment.
Why * here and not in mse?
Collaborator
Author
There was a problem hiding this comment.
Removed it from here.
Collaborator
Author
|
/ok to test 7d09c58 |
mnabian
enabled auto-merge
July 15, 2026 18:22
coreyjadams
disabled auto-merge
July 15, 2026 19:03
mnabian
added a commit
to mnabian/physicsnemo
that referenced
this pull request
Jul 15, 2026
…#1690) * Add experimental/nn/aerojepa subpackage skeleton Create an empty subpackage for the AeroJEPA reusable building blocks (attention blocks, geometry tokenizer, context/target encoders, decoder, predictor) that land in subsequent commits. Establishes the SPDX license header, module docstring, and ``__all__`` placeholder so that follow-up commits only need to register new public symbols. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add TokenSet and EncoderOutput dataclasses for AeroJEPA TokenSet bundles token features with their geometric coordinates and optional mask, global token, and auxiliary side data; EncoderOutput is a thin wrapper used by context and target encoders to surface a global summary alongside the per-token output. Includes raw-string docstrings with Parameters/Examples sections (three executable doctests), modern union syntax, and the SPDX header. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add FourierPositionalEncoding layer for AeroJEPA A deterministic log-frequency sinusoidal positional encoding used to lift continuous query coordinates into a high-dimensional feature space before the implicit decoder consumes them. Distinct from physicsnemo.nn.FourierEmbedding (random Gaussian frequencies on scalar timesteps); this variant uses fixed log-powers of pi on multi-dim coordinates with the standard sin/cos band layout. Includes an out_dim property, jaxtyping on forward, and an executable doctest. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add private GPU-native kNN helpers under experimental/nn/aerojepa A private _gpu_knn module bundling chunked torch.cdist plus topk for building homogeneous (gpu_knn_self) and bipartite (gpu_knn_bipartite) k-NN graphs and inverse-distance interpolation (gpu_knn_interpolate). Pure PyTorch, no warp or custom CUDA — works on CPU too, just slower. The leading underscore on the filename makes the module package- private; callers live inside the aerojepa subpackage only. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add token batching, masking and kNN helpers for AeroJEPA A token_utils module with the helpers used by the AeroJEPA tokenizer, encoders and attention blocks: gather_rows, counts_to_mask, flatten_padded_batch / unflatten_to_padded, compute_batch_offset_step, flatten_batched_coords, chunked_knn_indices (CPU/GPU dispatcher with the AE_KNN_BACKEND env override), masked_mean, trim_batched_tokens, and pad_token_sets. Behavior preserved; types modernized and the TokenSet import is package-relative. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add ResidualMLP and local attention blocks for AeroJEPA A reusable trio of attention building blocks: ResidualMLP (pre-norm residual MLP with optional AdaLN / AdaLN-Zero conditioning), LocalPointTransformerBlock (local self-attention over a per-point k-NN graph with learned relative-position bias), and LocalTokenCrossAttentionBlock (cross-attention from queries to a per-query k-NN of context tokens, with a 5-way conditioning MLP that modulates query and key/value sides independently). Behavior preserved: zero-init conditioning MLPs give an identity transform at construction time, and the N<=1 / empty-input fallbacks short-circuit the same way they do upstream. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add PointCloudTokenizer for AeroJEPA A tokenizer module that reduces a raw point set to a bounded token budget before attention. Seven strategies: identity, random, FPS, random/FPS/voxel-FPS cluster pooling, and prototype-anchored clustering. The cluster strategies return the kNN indices that link each token center back to the source points, allowing a downstream encoder to replace the default feature mean with a learned pooling (e.g. the message-passing PointClusterGraphPool that lands with the encoders in PR NVIDIA#3). Behavior preserved including the non-persistent prototype_coords buffer and the per-sample loop used by the prototype strategy. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add prototype anchor build/load helpers for AeroJEPA Build the fixed k-means anchor set used by the data_prototype_cluster tokenizer strategy. The build pass walks a training dataset, tokenizes each sample to obtain candidate token coordinates, optionally subsamples, runs chunked Lloyd-iteration k-means with empty-cluster FPS refill, sorts the centers lexicographically, and serializes them with a JSON metadata blob. Two load functions (target / context - identical file layout) and two ensure_* helpers (load-if-exists else build) round out the public surface. Behavior preserved; the seed argument governs k-means initialization and candidate subsampling but not the tokenizer pass, which intentionally uses random sampling. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Populate experimental/nn/aerojepa __init__ with public re-exports Export the 23 public symbols from the six source modules at the package level: TokenSet/EncoderOutput dataclasses, FourierPositionalEncoding, ResidualMLP and the two local attention blocks, PointCloudTokenizer, the ten batching/mask/kNN helpers, and the six prototype anchor build/load functions. Module docstring tightened to reflect the actual contents (encoders / decoder / predictor land in physicsnemo.experimental.models.aerojepa in a later PR). The package-private _gpu_knn helpers remain accessible via their submodule path but are intentionally not re-exported. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Promote AeroJEPA layer classes to physicsnemo.experimental.nn Re-export the five AeroJEPA nn.Module layer classes (FourierPositionalEncoding, ResidualMLP, LocalPointTransformerBlock, LocalTokenCrossAttentionBlock, PointCloudTokenizer) at the experimental.nn parent namespace, alongside the existing FLARE and DiffusionUNet3D family. Data types (TokenSet, EncoderOutput), batching/mask helpers, and prototype-anchor builders stay scoped to the aerojepa subpackage to keep the parent namespace focused on actual layers. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add tests for AeroJEPA core dataclasses Tests for TokenSet and EncoderOutput covering construction (both batched and unbatched), the is_batched / token_dim properties, the with_updates immutability + selective-replacement contract, and the independence of the default aux dict across instances. Uses the shared device fixture so the CUDA path runs when available. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add constructor/attribute tests for the remaining AeroJEPA building blocks Six new test files covering positional_encoding, attention_blocks, point_tokenizer, token_utils, _gpu_knn, and prototype_anchors. 85 new tests covering: constructor validation paths, forward output shapes, edge cases (N<=1 LPT fallback, empty cross-attention, empty/single- point kNN, missing voxel_size, non-persistent prototype_coords buffer), identity-at-init of AdaLN-Zero conditioning MLPs, the AE_KNN_BACKEND env override, and build/load round-trips with a tiny fake dataset. All tests use the shared device fixture so CUDA runs when available; CPU run is 18 s wall. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add CHANGELOG entry for AeroJEPA building blocks Documents the new physicsnemo.experimental.nn.aerojepa subpackage contributed across the preceding 12 commits on this branch: token dataclasses, Fourier positional encoding, ResidualMLP, the two local attention blocks, PointCloudTokenizer, token batching/mask/kNN helpers, and prototype anchor utilities, plus the parent-namespace re-export of the five layer classes. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add experimental/metrics/jepa subpackage skeleton Create an empty subpackage for JEPA-style losses and regularizers (SIGReg, TokenLatentSIGReg, the padding-aware masking helpers, and the reconstruction loss family) that land in subsequent commits. Establishes the SPDX license header, module docstring, and ``__all__`` placeholder so that follow-up commits only need to register new public symbols. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add padding-aware masking helpers for the JEPA losses Two utilities used by SIGReg / TokenLatentSIGReg to flatten padded batched token features and reshape them into the (T, B, D) layout SIGReg expects. flatten_valid_token_features is a passthrough on rank-2 inputs and uses boolean masking on rank-3 inputs; reshape_token_features_for_sigreg adds the leading T=1 axis and emits a zero-element (1, 0, D) placeholder when the mask removes every row. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add SIGReg and TokenLatentSIGReg regularizers for AeroJEPA SIGReg pushes a learned latent toward N(0, I) by comparing the empirical Fourier characteristic function of random projections against the reference Gaussian one on a uniform knot grid (the LeWorldModel construction). Three non-learnable buffers cache the knot positions, the reference window, and the trapezoidal + window-weighted integration weights. TokenLatentSIGReg is a thin wrapper that accepts (B, N, D) or (N, D) features plus an optional mask, drops padded rows via the masking helpers, and short-circuits to a zero scalar when the mask removes every row. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add reconstruction loss family for AeroJEPA Four loss families exposed as functional and nn.Module variants: mse_loss / MSELoss (channel-weighted MSE with mask + point weights), relative_l2_loss / RelativeL2Loss (per-channel relative L2 averaged over channels), relative_mse_loss / RelativeMSELoss (relative MSE with selectable pointwise vs channel_max normalization), and the relative_l2_mse_loss / RelativeL2MSELoss hybrid that linearly combines the L2 and MSE terms. Channel weights are stored as a persistent float32 buffer on the Module variants when supplied, and as a non-persistent None buffer otherwise. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Rename and populate metrics/aerojepa with public re-exports Move the JEPA losses subpackage from physicsnemo.experimental.metrics.jepa to .metrics.aerojepa so it mirrors the nn.aerojepa naming. Populate the package __init__ with the 12 public re-exports from masking, sigreg, and reconstruction (flatten/reshape token helpers, SIGReg/TokenLatentSIGReg, and the four reconstruction loss families both functional and as nn.Module). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add tests for the AeroJEPA losses Three test files mirroring the source modules: test_masking, test_sigreg, test_reconstruction. 37 tests covering constructor validation, forward shape, edge cases (rank-1, empty batch, all-False mask), the SIGReg buffer layout, state_dict persistence of channel_weights on the reconstruction Module variants, both modes of relative_mse_loss, and the hybrid degenerating to either of its two sub-losses when the corresponding weight is zero. CPU run is 4 s wall; device fixture picks up CUDA automatically. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add CHANGELOG entry for AeroJEPA losses Documents the new physicsnemo.experimental.metrics.aerojepa subpackage contributed across the preceding 6 commits on this branch: SIGReg / TokenLatentSIGReg regularizers, masking helpers, and the four reconstruction loss families. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add experimental/models/aerojepa subpackage skeleton Create an empty subpackage for the top-level AeroJEPA model and its model-specific subcomponents (context/target/point encoders, decoder, predictor, trunk) that land in subsequent commits. Module docstring points readers to experimental.nn.aerojepa for the reusable building blocks the model is composed from. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add encoder ABCs for AeroJEPA Abstract base classes BaseContextEncoder and BaseTargetEncoder (plus the encoders subpackage init) define the contract concrete encoders must satisfy: a required forward returning an EncoderOutput and an optional forward_batched gated by a supports_batched_forward class flag. The context encoder's forward args are named context_pos / context_feat (these bundle the boundary and any volumetric samples in whole-domain models; the SDF channel in context_feat distinguishes the two halves at inference). The target encoder keeps the surface / volume split because training-time subsamplings for the two are intentionally decoupled. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Drop gen_params from the context-encoder API Context tokens are produced from geometry alone - operating conditions enter the model downstream at the predictor head, not at the context branch. Remove gen_params from BaseContextEncoder forward and forward_batched signatures. Class docstring spells out the intent. BaseTargetEncoder is untouched. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add ContextTransformer and PointTransformer encoders for AeroJEPA PointTransformer (point.py) is a point-cloud encoder building block: tokenizes the input via PointCloudTokenizer, embeds tokens with a Fourier positional encoding plus per-feature linear projection, optionally adds a conditioning vector, runs a stack of LocalPointTransformerBlock layers with configurable dilation, and emits an EncoderOutput. Two entry points - encode_single for unbatched inputs and forward_batched for padded batches with per-batch coordinate offsetting so the inner k-NN does not mix tokens across batch items. The same file carries the build_geometry_features helper (assembles per-point features from positions and optional SDF / normals / n-dot channels) and the message-passing PointClusterGraphPool used when tokenizer_cluster_pooling='graph'. ContextTransformer (context.py) is the concrete BaseContextEncoder. Takes context_pos and context_feat - no gen_params, since operating conditions enter the model downstream at the predictor head. Internally wraps PointTransformer with conditioning disabled. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Drop gen_params from BaseTargetEncoder Mirror the context-side change: target encoders take their inputs straight, with no gen_params threaded through. Operating conditions enter the model only at the predictor head. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Drop context_tokens from BaseTargetEncoder JEPA target encoders are self-attention only. Remove context_tokens from forward and forward_batched. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add TargetTransformer for AeroJEPA Concrete BaseTargetEncoder that wraps an inner PointTransformer. Forward concatenates surface and volume into one bundled point set; forward_batched weaves variable-length surface and volume halves per batch via counts_to_mask. Self-attention only. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add QueryTokenDecoder for AeroJEPA Implicit field decoder driven by cross-attention to target tokens. Per-query embedding is a Fourier positional encoding plus optional SDF channel and optional cond vector; cross-attention to the target token set refines it, a trunk MLP and head produce the output. Several optional behaviors wire in: wall-velocity gate, pressure split head (MLP or SIREN), final SIREN refinement, extra SDF features. Both forward (single) and forward_batched (padded) process queries in chunks of query_chunk_size and return (pred, query_embeddings). SineLayer and SirenHead are the small SIREN building blocks used by the optional heads. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add PrototypeTokenJEPAHead for AeroJEPA The JEPA predictor head. Maps a target-token coordinate set to predicted target-token features, given context tokens and a conditioning vector. Operating conditions enter the model here (via the cond argument), projected once and threaded into every self- and cross-attention block. Accepts both unbatched (rank-2 context features) and padded batched (rank-3) inputs; target_positions and cond are broadcast across the batch when their leading dim is 1. The forward signature uses target_positions as the parameter name (not target_coords) for consistency with the rest of the model API. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add AeroJEPATrunk Owns context encoder, target encoder, and decoder, and wires them together. encode_context runs both encoders and emits a dict with context tokens, target tokens, and the decoder-side cond_global. decode_queries decodes a target token set at supplied query positions, optionally producing a per-query mask logit when the mask head is enabled. forward_single and forward_batch are convenience wrappers chaining the two phases for unbatched and padded batched inputs respectively. Public args use context_pos / context_feat naming; gen_params is used to build cond_global but is not threaded into the encoders. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add top-level AeroJEPA model Composes AeroJEPATrunk and PrototypeTokenJEPAHead into a single physicsnemo.Module with full MOD-001 / MOD-006 / MOD-010 compliance: @DataClass AeroJEPAMetaData inheriting ModelMetaData, jaxtyping on all public methods, validation guarded by torch.compiler.is_compiling, constructor taking typed components (no cfg dict, no kwargs). The forward entry takes context_pos / context_feat / gen_params / query_pos / query_sdf and derives target-token positions internally via build_target_token_coords (the target encoder's tokenizer with a placeholder feature tensor); callers no longer supply target_coords. predict is a no-grad wrapper around forward. encode_geometry, encode_geometry_and_flow, predict_field_tokens, decode_field, and build_target_token_coords are exposed for training-loop callers and for latent-optimization workflows that want to cache target coordinates across many predictor evaluations. decode_field_chunked wraps the decoder in chunked + autocast + CPU-offload for memory-bounded inference on very large query sets. The class docstring ships an executable doctest that wires the whole chain and asserts a forward-pass shape. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Populate experimental/models/aerojepa __init__ with public re-exports Re-export AeroJEPA + AeroJEPAMetaData (top-level model), AeroJEPATrunk and PrototypeTokenJEPAHead (composable components), QueryTokenDecoder, BaseContextEncoder / BaseTargetEncoder (ABCs for custom encoders), and the three concrete encoders ContextTransformer / PointTransformer / TargetTransformer. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add tests for the AeroJEPA model Ten new test files mirroring the source modules: encoders/test_base, encoders/test_context, encoders/test_point, encoders/test_target, test_decoder, test_predictor, test_trunk, and test_aerojepa. 63 tests covering constructor validation, signature checks (drops of target_coords / gen_params / context_tokens per the API redesign), forward / forward_batched shapes, every optional decoder feature (pressure split head, SIREN refinement, wall gate, extra SDF features), predictor broadcasting paths, trunk wiring (mask head on/off), and the top-level model contract (physicsnemo Module subclass, plain-tensor forward, no-grad predict, single-arg build_target_token_coords, chunked CPU-offload decode). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Add CHANGELOG entry for the AeroJEPA model Documents the new physicsnemo.experimental.models.aerojepa subpackage: the AeroJEPA top-level model and its composable subcomponents (context/target encoders, decoder, predictor, trunk). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: scaffold AeroJEPA SuperWing recipe (README + requirements) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add Hydra config skeleton for AeroJEPA SuperWing recipe Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): give target encoder its own positions and rename surface_main_feat Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: port SuperWing datapipe + split/normalization builders Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add training helpers (EMA, runtime utils, optimizer/warmup builders) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add loss builders (sigreg, recon, latent) for AeroJEPA recipe Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add train.py entry point for the AeroJEPA SuperWing recipe Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add SuperWing surface-field plot helpers Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add inference.py + config for AeroJEPA SuperWing recipe Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add CL/CD postprocessing and extend inference dump Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: embed AeroJEPA SuperWing field + force-coefficient plots in README Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): move layers into models/aerojepa/layers Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): move losses into models/aerojepa/losses, add recipe to CHANGELOG Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: scale AeroJEPA SuperWing config toward paper reproduction Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: add per-channel field-error metrics for the SuperWing test split Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * changelog: move SuperWing recipe bullet to [2.2.0] after upstream version bump Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * examples/cfd: load AeroJEPA checkpoints with weights_only=True (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): fail fast on cluster_size=None for clustering strategies (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): guard empty input ahead of KNN in data_prototype_cluster (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): use physicsnemo TE-aware LayerNorm in place of torch.nn.LayerNorm (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * test/aerojepa: exercise include_geometry_global_in_decoder_cond=True branch (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * test/aerojepa: cover AeroJEPA.forward rank-2 and N-mismatch guards (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * test/aerojepa: cover decode_field_chunked autocast (fp16/bf16) path (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * test/aerojepa: cover build_target_token_coords defensive ValueError branch (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * test/aerojepa: cover point_weights in reconstruction losses (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * test/aerojepa: assert mse_loss with partial mask matches loss on kept rows (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): dedupe FPS, route prototype anchors through point_tokenizer FPS (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): build SirenHead from physicsnemo SirenLayer, drop local SineLayer (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): use physicsnemo.nn.Mlp in place of local Sequential MLPs (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): route k-NN through physicsnemo.nn.functional.knn, drop local backends (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): guard query_sdf=None when decoder.use_sdf=True (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): reuse trunk._build_cond_global_single in encode_geometry (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): disable AdaLN in predictor blocks when cond_dim=0 (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): reuse static-coords kNN across attention block stacks (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * aerojepa(experimental): inherit physicsnemo.core.Module throughout for save/load round-trip (PR NVIDIA#1690 review) Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Promote general-purpose AeroJEPA layers to experimental/nn The attention blocks, point-cloud tokenizer, Fourier positional encoding, and the batch/gather/mask/k-NN helpers are domain-agnostic and reusable, so they move from the model-specific layers package up to physicsnemo.experimental.nn. The two TokenSet-coupled helpers (pad_token_sets, trim_batched_tokens) stay in the layers package. Moved symbols are re-exported from the layers package __init__ for backward-compatible imports, and the tests are split to match. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Condense AeroJEPA CHANGELOG entries to house style Collapse the four verbose AeroJEPA bullets (model, losses, building blocks, recipe) into two concise entries - one for the library additions. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Route tokenizer FPS through physicsnemo.nn.functional.farthest_point_sampling Replace the local farthest-point-sampling implementation in the point-cloud tokenizer with the shared physicsnemo.nn.functional primitive, which gains the Warp-accelerated CUDA backend and torch.compile-safety for free and auto-dispatches between Warp and a torch baseline. Offline prototype-anchor generation keeps a small self-contained seeded FPS helper: it must be reproducible from a seed for the k-means initialization and empty-cluster refill, and the shared primitive does not expose a seed for its random start. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Use physicsnemo.nn.FourierPositionalEmbedding for query coordinate embedding Replace the local Fourier positional-encoding layer with the shared physicsnemo.nn.FourierPositionalEmbedding, which provides the same deterministic axis-wise (NeRF-style) embedding with a matching constructor surface (in_dim, num_bands, include_input) and out_dim. The query decoder, the JEPA predictor head, and the point encoder all construct it directly from physicsnemo.nn; the local layer and its dedicated test are removed (the shared layer carries its own tests). The shared layer lays its features out axis-major rather than band-major; the feature set is identical up to a fixed permutation that the immediately following linear projection absorbs, so model outputs are unaffected beyond a re-init of those projections. Model shape / property / checkpoint round-trip tests are unchanged and pass. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Revise AeroJEPA README for clarity and consistency Updated README.md for AeroJEPA tutorial for title case. And a few edits for clarity (cherry picked from commit b0c7e19) * refactor and move reconstruction losses * move reconstruction.py to src/losses * Use upstreamed point-transformer attention blocks (PR NVIDIA#1690 review) Replace the vendored experimental attention blocks with the upstreamed physicsnemo.nn LocalPointTransformerBlock, LocalTokenCrossAttentionBlock, and AdaLNResidualMLP. Add an optional precomputed_idx argument to both blocks so a stack sharing static coordinates runs one k-NN and threads the shared index through every layer, preserving the static-coords reuse. The caller-side precompute now uses physicsnemo.nn.functional.knn directly, so the per-block chunk-size plumbing is removed. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Point AeroJEPA recipe at the relocated reconstruction losses (PR NVIDIA#1690 review) The reconstruction loss wrapper now lives in the recipe at src/losses/reconstruction.py and builds on the core physicsnemo.metrics losses. Import the channel-weighted / hybrid loss modules from that local wrapper in src/losses/builders.py, and stop re-exporting the reconstruction family from the experimental losses package (which no longer ships it). Relocate the reconstruction-loss tests next to the recipe code, with a recipe-level conftest that puts the recipe root on sys.path and provides the device fixture. Note: the wrapper depends on physicsnemo.metrics.general.relative_error and the weights-enabled mse from PR NVIDIA#1746, which is not merged yet; the recipe imports will resolve once NVIDIA#1746 lands. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Fold token masking helpers into layers/token_utils (PR NVIDIA#1690 review) flatten_valid_token_features and reshape_token_features_for_sigreg are now moved from the losses package into the model's layers/token_utils.py. Repoint sigreg.py at the new location, drop them from the losses package surface, and merge their tests into the token_utils test. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Move SIGReg regularizers into the training recipe (PR NVIDIA#1690 review). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Move AeroJEPA doc images to docs/img/aerojepa (PR NVIDIA#1690 review). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Remove duplicate intro paragraph in AeroJEPA README (PR NVIDIA#1690 review). Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Run AeroJEPA validation under no_grad (PR NVIDIA#1690 review) _run_epoch wrapped the forward only in the autocast context, so the non-training (validation) path still built the autograd graph. Gate grad on is_train via torch.set_grad_enabled so validation runs without tracking gradients. Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> * Accumulate AeroJEPA gradients per sample (PR NVIDIA#1690 review) 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> * Save EMA shadow weights as the AeroJEPA checkpoint model (PR NVIDIA#1690 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> * Tidy AeroJEPA inference: top-level import, grid-order comment (PR NVIDIA#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> * Tidy AeroJEPA training config comments (PR NVIDIA#1690 review) 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> * Drop unused target_min/target_max stats (PR NVIDIA#1690 review) 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> * Fix decoder output-channel comment in AeroJEPA model config (PR NVIDIA#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> * Clarify single-GPU training in AeroJEPA README (PR NVIDIA#1690 review) 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> * Train the target encoder through the JEPA latent loss 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> * Add two-phase AeroJEPA training and optional context SIGReg 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> * Add checkpoint resume and init-from-checkpoint to AeroJEPA training 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> * Fix masked_mean rank-3 unmasked return shape (PR NVIDIA#1690 review) 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> * Pair fp16 autocast with a GradScaler in AeroJEPA training (PR NVIDIA#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> * distributed training and inference, bug fixes, and other improvements * simplify multi-gpu inference * fix te layernorm --------- Signed-off-by: fgiral000 <fa.giral@alumnos.upm.es> Co-authored-by: megnvidia <mmiranda@nvidia.com> Co-authored-by: Mohammad Amin Nabian <m.a.nabiyan@gmail.com>
peterdsharpe
added a commit
to peterdsharpe/physicsnemo
that referenced
this pull request
Jul 21, 2026
…to upstream NVIDIA#1746 (relative_mse/relative_l2), deprecated alias warns
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PhysicsNeMo Pull Request
Description
Adds relative MES and relative L2 to
physicsnemo.metrics.general, and gives mse/rmse an optional weights argument for a masked/weighted mean. Both relative metrics accept element weights (subsuming a validity mask and per-point weights) and a dim reduction, matching general.mse.Checklist
Dependencies
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.
AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.