From 968123e49f4ba85e4ea59ca14370a98180395c2d Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Tue, 30 Jun 2026 21:33:32 -0400 Subject: [PATCH 1/3] Promote field schemas and layouts to physicsnemo.mesh physicsnemo/mesh/fields.py introduces representation-neutral RankSpecDict helpers for flattening, validation, rank counting, TensorDict inference, and declared-data checks. It also adds ScalarVectorFields and a deterministic FieldLayout that packs and unpacks named rank-0 scalars and rank-1 polar vectors without mixing their transformation laws. physicsnemo/mesh/__init__.py exports the schema and layout API from the stable mesh namespace so mesh-based models can share one field contract. physicsnemo/experimental/models/globe/field_kernel.py and model.py now import these neutral utilities from physicsnemo.mesh. The old GLOBE-specific rank_spec.py module and its utilities/__init__.py re-exports are removed, while tensordict_utils.py points schema users to the mesh module. test/mesh/test_fields.py verifies nested schemas, deterministic channel ordering, scalar/vector shape validation, round trips, empty layouts, and useful errors for missing or incorrectly ranked data. Centralizing these semantics avoids model-specific duplication and gives future mesh layers a common, explicitly representation-aware field interface. --- .../experimental/models/globe/field_kernel.py | 8 +- .../experimental/models/globe/model.py | 12 +- .../models/globe/utilities/__init__.py | 10 - .../models/globe/utilities/rank_spec.py | 232 --------- .../globe/utilities/tensordict_utils.py | 3 +- physicsnemo/mesh/__init__.py | 25 + physicsnemo/mesh/fields.py | 469 ++++++++++++++++++ test/mesh/test_fields.py | 190 +++++++ 8 files changed, 693 insertions(+), 256 deletions(-) delete mode 100644 physicsnemo/experimental/models/globe/utilities/rank_spec.py create mode 100644 physicsnemo/mesh/fields.py create mode 100644 test/mesh/test_fields.py diff --git a/physicsnemo/experimental/models/globe/field_kernel.py b/physicsnemo/experimental/models/globe/field_kernel.py index 13641d598a..cfb256187d 100644 --- a/physicsnemo/experimental/models/globe/field_kernel.py +++ b/physicsnemo/experimental/models/globe/field_kernel.py @@ -30,16 +30,12 @@ from torch.utils.checkpoint import checkpoint from physicsnemo.core.module import Module -from physicsnemo.experimental.models.globe.utilities.rank_spec import ( - RankSpecDict, - flatten_rank_spec, - rank_counts, -) from physicsnemo.experimental.models.globe.utilities.tensordict_utils import ( concatenate_leaves, concatenated_length, split_by_leaf_rank, ) +from physicsnemo.mesh import RankSpecDict, flatten_rank_spec, rank_counts from physicsnemo.nn import Mlp, Pade from physicsnemo.nn.functional.equivariant_ops import ( legendre_polynomials, @@ -919,12 +915,12 @@ def forward( TensorDict[str, Float[torch.Tensor, "n_targets ..."]] Kernel output fields at target points. """ + from physicsnemo.mesh.spatial._ragged import _ragged_arange from physicsnemo.mesh.spatial.cluster_tree import ( ClusterTree, DualInteractionPlan, SourceAggregates, ) - from physicsnemo.mesh.spatial._ragged import _ragged_arange n_sources = source_points.shape[0] n_targets = target_points.shape[0] diff --git a/physicsnemo/experimental/models/globe/model.py b/physicsnemo/experimental/models/globe/model.py index b079c27019..bb6898cf1a 100644 --- a/physicsnemo/experimental/models/globe/model.py +++ b/physicsnemo/experimental/models/globe/model.py @@ -27,17 +27,17 @@ from physicsnemo.core.meta import ModelMetaData from physicsnemo.core.module import Module -from physicsnemo.mesh.spatial.cluster_tree import ( - ClusterTree, - DualInteractionPlan, -) from physicsnemo.experimental.models.globe.field_kernel import MultiscaleKernel -from physicsnemo.experimental.models.globe.utilities.rank_spec import ( +from physicsnemo.mesh import ( + Mesh, RankSpecDict, flatten_rank_spec, validate_data_contains_ranks, ) -from physicsnemo.mesh import Mesh +from physicsnemo.mesh.spatial.cluster_tree import ( + ClusterTree, + DualInteractionPlan, +) from physicsnemo.utils.logging import PythonLogger # allow_in_graph wraps these TensorDict methods as opaque graph nodes so that diff --git a/physicsnemo/experimental/models/globe/utilities/__init__.py b/physicsnemo/experimental/models/globe/utilities/__init__.py index bfc651d344..a850502265 100644 --- a/physicsnemo/experimental/models/globe/utilities/__init__.py +++ b/physicsnemo/experimental/models/globe/utilities/__init__.py @@ -14,12 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from physicsnemo.experimental.models.globe.utilities.rank_spec import ( - RankSpecDict, - flatten_rank_spec, - rank_counts, - ranks_from_tensordict, -) from physicsnemo.experimental.models.globe.utilities.tensordict_utils import ( combine_tensordicts, concatenate_leaves, @@ -28,12 +22,8 @@ ) __all__ = [ - "RankSpecDict", "combine_tensordicts", "concatenate_leaves", "concatenated_length", - "flatten_rank_spec", - "rank_counts", - "ranks_from_tensordict", "split_by_leaf_rank", ] diff --git a/physicsnemo/experimental/models/globe/utilities/rank_spec.py b/physicsnemo/experimental/models/globe/utilities/rank_spec.py deleted file mode 100644 index 0fa581c7ff..0000000000 --- a/physicsnemo/experimental/models/globe/utilities/rank_spec.py +++ /dev/null @@ -1,232 +0,0 @@ -# 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. - -r"""Rank specification types and utilities for GLOBE field kernels. - -A rank spec describes the tensor rank (0 = scalar, 1 = vector) of each -field in a kernel's input or output. It is a plain Python dict - not a -TensorDict - so that ``torch.compile`` can specialize on it as a -compile-time constant without graph breaks. - -Rank specs may be flat:: - - {"pressure": 0, "velocity": 1} - -or nested to mirror a hierarchical TensorDict structure:: - - {"group": {"field_a": 0, "field_b": 1}} - -Use :func:`flatten_rank_spec` to reduce a nested spec to dot-separated -flat keys, and :func:`ranks_from_tensordict` to derive a rank spec from -an existing data TensorDict. - -For *runtime* grouping of actual tensor data by observed rank, see -:func:`~physicsnemo.experimental.models.globe.utilities.tensordict_utils.split_by_leaf_rank`. -""" - -from collections import Counter -from typing import TypeAlias, Union - -from tensordict import TensorDict - -### Type definition -# TODO: replace with `type RankSpecDict = dict[str, int | RankSpecDict]` -# once Python 3.11 support is dropped (PEP 695). -RankSpecDict: TypeAlias = dict[str, Union[int, "RankSpecDict"]] - - -def flatten_rank_spec( - rank_spec: RankSpecDict, sep: str = ".", -) -> dict[str, int]: - r"""Flatten a possibly-nested rank spec to dot-separated string keys. - - Parameters - ---------- - rank_spec : RankSpecDict - Rank spec mapping field names to integer ranks or nested sub-specs. - sep : str - Separator for joining nested key segments. Defaults to ``"."``. - - Returns - ------- - dict[str, int] - Flat mapping from dot-separated field name to integer rank. - - Examples - -------- - >>> flatten_rank_spec({"pressure": 0, "velocity": 1}) - {'pressure': 0, 'velocity': 1} - - >>> flatten_rank_spec({"group": {"a": 0, "b": 1}}) - {'group.a': 0, 'group.b': 1} - """ - result: dict[str, int] = {} - for k, v in rank_spec.items(): - if isinstance(v, dict): - for sub_k, sub_v in flatten_rank_spec(v, sep).items(): - result[f"{k}{sep}{sub_k}"] = sub_v - else: - result[k] = v - return result - - -def rank_counts(rank_spec: RankSpecDict) -> Counter[int]: - r"""Count leaves by rank value in a rank spec. - - Parameters - ---------- - rank_spec : RankSpecDict - Rank spec mapping field names to integer ranks. - - Returns - ------- - Counter[int] - Mapping from rank value to the number of fields with that rank. - Missing ranks default to 0. - - Examples - -------- - >>> rank_counts({"a": 0, "b": 0, "c": 1}) - Counter({0: 2, 1: 1}) - """ - return Counter(flatten_rank_spec(rank_spec).values()) - - -def ranks_from_tensordict(td: TensorDict) -> RankSpecDict: - r"""Derive a :class:`RankSpecDict` from a data TensorDict's leaf shapes. - - Each leaf tensor is replaced by its rank (number of non-batch dimensions), - producing a plain dict whose structure mirrors the TensorDict nesting. - - Parameters - ---------- - td : TensorDict - Data TensorDict whose leaf ranks should be extracted. - - Returns - ------- - RankSpecDict - Rank spec with integer leaves. - - Examples - -------- - >>> import torch - >>> from tensordict import TensorDict - >>> td = TensorDict({ - ... "pressure": torch.randn(10), - ... "velocity": torch.randn(10, 3), - ... }, batch_size=[10]) - >>> ranks_from_tensordict(td) - {'pressure': 0, 'velocity': 1} - """ - result: RankSpecDict = {} - for k in td.keys(): - v = td[k] - if isinstance(v, TensorDict): - result[k] = ranks_from_tensordict(v) # ty: ignore[invalid-assignment] - else: - result[k] = v.ndim - td.batch_dims # ty: ignore[invalid-assignment] - return result - - -def validate_data_contains_ranks( - *, - data: TensorDict, - declared_ranks: RankSpecDict, - source_label: str, -) -> None: - r"""Raise :class:`ValueError` if ``data`` doesn't contain every leaf in ``declared_ranks``. - - Subset check: every ``(key, rank)`` pair in ``declared_ranks`` must - appear (with the same rank) in :func:`ranks_from_tensordict` of - ``data``. Extra leaves in ``data`` are allowed and silently ignored - by this validator -- callers that don't want extras to flow further - should ``select`` them out separately. - - Differences are reported one leaf per line, classified as missing or - rank-mismatch. - - Used by models that filter their incoming data down to declared keys - via :meth:`tensordict.TensorDict.select`: this validator catches - typos and shape regressions in user declarations with a friendlier - error than the bare :class:`KeyError` ``select(strict=True)`` would - raise on a missing leaf. - - Parameters - ---------- - data : TensorDict - Data TensorDict to validate. - declared_ranks : RankSpecDict - Rank spec the caller expects ``data`` to be a (possibly improper) - superset of. - source_label : str - Human-readable description of the ``data`` argument, used verbatim - in the error message (e.g. ``"`mesh.cell_data`"`` or - ``"`global_data`"``). - - Raises - ------ - ValueError - If any declared leaf is missing from ``data`` or has a different - rank than declared. - - Examples - -------- - Matching cases (and supersets) are silent: - - >>> import torch - >>> from tensordict import TensorDict - >>> td = TensorDict( - ... {"pressure": torch.zeros(10), "extra": torch.zeros(10)}, - ... batch_size=[10], - ... ) - >>> validate_data_contains_ranks( - ... data=td, declared_ranks={"pressure": 0}, source_label="`td`" - ... ) - - Missing leaves and rank mismatches raise with a per-leaf diff: - - >>> bad = TensorDict({"pressure": torch.zeros(10, 3)}, batch_size=[10]) - >>> try: - ... validate_data_contains_ranks( - ... data=bad, - ... declared_ranks={"pressure": 0, "velocity": 1}, - ... source_label="`bad`", - ... ) - ... except ValueError as e: - ... print(e) - `bad` does not contain its declared rank spec: - - missing leaf 'velocity' (declared rank 1) - - rank mismatch for 'pressure': declared 0, got 1 - """ - declared = flatten_rank_spec(declared_ranks) - actual = flatten_rank_spec(ranks_from_tensordict(data)) - - ### Single pass over declared leaves: missing, then rank mismatches. - lines: list[str] = [] - for k in sorted(declared.keys() - actual.keys()): - lines.append(f" - missing leaf {k!r} (declared rank {declared[k]})") - for k in sorted(declared.keys() & actual.keys()): - if declared[k] != actual[k]: - lines.append( - f" - rank mismatch for {k!r}: declared {declared[k]}, " - f"got {actual[k]}" - ) - if lines: - raise ValueError( - f"{source_label} does not contain its declared rank spec:\n" - + "\n".join(lines) - ) diff --git a/physicsnemo/experimental/models/globe/utilities/tensordict_utils.py b/physicsnemo/experimental/models/globe/utilities/tensordict_utils.py index 8e939ed975..ac3e78fdd6 100644 --- a/physicsnemo/experimental/models/globe/utilities/tensordict_utils.py +++ b/physicsnemo/experimental/models/globe/utilities/tensordict_utils.py @@ -21,8 +21,7 @@ splitting by tensor rank. For field-to-rank *schema metadata* (e.g. ``{"pressure": 0, "velocity": 1}``), -see :mod:`~physicsnemo.experimental.models.globe.utilities.rank_spec` and its -:class:`RankSpecDict` type. +see :mod:`physicsnemo.mesh.fields` and its :class:`RankSpecDict` type. """ from math import prod diff --git a/physicsnemo/mesh/__init__.py b/physicsnemo/mesh/__init__.py index 408dc5b035..fb6fcf3cc9 100644 --- a/physicsnemo/mesh/__init__.py +++ b/physicsnemo/mesh/__init__.py @@ -15,4 +15,29 @@ # limitations under the License. from physicsnemo.mesh.domain_mesh import DomainMesh +from physicsnemo.mesh.fields import ( + FieldLayout, + RankSpecDict, + ScalarVectorFields, + flatten_rank_spec, + rank_counts, + ranks_from_tensordict, + validate_data_contains_ranks, + validate_rank_spec, +) from physicsnemo.mesh.mesh import MESH_FIELD_ASSOCIATIONS, Mesh, MeshFieldAssociation + +__all__ = [ + "DomainMesh", + "FieldLayout", + "MESH_FIELD_ASSOCIATIONS", + "Mesh", + "MeshFieldAssociation", + "RankSpecDict", + "ScalarVectorFields", + "flatten_rank_spec", + "rank_counts", + "ranks_from_tensordict", + "validate_data_contains_ranks", + "validate_rank_spec", +] diff --git a/physicsnemo/mesh/fields.py b/physicsnemo/mesh/fields.py new file mode 100644 index 0000000000..a299f5c7e4 --- /dev/null +++ b/physicsnemo/mesh/fields.py @@ -0,0 +1,469 @@ +# 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. + +r"""Semantic scalar and vector field layout utilities. + +Mesh models commonly need to turn named physical fields into dense channel +tensors. A plain concatenation is not sufficient: a scalar field and a +Cartesian vector field transform differently under a change of frame. This +module keeps those two representations separate while giving them a stable, +name-based channel order. + +``RankSpecDict`` describes the semantic tensor rank of every named leaf. The +currently supported :class:`FieldLayout` leaves are + +* rank 0: one scalar per point, with shape ``(N,)``; +* rank 1: one polar vector per point, with shape ``(N, D)``. + +Consequently, a rank-1 leaf is not an arbitrary feature axis of length ``D``: +it is explicitly promised to transform as a polar vector. Axial vectors and +higher-order tensors require additional representation metadata and are not +silently treated as rank-1 fields here. +""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Collection +from copy import deepcopy +from typing import NamedTuple, TypeAlias, Union + +import torch +from tensordict import TensorDict + +# TODO: replace with ``type RankSpecDict = ...`` after Python 3.11 support is +# dropped (PEP 695). +RankSpecDict: TypeAlias = dict[str, Union[int, "RankSpecDict"]] + + +def flatten_rank_spec(rank_spec: RankSpecDict, sep: str = ".") -> dict[str, int]: + r"""Flatten a nested rank specification to separator-joined names. + + The insertion order of ``rank_spec`` is retained for compatibility. Code + that assigns positional channels should sort the returned names, as + :class:`FieldLayout` does. + + Parameters + ---------- + rank_spec : RankSpecDict + Mapping from field names to integer ranks or nested mappings. + sep : str, default="." + Separator used to join nested path components. + + Returns + ------- + dict[str, int] + Flat field-name to semantic-rank mapping. + + Examples + -------- + >>> flatten_rank_spec({"pressure": 0, "velocity": 1}) + {'pressure': 0, 'velocity': 1} + >>> flatten_rank_spec({"fluid": {"pressure": 0, "velocity": 1}}) + {'fluid.pressure': 0, 'fluid.velocity': 1} + """ + result: dict[str, int] = {} + for key, value in rank_spec.items(): + if isinstance(value, dict): + for sub_key, rank in flatten_rank_spec(value, sep=sep).items(): + result[f"{key}{sep}{sub_key}"] = rank + else: + result[key] = value + return result + + +def validate_rank_spec( + rank_spec: RankSpecDict, + *, + allowed_ranks: Collection[int] | None = None, + source_label: str = "rank_spec", +) -> None: + r"""Validate the structure and integer leaves of a rank specification. + + Parameters + ---------- + rank_spec : RankSpecDict + Rank specification to validate. + allowed_ranks : Collection[int] or None, optional + If supplied, every leaf must belong to this collection. Otherwise all + non-negative integer ranks are accepted. + source_label : str, default="rank_spec" + Human-readable name used in error messages. + + Raises + ------ + TypeError + If a key is not a string or a leaf is not an integer. + ValueError + If a rank is negative or is not in ``allowed_ranks``. + """ + allowed = None if allowed_ranks is None else frozenset(allowed_ranks) + + def _validate(spec: RankSpecDict, path: tuple[str, ...]) -> None: + for key, value in spec.items(): + if not isinstance(key, str): + raise TypeError( + f"All keys in {source_label} must be strings; got " + f"{key!r} at {'.'.join(path) or ''}" + ) + leaf_path = (*path, key) + if isinstance(value, dict): + _validate(value, leaf_path) + continue + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError( + f"Rank for {'.'.join(leaf_path)!r} in {source_label} must " + f"be an integer, got {value!r}" + ) + if value < 0: + raise ValueError( + f"Rank for {'.'.join(leaf_path)!r} in {source_label} must " + f"be non-negative, got {value}" + ) + if allowed is not None and value not in allowed: + raise ValueError( + f"Rank for {'.'.join(leaf_path)!r} in {source_label} must " + f"be one of {sorted(allowed)}, got {value}" + ) + + if not isinstance(rank_spec, dict): + raise TypeError( + f"{source_label} must be a dict, got {type(rank_spec).__name__}" + ) + _validate(rank_spec, ()) + + +def rank_counts(rank_spec: RankSpecDict) -> Counter[int]: + r"""Count leaves of each semantic rank in ``rank_spec``.""" + return Counter(flatten_rank_spec(rank_spec).values()) + + +def ranks_from_tensordict(td: TensorDict) -> RankSpecDict: + r"""Derive semantic-rank-shaped metadata from TensorDict leaf shapes. + + A leaf rank is its number of non-batch dimensions. For a point-field + TensorDict with batch size ``(N,)``, ``(N,)`` is therefore rank 0 and + ``(N, D)`` is rank 1. + """ + result: RankSpecDict = {} + for key in td.keys(): + value = td[key] + if isinstance(value, TensorDict): + result[key] = ranks_from_tensordict(value) # ty: ignore[invalid-assignment] + else: + result[key] = value.ndim - td.batch_dims # ty: ignore[invalid-assignment] + return result + + +def validate_data_contains_ranks( + *, + data: TensorDict, + declared_ranks: RankSpecDict, + source_label: str, +) -> None: + r"""Validate that ``data`` contains every declared leaf at its stated rank. + + Additional leaves in ``data`` are allowed. Missing leaves and rank + mismatches are reported together to make schema errors easier to diagnose. + """ + validate_rank_spec(declared_ranks, source_label="declared_ranks") + declared = flatten_rank_spec(declared_ranks) + actual = flatten_rank_spec(ranks_from_tensordict(data)) + + lines = [ + f" - missing leaf {key!r} (declared rank {declared[key]})" + for key in sorted(declared.keys() - actual.keys()) + ] + lines.extend( + [ + f" - rank mismatch for {key!r}: declared {declared[key]}, " + f"got {actual[key]}" + for key in sorted(declared.keys() & actual.keys()) + if declared[key] != actual[key] + ] + ) + if lines: + raise ValueError( + f"{source_label} does not contain its declared rank spec:\n" + + "\n".join(lines) + ) + + +class ScalarVectorFields(NamedTuple): + r"""Dense scalar and polar-vector channels for a collection of points. + + ``scalars`` has shape ``(N, C_s)`` and ``vectors`` has shape + ``(N, C_v, D)``. Shape validation belongs to the :class:`FieldLayout` + that defines ``C_s``, ``C_v``, and ``D``. + """ + + scalars: torch.Tensor + vectors: torch.Tensor + + +def _rank_entries( + rank_spec: RankSpecDict, + path: tuple[str, ...] = (), +) -> list[tuple[tuple[str, ...], int]]: + """Return rank leaves as unambiguous key paths.""" + entries: list[tuple[tuple[str, ...], int]] = [] + for key, value in rank_spec.items(): + leaf_path = (*path, key) + if isinstance(value, dict): + entries.extend(_rank_entries(value, leaf_path)) + else: + entries.append((leaf_path, value)) + return entries + + +def _td_key(path: tuple[str, ...]) -> str | tuple[str, ...]: + """Use a string for top-level TensorDict keys and tuples for nested keys.""" + return path[0] if len(path) == 1 else path + + +class FieldLayout: + r"""Deterministic packing layout for named scalar and polar-vector fields. + + Channel order is lexicographic by the flattened field name and therefore + does not depend on dictionary or TensorDict construction order. Scalars + and vectors are packed separately so that equivariant code cannot + accidentally mix their transformation laws. + + Parameters + ---------- + rank_spec : RankSpecDict + Nested or flat field schema. Only rank-0 and rank-1 leaves are + supported. + spatial_dim : int + Cartesian vector dimension ``D``. + sep : str, default="." + Separator used only for public flattened names and sorting. + + Notes + ----- + Each rank-0 leaf must have shape ``(N,)`` and each rank-1 leaf must have + shape ``(N, D)``. Multiple channels must be represented by multiple named + leaves rather than an extra unnamed feature dimension. + """ + + def __init__( + self, + rank_spec: RankSpecDict, + spatial_dim: int, + *, + sep: str = ".", + ) -> None: + validate_rank_spec( + rank_spec, + allowed_ranks=(0, 1), + source_label="rank_spec", + ) + if isinstance(spatial_dim, bool) or not isinstance(spatial_dim, int): + raise TypeError(f"spatial_dim must be an integer, got {spatial_dim!r}") + if spatial_dim < 1: + raise ValueError(f"spatial_dim must be positive, got {spatial_dim}") + if not isinstance(sep, str): + raise TypeError(f"sep must be a string, got {type(sep).__name__}") + if not sep: + raise ValueError("sep must not be empty") + + entries = [ + (sep.join(path), path, rank) for path, rank in _rank_entries(rank_spec) + ] + if not entries: + raise ValueError("rank_spec must contain at least one field leaf") + + # A literal key containing ``sep`` can otherwise collide with a nested + # path in public names and make deterministic ordering ambiguous. + names = [name for name, _, _ in entries] + if len(names) != len(set(names)): + raise ValueError( + f"rank_spec contains field paths that collide when flattened " + f"with separator {sep!r}" + ) + + entries.sort(key=lambda entry: entry[0]) + self._rank_spec = deepcopy(rank_spec) + self.spatial_dim = spatial_dim + self.sep = sep + self._entries = tuple(entries) + self._scalar_entries = tuple(entry for entry in entries if entry[2] == 0) + self._vector_entries = tuple(entry for entry in entries if entry[2] == 1) + + @property + def rank_spec(self) -> RankSpecDict: + """A defensive copy of the named-field schema.""" + return deepcopy(self._rank_spec) + + @property + def flat_rank_spec(self) -> dict[str, int]: + """Field ranks in deterministic packed-name order.""" + return {name: rank for name, _, rank in self._entries} + + @property + def scalar_names(self) -> tuple[str, ...]: + """Names of scalar channels in packed order.""" + return tuple(name for name, _, _ in self._scalar_entries) + + @property + def vector_names(self) -> tuple[str, ...]: + """Names of polar-vector channels in packed order.""" + return tuple(name for name, _, _ in self._vector_entries) + + @property + def n_scalars(self) -> int: + """Number of packed scalar channels.""" + return len(self._scalar_entries) + + @property + def n_vectors(self) -> int: + """Number of packed polar-vector channels.""" + return len(self._vector_entries) + + def pack(self, data: TensorDict) -> ScalarVectorFields: + r"""Pack named point fields into deterministic scalar/vector channels. + + Extra leaves in ``data`` are ignored. All selected leaves must share + dtype and device, in addition to satisfying the semantic shapes in this + layout. + """ + if data.batch_dims != 1: + raise ValueError( + "FieldLayout expects a point TensorDict with one batch dimension " + f"(N,), got batch_size={tuple(data.batch_size)}" + ) + validate_data_contains_ranks( + data=data, + declared_ranks=self._rank_spec, + source_label="data", + ) + + n_points = data.batch_size[0] + selected: list[torch.Tensor] = [] + scalar_tensors: list[torch.Tensor] = [] + vector_tensors: list[torch.Tensor] = [] + + for name, path, rank in self._entries: + tensor = data[_td_key(path)] + expected = (n_points,) if rank == 0 else (n_points, self.spatial_dim) + if tuple(tensor.shape) != expected: + kind = "scalar" if rank == 0 else "polar vector" + raise ValueError( + f"Field {name!r} is declared as a {kind} and must have " + f"shape {expected}, got {tuple(tensor.shape)}" + ) + selected.append(tensor) + if rank == 0: + scalar_tensors.append(tensor) + else: + vector_tensors.append(tensor) + + reference = selected[0] + for name, tensor in zip( + (name for name, _, _ in self._entries), selected, strict=True + ): + if tensor.device != reference.device: + raise ValueError( + f"All packed fields must share a device; field {name!r} is " + f"on {tensor.device}, expected {reference.device}" + ) + if tensor.dtype != reference.dtype: + raise ValueError( + f"All packed fields must share a dtype; field {name!r} has " + f"{tensor.dtype}, expected {reference.dtype}" + ) + + scalars = ( + torch.stack(scalar_tensors, dim=-1) + if scalar_tensors + else reference.new_empty((n_points, 0)) + ) + vectors = ( + torch.stack(vector_tensors, dim=-2) + if vector_tensors + else reference.new_empty((n_points, 0, self.spatial_dim)) + ) + return ScalarVectorFields(scalars=scalars, vectors=vectors) + + def unpack(self, fields: ScalarVectorFields) -> TensorDict: + r"""Unpack scalar/vector channels to a TensorDict matching ``rank_spec``.""" + scalars, vectors = fields + if scalars.ndim != 2: + raise ValueError( + f"scalars must have shape (N, {self.n_scalars}), got " + f"{tuple(scalars.shape)}" + ) + if vectors.ndim != 3: + raise ValueError( + f"vectors must have shape (N, {self.n_vectors}, " + f"{self.spatial_dim}), got {tuple(vectors.shape)}" + ) + + n_points = scalars.shape[0] + expected_scalars = (n_points, self.n_scalars) + expected_vectors = (n_points, self.n_vectors, self.spatial_dim) + if tuple(scalars.shape) != expected_scalars: + raise ValueError( + f"scalars must have shape {expected_scalars}, got " + f"{tuple(scalars.shape)}" + ) + if tuple(vectors.shape) != expected_vectors: + raise ValueError( + f"vectors must have shape {expected_vectors}, got " + f"{tuple(vectors.shape)}" + ) + if vectors.shape[0] != n_points: + raise ValueError( + "scalars and vectors must contain the same number of points" + ) + if vectors.device != scalars.device: + raise ValueError("scalars and vectors must be on the same device") + if vectors.dtype != scalars.dtype: + raise ValueError("scalars and vectors must have the same dtype") + + nested: dict[str, object] = {} + + def _assign(path: tuple[str, ...], tensor: torch.Tensor) -> None: + target = nested + for key in path[:-1]: + child = target.setdefault(key, {}) + # The schema validation guarantees that a path component is + # either always a group or always a leaf. + target = child # type: ignore[assignment] + target[path[-1]] = tensor + + for index, (_, path, _) in enumerate(self._scalar_entries): + _assign(path, scalars[:, index]) + for index, (_, path, _) in enumerate(self._vector_entries): + _assign(path, vectors[:, index, :]) + + return TensorDict( + nested, + batch_size=torch.Size([n_points]), + device=scalars.device, + ) + + +__all__ = [ + "FieldLayout", + "RankSpecDict", + "ScalarVectorFields", + "flatten_rank_spec", + "rank_counts", + "ranks_from_tensordict", + "validate_data_contains_ranks", + "validate_rank_spec", +] diff --git a/test/mesh/test_fields.py b/test/mesh/test_fields.py new file mode 100644 index 0000000000..48a8e21c1d --- /dev/null +++ b/test/mesh/test_fields.py @@ -0,0 +1,190 @@ +# 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. + +import pytest +import torch +from tensordict import TensorDict + +from physicsnemo.mesh import FieldLayout, ScalarVectorFields +from physicsnemo.mesh.fields import ( + flatten_rank_spec, + rank_counts, + ranks_from_tensordict, + validate_data_contains_ranks, + validate_rank_spec, +) + + +def _mixed_fields(n: int = 4) -> TensorDict: + return TensorDict( + { + # Deliberately not in lexicographic order. + "wind": torch.arange(n * 3, dtype=torch.float32).reshape(n, 3), + "state": {"temperature": torch.arange(n, dtype=torch.float32) + 20}, + "displacement": -torch.arange(n * 3, dtype=torch.float32).reshape(n, 3), + "pressure": torch.arange(n, dtype=torch.float32) + 100, + "ignored": torch.ones(n), + }, + batch_size=[n], + ) + + +def test_rank_spec_helpers(): + ranks = {"wind": 1, "state": {"temperature": 0}, "pressure": 0} + assert flatten_rank_spec(ranks) == { + "wind": 1, + "state.temperature": 0, + "pressure": 0, + } + assert rank_counts(ranks) == {0: 2, 1: 1} + assert ranks_from_tensordict(_mixed_fields()) == { + "wind": 1, + "state": {"temperature": 0}, + "displacement": 1, + "pressure": 0, + "ignored": 0, + } + +def test_validate_rank_spec_and_data_reports_schema_errors(): + validate_rank_spec({"scalar": 0, "tensor": 2}) + with pytest.raises(ValueError, match="must be one of .* got 2"): + validate_rank_spec({"tensor": 2}, allowed_ranks=(0, 1)) + with pytest.raises(TypeError, match="must be an integer"): + validate_rank_spec({"bad": True}) + with pytest.raises(ValueError, match="must be non-negative"): + validate_rank_spec({"bad": -1}) + + data = TensorDict({"pressure": torch.zeros(3, 2)}, batch_size=[3]) + with pytest.raises(ValueError) as error: + validate_data_contains_ranks( + data=data, + declared_ranks={"pressure": 0, "velocity": 1}, + source_label="boundary data", + ) + assert str(error.value) == ( + "boundary data does not contain its declared rank spec:\n" + " - missing leaf 'velocity' (declared rank 1)\n" + " - rank mismatch for 'pressure': declared 0, got 1" + ) + + +def test_field_layout_pack_is_deterministic_and_round_trips(): + ranks = { + "wind": 1, + "state": {"temperature": 0}, + "pressure": 0, + "displacement": 1, + } + reordered_ranks = { + "pressure": 0, + "displacement": 1, + "state": {"temperature": 0}, + "wind": 1, + } + data = _mixed_fields() + layout = FieldLayout(ranks, spatial_dim=3) + reordered_layout = FieldLayout(reordered_ranks, spatial_dim=3) + + assert layout.scalar_names == ("pressure", "state.temperature") + assert layout.vector_names == ("displacement", "wind") + assert layout.flat_rank_spec == { + "displacement": 1, + "pressure": 0, + "state.temperature": 0, + "wind": 1, + } + + packed = layout.pack(data) + reordered = reordered_layout.pack(data) + torch.testing.assert_close(packed.scalars, reordered.scalars) + torch.testing.assert_close(packed.vectors, reordered.vectors) + torch.testing.assert_close( + packed.scalars, + torch.stack((data["pressure"], data["state", "temperature"]), dim=-1), + ) + torch.testing.assert_close( + packed.vectors, + torch.stack((data["displacement"], data["wind"]), dim=-2), + ) + + unpacked = layout.unpack(packed) + assert set(unpacked.keys()) == {"displacement", "pressure", "state", "wind"} + for key in ("pressure", ("state", "temperature"), "displacement", "wind"): + torch.testing.assert_close(unpacked[key], data[key]) + + +def test_field_layout_preserves_polar_vector_transformation(): + data = _mixed_fields() + layout = FieldLayout( + {"pressure": 0, "wind": 1, "displacement": 1}, + spatial_dim=3, + ) + packed = layout.pack(data) + rotation = torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + rotated_data = data.clone() + rotated_data["wind"] = data["wind"] @ rotation.T + rotated_data["displacement"] = data["displacement"] @ rotation.T + rotated = layout.pack(rotated_data) + + torch.testing.assert_close(rotated.scalars, packed.scalars) + torch.testing.assert_close(rotated.vectors, packed.vectors @ rotation.T) + + +@pytest.mark.parametrize( + ("ranks", "expected_scalar_shape", "expected_vector_shape"), + [ + ({"pressure": 0}, (4, 1), (4, 0, 3)), + ({"wind": 1}, (4, 0), (4, 1, 3)), + ], +) +def test_field_layout_supports_scalar_or_vector_only( + ranks, expected_scalar_shape, expected_vector_shape +): + layout = FieldLayout(ranks, spatial_dim=3) + packed = layout.pack(_mixed_fields()) + assert packed.scalars.shape == expected_scalar_shape + assert packed.vectors.shape == expected_vector_shape + round_trip = layout.pack(layout.unpack(packed)) + torch.testing.assert_close(round_trip.scalars, packed.scalars) + torch.testing.assert_close(round_trip.vectors, packed.vectors) + + +def test_field_layout_rejects_non_vector_rank_and_wrong_vector_dimension(): + with pytest.raises(ValueError, match="must be one of .* got 2"): + FieldLayout({"stress": 2}, spatial_dim=3) + + layout = FieldLayout({"velocity": 1}, spatial_dim=3) + wrong_dim = TensorDict({"velocity": torch.zeros(5, 2)}, batch_size=[5]) + with pytest.raises(ValueError, match=r"must have shape \(5, 3\)"): + layout.pack(wrong_dim) + + +def test_field_layout_validates_packed_shapes_and_dtype(): + layout = FieldLayout({"pressure": 0, "velocity": 1}, spatial_dim=3) + with pytest.raises(ValueError, match="vectors must have shape"): + layout.unpack( + ScalarVectorFields( + scalars=torch.empty(5, 1), + vectors=torch.empty(4, 1, 3), + ) + ) + with pytest.raises(ValueError, match="must have the same dtype"): + layout.unpack( + ScalarVectorFields( + scalars=torch.empty(5, 1, dtype=torch.float32), + vectors=torch.empty(5, 1, 3, dtype=torch.float64), + ) + ) From e3a91e2692953f1bd59b078cdf7b073586a1456d Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Tue, 30 Jun 2026 22:23:51 -0400 Subject: [PATCH 2/3] Clean up FieldLayout shape validation Remove a redundant point-count branch from FieldLayout.unpack(): the preceding full vector-shape check already enforces the same leading dimension. Restore the required separation between top-level field-schema tests. --- physicsnemo/mesh/fields.py | 4 ---- test/mesh/test_fields.py | 1 + 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/physicsnemo/mesh/fields.py b/physicsnemo/mesh/fields.py index a299f5c7e4..a471526df2 100644 --- a/physicsnemo/mesh/fields.py +++ b/physicsnemo/mesh/fields.py @@ -425,10 +425,6 @@ def unpack(self, fields: ScalarVectorFields) -> TensorDict: f"vectors must have shape {expected_vectors}, got " f"{tuple(vectors.shape)}" ) - if vectors.shape[0] != n_points: - raise ValueError( - "scalars and vectors must contain the same number of points" - ) if vectors.device != scalars.device: raise ValueError("scalars and vectors must be on the same device") if vectors.dtype != scalars.dtype: diff --git a/test/mesh/test_fields.py b/test/mesh/test_fields.py index 48a8e21c1d..e5d82988e1 100644 --- a/test/mesh/test_fields.py +++ b/test/mesh/test_fields.py @@ -58,6 +58,7 @@ def test_rank_spec_helpers(): "ignored": 0, } + def test_validate_rank_spec_and_data_reports_schema_errors(): validate_rank_spec({"scalar": 0, "tensor": 2}) with pytest.raises(ValueError, match="must be one of .* got 2"): From 2e66c4416d34c7af74c19903fadf302adf43366d Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Tue, 30 Jun 2026 23:23:01 -0400 Subject: [PATCH 3/3] Limit mesh fields to migrated rank schemas Remove FieldLayout, ScalarVectorFields, validate_rank_spec, and their supporting tests and exports because they have no consumers on main. Keep this branch as a pure relocation of the RankSpecDict helpers already used by GLOBE; representation-aware packing can be introduced with the feature that needs it. --- physicsnemo/mesh/__init__.py | 6 - physicsnemo/mesh/fields.py | 367 +++-------------------------------- test/mesh/test_fields.py | 124 +----------- 3 files changed, 28 insertions(+), 469 deletions(-) diff --git a/physicsnemo/mesh/__init__.py b/physicsnemo/mesh/__init__.py index fb6fcf3cc9..ec02a19e27 100644 --- a/physicsnemo/mesh/__init__.py +++ b/physicsnemo/mesh/__init__.py @@ -16,28 +16,22 @@ from physicsnemo.mesh.domain_mesh import DomainMesh from physicsnemo.mesh.fields import ( - FieldLayout, RankSpecDict, - ScalarVectorFields, flatten_rank_spec, rank_counts, ranks_from_tensordict, validate_data_contains_ranks, - validate_rank_spec, ) from physicsnemo.mesh.mesh import MESH_FIELD_ASSOCIATIONS, Mesh, MeshFieldAssociation __all__ = [ "DomainMesh", - "FieldLayout", "MESH_FIELD_ASSOCIATIONS", "Mesh", "MeshFieldAssociation", "RankSpecDict", - "ScalarVectorFields", "flatten_rank_spec", "rank_counts", "ranks_from_tensordict", "validate_data_contains_ranks", - "validate_rank_spec", ] diff --git a/physicsnemo/mesh/fields.py b/physicsnemo/mesh/fields.py index a471526df2..24eb6c40a5 100644 --- a/physicsnemo/mesh/fields.py +++ b/physicsnemo/mesh/fields.py @@ -14,34 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -r"""Semantic scalar and vector field layout utilities. +r"""Field-rank schema metadata shared by mesh-based models. -Mesh models commonly need to turn named physical fields into dense channel -tensors. A plain concatenation is not sufficient: a scalar field and a -Cartesian vector field transform differently under a change of frame. This -module keeps those two representations separate while giving them a stable, -name-based channel order. +A rank specification maps field names to their tensor ranks. It is a plain +Python dictionary rather than a TensorDict so model construction can treat it +as static metadata. Specifications may be flat or nested to mirror a +hierarchical TensorDict structure. -``RankSpecDict`` describes the semantic tensor rank of every named leaf. The -currently supported :class:`FieldLayout` leaves are - -* rank 0: one scalar per point, with shape ``(N,)``; -* rank 1: one polar vector per point, with shape ``(N, D)``. - -Consequently, a rank-1 leaf is not an arbitrary feature axis of length ``D``: -it is explicitly promised to transform as a polar vector. Axial vectors and -higher-order tensors require additional representation metadata and are not -silently treated as rank-1 fields here. +This module contains the representation-neutral schema helpers originally +implemented by GLOBE. Runtime grouping and packing of tensor data belongs to +the consuming model rather than this metadata module. """ from __future__ import annotations from collections import Counter -from collections.abc import Collection -from copy import deepcopy -from typing import NamedTuple, TypeAlias, Union +from typing import TypeAlias, Union -import torch from tensordict import TensorDict # TODO: replace with ``type RankSpecDict = ...`` after Python 3.11 support is @@ -52,9 +41,7 @@ def flatten_rank_spec(rank_spec: RankSpecDict, sep: str = ".") -> dict[str, int]: r"""Flatten a nested rank specification to separator-joined names. - The insertion order of ``rank_spec`` is retained for compatibility. Code - that assigns positional channels should sort the returned names, as - :class:`FieldLayout` does. + The insertion order of ``rank_spec`` is retained for compatibility. Parameters ---------- @@ -85,67 +72,6 @@ def flatten_rank_spec(rank_spec: RankSpecDict, sep: str = ".") -> dict[str, int] return result -def validate_rank_spec( - rank_spec: RankSpecDict, - *, - allowed_ranks: Collection[int] | None = None, - source_label: str = "rank_spec", -) -> None: - r"""Validate the structure and integer leaves of a rank specification. - - Parameters - ---------- - rank_spec : RankSpecDict - Rank specification to validate. - allowed_ranks : Collection[int] or None, optional - If supplied, every leaf must belong to this collection. Otherwise all - non-negative integer ranks are accepted. - source_label : str, default="rank_spec" - Human-readable name used in error messages. - - Raises - ------ - TypeError - If a key is not a string or a leaf is not an integer. - ValueError - If a rank is negative or is not in ``allowed_ranks``. - """ - allowed = None if allowed_ranks is None else frozenset(allowed_ranks) - - def _validate(spec: RankSpecDict, path: tuple[str, ...]) -> None: - for key, value in spec.items(): - if not isinstance(key, str): - raise TypeError( - f"All keys in {source_label} must be strings; got " - f"{key!r} at {'.'.join(path) or ''}" - ) - leaf_path = (*path, key) - if isinstance(value, dict): - _validate(value, leaf_path) - continue - if isinstance(value, bool) or not isinstance(value, int): - raise TypeError( - f"Rank for {'.'.join(leaf_path)!r} in {source_label} must " - f"be an integer, got {value!r}" - ) - if value < 0: - raise ValueError( - f"Rank for {'.'.join(leaf_path)!r} in {source_label} must " - f"be non-negative, got {value}" - ) - if allowed is not None and value not in allowed: - raise ValueError( - f"Rank for {'.'.join(leaf_path)!r} in {source_label} must " - f"be one of {sorted(allowed)}, got {value}" - ) - - if not isinstance(rank_spec, dict): - raise TypeError( - f"{source_label} must be a dict, got {type(rank_spec).__name__}" - ) - _validate(rank_spec, ()) - - def rank_counts(rank_spec: RankSpecDict) -> Counter[int]: r"""Count leaves of each semantic rank in ``rank_spec``.""" return Counter(flatten_rank_spec(rank_spec).values()) @@ -154,7 +80,7 @@ def rank_counts(rank_spec: RankSpecDict) -> Counter[int]: def ranks_from_tensordict(td: TensorDict) -> RankSpecDict: r"""Derive semantic-rank-shaped metadata from TensorDict leaf shapes. - A leaf rank is its number of non-batch dimensions. For a point-field + A leaf rank is its number of non-batch dimensions. For a point-field TensorDict with batch size ``(N,)``, ``(N,)`` is therefore rank 0 and ``(N, D)`` is rank 1. """ @@ -176,10 +102,23 @@ def validate_data_contains_ranks( ) -> None: r"""Validate that ``data`` contains every declared leaf at its stated rank. - Additional leaves in ``data`` are allowed. Missing leaves and rank + Additional leaves in ``data`` are allowed. Missing leaves and rank mismatches are reported together to make schema errors easier to diagnose. + + Parameters + ---------- + data : TensorDict + Data TensorDict to validate. + declared_ranks : RankSpecDict + Rank specification that ``data`` must contain as a subset. + source_label : str + Human-readable description used in the error message. + + Raises + ------ + ValueError + If a declared field is missing or has a different rank. """ - validate_rank_spec(declared_ranks, source_label="declared_ranks") declared = flatten_rank_spec(declared_ranks) actual = flatten_rank_spec(ranks_from_tensordict(data)) @@ -202,264 +141,10 @@ def validate_data_contains_ranks( ) -class ScalarVectorFields(NamedTuple): - r"""Dense scalar and polar-vector channels for a collection of points. - - ``scalars`` has shape ``(N, C_s)`` and ``vectors`` has shape - ``(N, C_v, D)``. Shape validation belongs to the :class:`FieldLayout` - that defines ``C_s``, ``C_v``, and ``D``. - """ - - scalars: torch.Tensor - vectors: torch.Tensor - - -def _rank_entries( - rank_spec: RankSpecDict, - path: tuple[str, ...] = (), -) -> list[tuple[tuple[str, ...], int]]: - """Return rank leaves as unambiguous key paths.""" - entries: list[tuple[tuple[str, ...], int]] = [] - for key, value in rank_spec.items(): - leaf_path = (*path, key) - if isinstance(value, dict): - entries.extend(_rank_entries(value, leaf_path)) - else: - entries.append((leaf_path, value)) - return entries - - -def _td_key(path: tuple[str, ...]) -> str | tuple[str, ...]: - """Use a string for top-level TensorDict keys and tuples for nested keys.""" - return path[0] if len(path) == 1 else path - - -class FieldLayout: - r"""Deterministic packing layout for named scalar and polar-vector fields. - - Channel order is lexicographic by the flattened field name and therefore - does not depend on dictionary or TensorDict construction order. Scalars - and vectors are packed separately so that equivariant code cannot - accidentally mix their transformation laws. - - Parameters - ---------- - rank_spec : RankSpecDict - Nested or flat field schema. Only rank-0 and rank-1 leaves are - supported. - spatial_dim : int - Cartesian vector dimension ``D``. - sep : str, default="." - Separator used only for public flattened names and sorting. - - Notes - ----- - Each rank-0 leaf must have shape ``(N,)`` and each rank-1 leaf must have - shape ``(N, D)``. Multiple channels must be represented by multiple named - leaves rather than an extra unnamed feature dimension. - """ - - def __init__( - self, - rank_spec: RankSpecDict, - spatial_dim: int, - *, - sep: str = ".", - ) -> None: - validate_rank_spec( - rank_spec, - allowed_ranks=(0, 1), - source_label="rank_spec", - ) - if isinstance(spatial_dim, bool) or not isinstance(spatial_dim, int): - raise TypeError(f"spatial_dim must be an integer, got {spatial_dim!r}") - if spatial_dim < 1: - raise ValueError(f"spatial_dim must be positive, got {spatial_dim}") - if not isinstance(sep, str): - raise TypeError(f"sep must be a string, got {type(sep).__name__}") - if not sep: - raise ValueError("sep must not be empty") - - entries = [ - (sep.join(path), path, rank) for path, rank in _rank_entries(rank_spec) - ] - if not entries: - raise ValueError("rank_spec must contain at least one field leaf") - - # A literal key containing ``sep`` can otherwise collide with a nested - # path in public names and make deterministic ordering ambiguous. - names = [name for name, _, _ in entries] - if len(names) != len(set(names)): - raise ValueError( - f"rank_spec contains field paths that collide when flattened " - f"with separator {sep!r}" - ) - - entries.sort(key=lambda entry: entry[0]) - self._rank_spec = deepcopy(rank_spec) - self.spatial_dim = spatial_dim - self.sep = sep - self._entries = tuple(entries) - self._scalar_entries = tuple(entry for entry in entries if entry[2] == 0) - self._vector_entries = tuple(entry for entry in entries if entry[2] == 1) - - @property - def rank_spec(self) -> RankSpecDict: - """A defensive copy of the named-field schema.""" - return deepcopy(self._rank_spec) - - @property - def flat_rank_spec(self) -> dict[str, int]: - """Field ranks in deterministic packed-name order.""" - return {name: rank for name, _, rank in self._entries} - - @property - def scalar_names(self) -> tuple[str, ...]: - """Names of scalar channels in packed order.""" - return tuple(name for name, _, _ in self._scalar_entries) - - @property - def vector_names(self) -> tuple[str, ...]: - """Names of polar-vector channels in packed order.""" - return tuple(name for name, _, _ in self._vector_entries) - - @property - def n_scalars(self) -> int: - """Number of packed scalar channels.""" - return len(self._scalar_entries) - - @property - def n_vectors(self) -> int: - """Number of packed polar-vector channels.""" - return len(self._vector_entries) - - def pack(self, data: TensorDict) -> ScalarVectorFields: - r"""Pack named point fields into deterministic scalar/vector channels. - - Extra leaves in ``data`` are ignored. All selected leaves must share - dtype and device, in addition to satisfying the semantic shapes in this - layout. - """ - if data.batch_dims != 1: - raise ValueError( - "FieldLayout expects a point TensorDict with one batch dimension " - f"(N,), got batch_size={tuple(data.batch_size)}" - ) - validate_data_contains_ranks( - data=data, - declared_ranks=self._rank_spec, - source_label="data", - ) - - n_points = data.batch_size[0] - selected: list[torch.Tensor] = [] - scalar_tensors: list[torch.Tensor] = [] - vector_tensors: list[torch.Tensor] = [] - - for name, path, rank in self._entries: - tensor = data[_td_key(path)] - expected = (n_points,) if rank == 0 else (n_points, self.spatial_dim) - if tuple(tensor.shape) != expected: - kind = "scalar" if rank == 0 else "polar vector" - raise ValueError( - f"Field {name!r} is declared as a {kind} and must have " - f"shape {expected}, got {tuple(tensor.shape)}" - ) - selected.append(tensor) - if rank == 0: - scalar_tensors.append(tensor) - else: - vector_tensors.append(tensor) - - reference = selected[0] - for name, tensor in zip( - (name for name, _, _ in self._entries), selected, strict=True - ): - if tensor.device != reference.device: - raise ValueError( - f"All packed fields must share a device; field {name!r} is " - f"on {tensor.device}, expected {reference.device}" - ) - if tensor.dtype != reference.dtype: - raise ValueError( - f"All packed fields must share a dtype; field {name!r} has " - f"{tensor.dtype}, expected {reference.dtype}" - ) - - scalars = ( - torch.stack(scalar_tensors, dim=-1) - if scalar_tensors - else reference.new_empty((n_points, 0)) - ) - vectors = ( - torch.stack(vector_tensors, dim=-2) - if vector_tensors - else reference.new_empty((n_points, 0, self.spatial_dim)) - ) - return ScalarVectorFields(scalars=scalars, vectors=vectors) - - def unpack(self, fields: ScalarVectorFields) -> TensorDict: - r"""Unpack scalar/vector channels to a TensorDict matching ``rank_spec``.""" - scalars, vectors = fields - if scalars.ndim != 2: - raise ValueError( - f"scalars must have shape (N, {self.n_scalars}), got " - f"{tuple(scalars.shape)}" - ) - if vectors.ndim != 3: - raise ValueError( - f"vectors must have shape (N, {self.n_vectors}, " - f"{self.spatial_dim}), got {tuple(vectors.shape)}" - ) - - n_points = scalars.shape[0] - expected_scalars = (n_points, self.n_scalars) - expected_vectors = (n_points, self.n_vectors, self.spatial_dim) - if tuple(scalars.shape) != expected_scalars: - raise ValueError( - f"scalars must have shape {expected_scalars}, got " - f"{tuple(scalars.shape)}" - ) - if tuple(vectors.shape) != expected_vectors: - raise ValueError( - f"vectors must have shape {expected_vectors}, got " - f"{tuple(vectors.shape)}" - ) - if vectors.device != scalars.device: - raise ValueError("scalars and vectors must be on the same device") - if vectors.dtype != scalars.dtype: - raise ValueError("scalars and vectors must have the same dtype") - - nested: dict[str, object] = {} - - def _assign(path: tuple[str, ...], tensor: torch.Tensor) -> None: - target = nested - for key in path[:-1]: - child = target.setdefault(key, {}) - # The schema validation guarantees that a path component is - # either always a group or always a leaf. - target = child # type: ignore[assignment] - target[path[-1]] = tensor - - for index, (_, path, _) in enumerate(self._scalar_entries): - _assign(path, scalars[:, index]) - for index, (_, path, _) in enumerate(self._vector_entries): - _assign(path, vectors[:, index, :]) - - return TensorDict( - nested, - batch_size=torch.Size([n_points]), - device=scalars.device, - ) - - __all__ = [ - "FieldLayout", "RankSpecDict", - "ScalarVectorFields", "flatten_rank_spec", "rank_counts", "ranks_from_tensordict", "validate_data_contains_ranks", - "validate_rank_spec", ] diff --git a/test/mesh/test_fields.py b/test/mesh/test_fields.py index e5d82988e1..92e96fd347 100644 --- a/test/mesh/test_fields.py +++ b/test/mesh/test_fields.py @@ -18,20 +18,17 @@ import torch from tensordict import TensorDict -from physicsnemo.mesh import FieldLayout, ScalarVectorFields -from physicsnemo.mesh.fields import ( +from physicsnemo.mesh import ( flatten_rank_spec, rank_counts, ranks_from_tensordict, validate_data_contains_ranks, - validate_rank_spec, ) def _mixed_fields(n: int = 4) -> TensorDict: return TensorDict( { - # Deliberately not in lexicographic order. "wind": torch.arange(n * 3, dtype=torch.float32).reshape(n, 3), "state": {"temperature": torch.arange(n, dtype=torch.float32) + 20}, "displacement": -torch.arange(n * 3, dtype=torch.float32).reshape(n, 3), @@ -59,15 +56,7 @@ def test_rank_spec_helpers(): } -def test_validate_rank_spec_and_data_reports_schema_errors(): - validate_rank_spec({"scalar": 0, "tensor": 2}) - with pytest.raises(ValueError, match="must be one of .* got 2"): - validate_rank_spec({"tensor": 2}, allowed_ranks=(0, 1)) - with pytest.raises(TypeError, match="must be an integer"): - validate_rank_spec({"bad": True}) - with pytest.raises(ValueError, match="must be non-negative"): - validate_rank_spec({"bad": -1}) - +def test_validate_data_contains_ranks_reports_schema_errors(): data = TensorDict({"pressure": torch.zeros(3, 2)}, batch_size=[3]) with pytest.raises(ValueError) as error: validate_data_contains_ranks( @@ -80,112 +69,3 @@ def test_validate_rank_spec_and_data_reports_schema_errors(): " - missing leaf 'velocity' (declared rank 1)\n" " - rank mismatch for 'pressure': declared 0, got 1" ) - - -def test_field_layout_pack_is_deterministic_and_round_trips(): - ranks = { - "wind": 1, - "state": {"temperature": 0}, - "pressure": 0, - "displacement": 1, - } - reordered_ranks = { - "pressure": 0, - "displacement": 1, - "state": {"temperature": 0}, - "wind": 1, - } - data = _mixed_fields() - layout = FieldLayout(ranks, spatial_dim=3) - reordered_layout = FieldLayout(reordered_ranks, spatial_dim=3) - - assert layout.scalar_names == ("pressure", "state.temperature") - assert layout.vector_names == ("displacement", "wind") - assert layout.flat_rank_spec == { - "displacement": 1, - "pressure": 0, - "state.temperature": 0, - "wind": 1, - } - - packed = layout.pack(data) - reordered = reordered_layout.pack(data) - torch.testing.assert_close(packed.scalars, reordered.scalars) - torch.testing.assert_close(packed.vectors, reordered.vectors) - torch.testing.assert_close( - packed.scalars, - torch.stack((data["pressure"], data["state", "temperature"]), dim=-1), - ) - torch.testing.assert_close( - packed.vectors, - torch.stack((data["displacement"], data["wind"]), dim=-2), - ) - - unpacked = layout.unpack(packed) - assert set(unpacked.keys()) == {"displacement", "pressure", "state", "wind"} - for key in ("pressure", ("state", "temperature"), "displacement", "wind"): - torch.testing.assert_close(unpacked[key], data[key]) - - -def test_field_layout_preserves_polar_vector_transformation(): - data = _mixed_fields() - layout = FieldLayout( - {"pressure": 0, "wind": 1, "displacement": 1}, - spatial_dim=3, - ) - packed = layout.pack(data) - rotation = torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - rotated_data = data.clone() - rotated_data["wind"] = data["wind"] @ rotation.T - rotated_data["displacement"] = data["displacement"] @ rotation.T - rotated = layout.pack(rotated_data) - - torch.testing.assert_close(rotated.scalars, packed.scalars) - torch.testing.assert_close(rotated.vectors, packed.vectors @ rotation.T) - - -@pytest.mark.parametrize( - ("ranks", "expected_scalar_shape", "expected_vector_shape"), - [ - ({"pressure": 0}, (4, 1), (4, 0, 3)), - ({"wind": 1}, (4, 0), (4, 1, 3)), - ], -) -def test_field_layout_supports_scalar_or_vector_only( - ranks, expected_scalar_shape, expected_vector_shape -): - layout = FieldLayout(ranks, spatial_dim=3) - packed = layout.pack(_mixed_fields()) - assert packed.scalars.shape == expected_scalar_shape - assert packed.vectors.shape == expected_vector_shape - round_trip = layout.pack(layout.unpack(packed)) - torch.testing.assert_close(round_trip.scalars, packed.scalars) - torch.testing.assert_close(round_trip.vectors, packed.vectors) - - -def test_field_layout_rejects_non_vector_rank_and_wrong_vector_dimension(): - with pytest.raises(ValueError, match="must be one of .* got 2"): - FieldLayout({"stress": 2}, spatial_dim=3) - - layout = FieldLayout({"velocity": 1}, spatial_dim=3) - wrong_dim = TensorDict({"velocity": torch.zeros(5, 2)}, batch_size=[5]) - with pytest.raises(ValueError, match=r"must have shape \(5, 3\)"): - layout.pack(wrong_dim) - - -def test_field_layout_validates_packed_shapes_and_dtype(): - layout = FieldLayout({"pressure": 0, "velocity": 1}, spatial_dim=3) - with pytest.raises(ValueError, match="vectors must have shape"): - layout.unpack( - ScalarVectorFields( - scalars=torch.empty(5, 1), - vectors=torch.empty(4, 1, 3), - ) - ) - with pytest.raises(ValueError, match="must have the same dtype"): - layout.unpack( - ScalarVectorFields( - scalars=torch.empty(5, 1, dtype=torch.float32), - vectors=torch.empty(5, 1, 3, dtype=torch.float64), - ) - )