feat: automatic data nodes creation based on dataset entries by icedoom888 · Pull Request #968 · ecmwf/anemoi-core

icedoom888 · 2026-03-11T16:36:32Z

Description

This PR tackles automatric data nodes creation based on dataset created in the dataloader.

What problem does this change solve?

Reduces redundancy avoiding multiple specifications of the same datasets.
Realtes to #966

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.

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

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

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

for more information, see https://pre-commit.ci

…oi-core into feat/auto_data_graph_creation

for more information, see https://pre-commit.ci

…oi-core into feat/auto_data_graph_creation

mpvginde · 2026-03-18T14:10:53Z

Hi @icedoom888, one minor comment is that now (unless I looked over it) there no longer an example graph config that would allow you to do stand-alone graph building.

MeraX

Hi Alberto, thank you for this contribution. I think it helps simplifying the config. At the same time it makes the discoverability of alternative node builders from the configs alone harder, but I think that should be acceptable as long as the note builders are kept in the documentation.

I've put a few comments regarding the maintainability and user-friendliness, that I', happy to discuss.

MeraX · 2026-03-18T10:51:11Z

training/src/anemoi/training/train/train.py

 from pytorch_lightning.utilities.rank_zero import rank_zero_only
 from torch_geometric.data import HeteroData

+from anemoi.graphs.utils.config import integrate_data_nodes_in_config


Suggested change

from anemoi.graphs.utils.config import integrate_data_nodes_in_config

from anemoi.models.utils.config import get_multiple_datasets_config

Is it possible to re-use the existent copy of this function?

Ideally this was done to remove dependency of graphs from models!

MeraX · 2026-03-18T14:07:05Z

graphs/src/anemoi/graphs/utils/config.py

+DEFAULT_NODE_ATTR = {
+    "area_weight": {
+        "_target_": "anemoi.graphs.nodes.attributes.SphericalAreaWeights",
+        "norm": "unit-max",
+        "fill_value": 0,
+    }
+}


This PR also introduces yet another place where defaults are stored (DEFAULT_NODE_ATTR). Are these defaults actually necessary or would it be possible to set these defaults simply as the default python argument values of anemoi.graphs.nodes.AnemoiDatasetNodes?

I would agree, these global variables seem unecessary

At the time being "area_weight" is a necessary attribute for any data nodes, this isn't enforced in the code, and test will fail if not provided.
I could try and find another workaround by enforcing creation somewhere else, but not sure where. Suggestions?

MeraX · 2026-03-18T14:09:47Z

graphs/src/anemoi/graphs/utils/config.py

+    return OmegaConf.create({default_dataset_name: config})
+
+
+def integrate_data_nodes_in_config(config: DictConfig):


I'm not sure how it's handled elsewhere in the code, but I would suggest noting here that config is actually modified. Thus I would also recommend considering returning nothing in this function.

MeraX · 2026-03-18T14:12:23Z

training/src/anemoi/training/train/train.py

            graph_filename = None

+        # Introduce Data Nodes in graph config
+        self.config = integrate_data_nodes_in_config(self.config)


Are you manipulating the global config object here? Which version of the config is logged to MLflow?

Did you consider adding a flag to the configuration that explicitly enables the automatic generation of AnemoiDatasetNodes nodes? I think this would help catching config errors where the the user specification of custom note builders is dismissed due to a typo.

Very good point @MeraX, will look into this!

JPXKQX · 2026-03-18T15:03:45Z

Hi Alberto, thanks for this pull request. I think this work is in line with the direction we agreed with anemoi-core for the coming months, but I'm not sure if we should wait for the other pieces discussed in the discussion.

I think removing the graph dependency from the tasks should be done first. Once this has been done, we can remove the new default node attributes, as these will be moved to a standalone scaler. At that point, I would expect the code to use the latlons from the dataset object directly. The current change attempts to solve the problem of updating your graph configuration when you have a new dataset, but I think this would still be problematic because you would need to specify the edges.

I think getting the coordinates from the dataset is what we want, but I would wait a little while longer to do this directly from the dataset object rather than populating the graph config

dietervdb-meteo · 2026-03-19T08:53:35Z

To be very honest I find this idea a bit strange... Doesn't this reverse the logic? Why not base the dataloader config on that of the graph??
We can build graphs on their own, not dataloaders (so far). The graph part of the training config now becomes completely useless as a config to just build a graph with. It looks quite strange, in edge builders things like 'era5' and 'cerra' show up but there is no trace of them anymore in the nodes part.

anaprietonem · 2026-03-25T13:55:16Z

Are there any updates from this @icedoom888? Based on the comments is this feature something we still need or consider merging? To flag that there is some overlap with the work Simon started in #855 and it would be good to get that one in first, as it's already in advanced review state.

icedoom888 · 2026-03-30T13:06:35Z

Hey @anaprietonem, sorry I got sidetracked previous weeks and didn't have time to look into this.
Reading the previous discussions, i'm not completely sure we all agree on a solution for the upcoming changes.

This PR as is: introduces an in-place config change that grabs data nodes information from the declared datasets.
Mario's suggestion would be to deprecate completely anemoi.graphs.nodes.AnemoiDatasetNodes and generating datasets from the dataloader's lat-lon coordinates.

At this point i'm not sure how to proceed.
#855 seems to be quite unrelated but somehow acting on graph configs too?

It's quite hard to divide, parallelize work and try to help out if the roadmap is not clear!

JPXKQX and others added 30 commits February 13, 2026 14:16

update models

f235998

.

4a3e142

update workflow

9dba030

[pre-commit.ci] auto fixes from pre-commit.com hooks

129bdd4

for more information, see https://pre-commit.ci

Merge branch 'main' into feat/expand-graph-config

a910dec

pre-commit

9ede2cf

fix: unit tests

3652dcc

tests schemas

152c4b1

fix: configs

5d7ec3d

update hierarchical

203937d

update plotting

9b570f2

fix: integration test

d79a2e7

fix

1388b00

updated

bb11369

[pre-commit.ci] auto fixes from pre-commit.com hooks

0e24869

for more information, see https://pre-commit.ci

Merge branch 'main' into feat/expand-graph-config

09e37bc

Merge branch 'main' into feat/expand-graph-config

ae01d1e

move hidden nodes name to model config

abdabac

schemas

63608a0

fix

e9c96fc

update

6177a1e

fix: tests

0daf8e3

Merge branch 'main' into feat/expand-graph-config

6cac9c9

nodes_name

85a3769

update schema

d48f097

update schema

ac31369

Merge branch 'main' into feat/expand-graph-config

adbbdf4

Merge branch 'main' into feat/expand-graph-config

ea11dbf

update test

c97a569

Merge branch 'main' into feat/expand-graph-config

a8ea455

icedoom888 added 2 commits March 12, 2026 01:35

Updated node attributes

4de5106

Merge branch 'feat/auto_data_graph_creation' of github.com:ecmwf/anem…

ede7639

…oi-core into feat/auto_data_graph_creation

icedoom888 added the ATS Approval Not Needed No approval needed by ATS label Mar 12, 2026

github-actions bot added the bug Something isn't working label Mar 12, 2026

icedoom888 added enhancement New feature or request config Affects the config files to training contributor dependencies and removed bug Something isn't working labels Mar 12, 2026

github-actions bot added the bug Something isn't working label Mar 12, 2026

icedoom888 and others added 13 commits March 12, 2026 02:21

gpc passed

d46a85f

tried something

2062bfc

Revert

7932b26

Merge branch 'main' into feat/auto_data_graph_creation

04b06b6

Fixe dmissing import

8fba12f

Created graph utils

ce99a94

[pre-commit.ci] auto fixes from pre-commit.com hooks

02f7c85

for more information, see https://pre-commit.ci

tmp print

eed77ac

Merge branch 'feat/auto_data_graph_creation' of github.com:ecmwf/anem…

2919052

…oi-core into feat/auto_data_graph_creation

relaxed config access for attributes

219b41b

Fixed default

ca7dbb7

Fixed default

d2452b3

Merge branch 'main' into feat/auto_data_graph_creation

af48de8

icedoom888 requested a review from MicheleCattaneo March 18, 2026 12:49

MeraX reviewed Mar 18, 2026

View reviewed changes

	from anemoi.graphs.utils.config import integrate_data_nodes_in_config
	from anemoi.models.utils.config import get_multiple_datasets_config

		return OmegaConf.create({default_dataset_name: config})


		def integrate_data_nodes_in_config(config: DictConfig):

Conversation

icedoom888 commented Mar 11, 2026 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

What problem does this change solve?

Uh oh!

mpvginde commented Mar 18, 2026

Uh oh!

MeraX left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JPXKQX commented Mar 18, 2026

Uh oh!

dietervdb-meteo commented Mar 19, 2026

Uh oh!

anaprietonem commented Mar 25, 2026

Uh oh!

icedoom888 commented Mar 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

icedoom888 commented Mar 11, 2026 •

edited by github-actions bot

Loading