Refactor data step, inference step, Jupyter notebooks

[raehik]

This PR covers multiple "feature" changes.

• refactor data step: replace CLI, rewrite some library internals, add more documentation
  • See refactor data step script into library (API) and consumer (CLI) #85 for further details on work completed here.
• refactor train step: replace CLI, refactor & document various data processes
• add new infer step: load prepped data (low-res ocean velocities) and pre-trained model, predict forcings using model
• update Jupyter notebooks to use new interfaces (code de-duplication, documenting)

Extra:

• document parts of the model, training (mirroring some explanations from the paper)

Important to-dos:

• re-add "testing" script cli/testing.py (needs better name! check?)
  • currently broken due to MLflow param loading. needed to reproduce paper figures, where the 5% of input data we skip during training is predicted on. you can still produce those figures using predictions from cli/infer.py, but you would have to split the data again yourself for the same behaviour. we could rewrite the dataset splitting used in cli/train.py in some way that it is easy to obtain the correct data to use in cli/infer.py.
  • also does fine-tuning. in previous discussions Arthur was OK to remove this, which was convenient for removing extracting MLflow dependency. if we re-added fine-tuning to cli/infer.py too, we could remove cli/testing.py
• discuss making predictions (using pretrained model)
  • somewhat avoided this-- pointed to paper for discussion on model usage
• general code & user documentation on forcing dataset shape (co-ord, variable, dimension names)
• Document how to use new inference mode: download pre-trained model from HuggingFace, load in dataset, run. Link from readme
• finish last checked-over Jupyter notebook (test_global_control)
  • This notebook needs two datasets:
    • Computed forcings from the CM2.6 dataset which are then coarsened
    • Output of inference (testing) step
  • These are then compared
  • [ ] Provide option to use the ML flow style, or to load from file (e.g., xarray.open_zarr on the forcings and output of inference step).

Closes #87 , #90 , #98 , #4 .

Known bugs:

• Training fails with IndexError when performing train/test dataset splitting #104
  when training, IndexError: index x is out of bounds for axis 0 with size y where x > y
  • somehow, the dataloaders are being asked for samples past their end?
  • perhaps the custom indexing over xarray is problematic for short arrays?
  • perhaps I broke the indexing, and it's using a coord for index? (the size is correct at least, but the index seems ~0-310)
• Forcing generation is too memory hungry again #107
• while training, I sometimes get a PyTorch Runtime Error: input type (double) and bias type (float) should be the same. It seems random. I can't tell what it's triggered by. Model parameter and training data types seem to match just fine -- I don't know where the double is coming from.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[dorchard]

Can I request that the module in the cli directory have their _cli_desc put as the first definition after imports with a comment that says # Description of this module or something as they provide a good explanation of what the file is about.

[dorchard]

The section marked ## Data on HuggingFace probably needs to say something about what kind of data and trained model (i.e., that the moment it is low res) and we can update this later if we can get the high-res.

[MarionBWeinzierl]

We discussed the changes previously, just a couple of tidy-up comments/questions below.

[MarionBWeinzierl]

This is ok to be merged in