tiling and concatinating tensors

[mariusaurus]

Earth2Studio Pull Request

Description

added methods to concatenate tensors and to tile tensors. see example in docstring of tile_xx_to_yy

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.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[greptile-apps]

Greptile Summary

This PR refactors and moves tensor concatenation functionality from a recipe-specific utility to the core earth2studio.utils.coords module, making it more broadly available. It also adds a new tile_xx_to_yy() function for expanding tensor dimensions.

Key changes:

• Added tile_xx_to_yy() function to tile tensors to match leading dimensions of target tensors
• Added cat_coords() function to concatenate tensors along coordinate dimensions with validation
• Enhanced handshake_coords() to accept lists of dimensions while maintaining backwards compatibility
• Moved coordinate mapping logic (lat/lon) from cat_coords to caller in hens_ensemble.py for better separation of concerns
• Comprehensive test coverage for new functionality

The implementation is well-tested with edge cases covered. The refactoring properly handles the migration of functionality while preserving behavior.

Confidence Score: 4/5

• Safe to merge with minor considerations
• The code is well-implemented with comprehensive tests covering edge cases. The refactoring properly moves functionality to the core library. One previous feedback item about coordinate compatibility checking in cat_coords() was already addressed, though the key-order validation could still catch mismatched dimensions more gracefully before PyTorch concatenation.
• No files require special attention

Important Files Changed

Filename	Overview
earth2studio/utils/coords.py	Added tile_xx_to_yy() and cat_coords() functions for tensor tiling/concatenation, and modified handshake_coords() to accept list of dimensions
test/utils/test_coords.py	Added comprehensive unit tests for new tile_xx_to_yy() and cat_coords() functions, and updated handshake_coords() test
recipes/hens/src/hens_ensemble.py	Updated import to use cat_coords from earth2studio.utils.coords and added lat/lon mapping before concatenation

_{Last reviewed commit: 02ccfd0}

[greptile-apps]

_{5 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

test/utils/test_coords.py
GPU-only test param

@pytest.mark.parametrize("device", ["cpu", "cuda:0"]) will fail in environments without CUDA (common in CI), since .to("cuda:0") raises at runtime. If CI isn’t guaranteed to have GPUs, gate the CUDA case behind torch.cuda.is_available() (or mark it with @pytest.mark.skipif(...)).

[mariusaurus]

@greptileai

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}