diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fd9e860ad..f94fe97bf6 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/dataset_paths.yaml b/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/dataset_paths.yaml index 4bd4e4c7e0..4bd2c8ebf8 100644 --- a/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/dataset_paths.yaml +++ b/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: ??? diff --git a/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/drivaer_ml_surface_zarr.yaml b/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/drivaer_ml_surface_zarr.yaml new file mode 100644 index 0000000000..1bdd147266 --- /dev/null +++ b/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 diff --git a/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/drivaer_ml_volume_zarr.yaml b/examples/cfd/external_aerodynamics/unified_external_aero_recipe/datasets/drivaer_ml_volume_zarr.yaml new file mode 100644 index 0000000000..0e4febf323 --- /dev/null +++ b/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 diff --git a/examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/convert_pdmsh_to_zarr.py b/examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/convert_pdmsh_to_zarr.py new file mode 100644 index 0000000000..5908fdeafe --- /dev/null +++ b/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//``), +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 --output \ + [--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 .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()