make srdrwfy6 inference work

src/weathergen/datasets/data_reader_anemoi.py

@@ -132,8 +132,8 @@ def __init__(
             self.geoinfo_idx = self.select_geoinfo_channels(ds)
             self.geoinfo_channels = [ds.variables[i] for i in self.geoinfo_idx]
         else:
-            self.geoinfo_channels = stream_info.get("geoinfo_channels")
-            self.geoinfo_idx = [ds.variables.index(ch) for ch in self.geoinfo_channels]
+            self.geoinfo_idx = self.select_geoinfo_channels(ds)
+            self.geoinfo_channels = [ds.variables[i] for i in self.geoinfo_idx]
 
         # set geoinfo normalization statistics
         if len(self.geoinfo_idx) > 0:

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -7,7 +7,6 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-import dataclasses
 import logging
 import pathlib
 from collections.abc import Sequence
@@ -85,12 +84,6 @@ def collect_datasources(stream_datasets: list, idx: int, type: str, rng) -> IORe
     return IOReaderData.combine(rdatas)
 
 
-@dataclasses.dataclass
-class _Stream:
-    info: Config
-    readers: list[DataReaderBase]
-
-
 class MultiStreamDataSampler(torch.utils.data.IterableDataset):
     def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
         super(MultiStreamDataSampler, self).__init__()
@@ -228,12 +221,12 @@ 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) -> dict[StreamName, list[AnyDataReader]]:
         """Load dataset readers for all streams from config."""
-        streams_datasets: dict[StreamName, _Stream] = {}
+        streams_datasets: dict[StreamName, list[AnyDataReader]] = {}
         for stream_name, stream_info in cf.streams.items():
             # list of sources for current stream
-            streams_datasets[stream_name] = _Stream(stream_info, [])
+            streams_datasets[stream_name] = []
 
             kwargs = {
                 "tw_handler": self.time_window_handler,
@@ -281,7 +274,7 @@ def _init_stream_datasets(self, cf) -> dict[StreamName, _Stream]:
                     )
                 ds = dataset(filename=filename, **kwargs)
 
-                streams_datasets[stream_info["name"]] = [ds]
+                streams_datasets[stream_name] += [ds]
 
             stream_info[str(self._stage) + "_source_channels"] = ds.source_channels
             stream_info[str(self._stage) + "_target_channels"] = ds.target_channels
@@ -388,11 +381,11 @@ def get_targets_coords_size(self):
 
     def denormalize_source_channels(self, stream_name, data) -> torch.Tensor:
         # [0]: with multiple ds per stream we use the first one
-        return self.streams_datasets[stream_name].readers[0].denormalize_source_channels(data)
+        return self.streams_datasets[stream_name][0].denormalize_source_channels(data)
 
     def denormalize_target_channels(self, stream_name, data) -> torch.Tensor:
         # [0]: with multiple ds per stream we use the first one
-        return self.streams_datasets[stream_name].readers[0].denormalize_target_channels(data)
+        return self.streams_datasets[stream_name][0].denormalize_target_channels(data)
 
     def _build_stream_data_input(
         self,

src/weathergen/train/trainer.py

@@ -75,12 +75,10 @@ def _expand_targets_to_match_preds(preds, targets_and_auxs: dict) -> None:
     """
     Replicate per-fstep entries in each TargetAuxOutput so its ``physical`` and ``latent``
     lists match the number of forecast steps in ``preds``.
-
     Diffusion inference produces one ``preds`` fstep per ODE denoising step, but the
     physical target is identical across the trajectory. Without this expansion the loss
     calculator (which zips preds and targets with ``strict=True``) raises a length
     mismatch.
-
     The expansion replicates references — no tensor copies are made — and is a no-op when
     the lengths already agree.
     """
@@ -248,7 +246,6 @@ def inference(self, cf, devices, run_id_contd, mini_epoch_contd):
         device_type = torch.accelerator.current_accelerator()
         self.device = torch.device(f"{device_type}:{cf.local_rank}")
         self.ema_model = None
-        [stream.update({"max_num_targets": -1}) for stream in cf.streams]
 
         # create data loader
         # only one needed since we only run the validation code path

src/weathergen/utils/validation_io.py

@@ -57,13 +57,17 @@ def write_output(
     # forecast indices, so synthesize a contiguous run of indices starting at the
     # original first index to cover every entry in model_output / target_aux_out.
     n_pred_steps = len(model_output.physical)
-    if n_pred_steps > len(timestep_idxs):
+    is_denoising_trajectory = n_pred_steps > len(timestep_idxs)
+    if is_denoising_trajectory:
         timestep_idxs = list(range(forecast_offset, forecast_offset + n_pred_steps))
 
     targets_lens = []
 
     # TODO Maybe stopping at forecast_steps explained #1657
-    for t_idx in timestep_idxs:
+    for t_pos, t_idx in enumerate(timestep_idxs):
+        # A denoising trajectory is reindexed from zero in ModelOutput and TargetAuxOutput,
+        # while t_idx is the forecast step written to the output store.
+        output_idx = t_pos if is_denoising_trajectory else t_idx
         preds_all += [[]]
         targets_all += [[]]
         targets_coords_all += [[]]
@@ -79,17 +83,28 @@ def write_output(
             # empty per-stream slots to keep the per-stream array alignment used downstream.
             not_reconstructed = not is_stream_reconstructed(cf.streams[sname])
 
-            if not_reconstructed or target_aux_out.physical[t_idx][sname]["is_spoof"][0]:
-                targets = target_aux_out.physical[t_idx][sname]["target"]
+            target_data = target_aux_out.physical[output_idx].get(sname)
+            if target_data is None:
+                # A stream omitted from the physical loss has no TargetAuxOutput
+                # entry. Keep its output slots empty.
+                n_channels = len(cf.streams[sname].val_target_channels)
+                n_samples = len(batch.get_source_samples())
+                preds_s = [np.zeros((1, 0, n_channels)) for _ in range(n_samples)]
+                targets_s = [np.zeros((0, n_channels)) for _ in range(n_samples)]
+                t_coords_s = [np.zeros((0, 2)) for _ in range(n_samples)]
+                t_times_s = [np.array([]).astype("datetime64[ns]") for _ in range(n_samples)]
+
+            elif not_reconstructed or target_data["is_spoof"][0]:
+                targets = target_data["target"]
                 # for-loop to make sure we have a consistent number of samples
                 preds_s = [np.zeros((1, 0, t.shape[1])) for t in targets]
                 targets_s = [np.zeros((0, t.shape[1])) for t in targets]
                 t_coords_s = [np.zeros((0, 2)) for t in targets]
                 t_times_s = [np.array([]).astype("datetime64[ns]") for t in targets]
 
             else:
-                preds = model_output.get_physical_prediction(t_idx, sname)
-                targets = target_aux_out.physical[t_idx][sname]["target"]
+                preds = model_output.get_physical_prediction(output_idx, sname)
+                targets = target_data["target"]
 
                 preds_s, targets_s, t_coords_s, t_times_s = [], [], [], []
 
@@ -100,11 +115,10 @@ def write_output(
                     preds = [target.clone().unsqueeze(0) for target in targets]
 
                 for i_batch, (pred, target) in enumerate(zip(preds, targets, strict=True)):
-                    target_data = target_aux_out.physical[t_idx][sname]
                     t_coords = target_data["target_coords"][i_batch]
                     t_times = target_data["target_times"][i_batch]
 
-                    idxs_inv = target_aux_out.physical[t_idx][sname]["idxs_inv"][i_batch]
+                    idxs_inv = target_data["idxs_inv"][i_batch]
                     if idxs_inv is not None:
                         pred = pred[:, idxs_inv]
                         target = target[idxs_inv]