Update geotransolver for 2d and 3d use cases

[coreyjadams]

PhysicsNeMo Pull Request

This PR adds the same structured 2d and 3d convolutional steps to geotransolver that are available in transolver.

To make sure the code is not overly duplicating, it refactors a few places to ensure the implementation can reuse core components as much as possible.

Description

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 extends GeoTransolver with structured 2D and 3D convolutional paths, mirroring the existing Transolver structured mesh support. The refactoring extracts shared logic into reusable functions (_gale_forward_impl, _gale_compute_slice_attention_cross, _SliceToContextMixin) and adds GALEStructuredMesh2D, GALEStructuredMesh3D, and StructuredContextProjector classes.

• Adds structured_shape parameter to GeoTransolver, GALE_block, and GlobalContextBuilder to enable Conv2d/Conv3d slice projection on structured grids
• Refactors GALE forward pass and cross-attention into shared module-level functions to avoid code duplication between irregular and structured variants
• Extracts slice-to-context initialization and computation into _SliceToContextMixin shared by ContextProjector and the new StructuredContextProjector
• Adds _flatten_for_structured helper in geotransolver.py to handle both spatial (B,H,W,C) and flattened (B,N,C) inputs with proper unflatten on output
• Correctly disables ball-query local features when structured_shape is set (incompatible paths)
• New tests cover 2D/3D structured forward passes, spatial-vs-flattened input equivalence, structured+global-context, and rejection of include_local_features with structured_shape
• New public classes (GALEStructuredMesh2D, GALEStructuredMesh3D, StructuredContextProjector) have minimal docstrings compared to the project's coding standards

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/geotransolver/init.py	Exports new structured GALE variants (GALEStructuredMesh2D, GALEStructuredMesh3D, StructuredContextProjector). Clean and correct.
physicsnemo/experimental/models/geotransolver/context_projector.py	Refactors shared slice logic into _SliceToContextMixin and adds StructuredContextProjector with Conv2d/Conv3d encoding. Well-structured refactoring; StructuredContextProjector.forward lacks docstring.
physicsnemo/experimental/models/geotransolver/gale.py	Extracts shared GALE forward/cross-attention into module-level functions; adds GALEStructuredMesh2D/3D via mixin + PhysicsAttentionStructured* inheritance; GALE_block now supports spatial_shape routing. Clean refactoring with correct MRO.
physicsnemo/experimental/models/geotransolver/geotransolver.py	Adds structured_shape parameter to GeoTransolver with flatten/unflatten support for 2D/3D grids. Includes proper validation and output reshaping. Well-integrated with existing architecture.
test/models/geotransolver/test_geotransolver.py	Adds structured 2D/3D forward tests, rejects-local-features validation, and global-context-with-structured-grid test. Good coverage of new functionality.

_{Last reviewed commit: a0cd8b0}

[mnabian]

LGTM! Left a few minor comments.

[copy-pr-bot]

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]

/ok to test 57981f5

[coreyjadams]

/ok to test 93084eb

[loliverhennigh]

I think the way nn is setup is fine

[coreyjadams]

/blossom-ci

[coreyjadams]

/blossom-ci