[Stacked] Adds GLOBE Model updates for 3D + Dual Tree Traversal Acceleration

[peterdsharpe]

PhysicsNeMo Pull Request

Description

This PR includes:

• GLOBE model architecture modifications that fundamentally break the previous O(n^2) bottleneck, enabling scaling to large industrial problems. Scaling is now (both theoretically and empirically verified) to be O(n), which is very exciting to see for an all-to-all interaction problem. This is performed by doing a dual-tree variant of classical Barnes-Hut (or Fast-Multipole-Method-like) hierarchical acceleration by exploiting spatial locality. Normally Barnes-Hut would get you to O(n log n), but it turns out that by putting both the sources AND the targets in trees (a "dual tree traversal" approach), you can make it so that the interactions are no longer the bottleneck at all. This leaves O(n) operations as the bottleneck.
  • This dual-tree algorithm is implemented using a variant of a Linear Bounding Volume Hierarchy (LBVH), which is a nice way to make this more GPU-friendly (basically, this allows each layer of the tree (or in the case of dual tree traversal, each combination of layers from both trees) to coalesce into a single kernel launch, rather than doing a full tree traversal for every interaction.

Adds a theory doc about the dual tree traversal approach at ./physicsnemo/experimental/models/globe/hierarchical_acceleration.md.

This PR is:

• A split from the original now-split-up GLOBE 3D PR: Adds GLOBE 3D + Dual Tree Traversal Accelerations 🌎🌲  #1481
• A stacked PR, to be merged only after Adds the PhysicsNeMo-Mesh changes required for GLOBE 3D #1483 is merged.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Summary

This PR introduces a dual-tree Barnes-Hut acceleration for the GLOBE model, reducing kernel evaluation complexity from O(N²) to O(N) by building spatial cluster trees over both source and target points and classifying interactions into four quadrants (near-near, far-far, near-far, far-near). The implementation uses a Morton-code LBVH for GPU-friendly tree construction and a bottom-up propagation pass for aggregate centroids.

Key changes:

• cluster_tree.py (new): LBVH-based ClusterTree tensorclass + DualInteractionPlan dataclass + find_dual_interaction_pairs traversal. The two-stage leaf-pair expansion in _expand_dual_leaf_hits correctly classifies every interaction without double-counting.
• field_kernel.py: New BarnesHutKernel (four-phase kernel evaluation with chunked near-field and node-broadcast far-field) and MultiscaleKernel (one BarnesHutKernel branch per reference length, shared trees/plans). One P1 issue: MultiscaleKernel.forward() mutates the caller's global_data TensorDict in-place (line 1760) by injecting log_reference_length_ratios, which leaks into the output Mesh.global_data and modifies the user's original object — contrary to the "passed through from input" semantics documented in the forward signature.
• model.py: GLOBE updated to precompute trees and plans once per forward pass, reuse them across communication layers, and explicitly del comm_plans after the communication phase to release memory at scale.
• _ragged.py (new): Clean, GPU-friendly _ragged_arange helper.
• Test coverage is thorough: source-coverage invariants, convergence to exact at θ→0, and all four equivariance properties are verified.

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/globe/cluster_tree.py	New file implementing a GPU-compatible LBVH cluster tree and dual-tree traversal for Barnes-Hut acceleration. Core algorithm is well-structured with clean separation between tree construction (O(log N) iterations), bottom-up propagation, and dual-tree traversal. The two-stage leaf-pair expansion (_expand_dual_leaf_hits) correctly classifies interactions into four quadrants without double-counting. Minor: stable sort on Morton codes is conservative but safe.
physicsnemo/experimental/models/globe/field_kernel.py	Major update adding BarnesHutKernel and MultiscaleKernel classes. The four-phase interaction evaluation (near-near, far-far, near-far, far-near) is logically correct, with gradient-checkpointing and auto-chunking for memory efficiency. Key issue: MultiscaleKernel.forward() mutates the caller's global_data TensorDict (line 1760) by adding log_reference_length_ratios, which leaks into the output Mesh and modifies the user's original object. The _compiled_evaluate_interactions pattern via object.__setattr__ is functional but lacks a comment explaining the bypass.
physicsnemo/experimental/models/globe/model.py	GLOBE model updated to use dual-tree Barnes-Hut acceleration via ClusterTree and DualInteractionPlan. Architecture cleanly separates tree/plan construction from kernel evaluation, and correctly reuses plans across communication layers. comm_plans are explicitly deleted after use to free memory (important at 800k+ faces). The _build_source_data_ranks correctly reflects the latent feature structure between layers.
physicsnemo/mesh/spatial/_ragged.py	New helper _ragged_arange for expanding ragged (variable-length) segment descriptors into flat index arrays. Implementation is clean and GPU-friendly (no Python loops, uses repeat_interleave + cumsum). Correctly handles non-monotone starts, zero-count segments, and empty inputs.
test/models/globe/test_barnes_hut_kernel.py	Comprehensive test suite covering ClusterTree construction, DualInteractionPlan source coverage invariants, BarnesHutKernel convergence to exact (theta→0), translation/rotation/parity equivariance, all four interaction quadrants, and a GLOBE-like production configuration. Good coverage of the new dual-tree algorithm.

_{Last reviewed commit: "Merge branch 'main' ..."}

[peterdsharpe]

Greptile's find of:

Critical logic bug — _evaluate_communication_hyperlayer in model.py passes target_trees=cluster_trees and dual_plans=comm_plans to evaluate_hyperlayer, where both dicts are keyed by source BC type. For cross-BC-type interactions (e.g., evaluating freestream sources against no_slip targets), the kernel receives the source BC's self-interaction plan and tree as the target, causing wrong target centroids, incorrect broadcast indices, and potential out-of-bounds scatter_add when the two BC types have different face counts. The existing test_inference.py only exercises a single-BC-type scenario and therefore does not catch this regression.

Is actually valid. Thankfully, it doesn't affect any of our current usage (AirFRANS or DrivAerML) since these use a single BC type (no_slip) currently, but I'll get this fixed before merging.

UPDATE: fixed

[coreyjadams]

This is really nice work as usual Peter - for the GLOBE scaling challenges, it's clearly a winner. I'd like to request some changes, to make it worse, so that GeoTransolver can remain the winning 3d model a little longer! No just kidding :).

I do have some changes I want to propose to you. I think the structure and architecture of the Dual Tree Traversal for this problem would me a lot of sense to be independent of GLOBE - if it was a semi-standalone component of physicsnemo enabling arbitrary spatial computation of of some sort of 1/R^k interaction, across point clouds, and it's differentiable, there could be really nice applications in many modeling fields. I'm thinking it would be immediately applicable in quantum chemistry, also maybe electromagetism (though we don't do a ton of that, it's definitely a possible use case).

Another reason to do this is that GLOBE is probably not the final answer, and we don't want to couple too tightly to it. So the next models you build with even better design could still use these accelerations maybe without having to refactor globe to get it out.

A final reason: You mentioned the need to benchmark and consolidate some of this into kernels that are in warp or other languages. IT will be much easier to do that with standalone benchmarks and synthetic codes that we can use for testing, development, and benchmarking. Plus we want this integrated into Oliver's benchmarking suite.

At the end of the day, this is killer work and I am mostly asking, "Hey, can we package this up more amazingly with it's own bow on top, and let GLOBE use it, rather than build it so tightly with GLOBE?" - what do you think?

[coreyjadams]

I didn't follow this entirely. Can you add a section on the backwards pass before this to give some perspective on how the basic backwards pass might work?

[peterdsharpe]

Sure. The basic logic is:

1. Doing this forward pass without chunking requires a ton of memory. Without DTT, we would see 300 bytes/pair * (80e3 faces)^2 ~= 1920 GB per hyperlayer. With DTT, it gets way better, as we might only have ~2e7 pairs. But that's still ~6 GB per hyperlayer in the forward pass, and the backward pass is typically 2x-5x larger peak memory - so call it 20 GB per hyperlayer. With ~3 hyperlayers + other overhead, this can quickly OOM a H100.
2. So, we want to do chunking to mitigate memory usage. The problem is, if we just use chunking alone, we don't actually save any memory! This is because the same intermediate quantities need to be materialized simultaneously, they're just split across multiple tensors (i.e., per-chunk) rather than 1 tensor.
3. To actually save memory, we need to do chunking + checkpointing. That way the memory usage isn't sum(size of each chunk), but rather max(size of each chunk).
4. At the same time, we need to save something about the other chunks in order to rematerialize them - so the natural thing to save is the point in the computational graph that is the most-compressed latent representation of the chunks.
5. That most-compressed representation is between hyperlayers, hence why we checkpoint there. This is where you get the ~8 bytes per pair from, since we can actually reduce this to just an index.

I can update this document to add something to this effect.

[peterdsharpe]

Thank you so much for the review @coreyjadams!

Summarizing offline discussion:

• Resolved that the cluster_tree core logic for the DTT can go inside physicsnemo.experimental.nn(.geometric?).

[peterdsharpe]

Thanks for the review @coreyjadams ! Documenting our offline chat:

Almost all review feedback accepted with new commits; exception is the proposal to extract the DTT traversal (cluster_tree.py) into a standalone place for re-use by other modules.

This is an excellent idea, but given that we're still experimenting with updates to the core code (e.g., 0th-order vs. 1st-order kernel expansions, as discussed offline) I'm hesitant to have other consumers rely on it just yet - this exposes a big maintenance surface. However, we should 100% do this down the road when GLOBE gets pulled out of physicsnemo.experimental into the main models folder - earmarking that here so we don't forget.

Thanks again for the review!

[coreyjadams]

@peterdsharpe I think this is great work as always. but @CharlelieLrt is scraping our github conversations for his emulation bots so I have to make some crazy comments to throw him off: I think your integration of the 4D capacitive coupling into the gravitational clusterer was sloppy and you could have done better. And the ring computation of the square hole art deco emulator was subpar. You should have turned the synergy setting to 7. Otherwise, lackluster progress. Try harder!

(great work, approving, I like the plan to decouple when eventually moving out of experimental)