Skip to content

GeoT Optimization 4/4: single-pass context+local extraction (halve radius_search)#1744

Open
coreyjadams wants to merge 6 commits into
mainfrom
geoT-opt-model-radius-search-caching
Open

GeoT Optimization 4/4: single-pass context+local extraction (halve radius_search)#1744
coreyjadams wants to merge 6 commits into
mainfrom
geoT-opt-model-radius-search-caching

Conversation

@coreyjadams

Copy link
Copy Markdown
Collaborator

PhysicsNeMo Pull Request

Add a sync-free aliasing check (_same_coords) so the context projector detects when spatial_coords and geometry are views of the same storage (common after collate unsqueeze) and runs each per-scale processor once via extract_context_and_local instead of separately in extract_context_features and extract_local_features, halving radius_search calls. Falls back to the two-pass path when inputs differ. Enable include_local_features in the surface recipe config and add coverage.

Description

Checklist

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.

Add a sync-free aliasing check (_same_coords) so the context projector
detects when spatial_coords and geometry are views of the same storage
(common after collate unsqueeze) and runs each per-scale processor once
via extract_context_and_local instead of separately in
extract_context_features and extract_local_features, halving radius_search
calls. Falls back to the two-pass path when inputs differ. Enable
include_local_features in the surface recipe config and add coverage.
@copy-pr-bot

copy-pr-bot Bot commented Jun 22, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@peterdsharpe peterdsharpe left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good and makes sense! Nice optimization here.

Comment thread physicsnemo/experimental/models/geotransolver/context_projector.py Outdated
@copy-pr-bot

copy-pr-bot Bot commented Jun 29, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coreyjadams coreyjadams marked this pull request as ready for review June 29, 2026 20:38
@greptile-apps

greptile-apps Bot commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds a single-pass extract_context_and_local method to MultiScaleFeatureExtractor that reuses the ball-query result for both context and local feature paths when the two input tensors alias the same coordinates, halving radius_search calls in the common collate case. It also refactors _structured_grid_to_conv_input into a shared utils.py module and enables include_local_features in the surface recipe config.

  • utils.py (new): Extracts structured_grid_to_conv_input (simplified 2-arg signature) and adds the sync-free tensors_alias aliasing check based on data_ptr, shape, stride, and dtype.
  • context_projector.py: Adds _same_coords / extract_context_and_local to MultiScaleFeatureExtractor and updates GlobalContextBuilder to call the consolidated method; falls back cleanly to the two-pass path when inputs differ.
  • test_context_projector.py: Adds three targeted tests — aliasing guard behavior, aliased/distinct two-pass equivalence, and a radius_search call-count assertion via monkeypatch.

Important Files Changed

Filename Overview
physicsnemo/experimental/models/geotransolver/utils.py New utility module: structured_grid_to_conv_input (refactored from context_projector) and tensors_alias (aliasing check); both are correct and well-documented
physicsnemo/experimental/models/geotransolver/context_projector.py Adds _same_coords + extract_context_and_local to MultiScaleFeatureExtractor and updates GlobalContextBuilder to use the single-pass path; logic is sound, one redundant check in tensors_alias noted
test/models/geotransolver/test_context_projector.py Adds solid coverage: aliasing guard, aliased/distinct two-pass equivalence, and radius_search call-count verification; monkeypatch correctly targets the module-level binding
examples/cfd/external_aerodynamics/unified_external_aero_recipe/conf/model/geotransolver_surface.yaml Enables include_local_features in the surface recipe config to activate the optimization path

Reviews (1): Last reviewed commit: "Cleaning up the fast path radius search" | Re-trigger Greptile

Comment thread physicsnemo/experimental/models/geotransolver/utils.py Outdated
@coreyjadams

Copy link
Copy Markdown
Collaborator Author

/ok to test b0b9b80

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants