feat: Implemented CaFA model architecture #105

[glitch401]

Pull Request

Description

This PR introduces a modular and well-tested implementation of the CaFA (Climate-Aware Factorized Attention) model architecture, as requested in issue #105. In addition to adding support for h3-py 4.x and a Feature manager

The core innovation of the paper, the FactorizedTransformerBlock, has been implemented as a standalone component. The model also internally handles input tensors with odd spatial dimensions via padding and cropping, making it robust and user-friendly.

This contribution adds the following key components:

• FactorizedTransformerBlock: A standalone, reusable Transformer block implementing Axial Attention, located in graph_weather/models/cafa/factorize.py.
• Encoder: A convolutional encoder for channel projection and spatial downsampling.
• Processor: A processor that stacks multiple FactorizedTransformerBlock modules.
• Decoder: A transposed convolutional decoder for spatial upsampling and output projection.
• CaFAForecaster: The final, assembled end-to-end model that combines the above components.

No new external dependencies are required for this change.

Fixes #105

How Has This Been Tested?

A comprehensive suite of unit tests has been added in tests/models/cafa/test_cafa.py. These tests can be run via pytest and validate:

• The shape transformations of each individual component (Encoder, Processor, Decoder).
• That the full CaFAForecaster model works end-to-end.

That the model correctly handles inputs with odd spatial dimensions.

Additionally, the model was interactively tested in a Jupyter notebook against both randomly generated tensors and a real, single-level ERA5 dataset (2m_temperature) to ensure correctness and GPU compatibility.

• Yes

If your changes affect data processing, have you plotted any changes? i.e. have you done a quick sanity check?

• Yes

Checklist:

• My code follows OCF's coding style guidelines
• I have performed a self-review of my own code
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• I have checked my code and corrected any misspellings

[jacobbieker]

Hi, nice work on this! A few notes, just as you do this.

1. I wouldn't include redoing the feature selection/engineering in this PR, just keep it to the model implementation and tests for the implementation
2. Thanks for updating the H3 calls, but I would also split that out into its own PR, since its unrelated to the model architecture.

[glitch401]

@jacobbieker, thank you! seperated the other changes from this model architecture. let me know if this requires any more changes