Update geotransolver for 2d and 3d use cases#1502
Update geotransolver for 2d and 3d use cases#1502coreyjadams wants to merge 5 commits intoNVIDIA:mainfrom
Conversation
Greptile SummaryThis 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 (
Important Files Changed
Last reviewed commit: a0cd8b0 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| class StructuredContextProjector(_SliceToContextMixin, nn.Module): | ||
| r"""Context projector with Conv2d/Conv3d geometry encoding on structured grids. | ||
|
|
||
| Same output interface as :class:`ContextProjector`—slice tokens | ||
| :math:`(B, H, S, D)`—but projects per-cell geometry via spatial convolutions | ||
| aligned with structured GALE attention. | ||
| """ |
There was a problem hiding this comment.
Minimal class docstring for public API
StructuredContextProjector is exported in __init__.py as part of the public API, but its class docstring is much shorter than the sibling ContextProjector. Per the project's coding standards (MOD-003c), public class docstrings should include Parameters, Forward, Outputs, See Also, and Examples sections. The forward method (line 507) also lacks a docstring entirely (MOD-003d).
Consider adding documentation to match the level of detail in ContextProjector to help users understand the interface (e.g., expected spatial_shape, how it relates to ContextProjector, and a usage example).
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
| class GALEStructuredMesh2D(_GALEStructuredForwardMixin, PhysicsAttentionStructuredMesh2D): | ||
| r"""GALE with Conv2d slice projection for 2D structured grids (see :class:`GALE`).""" |
There was a problem hiding this comment.
Minimal docstrings on public classes
GALEStructuredMesh2D and GALEStructuredMesh3D (line 402) are exported in __init__.py but have single-line docstrings. Per the project's coding standards (MOD-003c), public classes should have Parameters, Forward, Outputs, See Also, and Examples sections. Consider expanding the docstrings to document the constructor parameters (especially spatial_shape, context_dim, and kernel) and include a usage example, similar to the GALE class above.
File Used: CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| assert y_none.shape == (B, H, W, 2) | ||
|
|
||
|
|
||
| def test_geotransolver_structured_3d_forward(device): |
There was a problem hiding this comment.
There are no tests for the plus=True path (not just for the new structured code, but across the entire GeoTransolver test suite.) Some of these tests could be parameterized to cover this path as well.
| self.dropout = nn.Dropout(dropout) | ||
| self._init_slice_components(dim_head, slice_num, heads, use_te, plus) | ||
|
|
||
| def _grid_project( |
There was a problem hiding this comment.
The conv setup here mirrors PhysicsAttentionStructuredMesh2D/3D.init. The output layout differs (B,H,N,D vs B,N,H,D) so they can't trivially share code, but worth noting for future maintenance.
| This module provides the GeoTransolver model and its components for learning | ||
| physics-based representations with geometry and global context awareness. | ||
|
|
||
| Classes |
There was a problem hiding this comment.
Update to include the new classes
| This module provides classes for projecting context features (geometry or global | ||
| embeddings) onto learned physical state spaces for use in GALE attention layers. | ||
|
|
||
| Classes |
There was a problem hiding this comment.
Add StructuredContextProjector
mnabian
left a comment
mnabian
left a comment
There was a problem hiding this comment.
LGTM! Left a few minor comments.
2 participants
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
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.