Datastructure for FE conditioning

config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml

@@ -0,0 +1,22 @@
+# (C) Copyright 2025 WeatherGenerator 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.
+
+day_of_year_conditioning:
+  type: time_conditioning
+  conditioning: True
+  filenames: null
+  timestep_conditioning: scalar
+  conditioning_type: time_based
+  value_type: day
+  embed:
+    net: linear
+    dim_embed: 256
+  conditioning_FE:
+    type: add
+    dim_embed: 256

config/streams/era5_1deg_forecasting_conditioning/era5.yml

@@ -0,0 +1,110 @@
+# (C) Copyright 2024 WeatherGenerator 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.
+
+ERA5 :
+  type : anemoi
+  filenames : ['aifs-ea-an-oper-0001-mars-o96-1979-2024-1h-v3-with-era51.zarr']
+  stream_id : 0
+  source_exclude : ['z', 'w_10', 'w_50', 'w_100', 'w_150', 'w_200', 'w_250', 'w_300', 'w_400', 'w_500', 'w_600', 'w_700', 'w_850', 'w_925', 'w_1000', 'skt', 'tcw', 'cp', 'tp', 'q_50', 'q_100']
+  target_exclude : ['z', 'w_10', 'w_50', 'w_100', 'w_150', 'w_200', 'w_250', 'w_300', 'w_400', 'w_500', 'w_600', 'w_700', 'w_850', 'w_925', 'w_1000', 'slor', 'sdor', 'tcw', 'cp', 'tp', 'q_50', 'q_100']
+  geoinfo_channels : ['z', 'lsm', 'slor', 'sdor', 'insolation', 'cos_local_time', 'sin_local_time', 'cos_julian_day', 'sin_julian_day']
+  loss_weight : 1.
+  location_weight : cosine_latitude
+  token_size : 8
+  tokenize_spacetime : True
+  max_num_targets: 20000
+  frequency : 06:00:00
+  embed : 
+    net : transformer
+    num_tokens : 1
+    num_heads : 8
+    dim_embed : 512
+    num_blocks : 2
+  embed_target_coords :
+    net : linear
+    dim_embed : 512
+  target_readout :
+    num_layers : 2
+    num_heads : 4
+    # sampling_rate : 0.2
+  pred_head :
+    ens_size : 1
+    num_layers : 1
+  channel_weights :
+    q_10: 0.2
+    q_50: 0.2
+    q_100: 0.23
+    q_150: 0.26
+    q_200: 0.29
+    q_250: 0.33
+    q_300: 0.36
+    q_400: 0.42
+    q_500: 0.48
+    q_600: 0.55
+    q_700: 0.61
+    q_850: 0.71
+    q_925: 0.75
+    q_1000: 0.8
+    t_10: 0.2
+    t_50: 0.2
+    t_100: 0.23
+    t_150: 0.26
+    t_200: 0.29
+    t_250: 0.33
+    t_300: 0.36
+    t_400: 0.42
+    t_500: 0.48
+    t_600: 0.55
+    t_700: 0.61
+    t_850: 0.71
+    t_925: 0.75
+    t_1000: 0.8
+    u_10: 0.2
+    u_50: 0.2
+    u_100: 0.23
+    u_150: 0.26
+    u_200: 0.29
+    u_250: 0.33
+    u_300: 0.36
+    u_400: 0.42
+    u_500: 0.48
+    u_600: 0.55
+    u_700: 0.61
+    u_850: 0.71
+    u_925: 0.75
+    u_1000: 0.8
+    v_10: 0.2
+    v_50: 0.2
+    v_100: 0.23
+    v_150: 0.26
+    v_200: 0.29
+    v_250: 0.33
+    v_300: 0.36
+    v_400: 0.42
+    v_500: 0.48
+    v_600: 0.55
+    v_700: 0.61
+    v_850: 0.71
+    v_925: 0.75
+    v_1000: 0.8
+    z_10: 0.2
+    z_50: 0.2
+    z_100: 0.23
+    z_150: 0.26
+    z_200: 0.29
+    z_250: 0.33
+    z_300: 0.36
+    z_400: 0.42
+    z_500: 0.48
+    z_600: 0.55
+    z_700: 0.61
+    z_850: 0.71
+    z_925: 0.75
+    z_1000: 0.8
+

config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml

@@ -0,0 +1,49 @@
+# (C) Copyright 2025 WeatherGenerator 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.
+
+# Example field conditioning stream: SST forcing for the atmosphere.
+# The conditioning encoder reads SST data at each forecast step and encodes it
+# into a latent representation that is injected into the ForecastingEngine.
+ERA5_SST:
+  type: anemoi
+  conditioning: True
+  filenames: ['ERA5-1deg-6h-mean-forcing-1978-2024.zarr']
+  timestep_conditioning: field
+  stream_id: 99
+  source: ['sea_surface_temperature']
+  token_size: 8
+  tokenize_spacetime: False
+  geoinfo_channels: []
+
+  # Conditioning encoder configuration.
+  # Option A — load a pre-trained encoder from a model checkpoint:
+  #   encoder:
+  #     model_id: XXX
+  #     model_epoch: XXX
+  #
+  # Option B — define an inline encoder (same parameters as the main encoder):
+  encoder:
+    ae_local_dim_embed: 512
+    ae_local_num_blocks: 2
+    ae_local_num_heads: 8
+    ae_local_num_queries: 1
+    ae_global_dim_embed: 512
+    ae_global_num_blocks: 2
+    ae_global_num_heads: 8
+
+  embed:
+    net: transformer
+    num_tokens: 1
+    num_heads: 8
+    dim_embed: 512
+    num_blocks: 2
+
+  conditioning_FE:
+    type: cross_attention
+    dim_embed: 256

config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml

@@ -0,0 +1,22 @@
+# (C) Copyright 2025 WeatherGenerator 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.
+
+time_of_day_conditioning:
+  type: time_conditioning
+  conditioning: True
+  filenames: null
+  timestep_conditioning: scalar
+  conditioning_type: time_based
+  value_type: hour
+  embed:
+    net: linear
+    dim_embed: 256
+  conditioning_FE:
+    type: add
+    dim_embed: 256

packages/readers_extra/src/weathergen/readers_extra/registry.py

@@ -4,6 +4,10 @@ def get_extra_reader(stream_type: str) -> object | None:
     # There is no sanity check on them, so they may fail at runtime during imports
 
     match stream_type:
+        case "time_conditioning":
+            from weathergen.datasets.data_reader_time_conditioning import DataReaderTimeConditioning
+
+            return DataReaderTimeConditioning
         case "iconart":
             from weathergen.readers_extra.data_reader_iconart import DataReaderIconArt
 

src/weathergen/datasets/batch.py

@@ -25,6 +25,10 @@ class SampleMetaData:
 
     global_params: dict | None = None
 
+    # Per-step scalar conditioning values, shape (num_output_steps, scalar_dim).
+    # Populated by MultiStreamDataSampler for streams with timestep_conditioning: scalar.
+    conditioning: np.typing.NDArray | None = None
+
 
 class Sample:
     # keys: stream name, values: SampleMetaData
@@ -67,6 +71,10 @@ def to_device(self, device) -> None:
                 if self.meta_info[key].mask is not None
                 else None
             )
+            if self.meta_info[key].conditioning is not None:
+                self.meta_info[key].conditioning = self.meta_info[key].conditioning.to(
+                    device, non_blocking=True
+                )
 
         for key, val in self.streams_data.items():
             if val is not None:
@@ -301,6 +309,7 @@ class ModelBatch:
     def __init__(
         self,
         stream_names: list[str],
+        conditioning_stream_names: list[str],
         num_source_samples: int,
         num_target_samples: int,
         output_offset,
@@ -320,6 +329,13 @@ def __init__(
             stream_names, num_target_samples, output_steps, self.output_idxs
         )
 
+        self.conditioning_samples = BatchSamples(
+            stream_names=conditioning_stream_names,
+            num_samples=1,
+            output_steps=output_steps,
+            output_idxs=self.output_idxs,
+        )
+
         self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32)
         self.target2source_matching_idxs = [[] for _ in range(num_target_samples)]
 
@@ -332,6 +348,9 @@ def pin_memory(self):
         # pin target samples
         self.target_samples.pin_memory()
 
+        # pin conditioning samples
+        self.conditioning_samples.pin_memory()
+
         return self
 
     def to_device(self, device):  # -> ModelBatch
@@ -341,6 +360,7 @@ def to_device(self, device):  # -> ModelBatch
 
         self.source_samples.to_device(device)
         self.target_samples.to_device(device)
+        self.conditioning_samples.to_device(device)
 
         self.device = device
 
@@ -391,6 +411,40 @@ def add_target_stream(
             )
         self.target2source_matching_idxs[target_sample_idx] = source_sample_idx
 
+    def add_scalar_conditioning_stream(self, stream_name, conditioning_values):
+        """
+        Add scalar conditioning values for all samples in the batch for a specific stream.
+        """
+        for sample in self.source_samples.samples:
+            if stream_name not in sample.meta_info:
+                sample.add_meta_info(stream_name, SampleMetaData(params={}))
+            sample.meta_info[stream_name].conditioning = conditioning_values
+
+    def add_field_conditioning_stream(self, stream_name, stream_data: StreamData):
+        """
+        Add field conditioning values for all samples in the batch for a specific stream.
+        The StreamData contains one source step per forecast step.
+        """
+        for sample in self.conditioning_samples.samples:
+            sample.streams_data[stream_name] = stream_data
+
+    def get_scalar_conditioning_values(
+        self, stream_name: str, step: int
+    ) -> np.typing.NDArray | None:
+        """
+        Get scalar conditioning values for all samples at a specific forecast step.
+
+        Returns np.ndarray of shape (num_samples, scalar_dim), or None if not available
+        for any sample.
+        """
+        values = []
+        for sample in self.source_samples.samples:
+            meta = sample.meta_info.get(stream_name)
+            if meta is None or meta.conditioning is None or step >= len(meta.conditioning):
+                return None
+            values.append(meta.conditioning[step])
+        return torch.stack(values, dim=0) if values else None
+
     def is_empty(self):
         """
         Check if batch is empty