Unexpected assert triggered in variable scaling

[gabrieloks]

What happened?

This variable scaling assert here was triggered in a situation where I would not expect it to be triggered.

What are the steps to reproduce the bug?

The bug can very easily be reproduced by changing the name of any variable in the dataloader config, as is done here:

training:
dataset:
- dataset: ${system.input.dataset}
start: 1999
end: 2022
frequency: ${data.frequency}
- dataset: ${system.input.dataset_land}
start: 1999
end: 2022
frequency: ${data.frequency}
select: [swvl1, swvl2, stl1, stl2, ssrd, strd, sf, tcc, mcc, hcc, lcc, 100u, 100v, ro]
- dataset: ${system.input.dataset_imerg}
start: 1999
end: 2022
frequency: ${data.frequency}
select: [tp]
rename:
tp: imerg
start: 1999
end: 2022
drop: []

and then change the scaling on that variable

training:
scalers:
general_variable:
target: anemoi.training.losses.scalers.GeneralVariableLossScaler
weights:

    imerg: 0.025

The code crashes with:

File "/leonardo/home/userexternal/mgabrie1/code_combined_new/anemoi-core/training/src/anemoi/training/train/train.py", line 201, in model
model = model_task(**kwargs) # GraphForecaster -> pl.LightningModule
^^^^^^^^^^^^^^^^^^^^
File "/leonardo/home/userexternal/mgabrie1/code_combined_new/anemoi-core/training/src/anemoi/training/train/tasks/rollout.py", line 67, in init
super().init(
File "/leonardo/home/userexternal/mgabrie1/code_combined_new/anemoi-core/training/src/anemoi/training/train/tasks/base.py", line 190, in init
self.scalers, self.updating_scalars = create_scalers(
^^^^^^^^^^^^^^^
File "/leonardo/home/userexternal/mgabrie1/code_combined_new/anemoi-core/training/src/anemoi/training/losses/scalers/scalers.py", line 31, in create_scalers
scalers[name] = scaler_builder.get_scaling()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/leonardo/home/userexternal/mgabrie1/code_combined_new/anemoi-core/training/src/anemoi/training/losses/scalers/base_scaler.py", line 87, in get_scaling
scaler_values = self.get_scaling_values()
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/leonardo/home/userexternal/mgabrie1/code_combined_new/anemoi-core/training/src/anemoi/training/losses/scalers/variable.py", line 100, in get_scaling_values
self.weights.get(
AssertionError: Variable imerg is not allowed to have a separate scaling besides tp.

Is this behaviour expected or maybe I'm missing something?

Version

recent branch of anemoi core

Platform (OS and architecture)

Linux Env

Relevant log output

Accompanying data

No response

Organisation

ECMWF

[anaprietonem]

good catch, @gabrieloks! I guess then it seems to be something going on between the select and the renaming. It rings a bell that this is something similar to what @JPXKQX spotted for z.

@HCookie @mc4117 - could this be related to how the https://github.com/ecmwf/anemoi-core/blob/main/training/src/anemoi/training/utils/variables_metadata.py works?

I would agree that from a 'user point of view' - defining 'imerg:0.025' should work

[mc4117]

Yes the scalers are not sensitive to the renaming of variables because they pick up the names in the metadata. When we implemented this, the idea was that this was a first step to using metadata would be implemented throughout anemoi and so this sort of thing wouldn't happen but to my knowledge this was never done anywhere else