feat: automatic data nodes creation based on dataset entries

graphs/src/anemoi/graphs/utils/__init__.py

@@ -0,0 +1,25 @@
+# (C) Copyright 2025 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+from .compile import NodesAxis
+from .compile import concat_edges
+from .compile import get_distributed_device
+from .compile import get_edge_attributes
+from .compile import get_grid_reference_distance
+from .compile import get_nearest_neighbour
+from .compile import haversine_distance
+
+__all__ = [
+    "get_distributed_device",
+    "get_nearest_neighbour",
+    "get_grid_reference_distance",
+    "concat_edges",
+    "haversine_distance",
+    "NodesAxis",
+    "get_edge_attributes",
+]

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

@@ -0,0 +1,84 @@
+# (C) Copyright 2026- Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import logging
+
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+
+LOGGER = logging.getLogger(__name__)
+
+DEFAULT_DATASET_NAME = "data"
+DEFAULT_NODE_ATTR = {
+    "area_weight": {
+        "_target_": "anemoi.graphs.nodes.attributes.SphericalAreaWeights",
+        "norm": "unit-max",
+        "fill_value": 0,
+    }
+}
+
+
+def get_multiple_datasets_config(config: DictConfig, default_dataset_name: str = DEFAULT_DATASET_NAME) -> dict:
+    """Get multiple datasets configuration for old configs.
+    Use /'data/' as the default dataset name.
+    """
+    if "datasets" in config:
+        if isinstance(config, dict):
+            return config["datasets"]
+        return config.datasets
+
+    return OmegaConf.create({default_dataset_name: config})
+
+
+def integrate_data_nodes_in_config(config: DictConfig):
+
+    # Introduce Data Nodes in graph config
+    train_configs = get_multiple_datasets_config(config.dataloader.training)
+
+    val_configs = {}
+    if hasattr(config.dataloader, "validation"):
+        val_configs = get_multiple_datasets_config(config.dataloader.validation)
+
+    test_configs = {}
+    if hasattr(config.dataloader, "test"):
+        test_configs = get_multiple_datasets_config(config.dataloader.test)
+
+    dataset_configs = {
+        **train_configs,
+        **val_configs,
+        **test_configs,
+    }
+
+    for dataset_name, dataset_config in dataset_configs.items():
+        if dataset_name not in config.graph.nodes:
+            LOGGER.info("Creating graph node entry for dataset '%s'", dataset_name)
+            dataset_reader_config = dataset_config.dataset_config
+            if isinstance(dataset_reader_config, (DictConfig, dict)):
+                if "dataset" not in dataset_reader_config:
+                    msg = f"Dataset '{dataset_name}' is missing 'dataset' key."
+                    raise ValueError(msg)
+                dataset_source = dataset_reader_config["dataset"]
+            else:
+                dataset_source = dataset_reader_config
+
+            if dataset_source is None:
+                msg = f"Dataset source is None for dataset '{dataset_name}'. Check dataloader.dataset_config.dataset."
+                raise ValueError(msg)
+
+            # Add dataset nodes from dataloader into graph recepe
+            config.graph.nodes[dataset_name] = {
+                "node_builder": {"_target_": "anemoi.graphs.nodes.AnemoiDatasetNodes", "dataset": dataset_source},
+                "attributes": (
+                    config.graph.attributes.nodes if hasattr(config.graph, "attributes") else DEFAULT_NODE_ATTR
+                ),
+            }
+        else:
+            LOGGER.info("Graph node entry for dataset '%s' is already specified in the config.", dataset_name)
+
+    return config

training/src/anemoi/training/config/graph/encoder_decoder_only.yaml

@@ -2,12 +2,6 @@
 overwrite: True
 
 nodes:
-  # Data nodes
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes # options: AnemoiDatasetNodes, NPZFileNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes} # options: l1, l2, unit-max, unit-sum, unit-std
   # Hidden nodes
   hidden:
     node_builder:

training/src/anemoi/training/config/graph/hierarchical_2level.yaml

@@ -3,12 +3,7 @@ overwrite: True
 
 
 nodes:
-  # Data nodes
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes}
+  # Hidden nodes
   hidden_1:
     node_builder:
       _target_: anemoi.graphs.nodes.TriNodes

training/src/anemoi/training/config/graph/hierarchical_2level_encoder_decoder_only.yaml

@@ -2,12 +2,7 @@
 overwrite: True
 
 nodes:
-  # Data nodes
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes}
+  # Hidden nodes
   hidden_1:
     node_builder:
       _target_: anemoi.graphs.nodes.TriNodes

training/src/anemoi/training/config/graph/hierarchical_3level.yaml

@@ -3,12 +3,7 @@ overwrite: True
 
 
 nodes:
-  # Data nodes
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes}
+  # Hidden nodes
   hidden_1:
     node_builder:
       _target_: anemoi.graphs.nodes.TriNodes

training/src/anemoi/training/config/graph/limited_area.yaml

@@ -2,12 +2,6 @@
 overwrite: True
 
 nodes:
-  # Data nodes
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes}
   # Hidden nodes
   hidden:
     node_builder:

training/src/anemoi/training/config/graph/multi.yaml

@@ -4,16 +4,7 @@
 overwrite: True
 
 nodes:
-  era5:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.era5.dataset_config}
-    attributes: ${graph.attributes.nodes}
-  cerra:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.cerra.dataset_config}
-    attributes: ${graph.attributes.nodes}
+  # Hidden Nodes
   hidden:
     node_builder:
       _target_: anemoi.graphs.nodes.TriNodes

training/src/anemoi/training/config/graph/multi_scale.yaml

@@ -2,12 +2,6 @@
 overwrite: True
 
 nodes:
-  # Data nodes
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes # options: AnemoiDatasetNodes, NPZFileNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes}
   # Hidden nodes
   hidden:
     node_builder:

training/src/anemoi/training/config/graph/stretched_grid.yaml

@@ -5,11 +5,7 @@ overwrite: True
 
 
 nodes:
-  data:
-    node_builder:
-      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
-      dataset: ${dataloader.training.datasets.data.dataset_config}
-    attributes: ${graph.attributes.nodes}
+  # Hidden Nodes
   hidden:
     node_builder:
       _target_: anemoi.graphs.nodes.StretchedTriNodes

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

@@ -29,6 +29,7 @@
 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
 from anemoi.models.utils.compile import mark_for_compilation
 from anemoi.training.data.datamodule import AnemoiDatasetsDataModule
 from anemoi.training.diagnostics.callbacks import get_callbacks
@@ -166,6 +167,9 @@ def graph_data(self) -> HeteroData:
         else:
             graph_filename = None
 
+        # Introduce Data Nodes in graph config
+        self.config = integrate_data_nodes_in_config(self.config)
+
         # Create new graph
         from anemoi.graphs.create import GraphCreator
 

training/tests/integration/conftest.py

@@ -162,6 +162,7 @@ def global_config(
         "node_weights",
         "output_steps",
     ]
+
     return cfg, url, model_architecture