Skip to content

Add xDeepONet family to experimental models#1576

Open
wdyab wants to merge 3 commits intoNVIDIA:mainfrom
wdyab:pr/xdeeponet
Open

Add xDeepONet family to experimental models#1576
wdyab wants to merge 3 commits intoNVIDIA:mainfrom
wdyab:pr/xdeeponet

Conversation

@wdyab
Copy link
Copy Markdown

@wdyab wdyab commented Apr 17, 2026

Summary

Introduces physicsnemo.experimental.models.xdeeponet — a config-driven,
unified implementation of eight DeepONet-based operator-learning
architectures for both 2D and 3D spatial domains:

  • deeponet, u_deeponet, fourier_deeponet, conv_deeponet,
    hybrid_deeponet — single-branch variants
  • mionet, fourier_mionet — two-branch multi-input variants
  • tno — Temporal Neural Operator (branch2 = previous solution)

This is the first of several PRs restructuring the Neural Operator
Factory per discussion with code owners. Subsequent PRs will upstream
xFNO (experimental) and refactor the reservoir-simulation NOF example
to consume these library models.

Closes Issue #1575

Key features

  • Composable spatial branches (Fourier, UNet, Conv in any combination)
  • Three decoder types: mlp, conv, temporal_projection
  • Automatic spatial padding
  • Automatic trunk coordinate extraction (time or grid)
  • Optional adaptive pooling for resolution-agnostic training

Design decisions (per @coreyjadams guidance)

  • Placed under experimental/ per the convention for new models.
  • Custom UNet dropped.
  • Tests live at test/experimental/models/ for CI coverage.

Checklist

  • I am familiar with the Contributing Guidelines
  • New tests cover these changes (29 tests under
    test/experimental/models/test_xdeeponet.py)
  • The documentation is up to date (package README, docstrings)
  • The CHANGELOG.md is up to date
  • An issue (Add xDeepONet family to experimental models #1575) is linked to this pull request
  • I have followed the Models Implementation Coding Standards

Test plan

  • 29 unit tests pass locally (branch shapes, wrappers, variants,
    decoder types, temporal projection, target_times override,
    gradient flow, adaptive pooling)
  • ruff check, ruff format, interrogate, markdownlint,
    license pre-commit hooks pass
  • No modifications to existing code — pure addition under
    physicsnemo/experimental/ and test/experimental/
  • All commits signed off per DCO

Related

Introduces physicsnemo.experimental.models.xdeeponet — a config-driven,
unified implementation of eight DeepONet-based operator-learning
architectures for both 2D and 3D spatial domains:

- deeponet, u_deeponet, fourier_deeponet, conv_deeponet, hybrid_deeponet
  (single-branch variants)
- mionet, fourier_mionet (two-branch multi-input variants)
- tno (Temporal Neural Operator; branch2 = previous solution)

Features:
- Composable spatial branches (Fourier, UNet, Conv in any combination)
- Three decoder types: mlp, conv, temporal_projection
- Automatic spatial padding to multiples of 8
- Automatic trunk coordinate extraction (time or grid)
- Optional adaptive pooling (internal_resolution) for
  resolution-agnostic training and inference

Uses physicsnemo.models.unet.UNet as the UNet sub-module; a small
internal adapter tiles a short time axis to reuse the library's 3D UNet
for 2D spatial branches.  Imports spectral, convolutional, and MLP
layers from physicsnemo.nn and physicsnemo.models.mlp.

Includes 29 unit tests covering all variants (2D/3D), decoder types,
temporal projection, target_times override, gradient flow, and
adaptive pooling.

Related discussion with code owners:
- Placed under experimental/ per PhysicsNeMo convention for new models.
- Custom UNet dropped in favour of library UNet.
- Tests under test/experimental/models/ for CI coverage.

Signed-off-by: wdyab <wdyab@nvidia.com>
Made-with: Cursor
@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps bot commented Apr 17, 2026

Greptile Summary

This PR adds a config-driven xDeepONet family (8 variants, 2D + 3D) under physicsnemo.experimental.models.xdeeponet. The overall architecture is well-structured, but there are four P1 defects that need to be addressed before merging:

  • DeepONetWrapper / DeepONet3DWrapper inherit from nn.Module instead of physicsnemo.Module, breaking physicsnemo's checkpoint system for the recommended entry points (MOD-001).
  • Multi-branch variants (mionet, fourier_mionet, tno) lack construction-time validation that branch2_config is provided — omitting it silently creates an incorrect single-branch model with no error raised.
  • set_output_window mutates temporal_head from None to nn.Linear, causing state_dict key inconsistency that breaks checkpoint save/load for temporal_projection models.
  • MLPBranch with num_layers=1 produces the same two-layer network as num_layers=2 because range(-1) is empty in Python.

Important Files Changed

Filename Overview
physicsnemo/experimental/models/xdeeponet/deeponet.py Core 2D/3D DeepONet architectures — multi-branch variant construction lacks a guard ensuring mionet/tno variants receive branch2_config, temporal_head state_dict mutation breaks checkpointing, and docstrings don't meet MOD-003 standards.
physicsnemo/experimental/models/xdeeponet/wrappers.py Recommended public entry points DeepONetWrapper and DeepONet3DWrapper incorrectly inherit from nn.Module instead of physicsnemo.Module, violating MOD-001 and preventing use of physicsnemo serialisation.
physicsnemo/experimental/models/xdeeponet/branches.py MLPBranch has a range(-1) edge case that makes num_layers=1 identical to num_layers=2; spatial branches and TrunkNet logic appear correct but all classes need MOD-003 docstring fixes.
physicsnemo/experimental/models/xdeeponet/padding.py Padding helpers are well-implemented with good edge-case handling; reshape trick for replicate padding on high-dimensional tensors is sound.
physicsnemo/experimental/models/xdeeponet/init.py Clean public API export; all eight variants and building blocks correctly re-exported.
test/experimental/models/test_xdeeponet.py 29 tests cover shapes, gradient flow, temporal projection, and adaptive pooling; however no non-regression tests with reference data (MOD-008b) or checkpoint-loading tests (MOD-008c) are present, which are required before promotion out of experimental.
physicsnemo/experimental/models/xdeeponet/README.md Comprehensive README with usage examples, variant table, branch config schema, and references; content is accurate relative to the implementation.
CHANGELOG.md Changelog entry correctly prepended under the current unreleased section.

Reviews (1): Last reviewed commit: "Add xDeepONet family to experimental mod..." | Re-trigger Greptile

Comment thread physicsnemo/experimental/models/xdeeponet/wrappers.py Outdated
Comment thread physicsnemo/experimental/models/xdeeponet/deeponet.py
Comment thread physicsnemo/experimental/models/xdeeponet/deeponet.py Outdated
Comment thread physicsnemo/experimental/models/xdeeponet/branches.py
Comment thread physicsnemo/experimental/models/xdeeponet/deeponet.py
Comment thread physicsnemo/experimental/models/xdeeponet/deeponet.py
wdyab added 2 commits April 17, 2026 11:35
Fix six issues flagged by the Greptile review:

- Make DeepONetWrapper / DeepONet3DWrapper inherit from
  physicsnemo.core.module.Module (MOD-001). Core DeepONet / DeepONet3D
  also pass proper MetaData dataclasses.
- Raise ValueError at __init__ when mionet / fourier_mionet / tno are
  constructed without branch2_config (prevents silent degradation to a
  single-branch model).
- Add optional output_window constructor parameter so the
  temporal_projection decoder registers temporal_head at __init__,
  producing a deterministic state_dict that round-trips cleanly.
  set_output_window is retained for backwards compatibility.
- Raise ValueError from MLPBranch when num_layers < 2.
- Convert public docstrings to r-prefixed raw strings with
  Parameters / Forward / Outputs sections and LaTeX shape notation
  per MOD-003.
- Add jaxtyping.Float annotations and torch.compiler.is_compiling()
  guarded shape validation to all public forward methods
  (MOD-005, MOD-006).

Signed-off-by: wdyab <wdyab@nvidia.com>
Made-with: Cursor
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant