Residual an non-residual variables and remappers for downscaling

[yoel-zerah]

Description

This PR is for merging ds-collab-residuals into ds-collab. This brings two new features for downscaling :

1. The possibility to train with residual and non-residual variables at the same time.
2. The management of remappers for downscaling.

Training with residual and non-residual variables

For a given input x (not forcings, and ignoring whether high_res or low_res), an output variable y and a target y_target:
Residual variables are variables for which the actual prediction y_res by the model is the difference between the output y and the input x, and must match the target residual y_target - x:

y_res = y - x  => y = y_res + x

Residual variables require that the same field exists in both the input and output data.

Conversely, non-residual variables are variables that are a direct prediction y by the model, and attempt to match the target y_target directly. These variables can either exist in both input and output datasets or in the output dataset only.

Note

Variables that are concerned by being residual or non-residual are output variables.

To specify whether variables are residual or not, two recipe fields are introduced : data.residual_fields and data.direct_prediction_fields. This is how they behave:

• if both direct_prediction_fields and residual_fields are undefined (or if they are set to null), all output variables are selected as non-residuals.
• if direct_prediction_fields: []  and residual_fields is undefined, then all variables are residuals. Conversely, if direct_prediction_fields is undefined, and residual_fields: [] , then all variables are non-residuals.
• if residual_fields is undefined and direct_prediction_fields contains a list of variables that don't account for all of the output variables, the remaining variables are residual. Conversely, if direct_prediction_fields is undefined and residual_fields contains a list of variables that don't account for all of the output variables, the remaining variables are non-residual.

Checks are performed, and the config will fail if :

• direct_prediction_fields  and residual_fields  are both defined as empty list
• direct_prediction_fields  and residual_fields  have variables in common
• The joint set of variables of direct_prediction_fields  and residual_fields is different from the set of output variables

Remappers for downscaling

Remapper are processors, meaning that they apply a transformation to the variables before they are fed to the model, and the inverse transformation when they are output by the model, just like a normalizer.
Like normalizers, remappers are defined in data.processors:

data:
  processors:
     normalizer:
        ...
     remapper: # First remapper
       ...
     remapper: # second remapper, etc...

All processors (normalizer, remapper and inputer) are applied sequentially in the same order as they are defined in the recipe (and in the reverse order for the inverse transformation) .

This PR adds the possibility for downscaling to use remappers with anemoi.models.preprocessing.ds_remapper:TopRemapper (inspired by anemoi.models.preprocessing.multi_dataset_normalizer:TopNormalizer).
Remapper are specified similarly to normalizers:

data:
   processors:
       remapper:
            _target_: anemoi.models.preprocessing.ds_remapper:TopRemapper
            _convert_: all
            config:
               default: "none"      # default remapper, applied to all variables unless they are listed in a method below
               boxcox: [tp]           # the boxcox mapping is applied to tp only
               method_kwargs:   # Defines additional keyword argument for mappers
                   boxcox:             # the mapping for which a keyword argument is specified
                      lambd: 0.1     # the kwarg for which the value is defined

The PR also introduces the possibility to supply kwargs to mapping functions used in remappers, such as the box-cox parameter (this feature doesn't exist yet in main !).

Other features

• The behaviour of default for inference noise parameters (scheduler and sampler) has been updated:
  Hardcoded inference defaults are applied if nothing else is supplied ; one of two sets of default parameters is selected depending on whether config.training.training_approach is set to probabilistic_low_noise or probabilistic_high_noise. If inference default parameters are supplied in the training recipe, they will take precedence. Finally, if parameters are supplied in the inference recipe, they will be taken into account instead.
• Two new recipe fields have been added to allow training on a single training data sample (overfitting): config.dataloader.overfit_on_index and config.dataloader.overfit_on_index. When neither overfit_on_index and overfit_on_date aren't defined or set to null, nothing changes, the training will occur normally and iterate through the training dataset. When overfit_on_index is set to an integer, the dataloader is overriden, and all training epochs are made of a single training date, for which the index in the training dataset is set in overfit_on_index.
  When overfit_on_date is set to a date (format YYYY-MM-DDThh:mm:ss), and if this date exist in the training dataset, the sample with this date will be selected as the only sample to iterate on. overfit_on_date overrides whatever is put into overfit_on_index .

Backward compatibility

There is no breaking change in this PR, other than a behaviour change in variable selection :
In a ds-collab recipe, direct_prediction_fields  and residual_fields wouldn't be defined, and all variables would be taken as residuals. On the contrary, with ds-collab-recipe, the same recipe would yield a different behaviour : all variables would be non-residual.
As detailed above, to make sure that all variables are residual again in a ds-collab recipe, just define config.data.direct_prediction_fields: [].

---

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

---

📚 Documentation preview 📚: https://anemoi-graphs--1005.org.readthedocs.build/en/1005/

---

📚 Documentation preview 📚: https://anemoi-models--1005.org.readthedocs.build/en/1005/

[OpheliaMiralles]

see #973 for kwargs in remapper (including boxcox), might be some duplication here. Would be nice to run the precommit hooks so the diff is clean

[yoel-zerah]

see #973 for kwargs in remapper (including boxcox), might be some duplication here. Would be nice to run the precommit hooks so the diff is clean

Hello,
I'm definitely expecting some duplication with feat(models)-atanh-transform, since #973 is based on the developments of this branch. However, this PR is not for merging with main, but with another downscaling branch that doesn't have the mapper kwargs features.

[Lun4m]

The clip_negative argument is missing? It seems like in the way this is implemented you need the same args for both forward and backward transforms? Maybe it'd make more sense to have a class for each mapper with forward and backward methods so you can share the arguments more explicitly? But probably that's outside the scope of this PR