Improve stdataset

[SarahAlidoost]

closes #75

In this PR:

• data preparation is moved from STDataset class to util data preparation function
• util data preparation function has an option to save preprocessed data to zarr --> but currently this is not used in our workflow
• STDataset class has an option to keep data lazy in init and instead load it in getitem: save memory but slower --> data of training should not be lazy
• loop over year is added to train and predict
• data classes are created grouping input arguments
• Scripts to run data preparation on HPC in scripts folder using dask distributed

I'm still testing the example_hourly notebook on HPC. but this PR can be reviewed.

TODO:

• after merging this, we can finalize 25 with 3 years of data
• after merging this PR, we can finalize Fix remaining issues with tuning #79

[rogerkuou]

Hi @SarahAlidoost , thanks for the reorganization! I have looked into the changes, and managed to run several epochs of the hourly notebook without OOM error.

The performance on my local for the hourly notebooks is ~7mins per epoch, with 4 threads computing and 2 workers reading. This is comparible with the timing profile in your notebook.

I only made one small comment about whether we should always loop in time. It is not critical. If your 1-year experiment on Levante works well, feel free to merge!

And PS. there is one linting issue raised by Ruff, I guess you also saw it.

[rogerkuou]

Do we always want to loop by year, or for some circumstance we want to allow multiple years?

[meiertgrootes]

We are training at a monthly aggregation anyway, so it doesn't really matter I think. The loop basically goes over all months, grouped by year (so allows for multiple years) within one epoch.

[meiertgrootes]

This is a duplication of the class methos also used in predict.py. Given the very nice work you've done in cleaning and structuring one could consider moving it to utils and reusing in both modules. However, this is a most a suggestion - feel free to ignore

[meiertgrootes]

Very nice and clean organization! One very minor comment and a response to @rogerkuou 's . Works locally. if/when levante test is complete feel free to merge!