Develop hclimrep forcings

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

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

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"))

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

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

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

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