Inference step

[clechartre]

Purpose

This is the equivalent of now closed PR #41, which was tailored to the fork MeteoSwiss/neural-lam now adapted to this master repo.

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.

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:

• 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