fix(configs,schemas): hierarchical schemas (ecmwf#221)

Author: JPXKQX

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

@@ -40,6 +40,8 @@ edges:
     edge_builders:
     - _target_: anemoi.graphs.edges.KNNEdges
       num_nearest_neighbours: 3
+      source_mask_attr_name: null
+      target_mask_attr_name: null
     attributes: ${graph.attributes.edges}
 
   # Hierarchical connections: downscale

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

@@ -32,55 +32,66 @@ edges:
   # Encoder configuration
   - source_name: "data"
     target_name: "hidden_1"
-    edge_builder:
-      _target_: anemoi.graphs.edges.CutOffEdges
+    edge_builders:
+    - _target_: anemoi.graphs.edges.CutOffEdges
       cutoff_factor: 0.6
+      source_mask_attr_name: null
+      target_mask_attr_name: null
     attributes: ${graph.attributes.edges}
 
     # Decoder configuration
   - source_name: "hidden_1"
     target_name: "data"
-    edge_builder:
-      _target_: anemoi.graphs.edges.KNNEdges
+    edge_builders:
+    - _target_: anemoi.graphs.edges.KNNEdges
       num_nearest_neighbours: 3
+      source_mask_attr_name: null
+      target_mask_attr_name: null
     attributes: ${graph.attributes.edges}
 
   # Hierarchical connections: downscale
   - source_name: "hidden_1"
     target_name: "hidden_2"
-    edge_builder: ${graph.edge_builders.downscale}
+    edge_builders:
+    - ${graph.edge_builders.downscale}
     attributes: ${graph.attributes.edges}
 
   - source_name: "hidden_2"
     target_name: "hidden_3"
-    edge_builder: ${graph.edge_builders.downscale}
+    edge_builders:
+    - ${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}
+    edge_builders:
+    - ${graph.edge_builders.upscale}
     attributes: ${graph.attributes.edges}
 
   - source_name: "hidden_2"
     target_name: "hidden_1"
-    edge_builder: ${graph.edge_builders.upscale}
+    edge_builders:
+    - ${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}
+    edge_builders:
+    - ${graph.edge_builders.process}
     attributes: ${graph.attributes.edges}
 
   - source_name: "hidden_2"
     target_name: "hidden_2"
-    edge_builder: ${graph.edge_builders.process}
+    edge_builders:
+    - ${graph.edge_builders.process}
     attributes: ${graph.attributes.edges}
 
   - source_name: "hidden_3"
     target_name: "hidden_3"
-    edge_builder: ${graph.edge_builders.process}
+    edge_builders:
+    - ${graph.edge_builders.process}
     attributes: ${graph.attributes.edges}
 
 

training/src/anemoi/training/schemas/graphs/base_graph.py

@@ -55,7 +55,7 @@ class BaseGraphSchema(PydanticBaseModel):
     post_processors: list[ProcessorSchemas] = Field(default_factory=list)
     data: str = Field(example="data")
     "Key name for the data nodes. Default to 'data'."
-    hidden: str = Field(example="hidden")
+    hidden: Union[str, list[str]] = Field(example="hidden")
     "Key name for the hidden nodes. Default to 'hidden'."
     # TODO(Helen): Needs to be adjusted for more complex graph setups
 

training/src/anemoi/training/schemas/models/models.py

@@ -35,8 +35,11 @@
 
 class DefinedModels(str, Enum):
     ANEMOI_MODEL_ENC_PROC_DEC = "anemoi.models.models.encoder_processor_decoder.AnemoiModelEncProcDec"
+    ANEMOI_MODEL_ENC_PROC_DEC_SHORT = "anemoi.models.models.AnemoiModelEncProcDec"
     ANEMOI_MODEL_ENC_HIERPROC_DEC = "anemoi.models.models.hierarchical.AnemoiModelEncProcDecHierarchical"
+    ANEMOI_MODEL_ENC_HIERPROC_DEC_SHORT = "anemoi.models.models.AnemoiModelEncProcDecHierarchical"
     ANEMOI_MODEL_INTERPENC_PROC_DEC = "anemoi.models.models.interpolator.AnemoiModelEncProcDecInterpolator"
+    ANEMOI_MODEL_INTERPENC_PROC_DEC_SHORT = "anemoi.models.models.AnemoiModelEncProcDecInterpolator"
 
 
 class Model(BaseModel):