StormCast training code improvements

[pzharrington]

Modulus Pull Request

Description

Refactors the StormCast training recipe to be more in-line with Modulus examples and conventions. Notable changes:

• Use a hydra-based config management system instead of YParams
• Use Modulus neural networks (SongUNet, EDMPrecond) and add a regression U-Net wrapper for StormCast (StormCastUNet)
• Use Modulus EDM loss and sampling routines
• Use Modulus utilities for checkpointing and logging

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.

Dependencies

[pzharrington]

/blossom-ci

[mnabian]

/blossom-ci

[mnabian]

Hey @pzharrington , I mostly looked at the changes to the sampler, preconditioner UNet, and metrics to make sure changes are backward compatible and are not breaking other codes. I also briefly looked at the isolated changes for the stormcast recipe and everything look good to me. Approving