From 623b68c0df992c5260b7d098799747ad55f539f6 Mon Sep 17 00:00:00 2001
From: ankitpatnala <ankitpatnala@gmail.com>
Date: Wed, 9 Sep 2026 22:37:28 +0200
Subject: [PATCH 1/6] added conditioning configs

---
 .../era5_1deg_forecasting/forecast.yml        | 16 +++++++
 config/streams/era5_1deg_forecasting/sst.yml  | 46 +++++++++++++++++++
 2 files changed, 62 insertions(+)
 create mode 100644 config/streams/era5_1deg_forecasting/forecast.yml
 create mode 100644 config/streams/era5_1deg_forecasting/sst.yml

diff --git a/config/streams/era5_1deg_forecasting/forecast.yml b/config/streams/era5_1deg_forecasting/forecast.yml
new file mode 100644
index 0000000000..f3bee129eb
--- /dev/null
+++ b/config/streams/era5_1deg_forecasting/forecast.yml
@@ -0,0 +1,16 @@
+# (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.
+
+Forecast_conditions :
+  type : condition
+  filenames : []
+  variables : ['start_day', 'start_time', 'end_day', 'end_time']
+  transform : "fourier"  # absolute, absolute_normalized, cos_sin, fourier
+  entry_point : "forecast_engine"
+  emb_dimension : 16
diff --git a/config/streams/era5_1deg_forecasting/sst.yml b/config/streams/era5_1deg_forecasting/sst.yml
new file mode 100644
index 0000000000..0fd7dcb74e
--- /dev/null
+++ b/config/streams/era5_1deg_forecasting/sst.yml
@@ -0,0 +1,46 @@
+# (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.
+
+# SST forcing: prescribed lower boundary, not predicted. Re-read at every forecast
+# step's valid time and injected into the FE. SST is NaN over land.
+SST :
+  type : forcing
+  #filenames : ['/e/data1/slmet/ml_training/ERA5-1deg-6h-mean-forcing-1978-2024.zarr']
+  filenames : ['/e/data1/slmet/ml_training/aifs-ea-an-oper-0001-mars-o96-1979-2023-6h-v2-s2s-predictors.zarr']
+  stream_id : 10
+  entry_point : "forecast_engine"
+  frequency : 06:00:00
+  # empty target channels -> is_stream_forcing -> no decoder
+  #train_source_channels : ['sea_surface_temperature', 'sea_ice_cover', 'land_sea_mask']
+  #val_source_channels   : ['sea_surface_temperature', 'sea_ice_cover', 'land_sea_mask']
+  train_source_channels : ['sst']
+  val_source_channels   : ['sst']
+  train_target_channels : []
+  val_target_channels   : []
+  geoinfo_channels : []
+  loss_weight : 0.
+  location_weight : cosine_latitude
+  token_size : 8
+  tokenize_spacetime : True
+  injection :
+    #   none        : disabled
+    #   additive    : forcing embedding added to FE latent tokens
+    #   cross_attn  : latent tokens cross-attend to forcing tokens
+    #   global      : pooled to a global vector -> AdaLN condition path
+    #   adaln_local : per-cell AdaLN conditioning
+    mode : adaln_local
+    # True: raw grid pooled by a learned segment-softmax (LearnedForcingPool).
+    # False: fixed scatter-mean in the dataloader, no gradient.
+    learned_pool : True
+  embed :
+    net : transformer
+    num_tokens : 1
+    num_heads : 8
+    dim_embed : 512
+    num_blocks : 2

From fe8561a198ea3ec5439f1384778670b262fb602a Mon Sep 17 00:00:00 2001
From: ankitpatnala <ankitpatnala@gmail.com>
Date: Wed, 9 Sep 2026 22:38:01 +0200
Subject: [PATCH 2/6] added condition data readers

---
 .../readers_extra/data_reader_condition.py    | 189 ++++++++++++++++++
 .../readers_extra/data_reader_forcing.py      |  39 ++++
 .../src/weathergen/readers_extra/registry.py  |   4 +
 3 files changed, 232 insertions(+)
 create mode 100644 packages/readers_extra/src/weathergen/readers_extra/data_reader_condition.py
 create mode 100644 packages/readers_extra/src/weathergen/readers_extra/data_reader_forcing.py

diff --git a/packages/readers_extra/src/weathergen/readers_extra/data_reader_condition.py b/packages/readers_extra/src/weathergen/readers_extra/data_reader_condition.py
new file mode 100644
index 0000000000..4104196c9e
--- /dev/null
+++ b/packages/readers_extra/src/weathergen/readers_extra/data_reader_condition.py
@@ -0,0 +1,189 @@
+# (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.
+
+import logging
+import math
+from pathlib import Path
+from typing import override
+
+import numpy as np
+
+from weathergen.datasets.data_reader_base import (
+    DataReaderTimestep,
+    DTRange,
+    ReaderData,
+    TimeWindowHandler,
+    TIndex,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+class DataReaderCondition(DataReaderTimestep):
+    "Wrapper for forecast condition variables derived from time window metadata"
+
+    def __init__(
+        self,
+        tw_handler: TimeWindowHandler,
+        filename: Path,
+        stream_info: dict,
+    ) -> None:
+        """
+        Construct data reader for forecast condition variables.
+
+        Parameters
+        ----------
+        tw_handler :
+            time window handler
+        filename :
+            unused; kept for interface compatibility (filenames should be empty)
+        stream_info :
+            information about stream; must include 'transform' and 'variables'
+
+        Returns
+        -------
+        None
+        """
+        self.source_idx = []
+        self.source_channels = []
+        self.target_channels = []
+        self.geoinfo_channels = []
+        self.target_idx = []
+        self.geoinfo_idx = []
+        self.target_channel_weights = []
+        self.condition_idx = []
+
+        self.transform: str = stream_info.get("transform", "absolute")
+        self.variables: list[str] = list(
+            stream_info.get("variables", ["start_day", "start_time", "end_day", "end_time"])
+        )
+        self.emb_dimension: int = stream_info.get("emb_dimension", 4)
+        self.num_channels: int = self._compute_num_channels(stream_info)
+        self.source_idx = []
+
+        super().__init__(
+            tw_handler,
+            stream_info,
+            tw_handler.t_start,
+            tw_handler.t_end,
+            tw_handler.t_window_step,
+        )
+
+        self.len = int((tw_handler.t_end - tw_handler.t_start) // tw_handler.t_window_step)
+
+    def _compute_num_channels(self, stream_info: dict) -> int:
+        if self.transform in ("absolute", "absolute_normalized"):
+            return len(self.variables)
+        elif self.transform == "cos_sin":
+            return 2 * len(self.variables)
+        elif self.transform == "fourier":
+            assert "emb_dimension" in stream_info, (
+                "Fourier transform requires 'emb_dimension' in stream_info"
+            )
+            return stream_info.get("emb_dimension") * len(self.variables)
+        else:
+            raise ValueError(f"Unknown transform: {self.transform!r}")
+
+    @override
+    def init_empty(self) -> None:
+        super().init_empty()
+        self.len = 0
+
+    @override
+    def length(self) -> int:
+        return self.len
+
+    @override
+    def _get(self, idx: TIndex, channels_idx: list[int]) -> ReaderData:
+        """
+        Compute condition variables for a given time window.
+
+        Parameters
+        ----------
+        idx : TIndex
+            Index of temporal window
+        channels_idx : list[int]
+            Selection of channels
+
+        Returns
+        -------
+        ReaderData providing coords, geoinfos, data, datetimes
+        """
+
+        dtr = self.time_window_handler.window(idx)
+        encoded_condtions = self._encode(dtr, self.variables)
+        return encoded_condtions
+
+    def _encode(self, dtr: DTRange, variables: list[str]):
+        """
+        Encode start/end datetimes into condition variable values.
+
+        Parameters
+        ----------
+        start_dt :
+            start of time window
+        end_dt :
+            end of time window
+
+        Returns
+        -------
+        np.ndarray of shape (num_channels,)
+        """
+
+        _periods: dict[str, float] = {
+            "start_day": 365.0,
+            "end_day": 365.0,
+            "start_time": 24.0,
+            "end_time": 24.0,
+        }
+        _raw: dict[str, float] = {
+            # fractional day: day 25 at 12:00 → 25.5, unique signal per 6h timestep
+            "start_day": _day_of_year(dtr.start) + _hour_of_day(dtr.start) / 24.0,
+            "start_time": _hour_of_day(dtr.start),
+            "end_day": _day_of_year(dtr.end) + _hour_of_day(dtr.end) / 24.0,
+            "end_time": _hour_of_day(dtr.end),
+        }
+
+        values: list[float] = []
+
+        if self.transform == "absolute":
+            for var in variables:
+                values.append(_raw[var])
+
+        if self.transform == "absolute_normalized":
+            for var in variables:
+                values.append(_raw[var] / _periods[var])
+
+        elif self.transform == "cos_sin":
+            for var in variables:
+                angle = 2.0 * math.pi * _raw[var] / _periods[var]
+                values.append(math.cos(angle))
+                values.append(math.sin(angle))
+
+        elif self.transform == "fourier":
+            num_freqs = self.emb_dimension // 2
+            for var in variables:
+                for k in range(1, num_freqs + 1):
+                    angle = 2.0 * math.pi * k * _raw[var] / _periods[var]
+                    values.append(math.cos(angle))
+                    values.append(math.sin(angle))
+
+        return values
+
+
+def _day_of_year(dt: np.datetime64) -> float:
+    """Return 1-indexed day of year for a numpy datetime64."""
+    jan1 = dt.astype("datetime64[Y]").astype("datetime64[D]")
+    return float((dt.astype("datetime64[D]") - jan1).astype(int) + 1)
+
+
+def _hour_of_day(dt: np.datetime64) -> float:
+    """Return fractional hour of day [0, 24) for a numpy datetime64."""
+    day = dt.astype("datetime64[D]")
+    return float((dt - day) / np.timedelta64(1, "h"))
diff --git a/packages/readers_extra/src/weathergen/readers_extra/data_reader_forcing.py b/packages/readers_extra/src/weathergen/readers_extra/data_reader_forcing.py
new file mode 100644
index 0000000000..67ed0c3043
--- /dev/null
+++ b/packages/readers_extra/src/weathergen/readers_extra/data_reader_forcing.py
@@ -0,0 +1,39 @@
+# (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.
+
+import logging
+
+from weathergen.datasets.data_reader_anemoi import DataReaderAnemoi
+
+_logger = logging.getLogger(__name__)
+
+
+class DataReaderForcing(DataReaderAnemoi):
+    """
+    Per-step forcing reader (e.g. SST) backed by an anemoi dataset.
+
+    A forcing stream is a prescribed field that is *not predicted* and is re-read at
+    every forecast step's valid time, then injected into the forecasting engine
+    (``entry_point: forecast_engine``). This differs from:
+
+      * a normal ``anemoi`` source stream, which is only assimilated in the analysis
+        window (t <= 0), and
+      * a ``condition`` stream, which is a global scalar per step (no spatial field).
+
+    All the actual reading (open_dataset, channel selection, normalisation, coord
+    handling) is inherited unchanged from :class:`DataReaderAnemoi`; this subclass
+    exists purely as a dedicated stream *type* so the sampler can route it through
+    the per-step forcing path and the model can select an injection mode. The reader
+    is constructed with the same signature as ``DataReaderAnemoi`` (tw_handler,
+    filename, stream_info, stage).
+    """
+
+    # Marker consumed by the sampler to route this stream to the per-step forcing
+    # path rather than the assimilation (t<=0 source) path.
+    is_per_step_forcing: bool = True
diff --git a/packages/readers_extra/src/weathergen/readers_extra/registry.py b/packages/readers_extra/src/weathergen/readers_extra/registry.py
index b13ee46550..b3f4ea516a 100644
--- a/packages/readers_extra/src/weathergen/readers_extra/registry.py
+++ b/packages/readers_extra/src/weathergen/readers_extra/registry.py
@@ -32,5 +32,9 @@ def get_extra_reader(stream_type: str) -> object | None:
             from weathergen.readers_extra.data_reader_fesom import DataReaderFesom
 
             return DataReaderFesom
+        case "forcing":
+            from weathergen.readers_extra.data_reader_forcing import DataReaderForcing
+
+            return DataReaderForcing
         case _:
             return None

From 868059495a86376d7bb5aa167b66d67a6623b58c Mon Sep 17 00:00:00 2001
From: ankitpatnala <ankitpatnala@gmail.com>
Date: Wed, 9 Sep 2026 22:38:39 +0200
Subject: [PATCH 3/6] integrated into data reader so conditions are also read
 in the batch

---
 src/weathergen/datasets/batch.py              |  21 +++
 src/weathergen/datasets/data_reader_base.py   |  18 ++
 .../datasets/multi_stream_data_sampler.py     | 156 ++++++++++++++++++
 3 files changed, 195 insertions(+)

diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py
index ea6a0b26ab..0feac58dce 100644
--- a/src/weathergen/datasets/batch.py
+++ b/src/weathergen/datasets/batch.py
@@ -176,6 +176,15 @@ def __init__(
         self.output_steps = output_steps
         self.output_idxs = output_idxs
         self.device = None
+        # per-forecast-step encoded condition values, filled by
+        # MultiStreamDataSampler._build_condition_data; a tensor after to_device
+        self.conditions = [[] for _ in range(output_steps)]
+        # per-forecast-step spatial forcing (e.g. SST). Either a (num_cells, 2*num_vars)
+        # array of [cell_values | cell_valid] (fixed scatter-mean, the default), or the raw
+        # (num_points, num_vars) grid when learned_pool is on. A tensor after to_device.
+        self.forcing = [[] for _ in range(output_steps)]
+        # fixed point -> HEALPix-cell index, only set when learned_pool is enabled
+        self.forcing_cell_idx: torch.Tensor | None = None
 
     def __len__(self) -> int:
         return len(self.samples)
@@ -188,6 +197,18 @@ def to_device(self, device):
             self.tokens_lens.to(device, non_blocking=True) if self.tokens_lens is not None else None
         )
 
+        self.conditions = [
+            torch.tensor(cond, dtype=torch.float32, device=device) if len(cond) > 0 else cond
+            for cond in self.conditions
+        ]
+
+        self.forcing = [
+            torch.as_tensor(f, dtype=torch.float32, device=device) if len(f) > 0 else f
+            for f in self.forcing
+        ]
+        if self.forcing_cell_idx is not None:
+            self.forcing_cell_idx = self.forcing_cell_idx.to(device, non_blocking=True)
+
         self.device = device
 
         return self
diff --git a/src/weathergen/datasets/data_reader_base.py b/src/weathergen/datasets/data_reader_base.py
index 2dccde9301..ca2c5012e5 100644
--- a/src/weathergen/datasets/data_reader_base.py
+++ b/src/weathergen/datasets/data_reader_base.py
@@ -371,6 +371,24 @@ def __len__(self) -> int:
 
         return self.length()
 
+    def get_condition(self, idx: TIndex) -> ReaderData:
+        """
+        Get condition data for idx.
+
+        Only condition readers define `condition_idx`.
+
+        Parameters
+        ----------
+        idx : int
+            Index of temporal window
+
+        Returns
+        -------
+        condition data
+        """
+
+        return self._get(idx, self.condition_idx)
+
     def get_source(self, idx: TIndex) -> ReaderData:
         """
         Get source data for idx
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index bd4a8ad87e..5d6f98de6e 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -32,6 +32,7 @@
 from weathergen.datasets.utils import (
     get_tokens_lens,
 )
+from weathergen.readers_extra.data_reader_condition import DataReaderCondition
 from weathergen.readers_extra.registry import get_extra_reader
 from weathergen.train.utils import Stage, get_batch_size_from_config
 from weathergen.utils.distributed import is_root
@@ -145,6 +146,8 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
         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.condition_datasets = self._init_condition_streams(cf)
+        self.forcing_datasets = self._init_forcing_streams(cf)
 
         # RNG seed setup
         rs = cf.data_loading.rng_seed
@@ -217,10 +220,79 @@ def _calc_baseperms(self, fsm: int) -> np.typing.NDArray:
 
         return np.arange(self.max_input_steps, perms_len)
 
+    def _init_condition_streams(self, cf) -> dict[StreamName, _Stream]:
+        """Instantiate and register a condition stream (no backing files required)."""
+        condition_datasets: dict[StreamName, _Stream] = {}
+        for stream_name, stream_info in cf.streams.items():
+            if stream_info["type"] != "condition":
+                continue
+            condition_datasets[stream_name] = _Stream(stream_info, [])
+            if is_root():
+                logger.info(f"Opening condition dataset from stream config {stream_info['name']}.")
+            ds = DataReaderCondition(
+                tw_handler=self.time_window_handler, stream_info=stream_info, filename=None
+            )
+            condition_datasets[stream_name].readers += [ds]
+        return condition_datasets
+
+    def _init_forcing_streams(self, cf) -> dict[StreamName, _Stream]:
+        """
+        Instantiate per-step forcing streams (type: forcing, e.g. SST).
+
+        Read at every forecast step's valid time and injected into the forecasting engine
+        (they do not go through the assimilation/source path). The forcing grid is fixed, so
+        the scatter index onto the HEALPix cells is precomputed once per reader.
+        """
+        from weathergen.model.forcing import build_forcing_cell_index
+
+        forcing_datasets: dict[StreamName, _Stream] = {}
+        for stream_name, stream_info in cf.streams.items():
+            if stream_info["type"] != "forcing":
+                continue
+            stream = _Stream(stream_info, [])
+            stream_info["data_paths"] = cf.get("data_paths", [])
+            dataset = get_extra_reader(stream_info["type"])
+
+            for fname in stream_info["filenames"]:
+                fname = pathlib.Path(fname)
+                if fname.exists():
+                    filename = fname
+                else:
+                    filenames = [pathlib.Path(path) / fname for path in cf.data_paths]
+                    if not any(f.exists() for f in filenames):
+                        raise FileNotFoundError(
+                            f"Did not find input data for forcing stream '{stream_name}': "
+                            f"{filenames}."
+                        )
+                    filename = filenames[0]
+
+                ds = dataset(
+                    tw_handler=self.time_window_handler,
+                    stream_info=stream_info,
+                    stage=self._stage,
+                    filename=filename,
+                )
+                # precompute the fixed forcing-grid -> HEALPix-cell scatter index
+                ds.forcing_cell_idx = build_forcing_cell_index(
+                    ds.latitudes, ds.longitudes, self.healpix_level
+                )
+                stream.readers += [ds]
+                if is_root():
+                    logger.info(
+                        f"Opening forcing dataset '{stream_name}' from {filename} "
+                        f"(scatter onto {self.num_healpix_cells} cells)."
+                    )
+            forcing_datasets[stream_name] = stream
+        return forcing_datasets
+
     def _init_stream_datasets(self, cf) -> 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():
+            # condition and per-step forcing streams feed the forecasting engine directly,
+            # not the assimilation/source path; they are initialised separately.
+            if stream_info["type"] in ("condition", "forcing"):
+                continue
             stream_info["data_paths"] = cf.get("data_paths", [])
             # list of sources for current stream
             streams_datasets[stream_name] = _Stream(stream_info, [])
@@ -360,6 +432,9 @@ def get_sources_size(self):
             for ds in self.streams_datasets.values()
         ]
 
+    def get_condition_num_channels(self):
+        return sum([ds.readers[0].num_channels for ds in self.condition_datasets.values()])
+
     def get_sources_num_channels(self):
         return [ds.readers[0].get_source_num_channels() for ds in self.streams_datasets.values()]
 
@@ -649,6 +724,79 @@ def _preprocess_model_batch(
 
         return batch
 
+    def _build_condition_data(
+        self,
+        batch: ModelBatch,
+        condition_ds: AnyDataReader,
+        base_idx: TIndex,
+        num_output_steps: int,
+    ) -> None:
+        """
+        Collect encoded condition values for each forecast step into the batch.
+
+        Parameters
+        ----------
+        condition_ds :
+            The condition reader (DataReaderCondition instance).
+        base_idx :
+            Base time index for this sample.
+        num_output_steps :
+            Total number of output/forecast steps.
+        """
+
+        for i in range(num_output_steps):
+            condition_data = condition_ds.get_condition(
+                base_idx + (self.time_step * i) // self.step_timedelta
+            )
+            batch.get_source_samples().conditions[i] += condition_data
+
+    def _build_forcing_data(
+        self,
+        batch: ModelBatch,
+        forcing_ds: AnyDataReader,
+        base_idx: TIndex,
+        num_output_steps: int,
+    ) -> None:
+        """
+        Collect the per-step forcing field (e.g. SST) for each forecast step.
+
+        Mirrors `_build_condition_data`, but spatial: the reader is read at each step's
+        valid time (raw, full grid, normalised by the reader stats).
+
+        Two modes, from `injection.learned_pool` in the forcing stream's config:
+          * False (default): scatter-averaged onto the HEALPix cells here (CPU, no gradient),
+            stored per step as (num_cells, 2*num_vars) = [cell_values | cell_valid].
+          * True: the raw (num_points, num_vars) grid is stored and pooling happens on the
+            model side (LearnedForcingPool), so the pooling itself can be trained.
+        """
+        from weathergen.model.forcing import scatter_to_cells
+
+        learned_pool = forcing_ds.stream_info.get("injection", {}).get("learned_pool", False)
+        cell_idx = forcing_ds.forcing_cell_idx
+        num_grid_pts = len(cell_idx)
+        source_samples = batch.get_source_samples()
+        if learned_pool:
+            source_samples.forcing_cell_idx = cell_idx
+
+        for i in range(num_output_steps):
+            idx = base_idx + (self.time_step * i) // self.step_timedelta
+            rdata = forcing_ds.get_source(idx)
+            # full-grid values in fixed grid order; if the reader stacked multiple input
+            # steps, keep only the most recent grid
+            values = forcing_ds.normalize_source_channels(rdata.data)
+            values = np.asarray(values)[-num_grid_pts:]
+
+            if learned_pool:
+                source_samples.forcing[i] = values
+            else:
+                values_t = torch.as_tensor(values, dtype=torch.float32)
+                cell_values, cell_valid = scatter_to_cells(
+                    values_t.unsqueeze(0), cell_idx, self.num_healpix_cells
+                )
+                source_samples.forcing[i] = torch.cat(
+                    [cell_values[0], cell_valid[0]], dim=-1
+                ).numpy()
+
     def _get_batch(self, idx: int, num_forecast_steps: int):
         """
         Assemble a batch using the sample corresponding to idx
@@ -751,6 +899,14 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
                 ]
                 batch.add_target_stream(tidx, student_indices, stream_name, sdata, target_metadata)
 
+        # for condition streams
+        for _, condition_data in self.condition_datasets.items():
+            self._build_condition_data(batch, condition_data.readers[0], idx, num_output_steps)
+
+        # for per-step forcing streams (e.g. SST)
+        for _, forcing_data in self.forcing_datasets.items():
+            self._build_forcing_data(batch, forcing_data.readers[0], idx, num_output_steps)
+
         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()

From fae6c67196ff9b4fbb474a6ca1d8330c955ad20e Mon Sep 17 00:00:00 2001
From: ankitpatnala <ankitpatnala@gmail.com>
Date: Fri, 11 Sep 2026 13:44:31 +0200
Subject: [PATCH 4/6] resolved merge conflict

---
 src/weathergen/model/blocks.py          |   7 +-
 src/weathergen/model/engines.py         |  26 +-
 src/weathergen/model/forcing.py         | 469 ++++++++++++++++++++++++
 src/weathergen/model/model.py           | 103 +++++-
 src/weathergen/model/model_interface.py |  27 +-
 5 files changed, 601 insertions(+), 31 deletions(-)
 create mode 100644 src/weathergen/model/forcing.py

diff --git a/src/weathergen/model/blocks.py b/src/weathergen/model/blocks.py
index f8b4facc97..c9256213b6 100644
--- a/src/weathergen/model/blocks.py
+++ b/src/weathergen/model/blocks.py
@@ -16,7 +16,7 @@
 )
 from weathergen.model.layers import MLP
 from weathergen.model.norms import AdaLayerNormLayer
-from weathergen.utils.utils import get_dtype
+from weathergen.utils.utils import get_dtype, is_stream_fe_only
 
 
 class SelfAttentionBlock(nn.Module):
@@ -201,7 +201,10 @@ def __init__(
 
         self.block = nn.ModuleList()
 
-        target_readout_num_heads = next(self.cf.streams.values())["target_readout"]["num_heads"]
+        # first real data stream (condition/forcing streams have no target_readout)
+        target_readout_num_heads = next(
+            s for s in self.cf.streams.values() if not is_stream_fe_only(s)
+        )["target_readout"]["num_heads"]
 
         # Multi-Cross Attention Head
         self.block.append(
diff --git a/src/weathergen/model/engines.py b/src/weathergen/model/engines.py
index fde31213b6..2e31c7c578 100644
--- a/src/weathergen/model/engines.py
+++ b/src/weathergen/model/engines.py
@@ -30,7 +30,7 @@
 )
 from weathergen.model.layers import MLP
 from weathergen.model.utils import ActivationFactory
-from weathergen.utils.utils import get_dtype
+from weathergen.utils.utils import get_dtype, is_stream_fe_only
 
 
 class EmbeddingEngine(torch.nn.Module):
@@ -49,8 +49,16 @@ def __init__(self, cf: Config, sources_size) -> None:
         self.sources_size = sources_size  # KCT:iss130, what is this?
         self.embeds = torch.nn.ModuleDict()
         self.streams = cf.streams
-
-        for i, (stream_name, si) in enumerate(self.streams.items()):
+        # condition/forcing streams feed the FE directly and are omitted from
+        # get_sources_size(), so they must be omitted here too or the indices misalign
+        self.data_stream_names = [
+            name for name, cfg in cf.streams.items() if not is_stream_fe_only(cfg)
+        ]
+        self.data_streams = [cfg for cfg in cf.streams.values() if not is_stream_fe_only(cfg)]
+
+        for i, (si, stream_name) in enumerate(
+            zip(self.data_streams, self.data_stream_names, strict=True)
+        ):
             if si.get("diagnostic", False) or self.sources_size[i] == 0:
                 self.embeds[stream_name] = torch.nn.Identity()
                 continue
@@ -88,7 +96,7 @@ def forward(self, batch, pe_embed):
 
         # iterate over all streams
         x_embeds = []
-        for stream_name in self.streams.keys():
+        for stream_name in self.data_stream_names:
             # collect all source tokens from all input_steps and all samples in the batch
             sdata = []
             for istep in range(num_steps_input):
@@ -620,14 +628,15 @@ def init_weights_final(m):
         for block in self.fe_blocks:
             block.apply(init_weights_final)
 
-    def forward(self, tokens, fstep, coords=None):
+    def forward(self, tokens, condition, coords=None):
         if self.training:
             # Impute noise to the latent state
             noise_std = self.cf.get("fe_impute_latent_noise_std", 0.0)
             if noise_std > 0.0:
                 tokens = tokens + torch.randn_like(tokens) * torch.norm(tokens) * noise_std
 
-        aux_info = None
+        # condition stream (and any forcing routed through it) drives the FE AdaLN
+        aux_info = None if condition is None or len(condition) == 0 else condition.to(tokens.dtype)
         for _b_idx, block in enumerate(self.fe_blocks):
             if isinstance(block, torch.nn.modules.normalization.LayerNorm):
                 tokens = checkpoint(block, tokens, use_reentrant=False)
@@ -864,7 +873,10 @@ def __init__(
         self.pos_embed = nn.Parameter(torch.zeros(1, 9, self.cf.ae_global_dim_embed))
         dim_aux = self.cf.ae_global_dim_embed
 
-        target_readout_num_heads = next(self.cf.streams.values())["target_readout"]["num_heads"]
+        # first real data stream (condition/forcing streams have no target_readout)
+        target_readout_num_heads = next(
+            s for s in self.cf.streams.values() if not is_stream_fe_only(s)
+        )["target_readout"]["num_heads"]
         for ith, dim in enumerate(self.dims_embed[:-1]):
             if self.cf.decoder_type == "PerceiverIO":
                 # a single cross attention layer as per https://arxiv.org/pdf/2107.14795
diff --git a/src/weathergen/model/forcing.py b/src/weathergen/model/forcing.py
new file mode 100644
index 0000000000..7fff83af9d
--- /dev/null
+++ b/src/weathergen/model/forcing.py
@@ -0,0 +1,469 @@
+# (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.
+
+"""
+Generic per-step spatial forcing for the forecasting engine.
+
+A *forcing* is a prescribed field (SST, sea ice, aerosols, a CO2 field, solar forcing, ...)
+that is re-read at every forecast step's valid time (see the ``type: forcing`` stream) and
+injected into the forecasting engine so the rolling latent state is continually reminded of
+an external boundary condition. Nothing here is specific to any particular field -- SST is
+simply the first instance.
+
+To inject a forcing it is first placed on the FE latent grid: the (fixed) forcing point cloud
+is binned (scatter-mean) onto the HEALPix cells the FE tokens live on, then embedded to the
+model width by a small MLP. The result -- a per-cell forcing embedding -- is shared by every
+injection mode (``additive`` / ``cross_attn`` / ``global`` / ``adaln_local``).
+
+The scatter index depends only on the (fixed) forcing grid and the HEALPix level, so it is
+precomputed once via :func:`build_forcing_cell_index`. Everything here is cheap: one scatter
+and one matmul per forecast step, independent of the assimilation engine.
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+from numpy.typing import NDArray
+
+from weathergen.datasets.utils import coords_to_hpyidxs
+
+
+def build_forcing_cell_index(
+    latitudes: NDArray, longitudes: NDArray, healpix_level: int
+) -> torch.Tensor:
+    """
+    Map each forcing grid point to its HEALPix latent cell (nested convention).
+
+    Uses the same convention (``ang2pix(..., nest=True)``) as the tokenizer, so the binned
+    forcing cells align one-to-one with the FE latent token cells.
+
+    Parameters
+    ----------
+    latitudes, longitudes :
+        Forcing grid point coordinates, degrees, shape ``(num_points,)``.
+    healpix_level :
+        HEALPix level of the FE latent grid (num_cells = 12 * 4**healpix_level).
+
+    Returns
+    -------
+    LongTensor of shape ``(num_points,)`` with the cell index in ``[0, num_cells)`` for
+    every forcing point.
+    """
+    cell_idx = coords_to_hpyidxs(healpix_level, np.asarray(latitudes), np.asarray(longitudes))
+    return torch.as_tensor(np.asarray(cell_idx), dtype=torch.long)
+
+
+def scatter_to_cells(
+    values: torch.Tensor, cell_idx: torch.Tensor, num_cells: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Scatter-mean per-point forcing values onto HEALPix cells, ignoring NaNs.
+
+    Parameters
+    ----------
+    values :
+        Forcing point values, shape ``(..., num_points, num_vars)``. NaNs (e.g. SST over
+        land) are treated as missing and excluded from the per-cell mean.
+    cell_idx :
+        Cell index per point, shape ``(num_points,)`` (from ``build_forcing_cell_index``).
+    num_cells :
+        Number of HEALPix cells.
+
+    Returns
+    -------
+    (cell_values, cell_valid)
+        cell_values : ``(..., num_cells, num_vars)`` per-cell mean (0 where no valid point
+            contributed to a cell/var).
+        cell_valid  : ``(..., num_cells, num_vars)`` fraction of valid points that
+            contributed (0 = no data, e.g. all-land cell), usable as an input feature/mask.
+    """
+    *batch_shape, num_points, num_vars = values.shape
+    flat = values.reshape(-1, num_points, num_vars)  # (B, P, V)
+    n_batch = flat.shape[0]
+    device = flat.device
+
+    valid = torch.isfinite(flat)  # (B, P, V)
+    filled = torch.where(valid, flat, torch.zeros_like(flat))
+
+    idx = cell_idx.to(device).view(1, num_points, 1).expand(n_batch, num_points, num_vars)
+
+    sums = torch.zeros(n_batch, num_cells, num_vars, device=device, dtype=flat.dtype)
+    sums.scatter_add_(1, idx, filled)
+    counts = torch.zeros(n_batch, num_cells, num_vars, device=device, dtype=flat.dtype)
+    counts.scatter_add_(1, idx, valid.to(flat.dtype))
+
+    cell_values = torch.where(counts > 0, sums / counts.clamp_min(1.0), torch.zeros_like(sums))
+    total = torch.zeros(n_batch, num_cells, num_vars, device=device, dtype=flat.dtype)
+    total.scatter_add_(1, idx, torch.ones_like(filled))
+    cell_valid = torch.where(total > 0, counts / total.clamp_min(1.0), torch.zeros_like(counts))
+
+    cell_values = cell_values.reshape(*batch_shape, num_cells, num_vars)
+    cell_valid = cell_valid.reshape(*batch_shape, num_cells, num_vars)
+    return cell_values, cell_valid
+
+
+class LearnedForcingPool(nn.Module):
+    """
+    Learnable replacement for the fixed scatter-mean in :func:`scatter_to_cells`.
+
+    HEALPix cells don't have a fixed number of forcing grid points each (coastal cells vs.
+    open-ocean cells differ), so a standard attention module (fixed-size Q/K/V + masking)
+    doesn't fit directly. This instead scores every point with a small per-point MLP and
+    normalises the scores *within each cell* via a segment-softmax (two ``scatter_add_``
+    calls, the same primitive ``scatter_to_cells`` already uses for the mean) -- so cells
+    with more points just have more terms in their softmax, no padding required.
+
+    Same NaN-masking discipline as ``scatter_to_cells``: invalid (e.g. land) points get
+    zero weight rather than corrupting the pooled value -- and, like ``scatter_to_cells``,
+    each variable is masked/normalised *independently*, since e.g. ``land_sea_mask`` is
+    defined everywhere while ``sst``/``sea_ice_cover`` are NaN over land at the same points.
+    """
+
+    def __init__(self, num_vars: int, hidden_factor: int = 4) -> None:
+        super().__init__()
+        self.num_vars = num_vars
+        hidden = max(hidden_factor * num_vars, 8)
+        # per-variable logits: shape (P, V) out, not a single shared per-point score, so each
+        # variable's segment-softmax only involves points valid for *that* variable.
+        self.score = nn.Sequential(
+            nn.Linear(num_vars, hidden),
+            nn.SiLU(),
+            nn.Linear(hidden, num_vars),
+        )
+
+    def forward(
+        self, values: torch.Tensor, cell_idx: torch.Tensor, num_cells: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters
+        ----------
+        values : ``(num_points, num_vars)`` raw per-point forcing values for this step.
+        cell_idx : ``(num_points,)`` cell index per point (from ``build_forcing_cell_index``).
+        num_cells : number of HEALPix cells.
+
+        Returns
+        -------
+        (cell_values, cell_valid), same shapes/semantics as ``scatter_to_cells``:
+        ``cell_values`` a ``(num_cells, num_vars)`` learned weighted mean (0 where no valid
+        point contributed), ``cell_valid`` the ``(num_cells, num_vars)`` valid-point fraction.
+        """
+        num_points, num_vars = values.shape
+        device = values.device
+        cell_idx = cell_idx.to(device)
+        idx = cell_idx.view(num_points, 1).expand(num_points, num_vars)  # (P, V)
+
+        valid = torch.isfinite(values)  # (P, V)
+        filled = torch.where(valid, values, torch.zeros_like(values))
+
+        logits = self.score(filled)  # (P, V) -- one logit per (point, variable)
+        logits = logits.masked_fill(~valid, -1e9)
+
+        # segment-softmax per variable: normalise exp(logits) within each cell, independently
+        # per column, using the same scatter_add_ primitive scatter_to_cells uses for the mean.
+        exp_logits = torch.exp(logits - logits.detach().amax(dim=0, keepdim=True))
+        denom = torch.zeros(num_cells, num_vars, device=device, dtype=values.dtype)
+        denom.scatter_add_(0, idx, exp_logits)
+        weights = exp_logits / denom.gather(0, idx).clamp_min(1e-12)  # (P, V)
+
+        cell_values = torch.zeros(num_cells, num_vars, device=device, dtype=values.dtype)
+        cell_values.scatter_add_(0, idx, weights * filled)
+
+        counts = torch.zeros(num_cells, num_vars, device=device, dtype=values.dtype)
+        counts.scatter_add_(0, idx, valid.to(values.dtype))
+        total = torch.zeros(num_cells, num_vars, device=device, dtype=values.dtype)
+        total.scatter_add_(0, idx, torch.ones_like(filled))
+        cell_valid = torch.where(total > 0, counts / total.clamp_min(1.0), torch.zeros_like(counts))
+
+        # cells with no valid points for a variable: weights there are garbage over masked-out
+        # (-1e9 logit) points; zero explicitly rather than rely on the exp(-1e9) underflow.
+        cell_values = torch.where(cell_valid > 0, cell_values, torch.zeros_like(cell_values))
+
+        return cell_values, cell_valid
+
+
+class ForcingEmbed(nn.Module):
+    """
+    Per-cell forcing embedding shared by all injection modes.
+
+    Normalises the per-cell forcing field with the dataset statistics (or identity if the
+    field is already normalised upstream), appends a validity feature (valid fraction per
+    cell/var) so the model can tell present from absent/no-data, and maps to the model width
+    with a small MLP.
+
+    Parameters
+    ----------
+    num_vars :
+        Number of forcing source channels (e.g. sea_surface_temperature, sea_ice_cover,
+        land_sea_mask -> 3).
+    dim_embed :
+        Output embedding width (typically the FE model dim for additive injection).
+    mean, stdev :
+        Per-channel normalisation statistics, shape ``(num_vars,)``. If None, the input is
+        assumed already normalised.
+    hidden_factor :
+        Width multiplier for the MLP hidden layer.
+    """
+
+    def __init__(
+        self,
+        num_vars: int,
+        dim_embed: int,
+        mean: NDArray | None = None,
+        stdev: NDArray | None = None,
+        hidden_factor: int = 2,
+    ) -> None:
+        super().__init__()
+        self.num_vars = num_vars
+        self.dim_embed = dim_embed
+
+        if mean is None:
+            mean = np.zeros(num_vars, dtype=np.float32)
+        if stdev is None:
+            stdev = np.ones(num_vars, dtype=np.float32)
+        # kept as plain numpy (not a tensor/buffer) so reset_parameters() can restore these
+        # exact values after an FSDP2 meta-device to_empty() wipes the buffers -- see
+        # ForcingInjection.reset_parameters() for why this is necessary. np.array(..., copy=True)
+        # (not asarray) so this is never aliased into the buffer below: torch.as_tensor shares
+        # memory with a same-dtype CPU numpy array rather than copying, so without an explicit
+        # copy here, in-place mutation of the buffer (e.g. the to_empty() corruption itself)
+        # would silently corrupt this "backup" too.
+        self._mean_init = np.array(mean, dtype=np.float32, copy=True)
+        self._stdev_init = np.array(stdev, dtype=np.float32, copy=True)
+        self.register_buffer("mean", torch.as_tensor(self._mean_init).clone())
+        self.register_buffer("stdev", torch.as_tensor(self._stdev_init).clone())
+
+        # input = normalised values (num_vars) concatenated with validity (num_vars)
+        in_dim = 2 * num_vars
+        hidden = hidden_factor * dim_embed
+        self.mlp = nn.Sequential(
+            nn.Linear(in_dim, hidden),
+            nn.SiLU(),
+            nn.Linear(hidden, dim_embed),
+        )
+
+    def reset_parameters(self) -> None:
+        """
+        Restore mean/stdev to their intended values and reinitialise the MLP.
+
+        FSDP2's meta-device build (``model.to_empty()`` + ``Model.reset_parameters()``)
+        only auto-resets ``nn.Linear``/``nn.LayerNorm`` submodules (see model.py); buffers
+        registered directly here (mean/stdev) are left as uninitialised memory otherwise,
+        which silently breaks the "identity if unspecified" normalisation this class
+        documents (observed in practice as stdev=0 -> division amplifying values ~1e6x).
+        """
+        self.mean.data.copy_(torch.as_tensor(self._mean_init, device=self.mean.device))
+        self.stdev.data.copy_(torch.as_tensor(self._stdev_init, device=self.stdev.device))
+        for module in self.mlp.modules():
+            if isinstance(module, nn.Linear):
+                module.reset_parameters()
+
+    def forward(self, cell_values: torch.Tensor, cell_valid: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        cell_values : ``(..., num_cells, num_vars)`` per-cell forcing values (from scatter).
+        cell_valid  : ``(..., num_cells, num_vars)`` per-cell validity fraction.
+
+        Returns
+        -------
+        ``(..., num_cells, dim_embed)`` per-cell forcing embedding.
+        """
+        mean = self.mean.to(cell_values.dtype)
+        stdev = self.stdev.to(cell_values.dtype).clamp_min(1e-6)
+        normed = (cell_values - mean) / stdev
+        # zero-out normalised value where there is no data so it reads as "absent"
+        normed = torch.where(cell_valid > 0, normed, torch.zeros_like(normed))
+        x = torch.cat([normed, cell_valid.to(normed.dtype)], dim=-1)
+        return self.mlp(x)
+
+
+class ForcingInjection(nn.Module):
+    """
+    Inject a per-step spatial forcing into the forecasting engine, one of four modes.
+
+    Applied once per forecast step, in the model rollout loop, just before the FE advances the
+    ``(B, num_tokens, dim)`` latent state. The auxiliary tokens (register/class) are left
+    untouched; only the per-cell patch tokens / the FE conditioning are affected. The forcing
+    field for a step is global across the batch (same forecast time), so a per-cell tensor
+    ``(num_cells, dim)`` broadcasts across the batch dimension.
+
+    Modes (all a **zero-initialised gated residual** -> exact no-op at init, so a model
+    fine-tuned from a pretrained checkpoint starts identically and learns to use the forcing):
+
+      * ``additive``    : add the per-cell forcing embedding to the cell tokens.
+      * ``cross_attn``  : cell tokens cross-attend to the per-cell forcing tokens (non-local
+                          coupling / teleconnections).
+      * ``global``      : pool the forcing over cells and add the projected vector to the
+                          (global) FE conditioning -- the cheap "index" flavour; the existing
+                          FE AdaLN is untouched.
+      * ``adaln_local`` : make the FE conditioning per-cell (daytime condition + projected
+                          per-cell forcing); every FE block's AdaLN then modulates each cell
+                          token by its local forcing. Requires a condition stream (dim_aux > 0).
+
+    ``forward`` returns ``(tokens, condition)``: token-nudge modes modify ``tokens``,
+    conditioning modes modify ``condition``.
+    """
+
+    _MODES = ("none", "additive", "cross_attn", "global", "adaln_local")
+
+    def __init__(
+        self,
+        num_vars: int,
+        dim_embed: int,
+        mode: str,
+        dim_aux: int = 0,
+        num_heads: int = 8,
+        learned_pool: bool = False,
+    ) -> None:
+        super().__init__()
+        assert mode in self._MODES, f"unknown forcing mode {mode!r}, expected {self._MODES}"
+        self.mode = mode
+        self.num_vars = num_vars
+        self.dim_embed = dim_embed
+        self.dim_aux = dim_aux
+        self.learned_pool_module = LearnedForcingPool(num_vars) if learned_pool else None
+
+        if mode == "none":
+            return
+
+        if mode in ("global", "adaln_local") and dim_aux <= 0:
+            raise ValueError(
+                f"forcing mode {mode!r} routes the forcing through the FE conditioning (AdaLN) "
+                "and needs a condition stream (dim_aux > 0); use additive/cross_attn otherwise."
+            )
+
+        # Forcing field is already normalised by the sampler (reader stats) -> identity here.
+        self.embed = ForcingEmbed(num_vars, dim_embed)
+        # zero-init gate -> injection starts as an exact no-op
+        self.gate = nn.Parameter(torch.zeros(1))
+
+        if mode == "cross_attn":
+            self.norm_q = nn.LayerNorm(dim_embed)
+            self.cross_attn = nn.MultiheadAttention(dim_embed, num_heads, batch_first=True)
+        elif mode in ("global", "adaln_local"):
+            self.cond_proj = nn.Linear(dim_embed, dim_aux)
+
+    def reset_parameters(self) -> None:
+        """
+        Restore the documented zero-init-gate no-op-at-init guarantee.
+
+        Model.reset_parameters()'s generic sweep (model.py) only auto-resets
+        nn.Linear/nn.LayerNorm submodules, so it's called explicitly on this module too --
+        self-contained (doesn't rely on that generic sweep for its own children) so it also
+        works standalone from model_interface.py's "new module not found in checkpoint" path
+        (module_to_init.reset_parameters()), which previously raised AttributeError here
+        since this method didn't exist.
+
+        Without this, gate/embed.mean/embed.stdev -- none of which are nn.Linear/nn.LayerNorm
+        -- are left as whatever uninitialised memory model.to_empty() produced (observed in
+        practice: stdev=0, amplifying forcing values ~1e6x once gate drifts off zero).
+        """
+        if self.mode == "none":
+            return
+        self.gate.data.zero_()
+        self.embed.reset_parameters()
+        if self.mode == "cross_attn":
+            self.norm_q.reset_parameters()
+            self.cross_attn._reset_parameters()  # nn.MultiheadAttention's own (underscored) API
+        elif self.mode in ("global", "adaln_local"):
+            self.cond_proj.reset_parameters()
+        if self.learned_pool_module is not None:
+            for module in self.learned_pool_module.score.modules():
+                if isinstance(module, nn.Linear):
+                    module.reset_parameters()
+
+    @staticmethod
+    def _is_empty(x) -> bool:
+        return x is None or not torch.is_tensor(x) or x.numel() == 0
+
+    def _cell_emb(
+        self,
+        forcing_field: torch.Tensor,
+        dtype,
+        cell_idx: torch.Tensor | None,
+        num_cells: int | None,
+    ) -> torch.Tensor:
+        """
+        Per-cell embedding (num_cells, dim).
+
+        Two input shapes for ``forcing_field``, selected by ``learned_pool_module``:
+          * fixed mean (default): already-scattered ``(num_cells, 2*num_vars)`` = [cell_values
+            | cell_valid], produced upstream in the dataloader (see ``scatter_to_cells``).
+          * learned pool: raw, unscattered ``(num_points, num_vars)`` per-point values --
+            pooled here, on the GPU, so ``learned_pool_module`` gets gradients.
+        """
+        if forcing_field.dim() == 3:  # tolerate a leading batch dim
+            forcing_field = forcing_field[0]
+
+        if self.learned_pool_module is not None:
+            assert cell_idx is not None and num_cells is not None, (
+                "learned_pool requires cell_idx/num_cells (raw per-point forcing_field)"
+            )
+            cell_values, cell_valid = self.learned_pool_module(forcing_field, cell_idx, num_cells)
+        else:
+            num_vars = forcing_field.shape[-1] // 2
+            cell_values = forcing_field[..., :num_vars]
+            cell_valid = forcing_field[..., num_vars:]
+
+        return self.embed(cell_values, cell_valid).to(dtype)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        condition: torch.Tensor,
+        forcing_field: torch.Tensor,
+        num_aux: int,
+        cell_idx: torch.Tensor | None = None,
+        num_cells: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters
+        ----------
+        tokens :
+            FE latent state ``(B, num_tokens, dim)`` (num_aux leading aux tokens then num_cells).
+        condition :
+            FE conditioning fed to the AdaLN, global ``(dim_aux,)`` vector (or empty).
+        forcing_field :
+            Forcing data for this step; empty when no forcing is available (then this is a
+            no-op). Shape depends on the pooling mode -- see ``_cell_emb``.
+        num_aux :
+            Number of leading auxiliary tokens to leave untouched.
+        cell_idx, num_cells :
+            Only required when ``learned_pool=True`` was passed to ``__init__`` -- the
+            point -> HEALPix-cell index and cell count needed to pool ``forcing_field`` here.
+
+        Returns
+        -------
+        (tokens, condition) with the forcing injected according to the mode.
+        """
+        if self.mode == "none" or self._is_empty(forcing_field):
+            return tokens, condition
+
+        cell_emb = self._cell_emb(
+            forcing_field, tokens.dtype, cell_idx, num_cells
+        )  # (num_cells, dim)
+
+        if self.mode == "additive":
+            patch = tokens[:, num_aux:] + self.gate * cell_emb
+            tokens = torch.cat([tokens[:, :num_aux], patch], dim=1)
+        elif self.mode == "cross_attn":
+            q = self.norm_q(tokens[:, num_aux:])  # (B, num_cells, dim)
+            kv = cell_emb.unsqueeze(0).expand(q.shape[0], -1, -1)
+            delta, _ = self.cross_attn(q, kv, kv, need_weights=False)
+            patch = tokens[:, num_aux:] + self.gate * delta
+            tokens = torch.cat([tokens[:, :num_aux], patch], dim=1)
+        elif self.mode == "global":
+            g = self.cond_proj(cell_emb.mean(dim=0))  # (dim_aux,)
+            condition = condition + self.gate * g
+        elif self.mode == "adaln_local":
+            # per-cell conditioning: broadcast the global daytime condition, add per-cell forcing
+            base = condition.reshape(1, -1) if condition.numel() > 0 else 0.0
+            condition = base + self.gate * self.cond_proj(cell_emb)  # (num_cells, dim_aux)
+
+        return tokens, condition
diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py
index 2d3bebb6de..b5f706f956 100644
--- a/src/weathergen/model/model.py
+++ b/src/weathergen/model/model.py
@@ -37,11 +37,12 @@
     TargetPredictionEngine,
     TargetPredictionEngineClassic,
 )
+from weathergen.model.forcing import ForcingInjection
 from weathergen.model.layers import MLP, NamedLinear
 from weathergen.model.utils import get_num_parameters
 from weathergen.train.loss_modules.utils import compute_cos_sim_to_prev
 from weathergen.utils.distributed import is_root
-from weathergen.utils.utils import get_dtype, is_stream_forcing
+from weathergen.utils.utils import get_dtype, is_stream_fe_only, is_stream_forcing
 
 logger = logging.getLogger(__name__)
 
@@ -302,7 +303,14 @@ class Model(torch.nn.Module):
         coordinates to its physical space.
     """
 
-    def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coords_size):
+    def __init__(
+        self,
+        cf: Config,
+        sources_size,
+        targets_num_channels,
+        targets_coords_size,
+        condition_num_channels: int = 0,
+    ):
         """
         Args:
             cf : Configuration with model parameters
@@ -329,6 +337,10 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord
         self.pred_heads = None
         self.q_cells: torch.Tensor | None = None
         self.streams: dict[str, typing.Any] = cf.streams
+        self.data_stream_names: list | None = None
+        self.data_streams: list | None = None
+        # width of the FE conditioning vector contributed by condition streams
+        self.forecast_aux_infos = condition_num_channels
         self.target_token_engines = None
 
         assert cf.get("forecast", {}).get("att_dense_rate", 1.0) == 1.0, (
@@ -380,7 +392,12 @@ def create(self) -> "Model":
 
         mode_cfg = cf.training_config
         if cf.fe_num_blocks > 0:
-            self.forecast_engine = ForecastingEngine(cf, mode_cfg, self.num_healpix_cells)
+            self.forecast_engine = ForecastingEngine(
+                cf,
+                mode_cfg,
+                self.num_healpix_cells,
+                self.forecast_aux_infos if self.forecast_aux_infos > 0 else None,
+            )
         else:
             self.forecast_engine = IdentityEngine()
 
@@ -390,7 +407,37 @@ def create(self) -> "Model":
         self.target_token_engines = torch.nn.ModuleDict()
         self.pred_heads = torch.nn.ModuleDict()
 
-        # determine stream names once so downstream components use consistent keys
+        # determine stream names once so downstream components use consistent keys.
+        # condition/forcing streams feed the FE directly, not the assimilation/decoder path.
+        self.data_stream_names = [
+            name for name, cfg in cf.streams.items() if not is_stream_fe_only(cfg)
+        ]
+        self.data_streams = [cfg for cfg in cf.streams.values() if not is_stream_fe_only(cfg)]
+
+        # per-step spatial forcing injected into the FE (SST etc.); None if no such stream
+        self.forcing_module = None
+        for stream_cfg in cf.streams.values():
+            if stream_cfg.get("type") != "forcing":
+                continue
+            inj = stream_cfg.get("injection", {})
+            mode = inj.get("mode", "none")
+            if mode == "none":
+                continue
+            source_channels = stream_cfg.get("train_source_channels") or stream_cfg.get(
+                "val_source_channels", []
+            )
+            # dim_embed must equal the FE token dim so additive/cross_attn line up; the
+            # global/adaln_local modes then project this down to the conditioning dim.
+            self.forcing_module = ForcingInjection(
+                num_vars=len(source_channels),
+                dim_embed=cf.ae_global_dim_embed,
+                mode=mode,
+                dim_aux=self.forecast_aux_infos,
+                num_heads=cf.fe_num_heads,
+                learned_pool=inj.get("learned_pool", False),
+            )
+            break  # one forcing stream supported for now
+
         loss_terms = [
             v.type for _, v in cf.training_config.losses.items() if v.get("enabled", True)
         ]
@@ -403,7 +450,8 @@ def create(self) -> "Model":
         self.compute_cos_sim_to_prev = "LossLatent" in loss_terms
 
         if "LossPhysical" in loss_terms:
-            for i_stream, (stream_name, si) in enumerate(self.streams.items()):
+            for i_stream, si in enumerate(self.data_streams):
+                stream_name = self.data_stream_names[i_stream]
                 # skip decoder if channels are empty
                 if is_stream_forcing(si):
                     continue
@@ -494,7 +542,8 @@ def create(self) -> "Model":
                     )
 
             # iterate again to setup shared spatial pred heads if specified in config
-            for i_stream, (stream_name, si) in enumerate(self.streams.items()):
+            for i_stream, si in enumerate(self.data_streams):
+                stream_name = self.data_stream_names[i_stream]
                 # skip decoder if channels are empty
                 if is_stream_forcing(si):
                     continue
@@ -593,12 +642,17 @@ def _reset_params(module):
 
         self.apply(_reset_params)
 
+        # forcing_module owns Parameters/buffers (gate, embed.mean/.stdev) that aren't
+        # nn.Linear/nn.LayerNorm, so the sweep above misses them -- reset explicitly.
+        if self.forcing_module is not None:
+            self.forcing_module.reset_parameters()
+
     def print_num_parameters(self) -> None:
         """Print number of parameters for entire model and each module used to build the model"""
 
         num_params_embed = [
             get_num_parameters(self.encoder.embed_engine.embeds[name])
-            for name in self.streams.keys()
+            for name in self.data_stream_names
         ]
         num_params_total = get_num_parameters(self)
         num_params_ae_local = get_num_parameters(self.encoder.ae_local_engine.ae_local_blocks)
@@ -621,17 +675,17 @@ def print_num_parameters(self) -> None:
         mdict = self.embed_target_coords
         num_params_embed_tcs = [
             get_num_parameters(mdict[name]) if mdict and name in mdict else 0
-            for name in self.streams.keys()
+            for name in self.data_stream_names
         ]
         mdict = self.target_token_engines
         num_params_tte = [
             get_num_parameters(mdict[name]) if mdict and name in mdict else 0
-            for name in self.streams.keys()
+            for name in self.data_stream_names
         ]
         mdict = self.pred_heads
         num_params_preds = [
             get_num_parameters(mdict[name]) if mdict and name in mdict else 0
-            for name in self.streams.keys()
+            for name in self.data_stream_names
         ]
 
         print("-----------------")
@@ -640,7 +694,7 @@ def print_num_parameters(self) -> None:
         print("  Embedding networks:")
         [
             print("    {} : {:,}".format(si["name"], np))
-            for si, np in zip(self.streams.values(), num_params_embed, strict=False)
+            for si, np in zip(self.data_streams, num_params_embed, strict=False)
         ]
         print(f" Local assimilation engine: {num_params_ae_local:,}")
         print(f" Local-global adapter: {num_params_ae_adapter:,}")
@@ -651,7 +705,7 @@ def print_num_parameters(self) -> None:
         print(f" Forecast engine: {num_params_fe:,}")
         print(" coordinate embedding, prediction networks and prediction heads:")
         zps = zip(
-            self.streams.keys(),
+            self.data_stream_names,
             num_params_embed_tcs,
             num_params_tte,
             num_params_preds,
@@ -701,6 +755,29 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput:
         fe_steps = 0
         for step in batch.get_output_idxs():
             without_grad = p_fwd and self.training and step != max(batch.get_output_idxs())
+
+            # FE conditioning for this step: the condition stream's encoded values, indexed
+            # by absolute forecast step (conditions[] is built for every step from 0).
+            condition = batch.conditions[step] if step < len(batch.conditions) else None
+            # inject the per-step spatial forcing (e.g. SST) into the latent state /
+            # conditioning before advancing; no-op when no forcing module or no field.
+            if self.forcing_module is not None:
+                forcing_list = getattr(batch, "forcing", None)
+                forcing_field = (
+                    forcing_list[step]
+                    if forcing_list is not None and step < len(forcing_list)
+                    else None
+                )
+                tokens, condition = self.forcing_module(
+                    tokens,
+                    condition,
+                    forcing_field,
+                    self.num_aux_tokens,
+                    # only used by learned_pool (raw field pooled on GPU); unused otherwise
+                    cell_idx=getattr(batch, "forcing_cell_idx", None),
+                    num_cells=self.num_healpix_cells,
+                )
+
             if without_grad:
                 # Pushforward mode: advance tokens without grad; no decoding with torch.no_grad():
                 tokens = self.forecast_engine(tokens, step, model_params.rope_coords)
@@ -791,7 +868,7 @@ def predict_decoders(
         tokens_nbors_lens[0] = 0
 
         # pair with tokens from assimilation engine to obtain target tokens
-        for stream_name in self.streams.keys():
+        for stream_name in self.data_stream_names:
             # extract target coords for current stream and fstep and convert to one tensor
             t_coords = [
                 batch.samples[i_b].streams_data[stream_name].target_coords[step]
diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py
index c0b475b156..1b07d4e99e 100644
--- a/src/weathergen/model/model_interface.py
+++ b/src/weathergen/model/model_interface.py
@@ -17,7 +17,7 @@
     MixedPrecisionPolicy,
     fully_shard,
 )
-from torch.distributed.tensor import distribute_tensor
+from torch.distributed.tensor import DTensor, distribute_tensor
 
 from weathergen.common.config import Config, get_path_model, merge_configs
 from weathergen.model.attention import (
@@ -200,13 +200,17 @@ def load_model(cf, model, device, run_id: str, mini_epoch=-1):
             if sharded_meta_param is None:
                 logger.warning(f"Parameter {param_name} from checkpoint not found in model.")
                 continue
-            sharded_tensor = distribute_tensor(
-                full_tensor,
-                sharded_meta_param.device_mesh,
-                sharded_meta_param.placements,
-            )
-            # maybe_sharded_sd[param_name.replace("module.", "")] = nn.Parameter(sharded_tensor)
-            maybe_sharded_sd[param_name] = torch.nn.Parameter(sharded_tensor)
+            if isinstance(sharded_meta_param, DTensor):
+                sharded_tensor = distribute_tensor(
+                    full_tensor,
+                    sharded_meta_param.device_mesh,
+                    sharded_meta_param.placements,
+                )
+                maybe_sharded_sd[param_name] = torch.nn.Parameter(sharded_tensor)
+            else:
+                # buffers (e.g. ForcingEmbed.mean/stdev) are not sharded by FSDP2 -- assign
+                # the full tensor as-is, moved to the target device.
+                maybe_sharded_sd[param_name] = full_tensor.to(device)
         # choose `assign=True` for sharded model since we cannot call `copy_` on meta tensor
         mkeys, ukeys = model.load_state_dict(maybe_sharded_sd, strict=False, assign=True)
 
@@ -274,8 +278,13 @@ def get_model(cf: Config, training_mode: TrainingMode, dataset, overrides):
     sources_size = dataset.get_sources_size()
     targets_num_channels = dataset.get_targets_num_channels()
     targets_coords_size = dataset.get_targets_coords_size()
+    condition_num_channels = dataset.get_condition_num_channels()
 
     cf_with_overrides = merge_configs(cf, overrides)
     return Model(
-        cf_with_overrides, sources_size, targets_num_channels, targets_coords_size
+        cf_with_overrides,
+        sources_size,
+        targets_num_channels,
+        targets_coords_size,
+        condition_num_channels,
     ).create()

From e1b7f65fa2c4344d4d268c2adb422de312ab8bf6 Mon Sep 17 00:00:00 2001
From: ankitpatnala <ankitpatnala@gmail.com>
Date: Fri, 11 Sep 2026 14:01:22 +0200
Subject: [PATCH 5/6] declared some variables inside init function as None and
 remove getattr function

---
 src/weathergen/model/model.py | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py
index b5f706f956..5daeb55893 100644
--- a/src/weathergen/model/model.py
+++ b/src/weathergen/model/model.py
@@ -334,6 +334,8 @@ def __init__(
         self.embed_target_coords = None
         self.encoder: EncoderModule | None = None
         self.forecast_engine: ForecastingEngine | IdentityEngine | None = None
+        self.forcing_module: ForcingInjection | None = None
+        self.compute_cos_sim_to_prev = False
         self.pred_heads = None
         self.q_cells: torch.Tensor | None = None
         self.streams: dict[str, typing.Any] = cf.streams
@@ -414,8 +416,7 @@ def create(self) -> "Model":
         ]
         self.data_streams = [cfg for cfg in cf.streams.values() if not is_stream_fe_only(cfg)]
 
-        # per-step spatial forcing injected into the FE (SST etc.); None if no such stream
-        self.forcing_module = None
+        # per-step spatial forcing injected into the FE (SST etc.); stays None if no such stream
         for stream_cfg in cf.streams.values():
             if stream_cfg.get("type") != "forcing":
                 continue
@@ -753,40 +754,38 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput:
         # roll-out in latent space, iterate and generate output over requested output steps
         # the first FE step leaves the encoder regime, so it is excluded from the cosine band
         fe_steps = 0
+        # conditions/forcing are filled by the sampler on the source samples, not on the batch
+        source_samples = batch.get_source_samples()
+        conditions = source_samples.conditions
+        forcing = source_samples.forcing
         for step in batch.get_output_idxs():
             without_grad = p_fwd and self.training and step != max(batch.get_output_idxs())
 
             # FE conditioning for this step: the condition stream's encoded values, indexed
             # by absolute forecast step (conditions[] is built for every step from 0).
-            condition = batch.conditions[step] if step < len(batch.conditions) else None
+            condition = conditions[step] if step < len(conditions) else None
             # inject the per-step spatial forcing (e.g. SST) into the latent state /
             # conditioning before advancing; no-op when no forcing module or no field.
             if self.forcing_module is not None:
-                forcing_list = getattr(batch, "forcing", None)
-                forcing_field = (
-                    forcing_list[step]
-                    if forcing_list is not None and step < len(forcing_list)
-                    else None
-                )
                 tokens, condition = self.forcing_module(
                     tokens,
                     condition,
-                    forcing_field,
+                    forcing[step] if step < len(forcing) else None,
                     self.num_aux_tokens,
                     # only used by learned_pool (raw field pooled on GPU); unused otherwise
-                    cell_idx=getattr(batch, "forcing_cell_idx", None),
+                    cell_idx=source_samples.forcing_cell_idx,
                     num_cells=self.num_healpix_cells,
                 )
 
             if without_grad:
                 # Pushforward mode: advance tokens without grad; no decoding with torch.no_grad():
-                tokens = self.forecast_engine(tokens, step, model_params.rope_coords)
+                tokens = self.forecast_engine(tokens, condition, model_params.rope_coords)
                 fe_steps += 1
                 continue
 
             capture_cos = self.compute_cos_sim_to_prev and fe_steps > 0
             prev_tokens = tokens if capture_cos else None
-            tokens = self.forecast_engine(tokens, step, model_params.rope_coords)
+            tokens = self.forecast_engine(tokens, condition, model_params.rope_coords)
             fe_steps += 1
 
             if capture_cos:

From 63d18b66ae1695f5a37b61ac949f04f2240551e4 Mon Sep 17 00:00:00 2001
From: ankitpatnala <ankitpatnala@gmail.com>
Date: Wed, 9 Sep 2026 22:40:57 +0200
Subject: [PATCH 6/6] replaced all streams to data streams i.e filtering
 conditional streams

---
 .../loss_modules/loss_module_physical.py      | 11 ++++++++++-
 src/weathergen/utils/utils.py                 | 10 ++++++++++
 src/weathergen/utils/validation_io.py         | 19 +++++++++++++++----
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/weathergen/train/loss_modules/loss_module_physical.py b/src/weathergen/train/loss_modules/loss_module_physical.py
index 9ca0a6912a..4a7a087ff2 100644
--- a/src/weathergen/train/loss_modules/loss_module_physical.py
+++ b/src/weathergen/train/loss_modules/loss_module_physical.py
@@ -20,6 +20,7 @@
 import weathergen.train.loss_modules.loss_functions as loss_fns
 from weathergen.train.loss_modules.loss_module_base import LossModuleBase, LossValues
 from weathergen.train.utils import TRAIN, VAL, Stage
+from weathergen.utils.utils import is_stream_fe_only
 
 _logger = logging.getLogger(__name__)
 
@@ -281,7 +282,15 @@ def compute_loss(self, preds: dict, targets: dict, metadata) -> LossValues:
         source2target_idxs, output_info, target2source_idxs, target_info = metadata
 
         # TODO: iterate over batch dimension
-        for stream_name, stream_info in self.cf.streams.items():
+        # condition/forcing streams feed the FE directly and produce no physical predictions
+        data_streams = [
+            stream_cfg
+            for stream_cfg in self.cf.streams.values()
+            if not is_stream_fe_only(stream_cfg)
+        ]
+
+        for stream_info in data_streams:
+            stream_name = stream_info["name"]
             # TODO: avoid this
             target_channels = (
                 stream_info.val_target_channels
diff --git a/src/weathergen/utils/utils.py b/src/weathergen/utils/utils.py
index 291ab1521a..298a4b2bf6 100644
--- a/src/weathergen/utils/utils.py
+++ b/src/weathergen/utils/utils.py
@@ -29,6 +29,16 @@ def get_dtype(value: str) -> torch.dtype:
         )
 
 
+def is_stream_fe_only(stream_cfg: dict) -> bool:
+    """
+    Stream that feeds the forecasting engine directly (a `condition` scalar or a per-step
+    spatial `forcing` such as SST) rather than the assimilation path. Excluded from the
+    per-data-stream lists, which must stay aligned with `get_sources_size()` (which omits
+    them).
+    """
+    return stream_cfg.get("type") in ("condition", "forcing")
+
+
 def is_stream_forcing(stream_cfg: dict, stage: Stage | None = None) -> bool:
     """
     Determine if stream is forcing, i.e. does not produce (physical) predictions
diff --git a/src/weathergen/utils/validation_io.py b/src/weathergen/utils/validation_io.py
index 6f989b2bb2..e339603ede 100644
--- a/src/weathergen/utils/validation_io.py
+++ b/src/weathergen/utils/validation_io.py
@@ -19,6 +19,7 @@
 from weathergen.common.io import TimeRange, zarrio_writer
 from weathergen.datasets.data_reader_base import TimeWindowHandler
 from weathergen.model.engines import LatentState
+from weathergen.utils.utils import is_stream_fe_only
 
 _logger = logging.getLogger(__name__)
 
@@ -55,6 +56,11 @@ def write_output(
     forecast_offset = timestep_idxs[0]
     targets_lens = []
 
+    data_streams = {}
+    for stream_name in cf.streams.keys():
+        if not is_stream_fe_only(cf.streams[stream_name]):
+            data_streams[stream_name] = {}
+
     # TODO Maybe stopping at forecast_steps explained #1657
     for t_idx in timestep_idxs:
         preds_all += [[]]
@@ -62,7 +68,7 @@ def write_output(
         targets_coords_all += [[]]
         targets_times_all += [[]]
         targets_lens += [[]]
-        for sname in cf.streams.keys():
+        for sname in data_streams.keys():
             # handle spoof data: do not write since it might corrupt validation (spoofing invisible
             # there)
             if target_aux_out.physical[t_idx][sname]["is_spoof"][0]:
@@ -146,10 +152,15 @@ def write_output(
     }
     _logger.debug(f"Using output streams: {output_streams} from streams: {stream_names}")
 
-    target_channels: list[list[str]] = [list(stream.val_target_channels) for stream in stream_infos]
-    source_channels: list[list[str]] = [list(stream.val_source_channels) for stream in stream_infos]
+    target_channels: list[list[str]] = [
+        list(stream.val_target_channels) for stream in stream_infos if not is_stream_fe_only(stream)
+    ]
+    source_channels: list[list[str]] = [
+        list(stream.val_source_channels) for stream in stream_infos if not is_stream_fe_only(stream)
+    ]
 
-    geoinfo_channels = [[] for _ in stream_infos]  # TODO obtain channels
+    # TODO obtain channels
+    geoinfo_channels = [[] for s in stream_infos if not is_stream_fe_only(s)]
 
     # calculate global sample indices for this batch by offsetting by sample_start
     sample_start = batch_idx * batch_size