Fea model refactor fengwu pengwu swinrnn

[loliverhennigh]

PhysicsNeMo Pull Request

[loliverhennigh]

/blossom-ci

[greptile-apps]

Greptile Summary

This PR refactors the Fengwu, Pangu, and SwinRNN weather forecasting models to align with the project's coding standards. Key changes include:

• Jaxtyping annotations (MOD-006): All public method signatures now use Float[torch.Tensor, ...] annotations for tensor arguments and return types. This is appropriate since all three models have jit: bool = False in their MetaData, avoiding the jaxtyping/TorchScript incompatibility.
• NumPy-style docstrings (MOD-003a-h): Class and method docstrings rewritten with proper Parameters, Forward, Outputs, and Returns sections using r""" raw strings and :math: LaTeX notation for tensor shapes.
• Input validation (MOD-005): All forward and prepare_input methods now validate tensor shapes at entry, guarded by torch.compiler.is_compiling() to avoid overhead during compiled execution.
• Instance attributes: Constructor parameters are now stored as instance attributes (e.g., self.img_size, self.embed_dim), enabling attribute access after construction and supporting checkpoint loading assertions.
• Named constants: Magic numbers for channel counts replaced with named attributes (e.g., self.surface_channels = 4, self.pressure_level, self.upper_air_channels), improving readability in channel slicing.
• Test coverage (MOD-008a, MOD-008c): New tests for constructor attribute verification, checkpoint save/load, and from_checkpoint loading with output comparison.
• Minor: Class docstrings are missing Examples sections (MOD-003j encourages these for CI-testable documentation).

Important Files Changed

Filename	Overview
physicsnemo/models/fengwu/fengwu.py	Refactored with jaxtyping annotations, NumPy-style docstrings, input validation, and parameterized channel slicing. Well-structured changes that follow coding standards (MOD-003, MOD-005, MOD-006). Missing Examples section in class docstring (MOD-003j).
physicsnemo/models/pangu/pangu.py	Refactored with jaxtyping annotations, NumPy-style docstrings, comprehensive input validation, and named constants for channel counts. prepare_input now supports both 3D and 4D surface_mask. Missing Examples section in class docstring (MOD-003j).
physicsnemo/models/swinvrnn/swinvrnn.py	Refactored with jaxtyping annotations, NumPy-style docstrings, input validation, and removal of unused local variables from forward method. Instance attributes now stored at top of __init__. Missing Examples section in class docstring (MOD-003j).
test/models/fengwu/test_fengwu.py	Added constructor attribute assertions (MOD-008a), checkpoint save/load test (MOD-008c), and from_checkpoint loading test. Well-structured tests following the coding standards.
test/models/pangu/test_pangu.py	Added constructor attribute assertions (MOD-008a), checkpoint save/load test, and from_checkpoint loading test (MOD-008c). Thorough coverage of new public attributes.
test/models/swinrnn/test_swinrnn.py	Added constructor attribute assertions (MOD-008a) and from_checkpoint loading test (MOD-008c). Consistent with the refactored model's new public attributes.

_{Last reviewed commit: 3b383f5}

[greptile-apps]

_{6 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[loliverhennigh]

/blossom-ci

[loliverhennigh]

/blossom-ci