Move calculation of metrics to dm_pix

[phinate]

Overhauls our metric calculations to use JAX and DeepMind's pix library. Working from a pipeline perspective, which is great!

Issues right now:

• SSIM is producing NaNs (I've added a test to verify this)
• I've forked the repo to replace mean -> nanmean, but I haven't worked out how to fix this for SSIM. Will require an explicit install:

pip install git+https://github.com/phinate/dm_pix@ignore-nans

or ideally, just clone my fork and make some edits if you want to help debug this. I've added @dfulu as a collaborator since he's been working on metrics the most with me, but any help is appreciated!

[phinate]

Problems are solved, so have marked this ready for review.

Still need to add install instructions for the fork of dm_pix. We could possibly add it as a submodule. Open to thoughts?

[dfulu]

Problems are solved, so have marked this ready for review.

Still need to add install instructions for the fork of dm_pix. We could possibly add it as a submodule. Open to thoughts?

From what I can tell, it seems like we are only dependent on this one module of dm_pix and its only around 250 lines of code. Should we just copy that into this repo rather than remaining dependent on that external fork?

[phinate]

Should we just copy that into this repo rather than remaining dependent on that external fork?

I also thought this. I guess I didn't know if we'd want to use more functionality from the lib. Happy to take that approach for now, since it works nicely for the medium-term.

[phinate]

Have added the relevant code to our repo, keeping the name metrics.py, but importing as dm_pix where used for forward compatibility.

TL;DR on what the problems were: our tests failed because the filter size for the convolution in metrics.ssim was bigger than the image; the nan in the top-left corner made the whole image be replaced with nans. This is solved by just reducing the filter size during tests (our images there are only 8x9 to speed up computations, so our filter size needs to be around 9 or less, instead of the default of 11).

[phinate]

Note the extra install step for running things on gpu (now in README):

python -m pip install --upgrade "jax[cuda12]"

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 82.60870% with 20 lines in your changes missing coverage. Please review.

Project coverage is 91.71%. Comparing base (cd29ad7) to head (c32d074).
Report is 13 commits behind head on main.

Files with missing lines	Patch %	Lines
src/cloudcasting/metrics.py	77.50%	18 Missing ⚠️
src/cloudcasting/validation.py	93.54%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #58      +/-   ##
==========================================
- Coverage   94.88%   91.71%   -3.17%     
==========================================
  Files          11       11              
  Lines         469      531      +62     
==========================================
+ Hits          445      487      +42     
- Misses         24       44      +20     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.