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 Feature: MSE metrics inside/outside regional domain for stretched grid models (KNMI) #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.

---

📚 Documentation preview 📚: https://anemoi-training--136.org.readthedocs.build/en/136/

[JPXKQX]

Thanks for the contributions!, I really like making the spatial scaling explicit. I think it may be a good opportunity to refactor this node_loss_weights into a new config field under training. I leave a few things open for discussion below:

• Merge the model.node_loss_weights and the training.loss_scaling.spatial and taking it outside the training.loss_scaling to training.node_loss_weights ??.
• I am thinking of something like the following (names to be discussed):

node_loss_weights:
      _target_: anemoi.training.losses.scaling.ReweightedGraphNodeAttribute
      target_nodes: ${graph.data}
      node_attribute: area_weight   
      scale_by_attribute: "cutout"
      weight_frac_of_total: 0.3     

• Make it explicit for the global model too, so we have something like:

node_loss_weights:
      _target_: anemoi.training.losses.scaling.GraphNodeAttribute
      target_nodes: ${graph.data}
      node_attribute: area_weight     

• Maybe the code fits better in the losses/ folder, what do you think?

[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

[HCookie]

Thanks for adding those tests, once we have ATS Approval this can likely be merged.

[mchantry]

Really nice contribution, thanks.

[havardhhaugen]

I added documentation in user-guide.training under loss function scaling.

[havardhhaugen]

This should be finished now. I can see that it is failing some of the pytests, but I think that is because these are being run with an older version of anemoi-graphs.

[JPXKQX]

PR #159 will bump the minimum required version of anemoi-graphs to 0.4.1.

[JPXKQX]

@havardhhaugen can you rebase develop branch? anemoi-graphs=0.4.1 should be a requirement now.

[JPXKQX]

Great contribution! Approved