Mh/diffusion full pipeline latent rmse curve

src/weathergen/datasets/batch.py

@@ -309,6 +309,23 @@ def __init__(
         self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32)
         self.target2source_matching_idxs = [[] for _ in range(num_target_samples)]
 
+        # Optional isolated samples carrying source-channel network input at the forecast
+        # times, used only by the latent rollout RMSE diagnostic to encode truth latents.
+        # None on every standard batch; the source/target samples above are never affected.
+        self.latent_rmse_source: BatchSamples | None = None
+
+    def init_latent_rmse_source(self, streams, num_samples: int) -> None:
+        """Create the isolated latent-RMSE truth samples (see ``latent_rmse_source``)."""
+        self.latent_rmse_source = BatchSamples(
+            streams, num_samples, self.output_steps, self.output_idxs
+        )
+
+    def add_latent_rmse_source_stream(
+        self, sample_idx: int, stream_name: str, stream_data: StreamData
+    ) -> None:
+        """Add one stream's forecast-time source data to the latent-RMSE truth samples."""
+        self.latent_rmse_source.samples[sample_idx].add_stream_data(stream_name, stream_data)
+
     def pin_memory(self):
         """Pin all tensors in this batch to CPU pinned memory"""
 
@@ -318,6 +335,9 @@ def pin_memory(self):
         # pin target samples
         self.target_samples.pin_memory()
 
+        if self.latent_rmse_source is not None:
+            self.latent_rmse_source.pin_memory()
+
         return self
 
     def to_device(self, device):  # -> ModelBatch
@@ -328,6 +348,9 @@ def to_device(self, device):  # -> ModelBatch
         self.source_samples.to_device(device)
         self.target_samples.to_device(device)
 
+        if self.latent_rmse_source is not None:
+            self.latent_rmse_source.to_device(device)
+
         self.device = device
 
         return self

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -113,6 +113,17 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
         steps = np.array(forecast_cfg["num_steps"], dtype=np.int32).reshape(-1)
         self.list_num_forecast_steps = np.array(steps, dtype=np.int32)
 
+        # Latent rollout RMSE diagnostic (inference only). Attaches an ISOLATED set of
+        # source-channel samples at the forecast times [t, t+1, ..., t+(K-1)] to each batch
+        # (batch.latent_rmse_source), used only to encode truth latents for the diagnostic.
+        # The base index, the standard source/target samples, losses and zarr output are left
+        # exactly as a normal run — this is a read-only side channel.
+        self.latent_rollout_rmse = mode_cfg.get("latent_rollout_rmse", False)
+        if self.latent_rollout_rmse:
+            assert self.output_offset == 0, (
+                f"latent_rollout_rmse requires forecast.offset == 0, got {self.output_offset}"
+            )
+
         # initialise fsm, but can change for future mini_epochs
         self.batch_size = get_batch_size_from_config(mode_cfg)
         self.shuffle = mode_cfg.shuffle
@@ -418,6 +429,43 @@ def _build_stream_data_input(
 
         return stream_data
 
+    def _build_latent_rmse_stream_data(
+        self,
+        stream_info: dict,
+        base_idx: TIndex,
+        num_forecast_steps: int,
+        forecast_input_data: list,
+        forecast_input_tokens: list,
+        mask: torch.Tensor | None,
+    ) -> StreamData:
+        """
+        Build an ISOLATED source-channel StreamData at the forecast times
+        [t, t+1, ..., t+(K-1)] for the latent rollout RMSE diagnostic.
+
+        Like ``_build_stream_data_input`` but walks *forward* over forecast steps: input step k
+        holds the source encoding of the true state at t+k. Used only to encode truth latents;
+        never fed to the model's conditioning, losses or zarr output.
+        """
+        num_output_steps = self._get_output_length(num_forecast_steps)
+        stream_data = StreamData(
+            base_idx, num_output_steps, num_output_steps, self.num_healpix_cells
+        )
+        for step, timestep_idx in enumerate(range(self.output_offset, num_output_steps)):
+            step_forecast_dt = base_idx + (self.time_step * timestep_idx) // self.step_timedelta
+            time_win = self.time_window_handler.window(step_forecast_dt)
+
+            rdata = forecast_input_data[step]
+            token_data = forecast_input_tokens[step]
+            if token_data[0] is None and token_data[1] is None:
+                continue
+
+            (source_cells, source_cells_lens) = self.tokenizer.get_source(
+                stream_info, rdata, token_data, (time_win.start, time_win.end), mask
+            )
+            stream_data.add_source(step, rdata, source_cells_lens, source_cells)
+
+        return stream_data
+
     def _build_stream_data_output(
         self,
         mode: str,
@@ -586,7 +634,27 @@ def _get_data_windows(self, base_idx, num_forecast_steps, num_steps_input_max, s
 
             output_data += [rdata]
 
-        return (input_data, output_data)
+        # source-channel data at the forecast times, for the latent RMSE diagnostic only.
+        # Isolated from input_data/output_data; consumed only via batch.latent_rmse_source.
+        forecast_input_data = []
+        if self.latent_rollout_rmse:
+            for timestep_idx in range(self.output_offset, num_output_steps):
+                step_forecast_dt = (
+                    base_idx + (self.time_step * timestep_idx) // self.step_timedelta
+                )
+                rdata = collect_datasources(stream_ds, step_forecast_dt, "source", self.rng)
+                if rdata.is_empty():
+                    time_win = self.time_window_handler.window(step_forecast_dt)
+                    rdata = spoof(
+                        self.healpix_level,
+                        time_win.start,
+                        stream_ds[0].get_geoinfo_size(),
+                        len(stream_ds[0].mean[stream_ds[0].source_idx]),
+                    )
+                    rdata.is_spoof = True
+                forecast_input_data += [rdata]
+
+        return (input_data, output_data, forecast_input_data)
 
     def _get_source_target_masks(self, training_mode):
         """
@@ -658,6 +726,9 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
             self.output_offset,
             num_output_steps,
         )
+        if self.latent_rollout_rmse:
+            # isolated truth samples: one per target sample, mirroring the target masks
+            batch.init_latent_rmse_source(self.streams, num_target_samples)
 
         # for all streams
         for stream_info, (stream_name, stream_ds) in zip(
@@ -675,14 +746,19 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
             # input_data and output_data is conceptually consecutive but differs
             # in source and target channels; overlap in one window when self.output_offset=0
             i_max = input_steps.max().item()
-            (input_data, output_data) = self._get_data_windows(
+            (input_data, output_data, forecast_input_data) = self._get_data_windows(
                 idx, num_forecast_steps, i_max, stream_ds
             )
 
             # tokenize windows
             # *_tokens = [ (cells_idx, cells_idx_lens), ... ] with length = #time_steps
             input_tokens = self.tokenizer.get_tokens_windows(stream_info, input_data, True)
             output_tokens = self.tokenizer.get_tokens_windows(stream_info, output_data, False)
+            forecast_input_tokens = (
+                self.tokenizer.get_tokens_windows(stream_info, forecast_input_data, True)
+                if self.latent_rollout_rmse
+                else None
+            )
 
             for sidx, source_mask in enumerate(source_masks.masks):
                 # Map each source to its target
@@ -737,10 +813,27 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
                 ]
                 batch.add_target_stream(tidx, student_indices, stream_name, sdata, target_metadata)
 
+                # Isolated latent-RMSE truth: forecast-time source under the same target mask
+                # (the mask the model is trained to predict), added to a separate sample set.
+                if self.latent_rollout_rmse:
+                    truth_sdata = self._build_latent_rmse_stream_data(
+                        stream_info,
+                        idx,
+                        num_forecast_steps,
+                        forecast_input_data,
+                        forecast_input_tokens,
+                        mask=target_mask,
+                    )
+                    batch.add_latent_rmse_source_stream(tidx, stream_name, truth_sdata)
+
         source_in_steps = input_steps.max().item()
         target_in_steps = np.array([tc.get("num_steps_input", 1) for _, tc in target_cfgs.items()])
         target_in_steps = 1 if len(target_in_steps) == 0 else target_in_steps.max().item()
         batch = self._preprocess_model_batch(batch, source_in_steps, target_in_steps)
+        if self.latent_rollout_rmse:
+            batch.latent_rmse_source.tokens_lens = get_tokens_lens(
+                self.streams, batch.latent_rmse_source, num_output_steps
+            )
 
         #add target times in source for diffusion model date/time conditioning
         if self.diffusion_model_conditioning in ["date_time", "date", "time"]:

src/weathergen/model/model.py

@@ -49,6 +49,35 @@
 type StreamName = str
 
 
+def _single_step_source_view(batch, step: int):
+    """
+    Shallow view of a source ``BatchSamples`` exposing only input ``step``.
+
+    The embedding engine concatenates the source tokens of *all* input steps into one tensor
+    before embedding, so encoding a sample that carries many input steps (as the latent
+    rollout RMSE truth samples do, K steps) costs Kx the memory and OOMs. Encoding step by
+    step through this view keeps peak memory at a single step. Only ``source_tokens_cells`` /
+    ``source_tokens_lens`` and ``tokens_lens`` are read per step; nothing is deep-copied, so
+    the views alias the original tensors.
+    """
+
+    view = copy.copy(batch)
+    view.tokens_lens = batch.tokens_lens[step : step + 1]
+    view.samples = []
+    for sample in batch.samples:
+        sample_view = copy.copy(sample)
+        sample_view.streams_data = {}
+        for stream_name, stream_data in sample.streams_data.items():
+            stream_view = copy.copy(stream_data)
+            stream_view.input_steps = 1
+            stream_view.source_tokens_cells = [stream_data.source_tokens_cells[step]]
+            stream_view.source_tokens_lens = [stream_data.source_tokens_lens[step]]
+            sample_view.streams_data[stream_name] = stream_view
+        view.samples += [sample_view]
+
+    return view
+
+
 class ModelOutput:
     """
     Representation of model output
@@ -350,6 +379,10 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord
         # One-shot flag to avoid log spam when warning about an unsupported
         # diffusion-inference + multi-step-rollout combination.
         self._warned_diffusion_multi_step = False
+        # Set by the trainer for the latent rollout RMSE diagnostic: when True, forward()
+        # records the rolled-out latent per step under the "latent_rollout_pred" key. Purely
+        # additive — the standard path is unaffected.
+        self.record_latent_rollout = False
 
     def _create_latent_pred_head(
         self, global_cfg, name, loss_cfg, use_class_token, use_patch_token
@@ -830,6 +863,15 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput:
                     output.add_physical_prediction(
                         step, sname, (torch.cat(list(pred_tuple), dim=0),)
                     )
+                # Latent rollout RMSE diagnostic: record the rolled-out latent under a
+                # dedicated key (never "latent_state") so the latent loss is unaffected. The
+                # trainer pairs this with the encoded truth latent. Off by default.
+                if self.record_latent_rollout:
+                    output.add_latent_prediction(
+                        step,
+                        "latent_rollout_pred",
+                        self.tokens_to_latent_state(None, member_final_tokens),
+                    )
                 # Store per-member conditioning for the next rollout step.
                 # conditioning_tokens holds (N, H, D) during ensemble rollout; inference_forward
                 # calls expand(N, ...) which is a no-op when the dim already matches.
@@ -847,6 +889,20 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput:
 
         return output
 
+    @torch.no_grad()
+    def encode_source_chunked(self, model_params: ModelParams, source_samples) -> torch.Tensor:
+        """
+        Encode a multi-input-step source ``BatchSamples`` one step at a time and stack the
+        latents to ``[B, T, H, D]``. Used by the latent rollout RMSE diagnostic to encode the
+        truth latents at the forecast times without the K-step concatenation OOM (see
+        _single_step_source_view). Read-only; does not touch model state.
+        """
+        step_tokens = []
+        for s in range(source_samples.get_num_steps()):
+            tok_s, _ = self.encoder(model_params, _single_step_source_view(source_samples, s))
+            step_tokens.append(tok_s)
+        return torch.stack(step_tokens, dim=1)
+
     @staticmethod
     def _reindex_output_for_trajectory(output: ModelOutput, n_steps: int) -> ModelOutput:
         """

src/weathergen/train/trainer.py

@@ -47,6 +47,7 @@
     get_target_idxs_from_cfg,
 )
 from weathergen.utils.distributed import is_root
+from weathergen.utils.latent_rmse import LatentRolloutRMSE
 from weathergen.utils.performance import NullThroughputTracker, ThroughputTracker
 from weathergen.utils.train_logger import TrainLogger, prepare_losses_for_logging
 from weathergen.utils.utils import get_dtype
@@ -625,9 +626,21 @@ def validate(self, mini_epoch, mode_cfg, batch_size):
         all_losses: dict[str, list] = {}
         all_stddev: dict[str, list] = {}
 
+        # Latent rollout RMSE diagnostic (isolated side channel): the model records rolled-out
+        # latents under "latent_rollout_pred" while record_latent_rollout is set; the truth
+        # latents are encoded here from the isolated batch.latent_rmse_source. Accumulate only
+        # over the first (random noise-level) pass. Standard preds/losses/zarr are untouched.
+        base_model = getattr(self.model, "module", self.model)
+        latent_rmse = (
+            LatentRolloutRMSE(self.cf, mode_cfg, self.device)
+            if mode_cfg.get("latent_rollout_rmse", False)
+            else None
+        )
+
         for noise_idx, noise_level in enumerate(noise_levels):
             if is_diffusion:
                 self._set_validation_noise_level(noise_level)
+            base_model.record_latent_rollout = latent_rmse is not None and noise_idx == 0
 
             if noise_level is None:
                 loss_suffix = ""
@@ -699,6 +712,31 @@ def validate(self, mini_epoch, mode_cfg, batch_size):
                             metadata=extract_batch_metadata(batch),
                         )
 
+                        # Latent rollout RMSE: encode the isolated truth latents and pair them
+                        # with the recorded predictions by lead step (both valid at t+j). Encode
+                        # under the same autocast as the forward so truth and pred share dtype.
+                        if (
+                            latent_rmse is not None
+                            and noise_idx == 0
+                            and batch.latent_rmse_source is not None
+                        ):
+                            with torch.autocast(
+                                device_type=f"cuda:{cf.local_rank}",
+                                dtype=self.mixed_precision_dtype,
+                                enabled=cf.with_mixed_precision,
+                            ):
+                                truth = base_model.encode_source_chunked(
+                                    self.model_params, batch.latent_rmse_source
+                                )  # [B, K, H, D]
+                            for j in range(truth.shape[1]):
+                                pl = (
+                                    preds.latent[j].get("latent_rollout_pred")
+                                    if j < len(preds.latent)
+                                    else None
+                                )
+                                if pl is not None:
+                                    latent_rmse.add(j, pl.z_pre_norm, truth[:, j])
+
                         # log output
                         if noise_idx == 0:
                             if bidx < num_samples_write:
@@ -753,6 +791,11 @@ def validate(self, mini_epoch, mode_cfg, batch_size):
         if is_diffusion:
             self._set_validation_noise_level(None)
 
+        # latent rollout RMSE: reduce across ranks and plot like the evaluate package's curves
+        base_model.record_latent_rollout = False
+        if latent_rmse is not None:
+            latent_rmse.plot(config.get_path_run(self.cf))
+
         # avoid that there is a systematic bias in the validation subset
         self.dataset_val.advance()