Integrate Fourier Neural Operator (FNO) work

[bnubald]

Build integration of @pyatsysh FNO work into the IceNet pipeline.

The goals are as follows:

• Get the current demo notebook running, figure out dependencies such that it runs.
• Get it running on BAS HPC, and create a reproducible environment (e.g. conda) that can be shared.
• Demonstrator notebook of JAX + UNet with the IceNet library/pipeline.
  • https://github.com/icenet-ai/icenet-notebooks/tree/main/jax
• Demonstrator notebook integration with the IceNet library/pipeline - using the IceNet dataloaders and output processing.
• Integration with the icenet-pipeline for end to end runs.

@JimCircadian

[bnubald]

First task complete.

The second one is also mostly complete. I have run through the notebook on BAS HPC.

Still need to run through some updated dependency API changes with Orbax for checkpointing before completing the second task.

Will do so later this week.

[bnubald]

Added BAS HPC conda env file + nb run, and pushed commit to the working repo.

First two tasks are now complete.

Moving onto third task: Creation of a demonstrator notebook using IceNet dataloader and post processing (Need to use weekly averaging instead of daily inputs).

For initial implementation, will probably add it simply as an additional step after icenet data processing - for now). Although, do at some point need to allow monthly averaging to recreate the original paper.

[bnubald]

As a start to integrating JAX model development usage with IceNet, a demonstrator of Jax derived (Flax) UNet implementation (daily forecasting) using IceNet data download and plotting can be found here (Using Pytorch dataloader - deriving from work for #54 ).

The notebook also covers the usage of common JAX libraries such as Flax, Optax, Orbax, CommonLoopUtils to develop the model, train/validation/predict loops, usage of best checkpointing (based on validation rmse in this example) to save training states and load best one, custom weighted metrics, and more for model development.

Next, will be focusing on interface for FNO work, including pre-processing stages that are specific to this project.

[bnubald]

Links in with updates made (including torch dataloading) for #54, and pipeline integration architecture design with icenet v0.3.0_dev.