Feature: Option to set / change area weighting outside of graph-creation

[havardhhaugen]

Desctiption of functionality Adds functionality to change the area weighting of nodes in the limited area of a stretched grid graph. Adds two new classes to scaling.py:

• BaseAreaWeights: Uses scipy.SphericalVoronoi to calculate area weights based on latlon input grid similarly to how this is currently done in anemoi-graphs. Currently this probably won't have many use-cases since the area weights are already calculated like this in graphs, but it serves as a good base class.
• StretchedGridCutoutAreaWeights: Re-calculates area-weights using BaseAreaWeights and then scales the area weights in the limited area (cutout) such that their sum is equal to a configurable fraction of the sum of global area weights.

Adds staticmethod get_node_weights to GraphForecaster which uses the new scaling methods if they are set in the config.

Configuration options Specified in config.training.loss_scaling.spatial, for instance:

training:
  loss_scaling:
    spatial:
      _target_: anemoi.training.data.scaling.StretchedGridCutoutAreaWeights
      target_nodes: ${graph.data}
      cutout_weight_frac_of_global: 0.3     

Further comments

• Based on the implementation of area-weights in the stretched grid branch of aifs-mono and https://github.com/ecmwf/anemoi-training/tree/feature/limited_area
• Functionality tested (together with stretched grid specific metrics from #130) https://mlflow.ecmwf.int/#/metric?runs=[%227c31b5334f524b78a74bbdd99d6f4bf6%22,%22e03701c4613745a884c9aae34e83d472%22]&metric=%22val_wmse_inside_lam_contribution_epoch%22&experiments=[%2230%22]&plot_metric_keys=%5B%22val_wmse_inside_lam_contribution_epoch%22%5D&plot_layout={%22autosize%22:true,%22xaxis%22:{%22range%22:[-3752.513985943475,19833.295466415468],%22autorange%22:false},%22yaxis%22:{%22range%22:[-0.01092616814456952,0.04189337908390672],%22autorange%22:false}}&x_axis=relative&y_axis_scale=linear&line_smoothness=1&show_point=false&deselected_curves=[]&last_linear_y_axis_range=[]

Thanks to Magnus Ingstad, Jasper Wijnands and Sophie Buurman and Mario Santa Cruz for their contributions to this PR.

[havardhhaugen]

I updated the implementation based on your feedback. The base class behavior is now to load weights from the graph with a fall-back of calculating area-weights on the fly. Similiarly ReweightedGraphNodeAttribute will also attempt to load the weights from the graph and then re-weights.

I also changed the re-weighting implementation slightly to be more in line with https://github.com/ecmwf/anemoi-training/blob/feature/limited_area/src/anemoi/training/train/forecaster.py. In aifs-mono we scaled the limited area nodes to be a fraction of the global (unmasked) sum, but since the function is now more generic it makes more sense to use a fraction of the total (masked + unmasked).

I didn't actually test if the code works as expected yet, will do that when we are happy with the intended behaviour.

[HCookie]

It's looking very good, I like the overall implementation, and think it will become very useful in the future. I've left a few initial comments to start with.

[havardhhaugen]

Fixed the issues addressed in the comments and made sure the weights are loaded correctly in training for the four cases: load from graph (+fallback) and reweight (+fallback).

[havardhhaugen]

Made docstrings and changed the try-except. Should I remove node_loss_weight: area_weight from the model config files as part of this pr?

[HCookie]

If it is not used anywhere else, yes