Promotes cluster_tree up to physicsnemo.mesh.spatial for re-use. (#1710)

* Promotes `cluster_tree` up to `physicsnemo.mesh.spatial` for re-use.

* lints

---------

Co-authored-by: Corey adams <6619961+coreyjadams@users.noreply.github.com>

Author: peterdsharpe

physicsnemo/experimental/models/globe/__init__.py

@@ -14,11 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from physicsnemo.experimental.models.globe.cluster_tree import (
-    ClusterTree,
-    DualInteractionPlan,
-    SourceAggregates,
-)
 from physicsnemo.experimental.models.globe.field_kernel import (
     BarnesHutKernel,
     Kernel,
@@ -31,7 +26,4 @@
     "Kernel",
     "BarnesHutKernel",
     "MultiscaleKernel",
-    "ClusterTree",
-    "DualInteractionPlan",
-    "SourceAggregates",
 ]

physicsnemo/experimental/models/globe/field_kernel.py

@@ -51,7 +51,7 @@
 logger = logging.getLogger("globe.field_kernel")
 
 if TYPE_CHECKING:
-    from physicsnemo.experimental.models.globe.cluster_tree import (
+    from physicsnemo.mesh.spatial.cluster_tree import (
         ClusterTree,
         DualInteractionPlan,
         SourceAggregates,
@@ -919,7 +919,7 @@ def forward(
         TensorDict[str, Float[torch.Tensor, "n_targets ..."]]
             Kernel output fields at target points.
         """
-        from physicsnemo.experimental.models.globe.cluster_tree import (
+        from physicsnemo.mesh.spatial.cluster_tree import (
             ClusterTree,
             DualInteractionPlan,
             SourceAggregates,
@@ -1757,7 +1757,7 @@ def forward(
         TensorDict[str, Float[torch.Tensor, "n_targets ..."]]
             Summed results from all kernel branches.
         """
-        from physicsnemo.experimental.models.globe.cluster_tree import ClusterTree
+        from physicsnemo.mesh.spatial.cluster_tree import ClusterTree
 
         n_sources: int = len(source_points)
         device = source_points.device

physicsnemo/experimental/models/globe/model.py

@@ -27,7 +27,7 @@
 
 from physicsnemo.core.meta import ModelMetaData
 from physicsnemo.core.module import Module
-from physicsnemo.experimental.models.globe.cluster_tree import (
+from physicsnemo.mesh.spatial.cluster_tree import (
     ClusterTree,
     DualInteractionPlan,
 )
@@ -428,7 +428,7 @@ def _build_trees_and_plans(
         comm_plans : dict[str, dict[str, DualInteractionPlan]]
             Communication plans indexed as ``comm_plans[dst_bc][src_bc]``.
         """
-        from physicsnemo.experimental.models.globe.cluster_tree import ClusterTree
+        from physicsnemo.mesh.spatial.cluster_tree import ClusterTree
 
         ### ``no_grad`` is safe: tree inputs (centroids, areas) carry no grad
         ### and the outputs are consumed downstream as integer indices and as
@@ -515,7 +515,7 @@ def _build_prediction_plans(
             Plans indexed by source BC type, each computed from that source
             BC's tree to ``pred_target_tree``.
         """
-        from physicsnemo.experimental.models.globe.cluster_tree import ClusterTree
+        from physicsnemo.mesh.spatial.cluster_tree import ClusterTree
 
         ### See ``_build_trees_and_plans`` for the ``no_grad`` + build-device
         ### rationale.  ``cluster_trees`` arrive on the caller's device from

physicsnemo/mesh/spatial/__init__.py

@@ -18,6 +18,12 @@
 
 This module provides data structures and algorithms for fast spatial queries:
 - BVH (Bounding Volume Hierarchy) for point-in-cell queries
+- ClusterTree for dual-tree Barnes-Hut acceleration of kernel/attention operators
 """
 
 from physicsnemo.mesh.spatial.bvh import BVH
+from physicsnemo.mesh.spatial.cluster_tree import (
+    ClusterTree,
+    DualInteractionPlan,
+    SourceAggregates,
+)