Update GraphCast for model standards

[pzharrington]

PhysicsNeMo Pull Request

Description

Updates GraphCast to comply with model standards, along with misc fixes:

• Detailed shape documentation and annotation in the model itself
• Removal of deprecated multimesh_level option
• Moving recipe-specific items to the example
• Updating example to use torch fused AdamW instead of apex
• Update tests to drop old dgl reference

Tests are passing; transformer_engine dependency still needs to be worked out.

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 Overview

Greptile Summary

This PR successfully updates GraphCast to comply with model standards while making necessary refactoring changes. The main model file now includes comprehensive docstrings with detailed shape annotations using jaxtyping (compatible since jit=False), and the deprecated multimesh_level parameter has been cleanly removed in favor of mesh_level.

Key improvements include:

• Enhanced documentation with shape annotations for all forward methods
• Moved recipe-specific code (loss functions, data utilities) from core library to example directory
• Migrated from apex FusedAdam to PyTorch's native fused AdamW optimizer
• Updated tests to remove deprecated dgl references and multimesh_level usage
• Reorganized dependencies to make transformer_engine and netCDF4 optional rather than required for all model installations

All tests are reportedly passing, and the changes maintain backward compatibility at the functional level while cleaning up the API surface.

Important Files Changed

Filename	Overview
physicsnemo/models/graphcast/graph_cast_net.py	Added comprehensive docstrings, shape annotations using jaxtyping, removed deprecated multimesh_level parameter, and added input validation methods
examples/weather/graphcast/train_graphcast.py	Migrated from apex FusedAdam to torch fused AdamW, updated imports to use local loss and data_utils modules
examples/weather/graphcast/data_utils.py	Moved from physicsnemo.models.graphcast.utils.data_utils to example directory for recipe-specific use
test/models/graphcast/test_graphcast.py	Updated tests to use mesh_level instead of deprecated multimesh_level, removed backward compatibility test
pyproject.toml	Moved transformer_engine and netCDF4 out of required model dependencies, reorganized dependency structure

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[mnabian]

LGTM! Left a couple of comments.

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[ktangsali]

Changes to the toml file look good!