Output predictions as zarr dataset

When running evaluation mode figures are being created. Currently, there is no option to output the predictions. I would like to implement the option to output predictions as a Zarr dataset. (WIP: https://github.com/khintz/neural-lam/tree/feature/output).

This should be possible by converting the predictions tensor from the function test_step to a Zarr dataset.

    def test_step(self, batch, batch_idx):
        """
        Run test on single batch
        """
        prediction, target, pred_std, batch_times = self.common_step(batch)

This could be done by utilising the existing function create_dataarray_from_tensor https://github.com/mllam/neural-lam/blob/main/neural_lam/weather_dataset.py#L509 and then simply using xarray to output to zarr.

This is a great idea. Some thoughts that come to mind:

Given we have implemented a flexible datastore setup for various input datatypes, it might make sense to give the same flexibility in the output. So basically the mdp datastore returns a zarr, the npy datastore returns a numpy-array, by default. The user can however always override this default and output to any given format.
There actually is a way to output predictions right now, it is hardcoded as pytorch tensors: https://github.com/mllam/neural-lam/blob/c3c1722f5130bfe374acdbad01e57fc76893e5cd/neural_lam/models/ar_model.py#L491 I don't think it's good to have this inside plot_example but rather this could be replaced by your new functionality. (i.e. Output as tensor without any metadata for debugging.)
We are currently running inference and auto-regressive rollouts within pytorch's on_test_... logic. But actually pytorch would provide a on_predict_... set of functions. This PR suggested changing this and also making sure that predict_step() is properly used. https://github.com/mllam/neural-lam/pull/45/files Currently, neural-lam overwrites predict_step and uses it during training.
Here a first attempt at prediction output: https://github.com/mllam/neural-lam/pull/44/files It is based on GRIB specifically, but maybe still somewhat useful.

So to summarize, this is clearly required. I suggest to make the output format flexible with sensible defaults. Also to adhere to pytorch's predict_... logic (maybe in a separate PR). And finally get rid of the two Draft PRs listed above.

mllam / neural-lam

Output predictions as zarr dataset #89