NVIDIA / earth2mip

Earth-2 Model Intercomparison Project (MIP) is a python framework that enables climate researchers and scientists to inter-compare AI models for weather and climate.
https://nvidia.github.io/earth2mip/
Apache License 2.0
187 stars 41 forks source link

Fix graphcast #134

Closed nbren12 closed 9 months ago

nbren12 commented 9 months ago

Earth-2 MIP Pull Request

Description

Resolves #130

I basically rewrote the implementation and based it more closely of graphcast's xarray interface rather than trying to remove it.

After 1 or two steps the output is nearly identical to the reference (I think up to machine precision of non-deterministic GPU ops) image

Some new abstractions were helpful. I defined a functional TimeStepper abstraction that allows hiding model-specific state from the general purpose inference routines.

Checklist

Dependencies

NickGeneva commented 9 months ago

Tested and rebased this MR in: https://github.com/NVIDIA/earth2mip/pull/145

Both myself and @yairchn tested it and it seems to be running. Had a few additional bug fixes added as well.

Closing.