chore(training): Add default config files for 2 and 3 level hierarchi…

…cal processors (#50)

* chore: hierarchical configs

* fix: update default config

* fix: update config files

* udpate changelog

training/CHANGELOG.md

@@ -35,6 +35,7 @@ Keep it human-readable, your future self will thank you!
 - Add supporting arrrays (numpy) to checkpoint
 - Support for masking out unconnected nodes in LAM [#171](https://github.com/ecmwf/anemoi-training/pull/171)
 - Improved validation metrics, allow 'all' to be scaled [#202](https://github.com/ecmwf/anemoi-training/pull/202)
+- Added default configuration files for hierarchical processor [175](https://github.com/ecmwf/anemoi-training/pull/175)
 
 ### Changed
 

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

@@ -0,0 +1,90 @@
+---
+overwrite: True
+
+data: "data"
+hidden:
+ - "hidden_1"
+ - "hidden_2"
+
+
+nodes:
+  # Data nodes
+  data:
+    node_builder:
+      _target_: anemoi.graphs.nodes.ZarrDatasetNodes
+      dataset: ${dataloader.dataset}
+    attributes: ${graph.attributes.nodes}
+  hidden_1:
+    node_builder:
+      _target_: anemoi.graphs.nodes.TriNodes
+      resolution: 5
+  hidden_2:
+    node_builder:
+      _target_: anemoi.graphs.nodes.TriNodes
+      resolution: 4
+
+edges:
+  # Encoder configuration
+  - source_name: "data"
+    target_name: "hidden_1"
+    edge_builder:
+      _target_: anemoi.graphs.edges.CutOffEdges
+      cutoff_factor: 0.6
+    attributes: ${graph.attributes.edges}
+
+    # Decoder configuration
+  - source_name: "hidden_1"
+    target_name: "data"
+    edge_builder:
+      _target_: anemoi.graphs.edges.KNNEdges
+      num_nearest_neighbours: 3
+    attributes: ${graph.attributes.edges}
+
+  # Hierarchical connections: downscale
+  - source_name: "hidden_1"
+    target_name: "hidden_2"
+    edge_builder: ${graph.edge_builders.downscale}
+    attributes: ${graph.attributes.edges}
+
+  # Hierarchical connections: upscale
+  - source_name: "hidden_2"
+    target_name: "hidden_1"
+    edge_builder: ${graph.edge_builders.upscale}
+    attributes: ${graph.attributes.edges}
+
+  # Hierarchical connections: same level
+  - source_name: "hidden_1"
+    target_name: "hidden_1"
+    edge_builder: ${graph.edge_builders.process}
+    attributes: ${graph.attributes.edges}
+
+  - source_name: "hidden_2"
+    target_name: "hidden_2"
+    edge_builder: ${graph.edge_builders.process}
+    attributes: ${graph.attributes.edges}
+
+
+#############
+edge_builders:
+  downscale:
+    _target_: anemoi.graphs.edges.CutOffEdges
+    cutoff_factor: 1.5
+  process:
+    _target_: anemoi.graphs.edges.MultiScaleEdges
+    x_hops: 1
+  upscale:
+    _target_: anemoi.graphs.edges.KNNEdges
+    num_nearest_neighbours: 5
+
+attributes:
+  nodes:
+    area_weight:
+      _target_: anemoi.graphs.nodes.attributes.AreaWeights # options: Area, Uniform
+      norm: unit-max  # options: l1, l2, unit-max, unit-sum, unit-std
+  edges:
+    edge_dirs:
+      _target_: anemoi.graphs.edges.attributes.EdgeDirection
+      norm: unit-std
+    edge_length:
+      _target_: anemoi.graphs.edges.attributes.EdgeLength
+      norm: unit-std

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

@@ -0,0 +1,110 @@
+---
+overwrite: True
+
+data: "data"
+hidden:
+ - "hidden_1"
+ - "hidden_2"
+ - "hidden_3"
+
+
+nodes:
+  # Data nodes
+  data:
+    node_builder:
+      _target_: anemoi.graphs.nodes.ZarrDatasetNodes
+      dataset: ${dataloader.dataset}
+    attributes: ${graph.attributes.nodes}
+  hidden_1:
+    node_builder:
+      _target_: anemoi.graphs.nodes.TriNodes
+      resolution: 5
+  hidden_2:
+    node_builder:
+      _target_: anemoi.graphs.nodes.TriNodes
+      resolution: 4
+  hidden_3:
+    node_builder:
+      _target_: anemoi.graphs.nodes.TriNodes
+      resolution: 3
+
+edges:
+  # Encoder configuration
+  - source_name: "data"
+    target_name: "hidden_1"
+    edge_builder:
+      _target_: anemoi.graphs.edges.CutOffEdges
+      cutoff_factor: 0.6
+    attributes: ${graph.attributes.edges}
+
+    # Decoder configuration
+  - source_name: "hidden_1"
+    target_name: "data"
+    edge_builder:
+      _target_: anemoi.graphs.edges.KNNEdges
+      num_nearest_neighbours: 3
+    attributes: ${graph.attributes.edges}
+
+  # Hierarchical connections: downscale
+  - source_name: "hidden_1"
+    target_name: "hidden_2"
+    edge_builder: ${graph.edge_builders.downscale}
+    attributes: ${graph.attributes.edges}
+
+  - source_name: "hidden_2"
+    target_name: "hidden_3"
+    edge_builder: ${graph.edge_builders.downscale}
+    attributes: ${graph.attributes.edges}
+
+  # Hierarchical connections: upscale
+  - source_name: "hidden_3"
+    target_name: "hidden_2"
+    edge_builder: ${graph.edge_builders.upscale}
+    attributes: ${graph.attributes.edges}
+
+  - source_name: "hidden_2"
+    target_name: "hidden_1"
+    edge_builder: ${graph.edge_builders.upscale}
+    attributes: ${graph.attributes.edges}
+
+  # Hierarchical connections: same level
+  - source_name: "hidden_1"
+    target_name: "hidden_1"
+    edge_builder: ${graph.edge_builders.process}
+    attributes: ${graph.attributes.edges}
+
+  - source_name: "hidden_2"
+    target_name: "hidden_2"
+    edge_builder: ${graph.edge_builders.process}
+    attributes: ${graph.attributes.edges}
+
+  - source_name: "hidden_3"
+    target_name: "hidden_3"
+    edge_builder: ${graph.edge_builders.process}
+    attributes: ${graph.attributes.edges}
+
+
+#############
+edge_builders:
+  downscale:
+    _target_: anemoi.graphs.edges.CutOffEdges
+    cutoff_factor: 1.5
+  process:
+    _target_: anemoi.graphs.edges.MultiScaleEdges
+    x_hops: 1
+  upscale:
+    _target_: anemoi.graphs.edges.KNNEdges
+    num_nearest_neighbours: 5
+
+attributes:
+  nodes:
+    area_weight:
+      _target_: anemoi.graphs.nodes.attributes.AreaWeights # options: Area, Uniform
+      norm: unit-max  # options: l1, l2, unit-max, unit-sum, unit-std
+  edges:
+    edge_dirs:
+      _target_: anemoi.graphs.edges.attributes.EdgeDirection
+      norm: unit-std
+    edge_length:
+      _target_: anemoi.graphs.edges.attributes.EdgeLength
+      norm: unit-std

training/src/anemoi/training/config/hierarchical.yaml

@@ -0,0 +1,23 @@
+defaults:
+- data: zarr
+- dataloader: native_grid
+- diagnostics: evaluation
+- hardware: example
+- graph: hierarchical_3level
+- model: graphtransformer
+- training: default
+- _self_
+
+
+### This file is for local experimentation.
+##  When you commit your changes, assign the new features and keywords
+##  to the correct defaults.
+# For example to change from default GPU count:
+# hardware:
+#   num_gpus_per_node: 1
+
+model:
+  model:
+    _target_: anemoi.models.models.AnemoiModelEncProcDecHierarchical
+  enable_hierarchical_level_processing: True
+  level_process_num_layers: 2

training/src/anemoi/training/config/model/gnn.yaml

@@ -2,7 +2,7 @@ activation: GELU
 num_channels: 512
 
 model:
-  _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec
+  _target_: anemoi.models.models.AnemoiModelEncProcDec
 
 processor:
   _target_: anemoi.models.layers.processor.GNNProcessor

training/src/anemoi/training/config/model/graphtransformer.yaml

@@ -2,7 +2,7 @@ activation: GELU
 num_channels: 1024
 
 model:
-  _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec
+  _target_: anemoi.models.models.AnemoiModelEncProcDec
 
 processor:
   _target_: anemoi.models.layers.processor.GraphTransformerProcessor

training/src/anemoi/training/config/model/transformer.yaml

@@ -2,7 +2,7 @@ activation: GELU
 num_channels: 1024
 
 model:
-  _target_: anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec
+  _target_: anemoi.models.models.AnemoiModelEncProcDec
 
 processor:
   _target_: anemoi.models.layers.processor.TransformerProcessor