Merge pull request #10 from ecmwf/activation_functions_config

added config for bounding

CHANGELOG.md

@@ -66,6 +66,8 @@ Keep it human-readable, your future self will thank you!
 
 - Introduction of remapper to anemoi-models leads to changes in the data indices. Some preprocessors cannot be applied in-place anymore.
 
+- Variable Bounding as configurable model layers [#13](https://github.com/ecmwf/anemoi-models/issues/13)
+
 #### Functionality
 
 - Enable the callback for plotting a histogram for variables containing NaNs

src/anemoi/training/config/data/zarr.yaml

@@ -30,6 +30,21 @@ remapped:
 
 normalizer:
   default: "mean-std"
+
+  # Remap cp statistics to those of tp when using FractionBounding. This ensures
+  # that cp, as a fraction of tp, remains consistent with tp's scale and statistics.
+  # NOTE: This remap should only be applied if FractionBounding is enabled for cp.
+  # remap:
+  #   cp: tp
+
+  # Standardization applied to tp and cp variables. Ensure that if cp is bounded
+  # as a fraction of tp, both variables are normalized using these shared statistics.
+  # "Std" normalization is preferred here over "mean-std" to avoid shifting of the
+  # zero value in the normalized space.
+  std:
+  - "tp"
+  # - "cp"
+
   min-max:
   max:
   - "sdor"

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

@@ -46,3 +46,26 @@ attributes:
   nodes: []
 
 node_loss_weight: area_weight
+
+# Bounding configuration
+bounding: #These are applied in order
+
+  # Bound tp (total precipitation) with a Relu bounding layer
+  # ensuring a range of [0, infinity) to avoid negative precipitation values.
+  - _target_: anemoi.models.layers.bounding.ReluBounding #[0, infinity)
+    variables:
+    - tp
+
+  # [OPTIONAL] Bound cp (convective precipitation) as a fraction of tp.
+  # This guarantees that cp is physically consistent with tp by restricting cp
+  # to a fraction of tp [0 to 1]. Uncomment the lines below to apply.
+  # NOTE: If this bounding strategy is used, the normalization of cp must be
+  # changed to "std" normalization, and the "cp" statistics should be remapped
+  # to those of tp to ensure consistency.
+
+  # - _target_: anemoi.models.layers.bounding.FractionBounding # fraction of tp
+  #   variables:
+  #   - cp
+  #   min_val: 0
+  #   max_val: 1
+  #   total_var: tp

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

@@ -51,3 +51,26 @@ attributes:
   nodes: []
 
 node_loss_weight: area_weight
+
+# Bounding configuration
+bounding: #These are applied in order
+
+  # Bound tp (total precipitation) with a Relu bounding layer
+  # ensuring a range of [0, infinity) to avoid negative precipitation values.
+  - _target_: anemoi.models.layers.bounding.ReluBounding #[0, infinity)
+    variables:
+    - tp
+
+  # [OPTIONAL] Bound cp (convective precipitation) as a fraction of tp.
+  # This guarantees that cp is physically consistent with tp by restricting cp
+  # to a fraction of tp [0 to 1]. Uncomment the lines below to apply.
+  # NOTE: If this bounding strategy is used, the normalization of cp must be
+  # changed to "std" normalization, and the "cp" statistics should be remapped
+  # to those of tp to ensure consistency.
+
+  # - _target_: anemoi.models.layers.bounding.FractionBounding # fraction of tp
+  #   variables:
+  #   - cp
+  #   min_val: 0
+  #   max_val: 1
+  #   total_var: tp

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

@@ -50,3 +50,26 @@ attributes:
   nodes: []
 
 node_loss_weight: area_weight
+
+# Bounding configuration
+bounding: #These are applied in order
+
+  # Bound tp (total precipitation) with a Relu bounding layer
+  # ensuring a range of [0, infinity) to avoid negative precipitation values.
+  - _target_: anemoi.models.layers.bounding.ReluBounding #[0, infinity)
+    variables:
+    - tp
+
+  # [OPTIONAL] Bound cp (convective precipitation) as a fraction of tp.
+  # This guarantees that cp is physically consistent with tp by restricting cp
+  # to a fraction of tp [0 to 1]. Uncomment the lines below to apply.
+  # NOTE: If this bounding strategy is used, the normalization of cp must be
+  # changed to "std" normalization, and the "cp" statistics should be remapped
+  # to those of tp to ensure consistency.
+
+  # - _target_: anemoi.models.layers.bounding.FractionBounding # fraction of tp
+  #   variables:
+  #   - cp
+  #   min_val: 0
+  #   max_val: 1
+  #   total_var: tp