Geotransolver Model

[coreyjadams]

PhysicsNeMo Pull Request

This PR brings GeoTransolver to physicsnemo and unifies the training recipe with Transolver.

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

• Brings GeoTransolver to physicsnemo as an experimental model and unifies its training recipe with Transolver under a shared transformer_models directory structure
• Introduces new model components including GALE attention mechanism, context projectors, and geometric feature processors alongside comprehensive test coverage
• Reorganizes external aerodynamics examples by consolidating Transolver and GeoTransolver configurations, refactoring datapipes to support combined surface/volume processing modes

Important Files Changed

Filename	Overview
physicsnemo/experimental/models/geotransolver/geotransolver.py	New GeoTransolver model class missing required tensor shape validation logic (MOD-005) and has parameter documentation inconsistencies
physicsnemo/experimental/models/geotransolver/gale.py	GALE attention mechanism implementation with missing base class inheritance, commented debug code, and incomplete shape validation
examples/cfd/external_aerodynamics/transformer_models/src/train.py	Unified training script with complex @tensorwise decorator logic and hard-coded model switching that may introduce maintenance issues
physicsnemo/datapipes/cae/transolver_datapipe.py	Modified to support combined surface/volume mode with proper backward compatibility for normalization factors
test/models/geotransolver/test_geotransolver.py	Comprehensive test suite with inconsistent tensor indexing and non-deterministic randomness in checkpoint tests

[greptile-apps]

Additional Comments (42)

1. examples/cfd/external_aerodynamics/transformer_models/src/conf/data/surface.yaml, line 23 (link)

   syntax: Typo: 'Surface-speficic' should be 'Surface-specific'

2. examples/cfd/external_aerodynamics/transformer_models/src/conf/transolver_volume.yaml, line 24 (link)

   style: Run ID references 'bfloat16' but precision is set to 'float32' on line 27. Should the run_id match the actual precision setting?

   _{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!}

3. examples/cfd/external_aerodynamics/transformer_models/src/conf/data/volume.yaml, line 23 (link)

   syntax: Typo: 'speficic' should be 'specific'

4. examples/cfd/external_aerodynamics/transformer_models/src/conf/data/volume.yaml, line 29 (link)

   style: Missing newline at end of file

   _{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!}

5. examples/cfd/external_aerodynamics/transformer_models/src/conf/transolver_surface.yaml, line 24 (link)

   style: Run ID contains 'bfloat16' but precision is set to float32 on line 27 - consider updating for consistency

   _{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!}

6. examples/cfd/external_aerodynamics/transformer_models/src/conf/transolver_surface.yaml, line 37 (link)

   style: Missing newline at end of file

   _{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!}

7. test/models/geotransolver/test_context_projector.py, line 60 (link)

   style: Consider adding CPU testing for plus mode as well - limiting to CUDA-only may miss CPU-specific issues. Is there a specific reason plus mode only needs CUDA testing?

   _{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!}

8. examples/cfd/external_aerodynamics/transformer_models/src/compute_normalizations.py, line 92-94 (link)

   logic: Bug in Welford's algorithm: N is incremented twice (lines 77 and 92), causing incorrect variance calculation. Remove line 92 since N is already updated on line 77.

9. examples/cfd/external_aerodynamics/transformer_models/src/compute_normalizations.py, line 77 (link)

   style: Variable n is assigned but never used. Consider removing this line since batch_n on line82serves the same purpose.

   _{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!}

10. test/models/geotransolver/test_geotransolver.py, line 192 (link)

    logic: Incorrect tensor indexing - outputs is a single tensor, not a tuple

11. test/models/geotransolver/test_geotransolver.py, line 411 (link)

    logic: Incorrect tensor indexing - outputs is a single tensor, not a tuple

12. test/models/geotransolver/test_geotransolver.py, line 460 (link)

    logic: Non-deterministic randomness breaks test reproducibility

13. test/models/geotransolver/test_geotransolver.py, line 740 (link)

    logic: Incorrect tensor indexing - outputs is a single tensor, not a tuple

14. examples/cfd/external_aerodynamics/transformer_models/src/metrics.py, line 172 (link)

    logic: Inconsistent tensor dimensionality - mae_num is 3D but indexed as 2D here

15. examples/cfd/external_aerodynamics/transformer_models/README.md, line 46 (link)

    style: The API reference mentions physicsnemo.experimental.models.typhon but the model is described as 'Typhon' throughout the document while the API suggests 'geotransolver'. Consider consistency in naming. Should the documentation consistently use 'Typhon' or 'GeoTransolver' as the model name to match the API structure?

    _{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!}

16. examples/cfd/external_aerodynamics/transformer_models/README.md, line 34 (link)

    logic: The configuration examples mention typhon_surface and typhon_volume but based on the file structure, these might actually be geotransolver_surface and geotransolver_volume. Do the configuration names match the actual config files in the codebase?

17. examples/cfd/external_aerodynamics/transformer_models/src/inference_on_vtk.py, line 662 (link)

    logic: Potential IndexError if no matching files found

18. examples/cfd/external_aerodynamics/transformer_models/src/inference_on_vtk.py, line 677 (link)

    logic: Potential IndexError if no matching files found

19. examples/cfd/external_aerodynamics/transformer_models/deprecated/inference_on_vtp.py, line 73-74 (link)

    logic: Bug: surface_normals is undefined. Should be normals on line 73.

20. examples/cfd/external_aerodynamics/transformer_models/src/inference_on_zarr.py, line 543 (link)

    logic: Potential directory creation issue - path may not exist before writing file

21. examples/cfd/external_aerodynamics/transformer_models/src/inference_on_zarr.py, line 543 (link)

    logic: Datetime formatting could create invalid filenames due to colons and spaces

22. examples/cfd/external_aerodynamics/transformer_models/src/inference_on_zarr.py, line 568 (link)

    logic: Same directory and datetime filename issues as surface mode

23. examples/cfd/external_aerodynamics/transformer_models/src/utils.py, line 41-48 (link)

    style: Missing required docstring sections per MOD-003d. Add Parameters and Returns sections following NumPy style.

    _{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!}

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

24. physicsnemo/experimental/models/geotransolver/context_projector.py, line 39 (link)

    logic: Class should inherit from physicsnemo.Module instead of nn.Module to follow MOD-001 coding standard

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

25. physicsnemo/experimental/models/geotransolver/context_projector.py, line 313-314 (link)

    style: Remove commented debug print statements before merging

26. physicsnemo/experimental/models/geotransolver/context_projector.py, line 389 (link)

    logic: Argument order inconsistency - spatial_coords and geometry order differs from line 398. Should the argument order be consistent between extract_context_features and extract_local_features method calls?

27. physicsnemo/experimental/models/geotransolver/context_projector.py, line 126-128 (link)

    style: Missing jaxtyping tensor shape annotations in method signature per MOD-006

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

28. physicsnemo/experimental/models/geotransolver/geotransolver.py, line 93-95 (link)

    syntax: Parameter type documentation is inconsistent. Docstring says 'int' but type annotation shows 'int | tuple[int, ...]'

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

29. physicsnemo/experimental/models/geotransolver/geotransolver.py, line 128 (link)

    syntax: Default value documented as 512 but line 221 shows 32

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

30. physicsnemo/experimental/models/geotransolver/geotransolver.py, line 329-381 (link)

    logic: Forward method missing tensor shape validation logic as required by MOD-005. Should validate shapes of local_embedding, geometry, and global_embedding tensors with torch.compiler.is_compiling() guard

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

31. physicsnemo/experimental/models/geotransolver/geotransolver.py, line 331-335 (link)

    syntax: Missing jaxtyping tensor annotations for forward method parameters as required by MOD-006

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

32. examples/cfd/external_aerodynamics/transformer_models/deprecated/datapipe.py, line 64-78 (link)

    style: Function missing required docstring sections per MOD-003d. Should have Parameters and Returns sections.

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

33. physicsnemo/datapipes/cae/transolver_datapipe.py, line 509 (link)

    logic: Potential KeyError if outputs_surf doesn't contain 'fx' key when air_density/stream_velocity are missing. Should there be a check to ensure both outputs_surf and outputs_vol have the same keys before accessing fx?

34. physicsnemo/datapipes/cae/transolver_datapipe.py, line 568-569 (link)

    logic: The ValueError message doesn't handle the 'combined' model type case in the auto factor selection

35. examples/cfd/external_aerodynamics/transformer_models/src/train.py, line 21-25 (link)

    syntax: Duplicate import - Sequence is imported from both typing (line 21) and collections.abc (line 25). Remove one of these imports.

36. examples/cfd/external_aerodynamics/transformer_models/src/train.py, line 320 (link)

    logic: Potential bug: loss.item() called on tensor that may be a list/multi-dimensional. This will fail if loss is not a scalar. Is loss guaranteed to be a scalar tensor here, or could it be multi-dimensional when handling multiple point clouds?

37. physicsnemo/experimental/models/geotransolver/gale.py, line 82 (link)

    syntax: Parameter dim missing type annotation. Should be dim: int per MOD-006 for jaxtyping tensor annotations in public function signatures.

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

38. physicsnemo/experimental/models/geotransolver/gale.py, line 105-107 (link)

    logic: Method signature doesn't match docstring. Docstring describes slice_tokens as single tensor but method expects list. Line 125 calls torch.cat(slice_tokens, dim=-2) which will fail if slice_tokens is a single tensor.

39. physicsnemo/experimental/models/geotransolver/gale.py, line 149-151 (link)

    syntax: Return type annotation is incorrect. Method returns outputs which is a list (line 211-213), not a single torch.Tensor as annotated.

40. physicsnemo/experimental/models/geotransolver/gale.py, line 156-158 (link)

    syntax: Docstring parameter description doesn't match signature. Docstring describes x as single tensor but signature expects tuple[torch.Tensor, ...].

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

41. physicsnemo/experimental/models/geotransolver/gale.py, line 315 (link)

    syntax: Return type annotation is incorrect. Method returns fx which is a list (line 338), not a single torch.Tensor as annotated.

42. physicsnemo/experimental/models/geotransolver/gale.py, line 320-322 (link)

    syntax: Docstring parameter descriptions don't match method signature. Both fx and global_context are described as single tensors but signature expects tuple[torch.Tensor, ...] for both.

    Context Used: File from greptile.json - CODING_STANDARDS/MODELS_IMPLEMENTATION.md (source)

_{40 files reviewed, 42 comments}

_{Edit Code Review Agent Settings | Greptile}

[RishikeshRanade]

There is a shape mismatch here in the global_parameters setting. Global parameters have an extra dimension (1, 1, 2) in this case. We need to modify the datapipe to squeeze the first dimension.