physicsnemo
physicsnemo copied to clipboard
⛰️[EPIC]: GraphCast Enhancements
To bring the GraphCast in Modulus in parity with Google's GraphCast in terms of features and prediction skills
### Improvements - P0
- [ ] GraphCast: Cosine zenith addition (Impact: high, effort: low)
- [ ] GraphCast: Loss weighting scheme (Impact: high, effort: moderate)
- [ ] GraphCast: Distributed GNNs (Impact: high, effort: moderate)
- [ ] GraphCast: Optimize gradient checkpointing (Impact: moderate, effort: moderate)
- [ ] GraphCast: Normalization technique (Impact: moderate, effort: moderate)
- [ ] GraphCast: Increase the # of channels to match original data (Impact: moderate, effort: high)
- [ ] GraphCast: Fine-tuning (Impact: moderate, effort: high)
- [ ] https://github.com/NVIDIA/modulus/issues/508
- [ ] GraphCast: Add time history (Impact: low, effort: high)
- [ ] https://github.com/NVIDIA/modulus/issues/486
- [ ] https://github.com/NVIDIA/modulus/issues/511
### Validation
- [x] GraphCast: Choose the baseline model for comparison (GraphCast, GraphCast_small, GraphCast_operational)
- [x] GraphCast: Produce baseline results with the chosen baseline model
- [x] GraphCast: Export the dataset used in producing the baseline results
- [x] GraphCast: Export the baseline results
- [ ] GraphCast: Identify the exact configs used in the baseline model
- [x] GraphCast: Export the normalization stats
- [ ] GraphCast: Train the Modulus GraphCast model with the baseline configs, dataset, and stats
- [ ] GraphCast: Produce side-by-side comparison between Modulus GraphCats and Google GraphCast
- [ ] https://github.com/NVIDIA/modulus/issues/509
### Bugs
- [ ] Number of parameters are close but don't match exactly (35.7M vs 36.7M in the paper)
### Data Engineering
- [ ] https://github.com/NVIDIA/modulus/issues/503
### Develop a 1-degree GraphCast model
- [ ] https://github.com/NVIDIA/modulus/issues/505
- [ ] https://github.com/NVIDIA/modulus/issues/506
### Improvements - P1
- [ ] GraphCast: Add the option to use transformer_engine.LayerNorm
- [ ] Graphcast: Check whether the current combination of MeshGraphNodeBlock and MeshGraphEdgeBlock actually matches the paper
- [ ] GraphCast: Split into a GraphCastNetERA5 and a GraphCastNet model