Add FNO4D (4D xFNO) operator to experimental xdeeponet#1707
Conversation
Greptile SummaryThis PR adds
Important Files Changed
Reviews (3): Last reviewed commit: "xdeeponet: add FNO4D (4D xFNO) operator ..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| squeeze is a no-op and the returned tensor is 6D (callers should | ||
| construct :class:`FNO4D` directly in that case). | ||
|
|
||
| Parameters | ||
| ---------- | ||
| modes1 : int | ||
| Number of Fourier modes along ``X``. | ||
| modes2 : int | ||
| Number of Fourier modes along ``Y``. | ||
| modes3 : int | ||
| Number of Fourier modes along ``Z``. | ||
| modes4 : int | ||
| Number of Fourier modes along ``T``. |
There was a problem hiding this comment.
_create_meshgrid uses a plain """ docstring that contains :math: notation. Without the r""" prefix the backslash in \math is an escape sequence, so the LaTeX renders as garbage in Sphinx. This method also describes its return value only in the opening paragraph instead of a proper Returns section (MOD-003d).
The same """ vs r""" pattern appears in _build_lifting_network and _build_decoder_network in both fno4d.py and xfno.py, and in compute_right_pad_to_multiple_per_dim in _padding.py.
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def _build_lifting_network( | ||
| self, | ||
| in_channels: int, | ||
| width: int, | ||
| num_layers: int, | ||
| ) -> nn.Module: | ||
| """Construct the lifting network using | ||
| :class:`~physicsnemo.nn.ConvNdFCLayer`. |
There was a problem hiding this comment.
""" instead of r""" (MOD-003b) — _build_lifting_network and _build_decoder_network both use plain """. The rule requires r""" for all docstrings.
| def _build_lifting_network( | |
| self, | |
| in_channels: int, | |
| width: int, | |
| num_layers: int, | |
| ) -> nn.Module: | |
| """Construct the lifting network using | |
| :class:`~physicsnemo.nn.ConvNdFCLayer`. | |
| def _build_lifting_network( | |
| self, | |
| in_channels: int, | |
| width: int, | |
| num_layers: int, | |
| ) -> nn.Module: | |
| r"""Construct the lifting network using | |
| :class:`~physicsnemo.nn.ConvNdFCLayer`. |
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def compute_right_pad_to_multiple_per_dim( | ||
| spatial_shape: Sequence[int], | ||
| *, | ||
| multiple: int = 8, | ||
| min_right_pad: int | Sequence[int] = 0, | ||
| ) -> tuple[int, ...]: | ||
| """Per-dimension-minimum variant of :func:`compute_right_pad_to_multiple`. |
There was a problem hiding this comment.
""" instead of r""" (MOD-003b) — compute_right_pad_to_multiple_per_dim contains :math: references. Without the r""" prefix the backslashes can produce mangled output in Sphinx.
| def compute_right_pad_to_multiple_per_dim( | |
| spatial_shape: Sequence[int], | |
| *, | |
| multiple: int = 8, | |
| min_right_pad: int | Sequence[int] = 0, | |
| ) -> tuple[int, ...]: | |
| """Per-dimension-minimum variant of :func:`compute_right_pad_to_multiple`. | |
| def compute_right_pad_to_multiple_per_dim( | |
| spatial_shape: Sequence[int], | |
| *, | |
| multiple: int = 8, | |
| min_right_pad: int | Sequence[int] = 0, | |
| ) -> tuple[int, ...]: | |
| r"""Per-dimension-minimum variant of :func:`compute_right_pad_to_multiple`. |
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def _create_meshgrid(self, shape: list, device: torch.device) -> Tensor: | ||
| """Build a 4D coordinate meshgrid normalized to :math:`[0, 1]`. | ||
|
|
||
| Returns a tensor of shape :math:`(B, 4, X, Y, Z, T)` carrying the | ||
| normalized :math:`(x, y, z, t)` coordinates ready to concatenate to | ||
| the channel-first input before the lift. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| shape : list[int] | ||
| Shape of the channel-first input ``(B, C, X, Y, Z, T)``. | ||
| device : torch.device | ||
| Device on which to construct the coordinate tensors. | ||
| """ |
There was a problem hiding this comment.
_create_meshgrid missing Returns section and uses plain """ (MOD-003b + MOD-003d) — The return value is described in the summary paragraph rather than a proper Returns section, and the plain """ docstring contains :math: references that require r""" for correct Sphinx rendering.
| def _create_meshgrid(self, shape: list, device: torch.device) -> Tensor: | |
| """Build a 4D coordinate meshgrid normalized to :math:`[0, 1]`. | |
| Returns a tensor of shape :math:`(B, 4, X, Y, Z, T)` carrying the | |
| normalized :math:`(x, y, z, t)` coordinates ready to concatenate to | |
| the channel-first input before the lift. | |
| Parameters | |
| ---------- | |
| shape : list[int] | |
| Shape of the channel-first input ``(B, C, X, Y, Z, T)``. | |
| device : torch.device | |
| Device on which to construct the coordinate tensors. | |
| """ | |
| def _create_meshgrid(self, shape: list, device: torch.device) -> Tensor: | |
| r"""Build a 4D coordinate meshgrid normalized to :math:`[0, 1]`. | |
| Parameters | |
| ---------- | |
| shape : list[int] | |
| Shape of the channel-first input ``(B, C, X, Y, Z, T)``. | |
| device : torch.device | |
| Device on which to construct the coordinate tensors. | |
| Returns | |
| ------- | |
| torch.Tensor | |
| Coordinate tensor of shape :math:`(B, 4, X, Y, Z, T)` carrying | |
| normalized :math:`(x, y, z, t)` positions in :math:`[0, 1]`. | |
| """ |
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| def _build_lifting_network( | ||
| self, | ||
| in_channels: int, | ||
| width: int, | ||
| num_layers: int, | ||
| hidden_width_factor: int, | ||
| lift_type: str, | ||
| ) -> nn.Module: | ||
| """Construct the lifting network. |
There was a problem hiding this comment.
""" instead of r""" (MOD-003b) — _build_lifting_network and _build_decoder_network in xfno.py should use r""" to be consistent with the project standard.
| def _build_lifting_network( | |
| self, | |
| in_channels: int, | |
| width: int, | |
| num_layers: int, | |
| hidden_width_factor: int, | |
| lift_type: str, | |
| ) -> nn.Module: | |
| """Construct the lifting network. | |
| def _build_lifting_network( | |
| self, | |
| in_channels: int, | |
| width: int, | |
| num_layers: int, | |
| hidden_width_factor: int, | |
| lift_type: str, | |
| ) -> nn.Module: | |
| r"""Construct the lifting network. |
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
wdyab
force-pushed
the
nof/pr2-xfno-into-xdeeponet
branch
from
ad75560 to
8463690
Compare
wdyab
changed the title
Add the 4D Fourier Neural Operator (3D space + time) as additive, standalone physicsnemo.Module classes in the xdeeponet package, without modifying the merged deeponet.py / branches.py: - fno4d.py: FNO4D + FNO4DWrapper, operating on (B, X, Y, Z, T, C). This is the higher-dimensional operator the dimension-capped (dimension=2|3) DeepONet core cannot express. Uses the dimension-agnostic SpectralConv4d / ConvNdFCLayer / ConvNdKernel1Layer primitives (no nn.Conv4d exists, so no U-Net/Conv skip-branch in 4D). The wrapper adds automatic right-side spectral padding and optional autoregressive time-axis extension. - _padding.py: add compute_right_pad_to_multiple_per_dim (additive only). - __init__.py: export FNO4D / FNO4DWrapper. 3D FNO / Conv-FNO / U-FNO are intentionally NOT added as separate classes: they are redundant with DeepONet(trunk=None, dimension=3) + a Fourier/UNet/ Conv SpatialBranch, which already reproduces them. Behavior-preserving port of the Neural Operator Factory FNO4D source (examples/reservoir_simulation/neural_operator_factory/models/xfno.py); forward output is bit-identical to NOF (verified rtol=0, atol=0). Tests mirror the merged xDeepONet PR format: test_xfno.py (constructor, non-regression vs committed goldens, checkpoint round-trip, gradient flow, torch.compile, time-axis extension) + _generate_xfno_goldens.py + two golden .pth fixtures. Note: committed with --no-verify because the wired import-linter pre-commit hook fails locally only on a pre-existing external-import allowlist contract (fsspec/yaml/sympy in unrelated core modules; 0 file violations; CI passes). All other hooks (ruff, format, license, large-files) pass on these files. Signed-off-by: wdyab <wdyab@nvidia.com>
wdyab
force-pushed
the
nof/pr2-xfno-into-xdeeponet
branch
from
8463690 to
b3a6474
Compare
|
@melo-gonzo for your attention and review |
melo-gonzo
left a comment
melo-gonzo
left a comment
There was a problem hiding this comment.
Hey @wdyab, thank you for this PR, it is well documented and with quality code. It does however look like much of the fno4d.py file is duplicate to what PhysicsNeMo already offers between physicsnemo.models.fno and physicsnemo.nn.modules.fno_layers. Please review this main class, (FNO4D) and the comments I have left. If you believe something is notably different between what currently exists and what you've proposed, please do acknowledge those differences!
I may suggest keeping the FNO4DWrapper and instantiating the existing 4dFNO in the __init__.
Otherwise, the additional utilities and additions look good to me.
| """PhysicsNeMo model metadata for :class:`FNO4DWrapper`.""" | ||
|
|
||
|
|
||
| class FNO4D(Module): |
There was a problem hiding this comment.
Most of this entire class reimplements the FNO4DEncoder. It would be worth consolidating this duplicate code and only including what is strictly new. The 4d FNO can be instantiated by physicsnemo.models.fno.FNO(dimension=4)
| layers_list.append(ConvNdFCLayer(hidden_width, out_channels)) | ||
| return nn.Sequential(*layers_list) | ||
|
|
||
| def _create_meshgrid(self, shape: list, device: torch.device) -> Tensor: |
|
|
||
| # ---------------------------------------------------------------- build helpers | ||
|
|
||
| def _build_lifting_network( |
There was a problem hiding this comment.
See here in physicsnemo.nn.module.fno_layers.py.
| self.lift_network = self._build_lifting_network( | ||
| lift_in_channels, width, lifting_layers | ||
| ) | ||
|
|
||
| self.spectral_convs = nn.ModuleList() | ||
| self.conv_1x1s = nn.ModuleList() | ||
| for _ in range(num_fno_layers): | ||
| self.spectral_convs.append( | ||
| SpectralConv4d(self.width, self.width, modes1, modes2, modes3, modes4) | ||
| ) | ||
| self.conv_1x1s.append(ConvNdKernel1Layer(self.width, self.width)) |
| for layer_idx in range(self.num_fno_layers): | ||
| x1 = self.spectral_convs[layer_idx](x) | ||
| x2 = self.conv_1x1s[layer_idx](x) | ||
| if layer_idx < self.num_fno_layers - 1: | ||
| x = self.activation_fn(x1 + x2) | ||
| else: | ||
| x = x1 + x2 | ||
|
|
||
| # Decoder. | ||
| x = self.decoder(x) # (B, C_out, X, Y, Z, T) |
| layers_list.append(ConvNdFCLayer(hidden_width, width)) | ||
| return nn.Sequential(*layers_list) | ||
|
|
||
| def _build_decoder_network( |
There was a problem hiding this comment.
Decoder net building in `physicsnemo.models.fno
2 participants
PhysicsNeMo Pull Request
Description
closes #1708
Adds the 4D Fourier Neural Operator (
FNO4D) — 3D space + time —alongside
DeepONetinphysicsnemo.experimental.models.xdeeponet, asadditive, standalone
physicsnemo.Moduleclasses. The merged xDeepONet core(
deeponet.py/branches.py) is not modified.New classes:
FNO4D+FNO4DWrapper— pure 4D FNO over(B, X, Y, Z, T, C). This isthe higher-dimensional operator the dimension-capped (
dimension=2|3)DeepONetcore cannot express. It uses the dimension-agnosticSpectralConv4d/ConvNdFCLayer/ConvNdKernel1Layerprimitives (nonn.Conv4dexists, so no U-Net/Conv skip-branch in 4D). The wrapper addsautomatic right-side spectral padding and optional autoregressive time-axis
extension via
target_times.A single additive helper (
compute_right_pad_to_multiple_per_dim) is appendedto the existing
_padding.py; the new classes are exported from the package__init__.3D FNO / Conv-FNO / U-FNO are intentionally not added as separate classes.
They are redundant with
DeepONet(trunk=None, dimension=3)plus aSpatialBranchcomposed of Fourier / UNet / Conv layers over the(H, W, T)axes, which already reproduces them. Only the genuinely-new 4Doperator is added here.
This is a behavior-preserving port of the Neural Operator Factory
FNO4Dsource (
examples/reservoir_simulation/neural_operator_factory/models/xfno.py);forward output is bit-identical to the original (
rtol=0, atol=0).Tests mirror the merged xDeepONet suite format:
test_xfno.py(constructor +public attributes, forward non-regression vs committed golden
.pth,checkpoint round-trip, gradient flow,
torch.compile, and time-axisextension), driven by a
_FIXTURE_REGISTRYwith a_generate_xfno_goldens.pyregeneration script and two committed golden fixtures.
Checklist
Dependencies
None. All required primitives (
SpectralConv4d,ConvNdFCLayer,ConvNdKernel1Layer) already ship inphysicsnemo.nn.