feat: flexible forecaster

[dietervdb-meteo]

Description

We introduce the FlexibleForecaster class, which uses configurable input_offsets and output_offsets as well as a configurable autoregressive rollout step through step_shift.

Although the amount of configurability increases significantly compared to the Forecaster class, configurations are validated to ensure they make sense as a forecaster and will work technically.

A forecaster predicts the future, and can do so arbitrarily far ahead through autoregressive rollout. Concretely this implies the following validity conditions:

1. Input offsets must not contain duplicate entries. Output offsets must not contain duplicate entries.
2. All input offsets are strictly earlier than all output offsets.
3. The shift applied per rollout step must be strictly positive and ensure:
   • (3a) The model can forecast autoregressively: shifted input offsets are included in the union of input and output offsets.
   • (3b) Each future time step is forecast at most once across all rollout steps: no pairwise difference of output offsets may be a nonzero integer multiple of the shift.

If any of the above conditions are not met, a ValueError is raised at initialisation of FlexibleForecaster.

If step_shift is not specified in the config it defaults to the largest valid shift. If no valid shift exists, an error is raised.

---

Examples

Forecaster as a FlexibleForecaster

task:
  _target_: anemoi.training.tasks.FlexibleForecaster
  input_offsets: ["-6H", "0H"]
  output_offsets: ["6H"]
  # step_shift inferred as 6H

-6H  0H │ 6H
     0H   6H │ 12H
            ...

---

Input and output at different frequencies

task:
  _target_: anemoi.training.tasks.FlexibleForecaster
  input_offsets: ["-6H", "0H"]
  output_offsets: ["3H", "6H"]
  # step_shift inferred as 6H

-6H    0H │ 3H  6H
       0H       6H │ 9H  12H
             ...

---

Multiple valid rollout step shifts

task:
  _target_: anemoi.training.tasks.FlexibleForecaster
  input_offsets: ["0H"]
  output_offsets: ["2H", "3H"]
  step_shift: "2H"

0H │    2H  3H
        2H │    4H  5H
                4H │    6H  7H
         ...

task:
  _target_: anemoi.training.tasks.FlexibleForecaster
  input_offsets: ["0H"]
  output_offsets: ["2H", "3H"]
  step_shift: "3H"

0H │ 2H  3H
         3H │ 5H  6H
                  6H │ 8H  9H
                       ...

Excluding repeated forecasts

One needs to exclude step shifts that would forecast the same future time at different rollout steps, e.g.

# this will throw ValueError
task:
  _target_: anemoi.training.tasks.FlexibleForecaster
  input_offsets: ["0H"]
  output_offsets: ["1H", "3H"]
  step_shift: "1H"

0H │ 1H        3H
     1H │ 2H       4H
          2H |*3H*     5H       
                       ...

More generally, labelling the outputs in the $n$-th rollout step as
$$o_i^n=o_i+ns$$
one sees one gets repeated forecasts, $o_i^n=o_j^m$, when
$$\exists k=(m-n),,\quad o_i-o_j=ks$$
So one can only allow step_shifts such that for no pair of output offsets the difference is a mutiple of it. This is taken into account by the step_shift validation when the FlexibleForecaster is initiated.

Code changes

Apart from introducing the FlexibleForecaster class this PR also

1. refactors the input_offset and output_offset attributes of the BaseTask into a BaseTaskOffsets class.
2. introduces the subclass ForecastOffsets, which includes step_shift and validates the offsets and shift as described above.
3. slightly generalizes the rollout procedure to allow rollout under the constraints above

What problem does this change solve?

See #941

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.