Merge global forecasting capabilities from prob_model_global

[joeloskarsson]

Global forecasting can be viewed as a special case of LAM, with no boundary conditions to consider. Therefore it is not a hard task to allow for global forecasting within the Neural-LAM framework. This has the benefit of being able to use any models in here also for global forecasting.

Global forecasting, as done in Oskarsson et al. (2024), is currently implemented on the branch prob_model_global. This issue is to keep track of merging of the prob_model_global branch into main.

The main changes in this branch include:

• Add script for creating of global forcing

  neural-lam/create_global_forcing.py

  Line 26 in f68a513

  def main():

  This will likely be taken care of through Creation of Static and Forcing Fields #34, rather than merging create_global_forcing.py directly.
• Add script for creating global static grid features

  neural-lam/create_global_grid_features.py

  Line 23 in f68a513

  def main():

  Also likely better handled in Creation of Static and Forcing Fields #34.
• Add script for creating global mesh graphs (multiscale and hierarchical)

  neural-lam/create_global_mesh.py

  Line 128 in f68a513

  def main():

  Will see if this should sit in this repo or in weather-model-graphs, and to what extent we keep the graphcast dependency for graph generation.
• Allow for 3d graph plotting of global graphs

  neural-lam/plot_global_graph.py

  Line 99 in f68a513

  def main():

  It is perhaps better to first move the 3d plotting functionality to weather-model-graphs and do this there. See point above.
• Enable generating variable weights based on pressure levels

  neural-lam/create_parameter_weights.py

  Line 53 in f68a513

  if global_ds:

  This is better handled by defining more general variable weighting setup where one can either set these weights manually per-variable in config files or use some pre-defined strategies.
• Add ERA5 dataset class

  neural-lam/neural_lam/era5_dataset.py

  Line 13 in f68a513

  class ERA5Dataset(torch.utils.data.Dataset):

  This should not be added as a separate new class, as all functionality in there should be possible through the generic dataset created in Multiple Zarr to Rule them All #54. But it needs to be checked that the result of Multiple Zarr to Rule them All #54 will work with global ERA5 data.
• Add functionality to save forecasts to zarr+xarray for later plotting or evaluation

  neural-lam/neural_lam/forecast_to_xarr.py

  Line 199 in f68a513

  def forecast_to_xarr(

• Add grid weighting (spatial points weighted differently) to loss and metrics

  neural-lam/neural_lam/metrics.py

  Line 41 in f68a513

  if grid_weights is not None:

• Make boundary forcing optional when rolling out forecasts

  neural-lam/neural_lam/models/ar_model.py

  Line 150 in f68a513

  def optional_boundary_forcing(self, pred_state, border_state):

• Plot example predictions for specific variables and lead times also during validation steps

  neural-lam/neural_lam/models/base_graph_model.py

  Line 183 in f68a513

  def validation_step(self, batch, *args):

• Allow for plotting of global data

  neural-lam/neural_lam/vis.py

  Line 79 in f68a513

  def plot_on_axis(

  Plotting functionality should be based on arbitrary projections rather than hard-coded to MEPS case. This should be done already by Multiple Zarr to Rule them All #54, but it needs to be checked that this works also for global forecasts.