Refactor (#1208)

* Move filesystems and version_check to core

* Fix version check tests

* Reorganize distributed, domain_parallel, and begin nn / utils cleanup.

* Move modules and meta to core.  Move registry to core.

No tests fixed yet.

* Add missing init files

* Update build system and specify some deps.

* Reorganize tests.

* Update init files

* Clean up neighbor tools.

* Update testing

* Fix compat tests

* Move core model tests to tests/core/

* Add import lint config

* Relocate layers

* Move graphcast utils into model directory

* Relocating util functionalities.

* Add FIGConvNet to crash example (#1207)

* Add FIGConvNet to crash example.

* Add FIGConvNet to crash example

* Update model config

* propose fix some typos (#1209)

Signed-off-by: John E <[email protected]>
Co-authored-by: Corey adams <[email protected]>

* Further clean up and organize tests.

* utils tests are passing now

* Cleaning up distributed tests

* Patching tests working again in nn

* Fix sdf test

* Fix zenith angle tests

* Some organization of tests.  Checkpoints is moved into utils.

* Remove launch.utils and launch.config.  Checkpointing is moved to
phsyicsnemo.utils, launch.config is just gone.  It was empty.

* Most nn tests are passing

* Further cleanup.  Getting there!

* Remove constants file

* Add import linting to pre-commit.

---------

Signed-off-by: John E <[email protected]>
Co-authored-by: Alexey Kamenev <[email protected]>
Co-authored-by: John Eismeier <[email protected]>

Author: 3 people

.importlinter

@@ -0,0 +1,75 @@
+[importlinter]
+root_package = physicsnemo
+include_external_packages = True
+
+[importlinter:contract:physicsnemo-modules]
+name = Prevent Upward Imports in the PhysicsNemo Structure
+type = layers
+containers=
+    physicsnemo
+layers = 
+    experimental
+    models  : registry : datapipes : metrics : domain_parallel
+    utils
+    distributed |  nn
+    core
+
+[importlinter:contract:physicsnemo-core]
+name = Control Dependencies in PhysicsNeMo core
+type = layers
+containers=
+    physicsnemo.core
+layers = 
+    module : registry
+    meta
+    warnings | version_check | filesystem
+
+
+[importlinter:contract:physicsnemo-distributed]
+name = Control Dependencies in PhysicsNeMo distributed
+type = layers
+containers=
+    physicsnemo.distributed
+layers = 
+    fft | autograd
+    mappings
+    utils
+    manager
+    config
+
+[importlinter:contract:physicsnemo-utils]
+name = Control Dependencies in PhysicsNeMo utils
+type = layers
+containers=
+    physicsnemo.utils
+layers = 
+    mesh
+    profiling
+    checkpoint
+    capture
+    logging | memory
+
+[importlinter:contract:physicsnemo-nn]
+name = Control Dependencies in PhysicsNeMo nn
+type = layers
+containers=
+    physicsnemo.nn
+layers = 
+    fourier_layers | transformer_layers
+    dgm_layers | mlp_layers | fully_connected_layers
+    activations | attention_layers | ball_query | conv_layers  | drop | fft  | fused_silu | insolation | interpolation | kan_layers | patching | resample_layers | sdf | siren_layers | spectral_layers | transformer_decoder  | weight_fact | weight_norm | zenith_angle
+    neighbors 
+    utils
+
+
+[importlinter:contract:physicsnemo-models]
+name = Prevent Imports between physicsnemo models
+type = layers
+containers=
+    physicsnemo.models
+layers = 
+    mesh_reduced
+    afno | dlwp | dlwp_healpix | domino | dpot | fengwu | figconvnet | fno | graphcast | meshgraphnet | pangu | pix2pix | rnn |  srrn | swinvrnn | topodiff | transolver  | vfgn 
+    unet | diffusion | gnn_layers | dlwp_healpix_layers 
+    utils
+

.pre-commit-config.yaml

@@ -53,3 +53,8 @@ repos:
     hooks:
     - id: check-added-large-files
       args: [--maxkb=5000]
+
+-   repo: https://github.com/seddonym/import-linter
+    rev: v2.5.2
+    hooks:
+    - id: import-linter

FAQ.md

@@ -11,15 +11,15 @@
 ## What is the recommended hardware for training using PhysicsNeMo framework?
 
 Please refer to the recommended hardware section:
-[System Requirments](https://docs.nvidia.com/deeplearning/physicsnemo/getting-started/index.html#system-requirements)
+[System Requirements](https://docs.nvidia.com/deeplearning/physicsnemo/getting-started/index.html#system-requirements)
 
 ## What model architectures are in PhysicsNeMo?
 
 Nvidia PhysicsNeMo is built on top of PyTorch and you can build and train any model
 architecture you want in PhysicsNeMo. PhysicsNeMo however has a catalog of models that
 have been packaged in a configurable form to make it easy to retrain with new data or certain
 config parameters. Examples include GNNs like MeshGraphNet or Neural Operators like FNO.
-PhysicsNeMo samples have more models that illustrate how a specific approach with a specifc
+PhysicsNeMo samples have more models that illustrate how a specific approach with a specific
 model architecture can be applied to a specific problem.
 These are reference starting points for users to get started.
 
@@ -47,7 +47,7 @@ that illustrates the concept.
 
 ## What can I do if I dont see a PDE in PhysicsNeMo?
 
-PhysicsNeMo Symbolic provides a well documeted
+PhysicsNeMo Symbolic provides a well documented
 [example](https://docs.nvidia.com/deeplearning/physicsnemo/physicsnemo-sym/user_guide/foundational/1d_wave_equation.html#writing-custom-pdes-and-boundary-initial-conditions)
 that walks you through how to define a custom PDE. Please see the source [here](https://github.com/NVIDIA/physicsnemo-sym/tree/main/physicsnemo/sym/eq/pdes)
 to see the built-in PDE implementation as an additional reference for your own implementation.

examples/structural_mechanics/crash/conf/model/figconvunet_time_conditional.yaml

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+_target_: rollout.FIGConvUNetTimeConditionalRollout
+_convert_: all
+
+# Input/output channels
+in_channels: 2                       # thickness + time
+out_channels: 3                      # displacement offset (xyz)
+
+# Architecture
+kernel_size: 3
+hidden_channels: [16, 16, 16]      # channels at each level
+num_levels: 2                        # number of down/up levels
+num_down_blocks: 1
+num_up_blocks: 1
+mlp_channels: [256, 256]
+
+# Spatial domain
+aabb_max: [2.0, 2.0, 2.0]
+aabb_min: [-2.0, -2.0, -2.0]
+voxel_size: null
+
+# Grid resolutions (factorized implicit grids)
+# Format: Uses res_mem_pair resolver (memory_format_enum, resolution_tuple)
+resolution_memory_format_pairs:
+  - [b_xc_y_z, [2, 64, 64]]
+  - [b_yc_x_z, [64, 2, 64]]
+  - [b_zc_x_y, [64, 64, 2]]
+
+# Position encoding
+use_rel_pos: true
+use_rel_pos_embed: true
+pos_encode_dim: 16
+
+# Communication and sampling
+communication_types: ["sum"]
+to_point_sample_method: "graphconv"
+neighbor_search_type: "knn"
+knn_k: 16
+reductions: ["mean"]
+
+# Pooling (for global features if needed)
+pooling_type: "max"
+pooling_layers: [2]
+
+# Rollout parameters
+num_time_steps: ${training.num_time_steps}

examples/structural_mechanics/crash/datapipe.py

@@ -350,13 +350,6 @@ def _normalize_node_tensor(
         )
         return (invar - mu.view(1, 1, -1)) / (std.view(1, 1, -1) + EPS)
 
-    @staticmethod
-    def _normalize_thickness_tensor(
-        thickness: torch.Tensor, mu: torch.Tensor, std: torch.Tensor
-    ):
-        # thickness: [N], mu/std: scalar tensors
-        return (thickness - mu) / (std + EPS)
-
 
 class CrashGraphDataset(CrashBaseDataset):
     """

examples/structural_mechanics/crash/rollout.py

@@ -15,11 +15,11 @@
 # limitations under the License.
 
 import torch
-import torch.nn as nn
 from torch.utils.checkpoint import checkpoint as ckpt
 
 from physicsnemo.models.transolver import Transolver
 from physicsnemo.models.meshgraphnet import MeshGraphNet
+from physicsnemo.models.figconvnet.figconvunet import FIGConvUNet
 
 from datapipe import SimSample
 
@@ -406,3 +406,64 @@ def step_fn(nf, ef, g):
                 y_t0, y_t1 = y_t1, y_t2_pred
 
         return torch.stack(outputs, dim=0)  # [T,N,3]
+
+
+class FIGConvUNetTimeConditionalRollout(FIGConvUNet):
+    """
+    FIGConvUNet with time-conditional rollout for crash simulation.
+
+    Predicts each time step independently, conditioned on normalized time.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.rollout_steps: int = kwargs.pop("num_time_steps") - 1
+        super().__init__(*args, **kwargs)
+
+    def forward(
+        self,
+        sample: SimSample,
+        data_stats: dict,
+    ) -> torch.Tensor:
+        """
+        Args:
+            Sample: SimSample containing node_features and node_target
+            data_stats: dict containing normalization stats
+        Returns:
+            [T, N, 3] rollout of predicted positions
+        """
+        inputs = sample.node_features
+        x = inputs["coords"]  # initial pos [N, 3]
+        features = inputs.get("features", x.new_zeros((x.size(0), 0)))  # [N, F]
+
+        outputs: list[torch.Tensor] = []
+        time_seq = torch.linspace(0.0, 1.0, self.rollout_steps, device=x.device)
+
+        for time_t in time_seq:
+            # Prepare vertices for FIGConvUNet: [1, N, 3]
+            vertices = x.unsqueeze(0)  # [1, N, 3]
+
+            # Prepare features: features + time [N, F+1]
+            time_expanded = time_t.expand(x.size(0), 1)  # [N, 1]
+            features_t = torch.cat([features, time_expanded], dim=-1)  # [N, F+1]
+            features_t = features_t.unsqueeze(0)  # [1, N, F+1]
+
+            def step_fn(verts, feats):
+                out, _ = super(FIGConvUNetTimeConditionalRollout, self).forward(
+                    vertices=verts, features=feats
+                )
+                return out
+
+            if self.training:
+                outf = ckpt(
+                    step_fn,
+                    vertices,
+                    features_t,
+                    use_reentrant=False,
+                ).squeeze(0)  # [N, 3]
+            else:
+                outf = step_fn(vertices, features_t).squeeze(0)  # [N, 3]
+
+            y_t = x + outf
+            outputs.append(y_t)
+
+        return torch.stack(outputs, dim=0)  # [T, N, 3]

examples/structural_mechanics/crash/train.py

@@ -30,7 +30,6 @@
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm
 
 from physicsnemo.distributed.manager import DistributedManager
 from physicsnemo.launch.logging import PythonLogger, RankZeroLoggingWrapper
@@ -64,6 +63,11 @@ def __init__(self, cfg: DictConfig, logger0: RankZeroLoggingWrapper):
                 f"Model {model_name} requires a point-cloud datapipe, "
                 f"but you selected {datapipe_name}."
             )
+        if "FIGConvUNet" in model_name and "PointCloudDataset" not in datapipe_name:
+            raise ValueError(
+                f"Model {model_name} requires a point-cloud datapipe, "
+                f"but you selected {datapipe_name}."
+            )
 
         # Dataset
         reader = instantiate(cfg.reader)
@@ -223,7 +227,7 @@ def main(cfg: DictConfig) -> None:
         for sample in trainer.dataloader:
             sample = sample[0].to(dist.device)  # SimSample .to()
             loss = trainer.train(sample)
-            total_loss += loss.item()
+            total_loss += loss.detach().item()
             num_batches += 1
 
         trainer.scheduler.step()

examples/structural_mechanics/crash/vtp_reader.py

@@ -169,7 +169,7 @@ def process_vtp_data(data_dir, num_samples=2, write_vtp=False, logger=None):
 
     if not vtp_files:
         if logger:
-            logger.error("No .vtp files found in:", base_data_dir)
+            logger.error(f"No .vtp files found in: {base_data_dir}")
         exit(1)
 
     for vtp_path in vtp_files:

physicsnemo/__init__.py

@@ -13,10 +13,26 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 
-from .datapipes.datapipe import Datapipe
-from .datapipes.meta import DatapipeMetaData
-from .models.meta import ModelMetaData
-from .models.module import Module
+# Backwards-compatibility is opt-in. Enable with env var or via enable_compat().
+if os.getenv("PHYSICSNEMO_ENABLE_COMPAT") in {
+    "1",
+    "true",
+    "True",
+    "YES",
+    "yes",
+    "on",
+    "ON",
+}:
+    from .compat import install as _compat_install
+
+    _compat_install()
+
+
+# from .datapipes.datapipe import Datapipe  # noqa E402
+# from .datapipes.meta import DatapipeMetaData  # noqa E402
+# from .core.meta import ModelMetaData  # noqa E402
+# from .core.module import Module  # noqa E402
 
 __version__ = "1.3.0a0"