Pressure Level Scalings only applied in specific circumstances

[mc4117]

What happened?

Pressure level scalings are only applied if the variable is specified in the variable loss scaling. This is because of this line https://github.com/ecmwf/anemoi-training/blob/develop/src/anemoi/training/train/forecaster.py#L303

What are the steps to reproduce the bug?

Specify for example

training:
pl:
t: 1

in the training config. You will then see that only the temperature is scaled by pressure level.

Version

0.3.1

Platform (OS and architecture)

Relevant log output

No response

Accompanying data

No response

Organisation

No response

[HCookie]

What should be the behaviour if not specified? It seems that not doing anything is the appropriate response

[mc4117]

I think whether you scale the loss across pressure levels should be independent of whether you use the default loss scaling for the variable.
My proposal would be something like

pressure_level_scaler:
  _target_: anemoi.training.data.scaling.ReluPressureLevelScaler
  minimum: 0.2
  slope: 0.001
  variables: [t, z, u, v, w, q]

[HCookie]

So, you want the pressure level scaling independant of loss_scaling?
I believe we already have the instantiate like you suggest.
So a user in the config should be able to:

loss:
    _target_: anemoi.training.losses.mse.WeightedMSE
   scalars: ['variable', 'pressure_level']
  # or 
   scalars: ['pressure_level']

[mc4117]

That's a cool functionality! But I'm not sure it currently works given https://github.com/ecmwf/anemoi-training/blob/develop/src/anemoi/training/train/forecaster.py#L303?

[HCookie]

But I'm not sure it currently works

It was a suggestion, yes, it won't currently work, and will require a refactor.