Proposal: Unified Abstract Base Class 🧱 GraphModule 🧱 for all tasks in anemoi-training

[icedoom888]

With the recent and upcoming additions to anemoi-core, including:

1. GraphForecaster
2. GraphInterpolator
3. GraphEnsForecaster
4. GraphAutoencoder (coming soon)
5. ...and more in the pipeline

I'd like to propose introducing a new abstract base class to unify all GraphModule implementations.

Currently, most new tasks inherit from GraphForecaster, even when their objectives differ significantly—for example, the GraphInterpolator (see implementation) does not perform forecasting per se.

The goal is to define a clean and flexible base class that all graph-based modules can inherit from, each implementing their own init and forward methods as needed.

I'd like to open the discussion on the desired features and structure of this base class. The aim is to minimize required changes when extending it for new tasks, support shared configuration parameters, and reduce boilerplate code while allowing task-specific customizations.

[JPXKQX]

Hi Alberto, thanks for starting this discussion. I will add my 2 cents below. As I understand anemoi, there will be different dimensions that can be treated independently

• Task: Forecast / Autoencoder / Time interpolation / Downscaling
• ???: Deterministic / CRPS / diffusion
• Architecture: enc-proc-dec / enc-hierproc-dec / multenc-proc-dec / enc-proc-multdec / ...
• Model: gnn / graphtransformer / transformer / ...
• Mappers: static / dynamic

thanks again for starting the discussion around this topic

[icedoom888]

Hey Mario, very nicely put.

I would actually integrate with the following:

• Task: Forecast / Autoencoder / Time interpolation / Downscaling
• Objective: Deterministic / CRPS / diffusion
• Architecture: enc-proc-dec / enc-hierproc-dec / multenc-proc-dec / enc-proc-multdec / ...
• Blocks: Encoder / Decoder / Processor / Sampler / ...
• Model: GNN/ Graphtransformer / Transformer / ...
• Mappers: static / dynamic

[HCookie]

Hi,
100% agree this will become very important.
This is something fast approaching on our internal roadmap, so love seeing others engaging with the topic, I think it will be great to open up the working group.
With some other big developments currently in progress, I wonder if it would be good to start building this only when multiple datasets for example is more established?

[icedoom888]

Hey Harrison,
I would be more than happy to pitch up a first draft PR for this change where we can all collab and come together.
I would say this can happen in parallel, as i don't really envision a huge change to the anemoi-core structure:
New base class BaseGraphModule and GraphForecaster inheriting from it (and other Modules inheriting from BaseGraphModule rather than GraphForecaster)

What do you think?

[HCookie]

I feel as if there are some changes coming in the datasets which we capatalise on in this refactor, or at least reduce duplicated work.
100% I agree that parts of this can happen in parallel.
It may be useful to have a short planning meeting to establish who has capacity to work on it, and how to split it, thoughts?
There will be some untangling to be down, as well as expansion of capabilities,
Particularly

• Redefining the forecast task to be a prediction task, so that rollout can be used where necessary, i.e. downscaling doesn't need rollout.

[HCookie]

Also, an additional requirement I'd like to chuck in is preperation for coupled training...

[einrone]

Great, good idea :)

Just want to add: If we consider to create an abstract base class for the forecasters, we should probably think the same for the datamodule/class in my opinion. Should be another issue maybe.

[icedoom888]

I assume this is linked to the developements mentioned by @HCookie with the multiple datasets.
Could be a good chance to squeeze this in too!

[einrone]

Yes, i think there is a lot of synergies here.

[ssmmnn11]

I think we want to keep the class hierarchy flat and simple. We can provide a minimal base class for the different modules (forecaster etc.) and then provide functions that different modules can use to provide the required functionality.

[OpheliaMiralles]

In my opinion, the below is a bit too much granularity. An encoder with multiple input datasets is still an encoder, a hierarchical processor is a processor, so there is actually 1 architecture enc-proc-dec. I would combine task and objective (e.g. deterministic forecast, probabilistic downscaling...) and argue that we can simply implement the step / forward method for every specific task the way the interpolator used to be.

- Task: Forecast / Autoencoder / Time interpolation / Downscaling
- Objective: Deterministic / CRPS / diffusion
- Architecture: enc-proc-dec / enc-hierproc-dec / multenc-proc-dec / enc-proc-multdec / ...
- Blocks: Encoder / Decoder / Processor / Sampler / ...
- Model: GNN/ Graphtransformer / Transformer / ...
- Mappers: static / dynamic

[HCookie]

@icedoom888. Multiple Datasets has made some headway, and we are considering incorporating this refactor into that working group. Can we have a meeting to propose this combination to you? When are you available in the coming weeks?

[icedoom888]

Hey @HCookie,

Awesome! Yes would love to join, I am available next week or the one after :)