Improve lagged ensemble performance and versatility

[nbren12]

Earth-2 MIP Pull Request

Description

I have been trying to score graphcast, but because it is big, I was first saving the forecasts to hindcast directory using earht2mip.time_collection instead of scoring it online with earth2mip.lagged_ensembles. Unfortunately, there is some bug such that at the lead time =0 forecast is not the same as the initial condition.

This error is not present in the lagged ensemble script so I would like to use it to score graphcast directly. The lagged ensemble script currently does not handle the grid information properly and the scoring is very slow. This PR addresses both issues.

I add a memory and compute efficient implementation of CRPS in torch. The code relies on xskillscore/properscoring, but when run on the GPU this used the O(n^2) kernel method and runs out of memory when there are many channels. I did not use the modulus version of CRPS because it lacks an efficient yet exact implementation of CRPS. The kernel method is inefficient and the binned method is inexact. A bonus is that I can remove the xskillscores/properscoring dep. This code could be upstreamed to modulus core.

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.
• [ ] An issue is linked to this pull request.

Dependencies

[nbren12]

/blossom-ci

[nbren12]

@dallasfoster Thanks for the discussion. I feel it is good knowledge transfer, but perhaps a bit out of scope. Not sure if you have seen these docs: https://nvidia.github.io/earth2mip/userguide/concepts.html#forecast. If not, I highly recommend reading them to understand the current design of e2mip. This PR is however not the place to revisit all aspects of the design.

A few comments, some of which may turn into issues if desired. I think there is in this PR a pattern of code isolation that should be addressed in order to make the repo more modular.

Yes, this PR is focused on fixing some bugs and optimizing the lagged ensemble script, and some of the pieces may be useful in general, but that is something we can address later I think.

Let's upstream the crps code to modulus. I am okay with duplicating a 1-line helper for weighted averaging a few places in the code.

[NickGeneva]

@nbren12 looks in general fine to me, although right now I'm not fully familiar with the lagged ensembles. But changes outside seem good enough to get things fixed / discuss more later.

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[NickGeneva]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

/blossom-ci

[nbren12]

The tests are failing when trying to use pytest-regtest: https://github.com/NVIDIA/earth2mip/actions/runs/7493428670/job/20399108089#step:2:1160

I'm not sure why, since the logs show it is being installed.

@NickGeneva Do you have any ideas about what the issue is?

[NickGeneva]

The tests are failing when trying to use pytest-regtest: https://github.com/NVIDIA/earth2mip/actions/runs/7493428670/job/20399108089#step:2:1160

I'm not sure why, since the logs show it is being installed.

@NickGeneva Do you have any ideas about what the issue is?

Not sure, when running locally I had to remove pytest-regtest from the TOML to get pytest to function at all. May need to force a lower version to like 1.5.1. Seems other are having similar issues with the latest release: https://gitlab.com/uweschmitt/pytest-regtest/-/issues/20

[nbren12]

/blossom-ci

[nbren12]

Thanks @NickGeneva for the tip, that fixed it. Merging now.