feat(examples): zarr dataset configs and .pdmsh converter for the unified external aero recipe

CHANGELOG.md

@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Zarr dataset configs for the unified external aerodynamics recipe
+  (`drivaer_ml_surface_zarr` / `drivaer_ml_volume_zarr`, reading
+  `physicsnemo-domainmesh-zarr` stores via `ZarrMeshReader` /
+  `ZarrDomainMeshReader`), plus a transitional `.pdmsh`-to-zarr converter CLI.
 - Adds a `global_shape` argument to `ShardTensor.from_local`, enabling the
   no-communication `sharding_shapes="chunk"` path.
 - Adds exact-boundary quality mesh generation to

examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/dataset_paths.yaml

@@ -21,5 +21,9 @@
 # path that is still ???.
 
 drivaer_ml: ???
+# Full-DomainMesh zarr conversion of DrivaerML (physicsnemo-domainmesh-zarr
+# schema; emitted by PhysicsNeMo-Curator's DomainMeshZarrSink, or via
+# src/convert_pdmsh_to_zarr.py for existing .pdmsh data)
+drivaer_ml_zarr: ???
 highlift_aero_ml: ???
 shift_suv: ???

examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/drivaer_ml_surface_zarr.yaml

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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.
+
+# DrivaerML surface dataset -- Zarr variant.
+# Reads the vehicle boundary out of full-DomainMesh zarr groups (the
+# physicsnemo-domainmesh-zarr schema emitted by PhysicsNeMo-Curator's
+# DomainMeshZarrSink, or by src/convert_pdmsh_to_zarr.py for existing
+# .pdmsh data). Mirrors the .pdmsh pipeline exactly: `subpath` navigates
+# to the boundary (like the .pdmsh reader's glob into
+# _tensordict/boundaries/vehicle) and `merge_global_data_from` merges the
+# case-level freestream global_data (U_inf, rho_inf, p_inf, nu, L_ref) at
+# read time -- single-sourced, nothing baked into the boundary data.
+# Boundaries are curated as a per-cell vertex soup (layout attr) so the
+# reader's contiguous block subsample maps to sequential, chunk-aligned
+# zarr reads with no cells read at all. Transform stack is identical to
+# drivaer_ml_surface.yaml.
+
+train_datadir: ${dataset_paths.drivaer_ml_zarr}
+
+pipeline:
+  reader:
+    _target_: ${dp:ZarrMeshReader}
+    path: ${train_datadir}
+    pattern: "*.zarr"
+    subpath: "boundaries/vehicle"
+    merge_global_data_from: "../../global_data"
+    subsample_n_cells: ${sampling_resolution}
+  augmentations:
+    - _target_: ${dp:RandomRotateMesh}
+      axes: ["z"]
+      transform_cell_data: true
+      transform_global_data: true
+    - _target_: ${dp:RandomTranslateMesh}
+      distribution:
+        _target_: torch.distributions.Uniform
+        low: [-1.0, -1.0, 0.0]
+        high: [1.0, 1.0, 0.0]
+  transforms:
+    - _target_: ${dp:DropMeshFields}
+      global_data: [TimeValue]
+    - _target_: ${dp:CenterMesh}
+      use_area_weighting: false
+    - _target_: ${dp:NonDimensionalizeByMetadata}
+      fields:
+        pMeanTrim: pressure
+        wallShearStressMeanTrim: stress
+      association: cell_data
+    - _target_: ${dp:RenameMeshFields}
+      cell_data:
+        pMeanTrim: pressure
+        wallShearStressMeanTrim: wss
+    - _target_: ${dp:NormalizeMeshFields}
+      association: cell_data
+      fields:
+        wss: {type: vector, mean: [0.0, 0.0, 0.0], std: 0.00313}
+    - _target_: ${dp:ComputeSurfaceNormals}
+      store_as: cell_data
+      field_name: normals
+    - _target_: ${dp:SubsampleMesh}
+      n_cells: ${sampling_resolution}
+    # Terminal transform: convert the surface Mesh into a DomainMesh per the
+    # recipe's prediction-vs-input contract (see drivaer_ml_surface.yaml).
+    - _target_: ${dp:MeshToDomainMesh}
+      interior_points: cell_centroids
+
+# Single source of truth for prediction field names + types. Read by both the
+# dataset builder (auto-injected into MeshToDomainMesh above) and the train
+# loop (used to build LossCalculator and MetricCalculator).
+targets:
+  pressure: scalar
+  wss: vector

examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/drivaer_ml_volume_zarr.yaml

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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.
+
+# DrivaerML volume dataset -- Zarr variant.
+# Reads full DomainMesh cases from physicsnemo-domainmesh-zarr groups
+# (emitted by PhysicsNeMo-Curator's DomainMeshZarrSink, or by
+# src/convert_pdmsh_to_zarr.py --extra-boundary
+# "stl_geometry=*_single_solid.stl.pmsh" for existing .pdmsh data).
+# ZarrDomainMeshReader mirrors DomainMeshReader: interior (point cloud)
+# subsampled to ${sampling_resolution} points, boundaries to
+# ${sampling_resolution} cells, and the stl_geometry boundary read at
+# full resolution for exact SDF queries (stored indexed in the group --
+# the zarr analog of DomainMeshReader's sibling-file extra_boundaries).
+# Transform stack is identical to drivaer_ml_volume.yaml.
+
+train_datadir: ${dataset_paths.drivaer_ml_zarr}
+
+pipeline:
+  reader:
+    _target_: ${dp:ZarrDomainMeshReader}
+    path: ${train_datadir}
+    pattern: "*.zarr"
+    subsample_n_points: ${sampling_resolution}
+    subsample_n_cells: ${sampling_resolution}
+    full_resolution_boundaries: [stl_geometry]
+  augmentations:
+    - _target_: ${dp:RandomRotateMesh}
+      axes: ["z"]
+      transform_point_data: true
+      transform_global_data: true
+    - _target_: ${dp:RandomTranslateMesh}
+      distribution:
+        _target_: torch.distributions.Uniform
+        low: [-1.0, -1.0, 0.0]
+        high: [1.0, 1.0, 0.0]
+  transforms:
+    - _target_: ${dp:DropMeshFields}
+      global_data: [TimeValue]
+    - _target_: ${dp:CenterMesh}
+      use_area_weighting: false
+    - _target_: ${dp:NonDimensionalizeByMetadata}
+      fields:
+        UMeanTrim: velocity
+        pMeanTrim: pressure
+        nutMeanTrim: identity
+      association: point_data
+    - _target_: ${dp:ComputeSDFFromBoundary}
+      boundary_name: stl_geometry
+      sdf_field: sdf
+      normals_field: sdf_normals
+      # Pseudo-normal sign (fed by the fast nearest-triangle kernel). The
+      # _single_solid STL is watertight, so this is exact and avoids the
+      # O(n_points * n_faces) winding-number sign pass.
+      use_winding_number: false
+    - _target_: ${dp:DropBoundary}
+      names: [stl_geometry]
+    - _target_: ${dp:RenameMeshFields}
+      point_data:
+        UMeanTrim: velocity
+        pMeanTrim: pressure
+        nutMeanTrim: nut
+    - _target_: ${dp:NormalizeMeshFields}
+      association: point_data
+      fields:
+        nut: {type: scalar, mean: 4.8e-4, std: 9.4e-4}
+    # No terminal MeshToDomainMesh: like drivaer_ml_volume.yaml, the reader
+    # already produces a DomainMesh natively (interior = volume point cloud
+    # with target fields in point_data). The recipe's collate consumes it
+    # directly via each model YAML's `forward_kwargs:` spec.
+
+targets:
+  velocity: vector
+  pressure: scalar
+  nut: scalar

examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/convert_pdmsh_to_zarr.py

@@ -0,0 +1,172 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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.
+
+"""
+Convert curated ``.pdmsh`` DomainMesh cases to the mesh-zarr schema.
+
+Thin CLI over :func:`physicsnemo.datapipes.save_domain_mesh_to_zarr`: each
+``run_*/domain_*.pdmsh`` case is written as one full-DomainMesh zarr group
+(case-level ``global_data/`` + ``interior/`` + ``boundaries/<name>/``),
+readable by ``ZarrMeshReader`` via the ``drivaer_ml_surface_zarr`` dataset
+config. Nothing is baked in: case metadata stays single-sourced at the
+group root and is merged at read time.
+
+This is a transitional tool for datasets that were already curated to
+``.pdmsh``. For new curation, prefer emitting this schema directly from
+PhysicsNeMo-Curator (``DomainMeshZarrSink``) so raw CFD data converts to
+zarr in one hop.
+
+With ``--soup-boundaries`` (recommended for DrivAerML surface training)
+boundary meshes are denormalized to a per-cell vertex soup so training-time
+block subsamples read contiguously; the verified ``layout`` attr lets the
+reader skip the ``cells`` arrays entirely.
+
+Usage:
+    python src/convert_pdmsh_to_zarr.py \
+        --input <dataset dir with run_*/> --output <zarr dataset dir> \
+        [--runs 8] [--soup-boundaries] [--chunk-cells 200000] [--no-compress]
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from physicsnemo.datapipes import save_domain_mesh_to_zarr, to_cell_soup
+from physicsnemo.mesh import DomainMesh, Mesh
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n")[1])
+    ap.add_argument(
+        "--input",
+        type=Path,
+        required=True,
+        help="dataset root containing run_*/domain_*.pdmsh",
+    )
+    ap.add_argument(
+        "--output",
+        type=Path,
+        required=True,
+        help="output directory for <run>.zarr groups",
+    )
+    ap.add_argument(
+        "--runs",
+        type=int,
+        default=None,
+        help="convert only the first N runs (numeric order)",
+    )
+    ap.add_argument(
+        "--soup-boundaries",
+        action="store_true",
+        help="denormalize boundary meshes to per-cell vertex "
+        "soup for contiguous block reads",
+    )
+    ap.add_argument(
+        "--extra-boundary",
+        action="append",
+        default=[],
+        metavar="NAME=GLOB",
+        help="add a sibling .pmsh mesh (glob relative to the run "
+        "dir) as an extra boundary, stored indexed at full "
+        "resolution and never souped -- e.g. "
+        "stl_geometry='*_single_solid.stl.pmsh' for exact "
+        "SDF queries (mirrors DomainMeshReader's "
+        "extra_boundaries)",
+    )
+    ap.add_argument(
+        "--chunk-cells",
+        type=int,
+        default=200_000,
+        help="zarr chunk length on the cell axis; align with "
+        "the reader's subsample_n_cells",
+    )
+    ap.add_argument("--no-compress", action="store_true")
+    args = ap.parse_args()
+
+    extra_specs: list[tuple[str, str]] = []
+    for spec in args.extra_boundary:
+        name, _, pattern = spec.partition("=")
+        if not name or not pattern:
+            raise SystemExit(f"--extra-boundary expects NAME=GLOB, got {spec!r}")
+        extra_specs.append((name, pattern))
+
+    def _run_key(d: Path) -> tuple:
+        # Numeric runs (run_42) sort numerically; anything else (run_ref,
+        # run_0001_baseline) sorts lexicographically after them instead of
+        # crashing the conversion.
+        tail = d.name.split("_", 1)[1] if "_" in d.name else d.name
+        return (0, int(tail), "") if tail.isdigit() else (1, 0, d.name)
+
+    run_dirs = sorted(
+        (d for d in args.input.glob("run_*") if d.is_dir()),
+        key=_run_key,
+    )
+    if args.runs is not None:
+        run_dirs = run_dirs[: args.runs]
+    if not run_dirs:
+        raise SystemExit(f"no run_* directories under {args.input}")
+
+    args.output.mkdir(parents=True, exist_ok=True)
+    for i, run in enumerate(run_dirs):
+        pdmsh_candidates = sorted(run.glob("domain_*.pdmsh"))
+        if not pdmsh_candidates:
+            print(f"  {run.name}: no domain_*.pdmsh, skipping")
+            continue
+        out_group = args.output / f"{run.name}.zarr"
+        if out_group.exists():
+            print(f"  {run.name}: {out_group.name} exists, skipping")
+            continue
+        domain = DomainMesh.load(str(pdmsh_candidates[0]))
+
+        boundaries = {
+            name: to_cell_soup(domain.boundaries[name])
+            if args.soup_boundaries
+            else domain.boundaries[name]
+            for name in domain.boundary_names
+        }
+        # Extra boundaries: full resolution, indexed (exact geometry for
+        # SDF-style queries), never souped.
+        for name, pattern in extra_specs:
+            matches = sorted(run.glob(pattern))
+            if not matches:
+                raise FileNotFoundError(
+                    f"no mesh matching {pattern!r} in {run} for extra boundary {name!r}"
+                )
+            boundaries[name] = Mesh.load(str(matches[0]))
+
+        domain = DomainMesh(
+            interior=domain.interior,
+            boundaries=boundaries,
+            global_data=domain.global_data,
+        )
+        save_domain_mesh_to_zarr(
+            domain,
+            out_group,
+            chunk_cells=args.chunk_cells,
+            chunk_points=3 * args.chunk_cells,
+            compress=not args.no_compress,
+        )
+        print(
+            f"  {run.name} -> {out_group.name} "
+            f"(interior {domain.interior.n_points:,} pts, boundaries: "
+            f"{', '.join(domain.boundary_names)}) [{i + 1}/{len(run_dirs)}]",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()