⛰️[EPIC]: GraphCast Enhancements

[mnabian]

To bring the GraphCast in Modulus in parity with Google's GraphCast in terms of features and prediction skills

Improvements - P0

Beta Give feedback

1. GraphCast: Cosine zenith addition (Impact: high, effort: low)
2. GraphCast: Loss weighting scheme (Impact: high, effort: moderate) #530
   2 - In Progress +1
   
3. GraphCast: Distributed GNNs (Impact: high, effort: moderate)
4. GraphCast: Normalization technique (Impact: moderate, effort: moderate)
5. GraphCast: Fine-tuning (Impact: moderate, effort: high) #560
   +0
   
6. GraphCast: Rotation of icosphere (pre-processing) (Impact: low, effort: low) #508
   +0
   
7. GraphCast: Add time history (Impact: low, effort: high) #515
   +0
   
8. 🚀[FEA]: Better handling of the icosphere construction/usage in GraphCast #486
   ? - Needs Triage enhancement +2
   
9. Vectorized operations in the ERA5 HDF5 synthetic dataset to improve performance #511
   +0
   
10. GraphCast: Better handling of the static datasets #517
    2 - In Progress +1
    

Validation

Beta Give feedback

1. GraphCast: Choose the baseline model for comparison (GraphCast, GraphCast_small, GraphCast_operational)
2. GraphCast: Produce baseline results with the chosen baseline model
3. GraphCast: Export the dataset used in producing the baseline results
4. GraphCast: Export the baseline results
5. GraphCast: Identify the exact configs used in the baseline model
6. GraphCast: Export the normalization stats
7. GraphCast: Train the Modulus GraphCast model with the baseline configs, dataset, and stats #513
   +0
   
8. GraphCast: Produce side-by-side comparison between Modulus GraphCats and Google GraphCast
9. GraphCast: Bi-linear interpolation of 0.25 degree ERA5 data to a 1 degree grid #509
   +0

Bugs

Beta Give feedback

1. GraphCast: Number of parameters are close but don't match exactly (35.7M vs 36.7M in the paper) #514
   +0
   

Data Engineering

Beta Give feedback

1. GraphCast: Prepare an IO-optimized dataset of 227 channels, with an optimized dataloader #503
   0 of 2
   +0
2. 🚀[FEA]: Add support for cos zenith angle in the ERA5 HDF5 dataloader #527
   ? - Needs Triage enhancement +2
   
3. 🚀[FEA]: Enable history > 0 in the ERA5 HDF5 datapipe #519
   ? - Needs Triage enhancement +2
   

Develop a 1-degree GraphCast model

Beta Give feedback

1. Support for resizing with interpolation in the ERA5 HDF5 datapipe #505
   +0
   
2. GraphCast: Support for arbitrary input shapes #506
   +0

Improvements - P1

Beta Give feedback

1. GraphCast: Add the option to use transformer_engine.LayerNorm #516
   +0
   
2. Graphcast: Check whether the current combination of MeshGraphNodeBlock and MeshGraphEdgeBlock actually matches the paper
3. GraphCast: Split into a GraphCastNetERA5 and a GraphCastNet model
4. Configurable static channels
5. GraphCast: Add clock variables #555
   +0

Extension to GenCast - P2

Beta Give feedback

1. GraphCast: Add a config-option specifying whether we want to merge meshes #558
   +0
   
2. GraphCast: Add k-hops neighborhoods explicitly for mesh nodes #559
   +0
   

Other improvements - P2

Beta Give feedback

1. GraphCast: Symmetric approach for constructing grid2mesh and mesh2grid graphs