inference

[clechartre]

Purpose

Adding the inference functionality to the model trainer. Allows to feed a prediction dataset onto which we can run the build-in prediction function and generate graphical output as well as numpy arrays.

Code changes:

neural_lam/models/ar_model.py: incorporated the predict_step which overwrites trainer.predict(), and incorporation of the on_predict_epoch_end() function to handle model outputs. Adapted the plotting function to be flexible across test and predict steps. Changed the handling of the global-rank specification to using the decorator rather than having if statements in the function.

neural_lam/models/base_graph_model.py: renamed the previous predict_step() for proper use by the trainer.

neural_lam/weather_dataset.py: added a predict dataset and its respective dataloader.

slurm_predict.sh: created this runner to be specific for the inference and to easily submit the inference process to the queue on Balfrin.

train_model.py: added the option to launch the inference from the trainer. It now has all multiple possibles branching actions as defined in Pytorch Lightning, train, test, eval, and predict.

Checklist

Before submitting this PR, please make sure:

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

Review

For the review process follow the guidelines at Checklist

[sadamov]

@clechartre this PR and #42 basically introduce the whole MeteoSwiss repository into upstream. We need to make sure that we introduce our developments piece-by-piece ontop of the existing codebase, without breaking the code. We have a meeting planned to dicsuss this soon, in the meantime I suggest to convert both #41 and #42 into a Draft.