Update sym to use 2.0 core apis

CHANGELOG.md

@@ -25,6 +25,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Dependencies
 
+- Updated package to use `physicsnemo` v2.0 APIs
+
 ## [2.3.0] - 2025-11-17
 
 ### Added

Makefile

@@ -1,6 +1,6 @@
 install:
 	pip install --upgrade pip && \
-		pip install -e . --no-build-isolation
+		pip install -e . --no-build-isolation --no-deps
 
 setup-ci:
 	pip install pre-commit && \

docs/user_guide/neural_operators/darcy_pino.rst

@@ -101,7 +101,7 @@ which is done via its ``.forward(...)`` method:
 
 .. code:: python
 
-    from physicsnemo.models.layers import fourier_derivatives # helper function for computing spectral derivatives
+    from physicsnemo.nn.module.spectral_layers import fourier_derivatives # helper function for computing spectral derivatives
     from ops import dx, ddx # helper function for computing finite difference derivatives
 
 .. literalinclude:: ../../../examples/darcy/darcy_PINO.py

examples/darcy/darcy_PINO.py

@@ -21,7 +21,7 @@
 import torch.nn.functional as F
 
 import physicsnemo.sym
-from physicsnemo.models.layers.spectral_layers import fourier_derivatives
+from physicsnemo.nn.module.spectral_layers import fourier_derivatives
 from physicsnemo.sym.hydra import instantiate_arch, PhysicsNeMoConfig
 from physicsnemo.sym.key import Key
 from physicsnemo.sym.node import Node

examples/reservoir_simulation/2D/src/Forward_problem_AFNOP.py

@@ -1456,8 +1456,8 @@ def run(cfg: PhysicsNeMoConfig) -> None:
     )
     nodes = (
         [darcyy]
-        + [afno_pressure.make_node("afnop_forward_model_pressure", jit=cfg.jit)]
-        + [afno_saturation.make_node("afnop_forward_model_saturation", jit=cfg.jit)]
+        + [afno_pressure.make_node("afnop_forward_model_pressure", jit=False)]
+        + [afno_saturation.make_node("afnop_forward_model_saturation", jit=False)]
     )
 
     # [constraint]

examples/reservoir_simulation/2D/src/Forward_problem_PINO.py

@@ -1458,8 +1458,8 @@ def run(cfg: PhysicsNeMoConfig) -> None:
     )
     nodes = (
         [darcyy]
-        + [fno_pressure.make_node("pino_forward_model_pressure", jit=cfg.jit)]
-        + [fno_saturation.make_node("pino_forward_model_saturation", jit=cfg.jit)]
+        + [fno_pressure.make_node("pino_forward_model_pressure", jit=False)]
+        + [fno_saturation.make_node("pino_forward_model_saturation", jit=False)]
     )
 
     # [constraint]

examples/reservoir_simulation/3D/src/Forward_problem_PINO.py

@@ -1603,8 +1603,8 @@ def run(cfg: PhysicsNeMoConfig) -> None:
     )
     nodes = (
         [darcyy]
-        + [fno_pressure.make_node("pino_forward_model_pressure", jit=cfg.jit)]
-        + [fno_saturation.make_node("pino_forward_model_saturation", jit=cfg.jit)]
+        + [fno_pressure.make_node("pino_forward_model_pressure", jit=False)]
+        + [fno_saturation.make_node("pino_forward_model_saturation", jit=False)]
     )
 
     # [constraint]

examples/reservoir_simulation/GenAI_workflows/2D/src/Forward_problem_AFNOP.py

@@ -1583,8 +1583,8 @@ def run(cfg: PhysicsNeMoConfig) -> None:
     )
     nodes = (
         [darcyy]
-        + [afno_pressure.make_node("afnop_forward_model_pressure", jit=cfg.jit)]
-        + [afno_saturation.make_node("afnop_forward_model_saturation", jit=cfg.jit)]
+        + [afno_pressure.make_node("afnop_forward_model_pressure", jit=False)]
+        + [afno_saturation.make_node("afnop_forward_model_saturation", jit=False)]
     )
 
     # [constraint]

examples/reservoir_simulation/GenAI_workflows/2D/src/Forward_problem_PINO.py

@@ -1587,8 +1587,8 @@ def run(cfg: PhysicsNeMoConfig) -> None:
     )
     nodes = (
         [darcyy]
-        + [fno_pressure.make_node("pino_forward_model_pressure", jit=cfg.jit)]
-        + [fno_saturation.make_node("pino_forward_model_saturation", jit=cfg.jit)]
+        + [fno_pressure.make_node("pino_forward_model_pressure", jit=False)]
+        + [fno_saturation.make_node("pino_forward_model_saturation", jit=False)]
     )
 
     # [constraint]

examples/reservoir_simulation/GenAI_workflows/3D/src/Forward_problem_PINO.py

@@ -1635,8 +1635,8 @@ def run(cfg: PhysicsNeMoConfig) -> None:
     )
     nodes = (
         [darcyy]
-        + [fno_pressure.make_node("pino_forward_model_pressure", jit=cfg.jit)]
-        + [fno_saturation.make_node("pino_forward_model_saturation", jit=cfg.jit)]
+        + [fno_pressure.make_node("pino_forward_model_pressure", jit=False)]
+        + [fno_saturation.make_node("pino_forward_model_saturation", jit=False)]
     )
 
     # [constraint]

physicsnemo/sym/geometry/tessellation.py

@@ -140,7 +140,7 @@ def sdf(invar, params, compute_sdf_derivatives=False):
                 if airtight:
                     # Import locally to avoid pickling issues with multiprocessing
                     import torch
-                    from physicsnemo.utils.sdf import signed_distance_field
+                    from physicsnemo.nn.functional import signed_distance_field
 
                     torch_vertices = torch.from_numpy(store_triangles)
                     torch_indices = torch.arange(

physicsnemo/sym/geometry/tessellation_warp.py

@@ -385,7 +385,7 @@ def sdf(triangles, airtight, invar, params, compute_sdf_derivatives=False):
         if airtight:
             # Import locally to avoid pickling issues with multiprocessing
             import torch
-            from physicsnemo.utils.sdf import signed_distance_field
+            from physicsnemo.nn.functional import signed_distance_field
 
             torch_vertices = torch.from_numpy(store_triangles)
             torch_indices = torch.arange(store_triangles.shape[0], dtype=torch.int32)

physicsnemo/sym/models/dgm.py

@@ -21,7 +21,7 @@
 from torch import Tensor
 
 
-from physicsnemo.models.layers import FCLayer, DGMLayer
+from physicsnemo.nn import FCLayer, DGMLayer
 from physicsnemo.sym.models.activation import Activation, get_activation_fn
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key

physicsnemo/sym/models/fno.py

@@ -22,15 +22,15 @@
 import torch.nn.functional as F
 import logging
 
-from physicsnemo.models.layers import (
+from physicsnemo.nn import (
     Conv1dFCLayer,
     Conv2dFCLayer,
     Conv3dFCLayer,
     SpectralConv1d,
     SpectralConv2d,
     SpectralConv3d,
 )
-from physicsnemo.models.layers.spectral_layers import (
+from physicsnemo.nn.module.spectral_layers import (
     calc_latent_derivatives,
     first_order_pino_grads,
     second_order_pino_grads,

physicsnemo/sym/models/fourier_net.py

@@ -19,7 +19,7 @@
 from typing import Dict, List, Tuple
 
 import physicsnemo.sym.models.fully_connected as fully_connected
-from physicsnemo.models.layers import FourierLayer
+from physicsnemo.nn import FourierLayer
 from physicsnemo.sym.models.activation import Activation
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key

physicsnemo/sym/models/fully_connected.py

@@ -25,7 +25,7 @@
     amp_manager_scaler_enabled_and_disable_autocast_firstlayer,
 )
 
-from physicsnemo.models.layers import FCLayer, Conv1dFCLayer
+from physicsnemo.nn import FCLayer, Conv1dFCLayer
 from physicsnemo.sym.models.activation import Activation, get_activation_fn
 from physicsnemo.sym.models.arch import Arch
 

physicsnemo/sym/models/hash_encoding_net.py

@@ -23,7 +23,7 @@
 
 import physicsnemo.sym.models.fully_connected as fully_connected
 from physicsnemo.sym.models.activation import Activation
-from physicsnemo.models.layers.interpolation import (
+from physicsnemo.nn.module.interpolation import (
     _hyper_cube_weighting,
     smooth_step_2,
 )

physicsnemo/sym/models/highway_fourier_net.py

@@ -20,7 +20,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from physicsnemo.models.layers import FCLayer, FourierLayer
+from physicsnemo.nn import FCLayer, FourierLayer
 from physicsnemo.sym.models.activation import Activation, get_activation_fn
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key

physicsnemo/sym/models/modified_fourier_net.py

@@ -20,7 +20,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from physicsnemo.models.layers import FCLayer, FourierLayer
+from physicsnemo.nn import FCLayer, FourierLayer
 from physicsnemo.sym.models.activation import Activation, get_activation_fn
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key

physicsnemo/sym/models/multiplicative_filter_net.py

@@ -20,7 +20,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from physicsnemo.models.layers import FCLayer, FourierFilter, GaborFilter
+from physicsnemo.nn import FCLayer, FourierFilter, GaborFilter
 from physicsnemo.sym.models.activation import Activation, get_activation_fn
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key

physicsnemo/sym/models/multiscale_fourier_net.py

@@ -20,7 +20,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from physicsnemo.models.layers import FCLayer, FourierLayer
+from physicsnemo.nn import FCLayer, FourierLayer
 from physicsnemo.sym.models.activation import Activation, get_activation_fn
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key

physicsnemo/sym/models/radial_basis.py

@@ -22,7 +22,7 @@
 from torch import Tensor
 
 
-from physicsnemo.models.layers import FCLayer
+from physicsnemo.nn import FCLayer
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key
 

physicsnemo/sym/models/siren.py

@@ -19,7 +19,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from physicsnemo.models.layers import SirenLayer, SirenLayerType
+from physicsnemo.nn import SirenLayer, SirenLayerType
 from physicsnemo.sym.models.arch import Arch
 from physicsnemo.sym.key import Key
 from physicsnemo.sym.constants import NO_OP_NORM

pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
     "notebook>=7.2.2",
     "numpoly<=1.3.4",
     "numpy-stl>=2.16",
-    "nvidia-physicsnemo>=1.0.0",
+    "nvidia-physicsnemo>=2.0.0",
     "opencv-python>=4.8.1.78",
     "pillow>=10.3",
     "pint>=0.19.2",
@@ -51,3 +51,6 @@ include = ["physicsnemo.*"]
 
 [tool.setuptools.package-data]
 "*" = ["*.yaml"]
+
+[tool.ruff]
+exclude = ["examples"]

test/test_models/test_fno.py

@@ -39,8 +39,8 @@ def test_fno_1d():
         fno_modes=4,
         padding=0,
     )
-    # Testing JIT
-    model.make_node(name="FNO1d", jit=True)
+    # TorchScript (jit) disabled: upstream spectral layers use jaxtyping, which TorchScript does not support.
+    model.make_node(name="FNO1d", jit=False)
 
     bsize = 5
     invar = {
@@ -68,8 +68,8 @@ def test_fno_2d():
         fno_modes=16,
     )
 
-    # Testing JIT
-    model.make_node(name="FNO2d", jit=True)
+    # TorchScript (jit) disabled: upstream spectral layers use jaxtyping, which TorchScript does not support.
+    model.make_node(name="FNO2d", jit=False)
 
     bsize = 5
     invar = {
@@ -99,8 +99,8 @@ def test_fno_3d():
         fno_modes=16,
     )
 
-    # Testing JIT
-    model.make_node(name="FNO3d", jit=True)
+    # TorchScript (jit) disabled: upstream spectral layers use jaxtyping, which TorchScript does not support.
+    model.make_node(name="FNO3d", jit=False)
 
     bsize = 5
     invar = {

test/test_spectral_convs.py

@@ -17,7 +17,7 @@
 import torch
 import torch.nn as nn
 
-from physicsnemo.models.layers import SpectralConv1d, SpectralConv2d, SpectralConv3d
+from physicsnemo.nn import SpectralConv1d, SpectralConv2d, SpectralConv3d
 
 
 class SpectralConv1d_old(nn.Module):