Resync verified files (1/1)

.gitattributes

@@ -1733,3 +1733,5 @@ aux_data/era5_tp_6hr_china.nc filter=lfs diff=lfs merge=lfs -text
 datasets/1.5_nc/2m_temperature/2m_temperature_2015_1.5deg.nc filter=lfs diff=lfs merge=lfs -text
 aux_data/climatology/total_precipitation_6hr_seeps_dry_fraction.nc filter=lfs diff=lfs merge=lfs -text
 aux_data/climatology/total_precipitation_6hr_seeps_threshold.nc filter=lfs diff=lfs merge=lfs -text
+datasets/1.5_nc/2m_temperature/2m_temperature_2012_1.5deg.nc filter=lfs diff=lfs merge=lfs -text
+datasets/1.5_nc/2m_temperature/2m_temperature_2016_1.5deg.nc filter=lfs diff=lfs merge=lfs -text

datasets/1.5_nc/2m_temperature/2m_temperature_2012_1.5deg.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f72d50b681a70895fb3c36b02b65d1c294e4e252f5bea0e4cccc36467c56f9a3
+size 115216361

datasets/1.5_nc/2m_temperature/2m_temperature_2016_1.5deg.nc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92a5a6e02d41079c19dccde429f1071d4e51bd60fa92d1af2670694031551d3c
+size 114847590

downscaling/engine.py

@@ -41,40 +41,119 @@ class TemporalCorrectionHead(nn.Module):
         hidden_surface: int = 256,
         hidden_upper: int = 512,
         dropout: float = 0.0,
+        head_mode: str = "shared3d",
+        num_time_steps: int = 7,
     ):
         super().__init__()
         upper_ch = num_upper_vars * num_levels
-
-        self.surface_net = nn.Sequential(
-            nn.Conv3d(num_surface_vars, hidden_surface, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
-            nn.GELU(),
-            nn.Dropout3d(dropout) if dropout > 0 else nn.Identity(),
-            nn.Conv3d(hidden_surface, num_surface_vars, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
-        )
-        self.upper_net = nn.Sequential(
-            nn.Conv3d(upper_ch, hidden_upper, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
-            nn.GELU(),
-            nn.Dropout3d(dropout) if dropout > 0 else nn.Identity(),
-            nn.Conv3d(hidden_upper, upper_ch, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
-        )
-
-        # Start from interpolation baseline: initial correction ~= 0.
-        nn.init.zeros_(self.surface_net[-1].weight)
-        nn.init.zeros_(self.surface_net[-1].bias)
-        nn.init.zeros_(self.upper_net[-1].weight)
-        nn.init.zeros_(self.upper_net[-1].bias)
-
-    def forward(self, base_surface, base_upper):
-        # base_surface: [B, T, V, H, W]
-        # base_upper:   [B, T, V, C, H, W]
-        x_surface = rearrange(base_surface, "b t v h w -> b v t h w")
-        corr_surface = self.surface_net(x_surface)
-        corr_surface = rearrange(corr_surface, "b v t h w -> b t v h w")
-
-        x_upper = rearrange(base_upper, "b t v c h w -> b (v c) t h w")
-        corr_upper = self.upper_net(x_upper)
-        corr_upper = rearrange(corr_upper, "b (v c) t h w -> b t v c h w", v=base_upper.shape[2], c=base_upper.shape[3])
-
+        self.head_mode = head_mode
+        self.num_time_steps = num_time_steps
+
+        if self.head_mode == "shared3d":
+            # Endpoint-conditioned temporal correction:
+            #   input channels = [endpoint_{t-6}, base_{t+h}, endpoint_t]
+            # for each variable channel group.
+            self.surface_net = nn.Sequential(
+                nn.Conv3d(num_surface_vars * 3, hidden_surface, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
+                nn.GELU(),
+                nn.Dropout3d(dropout) if dropout > 0 else nn.Identity(),
+                nn.Conv3d(hidden_surface, num_surface_vars, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
+            )
+            self.upper_net = nn.Sequential(
+                nn.Conv3d(upper_ch * 3, hidden_upper, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
+                nn.GELU(),
+                nn.Dropout3d(dropout) if dropout > 0 else nn.Identity(),
+                nn.Conv3d(hidden_upper, upper_ch, kernel_size=(3, 1, 1), padding=(1, 0, 0)),
+            )
+            # Start from interpolation baseline: initial correction ~= 0.
+            nn.init.zeros_(self.surface_net[-1].weight)
+            nn.init.zeros_(self.surface_net[-1].bias)
+            nn.init.zeros_(self.upper_net[-1].weight)
+            nn.init.zeros_(self.upper_net[-1].bias)
+        elif self.head_mode == "delta2d":
+            # Per-time-step heads with endpoint-conditioning.
+            # Each head sees [endpoint_{t-6}, base_{t+h}, endpoint_t] context.
+            self.surface_heads = nn.ModuleList([
+                nn.Sequential(
+                    nn.Conv2d(num_surface_vars * 3, hidden_surface, kernel_size=3, padding=1),
+                    nn.GELU(),
+                    nn.Dropout2d(dropout) if dropout > 0 else nn.Identity(),
+                    nn.Conv2d(hidden_surface, num_surface_vars, kernel_size=3, padding=1),
+                )
+                for _ in range(num_time_steps)
+            ])
+            self.upper_heads = nn.ModuleList([
+                nn.Sequential(
+                    nn.Conv2d(upper_ch * 3, hidden_upper, kernel_size=3, padding=1),
+                    nn.GELU(),
+                    nn.Dropout2d(dropout) if dropout > 0 else nn.Identity(),
+                    nn.Conv2d(hidden_upper, upper_ch, kernel_size=3, padding=1),
+                )
+                for _ in range(num_time_steps)
+            ])
+            for head in self.surface_heads:
+                nn.init.zeros_(head[-1].weight)
+                nn.init.zeros_(head[-1].bias)
+            for head in self.upper_heads:
+                nn.init.zeros_(head[-1].weight)
+                nn.init.zeros_(head[-1].bias)
+        else:
+            raise ValueError(f"Unsupported temporal head_mode: {self.head_mode}")
+
+    def forward(self, cond_surface_endpoints, cond_upper_endpoints, base_surface, base_upper):
+        # cond_surface_endpoints: [B, 2, V, H, W] (t-6, t) on the current grid
+        # cond_upper_endpoints:   [B, 2, V, C, H, W]
+        # base_surface:           [B, T, V, H, W]  (U_t baseline for target horizons)
+        # base_upper:             [B, T, V, C, H, W]
+        if cond_surface_endpoints is None or cond_upper_endpoints is None:
+            raise ValueError("TemporalCorrectionHead requires endpoint-conditioned inputs.")
+        if cond_surface_endpoints.shape[1] < 2 or cond_upper_endpoints.shape[1] < 2:
+            raise ValueError("TemporalCorrectionHead expects two endpoint frames in cond inputs.")
+
+        if self.head_mode == "shared3d":
+            b, t, v, h, w = base_surface.shape
+            s0 = cond_surface_endpoints[:, 0].unsqueeze(1).expand(-1, t, -1, -1, -1)
+            s1 = cond_surface_endpoints[:, -1].unsqueeze(1).expand(-1, t, -1, -1, -1)
+            x_surface = torch.cat((s0, base_surface, s1), dim=2)
+            x_surface = rearrange(x_surface, "b t v h w -> b v t h w")
+            corr_surface = self.surface_net(x_surface)
+            corr_surface = rearrange(corr_surface, "b v t h w -> b t v h w")
+
+            _, _, vu, c, hu, wu = base_upper.shape
+            u0 = cond_upper_endpoints[:, 0].unsqueeze(1).expand(-1, t, -1, -1, -1, -1)
+            u1 = cond_upper_endpoints[:, -1].unsqueeze(1).expand(-1, t, -1, -1, -1, -1)
+            x_upper = torch.cat((u0, base_upper, u1), dim=2)
+            x_upper = rearrange(x_upper, "b t v c h w -> b (v c) t h w")
+            corr_upper = self.upper_net(x_upper)
+            corr_upper = rearrange(
+                corr_upper, "b (v c) t h w -> b t v c h w", v=base_upper.shape[2], c=base_upper.shape[3]
+            )
+            return corr_surface, corr_upper
+
+        # delta2d mode
+        b, t, v, h, w = base_surface.shape
+        _, _, vu, c, hu, wu = base_upper.shape
+        s0 = cond_surface_endpoints[:, 0]
+        s1 = cond_surface_endpoints[:, -1]
+        u0 = cond_upper_endpoints[:, 0].reshape(b, vu * c, hu, wu)
+        u1 = cond_upper_endpoints[:, -1].reshape(b, vu * c, hu, wu)
+
+        corr_surface_steps = []
+        corr_upper_steps = []
+        for ti in range(t):
+            hi = min(ti, self.num_time_steps - 1)
+            s_cur = base_surface[:, ti]
+            s_in = torch.cat((s0, s_cur, s1), dim=1)
+            s_corr = self.surface_heads[hi](s_in)
+            corr_surface_steps.append(s_corr)
+
+            u_cur = base_upper[:, ti].reshape(b, vu * c, hu, wu)
+            u_in = torch.cat((u0, u_cur, u1), dim=1)
+            u_corr = self.upper_heads[hi](u_in).reshape(b, vu, c, hu, wu)
+            corr_upper_steps.append(u_corr)
+
+        corr_surface = torch.stack(corr_surface_steps, dim=1)
+        corr_upper = torch.stack(corr_upper_steps, dim=1)
         return corr_surface, corr_upper
 
 
@@ -147,23 +226,41 @@ def _default_spatial_interp(surface, upper, target_hw):
     return surface_up, upper_up
 
 
-def _default_temporal_interp(surface_endpoints, upper_endpoints, out_steps):
-    """Linear interpolation along time between the two endpoint frames."""
+def _default_temporal_interp(surface_endpoints, upper_endpoints, out_steps, offsets_hours=None):
+    """Linear temporal baseline from [t-6, t] to target horizons.
+
+    By default, this generates a prediction-style baseline for [t+1, ..., t+out_steps]
+    using linear extrapolation from the two 6h endpoints.
+    """
     if surface_endpoints.shape[1] < 2 or upper_endpoints.shape[1] < 2:
         raise ValueError("Temporal interpolation needs at least two endpoint frames.")
 
-    s0 = surface_endpoints[:, 0].unsqueeze(1)
-    s1 = surface_endpoints[:, -1].unsqueeze(1)
-    u0 = upper_endpoints[:, 0].unsqueeze(1)
-    u1 = upper_endpoints[:, -1].unsqueeze(1)
+    if offsets_hours is None:
+        offsets = torch.arange(
+            1, out_steps + 1, device=surface_endpoints.device, dtype=surface_endpoints.dtype
+        )
+    else:
+        offsets = offsets_hours
+        if not torch.is_tensor(offsets):
+            offsets = torch.as_tensor(offsets)
+        if offsets.dim() == 2:
+            offsets = offsets[0]
+        offsets = offsets.to(device=surface_endpoints.device, dtype=surface_endpoints.dtype).reshape(-1)
+        if offsets.numel() != out_steps:
+            raise ValueError(
+                f"offset length ({offsets.numel()}) must equal out_steps ({out_steps})."
+            )
+
+    x_prev_s = surface_endpoints[:, 0].unsqueeze(1)  # at t-6
+    x_curr_s = surface_endpoints[:, -1].unsqueeze(1)  # at t
+    x_prev_u = upper_endpoints[:, 0].unsqueeze(1)
+    x_curr_u = upper_endpoints[:, -1].unsqueeze(1)
 
-    alpha_s = torch.linspace(
-        0.0, 1.0, out_steps, device=surface_endpoints.device, dtype=surface_endpoints.dtype
-    ).view(1, out_steps, 1, 1, 1)
-    alpha_u = alpha_s.view(1, out_steps, 1, 1, 1, 1)
+    beta_s = (offsets / 6.0).view(1, out_steps, 1, 1, 1)
+    beta_u = beta_s.view(1, out_steps, 1, 1, 1, 1)
 
-    surface_interp = (1.0 - alpha_s) * s0 + alpha_s * s1
-    upper_interp = (1.0 - alpha_u) * u0 + alpha_u * u1
+    surface_interp = x_curr_s + beta_s * (x_curr_s - x_prev_s)
+    upper_interp = x_curr_u + beta_u * (x_curr_u - x_prev_u)
     return surface_interp, upper_interp
 
 
@@ -214,26 +311,33 @@ def forward_paths(
     target_hw_6h = batch["y_hr6h_surface"].shape[-2:]
     target_hw_1h = batch["y_hr1h_surface"].shape[-2:]
     out_steps_1h = batch["y_hr1h_surface"].shape[1]
+    target_offsets_hours = batch.get("target_offsets_hours", None)
 
     # Spatial path at 6h: S(x) = U_s(x) + C_s(x)
+    # Use pretrained Aurora 6h prior as base signal at current endpoint (centered prior),
+    # not a post-hoc overwrite.
     s_base_surface, s_base_upper = spatial_interp_fn(x_lr6h_surface, x_lr6h_upper, target_hw_6h)
+    s_base_prior_surface = s_base_surface
+    s_base_prior_upper = s_base_upper
+    if backbone_s_last_surface is not None:
+        s_base_prior_surface = s_base_prior_surface.clone()
+        s_base_prior_surface[:, -1] = backbone_s_last_surface
+    if backbone_s_last_upper is not None:
+        s_base_prior_upper = s_base_prior_upper.clone()
+        s_base_prior_upper[:, -1] = backbone_s_last_upper
+
     pred_s_hr6h_surface, pred_s_hr6h_upper = _apply_correction(
         spatial_correction_fn,
         x_lr6h_surface,
         x_lr6h_upper,
-        s_base_surface,
-        s_base_upper,
+        s_base_prior_surface,
+        s_base_prior_upper,
     )
-    # Optional Aurora-backbone replacement for S-path last endpoint (enables LoRA gradient path).
-    if backbone_s_last_surface is not None:
-        pred_s_hr6h_surface = pred_s_hr6h_surface.clone()
-        pred_s_hr6h_surface[:, -1] = backbone_s_last_surface
-    if backbone_s_last_upper is not None:
-        pred_s_hr6h_upper = pred_s_hr6h_upper.clone()
-        pred_s_hr6h_upper[:, -1] = backbone_s_last_upper
 
     # Temporal path at LR: T(x) = U_t(x) + C_t(x)
-    t_base_surface, t_base_upper = temporal_interp_fn(x_lr6h_surface, x_lr6h_upper, out_steps_1h)
+    t_base_surface, t_base_upper = temporal_interp_fn(
+        x_lr6h_surface, x_lr6h_upper, out_steps_1h, offsets_hours=target_offsets_hours
+    )
     pred_t_lr1h_surface, pred_t_lr1h_upper = _apply_correction(
         temporal_correction_fn,
         x_lr6h_surface,
@@ -253,7 +357,9 @@ def forward_paths(
     )
 
     # ST path: x -> S -> T
-    st_base_surface, st_base_upper = temporal_interp_fn(pred_s_hr6h_surface, pred_s_hr6h_upper, out_steps_1h)
+    st_base_surface, st_base_upper = temporal_interp_fn(
+        pred_s_hr6h_surface, pred_s_hr6h_upper, out_steps_1h, offsets_hours=target_offsets_hours
+    )
     pred_st_hr1h_surface, pred_st_hr1h_upper = _apply_correction(
         temporal_correction_fn,
         pred_s_hr6h_surface,
@@ -265,6 +371,8 @@ def forward_paths(
     return {
         "pred_s_hr6h_surface": pred_s_hr6h_surface,
         "pred_s_hr6h_upper": pred_s_hr6h_upper,
+        "s_base_prior_surface": s_base_prior_surface,
+        "s_base_prior_upper": s_base_prior_upper,
         "pred_t_lr1h_surface": pred_t_lr1h_surface,
         "pred_t_lr1h_upper": pred_t_lr1h_upper,
         "pred_ts_hr1h_surface": pred_ts_hr1h_surface,
@@ -317,6 +425,8 @@ def _downscale_style_term_loss(
     surf_vars=None,
     upper_vars=None,
     level=None,
+    surface_var_weights=None,
+    upper_var_weights=None,
 ):
     """
     Match train_one_epoch_downscale style:
@@ -329,8 +439,33 @@ def _downscale_style_term_loss(
     pred_upper_n = _normalise_upper_bt_vchw(pred_upper, upper_vars, level)
     target_upper_n = _normalise_upper_bt_vchw(target_upper, upper_vars, level)
 
-    loss_surface = F.mse_loss(pred_surface_n, target_surface_n)
-    loss_upper = F.mse_loss(pred_upper_n, target_upper_n)
+    if surface_var_weights is None:
+        loss_surface = F.mse_loss(pred_surface_n, target_surface_n)
+    else:
+        ws = torch.as_tensor(
+            surface_var_weights,
+            device=pred_surface_n.device,
+            dtype=pred_surface_n.dtype,
+        ).view(1, 1, -1, 1, 1)
+        se_s = (pred_surface_n - target_surface_n) ** 2
+        loss_surface = (se_s * ws).sum() / (ws.sum() * se_s.shape[0] * se_s.shape[1] * se_s.shape[3] * se_s.shape[4])
+
+    if upper_var_weights is None:
+        loss_upper = F.mse_loss(pred_upper_n, target_upper_n)
+    else:
+        wu = torch.as_tensor(
+            upper_var_weights,
+            device=pred_upper_n.device,
+            dtype=pred_upper_n.dtype,
+        )
+        if wu.dim() == 1:
+            wu = wu.view(1, 1, -1, 1, 1, 1)
+        else:
+            wu = wu.view(1, 1, wu.shape[0], wu.shape[1], 1, 1)
+        se_u = (pred_upper_n - target_upper_n) ** 2
+        loss_upper = (se_u * wu).sum() / (
+            wu.sum() * se_u.shape[0] * se_u.shape[1] * se_u.shape[4] * se_u.shape[5]
+        )
     return 0.25 * loss_surface + loss_upper
 
 
@@ -341,10 +476,13 @@ def loss_fn(
     lambda_t=1.0,
     lambda_st=1.0,
     lambda_comm=1.0,
+    lambda_backbone=0.0,
     surface_var_weights=None,
     upper_var_weights=None,
     upper_level_weights=None,
     interior_only_1h=False,
+    lt_supervision_source="lr_path",
+    lt_hybrid_alpha=0.5,
     surf_vars=None,
     upper_vars=None,
     level=None,
@@ -356,8 +494,9 @@ def loss_fn(
             return x[:, 1:-1]
         return x
 
-    # Keep signature compatibility; old downscale style does not use these explicit weights.
-    _ = surface_var_weights, upper_var_weights, upper_level_weights
+    # Backward compatibility: if explicit upper-level weights are provided, override upper_var_weights.
+    if upper_level_weights is not None:
+        upper_var_weights = upper_level_weights
 
     # L_S: S(x) supervised by HR-6h
     l_s = _downscale_style_term_loss(
@@ -368,10 +507,14 @@ def loss_fn(
         surf_vars=surf_vars,
         upper_vars=upper_vars,
         level=level,
+        surface_var_weights=surface_var_weights,
+        upper_var_weights=upper_var_weights,
     )
 
-    # L_T: T(x) supervised by LR-1h
-    l_t = _downscale_style_term_loss(
+    # L_T: temporal module supervision on LR 1h targets.
+    # "lr_path" is the strict factorized definition.
+    # "st_downscaled"/"hybrid" are optional variants.
+    l_t_lr = _downscale_style_term_loss(
         _trim_time(preds["pred_t_lr1h_surface"]),
         _trim_time(preds["pred_t_lr1h_upper"]),
         _trim_time(batch["y_lr1h_surface"]),
@@ -379,7 +522,33 @@ def loss_fn(
         surf_vars=surf_vars,
         upper_vars=upper_vars,
         level=level,
+        surface_var_weights=surface_var_weights,
+        upper_var_weights=upper_var_weights,
     )
+    has_st_lr = ("pred_st_lr1h_surface" in preds) and ("pred_st_lr1h_upper" in preds)
+    if has_st_lr:
+        l_t_st = _downscale_style_term_loss(
+            _trim_time(preds["pred_st_lr1h_surface"]),
+            _trim_time(preds["pred_st_lr1h_upper"]),
+            _trim_time(batch["y_lr1h_surface"]),
+            _trim_time(batch["y_lr1h_upper"]),
+            surf_vars=surf_vars,
+            upper_vars=upper_vars,
+            level=level,
+            surface_var_weights=surface_var_weights,
+            upper_var_weights=upper_var_weights,
+        )
+    else:
+        l_t_st = l_t_lr
+
+    if lt_supervision_source == "st_downscaled":
+        l_t = l_t_st
+    elif lt_supervision_source == "hybrid":
+        a = float(lt_hybrid_alpha)
+        a = max(0.0, min(1.0, a))
+        l_t = a * l_t_lr + (1.0 - a) * l_t_st
+    else:
+        l_t = l_t_lr
 
     # L_ST: ST path supervised by HR-1h
     l_st = _downscale_style_term_loss(
@@ -390,9 +559,11 @@ def loss_fn(
         surf_vars=surf_vars,
         upper_vars=upper_vars,
         level=level,
+        surface_var_weights=surface_var_weights,
+        upper_var_weights=upper_var_weights,
     )
 
-    # L_comm: TS and ST consistency
+    # L_comm: TS and ST consistency (commutativity regularizer)
     l_comm = _downscale_style_term_loss(
         _trim_time(preds["pred_ts_hr1h_surface"]),
         _trim_time(preds["pred_ts_hr1h_upper"]),
@@ -401,15 +572,47 @@ def loss_fn(
         surf_vars=surf_vars,
         upper_vars=upper_vars,
         level=level,
+        surface_var_weights=surface_var_weights,
+        upper_var_weights=upper_var_weights,
     )
 
-    total = lambda_s * l_s + lambda_t * l_t + lambda_st * l_st + lambda_comm * l_comm
+    l_backbone = l_comm.new_zeros(())
+    if (
+        lambda_backbone > 0
+        and ("backbone_last_surface" in preds)
+        and ("backbone_last_upper" in preds)
+    ):
+        # Optional prior-consistency regularizer:
+        # keep S current-endpoint prediction close to pretrained backbone prior.
+        bb_s = preds["backbone_last_surface"].detach().unsqueeze(1)
+        bb_u = preds["backbone_last_upper"].detach().unsqueeze(1)
+        l_backbone_s = _downscale_style_term_loss(
+            preds["pred_s_hr6h_surface"][:, -1:].contiguous(),
+            preds["pred_s_hr6h_upper"][:, -1:].contiguous(),
+            bb_s,
+            bb_u,
+            surf_vars=surf_vars,
+            upper_vars=upper_vars,
+            level=level,
+            surface_var_weights=surface_var_weights,
+            upper_var_weights=upper_var_weights,
+        )
+        l_backbone = l_backbone_s
+
+    total = (
+        lambda_s * l_s
+        + lambda_t * l_t
+        + lambda_st * l_st
+        + lambda_comm * l_comm
+        + lambda_backbone * l_backbone
+    )
     loss_dict = {
         "loss_total": total,
         "loss_s": l_s,
         "loss_t": l_t,
         "loss_st": l_st,
         "loss_comm": l_comm,
+        "loss_backbone": l_backbone,
     }
     return total, loss_dict
 
@@ -687,6 +890,21 @@ def _trim_to_hw(x, h, w):
     return x[..., :h, :w]
 
 
+def _sanitize_tensor(x, clip_value=1e6):
+    x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
+    if clip_value is not None and clip_value > 0:
+        x = torch.clamp(x, min=-clip_value, max=clip_value)
+    return x
+
+
+def _tensor_health(x):
+    finite_mask = torch.isfinite(x)
+    finite_ratio = float(finite_mask.float().mean().item())
+    x_safe = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
+    max_abs = float(x_safe.abs().max().item())
+    return finite_ratio, max_abs
+
+
 def _resize_surface_bt_vhw(x, out_hw):
     b, t, v, h, w = x.shape
     th, tw = out_hw
@@ -709,12 +927,123 @@ def _resize_upper_bt_vchw(x, out_hw):
     ).reshape(b, t, v, c, th, tw)
 
 
-def _prepare_factorized_batch(batch, hr_hw, lr_hw, device):
+def _cell_edges_from_centers(centers_1d: torch.Tensor) -> torch.Tensor:
+    centers_1d = centers_1d.float()
+    if centers_1d.numel() < 2:
+        c = centers_1d[0]
+        return torch.stack((c - 0.5, c + 0.5), dim=0)
+    mids = 0.5 * (centers_1d[:-1] + centers_1d[1:])
+    first = centers_1d[0] - 0.5 * (centers_1d[1] - centers_1d[0])
+    last = centers_1d[-1] + 0.5 * (centers_1d[-1] - centers_1d[-2])
+    return torch.cat((first[None], mids, last[None]), dim=0)
+
+
+def _interval_overlap_matrix(src_edges: torch.Tensor, tgt_edges: torch.Tensor) -> torch.Tensor:
+    # Returns [N_tgt, N_src] overlap lengths for 1D intervals.
+    src_lo = src_edges[:-1].unsqueeze(0)  # [1, Ns]
+    src_hi = src_edges[1:].unsqueeze(0)   # [1, Ns]
+    tgt_lo = tgt_edges[:-1].unsqueeze(1)  # [Nt, 1]
+    tgt_hi = tgt_edges[1:].unsqueeze(1)   # [Nt, 1]
+    return torch.clamp(torch.minimum(src_hi, tgt_hi) - torch.maximum(src_lo, tgt_lo), min=0.0)
+
+
+def _maybe_flip_last_dim(x: torch.Tensor, dim: int, need_flip: bool) -> torch.Tensor:
+    if need_flip:
+        return torch.flip(x, dims=(dim,))
+    return x
+
+
+def _conservative_remap_2d(
+    x: torch.Tensor,
+    lat_src: torch.Tensor,
+    lon_src: torch.Tensor,
+    lat_tgt: torch.Tensor,
+    lon_tgt: torch.Tensor,
+    eps: float = 1e-12,
+) -> torch.Tensor:
+    """
+    Conservative remap from source grid to target grid for tensors with trailing [..., H, W].
+    Uses area-overlap weights in (sin(lat), lon) coordinates (first-order conservative).
+    """
+    # Keep target orientation to restore at end.
+    src_lat_desc = bool((lat_src[0] > lat_src[-1]).item())
+    src_lon_desc = bool((lon_src[0] > lon_src[-1]).item())
+    tgt_lat_desc = bool((lat_tgt[0] > lat_tgt[-1]).item())
+    tgt_lon_desc = bool((lon_tgt[0] > lon_tgt[-1]).item())
+
+    x_work = _maybe_flip_last_dim(x, -2, src_lat_desc)
+    x_work = _maybe_flip_last_dim(x_work, -1, src_lon_desc)
+    lat_src_inc = _maybe_flip_last_dim(lat_src, 0, src_lat_desc)
+    lon_src_inc = _maybe_flip_last_dim(lon_src, 0, src_lon_desc)
+    lat_tgt_inc = _maybe_flip_last_dim(lat_tgt, 0, tgt_lat_desc)
+    lon_tgt_inc = _maybe_flip_last_dim(lon_tgt, 0, tgt_lon_desc)
+
+    lat_src_edges = _cell_edges_from_centers(lat_src_inc).clamp(-90.0, 90.0)
+    lat_tgt_edges = _cell_edges_from_centers(lat_tgt_inc).clamp(-90.0, 90.0)
+    lon_src_edges = _cell_edges_from_centers(lon_src_inc)
+    lon_tgt_edges = _cell_edges_from_centers(lon_tgt_inc)
+
+    # Area metric on a sphere is proportional to d(lon) * d(sin(lat)).
+    lat_src_m = torch.sin(torch.deg2rad(lat_src_edges))
+    lat_tgt_m = torch.sin(torch.deg2rad(lat_tgt_edges))
+    lat_overlap = _interval_overlap_matrix(lat_src_m, lat_tgt_m)  # [Ht, Hs]
+    lon_overlap = _interval_overlap_matrix(lon_src_edges, lon_tgt_edges)  # [Wt, Ws]
+
+    tmp = torch.einsum("...hw,ih->...iw", x_work, lat_overlap)
+    num = torch.einsum("...iw,jw->...ij", tmp, lon_overlap)
+
+    lat_tgt_width = (lat_tgt_m[1:] - lat_tgt_m[:-1]).clamp_min(eps)  # [Ht]
+    lon_tgt_width = (lon_tgt_edges[1:] - lon_tgt_edges[:-1]).clamp_min(eps)  # [Wt]
+    den = lat_tgt_width[:, None] * lon_tgt_width[None, :]
+    out = num / den
+
+    out = _maybe_flip_last_dim(out, -2, tgt_lat_desc)
+    out = _maybe_flip_last_dim(out, -1, tgt_lon_desc)
+    return out
+
+
+def _conservative_downscale_surface_bt_vhw(x, lat_hr, lon_hr, lat_lr, lon_lr):
+    # x: [B, T, V, H, W]
+    b, t, v, h, w = x.shape
+    y = _conservative_remap_2d(
+        x.reshape(b * t * v, h, w),
+        lat_src=lat_hr,
+        lon_src=lon_hr,
+        lat_tgt=lat_lr,
+        lon_tgt=lon_lr,
+    )
+    return y.reshape(b, t, v, lat_lr.numel(), lon_lr.numel())
+
+
+def _conservative_downscale_upper_bt_vchw(x, lat_hr, lon_hr, lat_lr, lon_lr):
+    # x: [B, T, V, C, H, W]
+    b, t, v, c, h, w = x.shape
+    y = _conservative_remap_2d(
+        x.reshape(b * t * v * c, h, w),
+        lat_src=lat_hr,
+        lon_src=lon_hr,
+        lat_tgt=lat_lr,
+        lon_tgt=lon_lr,
+    )
+    return y.reshape(b, t, v, c, lat_lr.numel(), lon_lr.numel())
+
+
+def _prepare_factorized_batch(
+    batch,
+    hr_hw,
+    lr_hw,
+    device,
+    lat_hr=None,
+    lon_hr=None,
+    lat_lr=None,
+    lon_lr=None,
+):
     h_hr, w_hr = hr_hw
     h_lr, w_lr = lr_hw
 
     def _to_device(x, h, w):
-        return _trim_to_hw(x.float(), h, w).to(device, non_blocking=True)
+        x = _trim_to_hw(x.float(), h, w).to(device, non_blocking=True)
+        return _sanitize_tensor(x)
 
     out = {}
     out["x_lr6h_surface"] = _to_device(batch["x_lr6h_surface"], h_lr, w_lr)
@@ -741,8 +1070,30 @@ def _prepare_factorized_batch(batch, hr_hw, lr_hw, device):
         out["y_lr1h_surface"] = _to_device(batch["y_lr1h_surface"], h_lr, w_lr)
         out["y_lr1h_upper"] = _to_device(batch["y_lr1h_upper"], h_lr, w_lr)
     else:
-        out["y_lr1h_surface"] = _resize_surface_bt_vhw(out["y_hr1h_surface"], (h_lr, w_lr))
-        out["y_lr1h_upper"] = _resize_upper_bt_vchw(out["y_hr1h_upper"], (h_lr, w_lr))
+        use_conservative = (
+            (lat_hr is not None) and (lon_hr is not None) and
+            (lat_lr is not None) and (lon_lr is not None)
+        )
+        if use_conservative:
+            lat_hr_d = lat_hr.to(device=device, dtype=out["y_hr1h_surface"].dtype)
+            lon_hr_d = lon_hr.to(device=device, dtype=out["y_hr1h_surface"].dtype)
+            lat_lr_d = lat_lr.to(device=device, dtype=out["y_hr1h_surface"].dtype)
+            lon_lr_d = lon_lr.to(device=device, dtype=out["y_hr1h_surface"].dtype)
+            out["y_lr1h_surface"] = _conservative_downscale_surface_bt_vhw(
+                out["y_hr1h_surface"], lat_hr_d, lon_hr_d, lat_lr_d, lon_lr_d
+            )
+            out["y_lr1h_upper"] = _conservative_downscale_upper_bt_vchw(
+                out["y_hr1h_upper"], lat_hr_d, lon_hr_d, lat_lr_d, lon_lr_d
+            )
+        else:
+            out["y_lr1h_surface"] = _resize_surface_bt_vhw(out["y_hr1h_surface"], (h_lr, w_lr))
+            out["y_lr1h_upper"] = _resize_upper_bt_vchw(out["y_hr1h_upper"], (h_lr, w_lr))
+
+    if "target_offsets_hours" in batch:
+        offsets = batch["target_offsets_hours"]
+        if not torch.is_tensor(offsets):
+            offsets = torch.as_tensor(offsets)
+        out["target_offsets_hours"] = offsets.to(device=device, non_blocking=True).long()
 
     out["time_unix"] = batch["time_unix"]
     return out
@@ -851,6 +1202,8 @@ def _build_temporal_correction_fn(
     upper_vars=None,
     level=None,
     correction_scale=1.0,
+    fix_endpoints=True,
+    target_offsets_hours=None,
 ):
     if temporal_corrector is None:
         def _temporal_correction_zeros(input_surface, input_upper, base_surface, base_upper):
@@ -859,12 +1212,35 @@ def _build_temporal_correction_fn(
         return _temporal_correction_zeros
 
     def _temporal_correction(input_surface, input_upper, base_surface, base_upper):
+        cond_surface_n = _normalise_surface_bt_vhw(input_surface, surf_vars)
+        cond_upper_n = _normalise_upper_bt_vchw(input_upper, upper_vars, level)
         base_surface_n = _normalise_surface_bt_vhw(base_surface, surf_vars)
         base_upper_n = _normalise_upper_bt_vchw(base_upper, upper_vars, level)
-        corr_surface_n, corr_upper_n = temporal_corrector(base_surface_n, base_upper_n)
+        corr_surface_n, corr_upper_n = temporal_corrector(
+            cond_surface_n, cond_upper_n, base_surface_n, base_upper_n
+        )
 
         corr_surface_n = torch.nan_to_num(corr_surface_n, nan=0.0, posinf=0.0, neginf=0.0)
         corr_upper_n = torch.nan_to_num(corr_upper_n, nan=0.0, posinf=0.0, neginf=0.0)
+        if fix_endpoints and corr_surface_n.shape[1] >= 1 and corr_upper_n.shape[1] >= 1:
+            corr_surface_n = corr_surface_n.clone()
+            corr_upper_n = corr_upper_n.clone()
+            fix_indices = []
+            if target_offsets_hours is None:
+                # Backward-compatibility fallback (legacy setting with endpoint-including sequence).
+                fix_indices = [0, corr_surface_n.shape[1] - 1] if corr_surface_n.shape[1] >= 2 else [0]
+            else:
+                offs = target_offsets_hours
+                if not torch.is_tensor(offs):
+                    offs = torch.as_tensor(offs)
+                if offs.dim() == 2:
+                    offs = offs[0]
+                offs = offs.to(device=corr_surface_n.device).reshape(-1).long()
+                if offs.numel() == corr_surface_n.shape[1]:
+                    fix_indices = torch.nonzero(offs == 0, as_tuple=False).view(-1).tolist()
+            for idx in fix_indices:
+                corr_surface_n[:, idx] = 0.0
+                corr_upper_n[:, idx] = 0.0
 
         s_scale = _surface_scale_tensor(surf_vars, base_surface)
         u_scale = _upper_scale_tensor(upper_vars, level, base_upper)
@@ -902,14 +1278,18 @@ def train_one_epoch_factorized(
     lambda_t=1.0,
     lambda_st=1.0,
     lambda_comm=1.0,
+    lambda_backbone=0.0,
     surface_var_weights=None,
     upper_var_weights=None,
     upper_level_weights=None,
     interior_only_1h=False,
+    lt_supervision_source="lr_path",
+    lt_hybrid_alpha=0.5,
     spatial_corrector=None,
     temporal_corrector=None,
     spatial_correction_scale=1.0,
     temporal_correction_scale=1.0,
+    temporal_fix_endpoints=True,
     enable_backbone_lora=False,
     accum_backward_mode="micro",
     use_ours=False,
@@ -937,6 +1317,9 @@ def train_one_epoch_factorized(
     backward_calls = 0
     pending_loss = None
 
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+
     for data_iter_step, batch in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
         step = data_iter_step // update_freq
         if step >= num_training_steps_per_epoch:
@@ -956,6 +1339,10 @@ def train_one_epoch_factorized(
             hr_hw=hr_hw,
             lr_hw=lr_hw,
             device=device,
+            lat_hr=lat_hr,
+            lon_hr=lon_hr,
+            lat_lr=lat_lr,
+            lon_lr=lon_lr,
         )
 
         spatial_correction_fn = _build_spatial_correction_fn(
@@ -971,6 +1358,8 @@ def train_one_epoch_factorized(
             upper_vars=upper_vars,
             level=level,
             correction_scale=temporal_correction_scale,
+            fix_endpoints=temporal_fix_endpoints,
+            target_offsets_hours=batch_f.get("target_offsets_hours", None),
         )
         backbone_last_surface = None
         backbone_last_upper = None
@@ -993,6 +1382,20 @@ def train_one_epoch_factorized(
             backbone_s_last_surface=backbone_last_surface,
             backbone_s_last_upper=backbone_last_upper,
         )
+        if lt_supervision_source in {"st_downscaled", "hybrid"}:
+            lat_hr_d = lat_hr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+            lon_hr_d = lon_hr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+            lat_lr_d = lat_lr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+            lon_lr_d = lon_lr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+            preds["pred_st_lr1h_surface"] = _conservative_downscale_surface_bt_vhw(
+                preds["pred_st_hr1h_surface"], lat_hr_d, lon_hr_d, lat_lr_d, lon_lr_d
+            )
+            preds["pred_st_lr1h_upper"] = _conservative_downscale_upper_bt_vchw(
+                preds["pred_st_hr1h_upper"], lat_hr_d, lon_hr_d, lat_lr_d, lon_lr_d
+            )
+        if backbone_last_surface is not None and backbone_last_upper is not None:
+            preds["backbone_last_surface"] = backbone_last_surface
+            preds["backbone_last_upper"] = backbone_last_upper
         loss, loss_dict = loss_fn(
             preds=preds,
             batch=batch_f,
@@ -1000,18 +1403,29 @@ def train_one_epoch_factorized(
             lambda_t=lambda_t,
             lambda_st=lambda_st,
             lambda_comm=lambda_comm,
+            lambda_backbone=lambda_backbone,
             surface_var_weights=surface_var_weights,
             upper_var_weights=upper_var_weights,
             upper_level_weights=upper_level_weights,
             interior_only_1h=interior_only_1h,
+            lt_supervision_source=lt_supervision_source,
+            lt_hybrid_alpha=lt_hybrid_alpha,
             surf_vars=surf_vars,
             upper_vars=upper_vars,
             level=level,
         )
-
         loss_value = loss.item()
         if not np.isfinite(loss_value):
             print(f"[factorized] non-finite loss at iter {data_iter_step}; skip update.")
+            try:
+                for k in ("x_lr6h_surface", "x_lr6h_upper", "y_hr1h_surface", "y_hr1h_upper"):
+                    fr, ma = _tensor_health(batch_f[k])
+                    print(f"[factorized][health] {k}: finite_ratio={fr:.6f}, max_abs={ma:.3e}")
+                for k in ("pred_s_hr6h_surface", "pred_t_lr1h_surface", "pred_st_hr1h_surface"):
+                    fr, ma = _tensor_health(preds[k])
+                    print(f"[factorized][health] {k}: finite_ratio={fr:.6f}, max_abs={ma:.3e}")
+            except Exception as health_err:
+                print(f"[factorized][health] failed to collect tensor health: {health_err}")
             optimizer.zero_grad()
             pending_loss = None
             continue
@@ -1064,6 +1478,7 @@ def train_one_epoch_factorized(
         metric_logger.update(loss_t=loss_dict["loss_t"].item())
         metric_logger.update(loss_st=loss_dict["loss_st"].item())
         metric_logger.update(loss_comm=loss_dict["loss_comm"].item())
+        metric_logger.update(loss_backbone=loss_dict["loss_backbone"].item())
         metric_logger.update(loss_scale=loss_scale_value)
 
         min_lr = 10.0
@@ -1090,6 +1505,7 @@ def train_one_epoch_factorized(
             log_writer.update(loss_t=loss_dict["loss_t"].item(), head="loss")
             log_writer.update(loss_st=loss_dict["loss_st"].item(), head="loss")
             log_writer.update(loss_comm=loss_dict["loss_comm"].item(), head="loss")
+            log_writer.update(loss_backbone=loss_dict["loss_backbone"].item(), head="loss")
             log_writer.update(loss_scale=loss_scale_value, head="opt")
             log_writer.update(lr=max_lr, head="opt")
             log_writer.update(min_lr=min_lr, head="opt")
@@ -1126,14 +1542,18 @@ def validation_one_epoch_factorized(
     lambda_t=1.0,
     lambda_st=1.0,
     lambda_comm=1.0,
+    lambda_backbone=0.0,
     surface_var_weights=None,
     upper_var_weights=None,
     upper_level_weights=None,
     interior_only_1h=False,
+    lt_supervision_source="lr_path",
+    lt_hybrid_alpha=0.5,
     spatial_corrector=None,
     temporal_corrector=None,
     spatial_correction_scale=1.0,
     temporal_correction_scale=1.0,
+    temporal_fix_endpoints=True,
     enable_backbone_lora=False,
     print_freq=20,
     val_max_steps=-1,
@@ -1142,6 +1562,7 @@ def validation_one_epoch_factorized(
     metric_logger.add_meter("prep_s", utils.SmoothedValue(window_size=20, fmt="{avg:.3f}"))
     metric_logger.add_meter("model_s", utils.SmoothedValue(window_size=20, fmt="{avg:.3f}"))
     metric_logger.add_meter("post_s", utils.SmoothedValue(window_size=20, fmt="{avg:.3f}"))
+    metric_logger.add_meter("loss_backbone", utils.SmoothedValue(window_size=20, fmt="{avg:.3f}"))
     header = "Val:"
 
     model.eval()
@@ -1159,14 +1580,6 @@ def validation_one_epoch_factorized(
         level=level,
         correction_scale=spatial_correction_scale,
     )
-    temporal_correction_fn = _build_temporal_correction_fn(
-        temporal_corrector=temporal_corrector,
-        surf_vars=surf_vars,
-        upper_vars=upper_vars,
-        level=level,
-        correction_scale=temporal_correction_scale,
-    )
-
     warned_short_seq = False
     for data_iter_step, batch in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
         if val_max_steps is not None and val_max_steps > 0 and data_iter_step >= val_max_steps:
@@ -1177,12 +1590,25 @@ def validation_one_epoch_factorized(
             hr_hw=hr_hw,
             lr_hw=lr_hw,
             device=device,
+            lat_hr=lat_hr,
+            lon_hr=lon_hr,
+            lat_lr=lat_lr,
+            lon_lr=lon_lr,
         )
         if torch.cuda.is_available():
             torch.cuda.synchronize()
         t1 = time.perf_counter()
 
         with torch.inference_mode():
+            temporal_correction_fn = _build_temporal_correction_fn(
+                temporal_corrector=temporal_corrector,
+                surf_vars=surf_vars,
+                upper_vars=upper_vars,
+                level=level,
+                correction_scale=temporal_correction_scale,
+                fix_endpoints=temporal_fix_endpoints,
+                target_offsets_hours=batch_f.get("target_offsets_hours", None),
+            )
             backbone_last_surface = None
             backbone_last_upper = None
             if enable_backbone_lora:
@@ -1203,6 +1629,20 @@ def validation_one_epoch_factorized(
                 backbone_s_last_surface=backbone_last_surface,
                 backbone_s_last_upper=backbone_last_upper,
             )
+            if lt_supervision_source in {"st_downscaled", "hybrid"}:
+                lat_hr_d = lat_hr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+                lon_hr_d = lon_hr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+                lat_lr_d = lat_lr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+                lon_lr_d = lon_lr.to(device=preds["pred_st_hr1h_surface"].device, dtype=preds["pred_st_hr1h_surface"].dtype)
+                preds["pred_st_lr1h_surface"] = _conservative_downscale_surface_bt_vhw(
+                    preds["pred_st_hr1h_surface"], lat_hr_d, lon_hr_d, lat_lr_d, lon_lr_d
+                )
+                preds["pred_st_lr1h_upper"] = _conservative_downscale_upper_bt_vchw(
+                    preds["pred_st_hr1h_upper"], lat_hr_d, lon_hr_d, lat_lr_d, lon_lr_d
+                )
+            if backbone_last_surface is not None and backbone_last_upper is not None:
+                preds["backbone_last_surface"] = backbone_last_surface
+                preds["backbone_last_upper"] = backbone_last_upper
             loss, loss_dict = loss_fn(
                 preds=preds,
                 batch=batch_f,
@@ -1210,10 +1650,13 @@ def validation_one_epoch_factorized(
                 lambda_t=lambda_t,
                 lambda_st=lambda_st,
                 lambda_comm=lambda_comm,
+                lambda_backbone=lambda_backbone,
                 surface_var_weights=surface_var_weights,
                 upper_var_weights=upper_var_weights,
                 upper_level_weights=upper_level_weights,
                 interior_only_1h=interior_only_1h,
+                lt_supervision_source=lt_supervision_source,
+                lt_hybrid_alpha=lt_hybrid_alpha,
                 surf_vars=surf_vars,
                 upper_vars=upper_vars,
                 level=level,
@@ -1227,7 +1670,8 @@ def validation_one_epoch_factorized(
         target_seq_surface = batch_f["y_hr1h_surface"]
         target_seq_upper = batch_f["y_hr1h_upper"]
 
-        # Default setting is include_endpoints=True -> sequence is [t0, t+1, ..., t+6] (len=7).
+        # Default setting is include_endpoints=False -> sequence is [t+1, ..., t+6] (len=6).
+        # With include_endpoints=True, sequence is [t+0, t+1, ..., t+6] (len=7).
         # We expose per-delta metrics for t+1..t+6 when available.
         seq_len = pred_seq_surface.shape[1]
         if seq_len >= 7:
@@ -1269,6 +1713,7 @@ def validation_one_epoch_factorized(
         metric_logger.update(loss_t=loss_dict["loss_t"].item())
         metric_logger.update(loss_st=loss_dict["loss_st"].item())
         metric_logger.update(loss_comm=loss_dict["loss_comm"].item())
+        metric_logger.update(loss_backbone=loss_dict["loss_backbone"].item())
         metric_logger.update(prep_s=t1 - t0)
         metric_logger.update(model_s=t2 - t1)
         metric_logger.update(post_s=t3 - t2)

downscaling/output/WeatherPEFT/checkpoint-19.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4df5855ef1cd3b158eb5b2724fb4e3486ee2e543fe0cb799f5b88f72be3cfca1
+size 5062400227

downscaling/output/WeatherPEFT/checkpoint-29.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f90edeba59a2bc62285624d7fb5c532e535adca27c59d756c22234660a8230c
+size 5062400227

downscaling/output/WeatherPEFT/checkpoint-9.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:78f21d453f4d85be909ec43cb4876e021acafd757e3668fc0617732a2eef4bd5
+size 5062398726

downscaling/output/backbone_anchor_smoke/log.txt

@@ -0,0 +1,4 @@
+"Eval:"
+{"val_prep_s": 0.07330790814012289, "val_model_s": 0.5823051929473877, "val_post_s": 0.010655493009835482, "val_valid_loss": 0.6934626698493958, "val_loss_s": 0.2821418046951294, "val_loss_t": 0.05917414277791977, "val_loss_st": 0.238005131483078, "val_loss_comm": 0.11414159834384918, "val_loss_backbone": 0.0, "val_rmse_2t_tplus1h": 2.2092418670654297, "val_mean_bias_2t_tplus1h": 0.01837158203125, "val_rmse_10u_tplus1h": 2.0027098655700684, "val_mean_bias_10u_tplus1h": -0.045052409172058105, "val_rmse_10v_tplus1h": 2.0557847023010254, "val_mean_bias_10v_tplus1h": 0.08185440301895142, "val_rmse_t850_tplus1h": 1.5737048387527466, "val_mean_bias_t850_tplus1h": 0.1566162109375, "val_rmse_z500_tplus1h": 235.0081024169922, "val_mean_bias_z500_tplus1h": 13.34765625, "val_rmse_2t_tplus2h": 2.311091661453247, "val_mean_bias_2t_tplus2h": -0.047149658203125, "val_rmse_10u_tplus2h": 2.2796661853790283, "val_mean_bias_10u_tplus2h": -0.10876637697219849, "val_rmse_10v_tplus2h": 2.4048244953155518, "val_mean_bias_10v_tplus2h": 0.15996196866035461, "val_rmse_t850_tplus2h": 1.8633790016174316, "val_mean_bias_t850_tplus2h": 0.289154052734375, "val_rmse_z500_tplus2h": 294.5515441894531, "val_mean_bias_z500_tplus2h": 25.08203125, "val_rmse_2t_tplus3h": 2.514113426208496, "val_mean_bias_2t_tplus3h": -0.05084228515625, "val_rmse_10u_tplus3h": 2.639458417892456, "val_mean_bias_10u_tplus3h": -0.17488084733486176, "val_rmse_10v_tplus3h": 2.870990514755249, "val_mean_bias_10v_tplus3h": 0.233941450715065, "val_rmse_t850_tplus3h": 2.23453950881958, "val_mean_bias_t850_tplus3h": 0.416259765625, "val_rmse_z500_tplus3h": 366.6177673339844, "val_mean_bias_z500_tplus3h": 35.8828125, "val_rmse_2t_tplus4h": 2.748657464981079, "val_mean_bias_2t_tplus4h": -0.073028564453125, "val_rmse_10u_tplus4h": 3.0420165061950684, "val_mean_bias_10u_tplus4h": -0.24148303270339966, "val_rmse_10v_tplus4h": 3.3935253620147705, "val_mean_bias_10v_tplus4h": 0.29964739084243774, "val_rmse_t850_tplus4h": 2.6447532176971436, "val_mean_bias_t850_tplus4h": 0.5400390625, "val_rmse_z500_tplus4h": 444.1191101074219, "val_mean_bias_z500_tplus4h": 46.0546875, "val_rmse_2t_tplus5h": 2.8865745067596436, "val_mean_bias_2t_tplus5h": -0.24896240234375, "val_rmse_10u_tplus5h": 3.4593324661254883, "val_mean_bias_10u_tplus5h": -0.30440056324005127, "val_rmse_10v_tplus5h": 3.940408706665039, "val_mean_bias_10v_tplus5h": 0.3687131106853485, "val_rmse_t850_tplus5h": 3.0733869075775146, "val_mean_bias_t850_tplus5h": 0.66290283203125, "val_rmse_z500_tplus5h": 522.9864501953125, "val_mean_bias_z500_tplus5h": 55.27734375, "val_rmse_2t_tplus6h": 3.113255262374878, "val_mean_bias_2t_tplus6h": -0.284454345703125, "val_rmse_10u_tplus6h": 3.8806846141815186, "val_mean_bias_10u_tplus6h": -0.3701193332672119, "val_rmse_10v_tplus6h": 4.501622200012207, "val_mean_bias_10v_tplus6h": 0.4297124147415161, "val_rmse_t850_tplus6h": 3.5143191814422607, "val_mean_bias_t850_tplus6h": 0.78155517578125, "val_rmse_z500_tplus6h": 600.989013671875, "val_mean_bias_z500_tplus6h": 63.31640625, "val_rmse_2t": 3.113255262374878, "val_mean_bias_2t": -0.284454345703125, "val_rmse_10u": 3.8806846141815186, "val_mean_bias_10u": -0.3701193332672119, "val_rmse_10v": 4.501622200012207, "val_mean_bias_10v": 0.4297124147415161, "val_rmse_t850": 3.5143191814422607, "val_mean_bias_t850": 0.78155517578125, "val_rmse_z500": 600.989013671875, "val_mean_bias_z500": 63.31640625}
+" &0.073 &0.582 &0.011 &0.282 &0.059 &0.238 &0.114 &0.0 &2.209 &0.018 &2.003 &-0.045 &2.056 &0.082 &1.574 &0.157 &235.008 &13.348 &2.311 &-0.047 &2.28 &-0.109 &2.405 &0.16 &1.863 &0.289 &294.552 &25.082 &2.514 &-0.051 &2.639 &-0.175 &2.871 &0.234 &2.235 &0.416 &366.618 &35.883 &2.749 &-0.073 &3.042 &-0.241 &3.394 &0.3 &2.645 &0.54 &444.119 &46.055 &2.887 &-0.249 &3.459 &-0.304 &3.94 &0.369 &3.073 &0.663 &522.986 &55.277 &3.113 &-0.284 &3.881 &-0.37 &4.502 &0.43 &3.514 &0.782 &600.989 &63.316 &3.113 &-0.284 &3.881 &-0.37 &4.502 &0.43 &3.514 &0.782 &600.989 &63.316"
+

downscaling/output/val_freq_smoke/log.txt

@@ -0,0 +1,4 @@
+"Eval:"
+{"val_prep_s": 0.05949366791173816, "val_model_s": 0.3861960484646261, "val_post_s": 0.005747961113229394, "val_valid_loss": 0.6520180404186249, "val_loss_s": 0.2656427472829819, "val_loss_t": 0.05454788729548454, "val_loss_st": 0.22411399334669113, "val_loss_comm": 0.10771343111991882, "val_rmse_2t_tplus1h": 2.1595544815063477, "val_mean_bias_2t_tplus1h": 0.0153045654296875, "val_rmse_10u_tplus1h": 1.9929717779159546, "val_mean_bias_10u_tplus1h": -0.06720291078090668, "val_rmse_10v_tplus1h": 2.0678257942199707, "val_mean_bias_10v_tplus1h": 0.04859817773103714, "val_rmse_t850_tplus1h": 1.5690162181854248, "val_mean_bias_t850_tplus1h": 0.1692962646484375, "val_rmse_z500_tplus1h": 229.7540283203125, "val_mean_bias_z500_tplus1h": 11.56640625, "val_rmse_2t_tplus2h": 2.325566053390503, "val_mean_bias_2t_tplus2h": -0.111175537109375, "val_rmse_10u_tplus2h": 2.259079933166504, "val_mean_bias_10u_tplus2h": -0.12979388236999512, "val_rmse_10v_tplus2h": 2.4088897705078125, "val_mean_bias_10v_tplus2h": 0.0989474430680275, "val_rmse_t850_tplus2h": 1.8684790134429932, "val_mean_bias_t850_tplus2h": 0.3101043701171875, "val_rmse_z500_tplus2h": 285.2650146484375, "val_mean_bias_z500_tplus2h": 23.6953125, "val_rmse_2t_tplus3h": 2.559150457382202, "val_mean_bias_2t_tplus3h": -0.1783599853515625, "val_rmse_10u_tplus3h": 2.5984840393066406, "val_mean_bias_10u_tplus3h": -0.1920362114906311, "val_rmse_10v_tplus3h": 2.8546905517578125, "val_mean_bias_10v_tplus3h": 0.15115559101104736, "val_rmse_t850_tplus3h": 2.2511355876922607, "val_mean_bias_t850_tplus3h": 0.44720458984375, "val_rmse_z500_tplus3h": 353.3114013671875, "val_mean_bias_z500_tplus3h": 34.533203125, "val_rmse_2t_tplus4h": 2.829373598098755, "val_mean_bias_2t_tplus4h": -0.268890380859375, "val_rmse_10u_tplus4h": 2.954393148422241, "val_mean_bias_10u_tplus4h": -0.2777942419052124, "val_rmse_10v_tplus4h": 3.3426637649536133, "val_mean_bias_10v_tplus4h": 0.18345022201538086, "val_rmse_t850_tplus4h": 2.6910040378570557, "val_mean_bias_t850_tplus4h": 0.591400146484375, "val_rmse_z500_tplus4h": 428.23114013671875, "val_mean_bias_z500_tplus4h": 46.494140625, "val_rmse_2t_tplus5h": 2.988680839538574, "val_mean_bias_2t_tplus5h": -0.3951263427734375, "val_rmse_10u_tplus5h": 3.360450506210327, "val_mean_bias_10u_tplus5h": -0.33785462379455566, "val_rmse_10v_tplus5h": 3.8764631748199463, "val_mean_bias_10v_tplus5h": 0.23951280117034912, "val_rmse_t850_tplus5h": 3.136207103729248, "val_mean_bias_t850_tplus5h": 0.7298583984375, "val_rmse_z500_tplus5h": 503.41949462890625, "val_mean_bias_z500_tplus5h": 58.720703125, "val_rmse_2t_tplus6h": 3.1790788173675537, "val_mean_bias_2t_tplus6h": -0.4730224609375, "val_rmse_10u_tplus6h": 3.7709269523620605, "val_mean_bias_10u_tplus6h": -0.39494484663009644, "val_rmse_10v_tplus6h": 4.418558120727539, "val_mean_bias_10v_tplus6h": 0.29149335622787476, "val_rmse_t850_tplus6h": 3.5930073261260986, "val_mean_bias_t850_tplus6h": 0.868896484375, "val_rmse_z500_tplus6h": 578.9581298828125, "val_mean_bias_z500_tplus6h": 69.849609375, "val_rmse_2t": 3.1790788173675537, "val_mean_bias_2t": -0.4730224609375, "val_rmse_10u": 3.7709269523620605, "val_mean_bias_10u": -0.39494484663009644, "val_rmse_10v": 4.418558120727539, "val_mean_bias_10v": 0.29149335622787476, "val_rmse_t850": 3.5930073261260986, "val_mean_bias_t850": 0.868896484375, "val_rmse_z500": 578.9581298828125, "val_mean_bias_z500": 69.849609375}
+" &0.059 &0.386 &0.006 &0.266 &0.055 &0.224 &0.108 &2.16 &0.015 &1.993 &-0.067 &2.068 &0.049 &1.569 &0.169 &229.754 &11.566 &2.326 &-0.111 &2.259 &-0.13 &2.409 &0.099 &1.868 &0.31 &285.265 &23.695 &2.559 &-0.178 &2.598 &-0.192 &2.855 &0.151 &2.251 &0.447 &353.311 &34.533 &2.829 &-0.269 &2.954 &-0.278 &3.343 &0.183 &2.691 &0.591 &428.231 &46.494 &2.989 &-0.395 &3.36 &-0.338 &3.876 &0.24 &3.136 &0.73 &503.419 &58.721 &3.179 &-0.473 &3.771 &-0.395 &4.419 &0.291 &3.593 &0.869 &578.958 &69.85 &3.179 &-0.473 &3.771 &-0.395 &4.419 &0.291 &3.593 &0.869 &578.958 &69.85"
+

downscaling/output/val_speed_after/log.txt

@@ -0,0 +1,4 @@
+"Eval:"
+{"val_prep_s": 0.084631056097957, "val_model_s": 0.6615325791062787, "val_post_s": 0.0011346005291367571, "val_sample_s": 0.7511531277947748, "val_valid_loss": 0.6753831406434377, "val_loss_s": 0.27440425256888074, "val_loss_t": 0.05877547711133957, "val_loss_st": 0.23192234834035239, "val_loss_comm": 0.11028106758991878, "val_rmse_2t_tplus1h": 2.120826482772827, "val_mean_bias_2t_tplus1h": -0.08579571545124054, "val_rmse_10u_tplus1h": 2.0343947410583496, "val_mean_bias_10u_tplus1h": -0.044072579592466354, "val_rmse_10v_tplus1h": 2.0728132724761963, "val_mean_bias_10v_tplus1h": 0.05038035660982132, "val_rmse_t850_tplus1h": 1.610738754272461, "val_mean_bias_t850_tplus1h": 0.1519871950149536, "val_rmse_z500_tplus1h": 237.20074462890625, "val_mean_bias_z500_tplus1h": 9.716571807861328, "val_rmse_2t_tplus2h": 2.293344020843506, "val_mean_bias_2t_tplus2h": -0.18076401948928833, "val_rmse_10u_tplus2h": 2.3025588989257812, "val_mean_bias_10u_tplus2h": -0.08629271388053894, "val_rmse_10v_tplus2h": 2.409632444381714, "val_mean_bias_10v_tplus2h": 0.102354496717453, "val_rmse_t850_tplus2h": 1.9087270498275757, "val_mean_bias_t850_tplus2h": 0.2832726240158081, "val_rmse_z500_tplus2h": 296.6208801269531, "val_mean_bias_z500_tplus2h": 20.60210418701172, "val_rmse_2t_tplus3h": 2.481692314147949, "val_mean_bias_2t_tplus3h": -0.2545323371887207, "val_rmse_10u_tplus3h": 2.6446375846862793, "val_mean_bias_10u_tplus3h": -0.12955179810523987, "val_rmse_10v_tplus3h": 2.8543546199798584, "val_mean_bias_10v_tplus3h": 0.15386676788330078, "val_rmse_t850_tplus3h": 2.282665491104126, "val_mean_bias_t850_tplus3h": 0.4141359329223633, "val_rmse_z500_tplus3h": 367.32757568359375, "val_mean_bias_z500_tplus3h": 30.99903678894043, "val_rmse_2t_tplus4h": 2.7042994499206543, "val_mean_bias_2t_tplus4h": -0.3467860221862793, "val_rmse_10u_tplus4h": 2.998748540878296, "val_mean_bias_10u_tplus4h": -0.19702103734016418, "val_rmse_10v_tplus4h": 3.3468921184539795, "val_mean_bias_10v_tplus4h": 0.20130321383476257, "val_rmse_t850_tplus4h": 2.7033567428588867, "val_mean_bias_t850_tplus4h": 0.5463706254959106, "val_rmse_z500_tplus4h": 444.17047119140625, "val_mean_bias_z500_tplus4h": 42.49504089355469, "val_rmse_2t_tplus5h": 2.890420913696289, "val_mean_bias_2t_tplus5h": -0.4551726281642914, "val_rmse_10u_tplus5h": 3.40653395652771, "val_mean_bias_10u_tplus5h": -0.23875072598457336, "val_rmse_10v_tplus5h": 3.876429319381714, "val_mean_bias_10v_tplus5h": 0.25390106439590454, "val_rmse_t850_tplus5h": 3.1348283290863037, "val_mean_bias_t850_tplus5h": 0.6817998886108398, "val_rmse_z500_tplus5h": 521.3743286132812, "val_mean_bias_z500_tplus5h": 54.1654052734375, "val_rmse_2t_tplus6h": 3.1032655239105225, "val_mean_bias_2t_tplus6h": -0.538508415222168, "val_rmse_10u_tplus6h": 3.821709394454956, "val_mean_bias_10u_tplus6h": -0.27899816632270813, "val_rmse_10v_tplus6h": 4.415755271911621, "val_mean_bias_10v_tplus6h": 0.3063257932662964, "val_rmse_t850_tplus6h": 3.5775928497314453, "val_mean_bias_t850_tplus6h": 0.8159565329551697, "val_rmse_z500_tplus6h": 599.5579833984375, "val_mean_bias_z500_tplus6h": 65.35942077636719, "val_rmse_2t": 3.1032655239105225, "val_mean_bias_2t": -0.538508415222168, "val_rmse_10u": 3.821709394454956, "val_mean_bias_10u": -0.27899816632270813, "val_rmse_10v": 4.415755271911621, "val_mean_bias_10v": 0.3063257932662964, "val_rmse_t850": 3.5775928497314453, "val_mean_bias_t850": 0.8159565329551697, "val_rmse_z500": 599.5579833984375, "val_mean_bias_z500": 65.35942077636719}
+" &0.085 &0.662 &0.001 &0.751 &0.274 &0.059 &0.232 &0.11 &2.121 &-0.086 &2.034 &-0.044 &2.073 &0.05 &1.611 &0.152 &237.201 &9.717 &2.293 &-0.181 &2.303 &-0.086 &2.41 &0.102 &1.909 &0.283 &296.621 &20.602 &2.482 &-0.255 &2.645 &-0.13 &2.854 &0.154 &2.283 &0.414 &367.328 &30.999 &2.704 &-0.347 &2.999 &-0.197 &3.347 &0.201 &2.703 &0.546 &444.17 &42.495 &2.89 &-0.455 &3.407 &-0.239 &3.876 &0.254 &3.135 &0.682 &521.374 &54.165 &3.103 &-0.539 &3.822 &-0.279 &4.416 &0.306 &3.578 &0.816 &599.558 &65.359 &3.103 &-0.539 &3.822 &-0.279 &4.416 &0.306 &3.578 &0.816 &599.558 &65.359"
+

downscaling/output/val_speed_baseline/log.txt

@@ -0,0 +1,4 @@
+"Eval:"
+{"val_prep_s": 0.05412517949783554, "val_model_s": 0.16398887091781944, "val_post_s": 0.003232622635550797, "val_valid_loss": 0.6753831406434377, "val_loss_s": 0.27440425256888074, "val_loss_t": 0.05877547711133957, "val_loss_st": 0.23192234834035239, "val_loss_comm": 0.11028106758991878, "val_rmse_2t_tplus1h": 2.120826482772827, "val_mean_bias_2t_tplus1h": -0.0857950896024704, "val_rmse_10u_tplus1h": 2.0343947410583496, "val_mean_bias_10u_tplus1h": -0.044072579592466354, "val_rmse_10v_tplus1h": 2.0728132724761963, "val_mean_bias_10v_tplus1h": 0.05038034915924072, "val_rmse_t850_tplus1h": 1.610738754272461, "val_mean_bias_t850_tplus1h": 0.1519927978515625, "val_rmse_z500_tplus1h": 237.20074462890625, "val_mean_bias_z500_tplus1h": 9.715169906616211, "val_rmse_2t_tplus2h": 2.293344020843506, "val_mean_bias_2t_tplus2h": -0.18076324462890625, "val_rmse_10u_tplus2h": 2.3025588989257812, "val_mean_bias_10u_tplus2h": -0.08629272878170013, "val_rmse_10v_tplus2h": 2.409632444381714, "val_mean_bias_10v_tplus2h": 0.1023545116186142, "val_rmse_t850_tplus2h": 1.9087270498275757, "val_mean_bias_t850_tplus2h": 0.28327688574790955, "val_rmse_z500_tplus2h": 296.6208801269531, "val_mean_bias_z500_tplus2h": 20.602214813232422, "val_rmse_2t_tplus3h": 2.481692314147949, "val_mean_bias_2t_tplus3h": -0.2545318603515625, "val_rmse_10u_tplus3h": 2.6446375846862793, "val_mean_bias_10u_tplus3h": -0.12955179810523987, "val_rmse_10v_tplus3h": 2.8543546199798584, "val_mean_bias_10v_tplus3h": 0.15386676788330078, "val_rmse_t850_tplus3h": 2.282665491104126, "val_mean_bias_t850_tplus3h": 0.4141286313533783, "val_rmse_z500_tplus3h": 367.32757568359375, "val_mean_bias_z500_tplus3h": 30.99837303161621, "val_rmse_2t_tplus4h": 2.7042994499206543, "val_mean_bias_2t_tplus4h": -0.34678396582603455, "val_rmse_10u_tplus4h": 2.998748540878296, "val_mean_bias_10u_tplus4h": -0.19702103734016418, "val_rmse_10v_tplus4h": 3.3468921184539795, "val_mean_bias_10v_tplus4h": 0.20130321383476257, "val_rmse_t850_tplus4h": 2.7033567428588867, "val_mean_bias_t850_tplus4h": 0.5463689565658569, "val_rmse_z500_tplus4h": 444.17047119140625, "val_mean_bias_z500_tplus4h": 42.49609375, "val_rmse_2t_tplus5h": 2.890420913696289, "val_mean_bias_2t_tplus5h": -0.45517224073410034, "val_rmse_10u_tplus5h": 3.40653395652771, "val_mean_bias_10u_tplus5h": -0.23875072598457336, "val_rmse_10v_tplus5h": 3.876429319381714, "val_mean_bias_10v_tplus5h": 0.25390106439590454, "val_rmse_t850_tplus5h": 3.1348283290863037, "val_mean_bias_t850_tplus5h": 0.6818059682846069, "val_rmse_z500_tplus5h": 521.3743286132812, "val_mean_bias_z500_tplus5h": 54.1650390625, "val_rmse_2t_tplus6h": 3.1032655239105225, "val_mean_bias_2t_tplus6h": -0.5385081171989441, "val_rmse_10u_tplus6h": 3.821709394454956, "val_mean_bias_10u_tplus6h": -0.2789981961250305, "val_rmse_10v_tplus6h": 4.415755271911621, "val_mean_bias_10v_tplus6h": 0.3063257932662964, "val_rmse_t850_tplus6h": 3.5775928497314453, "val_mean_bias_t850_tplus6h": 0.8159612417221069, "val_rmse_z500_tplus6h": 599.5579833984375, "val_mean_bias_z500_tplus6h": 65.3583984375, "val_rmse_2t": 3.1032655239105225, "val_mean_bias_2t": -0.5385081171989441, "val_rmse_10u": 3.821709394454956, "val_mean_bias_10u": -0.2789981961250305, "val_rmse_10v": 4.415755271911621, "val_mean_bias_10v": 0.3063257932662964, "val_rmse_t850": 3.5775928497314453, "val_mean_bias_t850": 0.8159612417221069, "val_rmse_z500": 599.5579833984375, "val_mean_bias_z500": 65.3583984375}
+" &0.054 &0.164 &0.003 &0.274 &0.059 &0.232 &0.11 &2.121 &-0.086 &2.034 &-0.044 &2.073 &0.05 &1.611 &0.152 &237.201 &9.715 &2.293 &-0.181 &2.303 &-0.086 &2.41 &0.102 &1.909 &0.283 &296.621 &20.602 &2.482 &-0.255 &2.645 &-0.13 &2.854 &0.154 &2.283 &0.414 &367.328 &30.998 &2.704 &-0.347 &2.999 &-0.197 &3.347 &0.201 &2.703 &0.546 &444.17 &42.496 &2.89 &-0.455 &3.407 &-0.239 &3.876 &0.254 &3.135 &0.682 &521.374 &54.165 &3.103 &-0.539 &3.822 &-0.279 &4.416 &0.306 &3.578 &0.816 &599.558 &65.358 &3.103 &-0.539 &3.822 &-0.279 &4.416 &0.306 &3.578 &0.816 &599.558 &65.358"
+

downscaling/run_downscaling.py

@@ -87,6 +87,7 @@ def _resolve_pretrained_ckpt(path_str: str) -> Path:
 
 def _attach_factorized_heads(model, surf_vars, upper_vars, level, args):
     """Attach factorized temporal/spatial correction heads."""
+    temporal_num_steps = 7 if args.include_endpoints else 6
     model.temporal_corrector = TemporalCorrectionHead(
         num_surface_vars=len(surf_vars),
         num_upper_vars=len(upper_vars),
@@ -94,6 +95,8 @@ def _attach_factorized_heads(model, surf_vars, upper_vars, level, args):
         hidden_surface=args.temporal_hidden_surface,
         hidden_upper=args.temporal_hidden_upper,
         dropout=args.temporal_dropout,
+        head_mode=args.temporal_head_mode,
+        num_time_steps=temporal_num_steps,
     )
     model.spatial_corrector = SpatialCorrectionHead(
         num_surface_vars=len(surf_vars),
@@ -327,6 +330,8 @@ def get_args():
                         help='validation progress print frequency (iterations)')
     parser.add_argument('--val_max_steps', default=-1, type=int,
                         help='if >0, limit validation iterations per epoch for faster feedback')
+    parser.add_argument('--val_freq', default=1, type=int,
+                        help='run validation every N epochs (1 means every epoch)')
     parser.add_argument('--persistent_workers', action='store_true', default=True,
                         help='keep dataloader workers persistent across epochs')
     parser.add_argument('--no_persistent_workers', action='store_false', dest='persistent_workers')
@@ -366,6 +371,8 @@ def get_args():
     parser.add_argument('--lora_mode', default='single', choices=['single', 'all'])
     parser.add_argument('--disable_backbone_lora_in_factorized', action='store_true',
                         help='If set in finetune_mode=lora, freeze LoRA adapters and train only correction heads.')
+    parser.add_argument('--disable_backbone_prior_in_factorized', action='store_true',
+                        help='If set, do not use Aurora pretrained 6h endpoint prediction as factorized prior/anchor.')
     parser.add_argument('--lr_data_dir', default='', type=str,
                         help='LR dataset folder under datasets (auto if empty), e.g. 5.625')
     parser.add_argument('--hr_data_dir', default='', type=str,
@@ -386,12 +393,19 @@ def get_args():
     parser.add_argument('--derive_lr1h_from_hr1h', action='store_true', default=True,
                         help='Build y_lr1h from y_hr1h in training loop to reduce dataset IO.')
     parser.add_argument('--no_derive_lr1h_from_hr1h', action='store_false', dest='derive_lr1h_from_hr1h')
-    parser.add_argument('--include_endpoints', action='store_true', default=True,
-                        help='Use 1h targets including endpoints t0 and t0+6h')
+    parser.add_argument('--include_endpoints', action='store_true', default=False,
+                        help='If set, use 1h targets [t+0..t+6]. Default uses [t+1..t+6].')
+    parser.add_argument('--no_include_endpoints', action='store_false', dest='include_endpoints')
     parser.add_argument('--lambda_s', default=1.0, type=float)
     parser.add_argument('--lambda_t', default=1.0, type=float)
     parser.add_argument('--lambda_st', default=1.0, type=float)
     parser.add_argument('--lambda_comm', default=1.0, type=float)
+    parser.add_argument('--lambda_backbone', default=0.1, type=float,
+                        help='Weight for Aurora-6h anchor loss on unanchored S-module 6h endpoint prediction.')
+    parser.add_argument('--lt_supervision_source', default='lr_path', choices=['lr_path', 'st_downscaled', 'hybrid'],
+                        help='Source for L_t supervision: direct LR temporal path or ST(HR) downscaled to LR.')
+    parser.add_argument('--lt_hybrid_alpha', default=0.5, type=float,
+                        help='When lt_supervision_source=hybrid, weight on direct lr_path term (0..1).')
     parser.add_argument('--interior_only_1h', action='store_true',
                         help='If set, compute 1h losses on interior frames only (exclude endpoints)')
     parser.add_argument('--ours_prompt_length', default=0, type=int,
@@ -402,6 +416,8 @@ def get_args():
                         help='Hidden channels for temporal correction head (upper vars x levels).')
     parser.add_argument('--temporal_dropout', default=0.0, type=float,
                         help='Dropout for temporal correction head.')
+    parser.add_argument('--temporal_head_mode', default='shared3d', choices=['shared3d', 'delta2d'],
+                        help='Temporal correction head type: shared 3D conv or per-delta 2D heads.')
     parser.add_argument('--spatial_hidden_surface', default=256, type=int,
                         help='Hidden channels for spatial correction head (surface vars).')
     parser.add_argument('--spatial_hidden_upper', default=512, type=int,
@@ -412,6 +428,8 @@ def get_args():
                         help='Residual scale multiplier for spatial correction outputs.')
     parser.add_argument('--temporal_correction_scale', default=1e-3, type=float,
                         help='Residual scale multiplier for temporal correction outputs.')
+    parser.add_argument('--allow_temporal_endpoint_correction', action='store_true',
+                        help='Allow temporal correction head on offset=0 frame when include_endpoints is used.')
     parser.add_argument('--disable_factorized_lr_scale', action='store_true',
                         help='Disable LR scaling by total_batch_size/256 for factorized training.')
     parser.add_argument('--disable_scale_balanced_loss', action='store_true',
@@ -563,6 +581,10 @@ def main(args, ds_init):
     )
     
     print("Sampler_train = %s" % str(sampler_train))
+    if args.distributed and (dataset_val is not None) and (not args.dist_eval):
+        args.dist_eval = True
+        print("[validation] Distributed training detected; enabling --dist_eval to shard validation across ranks.")
+
     if args.dist_eval:
         if len(dataset_val) % num_tasks != 0:
             print('Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
@@ -607,6 +629,8 @@ def main(args, ds_init):
         f"[factorized-data] derive_hr6h_from_hr1h={args.derive_hr6h_from_hr1h}, "
         f"derive_lr1h_from_hr1h={args.derive_lr1h_from_hr1h}"
     )
+    if args.derive_lr1h_from_hr1h:
+        print("[factorized-data] y_lr1h generation: conservative coarsen from y_hr1h (engine-side)")
 
     if dataset_val is not None:
         val_loader_kwargs = dict(
@@ -627,6 +651,10 @@ def main(args, ds_init):
             f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else 'n/a'}, "
             f"pin_mem={args.pin_mem}"
         )
+        print(
+            f"[validation] samples={len(dataset_val)}, batches={len(data_loader_val)}, "
+            f"val_freq={args.val_freq}, val_max_steps={args.val_max_steps}"
+        )
     else:
         data_loader_val = None
 
@@ -666,14 +694,28 @@ def main(args, ds_init):
 
     use_ours = finetune_mode == "ours"
     use_lora = finetune_mode == "lora"
-    enable_backbone_lora = finetune_mode == "lora" and (not args.disable_backbone_lora_in_factorized)
+    train_lora_adapters = finetune_mode == "lora" and (not args.disable_backbone_lora_in_factorized)
+    use_backbone_prior = not args.disable_backbone_prior_in_factorized
     if finetune_mode == "lora":
         if args.lora_mode == "all" and args.lora_steps > 1:
             print(
                 "[lora mode] NOTE: factorized path uses single-step Aurora forward (rollout_step=0), "
                 "so only step-0 LoRA branch is used even when lora_mode=all."
             )
-        print(f"[lora mode] backbone LoRA enabled in factorized loop: {enable_backbone_lora}")
+        print(f"[lora mode] train LoRA adapters: {train_lora_adapters}")
+    print(f"[factorized] use backbone 6h prior: {use_backbone_prior} (lambda_backbone={args.lambda_backbone})")
+    print(
+        "[factorized] loss lambdas: "
+        f"lambda_s={args.lambda_s}, "
+        f"lambda_t={args.lambda_t}, "
+        f"lambda_st={args.lambda_st}, "
+        f"lambda_comm={args.lambda_comm}, "
+        f"lambda_backbone={args.lambda_backbone}"
+    )
+    print(
+        f"[factorized] L_t supervision source: {args.lt_supervision_source} "
+        f"(hybrid alpha={args.lt_hybrid_alpha:.2f})"
+    )
 
     model = Aurora(
         use_ours=use_ours,
@@ -701,13 +743,14 @@ def main(args, ds_init):
     print(
         "[factorized] Attached correction heads "
         f"(temporal: {args.temporal_hidden_surface}/{args.temporal_hidden_upper}, "
+        f"temporal_mode={args.temporal_head_mode}, "
         f"spatial: {args.spatial_hidden_surface}/{args.spatial_hidden_upper})."
     )
 
     _set_trainable_parameters(
         model,
         finetune_mode,
-        train_lora_adapters=enable_backbone_lora,
+        train_lora_adapters=train_lora_adapters,
     )
     total_params, trainable_params = _count_parameters(model)
     print(f"[mode={finetune_mode}] trainable params: {trainable_params:,} / {total_params:,}")
@@ -875,15 +918,19 @@ def main(args, ds_init):
             lambda_t=args.lambda_t,
             lambda_st=args.lambda_st,
             lambda_comm=args.lambda_comm,
+            lambda_backbone=args.lambda_backbone,
             surface_var_weights=surface_var_weights,
             upper_var_weights=upper_var_level_weights,
             upper_level_weights=None,
             interior_only_1h=args.interior_only_1h,
+            lt_supervision_source=args.lt_supervision_source,
+            lt_hybrid_alpha=args.lt_hybrid_alpha,
             spatial_corrector=_get_spatial_corrector(model),
             temporal_corrector=_get_temporal_corrector(model),
             spatial_correction_scale=args.spatial_correction_scale,
             temporal_correction_scale=args.temporal_correction_scale,
-            enable_backbone_lora=enable_backbone_lora,
+            temporal_fix_endpoints=(not args.allow_temporal_endpoint_correction),
+            enable_backbone_lora=use_backbone_prior,
             print_freq=args.val_print_freq,
             val_max_steps=args.val_max_steps,
         )
@@ -908,6 +955,10 @@ def main(args, ds_init):
         f"[factorized] correction scales: spatial={args.spatial_correction_scale:.2e}, "
         f"temporal={args.temporal_correction_scale:.2e}"
     )
+    print(
+        f"[factorized] temporal endpoint correction enabled: "
+        f"{args.allow_temporal_endpoint_correction}"
+    )
     start_time = time.time()
     train_time_only = 0
     total_step = args.epochs * num_training_steps_per_epoch
@@ -929,15 +980,19 @@ def main(args, ds_init):
             surf_vars=surf_vars, upper_vars=upper_vars,
             lambda_s=args.lambda_s, lambda_t=args.lambda_t,
             lambda_st=args.lambda_st, lambda_comm=args.lambda_comm,
+            lambda_backbone=args.lambda_backbone,
             surface_var_weights=surface_var_weights,
             upper_var_weights=upper_var_level_weights,
             upper_level_weights=None,
             interior_only_1h=args.interior_only_1h,
+            lt_supervision_source=args.lt_supervision_source,
+            lt_hybrid_alpha=args.lt_hybrid_alpha,
             spatial_corrector=_get_spatial_corrector(model),
             temporal_corrector=_get_temporal_corrector(model),
             spatial_correction_scale=args.spatial_correction_scale,
             temporal_correction_scale=args.temporal_correction_scale,
-            enable_backbone_lora=enable_backbone_lora,
+            temporal_fix_endpoints=(not args.allow_temporal_endpoint_correction),
+            enable_backbone_lora=use_backbone_prior,
             accum_backward_mode=args.accum_backward_mode,
             use_ours=use_ours, total_step=total_step
         )
@@ -948,7 +1003,19 @@ def main(args, ds_init):
                     args=args, model=model, model_without_ddp=model_without_ddp, optimizer=optimizer,
                     loss_scaler=loss_scaler, epoch=epoch)
                 
-        if data_loader_val is not None:
+        should_validate = (
+            data_loader_val is not None
+            and (args.val_freq > 0)
+            and (((epoch + 1) % args.val_freq == 0) or (epoch + 1 == args.epochs))
+        )
+        test_stats = None
+        if should_validate:
+            val_start_time = time.time()
+            val_rank_samples = len(data_loader_val.sampler) if hasattr(data_loader_val, "sampler") else len(data_loader_val.dataset)
+            print(
+                f"[validation] epoch={epoch} start: rank_samples={val_rank_samples}, "
+                f"batches={len(data_loader_val)}"
+            )
             test_stats = validation_one_epoch_factorized(
                 data_loader_val,
                 model,
@@ -966,31 +1033,49 @@ def main(args, ds_init):
                 lambda_t=args.lambda_t,
                 lambda_st=args.lambda_st,
                 lambda_comm=args.lambda_comm,
+                lambda_backbone=args.lambda_backbone,
                 surface_var_weights=surface_var_weights,
                 upper_var_weights=upper_var_level_weights,
                 upper_level_weights=None,
                 interior_only_1h=args.interior_only_1h,
+                lt_supervision_source=args.lt_supervision_source,
+                lt_hybrid_alpha=args.lt_hybrid_alpha,
                 spatial_corrector=_get_spatial_corrector(model),
                 temporal_corrector=_get_temporal_corrector(model),
                 spatial_correction_scale=args.spatial_correction_scale,
                 temporal_correction_scale=args.temporal_correction_scale,
-                enable_backbone_lora=enable_backbone_lora,
+                temporal_fix_endpoints=(not args.allow_temporal_endpoint_correction),
+                enable_backbone_lora=use_backbone_prior,
                 print_freq=args.val_print_freq,
                 val_max_steps=args.val_max_steps,
             )
+            val_elapsed = time.time() - val_start_time
+            val_steps_done = min(len(data_loader_val), args.val_max_steps) if args.val_max_steps > 0 else len(data_loader_val)
+            print(
+                f"[validation] epoch={epoch} done in {val_elapsed:.1f}s "
+                f"(~{val_elapsed / max(val_rank_samples, 1):.3f}s/sample on this rank, "
+                f"{val_steps_done} batches)."
+            )
+        elif data_loader_val is not None:
+            print(f"[validation] epoch={epoch} skipped (val_freq={args.val_freq}).")
         log_stats = {'epoch': epoch,
                     **{f'train_{k}': v for k, v in train_stats.items()}}
-        test_log = {**{f'val_{k}': _to_float(v) for k, v in test_stats.items()}}
-        
-        copy_log = " &"+" &".join([str(round(_to_float(v), 3)) for k, v in test_stats.items() if k!="valid_loss"])
+        if test_stats is not None:
+            test_log = {**{f'val_{k}': _to_float(v) for k, v in test_stats.items()}}
+            copy_log = " &"+" &".join([str(round(_to_float(v), 3)) for k, v in test_stats.items() if k!="valid_loss"])
+        else:
+            test_log = None
+            copy_log = None
         
         if args.output_dir and utils.is_main_process():
             if log_writer is not None:
                 log_writer.flush()
             with open(os.path.join(args.output_dir, "log.txt"), mode="a", encoding="utf-8") as f:
                 f.write(json.dumps(log_stats) + "\n")
-                f.write(json.dumps(test_log) + "\n")
-                f.write(json.dumps(copy_log) + "\n" + "\n")
+                if test_log is not None:
+                    f.write(json.dumps(test_log) + "\n")
+                    f.write(json.dumps(copy_log) + "\n")
+                f.write("\n")
 
     total_time = time.time() - start_time
     total_time_str = str(datetime.timedelta(seconds=int(total_time)))

downscaling/script_run_downscaling.sh

@@ -10,6 +10,22 @@ VAL_END_DATE="2018"
 OUTPUT_DIR="output/WeatherPEFT" 
 GRAD_ACCUM_STEPS=8
 
+# Preserve-first recipe:
+# 1) keep pretrained 6h behavior strong at start
+# 2) learn 1h~5h corrections gradually
+LAMBDA_BACKBONE=1.0
+LAMBDA_S=1.0
+LAMBDA_T=1.0
+LAMBDA_ST=1.0
+LAMBDA_COMM=1.0
+LT_SUPERVISION_SOURCE="lr_path"
+SPATIAL_HIDDEN_SURFACE=64
+SPATIAL_HIDDEN_UPPER=128
+TEMPORAL_HIDDEN_SURFACE=64
+TEMPORAL_HIDDEN_UPPER=128
+TEMPORAL_HEAD_MODE="delta2d"
+# backbone prior is ON by default (do not pass --disable_backbone_prior_in_factorized)
+
 OMP_NUM_THREADS=1 python -m torch.distributed.run --nproc_per_node=2 \
     --master_port 12326 --nnodes=1  --node_rank=0 --master_addr="127.0.0.1" \
     run_downscaling.py \
@@ -21,18 +37,19 @@ OMP_NUM_THREADS=1 python -m torch.distributed.run --nproc_per_node=2 \
     --output_dir ${OUTPUT_DIR} \
     --batch_size 2 \
     --grad_accum_steps ${GRAD_ACCUM_STEPS} \
-    --val_batch_size 5 \
+    --val_batch_size 2 \
     --num_workers 8 \
     --persistent_workers \
     --prefetch_factor 4 \
     --save_ckpt_freq 10 \
     --opt adamw \
-    --lr 7e-4 \
+    --lr 5e-4 \
     --opt_betas 0.9 0.999 \
     --weight_decay 0.05 \
     --clip_grad 1.0 \
     --warmup_epochs 3 \
     --epochs 30 \
+    --val_freq 1 \
     --dist_eval \
     --mode lora \
     --finetune_mode lora \
@@ -44,9 +61,22 @@ OMP_NUM_THREADS=1 python -m torch.distributed.run --nproc_per_node=2 \
     --hr1h_data_dir 1.5_1h_nc \
     --lr_degree_tag 5.625 \
     --hr_degree_tag 1.5 \
+    --no_include_endpoints \
     --norm_mean_path ../aux_data/normalize_mean_wb2.npz \
     --norm_std_path ../aux_data/normalize_std_wb2.npz \
-    --norm_std_floor 1e-6 \
+    --norm_std_floor 1e-8 \
+    --lambda_s ${LAMBDA_S} \
+    --lambda_t ${LAMBDA_T} \
+    --lambda_st ${LAMBDA_ST} \
+    --lambda_comm ${LAMBDA_COMM} \
+    --lambda_backbone ${LAMBDA_BACKBONE} \
+    --lt_supervision_source ${LT_SUPERVISION_SOURCE} \
+    --disable_scale_balanced_loss \
+    --spatial_hidden_surface ${SPATIAL_HIDDEN_SURFACE} \
+    --spatial_hidden_upper ${SPATIAL_HIDDEN_UPPER} \
+    --temporal_hidden_surface ${TEMPORAL_HIDDEN_SURFACE} \
+    --temporal_hidden_upper ${TEMPORAL_HIDDEN_UPPER} \
+    --temporal_head_mode ${TEMPORAL_HEAD_MODE} \
     --no_auto_resume \
     --spatial_correction_scale 1.0 \
     --temporal_correction_scale 1.0 \

downscaling/script_smoke_feasibility.sh

@@ -48,6 +48,7 @@ for MODE in "${MODES[@]}"; do
     --hr1h_data_dir 1.5_1h_nc \
     --lr_degree_tag 5.625 \
     --hr_degree_tag 1.5 \
+    --no_include_endpoints \
     --norm_mean_path ../aux_data/normalize_mean_wb2.npz \
     --norm_std_path ../aux_data/normalize_std_wb2.npz \
     --norm_std_floor 1e-8 \