From b2fb8a7bef9e55aadfa47df2afb8fd8a49030cef Mon Sep 17 00:00:00 2001
From: Savvas Melidonis <s.melidonis@fz-juelich.de>
Date: Tue, 14 Jul 2026 10:51:06 +0200
Subject: [PATCH 01/15] changes to load conditioning data

---
 .../src/weathergen/readers_extra/registry.py  |   4 +
 src/weathergen/datasets/batch.py              |  67 +++++++++
 .../datasets/multi_stream_data_sampler.py     | 133 +++++++++++++++---
 src/weathergen/model/engines.py               |   2 +-
 src/weathergen/model/model.py                 |   5 +-
 src/weathergen/train/trainer.py               |   6 +
 6 files changed, 193 insertions(+), 24 deletions(-)

diff --git a/packages/readers_extra/src/weathergen/readers_extra/registry.py b/packages/readers_extra/src/weathergen/readers_extra/registry.py
index b13ee46550..b219bff613 100644
--- a/packages/readers_extra/src/weathergen/readers_extra/registry.py
+++ b/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 "synthetic":
+            from weathergen.datasets.data_reader_synthetic import DataReaderSynthetic
+
+            return DataReaderSynthetic
         case "iconart":
             from weathergen.readers_extra.data_reader_iconart import DataReaderIconArt
 
diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py
index ea6a0b26ab..a81d5001ad 100644
--- a/src/weathergen/datasets/batch.py
+++ b/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
@@ -43,6 +47,15 @@ def pin_memory(self):
                 if stream_data is not None and hasattr(stream_data, "pin_memory"):
                     stream_data.pin_memory()
 
+        # Pin StreamData objects in conditioning_streams_data
+        if hasattr(self, "conditioning_streams_data") and isinstance(
+            self.conditioning_streams_data, dict
+        ):
+            for _name, steps in self.conditioning_streams_data.items():
+                for sd in steps:
+                    if sd is not None and hasattr(sd, "pin_memory"):
+                        sd.pin_memory()
+
         # Pin tensors in meta_info
         if hasattr(self, "meta_info") and isinstance(self.meta_info, dict):
             for _key, meta_data in self.meta_info.items():
@@ -60,6 +73,10 @@ def __init__(self, stream_names: list[str]) -> None:
         for stream_name in stream_names:
             self.streams_data[stream_name] = None
 
+        # Field conditioning stream data, keyed by stream name then indexed by forecast step.
+        # Populated by MultiStreamDataSampler for streams with timestep_conditioning: field.
+        self.conditioning_streams_data: dict[str, list[StreamData | None]] = {}
+
     def to_device(self, device) -> None:
         for key in self.meta_info.keys():
             self.meta_info[key].mask = (
@@ -72,6 +89,11 @@ def to_device(self, device) -> None:
             if val is not None:
                 self.streams_data[key] = val.to_device(device)
 
+        for name, steps in self.conditioning_streams_data.items():
+            self.conditioning_streams_data[name] = [
+                sd.to_device(device) if sd is not None else None for sd in steps
+            ]
+
     def is_empty(self) -> bool:
         """
         Check if sample is empty
@@ -156,6 +178,27 @@ def get_num_target_steps(self) -> int:
         ]
         return min(lens) if len(lens) > 0 else 0
 
+    def add_conditioning_stream_data(
+        self, stream_name: str, step: int, stream_data: StreamData | None
+    ) -> None:
+        """
+        Add StreamData for field conditioning stream @stream_name at forecast step @step
+        """
+        if stream_name not in self.conditioning_streams_data:
+            self.conditioning_streams_data[stream_name] = []
+        steps = self.conditioning_streams_data[stream_name]
+        while len(steps) <= step:
+            steps.append(None)
+        steps[step] = stream_data
+
+    def get_conditioning_stream_data(self, stream_name: str, step: int) -> StreamData | None:
+        """
+        Get StreamData for field conditioning stream @stream_name at forecast step @step
+        """
+        steps = self.conditioning_streams_data.get(stream_name)
+        if steps is None or step >= len(steps):
+            return None
+        return steps[step]
 
 class BatchSamples:
     """
@@ -391,6 +434,30 @@ def add_target_stream(
             )
         self.target2source_matching_idxs[target_sample_idx] = source_sample_idx
 
+    def add_scalar_conditioning(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 get_scalar_conditioning(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.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 np.stack(values, axis=0) if values else None
+
     def is_empty(self):
         """
         Check if batch is empty
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index bd4a8ad87e..91ccd32938 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -144,7 +144,13 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
         # check samples per mini epoch
         self.samples_per_mini_epoch = mode_cfg.samples_per_mini_epoch
         self.check_samples(self._get_fsm())
-        self.streams_datasets = self._init_stream_datasets(cf)
+        (
+            self.streams_datasets,
+            self.scalar_conditioning_datasets,
+            self.field_conditioning_datasets,
+        ) = self._init_stream_datasets(cf)
+        self.scalar_conditioning_stream_names = list(self.scalar_conditioning_datasets.keys())
+        self.field_conditioning_stream_names = list(self.field_conditioning_datasets.keys())
 
         # RNG seed setup
         rs = cf.data_loading.rng_seed
@@ -217,7 +223,9 @@ def _calc_baseperms(self, fsm: int) -> np.typing.NDArray:
 
         return np.arange(self.max_input_steps, perms_len)
 
-    def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]:
+    def _init_stream_datasets(
+        self, cf
+    ) -> tuple[dict[StreamName, _Stream], dict[StreamName, _Stream], dict[StreamName, _Stream]]:
         """Load dataset readers for all streams from config."""
         streams_datasets: dict[StreamName, _Stream] = {}
         for stream_name, stream_info in cf.streams.items():
@@ -242,32 +250,44 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]:
                         f"for stream name '{stream_name}'."
                         raise ValueError(msg)
 
-            for fname in stream_info.get("filenames", [pathlib.Path()]):
-                fname = pathlib.Path(fname)
-                # dont check if file exists since zarr stores might be directories
-                if fname.exists():
-                    # check if fname is a valid path to allow for simple overwriting
-                    filename = fname
-                else:
-                    filenames = [pathlib.Path(path) / fname for path in cf.data_paths]
-
-                    filename = next((f for f in filenames if f.exists()), None)
-                    if filename is None:
-                        msg = (
-                            f"Did not find input data for {stream_info['type']} "
-                            f"stream '{stream_name}': {filenames}."
-                        )
-                        raise FileNotFoundError(msg)
-
+            filenames_cfg = stream_info.get("filenames")
+            if filenames_cfg is None:
+                # No file required (e.g. synthetic streams)
                 ds_type = stream_info["type"]
                 if is_root():
                     logger.info(
                         f"Opening dataset with type: {ds_type}"
                         + f" from stream config {stream_name}.",
                     )
-                ds = dataset(filename=filename, **kwargs)
-
+                ds = dataset(filename=None, **kwargs)
                 streams_datasets[stream_name].readers += [ds]
+            else:
+                for fname in filenames_cfg:
+                    fname = pathlib.Path(fname)
+                    # dont check if file exists since zarr stores might be directories
+                    if fname.exists():
+                        # check if fname is a valid path to allow for simple overwriting
+                        filename = fname
+                    else:
+                        filenames = [pathlib.Path(path) / fname for path in cf.data_paths]
+
+                        filename = next((f for f in filenames if f.exists()), None)
+                        if filename is None:
+                            msg = (
+                                f"Did not find input data for {stream_info['type']} "
+                                f"stream '{stream_name}': {filenames}."
+                            )
+                            raise FileNotFoundError(msg)
+
+                    ds_type = stream_info["type"]
+                    if is_root():
+                        logger.info(
+                            f"Opening dataset with type: {ds_type}"
+                            + f" from stream config {stream_name}.",
+                        )
+                    ds = dataset(filename=filename, **kwargs)
+
+                    streams_datasets[stream_name].readers += [ds]
 
             stream_info[str(self._stage) + "_source_channels"] = ds.source_channels
             stream_info[str(self._stage) + "_target_channels"] = ds.target_channels
@@ -277,7 +297,20 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]:
                 else [1.0 for _ in ds.target_channels]
             )
 
-        return streams_datasets
+        # Separate streams by timestep_conditioning type
+        regular: dict[StreamName, _Stream] = {}
+        scalar_conditioning: dict[StreamName, _Stream] = {}
+        field_conditioning: dict[StreamName, _Stream] = {}
+        for name, stream in streams_datasets.items():
+            tc = stream.info.get("timestep_conditioning")
+            if tc == "scalar":
+                scalar_conditioning[name] = stream
+            elif tc == "field":
+                field_conditioning[name] = stream
+            else:
+                regular[name] = stream
+
+        return regular, scalar_conditioning, field_conditioning
 
     def reset(self) -> tuple[Sequence[int], Sequence[int]]:
         """
@@ -756,6 +789,62 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
         target_in_steps = 1 if len(target_in_steps) == 0 else target_in_steps.max().item()
         batch = self._preprocess_model_batch(batch, source_in_steps, target_in_steps)
 
+        if self.scalar_conditioning_stream_names:
+            batch = self._build_scalar_conditioning_data(batch, idx, num_forecast_steps)
+
+        if self.field_conditioning_stream_names:
+            batch = self._build_field_conditioning_data(batch, idx, num_forecast_steps)
+
+        return batch
+
+    def _build_scalar_conditioning_data(
+        self, batch: ModelBatch, idx: int, num_forecast_steps: int
+    ) -> ModelBatch:
+        """Collect per-step scalar conditioning values and store in sample meta_info."""
+        num_output_steps = self._get_output_length(num_forecast_steps)
+        for stream_name, stream_ds in self.scalar_conditioning_datasets.items():
+            step_values = []
+            for timestep_idx in range(self.output_offset, num_output_steps):
+                step_dt = idx + (self.time_step * timestep_idx) // self.step_timedelta
+                rdata = stream_ds.readers[0].get_source(step_dt)
+                step_values.append(rdata.data.flatten().copy())
+            conditioning_values = (
+                np.stack(step_values, axis=0) if step_values else np.zeros((0, 1), dtype=np.float32)
+            )
+            batch.add_scalar_conditioning(stream_name, conditioning_values)
+        return batch
+
+    def _build_field_conditioning_data(
+        self, batch: ModelBatch, idx: int, num_forecast_steps: int
+    ) -> ModelBatch:
+        """Collect per-step field conditioning data and store in conditioning_streams_data."""
+        num_output_steps = self._get_output_length(num_forecast_steps)
+        for stream_name, stream_ds in self.field_conditioning_datasets.items():
+            stream_info = stream_ds.info
+            for step, timestep_idx in enumerate(range(self.output_offset, num_output_steps)):
+                step_dt = idx + (self.time_step * timestep_idx) // self.step_timedelta
+                rdata = collect_datasources(stream_ds.readers, step_dt, "source", self.rng)
+                if rdata.is_empty():
+                    stream_data = None
+                else:
+                    token_data_list = self.tokenizer.get_tokens_windows(stream_info, [rdata], False)
+                    token_data = token_data_list[0]
+                    if token_data[0] is None:
+                        stream_data = None
+                    else:
+                        time_win = self.time_window_handler.window(step_dt)
+                        src_cells, src_lens = self.tokenizer.get_source(
+                            stream_info,
+                            rdata,
+                            token_data,
+                            (time_win.start, time_win.end),
+                            None,
+                        )
+                        stream_data = StreamData(step_dt, 1, 1, self.num_healpix_cells)
+                        stream_data.add_source(self._stage, 0, rdata, src_lens, src_cells, False)
+                # Field conditioning is the same for all views within a batch step
+                for sample in batch.source_samples.samples:
+                    sample.add_conditioning_stream_data(stream_name, step, stream_data)
         return batch
 
     def __iter__(self) -> ModelBatch:
diff --git a/src/weathergen/model/engines.py b/src/weathergen/model/engines.py
index fde31213b6..0cde93d5fc 100644
--- a/src/weathergen/model/engines.py
+++ b/src/weathergen/model/engines.py
@@ -48,7 +48,7 @@ def __init__(self, cf: Config, sources_size) -> None:
         self.dtype = get_dtype(self.cf.mixed_precision_dtype)
         self.sources_size = sources_size  # KCT:iss130, what is this?
         self.embeds = torch.nn.ModuleDict()
-        self.streams = cf.streams
+        self.streams = {k: v for k, v in cf.streams.items() if not v.get("timestep_conditioning")}
 
         for i, (stream_name, si) in enumerate(self.streams.items()):
             if si.get("diagnostic", False) or self.sources_size[i] == 0:
diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py
index f4035ea467..0511bba91d 100644
--- a/src/weathergen/model/model.py
+++ b/src/weathergen/model/model.py
@@ -327,7 +327,10 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord
         self.forecast_engine: ForecastingEngine | IdentityEngine | None = None
         self.pred_heads = None
         self.q_cells: torch.Tensor | None = None
-        self.streams: dict[str, typing.Any] = cf.streams
+        self.streams: dict[str, typing.Any] = {
+            k: v for k, v in cf.streams.items() if not v.get("timestep_conditioning")
+        }
+
         self.target_token_engines = None
 
         assert cf.get("forecast", {}).get("att_dense_rate", 1.0) == 1.0, (
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 276da0bd67..0923f6fb59 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -11,6 +11,7 @@
 import contextlib
 import copy
 import logging
+import os
 import time
 from math import sqrt
 
@@ -342,6 +343,11 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
             self.grad_scaler = torch.amp.GradScaler("cuda")
         assert len(self.dataset) > 0, f"No data found in {self.dataset}"
 
+        if os.environ.get("DEBUG_DATASET"):
+            self.dataset.reset()
+            breakpoint()
+            batch = self.dataset._get_batch(0, 2)
+
         # lr is updated after each batch so account for this
         # TODO: conf should be read-only, do not modify the conf in flight
         len_ds = len(self.dataset)

From c503816725923fa35f73123974513044746b0b94 Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Fri, 31 Jul 2026 11:55:50 +0200
Subject: [PATCH 02/15] cleaning some code

---
 .../src/weathergen/readers_extra/registry.py  |  6 +-
 src/weathergen/datasets/batch.py              | 35 ++++++++-
 .../datasets/multi_stream_data_sampler.py     | 73 ++++++++++---------
 src/weathergen/datasets/tokenizer_masking.py  |  3 +
 src/weathergen/train/trainer.py               |  6 --
 5 files changed, 75 insertions(+), 48 deletions(-)

diff --git a/packages/readers_extra/src/weathergen/readers_extra/registry.py b/packages/readers_extra/src/weathergen/readers_extra/registry.py
index b219bff613..4d8f00ae8b 100644
--- a/packages/readers_extra/src/weathergen/readers_extra/registry.py
+++ b/packages/readers_extra/src/weathergen/readers_extra/registry.py
@@ -4,10 +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 "synthetic":
-            from weathergen.datasets.data_reader_synthetic import DataReaderSynthetic
+        case "time_conditioning":
+            from weathergen.datasets.data_reader_time_conditioning import DataReaderTimeConditioning
 
-            return DataReaderSynthetic
+            return DataReaderTimeConditioning
         case "iconart":
             from weathergen.readers_extra.data_reader_iconart import DataReaderIconArt
 
diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py
index a81d5001ad..85ab148db4 100644
--- a/src/weathergen/datasets/batch.py
+++ b/src/weathergen/datasets/batch.py
@@ -182,7 +182,7 @@ def add_conditioning_stream_data(
         self, stream_name: str, step: int, stream_data: StreamData | None
     ) -> None:
         """
-        Add StreamData for field conditioning stream @stream_name at forecast step @step
+        Add StreamData for field conditioning stream @stream_name at forecast step @step to sample
         """
         if stream_name not in self.conditioning_streams_data:
             self.conditioning_streams_data[stream_name] = []
@@ -193,13 +193,14 @@ def add_conditioning_stream_data(
 
     def get_conditioning_stream_data(self, stream_name: str, step: int) -> StreamData | None:
         """
-        Get StreamData for field conditioning stream @stream_name at forecast step @step
+        Get StreamData for field conditioning stream @stream_name at forecast step @step from sample
         """
         steps = self.conditioning_streams_data.get(stream_name)
         if steps is None or step >= len(steps):
             return None
         return steps[step]
 
+
 class BatchSamples:
     """
     Container for source or target samples
@@ -434,7 +435,7 @@ def add_target_stream(
             )
         self.target2source_matching_idxs[target_sample_idx] = source_sample_idx
 
-    def add_scalar_conditioning(self, stream_name, conditioning_values):
+    def add_scalar_conditioning_stream(self, stream_name, conditioning_values):
         """
         Add scalar conditioning values for all samples in the batch for a specific stream.
         """
@@ -443,7 +444,16 @@ def add_scalar_conditioning(self, stream_name, conditioning_values):
                 sample.add_meta_info(stream_name, SampleMetaData(params={}))
             sample.meta_info[stream_name].conditioning = conditioning_values
 
-    def get_scalar_conditioning(self, stream_name: str, step: int) -> np.typing.NDArray | None:
+    def add_field_conditioning_stream(self, stream_name, step: int, stream_data: StreamData):
+        """
+        Add field conditioning values for all samples in the batch for a specific stream.
+        """
+        for sample in self.source_samples.samples:
+            sample.add_conditioning_stream_data(stream_name, step, 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.
 
@@ -458,6 +468,23 @@ def get_scalar_conditioning(self, stream_name: str, step: int) -> np.typing.NDAr
             values.append(meta.conditioning[step])
         return np.stack(values, axis=0) if values else None
 
+    def get_field_conditioning_values(
+        self, stream_name: str, step: int
+    ) -> np.typing.NDArray | None:
+        """
+        Get field conditioning values for all samples at a specific forecast step.
+
+        Returns np.ndarray of shape (num_samples, ...) or None if not available
+        for any sample.
+        """
+        values = []
+        for sample in self.samples:
+            stream_data = sample.get_conditioning_stream_data(stream_name, step)
+            if stream_data is None:
+                return None
+            values.append(stream_data.data)
+        return np.stack(values, axis=0) if values else None
+
     def is_empty(self):
         """
         Check if batch is empty
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index 91ccd32938..03b5c36d84 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -250,44 +250,49 @@ def _init_stream_datasets(
                         f"for stream name '{stream_name}'."
                         raise ValueError(msg)
 
-            filenames_cfg = stream_info.get("filenames")
-            if filenames_cfg is None:
-                # No file required (e.g. synthetic streams)
-                ds_type = stream_info["type"]
-                if is_root():
-                    logger.info(
-                        f"Opening dataset with type: {ds_type}"
-                        + f" from stream config {stream_name}.",
-                    )
-                ds = dataset(filename=None, **kwargs)
-                streams_datasets[stream_name].readers += [ds]
-            else:
-                for fname in filenames_cfg:
-                    fname = pathlib.Path(fname)
-                    # dont check if file exists since zarr stores might be directories
-                    if fname.exists():
-                        # check if fname is a valid path to allow for simple overwriting
-                        filename = fname
-                    else:
-                        filenames = [pathlib.Path(path) / fname for path in cf.data_paths]
-
-                        filename = next((f for f in filenames if f.exists()), None)
-                        if filename is None:
-                            msg = (
-                                f"Did not find input data for {stream_info['type']} "
-                                f"stream '{stream_name}': {filenames}."
-                            )
-                            raise FileNotFoundError(msg)
+            filenames_cfg = stream_info.get("filenames", [pathlib.Path()])
+            conditioning_cfg = stream_info.get("conditioning", False)
 
+            if conditioning_cfg:
+                if filenames_cfg is None or len(filenames_cfg) == 0:
                     ds_type = stream_info["type"]
                     if is_root():
                         logger.info(
-                            f"Opening dataset with type: {ds_type}"
+                            f"Opening conditioning dataset with type: {ds_type}"
                             + f" from stream config {stream_name}.",
                         )
-                    ds = dataset(filename=filename, **kwargs)
-
+                    ds = dataset(filename=None, **kwargs)
                     streams_datasets[stream_name].readers += [ds]
+                    continue
+                else:
+                    pass
+
+            for fname in filenames_cfg:
+                fname = pathlib.Path(fname)
+                # dont check if file exists since zarr stores might be directories
+                if fname.exists():
+                    # check if fname is a valid path to allow for simple overwriting
+                    filename = fname
+                else:
+                    filenames = [pathlib.Path(path) / fname for path in cf.data_paths]
+
+                    filename = next((f for f in filenames if f.exists()), None)
+                    if filename is None:
+                        msg = (
+                            f"Did not find input data for {stream_info['type']} "
+                            f"stream '{stream_name}': {filenames}."
+                        )
+                        raise FileNotFoundError(msg)
+
+                ds_type = stream_info["type"]
+                if is_root():
+                    logger.info(
+                        f"Opening dataset with type: {ds_type}"
+                        + f" from stream config {stream_name}.",
+                    )
+                ds = dataset(filename=filename, **kwargs)
+
+                streams_datasets[stream_name].readers += [ds]
 
             stream_info[str(self._stage) + "_source_channels"] = ds.source_channels
             stream_info[str(self._stage) + "_target_channels"] = ds.target_channels
@@ -811,7 +816,7 @@ def _build_scalar_conditioning_data(
             conditioning_values = (
                 np.stack(step_values, axis=0) if step_values else np.zeros((0, 1), dtype=np.float32)
             )
-            batch.add_scalar_conditioning(stream_name, conditioning_values)
+            batch.add_scalar_conditioning_stream(stream_name, conditioning_values)
         return batch
 
     def _build_field_conditioning_data(
@@ -842,9 +847,7 @@ def _build_field_conditioning_data(
                         )
                         stream_data = StreamData(step_dt, 1, 1, self.num_healpix_cells)
                         stream_data.add_source(self._stage, 0, rdata, src_lens, src_cells, False)
-                # Field conditioning is the same for all views within a batch step
-                for sample in batch.source_samples.samples:
-                    sample.add_conditioning_stream_data(stream_name, step, stream_data)
+                batch.add_field_conditioning_stream(stream_name, step, stream_data)
         return batch
 
     def __iter__(self) -> ModelBatch:
diff --git a/src/weathergen/datasets/tokenizer_masking.py b/src/weathergen/datasets/tokenizer_masking.py
index 7c033e398f..f59d15916a 100644
--- a/src/weathergen/datasets/tokenizer_masking.py
+++ b/src/weathergen/datasets/tokenizer_masking.py
@@ -99,6 +99,9 @@ def cell_to_token_mask(self, idxs_cells, idxs_cells_lens, mask):
         if num_tokens == 0:
             return (mask_tokens, mask_channels)
 
+        if mask is None:
+            mask = [True] * len(idxs_cells_lens)
+
         # TODO, TODO, TODO: use np.repeat
         # https://stackoverflow.com/questions/26038778/repeat-each-values-of-an-array-different-times
         # build token level mask: for each cell replicate the keep flag across its tokens
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 0923f6fb59..276da0bd67 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -11,7 +11,6 @@
 import contextlib
 import copy
 import logging
-import os
 import time
 from math import sqrt
 
@@ -343,11 +342,6 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
             self.grad_scaler = torch.amp.GradScaler("cuda")
         assert len(self.dataset) > 0, f"No data found in {self.dataset}"
 
-        if os.environ.get("DEBUG_DATASET"):
-            self.dataset.reset()
-            breakpoint()
-            batch = self.dataset._get_batch(0, 2)
-
         # lr is updated after each batch so account for this
         # TODO: conf should be read-only, do not modify the conf in flight
         len_ds = len(self.dataset)

From 5080d9c53d5c06c7e229f26088bca5aee5f9a6ac Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Fri, 31 Jul 2026 12:09:59 +0200
Subject: [PATCH 03/15] add configs

---
 .../day_of_year_conditioning.yml              |  22 ++++
 .../era5.yml                                  | 110 ++++++++++++++++++
 .../era5_sst_conditioning.yml                 |  49 ++++++++
 .../time_of_day_conditioning.yml              |  22 ++++
 4 files changed, 203 insertions(+)
 create mode 100644 config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml
 create mode 100644 config/streams/era5_1deg_forecasting_conditioning/era5.yml
 create mode 100644 config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml
 create mode 100644 config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml

diff --git a/config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/day_of_year_conditioning.yml
new file mode 100644
index 0000000000..452e63dae1
--- /dev/null
+++ b/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
diff --git a/config/streams/era5_1deg_forecasting_conditioning/era5.yml b/config/streams/era5_1deg_forecasting_conditioning/era5.yml
new file mode 100644
index 0000000000..6d6078cae4
--- /dev/null
+++ b/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
+   
\ No newline at end of file
diff --git a/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml
new file mode 100644
index 0000000000..bcd1d29540
--- /dev/null
+++ b/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: ['aifs-ea-an-oper-0001-mars-o96-1979-2024-1h-v3-with-era51.zarr']
+  timestep_conditioning: field
+  stream_id: 99
+  source: ['sst']
+  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
diff --git a/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml
new file mode 100644
index 0000000000..84f6e1c0e2
--- /dev/null
+++ b/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: day
+  embed:
+    net: linear
+    dim_embed: 256
+  conditioning_FE:
+    type: add
+    dim_embed: 256

From 80651e83d6b7341d3b51bca4a0a5f1088f86a115 Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Fri, 31 Jul 2026 12:14:24 +0200
Subject: [PATCH 04/15] add data reader for time conditioning

---
 .../datasets/data_reader_time_conditioning.py | 108 ++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 src/weathergen/datasets/data_reader_time_conditioning.py

diff --git a/src/weathergen/datasets/data_reader_time_conditioning.py b/src/weathergen/datasets/data_reader_time_conditioning.py
new file mode 100644
index 0000000000..562ecb88b5
--- /dev/null
+++ b/src/weathergen/datasets/data_reader_time_conditioning.py
@@ -0,0 +1,108 @@
+# (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.
+
+"""
+Synthetic data reader for scalar conditioning.
+
+Generates synthetic scalar values like time_of_day, day_of_year, and noise_level
+that can be used for model conditioning.
+"""
+
+import logging
+
+import numpy as np
+
+from weathergen.datasets.data_reader_base import (
+    DataReaderBase,
+    ReaderData,
+    TimeWindowHandler,
+    check_reader_data,
+)
+from weathergen.train.utils import Stage
+
+_logger = logging.getLogger(__name__)
+
+
+class DataReaderTimeConditioning(DataReaderBase):
+    """
+    DataReader that generates synthetic scalar conditioning values.
+
+    Supports three types:
+    - time_based: extracts hour/day from timestamp
+    - constant: fixed value
+    - random: random uniform value (per timestep)
+    """
+
+    def __init__(
+        self,
+        tw_handler: TimeWindowHandler,
+        filename,
+        stream_info: dict,
+        stage: Stage,
+    ) -> None:
+        super().__init__(tw_handler, stream_info)
+
+        self.source_channels = []
+        self.target_channels = []
+        self.source_idx = []
+        self.target_idx = []
+        self.geoinfo_channels = []
+        self.geoinfo_idx = []
+
+        self.conditioning_type = stream_info.get("conditioning_type", "time_based")
+        self.value_type = stream_info.get("value_type", "hour")
+
+        self.constant_value = stream_info.get("value", 0.0)
+        self.min_val = stream_info.get("min", 0.0)
+        self.max_val = stream_info.get("max", 1.0)
+
+        self.target_channel_weights = []
+        self.mean = np.array([0.0], dtype=np.float32)
+        self.stdev = np.array([1.0], dtype=np.float32)
+        self.mean_geoinfo = np.zeros(0, dtype=np.float32)
+        self.stdev_geoinfo = np.ones(0, dtype=np.float32)
+
+    def length(self) -> int:
+        return self.time_window_handler.get_index_range().end
+
+    def _get(self, idx, channels_idx) -> ReaderData:
+        dt_range = self.time_window_handler.window(idx)
+        dt = dt_range.start
+
+        if self.conditioning_type == "time_based":
+            if self.value_type == "hour":
+                hours = dt.astype("datetime64[h]").astype(int) % 24
+                minutes = dt.astype("datetime64[m]").astype(int) % 60
+                total_hours = hours + minutes / 60.0
+                value = np.array([total_hours / 24.0], dtype=np.float32)
+            elif self.value_type == "day":
+                days = dt.astype("datetime64[D]").astype(int) % 365
+                value = np.array([days / 365.0], dtype=np.float32)
+            else:
+                raise ValueError(f"Unknown value_type: {self.value_type}")
+
+        else:
+            raise ValueError(f"Unknown conditioning_type: {self.conditioning_type}")
+
+        coords = np.zeros((1, 2), dtype=np.float32)
+        geoinfos = np.zeros((1, 0), dtype=np.float32)
+        datetimes = np.array([dt], dtype=np.datetime64)
+
+        rdata = ReaderData(
+            coords=coords,
+            geoinfos=geoinfos,
+            data=value.reshape(1, 1),
+            datetimes=datetimes,
+        )
+        check_reader_data(rdata, dt_range)
+
+        return rdata
+
+    def get_geoinfo_size(self) -> int:
+        return 0

From c0ef63307be4dc0de2de7e64b90a9ede4f387ffa Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Fri, 31 Jul 2026 12:59:05 +0200
Subject: [PATCH 05/15] allow for data reader to accept one time conditioning
 config (value_type as list)

---
 .../datasets/data_reader_time_conditioning.py | 37 ++++++++++++-------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/weathergen/datasets/data_reader_time_conditioning.py b/src/weathergen/datasets/data_reader_time_conditioning.py
index 562ecb88b5..b418001ff3 100644
--- a/src/weathergen/datasets/data_reader_time_conditioning.py
+++ b/src/weathergen/datasets/data_reader_time_conditioning.py
@@ -56,7 +56,14 @@ def __init__(
         self.geoinfo_idx = []
 
         self.conditioning_type = stream_info.get("conditioning_type", "time_based")
-        self.value_type = stream_info.get("value_type", "hour")
+        value_type = stream_info.get("value_type", "hour")
+        if isinstance(value_type, list):
+            self.value_types = value_type
+        else:
+            self.value_types = [value_type]
+
+        self.source_channels = [f"time_conditioning_{vt}" for vt in self.value_types]
+        self.target_channels = []
 
         self.constant_value = stream_info.get("value", 0.0)
         self.min_val = stream_info.get("min", 0.0)
@@ -75,20 +82,22 @@ def _get(self, idx, channels_idx) -> ReaderData:
         dt_range = self.time_window_handler.window(idx)
         dt = dt_range.start
 
-        if self.conditioning_type == "time_based":
-            if self.value_type == "hour":
-                hours = dt.astype("datetime64[h]").astype(int) % 24
-                minutes = dt.astype("datetime64[m]").astype(int) % 60
-                total_hours = hours + minutes / 60.0
-                value = np.array([total_hours / 24.0], dtype=np.float32)
-            elif self.value_type == "day":
-                days = dt.astype("datetime64[D]").astype(int) % 365
-                value = np.array([days / 365.0], dtype=np.float32)
+        values = []
+        for vt in self.value_types:
+            if self.conditioning_type == "time_based":
+                if vt == "hour":
+                    hours = dt.astype("datetime64[h]").astype(int) % 24
+                    minutes = dt.astype("datetime64[m]").astype(int) % 60
+                    total_hours = hours + minutes / 60.0
+                    value = total_hours / 24.0
+                elif vt == "day":
+                    days = dt.astype("datetime64[D]").astype(int) % 365
+                    value = days / 365.0
+                else:
+                    raise ValueError(f"Unknown value_type: {vt}")
+                values.append(value)
             else:
-                raise ValueError(f"Unknown value_type: {self.value_type}")
-
-        else:
-            raise ValueError(f"Unknown conditioning_type: {self.conditioning_type}")
+                raise ValueError(f"Unknown conditioning_type: {self.conditioning_type}")
 
         coords = np.zeros((1, 2), dtype=np.float32)
         geoinfos = np.zeros((1, 0), dtype=np.float32)

From eb01a56401b182296e3f639584ffff00709ce5ae Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Fri, 31 Jul 2026 12:59:22 +0200
Subject: [PATCH 06/15] typo error fix

---
 .../time_of_day_conditioning.yml                                | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml
index 84f6e1c0e2..8b355c67ba 100644
--- a/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml
+++ b/config/streams/era5_1deg_forecasting_conditioning/time_of_day_conditioning.yml
@@ -13,7 +13,7 @@ time_of_day_conditioning:
   filenames: null
   timestep_conditioning: scalar
   conditioning_type: time_based
-  value_type: day
+  value_type: hour
   embed:
     net: linear
     dim_embed: 256

From 613be637ac2342be088db99c4c9feb34a5daf5ea Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Tue, 4 Aug 2026 13:03:54 +0200
Subject: [PATCH 07/15] correct bug, tokens needs to be padded as in source
 tokens

---
 src/weathergen/datasets/multi_stream_data_sampler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index 03b5c36d84..cd025ddd51 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -832,7 +832,7 @@ def _build_field_conditioning_data(
                 if rdata.is_empty():
                     stream_data = None
                 else:
-                    token_data_list = self.tokenizer.get_tokens_windows(stream_info, [rdata], False)
+                    token_data_list = self.tokenizer.get_tokens_windows(stream_info, [rdata], True)
                     token_data = token_data_list[0]
                     if token_data[0] is None:
                         stream_data = None

From 581cb270b58d3c19ece8260745f4c5188fadb3df Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Tue, 4 Aug 2026 16:33:59 +0200
Subject: [PATCH 08/15] clean multidatasampler for 1) reading synthetic data 2)
 handle empty paths

---
 .../datasets/data_reader_time_conditioning.py | 23 ++++++++-----------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/src/weathergen/datasets/data_reader_time_conditioning.py b/src/weathergen/datasets/data_reader_time_conditioning.py
index b418001ff3..53851587e8 100644
--- a/src/weathergen/datasets/data_reader_time_conditioning.py
+++ b/src/weathergen/datasets/data_reader_time_conditioning.py
@@ -48,30 +48,27 @@ def __init__(
     ) -> None:
         super().__init__(tw_handler, stream_info)
 
-        self.source_channels = []
-        self.target_channels = []
         self.source_idx = []
         self.target_idx = []
-        self.geoinfo_channels = []
         self.geoinfo_idx = []
 
+        self.conditioning = stream_info.get("conditioning", True)
         self.conditioning_type = stream_info.get("conditioning_type", "time_based")
-        value_type = stream_info.get("value_type", "hour")
+        value_type = stream_info.get("value_type")
         if isinstance(value_type, list):
             self.value_types = value_type
-        else:
+        elif value_type is not None:
             self.value_types = [value_type]
+        else:
+            raise ValueError("value_type in time_conditioning must be specified in stream_info")
 
         self.source_channels = [f"time_conditioning_{vt}" for vt in self.value_types]
         self.target_channels = []
-
-        self.constant_value = stream_info.get("value", 0.0)
-        self.min_val = stream_info.get("min", 0.0)
-        self.max_val = stream_info.get("max", 1.0)
-
+        self.geoinfo_channels = []
         self.target_channel_weights = []
-        self.mean = np.array([0.0], dtype=np.float32)
-        self.stdev = np.array([1.0], dtype=np.float32)
+
+        self.mean = np.zeros(len(self.source_channels), dtype=np.float32)
+        self.stdev = np.ones(len(self.source_channels), dtype=np.float32)
         self.mean_geoinfo = np.zeros(0, dtype=np.float32)
         self.stdev_geoinfo = np.ones(0, dtype=np.float32)
 
@@ -106,7 +103,7 @@ def _get(self, idx, channels_idx) -> ReaderData:
         rdata = ReaderData(
             coords=coords,
             geoinfos=geoinfos,
-            data=value.reshape(1, 1),
+            data=np.array(values, dtype=np.float32).reshape(1, -1),
             datetimes=datetimes,
         )
         check_reader_data(rdata, dt_range)

From 42781493bfee51b4f6f4c9784b1de2f5a5b51dba Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Tue, 4 Aug 2026 16:40:46 +0200
Subject: [PATCH 09/15] commit changes to multidatasampler

---
 .../datasets/multi_stream_data_sampler.py     | 39 +++++++++----------
 1 file changed, 18 insertions(+), 21 deletions(-)

diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index cd025ddd51..93958a06ea 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -230,6 +230,7 @@ def _init_stream_datasets(
         streams_datasets: dict[StreamName, _Stream] = {}
         for stream_name, stream_info in cf.streams.items():
             stream_info["data_paths"] = cf.get("data_paths", [])
+            ds_type = stream_info["type"]
             # list of sources for current stream
             streams_datasets[stream_name] = _Stream(stream_info, [])
             kwargs = {
@@ -238,7 +239,7 @@ def _init_stream_datasets(
                 "stage": self._stage,
             }
             dataset: type[AnyDataReader] | None = None
-            match stream_info["type"]:
+            match ds_type:
                 case "obs":
                     dataset = DataReaderObs
                 case "anemoi":
@@ -250,25 +251,22 @@ def _init_stream_datasets(
                         f"for stream name '{stream_name}'."
                         raise ValueError(msg)
 
-            filenames_cfg = stream_info.get("filenames", [pathlib.Path()])
-            conditioning_cfg = stream_info.get("conditioning", False)
-
-            if conditioning_cfg:
-                if filenames_cfg is None or len(filenames_cfg) == 0:
-                    ds_type = stream_info["type"]
-                    if is_root():
-                        logger.info(
-                            f"Opening conditioning dataset with type: {ds_type}"
-                            + f" from stream config {stream_name}.",
-                        )
-                    ds = dataset(filename=None, **kwargs)
-                    streams_datasets[stream_name].readers += [ds]
-                    continue
-                else:
-                    pass
-
-            for fname in filenames_cfg:
+            for fname in stream_info.get("filenames", pathlib.Path()):
                 fname = pathlib.Path(fname)
+                # skip if explicitly pointing to current directory
+                if fname is None or fname == pathlib.Path():
+                    if dataset.conditioning:
+                        if is_root():
+                            logger.info(
+                                f"Opening conditioning dataset with type: {ds_type}"
+                                + f" from stream config {stream_name}.",
+                            )
+                        ds = dataset(filename=None, **kwargs)
+                        streams_datasets[stream_name].readers += [ds]
+                        continue
+                    else:
+                        msg = f"Did not find input data for {ds_type} stream '{stream_name}'."
+                        raise FileNotFoundError(msg)
                 # dont check if file exists since zarr stores might be directories
                 if fname.exists():
                     # check if fname is a valid path to allow for simple overwriting
@@ -279,12 +277,11 @@ def _init_stream_datasets(
                     filename = next((f for f in filenames if f.exists()), None)
                     if filename is None:
                         msg = (
-                            f"Did not find input data for {stream_info['type']} "
+                            f"Did not find input data for {ds_type} "
                             f"stream '{stream_name}': {filenames}."
                         )
                         raise FileNotFoundError(msg)
 
-                ds_type = stream_info["type"]
                 if is_root():
                     logger.info(
                         f"Opening dataset with type: {ds_type}"

From 58fa0feb23d4972e6d303ba3a85287b092487515 Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Tue, 4 Aug 2026 16:41:58 +0200
Subject: [PATCH 10/15] correct syntax error

---
 src/weathergen/datasets/multi_stream_data_sampler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index 93958a06ea..b48df00db2 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -251,7 +251,7 @@ def _init_stream_datasets(
                         f"for stream name '{stream_name}'."
                         raise ValueError(msg)
 
-            for fname in stream_info.get("filenames", pathlib.Path()):
+            for fname in stream_info.get("filenames", [pathlib.Path()]):
                 fname = pathlib.Path(fname)
                 # skip if explicitly pointing to current directory
                 if fname is None or fname == pathlib.Path():

From 1679e77d07657dc56808a455e34fb59cf07fd6d2 Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Tue, 4 Aug 2026 16:45:06 +0200
Subject: [PATCH 11/15] add comment for time conditioning values

---
 src/weathergen/datasets/data_reader_time_conditioning.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/weathergen/datasets/data_reader_time_conditioning.py b/src/weathergen/datasets/data_reader_time_conditioning.py
index 53851587e8..fb0b26dd31 100644
--- a/src/weathergen/datasets/data_reader_time_conditioning.py
+++ b/src/weathergen/datasets/data_reader_time_conditioning.py
@@ -81,6 +81,8 @@ def _get(self, idx, channels_idx) -> ReaderData:
 
         values = []
         for vt in self.value_types:
+            # The values for time conditioning below denote fractional values of the
+            # day or year, normalized to [0, 1].
             if self.conditioning_type == "time_based":
                 if vt == "hour":
                     hours = dt.astype("datetime64[h]").astype(int) % 24

From 85ca9e1d033bfc8ca2cefbf35b1b36beac2e93f6 Mon Sep 17 00:00:00 2001
From: SavvasMel <s.melidonis@fz-juelich.de>
Date: Wed, 5 Aug 2026 15:56:52 +0200
Subject: [PATCH 12/15] correct bugs

---
 .../datasets/data_reader_time_conditioning.py    |  1 -
 .../datasets/multi_stream_data_sampler.py        | 16 ++++++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/weathergen/datasets/data_reader_time_conditioning.py b/src/weathergen/datasets/data_reader_time_conditioning.py
index fb0b26dd31..0e52b8421f 100644
--- a/src/weathergen/datasets/data_reader_time_conditioning.py
+++ b/src/weathergen/datasets/data_reader_time_conditioning.py
@@ -52,7 +52,6 @@ def __init__(
         self.target_idx = []
         self.geoinfo_idx = []
 
-        self.conditioning = stream_info.get("conditioning", True)
         self.conditioning_type = stream_info.get("conditioning_type", "time_based")
         value_type = stream_info.get("value_type")
         if isinstance(value_type, list):
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index b48df00db2..fa03cd1e25 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -230,7 +230,7 @@ def _init_stream_datasets(
         streams_datasets: dict[StreamName, _Stream] = {}
         for stream_name, stream_info in cf.streams.items():
             stream_info["data_paths"] = cf.get("data_paths", [])
-            ds_type = stream_info["type"]
+            ds_type = stream_info["type"]    
             # list of sources for current stream
             streams_datasets[stream_name] = _Stream(stream_info, [])
             kwargs = {
@@ -251,11 +251,19 @@ def _init_stream_datasets(
                         f"for stream name '{stream_name}'."
                         raise ValueError(msg)
 
-            for fname in stream_info.get("filenames", [pathlib.Path()]):
+            filenames_cfg = stream_info.get("filenames", [pathlib.Path()])
+
+            if filenames_cfg is None:
+                filenames_cfg = [pathlib.Path()]
+            else:
+                pass
+
+        
+            for fname in filenames_cfg:
                 fname = pathlib.Path(fname)
                 # skip if explicitly pointing to current directory
-                if fname is None or fname == pathlib.Path():
-                    if dataset.conditioning:
+                if fname == pathlib.Path():
+                    if stream_info.get("conditioning", False):
                         if is_root():
                             logger.info(
                                 f"Opening conditioning dataset with type: {ds_type}"

From 0a8c8d88cacf2ab37e7ad1dd65e8dbf5512f6d23 Mon Sep 17 00:00:00 2001
From: melidonis1 <s.melidonis@fz-juelich.de>
Date: Fri, 28 Aug 2026 13:12:35 +0200
Subject: [PATCH 13/15] introduce conditioning batch samples

---
 .../era5_sst_conditioning.yml                 |   4 +-
 src/weathergen/datasets/batch.py              | 184 ++++++++++--------
 .../datasets/multi_stream_data_sampler.py     |  22 ++-
 3 files changed, 115 insertions(+), 95 deletions(-)

diff --git a/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml b/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml
index bcd1d29540..31ba987128 100644
--- a/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml
+++ b/config/streams/era5_1deg_forecasting_conditioning/era5_sst_conditioning.yml
@@ -13,10 +13,10 @@
 ERA5_SST:
   type: anemoi
   conditioning: True
-  filenames: ['aifs-ea-an-oper-0001-mars-o96-1979-2024-1h-v3-with-era51.zarr']
+  filenames: ['ERA5-1deg-6h-mean-forcing-1978-2024.zarr']
   timestep_conditioning: field
   stream_id: 99
-  source: ['sst']
+  source: ['sea_surface_temperature']
   token_size: 8
   tokenize_spacetime: False
   geoinfo_channels: []
diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py
index 85ab148db4..ec30c42935 100644
--- a/src/weathergen/datasets/batch.py
+++ b/src/weathergen/datasets/batch.py
@@ -11,7 +11,6 @@
 
 import numpy as np
 import torch
-
 from weathergen.common.config import Config
 from weathergen.datasets.stream_data import StreamData
 
@@ -47,21 +46,14 @@ def pin_memory(self):
                 if stream_data is not None and hasattr(stream_data, "pin_memory"):
                     stream_data.pin_memory()
 
-        # Pin StreamData objects in conditioning_streams_data
-        if hasattr(self, "conditioning_streams_data") and isinstance(
-            self.conditioning_streams_data, dict
-        ):
-            for _name, steps in self.conditioning_streams_data.items():
-                for sd in steps:
-                    if sd is not None and hasattr(sd, "pin_memory"):
-                        sd.pin_memory()
-
         # Pin tensors in meta_info
         if hasattr(self, "meta_info") and isinstance(self.meta_info, dict):
             for _key, meta_data in self.meta_info.items():
                 if isinstance(meta_data, SampleMetaData):
                     # Pin mask tensor
-                    if meta_data.mask is not None and isinstance(meta_data.mask, torch.Tensor):
+                    if meta_data.mask is not None and isinstance(
+                        meta_data.mask, torch.Tensor
+                    ):
                         meta_data.mask = meta_data.mask.pin_memory()
 
         return self
@@ -73,10 +65,6 @@ def __init__(self, stream_names: list[str]) -> None:
         for stream_name in stream_names:
             self.streams_data[stream_name] = None
 
-        # Field conditioning stream data, keyed by stream name then indexed by forecast step.
-        # Populated by MultiStreamDataSampler for streams with timestep_conditioning: field.
-        self.conditioning_streams_data: dict[str, list[StreamData | None]] = {}
-
     def to_device(self, device) -> None:
         for key in self.meta_info.keys():
             self.meta_info[key].mask = (
@@ -84,63 +72,80 @@ 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:
                 self.streams_data[key] = val.to_device(device)
 
-        for name, steps in self.conditioning_streams_data.items():
-            self.conditioning_streams_data[name] = [
-                sd.to_device(device) if sd is not None else None for sd in steps
-            ]
-
     def is_empty(self) -> bool:
         """
         Check if sample is empty
         """
-        empty = [s.empty() if s is not None else True for _, s in self.streams_data.items()]
+        empty = [
+            s.empty() if s is not None else True for _, s in self.streams_data.items()
+        ]
         return np.array(empty).all()
 
     def is_nan(self) -> bool:
         """
         Check if sample is all NaN
         """
-        is_nan = [s.nan() if s is not None else False for _, s in self.streams_data.items()]
+        is_nan = [
+            s.nan() if s is not None else False for _, s in self.streams_data.items()
+        ]
         return np.array(is_nan).all()
 
     def sources_empty(self) -> bool:
         """
         Check if sources for sample are empty
         """
-        empty = [s.source_empty() if s is not None else True for _, s in self.streams_data.items()]
+        empty = [
+            s.source_empty() if s is not None else True
+            for _, s in self.streams_data.items()
+        ]
         return np.array(empty).all()
 
     def sources_nan(self) -> bool:
         """
         Check if sources for sample are all NaN
         """
-        is_nan = [s.source_nan() if s is not None else False for _, s in self.streams_data.items()]
+        is_nan = [
+            s.source_nan() if s is not None else False
+            for _, s in self.streams_data.items()
+        ]
         return np.array(is_nan).all()
 
     def targets_empty(self) -> bool:
         """
         Check if targets for sample are empty
         """
-        empty = [s.target_empty() if s is not None else True for _, s in self.streams_data.items()]
+        empty = [
+            s.target_empty() if s is not None else True
+            for _, s in self.streams_data.items()
+        ]
         return np.array(empty).all()
 
     def targets_nan(self) -> bool:
         """
         Check if targets for sample are all NaN
         """
-        is_nan = [s.target_nan() if s is not None else False for _, s in self.streams_data.items()]
+        is_nan = [
+            s.target_nan() if s is not None else False
+            for _, s in self.streams_data.items()
+        ]
         return np.array(is_nan).all()
 
     def add_stream_data(self, stream_name: str, stream_data: StreamData) -> None:
         """
         Add data for stream @stream_name to sample
         """
-        assert self.streams_data.get(stream_name, -1) != -1, "stream name does not exist"
+        assert self.streams_data.get(stream_name, -1) != -1, (
+            "stream name does not exist"
+        )
         self.streams_data[stream_name] = stream_data
 
     def add_meta_info(self, stream_name: str, meta_info: SampleMetaData) -> None:
@@ -153,7 +158,9 @@ def get_stream_data(self, stream_name: str) -> StreamData:
         """
         Get data for stream @stream_name from sample
         """
-        assert self.streams_data.get(stream_name, -1) != -1, "stream name does not exist"
+        assert self.streams_data.get(stream_name, -1) != -1, (
+            "stream name does not exist"
+        )
         return self.streams_data[stream_name]
 
     def get_num_source_steps(self) -> int:
@@ -178,28 +185,6 @@ def get_num_target_steps(self) -> int:
         ]
         return min(lens) if len(lens) > 0 else 0
 
-    def add_conditioning_stream_data(
-        self, stream_name: str, step: int, stream_data: StreamData | None
-    ) -> None:
-        """
-        Add StreamData for field conditioning stream @stream_name at forecast step @step to sample
-        """
-        if stream_name not in self.conditioning_streams_data:
-            self.conditioning_streams_data[stream_name] = []
-        steps = self.conditioning_streams_data[stream_name]
-        while len(steps) <= step:
-            steps.append(None)
-        steps[step] = stream_data
-
-    def get_conditioning_stream_data(self, stream_name: str, step: int) -> StreamData | None:
-        """
-        Get StreamData for field conditioning stream @stream_name at forecast step @step from sample
-        """
-        steps = self.conditioning_streams_data.get(stream_name)
-        if steps is None or step >= len(steps):
-            return None
-        return steps[step]
-
 
 class BatchSamples:
     """
@@ -229,7 +214,9 @@ def to_device(self, device):
             sample.to_device(device)
 
         self.tokens_lens = (
-            self.tokens_lens.to(device, non_blocking=True) if self.tokens_lens is not None else None
+            self.tokens_lens.to(device, non_blocking=True)
+            if self.tokens_lens is not None
+            else None
         )
 
         self.device = device
@@ -247,7 +234,9 @@ def get_subset(self, subset: list | None = None):
             # create copy and then select subset for samples and tokens_lens
             bs = copy.deepcopy(self)
             bs.samples = [bs.samples[i] for i in subset]
-            torch_idxs = torch.tensor(subset, dtype=torch.long, device=bs.tokens_lens.device)
+            torch_idxs = torch.tensor(
+                subset, dtype=torch.long, device=bs.tokens_lens.device
+            )
             bs.tokens_lens = torch.index_select(bs.tokens_lens, 1, torch_idxs)
             return bs
 
@@ -285,25 +274,33 @@ def sources_empty(self) -> bool:
         """
         Check if sources for all samples are empty
         """
-        return np.array([s.sources_empty() if s is not None else True for s in self.samples]).all()
+        return np.array(
+            [s.sources_empty() if s is not None else True for s in self.samples]
+        ).all()
 
     def targets_empty(self) -> bool:
         """
         Check if targets for all samples are empty
         """
-        return np.array([s.targets_empty() if s is not None else True for s in self.samples]).all()
+        return np.array(
+            [s.targets_empty() if s is not None else True for s in self.samples]
+        ).all()
 
     def sources_nan(self) -> bool:
         """
         Check if sources for all samples are all NaN
         """
-        return np.array([s.sources_nan() if s is not None else False for s in self.samples]).all()
+        return np.array(
+            [s.sources_nan() if s is not None else False for s in self.samples]
+        ).all()
 
     def targets_nan(self) -> bool:
         """
         Check if targets for all samples are all NaN
         """
-        return np.array([s.targets_nan() if s is not None else False for s in self.samples]).all()
+        return np.array(
+            [s.targets_nan() if s is not None else False for s in self.samples]
+        ).all()
 
     def pin_memory(self):
         """Pin all tensors in this batch to CPU pinned memory"""
@@ -345,6 +342,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,
@@ -364,7 +362,16 @@ def __init__(
             stream_names, num_target_samples, output_steps, self.output_idxs
         )
 
-        self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32)
+        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)]
 
     def pin_memory(self):
@@ -376,6 +383,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
@@ -385,6 +395,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
 
@@ -401,12 +412,18 @@ def add_source_stream(
         """
         Add data for one stream to sample @source_sample_idx
         """
-        self.source_samples.samples[source_sample_idx].add_stream_data(stream_name, stream_data)
+        self.source_samples.samples[source_sample_idx].add_stream_data(
+            stream_name, stream_data
+        )
 
         # add the meta_info
-        self.source_samples.samples[source_sample_idx].add_meta_info(stream_name, source_meta_info)
+        self.source_samples.samples[source_sample_idx].add_meta_info(
+            stream_name, source_meta_info
+        )
 
-        assert target_sample_idx < len(self.target_samples), "invalid value for target_sample_idx"
+        assert target_sample_idx < len(self.target_samples), (
+            "invalid value for target_sample_idx"
+        )
         self.source2target_matching_idxs[source_sample_idx] = target_sample_idx
 
     def add_target_stream(
@@ -420,10 +437,14 @@ def add_target_stream(
         """
         Add data for one stream to sample @target_sample_idx
         """
-        self.target_samples.samples[target_sample_idx].add_stream_data(stream_name, stream_data)
+        self.target_samples.samples[target_sample_idx].add_stream_data(
+            stream_name, stream_data
+        )
 
         # add the meta_info -- for target we have different
-        self.target_samples.samples[target_sample_idx].add_meta_info(stream_name, target_meta_info)
+        self.target_samples.samples[target_sample_idx].add_meta_info(
+            stream_name, target_meta_info
+        )
 
         if isinstance(source_sample_idx, int):
             assert source_sample_idx < len(self.source_samples), (
@@ -444,12 +465,14 @@ def add_scalar_conditioning_stream(self, stream_name, conditioning_values):
                 sample.add_meta_info(stream_name, SampleMetaData(params={}))
             sample.meta_info[stream_name].conditioning = conditioning_values
 
-    def add_field_conditioning_stream(self, stream_name, step: int, stream_data: StreamData):
+    def add_field_conditioning_stream(
+        self, stream_name, step: int, stream_data: StreamData
+    ):
         """
         Add field conditioning values for all samples in the batch for a specific stream.
         """
-        for sample in self.source_samples.samples:
-            sample.add_conditioning_stream_data(stream_name, step, stream_data)
+        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
@@ -461,35 +484,24 @@ def get_scalar_conditioning_values(
         for any sample.
         """
         values = []
-        for sample in self.samples:
+        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):
+            if (
+                meta is None
+                or meta.conditioning is None
+                or step >= len(meta.conditioning)
+            ):
                 return None
             values.append(meta.conditioning[step])
-        return np.stack(values, axis=0) if values else None
-
-    def get_field_conditioning_values(
-        self, stream_name: str, step: int
-    ) -> np.typing.NDArray | None:
-        """
-        Get field conditioning values for all samples at a specific forecast step.
-
-        Returns np.ndarray of shape (num_samples, ...) or None if not available
-        for any sample.
-        """
-        values = []
-        for sample in self.samples:
-            stream_data = sample.get_conditioning_stream_data(stream_name, step)
-            if stream_data is None:
-                return None
-            values.append(stream_data.data)
-        return np.stack(values, axis=0) if values else None
+        return torch.stack(values, dim=0) if values else None
 
     def is_empty(self):
         """
         Check if batch is empty
         """
-        return self.source_samples.sources_empty() or self.target_samples.targets_empty()
+        return (
+            self.source_samples.sources_empty() or self.target_samples.targets_empty()
+        )
 
     def is_nan(self):
         """
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index fa03cd1e25..c57fcb5c46 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -28,7 +28,7 @@
 from weathergen.datasets.data_reader_obs import DataReaderObs
 from weathergen.datasets.masking import Masker
 from weathergen.datasets.stream_data import StreamData, spoof
-from weathergen.datasets.tokenizer_masking import TokenizerMasking
+from weathergen.datasets.tokenizer_masking import TokenizerMasking, readerdata_to_torch
 from weathergen.datasets.utils import (
     get_tokens_lens,
 )
@@ -230,7 +230,7 @@ def _init_stream_datasets(
         streams_datasets: dict[StreamName, _Stream] = {}
         for stream_name, stream_info in cf.streams.items():
             stream_info["data_paths"] = cf.get("data_paths", [])
-            ds_type = stream_info["type"]    
+            ds_type = stream_info["type"]
             # list of sources for current stream
             streams_datasets[stream_name] = _Stream(stream_info, [])
             kwargs = {
@@ -258,7 +258,6 @@ def _init_stream_datasets(
             else:
                 pass
 
-        
             for fname in filenames_cfg:
                 fname = pathlib.Path(fname)
                 # skip if explicitly pointing to current directory
@@ -683,12 +682,17 @@ def _preprocess_model_batch(
         Perform necessary pre-processing of model batch
         """
         stream_names = list(self.streams_datasets.keys())
+        conditioning_stream_names = list(self.field_conditioning_datasets.keys())
         batch.source_samples.tokens_lens = get_tokens_lens(
             stream_names, batch.source_samples, source_input_steps
         )
         batch.target_samples.tokens_lens = get_tokens_lens(
             stream_names, batch.target_samples, target_input_steps
         )
+        if conditioning_stream_names:
+            batch.conditioning_samples.tokens_lens = get_tokens_lens(
+                conditioning_stream_names, batch.conditioning_samples, input_steps=1
+            )
 
         return batch
 
@@ -719,6 +723,7 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
         num_output_steps = self._get_output_length(num_forecast_steps)
         batch = ModelBatch(
             list(self.streams_datasets.keys()),
+            list(self.field_conditioning_datasets.keys()),
             num_source_samples,
             num_target_samples,
             self.output_offset,
@@ -797,7 +802,6 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
         source_in_steps = input_steps.max().item()
         target_in_steps = np.array([tc.get("num_steps_input", 1) for _, tc in target_cfgs.items()])
         target_in_steps = 1 if len(target_in_steps) == 0 else target_in_steps.max().item()
-        batch = self._preprocess_model_batch(batch, source_in_steps, target_in_steps)
 
         if self.scalar_conditioning_stream_names:
             batch = self._build_scalar_conditioning_data(batch, idx, num_forecast_steps)
@@ -805,6 +809,8 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
         if self.field_conditioning_stream_names:
             batch = self._build_field_conditioning_data(batch, idx, num_forecast_steps)
 
+        batch = self._preprocess_model_batch(batch, source_in_steps, target_in_steps)
+
         return batch
 
     def _build_scalar_conditioning_data(
@@ -817,9 +823,11 @@ def _build_scalar_conditioning_data(
             for timestep_idx in range(self.output_offset, num_output_steps):
                 step_dt = idx + (self.time_step * timestep_idx) // self.step_timedelta
                 rdata = stream_ds.readers[0].get_source(step_dt)
-                step_values.append(rdata.data.flatten().copy())
+                step_values.append(readerdata_to_torch(rdata).data.flatten())
             conditioning_values = (
-                np.stack(step_values, axis=0) if step_values else np.zeros((0, 1), dtype=np.float32)
+                torch.stack(step_values, dim=0)
+                if step_values
+                else torch.zeros((0, 1), dtype=torch.float32)
             )
             batch.add_scalar_conditioning_stream(stream_name, conditioning_values)
         return batch
@@ -827,7 +835,7 @@ def _build_scalar_conditioning_data(
     def _build_field_conditioning_data(
         self, batch: ModelBatch, idx: int, num_forecast_steps: int
     ) -> ModelBatch:
-        """Collect per-step field conditioning data and store in conditioning_streams_data."""
+        """Collect per-step field conditioning data and store in conditioning_samples.streams_data."""
         num_output_steps = self._get_output_length(num_forecast_steps)
         for stream_name, stream_ds in self.field_conditioning_datasets.items():
             stream_info = stream_ds.info

From d74a43406e9552075ac6e8140cdfbbb6b8a53747 Mon Sep 17 00:00:00 2001
From: melidonis1 <s.melidonis@fz-juelich.de>
Date: Fri, 28 Aug 2026 13:17:23 +0200
Subject: [PATCH 14/15] linting-multisampler-cond

---
 src/weathergen/datasets/batch.py              | 103 +++++-------------
 .../datasets/multi_stream_data_sampler.py     |   3 +-
 2 files changed, 27 insertions(+), 79 deletions(-)

diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py
index ec30c42935..7eb459ef6d 100644
--- a/src/weathergen/datasets/batch.py
+++ b/src/weathergen/datasets/batch.py
@@ -11,6 +11,7 @@
 
 import numpy as np
 import torch
+
 from weathergen.common.config import Config
 from weathergen.datasets.stream_data import StreamData
 
@@ -51,9 +52,7 @@ def pin_memory(self):
             for _key, meta_data in self.meta_info.items():
                 if isinstance(meta_data, SampleMetaData):
                     # Pin mask tensor
-                    if meta_data.mask is not None and isinstance(
-                        meta_data.mask, torch.Tensor
-                    ):
+                    if meta_data.mask is not None and isinstance(meta_data.mask, torch.Tensor):
                         meta_data.mask = meta_data.mask.pin_memory()
 
         return self
@@ -85,67 +84,49 @@ def is_empty(self) -> bool:
         """
         Check if sample is empty
         """
-        empty = [
-            s.empty() if s is not None else True for _, s in self.streams_data.items()
-        ]
+        empty = [s.empty() if s is not None else True for _, s in self.streams_data.items()]
         return np.array(empty).all()
 
     def is_nan(self) -> bool:
         """
         Check if sample is all NaN
         """
-        is_nan = [
-            s.nan() if s is not None else False for _, s in self.streams_data.items()
-        ]
+        is_nan = [s.nan() if s is not None else False for _, s in self.streams_data.items()]
         return np.array(is_nan).all()
 
     def sources_empty(self) -> bool:
         """
         Check if sources for sample are empty
         """
-        empty = [
-            s.source_empty() if s is not None else True
-            for _, s in self.streams_data.items()
-        ]
+        empty = [s.source_empty() if s is not None else True for _, s in self.streams_data.items()]
         return np.array(empty).all()
 
     def sources_nan(self) -> bool:
         """
         Check if sources for sample are all NaN
         """
-        is_nan = [
-            s.source_nan() if s is not None else False
-            for _, s in self.streams_data.items()
-        ]
+        is_nan = [s.source_nan() if s is not None else False for _, s in self.streams_data.items()]
         return np.array(is_nan).all()
 
     def targets_empty(self) -> bool:
         """
         Check if targets for sample are empty
         """
-        empty = [
-            s.target_empty() if s is not None else True
-            for _, s in self.streams_data.items()
-        ]
+        empty = [s.target_empty() if s is not None else True for _, s in self.streams_data.items()]
         return np.array(empty).all()
 
     def targets_nan(self) -> bool:
         """
         Check if targets for sample are all NaN
         """
-        is_nan = [
-            s.target_nan() if s is not None else False
-            for _, s in self.streams_data.items()
-        ]
+        is_nan = [s.target_nan() if s is not None else False for _, s in self.streams_data.items()]
         return np.array(is_nan).all()
 
     def add_stream_data(self, stream_name: str, stream_data: StreamData) -> None:
         """
         Add data for stream @stream_name to sample
         """
-        assert self.streams_data.get(stream_name, -1) != -1, (
-            "stream name does not exist"
-        )
+        assert self.streams_data.get(stream_name, -1) != -1, "stream name does not exist"
         self.streams_data[stream_name] = stream_data
 
     def add_meta_info(self, stream_name: str, meta_info: SampleMetaData) -> None:
@@ -158,9 +139,7 @@ def get_stream_data(self, stream_name: str) -> StreamData:
         """
         Get data for stream @stream_name from sample
         """
-        assert self.streams_data.get(stream_name, -1) != -1, (
-            "stream name does not exist"
-        )
+        assert self.streams_data.get(stream_name, -1) != -1, "stream name does not exist"
         return self.streams_data[stream_name]
 
     def get_num_source_steps(self) -> int:
@@ -214,9 +193,7 @@ def to_device(self, device):
             sample.to_device(device)
 
         self.tokens_lens = (
-            self.tokens_lens.to(device, non_blocking=True)
-            if self.tokens_lens is not None
-            else None
+            self.tokens_lens.to(device, non_blocking=True) if self.tokens_lens is not None else None
         )
 
         self.device = device
@@ -234,9 +211,7 @@ def get_subset(self, subset: list | None = None):
             # create copy and then select subset for samples and tokens_lens
             bs = copy.deepcopy(self)
             bs.samples = [bs.samples[i] for i in subset]
-            torch_idxs = torch.tensor(
-                subset, dtype=torch.long, device=bs.tokens_lens.device
-            )
+            torch_idxs = torch.tensor(subset, dtype=torch.long, device=bs.tokens_lens.device)
             bs.tokens_lens = torch.index_select(bs.tokens_lens, 1, torch_idxs)
             return bs
 
@@ -274,33 +249,25 @@ def sources_empty(self) -> bool:
         """
         Check if sources for all samples are empty
         """
-        return np.array(
-            [s.sources_empty() if s is not None else True for s in self.samples]
-        ).all()
+        return np.array([s.sources_empty() if s is not None else True for s in self.samples]).all()
 
     def targets_empty(self) -> bool:
         """
         Check if targets for all samples are empty
         """
-        return np.array(
-            [s.targets_empty() if s is not None else True for s in self.samples]
-        ).all()
+        return np.array([s.targets_empty() if s is not None else True for s in self.samples]).all()
 
     def sources_nan(self) -> bool:
         """
         Check if sources for all samples are all NaN
         """
-        return np.array(
-            [s.sources_nan() if s is not None else False for s in self.samples]
-        ).all()
+        return np.array([s.sources_nan() if s is not None else False for s in self.samples]).all()
 
     def targets_nan(self) -> bool:
         """
         Check if targets for all samples are all NaN
         """
-        return np.array(
-            [s.targets_nan() if s is not None else False for s in self.samples]
-        ).all()
+        return np.array([s.targets_nan() if s is not None else False for s in self.samples]).all()
 
     def pin_memory(self):
         """Pin all tensors in this batch to CPU pinned memory"""
@@ -369,9 +336,7 @@ def __init__(
             output_idxs=self.output_idxs,
         )
 
-        self.source2target_matching_idxs = np.full(
-            num_source_samples, -1, dtype=np.int32
-        )
+        self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32)
         self.target2source_matching_idxs = [[] for _ in range(num_target_samples)]
 
     def pin_memory(self):
@@ -412,18 +377,12 @@ def add_source_stream(
         """
         Add data for one stream to sample @source_sample_idx
         """
-        self.source_samples.samples[source_sample_idx].add_stream_data(
-            stream_name, stream_data
-        )
+        self.source_samples.samples[source_sample_idx].add_stream_data(stream_name, stream_data)
 
         # add the meta_info
-        self.source_samples.samples[source_sample_idx].add_meta_info(
-            stream_name, source_meta_info
-        )
+        self.source_samples.samples[source_sample_idx].add_meta_info(stream_name, source_meta_info)
 
-        assert target_sample_idx < len(self.target_samples), (
-            "invalid value for target_sample_idx"
-        )
+        assert target_sample_idx < len(self.target_samples), "invalid value for target_sample_idx"
         self.source2target_matching_idxs[source_sample_idx] = target_sample_idx
 
     def add_target_stream(
@@ -437,14 +396,10 @@ def add_target_stream(
         """
         Add data for one stream to sample @target_sample_idx
         """
-        self.target_samples.samples[target_sample_idx].add_stream_data(
-            stream_name, stream_data
-        )
+        self.target_samples.samples[target_sample_idx].add_stream_data(stream_name, stream_data)
 
         # add the meta_info -- for target we have different
-        self.target_samples.samples[target_sample_idx].add_meta_info(
-            stream_name, target_meta_info
-        )
+        self.target_samples.samples[target_sample_idx].add_meta_info(stream_name, target_meta_info)
 
         if isinstance(source_sample_idx, int):
             assert source_sample_idx < len(self.source_samples), (
@@ -465,9 +420,7 @@ def add_scalar_conditioning_stream(self, stream_name, conditioning_values):
                 sample.add_meta_info(stream_name, SampleMetaData(params={}))
             sample.meta_info[stream_name].conditioning = conditioning_values
 
-    def add_field_conditioning_stream(
-        self, stream_name, step: int, stream_data: StreamData
-    ):
+    def add_field_conditioning_stream(self, stream_name, step: int, stream_data: StreamData):
         """
         Add field conditioning values for all samples in the batch for a specific stream.
         """
@@ -486,11 +439,7 @@ def get_scalar_conditioning_values(
         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)
-            ):
+            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
@@ -499,9 +448,7 @@ def is_empty(self):
         """
         Check if batch is empty
         """
-        return (
-            self.source_samples.sources_empty() or self.target_samples.targets_empty()
-        )
+        return self.source_samples.sources_empty() or self.target_samples.targets_empty()
 
     def is_nan(self):
         """
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index c57fcb5c46..ac35ede6d9 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -835,7 +835,8 @@ def _build_scalar_conditioning_data(
     def _build_field_conditioning_data(
         self, batch: ModelBatch, idx: int, num_forecast_steps: int
     ) -> ModelBatch:
-        """Collect per-step field conditioning data and store in conditioning_samples.streams_data."""
+        """Collect per-step field conditioning data and
+        store in conditioning_samples.streams_data."""
         num_output_steps = self._get_output_length(num_forecast_steps)
         for stream_name, stream_ds in self.field_conditioning_datasets.items():
             stream_info = stream_ds.info

From 1385214eb6bbed0819163b5406e38350c5379d15 Mon Sep 17 00:00:00 2001
From: melidonis1 <s.melidonis@fz-juelich.de>
Date: Thu, 3 Sep 2026 17:53:19 +0200
Subject: [PATCH 15/15] conditioning data will be parsed in source per forecast
 step

---
 src/weathergen/datasets/batch.py              |  3 +-
 .../datasets/multi_stream_data_sampler.py     | 35 ++++++++++++-------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py
index 7eb459ef6d..c28dc6485f 100644
--- a/src/weathergen/datasets/batch.py
+++ b/src/weathergen/datasets/batch.py
@@ -420,9 +420,10 @@ def add_scalar_conditioning_stream(self, stream_name, conditioning_values):
                 sample.add_meta_info(stream_name, SampleMetaData(params={}))
             sample.meta_info[stream_name].conditioning = conditioning_values
 
-    def add_field_conditioning_stream(self, stream_name, step: int, stream_data: StreamData):
+    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
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index ac35ede6d9..6a0c55cfe3 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -835,14 +835,23 @@ def _build_scalar_conditioning_data(
     def _build_field_conditioning_data(
         self, batch: ModelBatch, idx: int, num_forecast_steps: int
     ) -> ModelBatch:
-        """Collect per-step field conditioning data and
+        """Collect per-step field conditioning data with sliding window and
         store in conditioning_samples.streams_data."""
         num_output_steps = self._get_output_length(num_forecast_steps)
+
         for stream_name, stream_ds in self.field_conditioning_datasets.items():
             stream_info = stream_ds.info
-            for step, timestep_idx in enumerate(range(self.output_offset, num_output_steps)):
-                step_dt = idx + (self.time_step * timestep_idx) // self.step_timedelta
-                rdata = collect_datasources(stream_ds.readers, step_dt, "source", self.rng)
+
+            # Create ONE StreamData with input_steps=num_output_steps
+            # Each source step corresponds to one forecast step's conditioning window
+            stream_data = StreamData(idx, num_forecast_steps, 1, self.num_healpix_cells)
+
+            # Collect data for each forecast step
+            for step, _ in enumerate(range(self.output_offset, num_output_steps)):
+                # Conditioning window for forecast step N is at idx + (N - output_offset)
+                # For output_offset=1: forecast step 1 → idx+0, step 2 → idx+1, etc.
+                rdata = collect_datasources(stream_ds.readers, idx + step, "source", self.rng)
+
                 if rdata.is_empty():
                     stream_data = None
                 else:
@@ -851,17 +860,17 @@ def _build_field_conditioning_data(
                     if token_data[0] is None:
                         stream_data = None
                     else:
-                        time_win = self.time_window_handler.window(step_dt)
+                        time_win = self.time_window_handler.window(idx + step)
                         src_cells, src_lens = self.tokenizer.get_source(
-                            stream_info,
-                            rdata,
-                            token_data,
-                            (time_win.start, time_win.end),
-                            None,
+                            stream_info, rdata, token_data, (time_win.start, time_win.end), None
                         )
-                        stream_data = StreamData(step_dt, 1, 1, self.num_healpix_cells)
-                        stream_data.add_source(self._stage, 0, rdata, src_lens, src_cells, False)
-                batch.add_field_conditioning_stream(stream_name, step, stream_data)
+
+                        # Add to the correct step index (step = 0, 1, 2, ...)
+                        stream_data.add_source(self._stage, step, rdata, src_lens, src_cells, False)
+
+            # Add the complete StreamData once (not in the loop)
+            batch.add_field_conditioning_stream(stream_name, stream_data)
+
         return batch
 
     def __iter__(self) -> ModelBatch: