Merge Graph-EFM model from `prob_model_lam` branch

[joeloskarsson]

This issue is to keep track of merging of the prob_model_lam branch into main. This branch contains the Graph-EFM model from Oskarsson et al. (2024) for LAM.

The main changes in this branch include:

• Addition of PropagationNet GNN layers

  neural-lam/neural_lam/interaction_net.py

  Line 134 in a12bceb

  class PropagationNet(InteractionNet):

  Make these optionally usable also in existing deterministic models.
• Add new metric: Sample-based CRPS

  neural-lam/neural_lam/metrics.py

  Line 235 in a12bceb

  def crps_ens(

• Add new metric: Squared ensemble spread

  neural-lam/neural_lam/metrics.py

  Line 328 in a12bceb

  def spread_squared(

  Remove square to make just ensemble spread before logging.
• Add the Graph-EFM model itself

  neural-lam/neural_lam/models/graph_efm.py

  Line 17 in a12bceb

  class GraphEFM(ARModel):

  This includes the different sub-modules (latent encoders and decoders).
• (Potentially) rename the Hi-LAM model to Graph-FM

  neural-lam/neural_lam/models/graph_fm.py

  Line 9 in a12bceb

  class GraphFM(BaseHiGraphModel):

• Add plotting functions for ensemble forecasts

  neural-lam/neural_lam/vis.py

  Line 116 in a12bceb

  def plot_ensemble_prediction(

I think it is best to do this merge (or at least the addition of the Graph-EFM model) after #49 is done, so the new model can fit into the new class hierarchy.

[ErikLarssonDev]

I've implemented all of these in https://github.com/ErikLarssonDev/real-prob-lam except the implementation of the actual Graph-EFM model. That is in progress.

If you want to look at the code, let me know, and I can give you access.

This work builds upon https://github.com/joeloskarsson/neural-lam-dev/releases/tag/building-ml-lams by adding probabilistic modelling.

On top of this, I have done:

• Some work related to Refactor model class hierarchy #49 and created some general probabilistic classes to make it easier to add new probabilistic models.
• Implemented cropping for a rectangular mesh grid (enabling processing with a non-graph-based processor) https://github.com/ErikLarssonDev/mllam-data-prep-real/tree/rect-boundary
• A graph/U-Net model that encodes the grid to a rectangular mesh with a GNN, processes the mesh with a U-Net processor, and then decodes the mesh to the grid with a GNN. This could be extended to any processor architecture that works on a regular grid.
• Implemented Diffusion-LAM with the Graph-U-Net architecture, but it could also be run with a Graph-FM model if the user wants a purely graph-based model.

For now, I am working on the research project and finishing the implementation of Graph-EFM, but when I am done, I think it would be nice to merge it into the main branch. If you want to start taking part now, that is fine as well.