⛰️[EPIC]: CorrDiff Usability Enhancements

[mnabian]

Tasks - P0

Beta Give feedback

1. CorrDiff: Split generate.py, train.py and training_loop.py into functions/classes that handle more well-defined tasks #639
   +0
   
2. CorrDiff: Reduce redundancy between generate.py, train.py and training_loop.py #633
   +0
   
3. CorrDiff: Comprehensive Documentation #632
   +0
   
4. CorrDiff: Improve config handling, separate configs for different dataloaders, reduce configurable parameters and use the defaults instead #614
   +0
   
5. CorrDiff: Remove unnecessary EDM abstractions #613
   +0
   
6. CorrDiff: Generalized dataloader
7. CorrDiff: Remove unused modules/code #612
   +0
   
8. CorDiff: Unit tests for training/inference, put it in the examples folder
9. Unit tests for the CorrDiff dataloader #601
   +0
   
10. CorrDiff: Build all the objects outside of the training loop #611
    +0
    
11. CorrDiff: Use more descriptive parameter names #600
    +0
    
12. Corrdiff: cleanup CWB dataloader. #599
    +0
    
13. CorrDiff: Better handling of the downsampling factor #634
    +0
    
14. Exclude examples from pytest #594
    +0
    
15. 🐛[BUG]: EDMPrecondSRV2 called with wrong arguments order #541
    ? - Needs Triage bug +2
16. 🐛[BUG]: CorrDiff - EDMLossSR.__init__() got an unexpected keyword argument 'regression_net' #537
    ? - Needs Triage bug +2
17. 🚀[FEA]: Modularize CorrDiff inference and move into the main package #313
    ? - Needs Triage enhancement +2
    
18. 🐛[BUG]: Running corrdiff/generate.py raises a shape exception #542
    ? - Needs Triage bug +2
    
19. CorrDiff: Better handling of the dataset attributes #415
    1 - On Deck +1
    
20. CorrDiff training and inference validation #655
    +0
    

Tasks - P1

Beta Give feedback

1. CorrDiff: Implement a synthetic datapipe #598
   +0
   
2. CorrDiff: Support distributing generation over time steps as well as ensemble members
3. CorrDiff: Don't make models fully configurable, but make different versions (UNet_S, UNet_XS, etc.) #602
   +0
   
4. CorrDiff-Lite recipe #635
   +0
   
5. CorrDiff: Better handling of EMA #623
   +0
   
6. CorrDiff: Avoid hard-coding the validation year in the datapipe
7. CorrDiff: Switch to Tensorboard logging #622
   +0
   
8. 🚀[FEA]: CorrDiff training/inference without regression #478
   ? - Needs Triage enhancement +2
   

[ram-cherukuri]

@mnabian So looking through the task list, we are rearchitecting the code to make it cleaner. From a user perspective, what will they be able to do more easily in terms of say what types of data formats can they ingest, does our training code provide configs that make it easy to train with the supported data formats?

[mnabian]

Scope of the enhancements:

• Re-architect Training and Generation: Significantly refactor training and generation processes to hide unnecessary details and configurations, improving modularity, and removing non-intuitive abstractions for a more user-friendly experience.

• Enhance Configuration Management: Improve the handling of configurations for more efficient, generalizable, and user-friendly setup. Also hide unnecessary configurations from the users to improve usability.

• Simplify Metrics Logging: Streamline the logging process for training and validation metrics to provide clear and accessible performance tracking.

• Develop CorrDiff-Lite: Create a simplified, educational version of CorrDiff training and generation, demonstrated on the HRRR dataset, to facilitate learning and accessibility.

• Refactor and Generalize Dataloader: Enhance and generalize the dataloader to be served as a template for the users to implement their own customized dataloaders.

• Implement Regression Tests: Introduce regression tests for CorrDiff training and generation to ensure accuracy and reliability in future code updates.

• Develop Comprehensive Documentation: Create thorough and user-friendly documentation to support all aspects of CorrDiff, ensuring users have clear guidance and resources. This includes both README documentation as well as API documentation.

[mnabian]

Closing as complete.

#604
#629
#636
#676