Promote GLOBE rank schema utilities to physicsnemo.mesh

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]

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

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",
 ]

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

physicsnemo/mesh/__init__.py

@@ -15,4 +15,23 @@
 # limitations under the License.
 
 from physicsnemo.mesh.domain_mesh import DomainMesh
+from physicsnemo.mesh.fields import (
+    RankSpecDict,
+    flatten_rank_spec,
+    rank_counts,
+    ranks_from_tensordict,
+    validate_data_contains_ranks,
+)
 from physicsnemo.mesh.mesh import MESH_FIELD_ASSOCIATIONS, Mesh, MeshFieldAssociation
+
+__all__ = [
+    "DomainMesh",
+    "MESH_FIELD_ASSOCIATIONS",
+    "Mesh",
+    "MeshFieldAssociation",
+    "RankSpecDict",
+    "flatten_rank_spec",
+    "rank_counts",
+    "ranks_from_tensordict",
+    "validate_data_contains_ranks",
+]