offline plotting of prediction and ground truth

[clechartre]

Purpose

Offline plotting function to handle the plotting of the prediction output as a numpy array against a single .zarr archive. Useful to verify the accuracy of the inference

Code changes:

offline.py: created a module for this purpose
slum_offline.sh: automation of the offline.py module through batch system

Checklist

Before submitting this PR, please make sure:

• You have followed the coding standards guidelines established at Code Review Checklist.
• Docstrings and type hints are added to new and updated routines, as appropriate
• All relevant documentation has been updated or added (e.g. README)

Review

For the review process follow the guidelines at Checklist

[sadamov]

Hi @clechartre and thanks for the PR. I have some high-level comments before a deeper code-review. We have a new zarr-archive on Tsa that will be the default MeteoSwiss archive after tomorrow (when creation is done). I have already merged and updated this PR-branch where possible with latest main-branch.

• Path: /scratch/cosuna/neural-lam/zarr/cosmo_ml_data.zarr2
• The new lat/lon don't require unrotation anymore and that module was removed from the repo, you can remove it as well

One more question: Do we still need cli_plotting.py or can we combine these two workflows?

[sadamov]

I have merged this branch with #16, added some prints to the log file and some smaller comments. Before merge we should:

• Merge Updated code for latest single zarr archive #16 to have a clean history
• Discuss whether cli_plotting.py + vis.py workflow is now redundant
• Fix minor issues

[sadamov]

Suggested change

	for i in range(22):
	for i in range(time_range):

The length of the prediction should define the plotting range, no?

[sadamov]

Suggested change

	time_range = len(predictions[1, :, 1, 1]) # Number of time steps
	time_range = len(predictions[0, :, 0, 0]) # Number of time steps

use first indices to avoid issues with edge cases

[sadamov]

Suggested change

	default="TQV",
	default="T_2M",

Use more common variable as default. Future checkpoints might not contain TQV

[sadamov]

• I have added a template of the latest zarr archive in data/cosmo/templates, to make sure we are using the latest set of variables and also the unrotated grid. For that reason I removed the unrotate functionality from offline.py file.
• Further, a template for a inference.npy files is also available now in data/cosmo/templates
• After addressing small comments above we can merge. Note that the scripts are not backward compatible with older input datasets that will be deleted soon anyways.
• Note that the latest checkpoint is hi_lam, we need to re-train the model and create one checkpoint for each model soon
• Since I merged this PR branch with latest main-branch you should now again only see changes relevant to offline-plotting (which was the goal of the whole exercise).

--> Maybe add a few words in the README.md about the new data/ folder structure