Add writing to zarr dataset for eval-mode of trained models

[leifdenby]

Describe your changes

Adds new CLI flag to neural_lam.train_model called --save-eval-to-zarr-path <path-to-dataset> which can be added when running neural-lam in eval mode (i.e. neural_lan.train_model --eval ...) to write the predictions to a zarr dataset stored in <path-to-dataset>. This functionality is motivated by our want to be able to store model predictions for later verification.

Example of usage:

Model trained with

$> pdm run python -m neural_lam.train_model --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml --hidden_dim 2 --epochs 1 --ar_steps_train 3 --ar_steps_eval 3 --graph 1level

used for inference with

$> pdm run python -m neural_lam.train_model \
   --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml \
   --hidden_dim 2 --epochs 1 --ar_steps_train 1 --ar_steps_eval 3 --eval val \
   --load saved_models/train-graph_lam-4x2-01_24_17-2502/min_val_loss.ckpt \
   --val_steps_to_log 3 --graph 1level --save-eval-to-zarr-path state_predictions.zarr/

results in:

$> zarrdump state_predictions.zarr
<xarray.Dataset> Size: 123MB
Dimensions:                    (elapsed_forecast_duration: 3, start_time: 11,
                                state_feature: 2, x: 789, y: 589)
Coordinates:
  * elapsed_forecast_duration  (elapsed_forecast_duration) timedelta64[ns] 24B ...
  * start_time                 (start_time) datetime64[ns] 88B 1990-09-07T06:...
  * state_feature              (state_feature) <U5 40B 'u100m' 'v100m'
    time                       (start_time, elapsed_forecast_duration) datetime64[ns] 264B dask.array<chunksize=(4, 3), meta=np.ndarray>
  * x                          (x) float64 6kB -1.999e+06 ... -2.925e+04
  * y                          (y) float64 5kB -6.095e+05 ... 8.605e+05
Data variables:
    state                      (start_time, elapsed_forecast_duration, state_feature, x, y) float32 123MB dask.array<chunksize=(4, 3, 2, 789, 589), meta=np.ndarray>

NB: This does not implement the inversion of the transformations that take place in mllam-data-prep (e.g. splitting individual features back into separate variables and levels. Also, the zarr datasets store time as [start_time, elapsed_forecast_duration] rather than [start_time, sample] to avoid producing a large array with many empty-values (NaNs) which would otherwise happen because each sample has a different start time. In the snippet below I have demonstrated how one could return to absolute time (probably there is a better way to do this...):

import xarray as xr
import matplotlib.pyplot as plt
ds = xr.open_zarr("state_predictions.zarr/", chunks={})

ds.state.isel(x=0, y=0, start_time=slice(0, 4)).plot(hue="start_time", col="state_feature")
plt.savefig("state_predictions_relative_time.png")

ds_abs_time = xr.concat([
    ds.isel(start_time=i).swap_dims(dict(elapsed_forecast_duration="time")) for i in range(len(ds.start_time))
], dim="sample")
ds_abs_time.state.isel(x=0, y=0).plot(hue="sample", col="state_feature")
plt.savefig("state_predictions_absolute_time.png")

Example plot with time-axis showing elapsed time:

Example plot with time-axis showing absolute time:

This probably needs more work, but I think it is ready for people to try it out and let me know what they think 😄

Issue Link

Implements #89

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.