From 8d35753e7385748ebb0f2d780cf492b839ce603a Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Thu, 30 Apr 2026 15:24:56 -0600
Subject: [PATCH 01/66] KN loderunner changes

---
 src/yoke/models/vit/patch_embed.py            |   3 +
 src/yoke/models/vit/swin/bomberman.py         |   3 +
 .../utils/training/datastep/loderunner.py     | 415 +++++++++++++++++-
 src/yoke/utils/training/epoch/loderunner.py   | 235 +++++++++-
 4 files changed, 642 insertions(+), 14 deletions(-)

diff --git a/src/yoke/models/vit/patch_embed.py b/src/yoke/models/vit/patch_embed.py
index e73809e3..0aa7e234 100644
--- a/src/yoke/models/vit/patch_embed.py
+++ b/src/yoke/models/vit/patch_embed.py
@@ -141,6 +141,9 @@ def forward(self, x: torch.Tensor, in_vars: torch.Tensor) -> torch.Tensor:
         groups = in_vars.shape[0]
         proj = F.conv2d(x, weights, biases, groups=groups, stride=self.patch_size)
 
+        print("proj:", proj.shape)
+        print("tokens per embed:", proj.shape[2] * proj.shape[3])
+
         # Flatten the patch arrays and separate the variables and embeddings.
         proj = rearrange(
             proj, "b (v e) h1 h2 -> b v (h1 h2) e", v=groups, e=self.embed_dim
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index e3c6af63..d30f6aff 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -187,6 +187,9 @@ def forward(
         # Aggregate variables
         x = self.agg_vars(x)
 
+        print("x before pos_embed:", x.shape)                 # expect [B, L, D]
+        print("pos_embed param:", self.pos_embed.pos_embed.shape)  # likely [1, L0, D]
+
         # Encode patch positions, spatial information
         x = self.pos_embed(x)
 
diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py
index 600996ab..8cf7a866 100644
--- a/src/yoke/utils/training/datastep/loderunner.py
+++ b/src/yoke/utils/training/datastep/loderunner.py
@@ -40,11 +40,18 @@ def train_loderunner_datastep(
     model.train()
 
     # Extract data
-    (start_img, end_img, Dt) = data
+    #(start_img, end_img, Dt) = data
+    img_seq, Dt = data
 
-    start_img = start_img.to(device, non_blocking=True)
+    #start_img = start_img.to(device, non_blocking=True)
+    #Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+    #end_img = end_img.to(device, non_blocking=True)
+
+    img_seq = img_seq.to(device, non_blocking=True)
     Dt = Dt.to(torch.float32).to(device, non_blocking=True)
-    end_img = end_img.to(device, non_blocking=True)
+
+    start_img = img_seq[:, 0]
+    end_img = img_seq[:, -1]
 
     # For our first LodeRunner training on the lsc240420 dataset the input and
     # output prediction variables are fixed.
@@ -66,6 +73,9 @@ def train_loderunner_datastep(
     # Perform a forward pass
     # NOTE: If training on GPU model should have already been moved to GPU
     # prior to initalizing optimizer.
+    print("start_img entering model:", start_img.shape)   # expect [B, 8, 1120, 800]
+    print("len(in_vars):", len(in_vars))
+    print("in_vars:", in_vars)
     pred_img = model(start_img, in_vars, out_vars, Dt)
 
     # Expecting to use a *reduction="none"* loss function so we can track loss
@@ -224,6 +234,193 @@ def train_DDP_loderunner_datastep(
     return end_img, pred_img, all_losses
 
 
+def train_DDP_loderunner_seq_channel_datastep(
+    data,
+    model,
+    optimizer,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.train()
+
+    start_img, end_img, Dt = data
+
+    start_img = start_img.to(device, non_blocking=True)
+    end_img = end_img.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    pred_img = model(start_img, in_vars, out_vars, Dt)
+
+    loss = loss_fn(pred_img, end_img)
+    per_sample_loss = loss.mean(dim=[1, 2, 3])
+
+    optimizer.zero_grad(set_to_none=True)
+    per_sample_loss.mean().backward()
+    optimizer.step()
+
+    return end_img, pred_img, per_sample_loss.detach()
+
+
+def train_DDP_loderunner_seq_datastep(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+    scheduled_prob: float = 1.0,
+    channel_map: list[int] | None = None,
+):
+    """
+    DDP training step for autoregressive sequence training.
+
+    Expected data:
+        img_seq, Dt = data
+
+    Shapes:
+        img_seq: [B, S, C, H, W]
+        Dt:
+            either [B] / [B, 1] for a single constant Dt reused at every step,
+            or [B, S-1] / [B, S-1, 1] for per-step lead times.
+
+    Returns:
+        gt_seq:      [B, S-1, C, H, W]
+        pred_seq:    [B, S-1, C, H, W]
+        all_losses:  concatenated per-sample losses on rank 0, else None
+    """
+    model.train()
+
+    img_seq, Dt = data
+    img_seq = img_seq.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    if channel_map is None:
+        channel_map = [0, 1, 2, 3, 4, 5, 6, 7]
+
+    in_vars = torch.tensor(channel_map, device=device)
+    out_vars = torch.tensor(channel_map, device=device)
+
+    B, S, C, H, W = img_seq.shape
+    assert S >= 2, "Sequence length must be at least 2."
+
+    pred_seq = []
+
+    # initial input is first frame
+    current_input = img_seq[:, 0]
+
+    for k in range(S - 1):
+        # support either one Dt for all steps or one Dt per step
+        if Dt.ndim == 1 or (Dt.ndim == 2 and Dt.shape[-1] == 1):
+            Dt_k = Dt
+        elif Dt.ndim == 2:
+            Dt_k = Dt[:, k].unsqueeze(-1)
+        elif Dt.ndim == 3:
+            Dt_k = Dt[:, k]
+        else:
+            raise ValueError(f"Unsupported Dt shape: {Dt.shape}")
+
+        pred_img = model(current_input, in_vars, out_vars, Dt_k)
+        pred_seq.append(pred_img)
+
+        if k < S - 2:
+            if random.random() < scheduled_prob:
+                current_input = img_seq[:, k + 1]   # teacher forcing
+            else:
+                current_input = pred_img.detach()   # autoregressive rollout
+
+    pred_seq = torch.stack(pred_seq, dim=1)   # [B, S-1, C, H, W]
+    gt_seq = img_seq[:, 1:]                   # [B, S-1, C, H, W]
+
+    loss = loss_fn(pred_seq, gt_seq)
+    per_sample_loss = loss.mean(dim=[1, 2, 3, 4])
+
+    optimizer.zero_grad(set_to_none=True)
+    loss.mean().backward()
+    optimizer.step()
+
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)
+    else:
+        all_losses = None
+
+    return gt_seq, pred_seq, all_losses
+
+
+def train_DDP_loderunner_datastep_seq_old(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """A DDP-compatible training step for multi-input, multi-output data.
+
+    Args:
+        data (tuple): tuple of model input, corresponding ground truth, and lead time
+        model (loaded pytorch model): model to train
+        optimizer (torch.optim): optimizer for training set
+        loss_fn (torch.nn Loss Function): loss function for training set
+        device (torch.device): device index to select
+        rank (int): Rank of device
+        world_size (int): Number of total DDP processes
+
+    Returns:
+        end_img (torch.Tensor): Ground truth end image
+        pred_img (torch.Tensor): Predicted end image
+        all_losses (torch.Tensor): Concatenated per-sample losses from all processes
+    """
+    # Set model to train mode
+    model.train()
+
+    # Extract data
+    #start_img, end_img, Dt = data
+    img_seq, Dt = data
+    #for img in img_seq:
+    #    # ...
+    start_img = start_img.to(device, non_blocking=True)
+    Dt = Dt.to(device, non_blocking=True)
+    end_img = end_img.to(device, non_blocking=True)
+
+    # Fixed input and output variable indices
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
+
+    # Forward pass
+    pred_img = model(start_img, in_vars, out_vars, Dt)
+
+    # Compute loss
+    loss = loss_fn(pred_img, end_img)
+    per_sample_loss = loss.mean(dim=[1, 2, 3])  # Per-sample loss
+
+    # Backward pass and optimization
+    optimizer.zero_grad(set_to_none=True)
+    loss.mean().backward()
+    optimizer.step()
+
+    # Gather per-sample losses from all processes
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    # Rank 0 concatenates and saves or returns all losses
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)  # Shape: (total_batch_size,)
+    else:
+        all_losses = None
+
+    return end_img, pred_img, all_losses
+
+
 ####################################
 # Evaluating on a Datastep
 ####################################
@@ -294,6 +491,87 @@ def eval_loderunner_datastep(
     return end_img, pred_img, per_sample_loss
 
 
+def eval_DDP_loderunner_seq_channel_datastep(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.eval()
+
+    start_img, end_img, Dt = data
+
+    start_img = start_img.to(device, non_blocking=True)
+    end_img = end_img.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    with torch.no_grad():
+        pred_img = model(start_img, in_vars, out_vars, Dt)
+
+    loss = loss_fn(pred_img, end_img)
+    per_sample_loss = loss.mean(dim=[1, 2, 3])
+
+    # No manual all_gather during eval.
+    return end_img, pred_img, per_sample_loss.detach()
+
+
+def eval_DDP_loderunner_seq_datastep(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.eval()
+
+    img_seq, Dt = data
+    img_seq = img_seq.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    in_vars = torch.arange(img_seq.shape[2], device=device)
+    out_vars = torch.arange(img_seq.shape[2], device=device)
+
+    B, S, C, H, W = img_seq.shape
+    current_input = img_seq[:, 0]
+    preds = []
+
+    for k in range(S - 1):
+        if Dt.ndim == 1:
+            Dt_k = Dt
+        elif Dt.ndim == 2:
+            Dt_k = Dt[:, k]
+        else:
+            Dt_k = Dt[:, k].squeeze(-1)
+
+        pred_img = model(current_input, in_vars, out_vars, Dt_k)
+        preds.append(pred_img)
+
+        if k < S - 2:
+            current_input = img_seq[:, k + 1]
+
+    pred_seq = torch.stack(preds, dim=1)
+    gt_seq = img_seq[:, 1:]
+
+    loss = loss_fn(pred_seq, gt_seq)
+    per_sample_loss = loss.mean(dim=[1, 2, 3, 4])
+
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)
+    else:
+        all_losses = None
+
+    return gt_seq, pred_seq, all_losses
+
+
 def eval_scheduled_loderunner_datastep(
     data: tuple,
     model: torch.nn.Module,
@@ -419,3 +697,134 @@ def eval_DDP_loderunner_datastep(
         all_losses = None
 
     return end_img, pred_img, all_losses
+
+
+def train_DDP_temporal_loderunner_datastep(
+    data,
+    model,
+    optimizer,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.train()
+
+    context_seq, target_img, Dt = data
+
+    context_seq = context_seq.to(device, non_blocking=True)   # [B, K, C, H, W]
+    target_img = target_img.to(device, non_blocking=True)     # [B, C, H, W]
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    C = target_img.shape[1]
+    in_vars = torch.arange(C, device=device)
+    out_vars = torch.arange(C, device=device)
+
+    pred_img = model(context_seq, in_vars, out_vars, Dt)
+
+    loss = loss_fn(pred_img, target_img)              # [B, C, H, W] if reduction="none"
+    per_sample_loss = loss.mean(dim=[1, 2, 3])
+
+    optimizer.zero_grad(set_to_none=True)
+    per_sample_loss.mean().backward()
+    optimizer.step()
+
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)
+    else:
+        all_losses = None
+
+    return target_img, pred_img, all_losses
+
+
+def eval_DDP_temporal_loderunner_datastep(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.eval()
+
+    context_seq, target_img, Dt = data
+
+    context_seq = context_seq.to(device, non_blocking=True)
+    target_img = target_img.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    C = target_img.shape[1]
+    in_vars = torch.arange(C, device=device)
+    out_vars = torch.arange(C, device=device)
+
+    pred_img = model(context_seq, in_vars, out_vars, Dt)
+
+    loss = loss_fn(pred_img, target_img)
+    per_sample_loss = loss.mean(dim=[1, 2, 3])
+
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)
+    else:
+        all_losses = None
+
+    return target_img, pred_img, all_losses
+
+
+def eval_DDP_loderunner_seq_context_datastep(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    """
+    DDP eval step for TemporalLodeRunner / channel-stacked context model.
+
+    Expected data:
+        context_seq, target_img, Dt = data
+
+    Shapes:
+        context_seq: [B, K, C, H, W]
+        target_img:  [B, C, H, W]
+        Dt:          [B] or [B, 1]
+    """
+    import torch
+    import torch.distributed as dist
+
+    model.eval()
+
+    context_seq, target_img, Dt = data
+
+    context_seq = context_seq.to(device, non_blocking=True)
+    target_img = target_img.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    C = target_img.shape[1]
+    in_vars = torch.arange(C, device=device)
+    out_vars = torch.arange(C, device=device)
+
+    pred_img = model(context_seq, in_vars, out_vars, Dt)
+
+    loss = loss_fn(pred_img, target_img)
+
+    if loss.ndim == 0:
+        per_sample_loss = loss.repeat(target_img.shape[0])
+    else:
+        per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
+
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)
+    else:
+        all_losses = None
+
+    return target_img, pred_img, all_losses
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 26629e4e..9d26af3e 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -11,7 +11,14 @@
     train_scheduled_loderunner_datastep,
     eval_scheduled_loderunner_datastep,
     train_DDP_loderunner_datastep,
+    train_DDP_loderunner_seq_datastep,
+    train_DDP_loderunner_seq_channel_datastep,
+    train_DDP_temporal_loderunner_datastep,
     eval_DDP_loderunner_datastep,
+    eval_DDP_loderunner_seq_datastep,
+    eval_DDP_loderunner_seq_context_datastep,
+    eval_DDP_loderunner_seq_channel_datastep,
+
 )
 
 
@@ -346,6 +353,7 @@ def train_DDP_loderunner_epoch(
     device: torch.device,
     rank: int,
     world_size: int,
+    seq: bool = False
 ) -> None:
     """Distributed data-parallel LodeRunner Epoch.
 
@@ -387,10 +395,22 @@ def train_DDP_loderunner_epoch(
             if trainbatch_ID >= num_train_batches:
                 break
 
-            # Perform a single training step
-            truth, pred, train_losses = train_DDP_loderunner_datastep(
-                traindata, model, optimizer, loss_fn, device, rank, world_size
-            )
+            if seq:  #all_losses??
+                gt_seq, pred_seq, train_losses = train_DDP_loderunner_seq_datastep(
+                    data=traindata,
+                    model=model,
+                    optimizer=optimizer,
+                    loss_fn=loss_fn,
+                    device=device,
+                    rank=rank,
+                    world_size=world_size,
+                    scheduled_prob=1.0,   # start with pure teacher forcing
+                )
+            else:
+                # Perform a single training step
+                truth, pred, train_losses = train_DDP_loderunner_seq_channel_datastep(
+                    traindata, model, optimizer, loss_fn, device, rank, world_size
+                )
 
             # Increment the learning-rate scheduler
             LRsched.step()
@@ -420,15 +440,207 @@ def train_DDP_loderunner_epoch(
                     if valbatch_ID >= num_val_batches:
                         break
 
-                    end_img, pred_img, val_losses = eval_DDP_loderunner_datastep(
-                        valdata,
-                        model,
-                        loss_fn,
-                        device,
-                        rank,
-                        world_size,
+                    if seq:  #all_losses??
+                        gt_seq, pred_seq, val_losses = eval_DDP_loderunner_seq_datastep(
+                            data=valdata,
+                            model=model,
+                            loss_fn=loss_fn,
+                            device=device,
+                            rank=rank,
+                            world_size=world_size,
+                        )
+                    else:
+                        # Perform a single training step
+                        end_img, pred_img, val_losses = eval_DDP_loderunner_seq_channel_datastep(
+                            valdata, model, loss_fn, device, rank, world_size
+                        )
+
+
+                    # Save validation record (rank 0 only)
+                    if rank == 0:
+                        batch_records = np.column_stack(
+                            [
+                                np.full(len(val_losses), epochIDX),
+                                np.full(len(val_losses), valbatch_ID),
+                                val_losses.cpu().numpy().flatten(),
+                            ]
+                        )
+                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+
+def train_DDP_loderunner_epoch_seq_context(
+    training_data: torch.utils.data.DataLoader,
+    validation_data: torch.utils.data.DataLoader,
+    num_train_batches: int,
+    num_val_batches: int,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    LRsched: torch.optim.lr_scheduler._LRScheduler,
+    epochIDX: int,
+    train_per_val: int,
+    train_rcrd_filename: str,
+    val_rcrd_filename: str,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+    seq: bool = False
+) -> None:
+    """Distributed data-parallel LodeRunner Epoch.
+
+    Function to complete a training epoch on the LodeRunner architecture with
+    fixed channels in the input and output. Training and validation information
+    is saved to successive CSV files.
+
+    Args:
+        training_data (torch.utils.data.DataLoader): training dataloader
+        validation_data (torch.utils.data.DataLoader): validation dataloader
+        num_train_batches (int): Number of batches in training epoch
+        num_val_batches (int): Number of batches in validation epoch
+        model (torch.nn.Module): model to train
+        optimizer (torch.optim.Optimizer): optimizer for training set
+        loss_fn (torch.nn.Module): loss function for training set
+        LRsched (torch.optim.lr_scheduler._LRScheduler): Learning-rate scheduler called
+                                                         every training step.
+        epochIDX (int): Index of current training epoch
+        train_per_val (int): Number of Training epochs between each validation
+        train_rcrd_filename (str): Name of CSV file to save training sample stats to
+        val_rcrd_filename (str): Name of CSV file to save validation sample stats to
+        device (torch.device): device index to select
+        rank (int): rank of process
+        world_size (int): number of total processes
+
+    """
+    # Initialize things to save
+    trainbatch_ID = 0
+    valbatch_ID = 0
+
+    # Training loop
+    model.train()
+    train_rcrd_filename = train_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
+    max_train_batches = min(num_train_batches, len(training_data))
+
+    with (
+        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
+    ) as train_rcrd_file:
+        for trainbatch_ID, traindata in enumerate(training_data):
+            # Stop when number of training batches is reached
+            if trainbatch_ID >= max_train_batches:
+                break
+
+            if seq:  #all_losses??
+                gt_seq, pred_seq, train_losses = train_DDP_temporal_loderunner_datastep(
+                    data=traindata,
+                    model=model,
+                    optimizer=optimizer,
+                    loss_fn=loss_fn,
+                    device=device,
+                    rank=rank,
+                    world_size=world_size,
+                    #scheduled_prob=1.0,   # start with pure teacher forcing
+                )
+            else:
+                # Perform a single training step
+                truth, pred, train_losses = train_DDP_loderunner_seq_channel_datastep(
+                    traindata, model, optimizer, loss_fn, device, rank, world_size
+                )
+
+            # Increment the learning-rate scheduler
+            LRsched.step()
+
+            # Save training record (rank 0 only)
+            if rank == 0:
+                batch_records = np.column_stack(
+                    [
+                        np.full(len(train_losses), epochIDX),
+                        np.full(len(train_losses), trainbatch_ID),
+                        train_losses.cpu().numpy().flatten(),
+                    ]
+                )
+                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+    # Validation loop
+    if epochIDX % train_per_val == 0:
+        print("Validating...", epochIDX)
+        val_rcrd_filename = val_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
+        model.eval()
+        max_val_batches = min(num_val_batches, len(val_data))
+
+        with (open(val_rcrd_filename, "a") if rank == 0 else nullcontext()) as val_rcrd_file:
+            with torch.no_grad():
+                for valbatch_ID, valdata in enumerate(validation_data):
+                    if valbatch_ID >= max_val_batches:
+                        break
+
+                    if seq:
+                        end_img, pred_img, val_losses = eval_DDP_loderunner_seq_context_datastep(
+                            data=valdata,
+                            model=model,
+                            loss_fn=loss_fn,
+                            device=device,
+                            rank=rank,
+                            world_size=world_size,
+                        )
+                    else:
+                        end_img, pred_img, val_losses = eval_DDP_loderunner_seq_channel_datastep(
+                            valdata,
+                            model,
+                            loss_fn,
+                            device,
+                            rank,
+                            world_size,
+                        )
+
+                    if rank == 0:
+                        batch_records = np.column_stack(
+                            [
+                                np.full(len(val_losses), epochIDX),
+                                np.full(len(val_losses), valbatch_ID),
+                                val_losses.cpu().numpy().flatten(),
+                            ]
+                        )
+                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+    '''
+    # Validation loop
+    if epochIDX % train_per_val == 0:
+        print("Validating...", epochIDX)
+        val_rcrd_filename = val_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
+        model.eval()
+        with (
+            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
+        ) as val_rcrd_file:
+            with torch.no_grad():
+                for valbatch_ID, valdata in enumerate(validation_data):
+                    # Stop when number of training batches is reached
+                    if valbatch_ID >= num_val_batches:
+                        break
+
+                if seq:  #all_losses??
+                    gt_seq, pred_seq, train_losses = eval_DDP_loderunner_seq_datastep(
+                        data=valdata,
+                        model=model,
+                        loss_fn=loss_fn,
+                        device=device,
+                        rank=rank,
+                        world_size=world_size,
+                    )
+                else:
+                    # Perform a single training step
+                    emd_img, pred_img, val_losses = eval_DDP_loderunner_datastep(
+                        valdata, model, loss_fn, device, rank, world_size
                     )
 
+
+                    #end_img, pred_img, val_losses = eval_DDP_loderunner_datastep(
+                    #    valdata,
+                    #    model,
+                    #    loss_fn,
+                    #    device,
+                    #    rank,
+                    #    world_size,
+                    #)
+
                     # Save validation record (rank 0 only)
                     if rank == 0:
                         batch_records = np.column_stack(
@@ -439,3 +651,4 @@ def train_DDP_loderunner_epoch(
                             ]
                         )
                         np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+        '''

From c86f9d222271c7d48d8559169e181faf49b4b3dd Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Thu, 30 Apr 2026 15:29:02 -0600
Subject: [PATCH 02/66] temporary plotting scripts

---
 .../utils/KN_tmp/plot_loss_curves_channel.py  | 123 +++
 src/yoke/utils/KN_tmp/plot_pred_channel.py    | 214 +++++
 src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py | 844 ++++++++++++++++++
 3 files changed, 1181 insertions(+)
 create mode 100644 src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
 create mode 100644 src/yoke/utils/KN_tmp/plot_pred_channel.py
 create mode 100644 src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py

diff --git a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
new file mode 100644
index 00000000..919f33d2
--- /dev/null
+++ b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
@@ -0,0 +1,123 @@
+import argparse
+import glob
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def load_records(pattern):
+    files = sorted(glob.glob(pattern))
+
+    if len(files) == 0:
+        raise FileNotFoundError(f"No files matched pattern: {pattern}")
+
+    arrays = []
+
+    for fn in files:
+        try:
+            arr = np.loadtxt(fn, delimiter=",")
+        except Exception as e:
+            print(f"Skipping {fn}: {e}")
+            continue
+
+        if arr.size == 0:
+            continue
+
+        if arr.ndim == 1:
+            arr = arr[None, :]
+
+        arrays.append(arr)
+
+    if len(arrays) == 0:
+        raise RuntimeError(f"No valid data found for pattern: {pattern}")
+
+    data = np.vstack(arrays)
+
+    # columns: epoch, batch, loss
+    epochs = data[:, 0].astype(int)
+    batches = data[:, 1].astype(int)
+    losses = data[:, 2]
+
+    return epochs, batches, losses, files
+
+
+def epoch_means(epochs, losses):
+    unique_epochs = np.array(sorted(set(epochs)))
+    mean_losses = np.array([losses[epochs == e].mean() for e in unique_epochs])
+    std_losses = np.array([losses[epochs == e].std() for e in unique_epochs])
+    return unique_epochs, mean_losses, std_losses
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+    "--train_pattern",
+    type=str,
+    #default="runs/study_010/training_study010_epoch*.csv",
+    default="runs/study_011/training_study011_epoch*.csv",
+    )
+
+    parser.add_argument(
+        "--val_pattern",
+        type=str,
+        #default="runs/study_010/validation_study010_epoch*.csv",
+        default="runs/study_011/validation_study011_epoch*.csv",
+    )
+
+    parser.add_argument(
+        "--out",
+        type=str,
+        default="loss_curves_study011.png",
+    )
+
+    parser.add_argument(
+        "--logy",
+        action="store_true",
+        default=True,
+        help="Use log scale on y-axis.",
+    )
+
+    args = parser.parse_args()
+
+    train_epochs, train_batches, train_losses, train_files = load_records(args.train_pattern)
+
+    print("Loaded training files:")
+    for f in train_files:
+        print("  ", f)
+
+    train_ep, train_mean, train_std = epoch_means(train_epochs, train_losses)
+
+    plt.figure(figsize=(8, 5))
+    plt.plot(train_ep, train_mean, marker="o", label="Train")
+
+    # Try validation, but do not fail if absent
+    try:
+        val_epochs, val_batches, val_losses, val_files = load_records(args.val_pattern)
+
+        print("Loaded validation files:")
+        for f in val_files:
+            print("  ", f)
+
+        val_ep, val_mean, val_std = epoch_means(val_epochs, val_losses)
+        plt.plot(val_ep, val_mean, marker="s", label="Validation")
+
+    except Exception as e:
+        print(f"No validation curve plotted: {e}")
+
+    plt.xlabel("Epoch")
+    plt.ylabel("Mean loss")
+    plt.title("Loss curves")
+    plt.grid(True, alpha=0.3)
+    plt.legend()
+
+    if args.logy:
+        plt.yscale("log")
+
+    plt.tight_layout()
+    plt.savefig(args.out, dpi=200)
+    print(f"Saved {args.out}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/yoke/utils/KN_tmp/plot_pred_channel.py b/src/yoke/utils/KN_tmp/plot_pred_channel.py
new file mode 100644
index 00000000..5bc72adf
--- /dev/null
+++ b/src/yoke/utils/KN_tmp/plot_pred_channel.py
@@ -0,0 +1,214 @@
+import argparse
+import numpy as np
+import torch
+import matplotlib
+import matplotlib.pyplot as plt
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from torch.utils.data import DataLoader
+
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (6, 6)
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        #default="runs/study_007/study007_modelState_epoch0100.pth",
+        #default="runs/study_010/study010_modelState_epoch0100.pth",
+        default="runs/study_011/study011_modelState_epoch0100.pth",
+
+    )
+    parser.add_argument("--N_imgs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--n_future_steps", type=int, default=10)
+
+    return parser.parse_args()
+
+
+def load_channel_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    model = LodeRunner(**model_args)
+    model.to(device)
+
+    state_dict = ckpt["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+    print("Loaded model_args:", model_args)
+
+    model.noise_scale = noise_scale
+    model.eval()
+
+    return model
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    context_len = 3
+    model = load_channel_model(args.ckpt, device)
+
+    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    # ------------------------------------------------------------
+    # One-step predictions using true context windows
+    # ------------------------------------------------------------
+    preds = []
+    targets = []
+    idxs = []
+    prefix = []
+
+    for idx, (context_img, target, Dt) in enumerate(loader):
+        context_img = context_img.to(device)
+        if idx == 0:
+            context_means = context_img.mean(dim=(2, 3))[0].detach().cpu().numpy()
+            for context in context_means:
+                prefix.append(context.mean().item())
+        target = target.to(device)
+        Dt = Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_image = model(context_img, in_vars, out_vars, Dt)
+
+        preds.append(pred_image.mean().item())
+        targets.append(target.mean().item())
+        idxs.append(idx)
+
+    plt.figure()
+    plt.scatter(idxs, preds, label="Predictions")
+    plt.scatter(idxs, targets, label="Truth")
+    plt.scatter(np.arange(len(prefix))-(len(prefix)), prefix, label='Initial Context Window')
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Sample index")
+    plt.ylabel("Mean magnitude/image value")
+    plt.tight_layout()
+    plt.savefig("pred_vs_truth_channel.png", dpi=200)
+
+
+
+    context_seq, target, Dt = next(iter(loader))
+
+    context_seq = context_seq.to(device)
+    Dt = Dt.to(torch.float32).to(device)
+
+    x = context_seq
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    future_iter = iter(loader)
+
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target = future_target.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_image = model(x, in_vars, out_vars, future_Dt)
+
+        preds_seq.append(pred_image.mean().item())
+        truth_seq.append(future_target.mean().item())
+        idxs_seq.append(step)
+
+        # autoregressive update: append prediction
+        x = torch.cat([x[:, 1:], pred_image[:, -1:].detach()], dim=1)
+
+    plt.figure()
+
+    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
+    plt.scatter(idxs_seq, truth_seq, label="Truth")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Mean magnitude/image value")
+    plt.tight_layout()
+    plt.savefig("pred_vs_truth_channel_autoreg.png", dpi=200)
+
+
+    # ------------------------------------------------------------
+    # Image comparison for the final one-step batch above
+    # ------------------------------------------------------------
+    pred_plot = pred_image.squeeze().mean(dim=0).detach().cpu().numpy()
+    true_plot = target.squeeze().mean(dim=0).detach().cpu().numpy()
+    error_plot = pred_plot - true_plot
+
+    vmin = min(pred_plot.min(), true_plot.min())
+    vmax = max(pred_plot.max(), true_plot.max())
+    err_max = np.max(np.abs(error_plot))
+
+    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
+
+    im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
+    ax1.set_title("Prediction")
+
+    ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
+    ax2.set_title("Truth")
+
+    im3 = ax3.imshow(error_plot, origin="lower", vmin=-err_max, vmax=err_max)
+    ax3.set_title("Error (Pred - Truth)")
+
+    cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
+    cbar.set_label("Field value")
+
+    cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
+    cbar_err.set_label("Error")
+
+    for ax in (ax1, ax2, ax3):
+        ax.axis("off")
+
+    plt.tight_layout()
+    plt.savefig("img_comp_channel.png", bbox_inches="tight", dpi=200)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
new file mode 100644
index 00000000..5bf82868
--- /dev/null
+++ b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
@@ -0,0 +1,844 @@
+import os
+import time
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from yoke.datasets.lsc_dataset import LSC_rho2rho_temporal_DataSet
+from yoke.utils.training.epoch.loderunner import train_DDP_loderunner_epoch
+from yoke.utils.training.epoch.loderunner import train_DDP_loderunner_epoch_seq_context
+from yoke.utils.restart import continuation_setup
+from yoke.utils.dataload import make_distributed_dataloader
+from yoke.utils.checkpointing import load_model_and_optimizer
+from yoke.utils.checkpointing import save_model_and_optimizer
+from yoke.lr_schedulers import CosineWithWarmupScheduler
+from yoke.helpers import cli
+
+# FIXME remove if restructure
+from torch.utils.data import Dataset, DataLoader, random_split
+import glob
+import random
+
+
+#############################################
+# Inputs
+#############################################
+descr_str = (
+    "Uses DDP to train LodeRunner architecture on single-timstep input and output "
+    "of the lsc240420 per-material density fields."
+)
+parser = argparse.ArgumentParser(
+    prog="DDP LodeRunner Training", description=descr_str, fromfile_prefix_chars="@"
+)
+parser = cli.add_default_args(parser=parser)
+parser = cli.add_filepath_args(parser=parser)
+parser = cli.add_computing_args(parser=parser)
+parser = cli.add_model_args(parser=parser)
+parser = cli.add_training_args(parser=parser)
+parser = cli.add_cosine_lr_scheduler_args(parser=parser)
+
+# DPOT‐style noise parameter
+parser.add_argument(
+    "--noise_scale",
+    type=float,
+    default=0.0,
+    help="Relative magnitude ε for Gaussian noise injection (e.g. 5e-5).",
+)
+
+# Change some default filepaths.
+parser.set_defaults(
+    train_filelist="lsc240420_prefixes_train_80pct.txt",
+    validation_filelist="lsc240420_prefixes_validation_10pct.txt",
+    test_filelist="lsc240420_prefixes_test_10pct.txt",
+)
+
+
+class Kilonova_lc_img_DataSet(Dataset):
+    def __init__(self, half_image=False, N_imgs=0):
+        file_prefix_list = sorted(
+            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
+
+        random.shuffle(self.file_prefix_list)
+
+        #self.max_timeIDX_offset = max_timeIDX_offset
+        self.half_image = half_image
+
+        # Build a global index: one entry per usable (file, startIDX)
+        self.samples = []
+        seqLen = 1
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+            mjd = data["arr_ztfg"][:, 0]
+            n_times = len(mjd)
+            data.close()
+
+            max_start = n_times - seqLen - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((file_idx, startIDX))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        file_idx, startIDX = self.samples[index]
+        fn = self.file_prefix_list[file_idx]
+
+        data = np.load(fn, allow_pickle=True)
+
+        mjd = data["arr_ztfg"][:, 0]
+        t0 = mjd.min()
+        t_obs = mjd - t0
+        g_mag = data["arr_ztfg"][:, 1]
+
+        seqLen = 1
+        endIDX = startIDX + seqLen
+
+        start_mag = g_mag[startIDX]
+        end_mag = g_mag[endIDX]
+        start_t = t_obs[startIDX]
+        end_t = t_obs[endIDX]
+
+        Dt = torch.tensor(end_t - start_t, dtype=torch.float32)
+
+        H, W = 1120, 400
+
+        s = torch.tensor(start_mag, dtype=torch.float32)
+        start_img = s.view(1, 1, 1).expand(8, H, W)
+
+        e = torch.tensor(end_mag, dtype=torch.float32)
+        end_img = e.view(1, 1, 1).expand(8, H, W)
+
+        data.close()
+        return start_img, end_img, Dt
+
+
+class Kilonova_lc_img_DataSet_seq(Dataset):
+    def __init__(self, half_image=False, N_imgs=0):
+        file_prefix_list = sorted(
+            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
+
+        random.shuffle(self.file_prefix_list)
+
+        #self.max_timeIDX_offset = max_timeIDX_offset
+        self.half_image = half_image
+
+        # Build a global index: one entry per usable (file, startIDX)
+        self.samples = []
+        seqLen = 3
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+            mjd = data["arr_ztfg"][:, 0]
+            n_times = len(mjd)
+            data.close()
+
+            max_start = n_times - seqLen - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((file_idx, startIDX))
+
+    def __len__(self):
+        return len(self.samples)
+
+    
+    def __getitem__(self, index):
+        file_idx, startIDX = self.samples[index]
+        fn = self.file_prefix_list[file_idx]
+
+        frames = []
+        seqLen = 3
+        H, W = 1120, 400
+
+        data = np.load(fn, allow_pickle=True)
+
+        mjd = data["arr_ztfg"][:, 0]
+        t0 = mjd.min()
+        t_obs = mjd - t0
+        g_mag = data["arr_ztfg"][:, 1]
+
+        endIDX = startIDX + seqLen
+
+        for i in range(seqLen):
+            seq_mag = g_mag[startIDX + i]
+            s = torch.tensor(seq_mag, dtype=torch.float32)
+            seq_img = s.view(1, 1, 1).expand(8, H, W)
+            frames.append(seq_img)
+
+        end_mag = g_mag[endIDX]
+        e = torch.tensor(end_mag, dtype=torch.float32)
+        end_img = e.view(1, 1, 1).expand(8, H, W)
+        frames.append(end_img)
+
+        start_t = t_obs[startIDX]
+        end_t = t_obs[endIDX]
+        Dt = torch.tensor(end_t - start_t, dtype=torch.float32)
+
+        data.close()
+
+        img_seq = torch.stack(frames, dim=0)
+        return img_seq, Dt
+
+
+class Kilonova_lc_img_DataSet_channels_context(Dataset):
+    def __init__(
+        self,
+        half_image=False,
+        N_imgs=0,
+        context_len=3,
+        H=1120,
+        W=400,
+        n_channels=8,
+    ):
+        assert context_len <= n_channels
+
+        file_prefix_list = sorted(
+            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
+
+        random.shuffle(self.file_prefix_list)
+
+        self.context_len = context_len
+        self.H = H
+        self.W = W
+        self.n_channels = n_channels
+        self.samples = []
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+            mjd = data["arr_ztfg"][:, 0]
+            n_times = len(mjd)
+            data.close()
+
+            max_start = n_times - context_len - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((file_idx, startIDX))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        file_idx, startIDX = self.samples[index]
+        fn = self.file_prefix_list[file_idx]
+
+        data = np.load(fn, allow_pickle=True)
+        arr = data["arr_ztfg"]
+
+        mjd = arr[:, 0]
+        g_mag = arr[:, 1]
+
+        t0 = mjd.min()
+        t_obs = mjd - t0
+
+        target_idx = startIDX + self.context_len
+
+        '''
+        # Input: [8, H, W], where channels encode previous timesteps.
+        # Right-align context. Unused earlier channels repeat earliest value.
+        context_img = torch.empty(self.n_channels, self.H, self.W, dtype=torch.float32)
+
+        earliest_mag = float(g_mag[startIDX])
+        context_img[:] = earliest_mag
+
+        offset = self.n_channels - self.context_len
+        for i in range(self.context_len):
+            ch = offset + i
+            context_img[ch] = float(g_mag[startIDX + i])
+
+        # Target: next scalar copied across all 8 channels.
+        target_mag = float(g_mag[target_idx])
+        target_img = torch.empty(self.n_channels, self.H, self.W, dtype=torch.float32)
+        target_img[:] = target_mag
+        '''
+
+        context_vals = torch.empty(self.n_channels, dtype=torch.float32)
+
+        earliest_mag = float(g_mag[startIDX])
+        context_vals[:] = earliest_mag
+
+        offset = self.n_channels - self.context_len
+        for i in range(self.context_len):
+            ch = offset + i
+            context_vals[ch] = float(g_mag[startIDX + i])
+
+        # expand() better for memory
+        context_img = context_vals.view(self.n_channels, 1, 1).expand(
+            self.n_channels,
+            self.H,
+            self.W,
+        )
+
+        target_mag = float(g_mag[target_idx])
+        target_val = torch.tensor(target_mag, dtype=torch.float32)
+
+        target_img = target_val.view(1, 1, 1).expand(
+            self.n_channels,
+            self.H,
+            self.W,
+        )
+
+        Dt = torch.tensor(
+            t_obs[target_idx] - t_obs[target_idx - 1],
+            dtype=torch.float32,
+        )
+
+        data.close()
+
+        return context_img, target_img, Dt
+
+
+class ChannelStackAdapter(nn.Module):
+    """
+    Converts a sequence [B, K, C, H, W] into a fused image [B, C, H, W].
+    """
+
+    def __init__(self, in_channels: int, context_len: int, hidden_channels: int = 64):
+        super().__init__()
+        self.in_channels = in_channels
+        self.context_len = context_len
+        stacked_channels = in_channels * context_len
+
+        self.adapter = nn.Sequential(
+            nn.Conv2d(stacked_channels, hidden_channels, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Conv2d(hidden_channels, in_channels, kernel_size=1),
+        )
+
+    def forward(self, x_seq: torch.Tensor) -> torch.Tensor:
+        """
+        x_seq: [B, K, C, H, W]
+        returns: [B, C, H, W]
+        """
+        if x_seq.ndim != 5:
+            raise ValueError(f"Expected x_seq to have shape [B, K, C, H, W], got {x_seq.shape}")
+
+        B, K, C, H, W = x_seq.shape
+        if C != self.in_channels:
+            raise ValueError(f"Expected {self.in_channels} channels, got {C}")
+        if K != self.context_len:
+            raise ValueError(f"Expected context_len={self.context_len}, got K={K}")
+
+        x = x_seq.reshape(B, K * C, H, W)
+        return self.adapter(x)
+
+
+class TemporalLodeRunner(nn.Module):
+    """
+    Wraps a pretrained one-step LodeRunner with a temporal adapter.
+
+    Input:
+        x_seq: [B, K, C, H, W]
+        in_vars, out_vars, Dt: same as original LodeRunner API
+
+    Output:
+        pred: [B, C, H, W]
+    """
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        in_channels: int = 8,
+        context_len: int = 3,
+        hidden_channels: int = 64,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.temporal_adapter = ChannelStackAdapter(
+            in_channels=in_channels,
+            context_len=context_len,
+            hidden_channels=hidden_channels,
+        )
+
+    def forward(
+        self,
+        x_seq: torch.Tensor,
+        in_vars: torch.Tensor,
+        out_vars: torch.Tensor,
+        Dt: torch.Tensor,
+    ) -> torch.Tensor:
+        fused_x = self.temporal_adapter(x_seq)          # [B, C, H, W]
+        pred = self.backbone(fused_x, in_vars, out_vars, Dt)
+        return pred
+
+
+def load_direct_loderunner_checkpoint(
+    checkpoint_path,
+    model_args,
+    optimizer_kwargs,
+    device,
+):
+    checkpoint_data = torch.load(
+        checkpoint_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    saved_model_args = checkpoint_data.get("model_args", model_args)
+
+    model = LodeRunner(**saved_model_args)
+    model.to(device)
+
+    state_dict = checkpoint_data["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    model.load_state_dict(state_dict, strict=True)
+
+    noise_scale = checkpoint_data.get("noise_scale", 0.0)
+    model.noise_scale = noise_scale
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        **optimizer_kwargs,
+    )
+
+    if "optimizer_state_dict" in checkpoint_data:
+        optimizer.load_state_dict(checkpoint_data["optimizer_state_dict"])
+
+        for state in optimizer.state.values():
+            for key, value in state.items():
+                if isinstance(value, torch.Tensor):
+                    state[key] = value.to(device)
+
+    starting_epoch = checkpoint_data["epoch"]
+
+    return model, optimizer, starting_epoch
+
+
+def setup_distributed():
+    # ----- 1) Basic setup & environment variables -----
+    # Rely on Slurm variables: SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, etc.
+    rank = int(os.environ["SLURM_PROCID"])  # global rank
+    world_size = int(os.environ["SLURM_NTASKS"])  # total number of processes
+    local_rank = int(os.environ["SLURM_LOCALID"])  # local rank (GPU index on this node)
+
+    master_addr = os.environ["MASTER_ADDR"]
+    master_port = os.environ["MASTER_PORT"]
+
+    # ----- 2) Set the current GPU device for this process -----
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+
+    # ----- 3) Initialize the process group -----
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://{master_addr}:{master_port}",
+        world_size=world_size,
+        rank=rank,
+    )
+
+    return rank, world_size, local_rank, device
+
+
+def cleanup_distributed():
+    # ----- 8) Clean up (optional) -----
+    dist.destroy_process_group()
+
+
+def main(args, rank, world_size, local_rank, device):
+    #############################################
+    # Process Inputs
+    #############################################
+    # Study ID
+    studyIDX = args.studyIDX
+
+    # Resources
+    Ngpus = args.Ngpus
+    Knodes = args.Knodes
+
+    # Data Paths
+    train_filelist = args.FILELIST_DIR + args.train_filelist
+    validation_filelist = args.FILELIST_DIR + args.validation_filelist
+
+    # Model Parameters
+    embed_dim = args.embed_dim
+    block_structure = tuple(args.block_structure)
+
+    # Training Parameters
+    anchor_lr = args.anchor_lr
+    num_cycles = args.num_cycles
+    min_fraction = args.min_fraction
+    terminal_steps = args.terminal_steps
+    warmup_steps = args.warmup_steps
+    noise_scale = args.noise_scale
+
+    # Number of workers controls how batches of data are prefetched and,
+    # possibly, pre-loaded onto GPUs. If the number of workers is large they
+    # will swamp memory and jobs will fail.
+    num_workers = args.num_workers
+
+    # Epoch Parameters
+    batch_size = args.batch_size
+    total_epochs = args.total_epochs
+    cycle_epochs = args.cycle_epochs
+    train_batches = args.train_batches
+    val_batches = args.val_batches
+    train_per_val = args.TRAIN_PER_VAL
+    trn_rcrd_filename = args.trn_rcrd_filename
+    val_rcrd_filename = args.val_rcrd_filename
+    CONTINUATION = args.continuation
+    checkpoint = args.checkpoint
+
+    #############################################
+    # Model Arguments for Dynamic Reconstruction
+    #############################################
+    # Dictionary of available models.
+    available_models = {
+        "LodeRunner": LodeRunner
+    }
+
+    # Model arguments for LodeRunner.
+    model_args = {
+        "default_vars": [
+            "density_case",
+            "density_cushion",
+            "density_maincharge",
+            "density_outside_air",
+            "density_striker",
+            "density_throw",
+            "Uvelocity",
+            "Wvelocity",
+        ],
+        "image_size": (1120, 400),
+        "patch_size": (10, 5),
+        "embed_dim": embed_dim,
+        "emb_factor": 2,
+        "num_heads": 8,
+        "block_structure": block_structure,
+        "window_sizes": [(8, 8), (8, 8), (4, 4), (2, 2)],
+        "patch_merge_scales": [(2, 2), (2, 2), (2, 2)],
+        #"noise_scale": noise_scale,
+    }
+
+
+    CONTEXT_LEN = 5 #3
+    HIDDEN_CHANNELS = 64
+
+    optimizer_kwargs = {
+        "lr": 1e-5,
+        "betas": (0.9, 0.999),
+        "eps": 1e-08,
+        "weight_decay": 0.01,
+    }
+
+
+    if CONTINUATION:
+        model, optimizer, starting_epoch = load_direct_loderunner_checkpoint(
+            checkpoint_path=checkpoint,
+            model_args=model_args,
+            optimizer_kwargs=optimizer_kwargs,
+            device=device,
+        )
+
+        if rank == 0:
+            print(f"Loaded direct checkpoint from {checkpoint}")
+            print(f"Continuing from epoch {starting_epoch}")
+
+        '''  # FIXME block should be unindented if uncommented 
+        if CONTINUATION:
+            model, optimizer, starting_epoch = load_model_and_optimizer(
+                checkpoint,
+                optimizer_class=torch.optim.AdamW,
+                optimizer_kwargs=optimizer_kwargs,
+                available_models=available_models,
+                device=device,
+            )
+
+            if rank == 0:
+                print(f"Loaded temporal checkpoint from {checkpoint}")
+                print(f"Continuing from epoch {starting_epoch}")
+        '''
+
+    else:
+        starting_epoch = 0
+
+        model = LodeRunner(**model_args)
+        model.to(device)
+
+        manual_checkpoint = "/usr/projects/artimis/mpmm/pretrained_models/ddp_ldr_prod_250721/study005_modelState_epoch0100.pth"
+
+        checkpoint_data = torch.load(
+            manual_checkpoint,
+            map_location=device,
+            weights_only=False,
+        )
+
+        state_dict = checkpoint_data["model_state_dict"]
+
+        if all(k.startswith("module.") for k in state_dict.keys()):
+            state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+        missing_keys, unexpected_keys = model.load_state_dict(
+            state_dict,
+            strict=False,
+        )
+
+        if rank == 0:
+            print("Loaded pretrained backbone weights.")
+            print("Missing keys:", missing_keys)
+            print("Unexpected keys:", unexpected_keys)
+
+        model.noise_scale = noise_scale
+
+        # End-to-end fine-tuning: train adapter + backbone
+        for p in model.parameters():
+            p.requires_grad = True
+
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            **optimizer_kwargs,
+        )
+
+    loss_fn = nn.MSELoss(reduction="none")
+
+    model = DDP(model, device_ids=[local_rank], output_device=local_rank)
+
+    #############################################
+    # Learning Rate Scheduler
+    #############################################
+    print("Starting epoch: ", starting_epoch)
+    if starting_epoch == 0:
+        last_epoch = -1
+    else:
+        last_epoch = train_batches * (starting_epoch - 1)
+
+    # Scale the anchor LR by global batchsize
+    #
+    # # For multi-node
+    lr_scale = np.sqrt(float(Ngpus) * float(Knodes) * float(batch_size))
+    original_batchsize = 40.0  # 1 node, 4 gpus, 10 samples/gpu
+    ddp_anchor_lr = anchor_lr * lr_scale / original_batchsize
+    #
+    # For single node
+    # ddp_anchor_lr = anchor_lr
+
+    LRsched = CosineWithWarmupScheduler(
+        optimizer,
+        anchor_lr=ddp_anchor_lr,
+        terminal_steps=terminal_steps,
+        warmup_steps=warmup_steps,
+        num_cycles=num_cycles,
+        min_fraction=min_fraction,
+        last_epoch=last_epoch,
+    )
+
+    #############################################
+    # Data Initialization (Distributed Dataloader)
+    #############################################
+    #train_dataset = LSC_rho2rho_temporal_DataSet(
+    #    args.LSC_NPZ_DIR,
+    #    file_prefix_list=train_filelist,
+    #    max_timeIDX_offset=2,
+    #    max_file_checks=10,
+    #    half_image=True,
+    #)
+    #val_dataset = LSC_rho2rho_temporal_DataSet(
+    #    args.LSC_NPZ_DIR,
+    #    file_prefix_list=validation_filelist,
+    #    max_timeIDX_offset=2,
+    #    max_file_checks=10,
+    #    half_image=True,
+    #)
+
+    '''
+    train_dataset = Kilonova_lc_img_DataSet_seq(
+        half_image=False,
+    )
+    val_dataset = Kilonova_lc_img_DataSet_seq(
+        half_image=False,
+    )
+    '''
+
+    train_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        context_len=CONTEXT_LEN,
+        #N_imgs=100,
+    )
+
+    val_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        context_len=CONTEXT_LEN,
+        #N_imgs=20, #100,
+    )
+
+    # NOTE: For DDP the batch_size is the per-GPU batch_size!!!
+    train_dataloader = make_distributed_dataloader(
+        train_dataset,
+        batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        rank=rank,
+        world_size=world_size,
+    )
+    val_dataloader = make_distributed_dataloader(
+        val_dataset,
+        batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    #############################################
+    # Training Loop (Modified for DDP)
+    #############################################
+    # Train Model
+    print("Training Model . . .")
+    starting_epoch += 1
+    ending_epoch = min(starting_epoch + cycle_epochs, total_epochs + 1)
+
+    TIME_EPOCH = True
+    for epochIDX in range(starting_epoch, ending_epoch):
+        print('%%%%%%%%%%%%%')
+        print(epochIDX)
+        print('%%%%%%%%%%%%%')
+        train_sampler = train_dataloader.sampler
+        train_sampler.set_epoch(epochIDX)
+
+        # For timing epochs
+        if TIME_EPOCH:
+            # Synchronize before starting the timer
+            #dist.barrier()  # Ensure that all nodes sync
+            torch.cuda.synchronize(device)  # Ensure GPUs on each node sync
+            # Time each epoch and print to stdout
+            startTime = time.time()
+
+
+        train_DDP_loderunner_epoch(
+            training_data=train_dataloader,
+            validation_data=val_dataloader,
+            num_train_batches=train_batches,
+            num_val_batches=val_batches,
+            model=model,
+            optimizer=optimizer,
+            loss_fn=loss_fn,
+            LRsched=LRsched,
+            epochIDX=epochIDX,
+            train_per_val=train_per_val,
+            train_rcrd_filename=trn_rcrd_filename,
+            val_rcrd_filename=val_rcrd_filename,
+            device=device,
+            rank=rank,
+            world_size=world_size,
+            seq=False,
+        )
+
+        print(f"[rank {rank}] finished epoch", flush=True)
+
+
+        if TIME_EPOCH:
+            # Synchronize before stopping the timer
+            torch.cuda.synchronize(device)  # Ensure GPUs on each node sync
+            #dist.barrier()  # Ensure that all nodes sync
+            # Time each epoch and print to stdout
+            endTime = time.time()
+
+        epoch_time = (endTime - startTime) / 60
+
+        # Print Summary Results
+        if rank == 0:
+            print(f"Completed epoch {epochIDX}...", flush=True)
+            print(f"Epoch time (minutes): {epoch_time:.2f}", flush=True)
+
+        # Save model and optimizer
+        #chkpt_name_str = f'study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth'
+        #new_chkpt_path = os.path.join("./", chkpt_name_str)
+
+        if rank == 0:
+            chkpt_name_str = f"study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth"
+            new_chkpt_path = os.path.join("./", chkpt_name_str)
+
+            print(f"Saving checkpoint: {new_chkpt_path}", flush=True)
+
+            torch.save(
+                {
+                    "epoch": epochIDX,
+                    "model_class": "LodeRunner",
+                    "model_args": model_args,
+                    "model_state_dict": model.module.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "noise_scale": noise_scale,
+                },
+                new_chkpt_path,
+            )
+
+            '''
+            save_model_and_optimizer(
+                model.module,
+                optimizer,
+                epochIDX,
+                new_chkpt_path,
+                model_class=LodeRunner,
+                model_args=model_args,
+            )
+            '''
+
+            print(f"Saved checkpoint: {new_chkpt_path}", flush=True)
+
+        '''
+        if rank == 0:
+            chkpt_name_str = f"study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth"
+            new_chkpt_path = os.path.join("./", chkpt_name_str)
+
+            #save_model_and_optimizer(
+            #    model,
+            #    optimizer,
+            #    epochIDX,
+            #    new_chkpt_path,
+            #    model_class=LodeRunner,
+            #    model_args=model_args,
+            #)
+
+            print(f"Saved checkpoint: {new_chkpt_path}", flush=True)
+        '''
+    '''
+        save_model_and_optimizer(
+            model,
+            optimizer,
+            epochIDX,
+            new_chkpt_path,
+            model_class=LodeRunner,
+            model_args=model_args,
+        )
+    '''
+    if rank == 0:
+        #############################################
+        # Continue if Necessary
+        #############################################
+        FINISHED_TRAINING = epochIDX + 1 > total_epochs
+        if not FINISHED_TRAINING:
+            new_slurm_file = continuation_setup(
+                new_chkpt_path, studyIDX, last_epoch=epochIDX
+            )
+            os.system(f"sbatch {new_slurm_file}")
+
+if __name__ == "__main__":
+    print('running main')
+    args = parser.parse_args()
+
+    rank, world_size, local_rank, device = setup_distributed()
+
+    main(args, rank, world_size, local_rank, device)
+
+    cleanup_distributed()

From 60da54aabb249cddd682a02e72e3a1239e213422 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Tue, 5 May 2026 09:38:48 -0600
Subject: [PATCH 03/66] training script

---
 src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
index 5bf82868..7f8d1398 100644
--- a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
+++ b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
@@ -23,6 +23,12 @@
 import glob
 import random
 
+#MEAN = 24.694652705328807
+#STD = 4.67030961432848
+
+GLOBAL_GMAG_MEAN = 24.694652705328807
+GLOBAL_GMAG_STD  = 4.67030961432848
+EPS = 1e-6
 
 #############################################
 # Inputs
@@ -244,8 +250,14 @@ def __getitem__(self, index):
         data = np.load(fn, allow_pickle=True)
         arr = data["arr_ztfg"]
 
+        #mjd = arr[:, 0]
+        #g_mag = arr[:, 1]
+
         mjd = arr[:, 0]
-        g_mag = arr[:, 1]
+        g_mag = arr[:, 1].astype(np.float32)
+
+        # GLOBAL NORMALIZATION
+        g_mag = (g_mag - GLOBAL_GMAG_MEAN) / (GLOBAL_GMAG_STD + EPS)
 
         t0 = mjd.min()
         t_obs = mjd - t0

From 8936af2f0693a2a95f2ade51aa5d7f9285fcf68a Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Tue, 5 May 2026 09:40:56 -0600
Subject: [PATCH 04/66] plotting scripts

---
 src/yoke/utils/KN_tmp/plot_loss_curves_channel.py |  6 +++---
 src/yoke/utils/KN_tmp/plot_pred_channel.py        | 10 +++++-----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
index 919f33d2..3ea8f1e1 100644
--- a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
+++ b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
@@ -55,20 +55,20 @@ def main():
     "--train_pattern",
     type=str,
     #default="runs/study_010/training_study010_epoch*.csv",
-    default="runs/study_011/training_study011_epoch*.csv",
+    default="runs/study_012/training_study012_epoch*.csv",
     )
 
     parser.add_argument(
         "--val_pattern",
         type=str,
         #default="runs/study_010/validation_study010_epoch*.csv",
-        default="runs/study_011/validation_study011_epoch*.csv",
+        default="runs/study_012/validation_study012_epoch*.csv",
     )
 
     parser.add_argument(
         "--out",
         type=str,
-        default="loss_curves_study011.png",
+        default="loss_curves_study012.png",
     )
 
     parser.add_argument(
diff --git a/src/yoke/utils/KN_tmp/plot_pred_channel.py b/src/yoke/utils/KN_tmp/plot_pred_channel.py
index 5bc72adf..d126bbcb 100644
--- a/src/yoke/utils/KN_tmp/plot_pred_channel.py
+++ b/src/yoke/utils/KN_tmp/plot_pred_channel.py
@@ -23,7 +23,7 @@ def get_args():
         type=str,
         #default="runs/study_007/study007_modelState_epoch0100.pth",
         #default="runs/study_010/study010_modelState_epoch0100.pth",
-        default="runs/study_011/study011_modelState_epoch0100.pth",
+        default="runs/study_012/study012_modelState_epoch0100.pth",
 
     )
     parser.add_argument("--N_imgs", type=int, default=1)
@@ -70,7 +70,7 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print("Using device:", device)
 
-    context_len = 3
+    context_len = 5
     model = load_channel_model(args.ckpt, device)
 
     eval_dataset = Kilonova_lc_img_DataSet_channels_context(
@@ -121,7 +121,7 @@ def main():
     plt.xlabel("Sample index")
     plt.ylabel("Mean magnitude/image value")
     plt.tight_layout()
-    plt.savefig("pred_vs_truth_channel.png", dpi=200)
+    plt.savefig("pred_vs_truth_channel_norm.png", dpi=200)
 
 
 
@@ -172,7 +172,7 @@ def main():
     plt.xlabel("Autoregressive step")
     plt.ylabel("Mean magnitude/image value")
     plt.tight_layout()
-    plt.savefig("pred_vs_truth_channel_autoreg.png", dpi=200)
+    plt.savefig("pred_vs_truth_channel_norm_autoreg.png", dpi=200)
 
 
     # ------------------------------------------------------------
@@ -207,7 +207,7 @@ def main():
         ax.axis("off")
 
     plt.tight_layout()
-    plt.savefig("img_comp_channel.png", bbox_inches="tight", dpi=200)
+    plt.savefig("img_comp_channel_norm.png", bbox_inches="tight", dpi=200)
 
 
 if __name__ == "__main__":

From 87532c7efc3d3614267fa80cad08cfe422fc370e Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Fri, 8 May 2026 11:29:55 -0600
Subject: [PATCH 05/66] wip

---
 .../utils/KN_tmp/plot_loss_curves_channel.py  |   6 +-
 .../utils/KN_tmp/plot_pred_channel_delta.py   | 254 ++++++++++++++++++
 src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py |   9 +-
 3 files changed, 264 insertions(+), 5 deletions(-)
 create mode 100644 src/yoke/utils/KN_tmp/plot_pred_channel_delta.py

diff --git a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
index 3ea8f1e1..008d95b1 100644
--- a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
+++ b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
@@ -55,20 +55,20 @@ def main():
     "--train_pattern",
     type=str,
     #default="runs/study_010/training_study010_epoch*.csv",
-    default="runs/study_012/training_study012_epoch*.csv",
+    default="runs/study_013/training_study013_epoch*.csv",
     )
 
     parser.add_argument(
         "--val_pattern",
         type=str,
         #default="runs/study_010/validation_study010_epoch*.csv",
-        default="runs/study_012/validation_study012_epoch*.csv",
+        default="runs/study_013/validation_study013_epoch*.csv",
     )
 
     parser.add_argument(
         "--out",
         type=str,
-        default="loss_curves_study012.png",
+        default="loss_curves_study013.png",
     )
 
     parser.add_argument(
diff --git a/src/yoke/utils/KN_tmp/plot_pred_channel_delta.py b/src/yoke/utils/KN_tmp/plot_pred_channel_delta.py
new file mode 100644
index 00000000..bf1874a0
--- /dev/null
+++ b/src/yoke/utils/KN_tmp/plot_pred_channel_delta.py
@@ -0,0 +1,254 @@
+import argparse
+import numpy as np
+import torch
+import matplotlib
+import matplotlib.pyplot as plt
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from torch.utils.data import DataLoader
+
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (6, 6)
+
+# ============================================================
+# RUN IDENTIFIER
+# ============================================================
+RUN_ID = "013"
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default=f"runs/study_{RUN_ID}/study{RUN_ID}_modelState_epoch0100.pth",
+    )
+    parser.add_argument("--N_imgs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--n_future_steps", type=int, default=15)
+
+    return parser.parse_args()
+
+
+def load_channel_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+    context_len = ckpt.get("context_len", 5)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("predicts_delta:", ckpt.get("predicts_delta", False))
+    print("target_type:", ckpt.get("target_type", "absolute"))
+    print("context_len:", context_len)
+
+    model = LodeRunner(**model_args)
+    model.to(device)
+
+    state_dict = ckpt["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.noise_scale = noise_scale
+    model.eval()
+
+    return model, context_len
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_channel_model(args.ckpt, device)
+
+    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    # ------------------------------------------------------------
+    # One-step predictions
+    # ------------------------------------------------------------
+    preds = []
+    targets = []
+    idxs = []
+    prefix = []
+
+    for idx, (context_img, target_delta, Dt) in enumerate(loader):
+        context_img = context_img.to(device)
+        target_delta = target_delta.to(device)
+        Dt = Dt.to(torch.float32).to(device)
+
+        if idx == 0:
+            context_means = context_img.mean(dim=(2, 3))[0].detach().cpu().numpy()
+            for context in context_means:
+                prefix.append(context.mean().item())
+
+        with torch.no_grad():
+            pred_delta_img = model(context_img, in_vars, out_vars, Dt)
+
+        last_mag_img = context_img[:, -1:]
+
+        pred_next_img = last_mag_img + pred_delta_img[:, -1:]
+        true_next_img = last_mag_img + target_delta[:, -1:]
+
+        preds.append(pred_next_img.mean().item())
+        targets.append(true_next_img.mean().item())
+        idxs.append(idx)
+
+    plt.figure()
+    plt.scatter(idxs, preds, label="Predicted next magnitude")
+    plt.scatter(idxs, targets, label="True next magnitude")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Sample index")
+    plt.ylabel("Normalized magnitude")
+    plt.tight_layout()
+
+    plt.savefig(
+        f"study{RUN_ID}_pred_vs_truth_channel_delta_onestep.png",
+        dpi=200,
+    )
+
+    # ------------------------------------------------------------
+    # Autoregressive rollout
+    # ------------------------------------------------------------
+    context_seq, target_delta, Dt = next(iter(loader))
+
+    x = context_seq.to(device)
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    future_iter = iter(loader)
+
+    pred_next_img = None
+    true_next_img = None
+
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target_delta, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target_delta = future_target_delta.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_delta_img = model(x, in_vars, out_vars, future_Dt)
+
+        last_mag_img = x[:, -1:]
+
+        pred_next_img = last_mag_img + pred_delta_img[:, -1:]
+        true_next_img = last_mag_img + future_target_delta[:, -1:]
+
+        preds_seq.append(pred_next_img.mean().item())
+        truth_seq.append(true_next_img.mean().item())
+        idxs_seq.append(step)
+
+        # Append predicted absolute next magnitude
+        x = torch.cat([x[:, 1:], pred_next_img.detach()], dim=1)
+
+    plt.figure()
+
+    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
+    plt.scatter(idxs_seq, truth_seq, label="Truth")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Normalized magnitude")
+    plt.tight_layout()
+
+    plt.savefig(
+        f"study{RUN_ID}_pred_vs_truth_channel_delta_autoreg.png",
+        dpi=200,
+    )
+
+    # ------------------------------------------------------------
+    # Image comparison
+    # ------------------------------------------------------------
+    if pred_next_img is not None and true_next_img is not None:
+        pred_plot = pred_next_img[0, 0].detach().cpu().numpy()
+        true_plot = true_next_img[0, 0].detach().cpu().numpy()
+        error_plot = pred_plot - true_plot
+
+        vmin = min(pred_plot.min(), true_plot.min())
+        vmax = max(pred_plot.max(), true_plot.max())
+        err_max = np.max(np.abs(error_plot))
+
+        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
+
+        im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
+        ax1.set_title("Prediction")
+
+        ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
+        ax2.set_title("Truth")
+
+        im3 = ax3.imshow(
+            error_plot,
+            origin="lower",
+            vmin=-err_max,
+            vmax=err_max,
+        )
+        ax3.set_title("Error (Pred - Truth)")
+
+        cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
+        cbar.set_label("Normalized magnitude")
+
+        cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
+        cbar_err.set_label("Error")
+
+        for ax in (ax1, ax2, ax3):
+            ax.axis("off")
+
+        plt.tight_layout()
+
+        plt.savefig(
+            f"study{RUN_ID}_img_comp_channel_delta.png",
+            bbox_inches="tight",
+            dpi=200,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
index 7f8d1398..ffb79712 100644
--- a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
+++ b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
@@ -300,8 +300,10 @@ def __getitem__(self, index):
             self.W,
         )
 
-        target_mag = float(g_mag[target_idx])
-        target_val = torch.tensor(target_mag, dtype=torch.float32)
+        # Predict DELTA (next - previous) rather than absolute next value
+        prev_mag = float(g_mag[target_idx - 1])
+        delta_mag = float(g_mag[target_idx] - g_mag[target_idx - 1])
+        target_val = torch.tensor(delta_mag, dtype=torch.float32)
 
         target_img = target_val.view(1, 1, 1).expand(
             self.n_channels,
@@ -791,6 +793,9 @@ def main(args, rank, world_size, local_rank, device):
                     "model_state_dict": model.module.state_dict(),
                     "optimizer_state_dict": optimizer.state_dict(),
                     "noise_scale": noise_scale,
+                    "predicts_delta": True,
+                    "target_type": "delta",
+                    "context_len": CONTEXT_LEN,
                 },
                 new_chkpt_path,
             )

From 2145fa6649ba2122eddd966a32288b7171c0f552 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Fri, 8 May 2026 11:32:37 -0600
Subject: [PATCH 06/66] remove print statements

---
 src/yoke/models/vit/patch_embed.py    | 3 ---
 src/yoke/models/vit/swin/bomberman.py | 3 ---
 2 files changed, 6 deletions(-)

diff --git a/src/yoke/models/vit/patch_embed.py b/src/yoke/models/vit/patch_embed.py
index 0aa7e234..e73809e3 100644
--- a/src/yoke/models/vit/patch_embed.py
+++ b/src/yoke/models/vit/patch_embed.py
@@ -141,9 +141,6 @@ def forward(self, x: torch.Tensor, in_vars: torch.Tensor) -> torch.Tensor:
         groups = in_vars.shape[0]
         proj = F.conv2d(x, weights, biases, groups=groups, stride=self.patch_size)
 
-        print("proj:", proj.shape)
-        print("tokens per embed:", proj.shape[2] * proj.shape[3])
-
         # Flatten the patch arrays and separate the variables and embeddings.
         proj = rearrange(
             proj, "b (v e) h1 h2 -> b v (h1 h2) e", v=groups, e=self.embed_dim
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index d30f6aff..e3c6af63 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -187,9 +187,6 @@ def forward(
         # Aggregate variables
         x = self.agg_vars(x)
 
-        print("x before pos_embed:", x.shape)                 # expect [B, L, D]
-        print("pos_embed param:", self.pos_embed.pos_embed.shape)  # likely [1, L0, D]
-
         # Encode patch positions, spatial information
         x = self.pos_embed(x)
 

From c138fcf68b46e5960b59ab93331bfd570f6719d2 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Tue, 9 Jun 2026 13:33:04 -0600
Subject: [PATCH 07/66] wip, wokring temporal adapter

---
 applications/evaluation/TandV.input           |  6 +-
 applications/evaluation/TandVplot.py          |  6 +-
 .../evaluation/img_compare_loderunner.input   |  4 +-
 .../utils/training/datastep/loderunner.py     | 81 +++++++++++++++++
 src/yoke/utils/training/epoch/loderunner.py   | 89 ++++++++++++++++++-
 5 files changed, 177 insertions(+), 9 deletions(-)

diff --git a/applications/evaluation/TandV.input b/applications/evaluation/TandV.input
index a45fbbf2..76629c3f 100644
--- a/applications/evaluation/TandV.input
+++ b/applications/evaluation/TandV.input
@@ -1,7 +1,7 @@
 --basedir
-/usr/projects/artimis/mpmm/hickmank/chicoma/ch_yoke/applications/harnesses/ch_DDP_loderunner/ddp_lrstudy1/
+/net/sescratch1/atoivonen/projects/yoke_runs/KN_loderunner_fine_tune/runs
 --IDX
-10
+5
 --Nsamps_per_trn_pt
 200
 --Nsamps_per_val_pt
@@ -9,5 +9,5 @@
 --ylim
 0.25
 --savedir
-./ch_ddp_lrstudy1/
+./KN_se_fine_tune/
 --savefig
diff --git a/applications/evaluation/TandVplot.py b/applications/evaluation/TandVplot.py
index b249efbb..ffb19ea8 100644
--- a/applications/evaluation/TandVplot.py
+++ b/applications/evaluation/TandVplot.py
@@ -153,8 +153,8 @@
 # trn_idxlist = trn_DF.index.values
 # val_idxlist = val_DF.index.values
 
-trn_csv_list = ["train.csv"]
-val_csv_list = ["valcsv"]
+#trn_csv_list = [] #["train.csv"]
+#val_csv_list = [] #["valcsv"]
 
 # Plot loss for training over all steps and epochs
 fig1 = plt.figure(num=1, figsize=(6, 6))
@@ -275,5 +275,5 @@
     plt.figure(fig1.number)
     filenameA = f"{savedir}/study{IDX:03d}_TandV_curve.png"
     plt.savefig(filenameA, bbox_inches="tight")
-else:
+    #else:
     plt.show()
diff --git a/applications/evaluation/img_compare_loderunner.input b/applications/evaluation/img_compare_loderunner.input
index 097039f3..2802dd13 100644
--- a/applications/evaluation/img_compare_loderunner.input
+++ b/applications/evaluation/img_compare_loderunner.input
@@ -9,7 +9,7 @@ design_lsc240420_MASTER.csv
 --LSC_NPZ_DIR
 /data2/lsc240420/
 --checkpoint
-study011_modelState_epoch0090.hdf5
+/net/sescratch1/atoivonen/projects/yoke_runs/KN_loderunner_fine_tune/runs/study_005/study005_modelState_epoch0100.pth
 --sampIDX
-3914
+10
 -S
diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py
index 8cf7a866..abec3d8f 100644
--- a/src/yoke/utils/training/datastep/loderunner.py
+++ b/src/yoke/utils/training/datastep/loderunner.py
@@ -740,6 +740,49 @@ def train_DDP_temporal_loderunner_datastep(
     return target_img, pred_img, all_losses
 
 
+def train_DDP_scalar_temporal_loderunner_datastep(
+    data,
+    model,
+    optimizer,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.train()
+
+    x, target, Dt = data
+
+    x = x.to(torch.float32).to(device, non_blocking=True)          # [B, 2 * context_len]
+    target = target.to(torch.float32).to(device, non_blocking=True) # [B]
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)         # [B]
+
+    C = 8
+    in_vars = torch.arange(C, device=device)
+    out_vars = torch.arange(C, device=device)
+
+    pred = model(x, in_vars, out_vars, Dt)  # [B]
+
+    pred = pred.view_as(target)
+
+    loss = loss_fn(pred, target)            # [B] if reduction="none"
+    per_sample_loss = loss.view(loss.shape[0], -1).mean(dim=1)
+
+    optimizer.zero_grad(set_to_none=True)
+    per_sample_loss.mean().backward()
+    optimizer.step()
+
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)
+    else:
+        all_losses = None
+
+    return target, pred, all_losses
+
+
 def eval_DDP_temporal_loderunner_datastep(
     data,
     model,
@@ -776,6 +819,44 @@ def eval_DDP_temporal_loderunner_datastep(
     return target_img, pred_img, all_losses
 
 
+def eval_DDP_scalar_temporal_loderunner_datastep(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    model.eval()
+
+    with torch.no_grad():
+        x, target, Dt = data
+
+        x = x.to(torch.float32).to(device, non_blocking=True)          # [B, 2 * context_len]
+        target = target.to(torch.float32).to(device, non_blocking=True) # [B]
+        Dt = Dt.to(torch.float32).to(device, non_blocking=True)         # [B]
+
+        C = 8
+        in_vars = torch.arange(C, device=device)
+        out_vars = torch.arange(C, device=device)
+
+        pred = model(x, in_vars, out_vars, Dt)  # [B]
+        pred = pred.view_as(target)
+
+        loss = loss_fn(pred, target)            # [B] if reduction="none"
+        per_sample_loss = loss.view(loss.shape[0], -1).mean(dim=1)
+
+        gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+        dist.all_gather(gathered_losses, per_sample_loss)
+
+        if rank == 0:
+            all_losses = torch.cat(gathered_losses, dim=0)
+        else:
+            all_losses = None
+
+    return target, pred, all_losses
+
+
 def eval_DDP_loderunner_seq_context_datastep(
     data,
     model,
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 9d26af3e..ce81a894 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -14,11 +14,12 @@
     train_DDP_loderunner_seq_datastep,
     train_DDP_loderunner_seq_channel_datastep,
     train_DDP_temporal_loderunner_datastep,
+    train_DDP_scalar_temporal_loderunner_datastep,
     eval_DDP_loderunner_datastep,
     eval_DDP_loderunner_seq_datastep,
     eval_DDP_loderunner_seq_context_datastep,
     eval_DDP_loderunner_seq_channel_datastep,
-
+    eval_DDP_scalar_temporal_loderunner_datastep,
 )
 
 
@@ -337,6 +338,92 @@ def train_LRsched_loderunner_epoch(
                     np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
 
 
+def train_DDP_scalar_temporal_loderunner_epoch(
+    training_data: torch.utils.data.DataLoader,
+    validation_data: torch.utils.data.DataLoader,
+    num_train_batches: int,
+    num_val_batches: int,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    LRsched: torch.optim.lr_scheduler._LRScheduler,
+    epochIDX: int,
+    train_per_val: int,
+    train_rcrd_filename: str,
+    val_rcrd_filename: str,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> None:
+    trainbatch_ID = 0
+    valbatch_ID = 0
+
+    model.train()
+
+    train_rcrd_filename = train_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
+
+    with (
+        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
+    ) as train_rcrd_file:
+        for trainbatch_ID, traindata in enumerate(training_data):
+            if trainbatch_ID >= num_train_batches:
+                break
+
+            truth, pred, train_losses = train_DDP_scalar_temporal_loderunner_datastep(
+                traindata,
+                model,
+                optimizer,
+                loss_fn,
+                device,
+                rank,
+                world_size,
+            )
+
+            LRsched.step()
+
+            if rank == 0:
+                batch_records = np.column_stack(
+                    [
+                        np.full(len(train_losses), epochIDX),
+                        np.full(len(train_losses), trainbatch_ID),
+                        train_losses.cpu().numpy().flatten(),
+                    ]
+                )
+                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+    if epochIDX % train_per_val == 0:
+        print("Validating...", epochIDX)
+
+        val_rcrd_filename = val_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
+        model.eval()
+
+        with (
+            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
+        ) as val_rcrd_file:
+            for valbatch_ID, valdata in enumerate(validation_data):
+                if valbatch_ID >= num_val_batches:
+                    break
+
+                truth, pred, val_losses = eval_DDP_scalar_temporal_loderunner_datastep(
+                    valdata,
+                    model,
+                    loss_fn,
+                    device,
+                    rank,
+                    world_size,
+                )
+
+                if rank == 0:
+                    batch_records = np.column_stack(
+                        [
+                            np.full(len(val_losses), epochIDX),
+                            np.full(len(val_losses), valbatch_ID),
+                            val_losses.cpu().numpy().flatten(),
+                        ]
+                    )
+                    np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+
 def train_DDP_loderunner_epoch(
     training_data: torch.utils.data.DataLoader,
     validation_data: torch.utils.data.DataLoader,

From 61b1b60dc765f2f4aebaa363fd64c77759f72246 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Tue, 28 Jul 2026 09:03:24 -0600
Subject: [PATCH 08/66] Include multiple filters

---
 .../utils/training/datastep/loderunner.py     | 176 ++++++++++++++++++
 src/yoke/utils/training/epoch/loderunner.py   | 174 ++++++++++++++++-
 2 files changed, 340 insertions(+), 10 deletions(-)

diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py
index abec3d8f..6077e4ff 100644
--- a/src/yoke/utils/training/datastep/loderunner.py
+++ b/src/yoke/utils/training/datastep/loderunner.py
@@ -10,6 +10,182 @@
 import random
 
 
+def train_DDP_scalar_temporal_loderunner_datastep_gri(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    DDP training datastep for scalar temporal LodeRunner wrapper.
+
+    Expected data:
+        x:      [B, input_dim]
+                for GRI with context_len=5, input_dim = 20
+
+        target: [B, 3]
+                normalized [delta_g, delta_r, delta_i]
+
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, 3]
+    """
+
+    model.train()
+
+    x, target, Dt = data
+
+    x = x.to(device, non_blocking=True)
+    target = target.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    # Kept for LodeRunner-style API compatibility.
+    # The scalar wrapper can ignore these or internally override them.
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    pred = model(x, in_vars, out_vars, Dt)
+
+    if pred.shape != target.shape:
+        raise RuntimeError(
+            f"Prediction and target shapes do not match: "
+            f"pred.shape={pred.shape}, target.shape={target.shape}"
+        )
+
+    loss = loss_fn(pred, target)
+
+    # For HuberLoss/MSELoss with reduction='none':
+    # loss.shape == [B, 3]
+    #
+    # Reduce over output channels only, leaving one scalar loss per sample.
+    if loss.ndim == 1:
+        per_sample_loss = loss
+    else:
+        per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
+
+    optimizer.zero_grad(set_to_none=True)
+    per_sample_loss.mean().backward()
+    optimizer.step()
+
+    # Do not all_gather here. DDP already synchronizes gradients.
+    # Returning local-rank losses avoids the all_gather hangs you saw before.
+    return target, pred, per_sample_loss.detach()
+
+
+def eval_DDP_scalar_temporal_loderunner_datastep_gri(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    """
+    Evaluation datastep for ScalarTemporalConditionedLodeRunner.
+
+    Expected dataset output:
+        x:      [B, input_dim]
+                e.g. [B, 20] for context_len=5 with g/r/i + relative times
+
+        target: [B, 3]
+                normalized [delta_g, delta_r, delta_i]
+
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, 3]
+    """
+
+    model.eval()
+
+    x, target, Dt = data
+
+    x = x.to(device, non_blocking=True)
+    target = target.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    # Kept for compatibility with the LodeRunner-style model call.
+    # The scalar wrapper internally uses 8 backbone channels.
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    with torch.no_grad():
+        pred = model(x, in_vars, out_vars, Dt)
+
+    if pred.shape != target.shape:
+        raise RuntimeError(
+            "Prediction and target shapes do not match in eval datastep: "
+            f"pred.shape={pred.shape}, target.shape={target.shape}"
+        )
+
+    loss = loss_fn(pred, target)
+
+    # For HuberLoss/MSELoss with reduction='none':
+    #   loss.shape == [B, 3]
+    #
+    # Reduce over output bands, leaving one scalar loss per sample.
+    if loss.ndim == 1:
+        per_sample_loss = loss
+    else:
+        per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
+
+    return target, pred, per_sample_loss.detach()
+
+
+def eval_DDP_scalar_temporal_loderunner_datastep(
+    data: tuple,
+    model: torch.nn.Module,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    DDP evaluation datastep for scalar temporal LodeRunner wrapper.
+
+    Expected data:
+        x:      [B, input_dim]
+        target: [B, 3]
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, 3]
+    """
+
+    model.eval()
+
+    x, target, Dt = data
+
+    x = x.to(device, non_blocking=True)
+    target = target.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    with torch.no_grad():
+        pred = model(x, in_vars, out_vars, Dt)
+
+    if pred.shape != target.shape:
+        raise RuntimeError(
+            f"Validation prediction and target shapes do not match: "
+            f"pred.shape={pred.shape}, target.shape={target.shape}"
+        )
+
+    loss = loss_fn(pred, target)
+
+    if loss.ndim == 1:
+        per_sample_loss = loss
+    else:
+        per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
+
+    return target, pred, per_sample_loss.detach()
+
+
 def train_loderunner_datastep(
     data: tuple,
     model: torch.nn.Module,
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index ce81a894..18c7ca27 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -14,12 +14,12 @@
     train_DDP_loderunner_seq_datastep,
     train_DDP_loderunner_seq_channel_datastep,
     train_DDP_temporal_loderunner_datastep,
-    train_DDP_scalar_temporal_loderunner_datastep,
+    train_DDP_scalar_temporal_loderunner_datastep_gri,
     eval_DDP_loderunner_datastep,
     eval_DDP_loderunner_seq_datastep,
     eval_DDP_loderunner_seq_context_datastep,
     eval_DDP_loderunner_seq_channel_datastep,
-    eval_DDP_scalar_temporal_loderunner_datastep,
+    eval_DDP_scalar_temporal_loderunner_datastep_gri,
 )
 
 
@@ -338,6 +338,159 @@ def train_LRsched_loderunner_epoch(
                     np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
 
 
+def train_DDP_scalar_temporal_loderunner_epoch_gri(
+    training_data: torch.utils.data.DataLoader,
+    validation_data: torch.utils.data.DataLoader,
+    num_train_batches: int,
+    num_val_batches: int,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    LRsched: torch.optim.lr_scheduler._LRScheduler,
+    epochIDX: int,
+    train_per_val: int,
+    train_rcrd_filename: str,
+    val_rcrd_filename: str,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> None:
+    """
+    DDP epoch function for scalar temporal LodeRunner training.
+
+    Expected dataset output:
+        x:      [B, input_dim]
+        target: [B, n_outputs]
+        Dt:     [B]
+
+    For the 3-band kilonova case:
+        x:      [B, 4 * context_len]
+                flattened [g, r, i] context plus relative times
+        target: [B, 3]
+                normalized delta_g, delta_r, delta_i
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, 3]
+    """
+
+    train_rcrd_filename = train_rcrd_filename.replace(
+        "<epochIDX>",
+        f"{epochIDX:04d}",
+    )
+
+    model.train()
+
+    with (
+        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
+    ) as train_rcrd_file:
+
+        for trainbatch_ID, data in enumerate(training_data):
+            if trainbatch_ID >= num_train_batches:
+                break
+
+            x, target, Dt = data
+
+            x = x.to(device, non_blocking=True)
+            target = target.to(device, non_blocking=True)
+            Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+            optimizer.zero_grad(set_to_none=True)
+
+            # These are kept for API compatibility with LodeRunner-style wrappers.
+            # ScalarTemporalConditionedLodeRunner may ignore them internally,
+            # or pass them to the backbone.
+            in_vars = torch.arange(8, device=device)
+            out_vars = torch.arange(8, device=device)
+
+            pred = model(x, in_vars, out_vars, Dt)
+
+            if pred.shape != target.shape:
+                raise RuntimeError(
+                    f"Prediction and target shapes do not match: "
+                    f"pred.shape={pred.shape}, target.shape={target.shape}"
+                )
+
+            loss = loss_fn(pred, target)
+
+            # Huber/MSE with reduction='none' gives [B, 3].
+            # Reduce over output channels, leaving one loss per sample.
+            if loss.ndim == 1:
+                per_sample_loss = loss
+            else:
+                per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
+
+            batch_loss = per_sample_loss.mean()
+
+            batch_loss.backward()
+            optimizer.step()
+            LRsched.step()
+
+            if rank == 0:
+                batch_records = np.column_stack(
+                    [
+                        np.full(len(per_sample_loss), epochIDX),
+                        np.full(len(per_sample_loss), trainbatch_ID),
+                        per_sample_loss.detach().cpu().numpy().flatten(),
+                    ]
+                )
+                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+    if epochIDX % train_per_val == 0:
+        if rank == 0:
+            print("Validating...", epochIDX, flush=True)
+
+        val_rcrd_filename = val_rcrd_filename.replace(
+            "<epochIDX>",
+            f"{epochIDX:04d}",
+        )
+
+        model.eval()
+
+        with (
+            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
+        ) as val_rcrd_file:
+
+            with torch.no_grad():
+                for valbatch_ID, data in enumerate(validation_data):
+                    if valbatch_ID >= num_val_batches:
+                        break
+
+                    x, target, Dt = data
+
+                    x = x.to(device, non_blocking=True)
+                    target = target.to(device, non_blocking=True)
+                    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+                    in_vars = torch.arange(8, device=device)
+                    out_vars = torch.arange(8, device=device)
+
+                    pred = model(x, in_vars, out_vars, Dt)
+
+                    if pred.shape != target.shape:
+                        raise RuntimeError(
+                            f"Validation prediction and target shapes do not match: "
+                            f"pred.shape={pred.shape}, target.shape={target.shape}"
+                        )
+
+                    loss = loss_fn(pred, target)
+
+                    if loss.ndim == 1:
+                        per_sample_loss = loss
+                    else:
+                        per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
+
+                    if rank == 0:
+                        batch_records = np.column_stack(
+                            [
+                                np.full(len(per_sample_loss), epochIDX),
+                                np.full(len(per_sample_loss), valbatch_ID),
+                                per_sample_loss.detach().cpu().numpy().flatten(),
+                            ]
+                        )
+                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+
 def train_DDP_scalar_temporal_loderunner_epoch(
     training_data: torch.utils.data.DataLoader,
     validation_data: torch.utils.data.DataLoader,
@@ -369,7 +522,7 @@ def train_DDP_scalar_temporal_loderunner_epoch(
             if trainbatch_ID >= num_train_batches:
                 break
 
-            truth, pred, train_losses = train_DDP_scalar_temporal_loderunner_datastep(
+            truth, pred, train_losses = train_DDP_scalar_temporal_loderunner_datastep_gri(
                 traindata,
                 model,
                 optimizer,
@@ -404,13 +557,14 @@ def train_DDP_scalar_temporal_loderunner_epoch(
                 if valbatch_ID >= num_val_batches:
                     break
 
-                truth, pred, val_losses = eval_DDP_scalar_temporal_loderunner_datastep(
-                    valdata,
-                    model,
-                    loss_fn,
-                    device,
-                    rank,
-                    world_size,
+
+                truth, pred, val_losses = eval_DDP_scalar_temporal_loderunner_datastep_gri(
+                    data=data,
+                    model=model,
+                    loss_fn=loss_fn,
+                    device=device,
+                    rank=rank,
+                    world_size=world_size,
                 )
 
                 if rank == 0:

From 898dfa1c9056cb1f8145c5f30ff73b5f8a9758ab Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Tue, 28 Jul 2026 16:02:58 -0600
Subject: [PATCH 09/66] Harness files

---
 .../KN_loderunner/ddp_production.csv          |   29 +
 .../KN_loderunner/plot_loss_curves_channel.py |  124 ++
 .../KN_loderunner/plot_loss_curves_gri.py     |  290 +++++
 .../harnesses/KN_loderunner/plot_pred.py      |  369 ++++++
 .../KN_loderunner/plot_pred_channel.py        |  214 ++++
 .../KN_loderunner/plot_pred_channel_delta.py  |  396 ++++++
 .../KN_loderunner/plot_pred_diagnostics.py    |  382 ++++++
 .../plot_pred_diagnostics_gri.py              |  635 ++++++++++
 .../plot_pred_diagnostics_new.py              |  574 +++++++++
 .../harnesses/KN_loderunner/plot_pred_gri.py  |  518 ++++++++
 .../KN_loderunner/plot_pred_seq_context.py    |  229 ++++
 .../KN_loderunner/train_LodeRunner_ddp.py     | 1086 +++++++++++++++++
 .../KN_loderunner/training_START.input        |   53 +
 .../KN_loderunner/training_START.slurm        |   73 ++
 .../KN_loderunner/training_input.tmpl         |   56 +
 .../KN_loderunner/training_slurm.tmpl         |   73 ++
 16 files changed, 5101 insertions(+)
 create mode 100644 applications/harnesses/KN_loderunner/ddp_production.csv
 create mode 100644 applications/harnesses/KN_loderunner/plot_loss_curves_channel.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_channel.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_channel_delta.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_diagnostics.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_diagnostics_gri.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_diagnostics_new.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_gri.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_seq_context.py
 create mode 100644 applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
 create mode 100644 applications/harnesses/KN_loderunner/training_START.input
 create mode 100644 applications/harnesses/KN_loderunner/training_START.slurm
 create mode 100644 applications/harnesses/KN_loderunner/training_input.tmpl
 create mode 100644 applications/harnesses/KN_loderunner/training_slurm.tmpl

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
new file mode 100644
index 00000000..a6695405
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -0,0 +1,29 @@
+studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZE,NTRN_BATCH,NVAL_BATCH,ANCHOR_LR,NUM_CYCLES,MIN_FRACTION,TERMINAL_STEPS,WARMUP_STEPS,NOISE_SCALE,train_script
+# This is a longer epoch production run after tuning.
+# Single epoch with validation should be ~30 mins
+#1,torch_ch_gpu_241112,4,4,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#2,torch_ch_gpu_241112,1,4,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#3,torch_ch_gpu_241112,1,4,128,1,1,9,1,2,10,200,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#4,yoke311,1,4,128,1,1,9,1,2,10,200,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#5,yoke311,1,4,128,1,1,9,1,2,10,200,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#6,yoke311,1,4,128,1,1,9,1,2,10,200,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#7,yoke311,1,4,128,1,1,9,1,2,10,200,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#8,yoke311,1,4,128,1,1,9,1,2,10,200,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#9,yoke311,1,4,128,1,1,9,1,1,1,5,0,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.#py
+#10,yoke311,1,1,128,1,1,9,1,2,10,200,50,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#11,yoke311,1,4,128,1,1,9,1,2,10,1000,500,1.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#12,yoke311,1,4,128,1,1,9,1,2,10,1000,500,1.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#13,yoke311,1,4,128,1,1,9,1,2,10,1000,500,1.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#14,yoke311,1,4,128,1,1,9,1,2,10,1000,500,1.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#15,yoke311,1,4,128,1,1,9,1,2,10,1000,500,2.0e-4,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#16,yoke311,1,4,128,1,1,9,1,2,10,1000,500,2.0e-4,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#17,yoke311,1,4,128,1,1,9,1,2,10,1000,500,2.0e-4,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#18,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-4,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#19,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#20,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#21,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#22,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#23,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+24,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+
+
diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_channel.py b/applications/harnesses/KN_loderunner/plot_loss_curves_channel.py
new file mode 100644
index 00000000..507a067e
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_channel.py
@@ -0,0 +1,124 @@
+import argparse
+import glob
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+
+study = '21'
+
+def load_records(pattern):
+    files = sorted(glob.glob(pattern))
+
+    if len(files) == 0:
+        raise FileNotFoundError(f"No files matched pattern: {pattern}")
+
+    arrays = []
+
+    for fn in files:
+        try:
+            arr = np.loadtxt(fn, delimiter=",")
+        except Exception as e:
+            print(f"Skipping {fn}: {e}")
+            continue
+
+        if arr.size == 0:
+            continue
+
+        if arr.ndim == 1:
+            arr = arr[None, :]
+
+        arrays.append(arr)
+
+    if len(arrays) == 0:
+        raise RuntimeError(f"No valid data found for pattern: {pattern}")
+
+    data = np.vstack(arrays)
+
+    # columns: epoch, batch, loss
+    epochs = data[:, 0].astype(int)
+    batches = data[:, 1].astype(int)
+    losses = data[:, 2]
+
+    return epochs, batches, losses, files
+
+
+def epoch_means(epochs, losses):
+    unique_epochs = np.array(sorted(set(epochs)))
+    mean_losses = np.array([losses[epochs == e].mean() for e in unique_epochs])
+    std_losses = np.array([losses[epochs == e].std() for e in unique_epochs])
+    return unique_epochs, mean_losses, std_losses
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+    "--train_pattern",
+    type=str,
+    #default="runs/study_010/training_study010_epoch*.csv",
+    default=f"runs/study_0{study}/training_study0{study}_epoch*.csv",
+    )
+
+    parser.add_argument(
+        "--val_pattern",
+        type=str,
+        #default="runs/study_010/validation_study010_epoch*.csv",
+        default=f"runs/study_0{study}/validation_study0{study}_epoch*.csv",
+    )
+
+    parser.add_argument(
+        "--out",
+        type=str,
+        default=f"loss_curves_study0{study}.png",
+    )
+
+    parser.add_argument(
+        "--logy",
+        action="store_true",
+        default=True,
+        help="Use log scale on y-axis.",
+    )
+
+    args = parser.parse_args()
+
+    train_epochs, train_batches, train_losses, train_files = load_records(args.train_pattern)
+
+    print("Loaded training files:")
+    for f in train_files:
+        print("  ", f)
+
+    train_ep, train_mean, train_std = epoch_means(train_epochs, train_losses)
+
+    plt.figure(figsize=(8, 5))
+    plt.plot(train_ep, train_mean, marker="o", label="Train")
+
+    # Try validation, but do not fail if absent
+    try:
+        val_epochs, val_batches, val_losses, val_files = load_records(args.val_pattern)
+
+        print("Loaded validation files:")
+        for f in val_files:
+            print("  ", f)
+
+        val_ep, val_mean, val_std = epoch_means(val_epochs, val_losses)
+        plt.plot(val_ep, val_mean, marker="s", label="Validation")
+
+    except Exception as e:
+        print(f"No validation curve plotted: {e}")
+
+    plt.xlabel("Epoch")
+    plt.ylabel("Mean loss")
+    plt.title("Loss curves")
+    plt.grid(True, alpha=0.3)
+    plt.legend()
+
+    if args.logy:
+        plt.yscale("log")
+
+    plt.tight_layout()
+    plt.savefig(args.out, dpi=200)
+    print(f"Saved {args.out}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
new file mode 100644
index 00000000..4d97aabb
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
@@ -0,0 +1,290 @@
+import argparse
+import glob
+from pathlib import Path
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+DEFAULT_STUDY = 24
+DEFAULT_RUNS_ROOT = "runs"
+DEFAULT_COLUMNS = "epoch,batch,loss"
+
+
+def format_study(study):
+    """Return both integer and zero-padded study strings."""
+    study_int = int(study)
+    return study_int, f"{study_int:03d}"
+
+
+def default_patterns(study, runs_root):
+    _, study_tag = format_study(study)
+    run_dir = Path(runs_root) / f"study_{study_tag}"
+    return {
+        "train": str(run_dir / f"training_study{study_tag}_epoch*.csv"),
+        "val": str(run_dir / f"validation_study{study_tag}_epoch*.csv"),
+        "out": f"loss_curves_study{study_tag}.png",
+    }
+
+
+def parse_column_names(columns_arg, n_cols):
+    """Build display names for columns in the CSV record files."""
+    provided = [c.strip() for c in columns_arg.split(",") if c.strip()]
+
+    if len(provided) < n_cols:
+        provided.extend([f"col{i}" for i in range(len(provided), n_cols)])
+
+    return provided[:n_cols]
+
+
+def load_records(pattern, columns_arg=DEFAULT_COLUMNS):
+    files = sorted(glob.glob(pattern))
+
+    if len(files) == 0:
+        raise FileNotFoundError(f"No files matched pattern: {pattern}")
+
+    arrays = []
+    skipped = []
+
+    for fn in files:
+        try:
+            arr = np.loadtxt(fn, delimiter=",")
+        except Exception as exc:
+            skipped.append((fn, str(exc)))
+            continue
+
+        if arr.size == 0:
+            skipped.append((fn, "empty file"))
+            continue
+
+        if arr.ndim == 1:
+            arr = arr[None, :]
+
+        if arr.shape[1] < 3:
+            skipped.append((fn, f"expected at least 3 columns, found {arr.shape[1]}"))
+            continue
+
+        arrays.append(arr)
+
+    if len(arrays) == 0:
+        details = "\n".join(f"  {fn}: {reason}" for fn, reason in skipped)
+        raise RuntimeError(f"No valid data found for pattern: {pattern}\n{details}")
+
+    # Keep only the common width if an interrupted run left mixed-width records.
+    n_cols = min(arr.shape[1] for arr in arrays)
+    if any(arr.shape[1] != n_cols for arr in arrays):
+        print(f"Warning: mixed CSV widths found; using first {n_cols} columns.")
+        arrays = [arr[:, :n_cols] for arr in arrays]
+
+    data = np.vstack(arrays)
+
+    # Sort by epoch, then batch.
+    data = data[np.lexsort((data[:, 1], data[:, 0]))]
+
+    names = parse_column_names(columns_arg, n_cols)
+
+    epochs = data[:, 0].astype(int)
+    batches = data[:, 1].astype(int)
+    losses = data[:, 2:]
+    loss_names = names[2:]
+
+    return {
+        "epochs": epochs,
+        "batches": batches,
+        "losses": losses,
+        "loss_names": loss_names,
+        "files": files,
+        "skipped": skipped,
+        "data": data,
+    }
+
+
+def epoch_stats(epochs, losses):
+    unique_epochs = np.array(sorted(set(epochs)))
+    mean_losses = np.vstack(
+        [losses[epochs == e].mean(axis=0) for e in unique_epochs]
+    )
+    std_losses = np.vstack(
+        [losses[epochs == e].std(axis=0) for e in unique_epochs]
+    )
+    return unique_epochs, mean_losses, std_losses
+
+
+def infer_loss_labels(loss_names, n_loss_cols):
+    """Make nicer labels for common scalar/GRI cases."""
+    if n_loss_cols == 1:
+        return [loss_names[0] if loss_names else "loss"]
+
+    if n_loss_cols == 3 and loss_names == ["loss", "col3", "col4"]:
+        return ["g", "r", "i"]
+
+    if n_loss_cols == 4 and loss_names == ["loss", "col3", "col4", "col5"]:
+        return ["total", "g", "r", "i"]
+
+    return loss_names
+
+
+def plot_epoch_curves(train, val, args):
+    train_ep, train_mean, train_std = epoch_stats(
+        train["epochs"],
+        train["losses"],
+    )
+
+    n_loss_cols = train["losses"].shape[1]
+    loss_labels = infer_loss_labels(train["loss_names"], n_loss_cols)
+
+    if val is not None:
+        val_ep, val_mean, val_std = epoch_stats(
+            val["epochs"],
+            val["losses"],
+        )
+
+        if val_mean.shape[1] != n_loss_cols:
+            print(
+                "Validation has a different number of loss columns "
+                f"({val_mean.shape[1]}) than training ({n_loss_cols}); "
+                "skipping validation."
+            )
+            val = None
+            val_ep = None
+            val_mean = None
+            val_std = None
+    else:
+        val_ep = None
+        val_mean = None
+        val_std = None
+
+    plt.figure(figsize=(9, 5.5))
+
+    for idx, label in enumerate(loss_labels):
+        suffix = "" if n_loss_cols == 1 else f" {label}"
+
+        plt.plot(
+            train_ep,
+            train_mean[:, idx],
+            marker="o",
+            label=f"Train{suffix}",
+        )
+
+        if args.show_std:
+            lo = np.maximum(train_mean[:, idx] - train_std[:, idx], 1e-30)
+            hi = train_mean[:, idx] + train_std[:, idx]
+            plt.fill_between(train_ep, lo, hi, alpha=0.15)
+
+        if val is not None:
+            plt.plot(
+                val_ep,
+                val_mean[:, idx],
+                marker="s",
+                linestyle="--",
+                label=f"Validation{suffix}",
+            )
+
+            if args.show_std:
+                lo = np.maximum(val_mean[:, idx] - val_std[:, idx], 1e-30)
+                hi = val_mean[:, idx] + val_std[:, idx]
+                plt.fill_between(val_ep, lo, hi, alpha=0.10)
+
+    plt.xlabel("Epoch")
+    plt.ylabel("Mean loss")
+    plt.title(args.title)
+    plt.grid(True, alpha=0.3)
+    plt.legend()
+
+    if args.logy:
+        plt.yscale("log")
+
+    plt.tight_layout()
+    plt.savefig(args.out, dpi=args.dpi)
+    print(f"Saved {args.out}")
+
+
+def print_loaded(label, record):
+    print(f"Loaded {label} files:")
+    for fn in record["files"]:
+        print("  ", fn)
+
+    for fn, reason in record["skipped"]:
+        print(f"  skipped {fn}: {reason}")
+
+    print(f"{label} rows: {record['data'].shape[0]}")
+    print(f"{label} columns: {record['data'].shape[1]}")
+    print(f"{label} loss columns: {', '.join(record['loss_names'])}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Plot training/validation loss curves for scalar temporal LodeRunner GRI runs."
+    )
+
+    parser.add_argument("--study", type=int, default=DEFAULT_STUDY)
+    parser.add_argument("--runs_root", type=str, default=DEFAULT_RUNS_ROOT)
+    parser.add_argument("--train_pattern", type=str, default=None)
+    parser.add_argument("--val_pattern", type=str, default=None)
+    parser.add_argument("--out", type=str, default=None)
+
+    parser.add_argument(
+        "--columns",
+        type=str,
+        default=DEFAULT_COLUMNS,
+        help=(
+            "Comma-separated CSV column names. The first two must be epoch,batch. "
+            "Examples: epoch,batch,loss or epoch,batch,total,g,r,i."
+        ),
+    )
+
+    parser.add_argument("--title", type=str, default="GRI loss curves")
+    parser.add_argument("--dpi", type=int, default=200)
+
+    parser.add_argument(
+        "--logy",
+        dest="logy",
+        action="store_true",
+        default=True,
+        help="Use log scale on y-axis. This is the default.",
+    )
+
+    parser.add_argument(
+        "--linear",
+        dest="logy",
+        action="store_false",
+        help="Use linear y-axis.",
+    )
+
+    parser.add_argument(
+        "--show_std",
+        action="store_true",
+        help="Shade +/- one epoch standard deviation.",
+    )
+
+    parser.add_argument(
+        "--require_val",
+        action="store_true",
+        help="Fail instead of continuing when validation records are missing or invalid.",
+    )
+
+    args = parser.parse_args()
+
+    defaults = default_patterns(args.study, args.runs_root)
+
+    args.train_pattern = args.train_pattern or defaults["train"]
+    args.val_pattern = args.val_pattern or defaults["val"]
+    args.out = args.out or defaults["out"]
+
+    train = load_records(args.train_pattern, args.columns)
+    print_loaded("training", train)
+
+    try:
+        val = load_records(args.val_pattern, args.columns)
+        print_loaded("validation", val)
+    except Exception as exc:
+        if args.require_val:
+            raise
+        print(f"No validation curve plotted: {exc}")
+        val = None
+
+    plot_epoch_curves(train, val, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred.py b/applications/harnesses/KN_loderunner/plot_pred.py
new file mode 100644
index 00000000..31168f02
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred.py
@@ -0,0 +1,369 @@
+import os
+import time
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+#from yoke.datasets.lsc_dataset import LSC_rho2rho_temporal_DataSet
+#from yoke.utils.training.epoch.loderunner import train_DDP_loderunner_epoch
+#from yoke.utils.restart import continuation_setup
+#from yoke.utils.dataload import make_distributed_dataloader
+#from yoke.utils.checkpointing import load_model_and_optimizer
+#from yoke.utils.checkpointing import save_model_and_optimizer
+#from yoke.lr_schedulers import CosineWithWarmupScheduler
+#from yoke.helpers import cli
+
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet
+from torch.utils.data import DataLoader
+
+# FIXME remove if restructure
+#from torch.utils.data import Dataset, DataLoader, random_split
+import glob
+import random
+
+import torch
+
+import matplotlib
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+import pdb
+
+# matplotlib.use('MacOSX')
+# matplotlib.use('pdf')
+# Get rid of type 3 fonts in figures
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+# Ensure LaTeX font
+font = {"family": "serif"}
+plt.rc("font", **font)
+plt.rcParams["figure.figsize"] = (6, 6)
+
+#ckpt = torch.load(
+#    "runs/study_005/study005_modelState_epoch0100.pth",
+#    map_location="cpu",
+#    weights_only=False,   # trusted checkpoint
+#)
+
+#model = LodeRunner(**ckpt["model_args"])
+#model.load_state_dict(ckpt["model_state_dict"])
+#model.eval()
+
+#file_prefix_list = sorted(glob.glob(f"/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz"))
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#ckpt = torch.load("runs/study_005/study005_modelState_epoch0100.pth", map_location=device, weights_only=False)
+ckpt = torch.load("runs/study_006/study006_modelState_epoch0028.pth", map_location=device, weights_only=False)
+
+model = LodeRunner(**ckpt["model_args"])
+model.load_state_dict(ckpt["model_state_dict"])
+model.to(device)
+model.eval()
+#eval_dataset = Kilonova_lc_img_DataSet(max_timeIDX_offset=2, half_image=False, N_imgs=1)
+eval_dataset = Kilonova_lc_img_DataSet(half_image=False, N_imgs=1)
+
+xs, targets, Dts = [], [], []
+idxs = []
+
+loader = DataLoader(eval_dataset, batch_size=1, shuffle=False)
+for idx, (x, target, Dt) in enumerate(loader):
+    xs.append(x.mean().item())
+    targets.append(target.mean().item())
+    idxs.append(idx)
+
+plt.figure()
+plt.plot(idxs, xs)
+plt.savefig('plot_val_lc_seq_new.png')
+
+xs_pred, targets_pred, Dts_pred = [], [], []
+preds, idxs_pred = [], []
+
+in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+loader = DataLoader(eval_dataset, batch_size=1, shuffle=False)
+for idx, (x, target, Dt) in enumerate(loader):
+    xs_pred.append(x.mean().item())
+    targets_pred.append(target.mean().item())
+    idxs_pred.append(idx)
+
+    with torch.no_grad():
+        pred_image = model(x, in_vars, out_vars, Dt)   # [1, 8, 1120, 400]
+    pred_plot = pred_image.squeeze().detach().cpu().numpy().mean()
+    preds.append(pred_plot)
+
+print(idxs_pred, len(idxs_pred))
+print(preds, len(preds))
+print(targets_pred, len(targets_pred))
+
+plt.figure()
+plt.scatter(idxs_pred, preds, label='Predictions')
+plt.scatter(idxs_pred, targets_pred, label='Truth')
+plt.legend()
+plt.gca().invert_yaxis()
+plt.savefig('pred_vs_truth_seq_new.png')
+
+
+loader = DataLoader(eval_dataset, batch_size=1, shuffle=False)
+it = iter(loader)
+
+x0, target0, Dt0 = next(it)
+prev = x0
+
+xs_seq = []
+preds_seq = []
+idxs_seq = []
+
+for idx, (x, target, Dt) in enumerate(loader):
+    xs_seq.append(x.mean().item())
+    idxs_seq.append(idx)
+
+    with torch.no_grad():
+        pred_image = model(prev, in_vars, out_vars, Dt)
+        prev = pred_image
+
+    preds_seq.append(pred_image.mean().item())
+
+plt.figure()
+plt.scatter(idxs_seq, preds_seq, label='Predictions')
+plt.scatter(idxs_seq, xs_seq, label='Truth')
+plt.legend()
+plt.gca().invert_yaxis()
+plt.savefig('pred_vs_truth_seq_new2.png')
+
+pred_plot = pred_image.squeeze().mean(dim=0).detach().cpu().numpy() #.mean(dim=0)
+true_plot = target.squeeze().mean(dim=0).detach().cpu().numpy()
+error_plot = pred_plot - true_plot
+
+# --- consistent color scale for pred + truth ---
+vmin = min(pred_plot.min(), true_plot.min())
+vmax = max(pred_plot.max(), true_plot.max())
+
+# --- error scale (symmetric around 0 looks better) ---
+err_max = np.max(np.abs(error_plot))
+
+fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
+
+# Prediction
+im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
+ax1.set_title(f"Prediction")
+
+# Truth
+im2 = ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
+ax2.set_title(f"Truth")
+
+# Error
+im3 = ax3.imshow(error_plot, origin="lower", vmin=-err_max, vmax=err_max)
+ax3.set_title("Error (Pred - Truth)")
+
+# --- shared colorbar for pred + truth ---
+cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
+cbar.set_label("Field value")
+
+# --- separate colorbar for error ---
+cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
+cbar_err.set_label("Error")
+
+# Clean up axes
+for ax in (ax1, ax2, ax3):
+    ax.axis("off")
+
+plt.tight_layout()
+savefile = 'img_comp_test_new.png'
+plt.savefig(savefile, bbox_inches="tight")
+
+n_times = 5   # number of frames/samples to plot
+
+pred_means = []
+true_means = []
+dts = []
+
+'''
+for idx in range(n_times):
+    x, target, Dt = eval_dataset[idx]
+
+    x = x.unsqueeze(0).to(device)              # [1, 8, 1120, 400]
+    lead_times = Dt.unsqueeze(0).to(device)
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    with torch.no_grad():
+        pred_image = model(x, in_vars, out_vars, lead_times)   # [1, 8, 1120, 400]
+
+    # select one channel and average over the whole image
+    pred_mean = pred_image[0, channel].mean().item()
+    true_mean = target[channel].mean().item()
+
+    pred_means.append(pred_mean)
+    true_means.append(true_mean)
+    dts.append(Dt.item())
+
+times = np.arange(n_times)
+
+plt.figure(figsize=(10, 5))
+plt.plot(times, pred_means, marker="o", label="Prediction")
+plt.plot(times, true_means, marker="s", label="Truth")
+plt.xlabel("Sample index")
+plt.ylabel(f"Mean image value, channel {channel}")
+plt.title("Mean-value time series")
+plt.grid(True)
+plt.legend()
+plt.tight_layout()
+plt.savefig('series_test_dt.png')
+plt.show()
+
+times = np.cumsum([0.0] + dts[:-1])
+
+plt.figure(figsize=(10, 5))
+plt.plot(times, pred_means, marker="o", label="Prediction")
+plt.plot(times, true_means, marker="s", label="Truth")
+plt.xlabel("Time")
+plt.ylabel(f"Mean image value, channel {channel}")
+plt.title("Mean-value time series")
+plt.grid(True)
+plt.legend()
+plt.tight_layout()
+plt.savefig('series_test.png')
+plt.show()
+'''
+
+
+
+'''
+idx = 0
+x, target, Dt = eval_dataset[idx]
+
+x = x.unsqueeze(0).to(device)          # [1, 8, 1120, 400]
+lead_times = Dt.unsqueeze(0).to(device)
+
+in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+with torch.no_grad():
+    pred_image = model(x, in_vars, out_vars, lead_times)   # [1, 8, 1120, 400]
+
+print("x shape:", x.shape)
+print("target shape:", target.shape)
+print("pred_image shape:", pred_image.shape)
+
+channel = 0
+
+pred_plot = pred_image[0, channel].detach().cpu().numpy()
+true_plot = target[channel].detach().cpu().numpy()
+
+fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
+
+img1 = ax1.imshow(pred_plot, origin="lower")
+img2 = ax2.imshow(true_plot, origin="lower")
+
+ax1.set_title(f"Prediction channel {channel}")
+ax2.set_title(f"Truth channel {channel}")
+
+plt.colorbar(img1, ax=ax1)
+plt.colorbar(img2, ax=ax2)
+plt.show()
+
+idx = 0
+sample = eval_dataset[idx]
+x, target, Dt = sample   # adapt this to your dataset's actual return format
+
+x = x.unsqueeze(0).to(device)        # add batch dimension
+lead_times = Dt.unsqueeze(0).to(device)
+
+in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+with torch.no_grad():
+    pred_image = model(x, in_vars, out_vars, lead_times)
+
+print(pred_image)
+
+#with torch.no_grad():
+#    out = model(x)   # x should already be on the same device
+
+#fig1, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 6))
+fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
+
+sim_params = x
+true_image = target
+
+#pred_plot = pred_image[0, 0].detach().cpu().numpy()
+#true_plot = true_image[0, 0].detach().cpu().numpy()
+
+pred_plot = pred_image[0].detach().cpu().numpy()
+true_plot = target.detach().cpu().numpy()
+
+print("x shape:", x.shape)
+print("target shape:", target.shape)
+print("pred_image shape:", pred_image.shape)
+
+img1 = ax1.imshow(pred_plot, origin="lower")
+img2 = ax2.imshow(true_plot, origin="lower")
+
+# Reshape for plotting
+sim_params = sim_params.numpy()
+true_image = np.squeeze(true_image.numpy())
+# Predictions from network must be detached from gradients in order to be
+# written to numpy arrays.
+pred_image = np.squeeze(pred_image.detach().numpy())
+# print('Shape of image prediction:', pred_image.shape)
+
+# Plot Truth/Prediction/Discrepancy panel.
+fig1, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 6))
+#fig1.suptitle(f"Time={sim_params[-1]:.3f}us", fontsize=18)
+img1 = ax1.imshow(
+    true_image,
+    aspect="equal",
+    origin="lower",
+    cmap="jet",
+    vmin=true_image.min(),
+    vmax=true_image.max(),
+)
+ax1.set_ylabel("Z-axis", fontsize=16)
+ax1.set_xlabel("R-axis", fontsize=16)
+ax1.set_title("True", fontsize=18)
+
+# divider1 = make_axes_locatable(ax1)
+# cax1 = divider1.append_axes('right', size='10%', pad=0.1)
+# fig1.colorbar(img1,
+#               cax=cax1).set_label('Density',
+#                                   fontsize=14)
+
+
+img2 = ax2.imshow(
+    pred_image,
+    aspect="equal",
+    origin="lower",
+    cmap="jet",
+    vmin=true_image.min(),
+    vmax=true_image.max(),
+)
+ax2.set_title("Predicted", fontsize=18)
+ax2.tick_params(axis="y", which="both", left=False, labelleft=False)
+
+divider2 = make_axes_locatable(ax2)
+cax2 = divider2.append_axes("right", size="10%", pad=0.1)
+fig1.colorbar(img2, cax=cax2).set_label("Density (g/cc)", fontsize=14)
+
+discrepancy = np.abs(true_image - pred_image)
+img3 = ax3.imshow(
+    discrepancy,
+    aspect="equal",
+    origin="lower",
+    cmap="hot",
+    vmin=discrepancy.min(),
+    vmax=dscale * discrepancy.max(),
+)
+ax3.set_title("Discrepancy", fontsize=18)
+ax3.tick_params(axis="y", which="both", left=False, labelleft=False)
+
+divider3 = make_axes_locatable(ax3)
+cax3 = divider3.append_axes("right", size="10%", pad=0.1)
+fig1.colorbar(img3, cax=cax3).set_label("Discrepancy", fontsize=14)
+'''
+
diff --git a/applications/harnesses/KN_loderunner/plot_pred_channel.py b/applications/harnesses/KN_loderunner/plot_pred_channel.py
new file mode 100644
index 00000000..d126bbcb
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_channel.py
@@ -0,0 +1,214 @@
+import argparse
+import numpy as np
+import torch
+import matplotlib
+import matplotlib.pyplot as plt
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from torch.utils.data import DataLoader
+
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (6, 6)
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        #default="runs/study_007/study007_modelState_epoch0100.pth",
+        #default="runs/study_010/study010_modelState_epoch0100.pth",
+        default="runs/study_012/study012_modelState_epoch0100.pth",
+
+    )
+    parser.add_argument("--N_imgs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--n_future_steps", type=int, default=10)
+
+    return parser.parse_args()
+
+
+def load_channel_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    model = LodeRunner(**model_args)
+    model.to(device)
+
+    state_dict = ckpt["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+    print("Loaded model_args:", model_args)
+
+    model.noise_scale = noise_scale
+    model.eval()
+
+    return model
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    context_len = 5
+    model = load_channel_model(args.ckpt, device)
+
+    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    # ------------------------------------------------------------
+    # One-step predictions using true context windows
+    # ------------------------------------------------------------
+    preds = []
+    targets = []
+    idxs = []
+    prefix = []
+
+    for idx, (context_img, target, Dt) in enumerate(loader):
+        context_img = context_img.to(device)
+        if idx == 0:
+            context_means = context_img.mean(dim=(2, 3))[0].detach().cpu().numpy()
+            for context in context_means:
+                prefix.append(context.mean().item())
+        target = target.to(device)
+        Dt = Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_image = model(context_img, in_vars, out_vars, Dt)
+
+        preds.append(pred_image.mean().item())
+        targets.append(target.mean().item())
+        idxs.append(idx)
+
+    plt.figure()
+    plt.scatter(idxs, preds, label="Predictions")
+    plt.scatter(idxs, targets, label="Truth")
+    plt.scatter(np.arange(len(prefix))-(len(prefix)), prefix, label='Initial Context Window')
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Sample index")
+    plt.ylabel("Mean magnitude/image value")
+    plt.tight_layout()
+    plt.savefig("pred_vs_truth_channel_norm.png", dpi=200)
+
+
+
+    context_seq, target, Dt = next(iter(loader))
+
+    context_seq = context_seq.to(device)
+    Dt = Dt.to(torch.float32).to(device)
+
+    x = context_seq
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    future_iter = iter(loader)
+
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target = future_target.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_image = model(x, in_vars, out_vars, future_Dt)
+
+        preds_seq.append(pred_image.mean().item())
+        truth_seq.append(future_target.mean().item())
+        idxs_seq.append(step)
+
+        # autoregressive update: append prediction
+        x = torch.cat([x[:, 1:], pred_image[:, -1:].detach()], dim=1)
+
+    plt.figure()
+
+    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
+    plt.scatter(idxs_seq, truth_seq, label="Truth")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Mean magnitude/image value")
+    plt.tight_layout()
+    plt.savefig("pred_vs_truth_channel_norm_autoreg.png", dpi=200)
+
+
+    # ------------------------------------------------------------
+    # Image comparison for the final one-step batch above
+    # ------------------------------------------------------------
+    pred_plot = pred_image.squeeze().mean(dim=0).detach().cpu().numpy()
+    true_plot = target.squeeze().mean(dim=0).detach().cpu().numpy()
+    error_plot = pred_plot - true_plot
+
+    vmin = min(pred_plot.min(), true_plot.min())
+    vmax = max(pred_plot.max(), true_plot.max())
+    err_max = np.max(np.abs(error_plot))
+
+    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
+
+    im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
+    ax1.set_title("Prediction")
+
+    ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
+    ax2.set_title("Truth")
+
+    im3 = ax3.imshow(error_plot, origin="lower", vmin=-err_max, vmax=err_max)
+    ax3.set_title("Error (Pred - Truth)")
+
+    cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
+    cbar.set_label("Field value")
+
+    cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
+    cbar_err.set_label("Error")
+
+    for ax in (ax1, ax2, ax3):
+        ax.axis("off")
+
+    plt.tight_layout()
+    plt.savefig("img_comp_channel_norm.png", bbox_inches="tight", dpi=200)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_channel_delta.py b/applications/harnesses/KN_loderunner/plot_pred_channel_delta.py
new file mode 100644
index 00000000..0a3394bb
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_channel_delta.py
@@ -0,0 +1,396 @@
+import argparse
+import numpy as np
+import torch
+import matplotlib
+import matplotlib.pyplot as plt
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from torch.utils.data import DataLoader
+
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (6, 6)
+
+# ============================================================
+# RUN IDENTIFIER
+# ============================================================
+RUN_ID = "015"
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default=f"runs/study_{RUN_ID}/study{RUN_ID}_modelState_epoch0089.pth",
+    )
+    parser.add_argument("--N_imgs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--n_future_steps", type=int, default=15)
+
+    return parser.parse_args()
+
+
+def load_channel_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+    context_len = ckpt.get("context_len", 5)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("predicts_delta:", ckpt.get("predicts_delta", False))
+    print("target_type:", ckpt.get("target_type", "absolute"))
+    print("context_len:", context_len)
+
+    model = LodeRunner(**model_args)
+    model.to(device)
+
+    state_dict = ckpt["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.noise_scale = noise_scale
+    model.eval()
+
+    return model, context_len
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_channel_model(args.ckpt, device)
+
+    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    # ------------------------------------------------------------
+    # One-step predictions
+    # ------------------------------------------------------------
+    preds = []
+    targets = []
+    idxs = []
+    prefix = []
+
+    for idx, (context_img, target_delta, Dt) in enumerate(loader):
+        context_img = context_img.to(device)
+        target_delta = target_delta.to(device)
+        Dt = Dt.to(torch.float32).to(device)
+
+        if idx == 0:
+            context_means = context_img.mean(dim=(2, 3))[0].detach().cpu().numpy()
+            for context in context_means:
+                prefix.append(context.mean().item())
+
+        with torch.no_grad():
+            pred_delta_img = model(context_img, in_vars, out_vars, Dt)
+
+        last_mag_img = context_img[:, -1:]
+
+        pred_next_img = last_mag_img + pred_delta_img[:, -1:]
+        true_next_img = last_mag_img + target_delta[:, -1:]
+
+        preds.append(pred_next_img.mean().item())
+        targets.append(true_next_img.mean().item())
+        idxs.append(idx)
+
+    plt.figure()
+    plt.scatter(idxs, preds, label="Predicted next magnitude")
+    plt.scatter(idxs, targets, label="True next magnitude")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Sample index")
+    plt.ylabel("Normalized magnitude")
+    plt.tight_layout()
+
+    plt.savefig(
+        f"study{RUN_ID}_pred_vs_truth_channel_delta_onestep.png",
+        dpi=200,
+    )
+
+    '''
+    # ------------------------------------------------------------
+    # Autoregressive rollout
+    # ------------------------------------------------------------
+    context_seq, target_delta, Dt = next(iter(loader))
+
+    x = context_seq.to(device)
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    future_iter = iter(loader)
+
+    pred_next_img = None
+    true_next_img = None
+
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target_delta, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target_delta = future_target_delta.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            #pred_delta_img = model(x, in_vars, out_vars, future_Dt)
+            noise_scale = 100
+            x_noisy = x + noise_scale * torch.randn_like(x)
+            pred_delta_img = model(x_noisy, in_vars, out_vars, future_Dt)
+
+        last_mag_img = x[:, -1:]
+
+        pred_next_img = last_mag_img + pred_delta_img[:, -1:]
+        true_next_img = last_mag_img + future_target_delta[:, -1:]
+
+        preds_seq.append(pred_next_img.mean().item())
+        truth_seq.append(true_next_img.mean().item())
+        idxs_seq.append(step)
+
+        # Append predicted absolute next magnitude
+        x = torch.cat([x[:, 1:], pred_next_img.detach()], dim=1)
+    '''
+
+    # ------------------------------------------------------------
+    # Autoregressive rollout
+    # ------------------------------------------------------------
+    context_seq, target_delta, Dt = next(iter(loader))
+
+    x_pred = context_seq.to(device)
+    x_true = context_seq.to(device)
+
+    # Initial context window for plotting
+    context_means = context_seq.mean(dim=(2, 3))[0].detach().cpu().numpy()
+    prefix = [context.mean().item() for context in context_means]
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    future_iter = iter(loader)
+
+    pred_next_img = None
+    true_next_img = None
+
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target_delta, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target_delta = future_target_delta.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        # --------------------------------------------------------
+        # Hardcoded nonsense test for autoregressive dependence
+        # --------------------------------------------------------
+        DEBUG_AUTOREG_NONSENSE = False
+        DEBUG_START_STEP = 1
+        DEBUG_VALUE = 10.0
+
+        x_model_input = x_pred
+
+        if DEBUG_AUTOREG_NONSENSE and step >= DEBUG_START_STEP:
+            x_model_input = x_pred.clone()
+            x_model_input[:, -1:] = DEBUG_VALUE
+
+        with torch.no_grad():
+            pred_delta_img = model(x_model_input, in_vars, out_vars, future_Dt)
+
+        pred_last_img = x_pred[:, -1:]
+        true_last_img = x_true[:, -1:]
+
+        pred_next_img = pred_last_img + pred_delta_img[:, -1:]
+        true_next_img = true_last_img + future_target_delta[:, -1:]
+
+        preds_seq.append(pred_next_img.mean().item())
+        truth_seq.append(true_next_img.mean().item())
+        idxs_seq.append(step)
+
+        x_pred = torch.cat(
+            [x_pred[:, 1:], pred_next_img.detach()],
+            dim=1,
+        )
+
+        x_true = torch.cat(
+            [x_true[:, 1:], true_next_img.detach()],
+            dim=1,
+        )
+
+    '''
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target_delta, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target_delta = future_target_delta.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_delta_img = model(x_pred, in_vars, out_vars, future_Dt)
+
+        pred_last_img = x_pred[:, -1:]
+        true_last_img = x_true[:, -1:]
+
+        pred_next_img = pred_last_img + pred_delta_img[:, -1:]
+        true_next_img = true_last_img + future_target_delta[:, -1:]
+
+        preds_seq.append(pred_next_img.mean().item())
+        truth_seq.append(true_next_img.mean().item())
+        idxs_seq.append(step)
+
+        # Update autoregressive prediction context
+        x_pred = torch.cat(
+            [x_pred[:, 1:], pred_next_img.detach()],
+            dim=1,
+        )
+
+        # Update true context separately
+        x_true = torch.cat(
+            [x_true[:, 1:], true_next_img.detach()],
+            dim=1,
+        )
+    '''
+    # ------------------------------------------------------------
+    # Plot rollout
+    # ------------------------------------------------------------
+    plt.figure()
+
+    plt.scatter(
+        idxs_seq,
+        preds_seq,
+        label="Autoregressive predictions",
+    )
+
+    plt.scatter(
+        idxs_seq,
+        truth_seq,
+        label="Truth",
+    )
+
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Normalized magnitude")
+    plt.tight_layout()
+
+    plt.savefig(
+        f"study{RUN_ID}_pred_vs_truth_channel_delta_autoreg_clean.png",
+        dpi=200,
+    )
+
+    '''
+    plt.figure()
+
+    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
+    plt.scatter(idxs_seq, truth_seq, label="Truth")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Normalized magnitude")
+    plt.tight_layout()
+
+    plt.savefig(
+        f"study{RUN_ID}_pred_vs_truth_channel_delta_autoreg.png",
+        dpi=200,
+    )
+    '''
+
+    # ------------------------------------------------------------
+    # Image comparison
+    # ------------------------------------------------------------
+    if pred_next_img is not None and true_next_img is not None:
+        pred_plot = pred_next_img[0, 0].detach().cpu().numpy()
+        true_plot = true_next_img[0, 0].detach().cpu().numpy()
+        error_plot = pred_plot - true_plot
+
+        vmin = min(pred_plot.min(), true_plot.min())
+        vmax = max(pred_plot.max(), true_plot.max())
+        err_max = np.max(np.abs(error_plot))
+
+        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
+
+        im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
+        ax1.set_title("Prediction")
+
+        ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
+        ax2.set_title("Truth")
+
+        im3 = ax3.imshow(
+            error_plot,
+            origin="lower",
+            vmin=-err_max,
+            vmax=err_max,
+        )
+        ax3.set_title("Error (Pred - Truth)")
+
+        cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
+        cbar.set_label("Normalized magnitude")
+
+        cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
+        cbar_err.set_label("Error")
+
+        for ax in (ax1, ax2, ax3):
+            ax.axis("off")
+
+        plt.tight_layout()
+
+        plt.savefig(
+            f"study{RUN_ID}_img_comp_channel_delta.png",
+            bbox_inches="tight",
+            dpi=200,
+        )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics.py
new file mode 100644
index 00000000..8c206db1
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics.py
@@ -0,0 +1,382 @@
+import argparse
+import csv
+import os
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+RUN_ID = "021"
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default=f"runs/study_{RUN_ID}/study{RUN_ID}_modelState_epoch0300.pth",
+    )
+    parser.add_argument("--N_imgs", type=int, default=10)
+    parser.add_argument("--n_future_steps", type=int, default=15)
+    parser.add_argument("--n_series", type=int, default=10)
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        default=f"runs/study_{RUN_ID}/autoreg_diagnostics",
+    )
+
+    return parser.parse_args()
+
+
+def load_channel_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+    context_len = ckpt.get("context_len", 5)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("predicts_delta:", ckpt.get("predicts_delta", False))
+    print("target_type:", ckpt.get("target_type", "absolute"))
+    print("context_len:", context_len)
+
+    model = LodeRunner(**model_args).to(device)
+
+    state_dict = ckpt["model_state_dict"]
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.noise_scale = noise_scale
+    model.eval()
+
+    return model, context_len
+
+
+def ensure_batch(x):
+    """
+    Dataset item usually has shape:
+        [T, H, W]
+
+    Model expects:
+        [B, T, H, W]
+    """
+    if x.ndim == 3:
+        return x.unsqueeze(0)
+
+    return x
+
+
+def tensor_time_means(x):
+    """
+    Convert image sequence tensor to scalar light curve.
+
+    Supports:
+        [T, H, W]
+        [B, T, H, W]
+        [B, T, C, H, W]
+    """
+    if x.ndim == 3:
+        return x.mean(dim=(1, 2)).detach().cpu().numpy().squeeze()
+
+    if x.ndim == 4:
+        return x.mean(dim=(2, 3)).detach().cpu().numpy().squeeze()
+
+    if x.ndim == 5:
+        return x.mean(dim=(2, 3, 4)).detach().cpu().numpy().squeeze()
+
+    raise ValueError(f"Unexpected tensor shape: {x.shape}")
+
+
+def get_rollout_from_start(
+    dataset,
+    model,
+    device,
+    start_idx,
+    n_future_steps,
+    in_vars,
+    out_vars,
+):
+    """
+    Clean autoregressive rollout.
+
+    x_pred:
+        model-generated autoregressive context
+
+    x_true:
+        ground-truth context used only to reconstruct true future values
+
+    This avoids the bug where truth was reconstructed using the predicted
+    previous frame.
+    """
+
+    context_img, _, _ = dataset[start_idx]
+
+    x_pred = ensure_batch(context_img).to(device)
+    x_true = ensure_batch(context_img).to(device)
+
+    context_curve = tensor_time_means(x_true)
+
+    pred_curve = []
+    truth_curve = []
+    residual_curve = []
+    step_mses = []
+
+    with torch.no_grad():
+        for step in range(n_future_steps):
+            future_idx = start_idx + step
+
+            if future_idx >= len(dataset):
+                break
+
+            _, future_target_delta, future_Dt = dataset[future_idx]
+
+            future_target_delta = ensure_batch(future_target_delta).to(device)
+
+            future_Dt = torch.as_tensor(
+                future_Dt,
+                dtype=torch.float32,
+                device=device,
+            )
+
+            if future_Dt.ndim == 0:
+                future_Dt = future_Dt.unsqueeze(0)
+
+            pred_delta_img = model(x_pred, in_vars, out_vars, future_Dt)
+
+            pred_last_img = x_pred[:, -1:]
+            true_last_img = x_true[:, -1:]
+
+            pred_next_img = pred_last_img + pred_delta_img[:, -1:]
+            true_next_img = true_last_img + future_target_delta[:, -1:]
+
+            pred_scalar = pred_next_img.mean().item()
+            true_scalar = true_next_img.mean().item()
+            residual_scalar = pred_scalar - true_scalar
+
+            step_mse = torch.mean((pred_next_img - true_next_img) ** 2).item()
+
+            pred_curve.append(pred_scalar)
+            truth_curve.append(true_scalar)
+            residual_curve.append(residual_scalar)
+            step_mses.append(step_mse)
+
+            # Autoregressive model context gets the prediction.
+            x_pred = torch.cat(
+                [x_pred[:, 1:], pred_next_img.detach()],
+                dim=1,
+            )
+
+            # Truth context gets the independently reconstructed truth.
+            x_true = torch.cat(
+                [x_true[:, 1:], true_next_img.detach()],
+                dim=1,
+            )
+
+    pred_curve = np.asarray(pred_curve)
+    truth_curve = np.asarray(truth_curve)
+    residual_curve = np.asarray(residual_curve)
+    step_mses = np.asarray(step_mses)
+
+    total_mse = np.mean(step_mses) if len(step_mses) > 0 else np.nan
+
+    return {
+        "start_idx": start_idx,
+        "context": context_curve,
+        "pred": pred_curve,
+        "truth": truth_curve,
+        "residual": residual_curve,
+        "step_mses": step_mses,
+        "mse": total_mse,
+    }
+
+
+def plot_residuals_vs_step(rollouts, outpath):
+    plt.figure(figsize=(8, 5))
+
+    for rollout in rollouts:
+        steps = np.arange(len(rollout["residual"]))
+        plt.plot(
+            steps,
+            rollout["residual"],
+            marker="o",
+            alpha=0.75,
+            label=f"start {rollout['start_idx']}",
+        )
+
+    plt.axhline(0.0, linestyle="--", linewidth=1)
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Residual: prediction - truth")
+    plt.title("Autoregressive residuals vs time step")
+    plt.legend(fontsize=8, ncol=2)
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def plot_multiple_series_predictions(rollouts, outpath):
+    plt.figure(figsize=(9, 6))
+
+    for rollout in rollouts:
+        start_idx = rollout["start_idx"]
+
+        context_steps = np.arange(-len(rollout["context"]), 0)
+        future_steps = np.arange(len(rollout["pred"]))
+
+        plt.plot(
+            context_steps,
+            rollout["context"],
+            linestyle=":",
+            alpha=0.45,
+        )
+
+        plt.plot(
+            future_steps,
+            rollout["truth"],
+            linewidth=1.5,
+            alpha=0.75,
+            label=f"truth start {start_idx}",
+        )
+
+        plt.plot(
+            future_steps,
+            rollout["pred"],
+            linestyle="--",
+            linewidth=1.5,
+            alpha=0.75,
+            label=f"pred start {start_idx}",
+        )
+
+    plt.axvline(-0.5, linewidth=1, alpha=0.5)
+    plt.gca().invert_yaxis()
+    plt.xlabel("Time step relative to forecast start")
+    plt.ylabel("Normalized magnitude")
+    plt.title("Autoregressive predictions for validation curves")
+    plt.legend(fontsize=7, ncol=2)
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def plot_mse_histogram(rollouts, outpath):
+    mses = np.asarray([r["mse"] for r in rollouts])
+    mses = mses[np.isfinite(mses)]
+
+    plt.figure(figsize=(7, 5))
+    plt.hist(mses, bins=min(10, max(1, len(mses))))
+    plt.xlabel("Mean autoregressive MSE per validation curve")
+    plt.ylabel("Count")
+    plt.title("Distribution of autoregressive rollout MSEs")
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def save_mse_csv(rollouts, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["start_idx", "mse", "n_steps"])
+
+        for rollout in rollouts:
+            writer.writerow(
+                [
+                    rollout["start_idx"],
+                    rollout["mse"],
+                    len(rollout["residual"]),
+                ]
+            )
+
+
+def main():
+    args = get_args()
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_channel_model(args.ckpt, device)
+
+    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    max_start = max(0, len(eval_dataset) - args.n_future_steps)
+    n_series = min(args.n_series, max_start + 1)
+
+    print("Dataset length:", len(eval_dataset))
+    print("Number of rollout series:", n_series)
+    print("Autoregressive future steps:", args.n_future_steps)
+
+    start_indices = np.linspace(0, max_start, n_series, dtype=int)
+
+    rollouts = []
+
+    for start_idx in start_indices:
+        print(f"Rolling out validation curve starting at index {start_idx}")
+
+        rollout = get_rollout_from_start(
+            dataset=eval_dataset,
+            model=model,
+            device=device,
+            start_idx=int(start_idx),
+            n_future_steps=args.n_future_steps,
+            in_vars=in_vars,
+            out_vars=out_vars,
+        )
+
+        rollouts.append(rollout)
+
+    residual_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_residuals_vs_step.png",
+    )
+    series_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_multi_series_predictions.png",
+    )
+    hist_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_mse_histogram.png",
+    )
+    csv_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_mse_by_curve.csv",
+    )
+
+    plot_residuals_vs_step(rollouts, residual_path)
+    plot_multiple_series_predictions(rollouts, series_path)
+    plot_mse_histogram(rollouts, hist_path)
+    save_mse_csv(rollouts, csv_path)
+
+    print("Saved:")
+    print(" ", residual_path)
+    print(" ", series_path)
+    print(" ", hist_path)
+    print(" ", csv_path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_gri.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_gri.py
new file mode 100644
index 00000000..2381781a
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_gri.py
@@ -0,0 +1,635 @@
+import argparse
+import csv
+import os
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+
+from train_LodeRunner_ddp import (
+    Kilonova_lc_scalar_context_DataSet_gri,
+    ScalarTemporalConditionedLodeRunner_gri,
+    load_or_compute_band_normalization,
+)
+
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+
+BAND_KEYS = ("arr_ztfg", "arr_ztfr", "arr_ztfi")
+BAND_NAMES = ("g", "r", "i")
+VALUE_COL = 1
+
+
+def study_tag(study):
+    return f"{int(study):03d}"
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description="Autoregressive prediction diagnostics for scalar temporal LodeRunner GRI runs."
+    )
+
+    parser.add_argument("--study", type=int, default=24)
+    parser.add_argument("--epoch", type=int, default=500)
+    parser.add_argument("--ckpt", type=str, default=None)
+
+    parser.add_argument("--N_imgs", type=int, default=10)
+    parser.add_argument("--n_future_steps", type=int, default=15)
+    parser.add_argument("--n_series", type=int, default=10)
+
+    parser.add_argument("--outdir", type=str, default=None)
+    parser.add_argument(
+        "--norm_stats_path",
+        type=str,
+        default="kilonova_gri_norm_stats.npz",
+    )
+
+    parser.add_argument(
+        "--plot_all_series",
+        action="store_true",
+        help="Plot every rollout series. By default, plots all selected series too, but this flag is kept for compatibility.",
+    )
+
+    return parser.parse_args()
+
+
+def resolve_paths(args):
+    tag = study_tag(args.study)
+
+    if args.ckpt is None:
+        args.ckpt = (
+            f"runs/study_{tag}/study{tag}_modelState_epoch{args.epoch:04d}.pth"
+        )
+
+    if args.outdir is None:
+        args.outdir = f"runs/study_{tag}/autoreg_diagnostics_gri"
+
+    return tag
+
+
+def strip_ddp_prefix(state_dict):
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        return {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+    return state_dict
+
+
+def load_gri_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    context_len = ckpt.get("context_len", 5)
+
+    n_input_channels = ckpt.get("n_input_channels", 3)
+    n_output_channels = ckpt.get("n_output_channels", 3)
+    backbone_channels = ckpt.get("backbone_channels", 8)
+    hidden = ckpt.get("hidden", 64)
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("model_class:", ckpt.get("model_class", "unknown"))
+    print("backbone_class:", ckpt.get("backbone_class", "LodeRunner"))
+    print("predicts_delta:", ckpt.get("predicts_delta", False))
+    print("target_type:", ckpt.get("target_type", "unknown"))
+    print("context_len:", context_len)
+    print("n_input_channels:", n_input_channels)
+    print("n_output_channels:", n_output_channels)
+    print("backbone_channels:", backbone_channels)
+    print("hidden:", hidden)
+
+    backbone = LodeRunner(**model_args).to(device)
+    backbone.noise_scale = noise_scale
+
+    model = ScalarTemporalConditionedLodeRunner_gri(
+        backbone=backbone,
+        context_len=context_len,
+        n_input_channels=n_input_channels,
+        n_output_channels=n_output_channels,
+        image_size=model_args["image_size"],
+        backbone_channels=backbone_channels,
+        hidden=hidden,
+    ).to(device)
+
+    state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded ScalarTemporalConditionedLodeRunner_gri checkpoint")
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.eval()
+
+    return model, context_len
+
+
+def make_eval_dataset(args, context_len):
+    band_means, band_stds = load_or_compute_band_normalization(
+        stats_path=args.norm_stats_path,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+    )
+
+    print("Using band normalization:")
+    print("band_means:", band_means)
+    print("band_stds:", band_stds)
+
+    dataset = Kilonova_lc_scalar_context_DataSet_gri(
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+        means=band_means,
+        stds=band_stds,
+        predicts_delta=True,
+    )
+
+    return dataset
+
+
+def split_gri_context(x, context_len, n_bands=3):
+    """
+    Dataset x layout:
+        [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK, t0, t1, ..., tK]
+
+    Returns:
+        values: [context_len, 3]
+        rel_t:  [context_len]
+    """
+    x = torch.as_tensor(x, dtype=torch.float32)
+
+    value_count = context_len * n_bands
+    values = x[:value_count].detach().cpu().numpy().reshape(context_len, n_bands)
+    rel_t = x[value_count:].detach().cpu().numpy()
+
+    return values.astype(np.float32), rel_t.astype(np.float32)
+
+
+def build_gri_input(values, rel_t, device):
+    """
+    values:
+        [context_len, 3]
+    rel_t:
+        [context_len]
+
+    Returns:
+        x: [1, context_len * 3 + context_len]
+    """
+    values = np.asarray(values, dtype=np.float32)
+    rel_t = np.asarray(rel_t, dtype=np.float32)
+
+    x = np.concatenate(
+        [
+            values.reshape(-1),
+            rel_t,
+        ],
+        axis=0,
+    )
+
+    return torch.tensor(
+        x,
+        dtype=torch.float32,
+        device=device,
+    ).unsqueeze(0)
+
+
+def get_rollout_from_start_gri(
+    dataset,
+    model,
+    device,
+    start_idx,
+    n_future_steps,
+    context_len,
+):
+    x0, _, _ = dataset[start_idx]
+
+    context_vals, _ = split_gri_context(
+        x0,
+        context_len=context_len,
+        n_bands=3,
+    )
+
+    pred_vals = [row.copy() for row in context_vals]
+    true_vals = [row.copy() for row in context_vals]
+
+    pred_curve = []
+    truth_curve = []
+    residual_curve = []
+    step_mses = []
+    step_band_mses = []
+
+    with torch.no_grad():
+        for step in range(n_future_steps):
+            future_idx = start_idx + step
+
+            if future_idx >= len(dataset):
+                break
+
+            x_true_step, target_delta, future_Dt = dataset[future_idx]
+
+            x_true_step = torch.as_tensor(x_true_step, dtype=torch.float32)
+            target_delta = torch.as_tensor(
+                target_delta,
+                dtype=torch.float32,
+                device=device,
+            )
+            future_Dt = torch.as_tensor(
+                future_Dt,
+                dtype=torch.float32,
+                device=device,
+            )
+
+            if target_delta.ndim == 0:
+                raise ValueError(
+                    "Expected GRI target_delta shape [3], but got scalar target."
+                )
+
+            target_delta = target_delta.reshape(3)
+
+            if future_Dt.ndim == 0:
+                future_Dt = future_Dt.unsqueeze(0)
+
+            _, current_rel_t = split_gri_context(
+                x_true_step,
+                context_len=context_len,
+                n_bands=3,
+            )
+
+            current_pred_vals = np.asarray(
+                pred_vals[-context_len:],
+                dtype=np.float32,
+            )
+
+            x_pred = build_gri_input(
+                values=current_pred_vals,
+                rel_t=current_rel_t,
+                device=device,
+            )
+
+            pred_delta = model(
+                x_pred,
+                in_vars=None,
+                out_vars=None,
+                Dt=future_Dt,
+            )
+
+            pred_delta = pred_delta.reshape(3)
+
+            pred_next = (
+                torch.as_tensor(pred_vals[-1], dtype=torch.float32, device=device)
+                + pred_delta
+            )
+            true_next = (
+                torch.as_tensor(true_vals[-1], dtype=torch.float32, device=device)
+                + target_delta
+            )
+
+            residual = pred_next - true_next
+            band_mse = residual.pow(2)
+            total_mse = band_mse.mean()
+
+            pred_next_np = pred_next.detach().cpu().numpy()
+            true_next_np = true_next.detach().cpu().numpy()
+            residual_np = residual.detach().cpu().numpy()
+            band_mse_np = band_mse.detach().cpu().numpy()
+
+            pred_curve.append(pred_next_np)
+            truth_curve.append(true_next_np)
+            residual_curve.append(residual_np)
+            step_band_mses.append(band_mse_np)
+            step_mses.append(float(total_mse.detach().cpu()))
+
+            pred_vals.append(pred_next_np)
+            true_vals.append(true_next_np)
+
+    pred_curve = np.asarray(pred_curve, dtype=np.float32)
+    truth_curve = np.asarray(truth_curve, dtype=np.float32)
+    residual_curve = np.asarray(residual_curve, dtype=np.float32)
+    step_mses = np.asarray(step_mses, dtype=np.float32)
+    step_band_mses = np.asarray(step_band_mses, dtype=np.float32)
+
+    total_mse = np.mean(step_mses) if len(step_mses) > 0 else np.nan
+
+    if len(step_band_mses) > 0:
+        band_mse = np.mean(step_band_mses, axis=0)
+    else:
+        band_mse = np.full(3, np.nan, dtype=np.float32)
+
+    return {
+        "start_idx": start_idx,
+        "context": np.asarray(context_vals, dtype=np.float32),
+        "pred": pred_curve,
+        "truth": truth_curve,
+        "residual": residual_curve,
+        "step_mses": step_mses,
+        "step_band_mses": step_band_mses,
+        "mse": total_mse,
+        "band_mse": band_mse,
+    }
+
+
+def plot_residuals_vs_step(rollouts, outpath):
+    fig, axes = plt.subplots(3, 1, figsize=(8, 9), sharex=True)
+
+    for band_idx, band_name in enumerate(BAND_NAMES):
+        ax = axes[band_idx]
+
+        for rollout in rollouts:
+            residual = rollout["residual"]
+            steps = np.arange(len(residual))
+
+            ax.plot(
+                steps,
+                residual[:, band_idx],
+                marker="o",
+                alpha=0.70,
+                label=f"start {rollout['start_idx']}",
+            )
+
+        ax.axhline(0.0, linestyle="--", linewidth=1)
+        ax.set_ylabel(f"{band_name} residual")
+        ax.set_title(f"{band_name}-band residuals")
+
+    axes[-1].set_xlabel("Autoregressive step")
+
+    handles, labels = axes[0].get_legend_handles_labels()
+    fig.legend(
+        handles,
+        labels,
+        fontsize=7,
+        ncol=2,
+        loc="upper center",
+        bbox_to_anchor=(0.5, 1.02),
+    )
+
+    fig.suptitle("Autoregressive residuals vs time step", y=1.06)
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_multiple_series_predictions(rollouts, outpath):
+    fig, axes = plt.subplots(3, 1, figsize=(9, 10), sharex=True)
+
+    for band_idx, band_name in enumerate(BAND_NAMES):
+        ax = axes[band_idx]
+
+        for rollout in rollouts:
+            start_idx = rollout["start_idx"]
+
+            context_steps = np.arange(-len(rollout["context"]), 0)
+            future_steps = np.arange(len(rollout["pred"]))
+
+            ax.plot(
+                context_steps,
+                rollout["context"][:, band_idx],
+                linestyle=":",
+                alpha=0.45,
+            )
+
+            ax.plot(
+                future_steps,
+                rollout["truth"][:, band_idx],
+                linewidth=1.5,
+                alpha=0.75,
+                label=f"truth start {start_idx}",
+            )
+
+            ax.plot(
+                future_steps,
+                rollout["pred"][:, band_idx],
+                linestyle="--",
+                linewidth=1.5,
+                alpha=0.75,
+                label=f"pred start {start_idx}",
+            )
+
+        ax.axvline(-0.5, linewidth=1, alpha=0.5)
+        ax.invert_yaxis()
+        ax.set_ylabel(f"{band_name} norm mag")
+        ax.set_title(f"{band_name}-band autoregressive predictions")
+
+    axes[-1].set_xlabel("Time step relative to forecast start")
+
+    handles, labels = axes[0].get_legend_handles_labels()
+    fig.legend(
+        handles,
+        labels,
+        fontsize=6,
+        ncol=2,
+        loc="upper center",
+        bbox_to_anchor=(0.5, 1.02),
+    )
+
+    fig.suptitle("Autoregressive predictions for validation curves", y=1.06)
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_mse_histogram(rollouts, outpath):
+    mses = np.asarray([r["mse"] for r in rollouts], dtype=np.float32)
+    mses = mses[np.isfinite(mses)]
+
+    plt.figure(figsize=(7, 5))
+    plt.hist(mses, bins=min(10, max(1, len(mses))))
+    plt.xlabel("Mean autoregressive MSE per validation curve")
+    plt.ylabel("Count")
+    plt.title("Distribution of autoregressive rollout MSEs")
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def plot_band_mse_histograms(rollouts, outpath):
+    band_mses = np.asarray([r["band_mse"] for r in rollouts], dtype=np.float32)
+
+    fig, axes = plt.subplots(3, 1, figsize=(7, 9), sharex=False)
+
+    for band_idx, band_name in enumerate(BAND_NAMES):
+        vals = band_mses[:, band_idx]
+        vals = vals[np.isfinite(vals)]
+
+        axes[band_idx].hist(vals, bins=min(10, max(1, len(vals))))
+        axes[band_idx].set_xlabel(f"{band_name} mean autoregressive MSE")
+        axes[band_idx].set_ylabel("Count")
+        axes[band_idx].set_title(f"{band_name}-band rollout MSE distribution")
+
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200)
+    plt.close(fig)
+
+
+def save_mse_csv(rollouts, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(
+            [
+                "start_idx",
+                "mse_total",
+                "mse_g",
+                "mse_r",
+                "mse_i",
+                "n_steps",
+            ]
+        )
+
+        for rollout in rollouts:
+            writer.writerow(
+                [
+                    rollout["start_idx"],
+                    rollout["mse"],
+                    rollout["band_mse"][0],
+                    rollout["band_mse"][1],
+                    rollout["band_mse"][2],
+                    len(rollout["residual"]),
+                ]
+            )
+
+
+def save_step_csv(rollouts, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(
+            [
+                "start_idx",
+                "step",
+                "pred_g",
+                "pred_r",
+                "pred_i",
+                "truth_g",
+                "truth_r",
+                "truth_i",
+                "residual_g",
+                "residual_r",
+                "residual_i",
+                "mse_total",
+                "mse_g",
+                "mse_r",
+                "mse_i",
+            ]
+        )
+
+        for rollout in rollouts:
+            for step in range(len(rollout["pred"])):
+                writer.writerow(
+                    [
+                        rollout["start_idx"],
+                        step,
+                        rollout["pred"][step, 0],
+                        rollout["pred"][step, 1],
+                        rollout["pred"][step, 2],
+                        rollout["truth"][step, 0],
+                        rollout["truth"][step, 1],
+                        rollout["truth"][step, 2],
+                        rollout["residual"][step, 0],
+                        rollout["residual"][step, 1],
+                        rollout["residual"][step, 2],
+                        rollout["step_mses"][step],
+                        rollout["step_band_mses"][step, 0],
+                        rollout["step_band_mses"][step, 1],
+                        rollout["step_band_mses"][step, 2],
+                    ]
+                )
+
+
+def main():
+    args = get_args()
+    run_id = resolve_paths(args)
+
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_gri_model(args.ckpt, device)
+
+    eval_dataset = make_eval_dataset(
+        args=args,
+        context_len=context_len,
+    )
+
+    max_start = max(0, len(eval_dataset) - args.n_future_steps)
+    n_series = min(args.n_series, max_start + 1)
+
+    print("Dataset length:", len(eval_dataset))
+    print("Number of rollout series:", n_series)
+    print("Autoregressive future steps:", args.n_future_steps)
+
+    if n_series <= 0:
+        raise RuntimeError("No rollout series available. Check dataset size and n_future_steps.")
+
+    start_indices = np.linspace(0, max_start, n_series, dtype=int)
+
+    rollouts = []
+
+    for start_idx in start_indices:
+        print(f"Rolling out validation curve starting at index {start_idx}")
+
+        rollout = get_rollout_from_start_gri(
+            dataset=eval_dataset,
+            model=model,
+            device=device,
+            start_idx=int(start_idx),
+            n_future_steps=args.n_future_steps,
+            context_len=context_len,
+        )
+
+        rollouts.append(rollout)
+
+    residual_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_autoreg_gri_residuals_vs_step.png",
+    )
+    series_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_autoreg_gri_multi_series_predictions.png",
+    )
+    hist_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_autoreg_gri_mse_histogram.png",
+    )
+    band_hist_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_autoreg_gri_band_mse_histograms.png",
+    )
+    csv_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_autoreg_gri_mse_by_curve.csv",
+    )
+    step_csv_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_autoreg_gri_step_predictions.csv",
+    )
+
+    plot_residuals_vs_step(rollouts, residual_path)
+    plot_multiple_series_predictions(rollouts, series_path)
+    plot_mse_histogram(rollouts, hist_path)
+    plot_band_mse_histograms(rollouts, band_hist_path)
+    save_mse_csv(rollouts, csv_path)
+    save_step_csv(rollouts, step_csv_path)
+
+    print("Saved:")
+    print(" ", residual_path)
+    print(" ", series_path)
+    print(" ", hist_path)
+    print(" ", band_hist_path)
+    print(" ", csv_path)
+    print(" ", step_csv_path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_new.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_new.py
new file mode 100644
index 00000000..e757d395
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_new.py
@@ -0,0 +1,574 @@
+import argparse
+import csv
+import os
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
+from train_LodeRunner_ddp import (
+    Kilonova_lc_scalar_context_DataSet,
+    ScalarTemporalConditionedLodeRunner,
+)
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+RUN_ID = "021"
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default=f"runs/study_{RUN_ID}/study{RUN_ID}_modelState_epoch0300.pth",
+    )
+    parser.add_argument("--N_imgs", type=int, default=10)
+    parser.add_argument("--n_future_steps", type=int, default=15)
+    parser.add_argument("--n_series", type=int, default=10)
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        default=f"runs/study_{RUN_ID}/autoreg_diagnostics",
+    )
+
+    return parser.parse_args()
+
+class ScalarTemporalConditionedLodeRunner(nn.Module):
+    def __init__(
+        self,
+        backbone: nn.Module,
+        context_len: int = 5,
+        image_size=(1120, 400),
+        n_channels: int = 8,
+        hidden: int = 64,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.context_len = context_len
+        self.image_size = image_size
+        self.n_channels = n_channels
+
+        self.conditioner = nn.Sequential(
+            nn.Linear(2 * context_len, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, n_channels),
+        )
+
+    def forward(self, x, in_vars, out_vars, Dt):
+        B = x.shape[0]
+        H, W = self.image_size
+
+        channel_vals = self.conditioner(x)
+
+        pseudo_img = channel_vals.view(B, self.n_channels, 1, 1).expand(
+            B,
+            self.n_channels,
+            H,
+            W,
+        )
+
+        pred_img = self.backbone(pseudo_img, in_vars, out_vars, Dt)
+        pred_scalar = pred_img.mean(dim=(1, 2, 3))
+
+        return pred_scalar
+
+
+def load_channel_model_new(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    context_len = ckpt.get("context_len", 5)
+
+    backbone = LodeRunner(**model_args).to(device)
+
+    model = ScalarTemporalConditionedLodeRunner(
+        backbone=backbone,
+        context_len=context_len,
+        image_size=model_args["image_size"],
+        n_channels=8,
+        hidden=64,
+    ).to(device)
+
+    state_dict = ckpt["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded scalar-conditioned LodeRunner checkpoint")
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.eval()
+
+    return model, context_len
+
+
+def load_channel_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+    context_len = ckpt.get("context_len", 5)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("predicts_delta:", ckpt.get("predicts_delta", False))
+    print("target_type:", ckpt.get("target_type", "absolute"))
+    print("context_len:", context_len)
+
+    model = LodeRunner(**model_args).to(device)
+
+    state_dict = ckpt["model_state_dict"]
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.noise_scale = noise_scale
+    model.eval()
+
+    return model, context_len
+
+
+def ensure_batch(x):
+    """
+    Dataset item usually has shape:
+        [T, H, W]
+
+    Model expects:
+        [B, T, H, W]
+    """
+    if x.ndim == 3:
+        return x.unsqueeze(0)
+
+    return x
+
+
+def tensor_time_means(x):
+    """
+    Convert image sequence tensor to scalar light curve.
+
+    Supports:
+        [T, H, W]
+        [B, T, H, W]
+        [B, T, C, H, W]
+    """
+    if x.ndim == 3:
+        return x.mean(dim=(1, 2)).detach().cpu().numpy().squeeze()
+
+    if x.ndim == 4:
+        return x.mean(dim=(2, 3)).detach().cpu().numpy().squeeze()
+
+    if x.ndim == 5:
+        return x.mean(dim=(2, 3, 4)).detach().cpu().numpy().squeeze()
+
+    raise ValueError(f"Unexpected tensor shape: {x.shape}")
+
+
+def get_rollout_from_start_new(
+    dataset,
+    model,
+    device,
+    start_idx,
+    n_future_steps,
+    in_vars,
+    out_vars,
+):
+    x0, _, _ = dataset[start_idx]
+
+    x0 = x0.to(torch.float32)
+    context_len = x0.numel() // 2
+
+    mags0 = x0[:context_len].detach().cpu().numpy()
+    rel_t0 = x0[context_len:].detach().cpu().numpy()
+
+    pred_mags = list(mags0)
+    true_mags = list(mags0)
+
+    context_curve = np.asarray(mags0)
+
+    pred_curve = []
+    truth_curve = []
+    residual_curve = []
+    step_mses = []
+
+    with torch.no_grad():
+        for step in range(n_future_steps):
+            future_idx = start_idx + step
+
+            if future_idx >= len(dataset):
+                break
+
+            x_true_step, target_delta, future_Dt = dataset[future_idx]
+
+            x_true_step = x_true_step.to(torch.float32)
+            target_delta = torch.as_tensor(
+                target_delta,
+                dtype=torch.float32,
+                device=device,
+            )
+            future_Dt = torch.as_tensor(
+                future_Dt,
+                dtype=torch.float32,
+                device=device,
+            )
+
+            if target_delta.ndim == 0:
+                target_delta = target_delta.unsqueeze(0)
+
+            if future_Dt.ndim == 0:
+                future_Dt = future_Dt.unsqueeze(0)
+
+            # Build autoregressive scalar input from predicted history.
+            current_pred_mags = np.asarray(pred_mags[-context_len:], dtype=np.float32)
+
+            # Use the relative-time pattern from the dataset item for this step.
+            current_rel_t = x_true_step[context_len:].detach().cpu().numpy().astype(np.float32)
+
+            x_pred = np.concatenate([current_pred_mags, current_rel_t], axis=0)
+            x_pred = torch.tensor(
+                x_pred,
+                dtype=torch.float32,
+                device=device,
+            ).unsqueeze(0)
+
+            pred_delta = model(x_pred, in_vars, out_vars, future_Dt)
+            pred_delta = pred_delta.view_as(target_delta)
+
+            pred_next = pred_mags[-1] + pred_delta.item()
+            true_next = true_mags[-1] + target_delta.item()
+
+            residual_scalar = pred_next - true_next
+            step_mse = residual_scalar ** 2
+
+            pred_curve.append(pred_next)
+            truth_curve.append(true_next)
+            residual_curve.append(residual_scalar)
+            step_mses.append(step_mse)
+
+            pred_mags.append(pred_next)
+            true_mags.append(true_next)
+
+    pred_curve = np.asarray(pred_curve)
+    truth_curve = np.asarray(truth_curve)
+    residual_curve = np.asarray(residual_curve)
+    step_mses = np.asarray(step_mses)
+
+    total_mse = np.mean(step_mses) if len(step_mses) > 0 else np.nan
+
+    return {
+        "start_idx": start_idx,
+        "context": context_curve,
+        "pred": pred_curve,
+        "truth": truth_curve,
+        "residual": residual_curve,
+        "step_mses": step_mses,
+        "mse": total_mse,
+    }
+
+
+def get_rollout_from_start(
+    dataset,
+    model,
+    device,
+    start_idx,
+    n_future_steps,
+    in_vars,
+    out_vars,
+):
+    """
+    Clean autoregressive rollout.
+
+    x_pred:
+        model-generated autoregressive context
+
+    x_true:
+        ground-truth context used only to reconstruct true future values
+
+    This avoids the bug where truth was reconstructed using the predicted
+    previous frame.
+    """
+
+    context_img, _, _ = dataset[start_idx]
+
+    x_pred = ensure_batch(context_img).to(device)
+    x_true = ensure_batch(context_img).to(device)
+
+    context_curve = tensor_time_means(x_true)
+
+    pred_curve = []
+    truth_curve = []
+    residual_curve = []
+    step_mses = []
+
+    with torch.no_grad():
+        for step in range(n_future_steps):
+            future_idx = start_idx + step
+
+            if future_idx >= len(dataset):
+                break
+
+            _, future_target_delta, future_Dt = dataset[future_idx]
+
+            future_target_delta = ensure_batch(future_target_delta).to(device)
+
+            future_Dt = torch.as_tensor(
+                future_Dt,
+                dtype=torch.float32,
+                device=device,
+            )
+
+            if future_Dt.ndim == 0:
+                future_Dt = future_Dt.unsqueeze(0)
+
+            pred_delta_img = model(x_pred, in_vars, out_vars, future_Dt)
+
+            pred_last_img = x_pred[:, -1:]
+            true_last_img = x_true[:, -1:]
+
+            pred_next_img = pred_last_img + pred_delta_img[:, -1:]
+            true_next_img = true_last_img + future_target_delta[:, -1:]
+
+            pred_scalar = pred_next_img.mean().item()
+            true_scalar = true_next_img.mean().item()
+            residual_scalar = pred_scalar - true_scalar
+
+            step_mse = torch.mean((pred_next_img - true_next_img) ** 2).item()
+
+            pred_curve.append(pred_scalar)
+            truth_curve.append(true_scalar)
+            residual_curve.append(residual_scalar)
+            step_mses.append(step_mse)
+
+            # Autoregressive model context gets the prediction.
+            x_pred = torch.cat(
+                [x_pred[:, 1:], pred_next_img.detach()],
+                dim=1,
+            )
+
+            # Truth context gets the independently reconstructed truth.
+            x_true = torch.cat(
+                [x_true[:, 1:], true_next_img.detach()],
+                dim=1,
+            )
+
+    pred_curve = np.asarray(pred_curve)
+    truth_curve = np.asarray(truth_curve)
+    residual_curve = np.asarray(residual_curve)
+    step_mses = np.asarray(step_mses)
+
+    total_mse = np.mean(step_mses) if len(step_mses) > 0 else np.nan
+
+    return {
+        "start_idx": start_idx,
+        "context": context_curve,
+        "pred": pred_curve,
+        "truth": truth_curve,
+        "residual": residual_curve,
+        "step_mses": step_mses,
+        "mse": total_mse,
+    }
+
+
+def plot_residuals_vs_step(rollouts, outpath):
+    plt.figure(figsize=(8, 5))
+
+    for rollout in rollouts:
+        steps = np.arange(len(rollout["residual"]))
+        plt.plot(
+            steps,
+            rollout["residual"],
+            marker="o",
+            alpha=0.75,
+            label=f"start {rollout['start_idx']}",
+        )
+
+    plt.axhline(0.0, linestyle="--", linewidth=1)
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Residual: prediction - truth")
+    plt.title("Autoregressive residuals vs time step")
+    plt.legend(fontsize=8, ncol=2)
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def plot_multiple_series_predictions(rollouts, outpath):
+    plt.figure(figsize=(9, 6))
+
+    for rollout in rollouts:
+        start_idx = rollout["start_idx"]
+
+        context_steps = np.arange(-len(rollout["context"]), 0)
+        future_steps = np.arange(len(rollout["pred"]))
+
+        plt.plot(
+            context_steps,
+            rollout["context"],
+            linestyle=":",
+            alpha=0.45,
+        )
+
+        plt.plot(
+            future_steps,
+            rollout["truth"],
+            linewidth=1.5,
+            alpha=0.75,
+            label=f"truth start {start_idx}",
+        )
+
+        plt.plot(
+            future_steps,
+            rollout["pred"],
+            linestyle="--",
+            linewidth=1.5,
+            alpha=0.75,
+            label=f"pred start {start_idx}",
+        )
+
+    plt.axvline(-0.5, linewidth=1, alpha=0.5)
+    plt.gca().invert_yaxis()
+    plt.xlabel("Time step relative to forecast start")
+    plt.ylabel("Normalized magnitude")
+    plt.title("Autoregressive predictions for validation curves")
+    plt.legend(fontsize=7, ncol=2)
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def plot_mse_histogram(rollouts, outpath):
+    mses = np.asarray([r["mse"] for r in rollouts])
+    mses = mses[np.isfinite(mses)]
+
+    plt.figure(figsize=(7, 5))
+    plt.hist(mses, bins=min(10, max(1, len(mses))))
+    plt.xlabel("Mean autoregressive MSE per validation curve")
+    plt.ylabel("Count")
+    plt.title("Distribution of autoregressive rollout MSEs")
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def save_mse_csv(rollouts, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["start_idx", "mse", "n_steps"])
+
+        for rollout in rollouts:
+            writer.writerow(
+                [
+                    rollout["start_idx"],
+                    rollout["mse"],
+                    len(rollout["residual"]),
+                ]
+            )
+
+
+def main():
+    args = get_args()
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_channel_model_new(args.ckpt, device)
+
+    #eval_dataset = Kilonova_lc_img_DataSet_channels_context(
+    #    half_image=False,
+    #    N_imgs=args.N_imgs,
+    #    context_len=context_len,
+    #)
+
+    eval_dataset = Kilonova_lc_scalar_context_DataSet(
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    max_start = max(0, len(eval_dataset) - args.n_future_steps)
+    n_series = min(args.n_series, max_start + 1)
+
+    print("Dataset length:", len(eval_dataset))
+    print("Number of rollout series:", n_series)
+    print("Autoregressive future steps:", args.n_future_steps)
+
+    start_indices = np.linspace(0, max_start, n_series, dtype=int)
+
+    rollouts = []
+
+    for start_idx in start_indices:
+        print(f"Rolling out validation curve starting at index {start_idx}")
+
+        rollout = get_rollout_from_start_new(
+            dataset=eval_dataset,
+            model=model,
+            device=device,
+            start_idx=int(start_idx),
+            n_future_steps=args.n_future_steps,
+            in_vars=in_vars,
+            out_vars=out_vars,
+        )
+
+        rollouts.append(rollout)
+
+    residual_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_residuals_vs_step.png",
+    )
+    series_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_multi_series_predictions.png",
+    )
+    hist_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_mse_histogram.png",
+    )
+    csv_path = os.path.join(
+        args.outdir,
+        f"study{RUN_ID}_autoreg_mse_by_curve.csv",
+    )
+
+    plot_residuals_vs_step(rollouts, residual_path)
+    plot_multiple_series_predictions(rollouts, series_path)
+    plot_mse_histogram(rollouts, hist_path)
+    save_mse_csv(rollouts, csv_path)
+
+    print("Saved:")
+    print(" ", residual_path)
+    print(" ", series_path)
+    print(" ", hist_path)
+    print(" ", csv_path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_gri.py b/applications/harnesses/KN_loderunner/plot_pred_gri.py
new file mode 100644
index 00000000..20ca2d93
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_gri.py
@@ -0,0 +1,518 @@
+import argparse
+import os
+
+import numpy as np
+import torch
+import matplotlib
+import matplotlib.pyplot as plt
+from torch.utils.data import DataLoader
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+
+from train_LodeRunner_ddp import (
+    Kilonova_lc_scalar_context_DataSet_gri,
+    ScalarTemporalConditionedLodeRunner_gri,
+    load_or_compute_band_normalization,
+)
+
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+
+BAND_KEYS = ("arr_ztfg", "arr_ztfr", "arr_ztfi")
+BAND_NAMES = ("g", "r", "i")
+VALUE_COL = 1
+
+
+def study_tag(study):
+    return f"{int(study):03d}"
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description="Plot one-step and autoregressive predictions for scalar GRI LodeRunner."
+    )
+
+    parser.add_argument("--study", type=int, default=24)
+    parser.add_argument("--epoch", type=int, default=500)
+    parser.add_argument("--ckpt", type=str, default=None)
+
+    parser.add_argument("--N_imgs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--n_future_steps", type=int, default=15)
+
+    parser.add_argument(
+        "--norm_stats_path",
+        type=str,
+        default="kilonova_gri_norm_stats.npz",
+    )
+
+    parser.add_argument("--outdir", type=str, default=None)
+
+    return parser.parse_args()
+
+
+def resolve_paths(args):
+    run_id = study_tag(args.study)
+
+    if args.ckpt is None:
+        args.ckpt = (
+            f"runs/study_{run_id}/study{run_id}_modelState_epoch{args.epoch:04d}.pth"
+        )
+
+    if args.outdir is None:
+        args.outdir = f"runs/study_{run_id}/pred_plots_gri"
+
+    return run_id
+
+
+def strip_ddp_prefix(state_dict):
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        return {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+    return state_dict
+
+
+def load_gri_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    context_len = ckpt.get("context_len", 5)
+
+    n_input_channels = ckpt.get("n_input_channels", 3)
+    n_output_channels = ckpt.get("n_output_channels", 3)
+    backbone_channels = ckpt.get("backbone_channels", 8)
+    hidden = ckpt.get("hidden", 64)
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("model_class:", ckpt.get("model_class", "unknown"))
+    print("backbone_class:", ckpt.get("backbone_class", "LodeRunner"))
+    print("predicts_delta:", ckpt.get("predicts_delta", False))
+    print("target_type:", ckpt.get("target_type", "unknown"))
+    print("context_len:", context_len)
+    print("n_input_channels:", n_input_channels)
+    print("n_output_channels:", n_output_channels)
+    print("backbone_channels:", backbone_channels)
+    print("hidden:", hidden)
+
+    backbone = LodeRunner(**model_args).to(device)
+    backbone.noise_scale = noise_scale
+
+    model = ScalarTemporalConditionedLodeRunner_gri(
+        backbone=backbone,
+        context_len=context_len,
+        n_input_channels=n_input_channels,
+        n_output_channels=n_output_channels,
+        image_size=model_args["image_size"],
+        backbone_channels=backbone_channels,
+        hidden=hidden,
+    ).to(device)
+
+    state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded ScalarTemporalConditionedLodeRunner_gri checkpoint")
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.eval()
+
+    return model, context_len
+
+
+def make_eval_dataset(args, context_len):
+    band_means, band_stds = load_or_compute_band_normalization(
+        stats_path=args.norm_stats_path,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+    )
+
+    print("Using band normalization:")
+    print("band_means:", band_means)
+    print("band_stds:", band_stds)
+
+    dataset = Kilonova_lc_scalar_context_DataSet_gri(
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+        means=band_means,
+        stds=band_stds,
+        predicts_delta=True,
+    )
+
+    return dataset
+
+
+def split_gri_context(x, context_len, n_bands=3):
+    """
+    Dataset x layout:
+        [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK, t0, t1, ..., tK]
+
+    Returns:
+        values: [context_len, 3]
+        rel_t:  [context_len]
+    """
+    x = torch.as_tensor(x, dtype=torch.float32)
+
+    value_count = context_len * n_bands
+    values = x[:value_count].detach().cpu().numpy().reshape(context_len, n_bands)
+    rel_t = x[value_count:].detach().cpu().numpy()
+
+    return values.astype(np.float32), rel_t.astype(np.float32)
+
+
+def build_gri_input(values, rel_t, device):
+    """
+    values:
+        [context_len, 3]
+    rel_t:
+        [context_len]
+
+    Returns:
+        x: [1, context_len * 3 + context_len]
+    """
+    values = np.asarray(values, dtype=np.float32)
+    rel_t = np.asarray(rel_t, dtype=np.float32)
+
+    x = np.concatenate(
+        [
+            values.reshape(-1),
+            rel_t,
+        ],
+        axis=0,
+    )
+
+    return torch.tensor(
+        x,
+        dtype=torch.float32,
+        device=device,
+    ).unsqueeze(0)
+
+
+def plot_onestep_predictions(
+    idxs,
+    preds,
+    targets,
+    prefix,
+    outpath,
+):
+    preds = np.asarray(preds, dtype=np.float32)
+    targets = np.asarray(targets, dtype=np.float32)
+    prefix = np.asarray(prefix, dtype=np.float32)
+
+    fig, axes = plt.subplots(3, 1, figsize=(7, 10), sharex=True)
+
+    for band_idx, band_name in enumerate(BAND_NAMES):
+        ax = axes[band_idx]
+
+        ax.scatter(
+            idxs,
+            preds[:, band_idx],
+            label="Predicted next magnitude",
+        )
+
+        ax.scatter(
+            idxs,
+            targets[:, band_idx],
+            label="True next magnitude",
+        )
+
+        ax.scatter(
+            np.arange(len(prefix)) - len(prefix),
+            prefix[:, band_idx],
+            label="Initial context window",
+        )
+
+        ax.invert_yaxis()
+        ax.set_ylabel(f"{band_name} norm mag")
+        ax.set_title(f"{band_name}-band one-step prediction")
+
+    axes[-1].set_xlabel("Sample index")
+
+    handles, labels = axes[0].get_legend_handles_labels()
+    fig.legend(
+        handles,
+        labels,
+        fontsize=8,
+        ncol=1,
+        loc="upper center",
+        bbox_to_anchor=(0.5, 1.02),
+    )
+
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_autoreg_rollout(
+    idxs_seq,
+    preds_seq,
+    truth_seq,
+    prefix,
+    outpath,
+):
+    preds_seq = np.asarray(preds_seq, dtype=np.float32)
+    truth_seq = np.asarray(truth_seq, dtype=np.float32)
+    prefix = np.asarray(prefix, dtype=np.float32)
+
+    fig, axes = plt.subplots(3, 1, figsize=(7, 10), sharex=True)
+
+    for band_idx, band_name in enumerate(BAND_NAMES):
+        ax = axes[band_idx]
+
+        ax.scatter(
+            idxs_seq,
+            preds_seq[:, band_idx],
+            label="Autoregressive predictions",
+        )
+
+        ax.scatter(
+            idxs_seq,
+            truth_seq[:, band_idx],
+            label="Truth",
+        )
+
+        ax.scatter(
+            np.arange(len(prefix)) - len(prefix),
+            prefix[:, band_idx],
+            label="Initial context window",
+        )
+
+        ax.invert_yaxis()
+        ax.set_ylabel(f"{band_name} norm mag")
+        ax.set_title(f"{band_name}-band autoregressive rollout")
+
+    axes[-1].set_xlabel("Autoregressive step")
+
+    handles, labels = axes[0].get_legend_handles_labels()
+    fig.legend(
+        handles,
+        labels,
+        fontsize=8,
+        ncol=1,
+        loc="upper center",
+        bbox_to_anchor=(0.5, 1.02),
+    )
+
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def main():
+    args = get_args()
+    run_id = resolve_paths(args)
+
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_gri_model(args.ckpt, device)
+
+    eval_dataset = make_eval_dataset(
+        args=args,
+        context_len=context_len,
+    )
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+    )
+
+    # ------------------------------------------------------------
+    # One-step predictions
+    # ------------------------------------------------------------
+    preds = []
+    targets = []
+    idxs = []
+    prefix = None
+
+    for idx, (x, target_delta, Dt) in enumerate(loader):
+        if x.shape[0] != 1:
+            raise ValueError(
+                "This plotting script assumes batch_size=1 so each point maps "
+                "cleanly to a single light curve window."
+            )
+
+        x = x.to(torch.float32).to(device)
+        target_delta = target_delta.to(torch.float32).to(device)
+        Dt = Dt.to(torch.float32).to(device)
+
+        if prefix is None:
+            context_vals, _ = split_gri_context(
+                x[0].detach().cpu(),
+                context_len=context_len,
+                n_bands=3,
+            )
+            prefix = context_vals
+
+        with torch.no_grad():
+            pred_delta = model(
+                x,
+                in_vars=None,
+                out_vars=None,
+                Dt=Dt,
+            )
+
+        context_vals, _ = split_gri_context(
+            x[0].detach().cpu(),
+            context_len=context_len,
+            n_bands=3,
+        )
+
+        last_vals = torch.tensor(
+            context_vals[-1],
+            dtype=torch.float32,
+            device=device,
+        )
+
+        pred_next = last_vals + pred_delta[0].reshape(3)
+        true_next = last_vals + target_delta[0].reshape(3)
+
+        preds.append(pred_next.detach().cpu().numpy())
+        targets.append(true_next.detach().cpu().numpy())
+        idxs.append(idx)
+
+    onestep_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_pred_vs_truth_gri_delta_onestep.png",
+    )
+
+    plot_onestep_predictions(
+        idxs=idxs,
+        preds=preds,
+        targets=targets,
+        prefix=prefix,
+        outpath=onestep_path,
+    )
+
+    # ------------------------------------------------------------
+    # Autoregressive rollout
+    # ------------------------------------------------------------
+    if len(eval_dataset) == 0:
+        raise RuntimeError("Evaluation dataset is empty.")
+
+    x0, _, _ = eval_dataset[0]
+
+    context_vals, rel_t0 = split_gri_context(
+        x0,
+        context_len=context_len,
+        n_bands=3,
+    )
+
+    pred_vals = [row.copy() for row in context_vals]
+    true_vals = [row.copy() for row in context_vals]
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    with torch.no_grad():
+        for step in range(args.n_future_steps):
+            if step >= len(eval_dataset):
+                break
+
+            x_true_step, target_delta, future_Dt = eval_dataset[step]
+
+            x_true_step = torch.as_tensor(x_true_step, dtype=torch.float32)
+            target_delta = torch.as_tensor(
+                target_delta,
+                dtype=torch.float32,
+                device=device,
+            ).reshape(3)
+
+            future_Dt = torch.as_tensor(
+                future_Dt,
+                dtype=torch.float32,
+                device=device,
+            )
+
+            if future_Dt.ndim == 0:
+                future_Dt = future_Dt.unsqueeze(0)
+
+            _, current_rel_t = split_gri_context(
+                x_true_step,
+                context_len=context_len,
+                n_bands=3,
+            )
+
+            current_pred_vals = np.asarray(
+                pred_vals[-context_len:],
+                dtype=np.float32,
+            )
+
+            x_pred = build_gri_input(
+                values=current_pred_vals,
+                rel_t=current_rel_t,
+                device=device,
+            )
+
+            pred_delta = model(
+                x_pred,
+                in_vars=None,
+                out_vars=None,
+                Dt=future_Dt,
+            ).reshape(3)
+
+            pred_next = (
+                torch.tensor(
+                    pred_vals[-1],
+                    dtype=torch.float32,
+                    device=device,
+                )
+                + pred_delta
+            )
+
+            true_next = (
+                torch.tensor(
+                    true_vals[-1],
+                    dtype=torch.float32,
+                    device=device,
+                )
+                + target_delta
+            )
+
+            pred_next_np = pred_next.detach().cpu().numpy()
+            true_next_np = true_next.detach().cpu().numpy()
+
+            preds_seq.append(pred_next_np)
+            truth_seq.append(true_next_np)
+            idxs_seq.append(step)
+
+            pred_vals.append(pred_next_np)
+            true_vals.append(true_next_np)
+
+    autoreg_path = os.path.join(
+        args.outdir,
+        f"study{run_id}_pred_vs_truth_gri_delta_autoreg_clean.png",
+    )
+
+    plot_autoreg_rollout(
+        idxs_seq=idxs_seq,
+        preds_seq=preds_seq,
+        truth_seq=truth_seq,
+        prefix=context_vals,
+        outpath=autoreg_path,
+    )
+
+    print("Saved:")
+    print(" ", onestep_path)
+    print(" ", autoreg_path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_seq_context.py b/applications/harnesses/KN_loderunner/plot_pred_seq_context.py
new file mode 100644
index 00000000..15e6fea7
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_seq_context.py
@@ -0,0 +1,229 @@
+import argparse
+import numpy as np
+import torch
+import matplotlib
+import matplotlib.pyplot as plt
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from torch.utils.data import DataLoader
+
+from train_LodeRunner_ddp import (
+    Kilonova_lc_img_DataSet_seq_context,
+    TemporalLodeRunner,
+)
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (6, 6)
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        #default="runs/study_007/study007_modelState_epoch0100.pth",
+        default="runs/study_008/study008_modelState_epoch0100.pth",
+    )
+    parser.add_argument("--N_imgs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--n_future_steps", type=int, default=10)
+
+    return parser.parse_args()
+
+
+def load_temporal_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    context_len = ckpt["context_len"]
+    hidden_channels = ckpt["hidden_channels"]
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    base_model = LodeRunner(**model_args)
+
+    model = TemporalLodeRunner(
+        backbone=base_model,
+        in_channels=8,
+        context_len=context_len,
+        hidden_channels=hidden_channels,
+    )
+
+    model.to(device)
+
+    state_dict = ckpt["model_state_dict"]
+
+    if all(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+    print("Loaded model_args:", model_args)
+    print("context_len:", context_len)
+    print("hidden_channels:", hidden_channels)
+
+    model.backbone.noise_scale = noise_scale
+    model.eval()
+
+    return model, context_len
+
+
+def main():
+    args = get_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len = load_temporal_model(args.ckpt, device)
+
+    eval_dataset = Kilonova_lc_img_DataSet_seq_context(
+        half_image=False,
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+    )
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+    )
+
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
+
+    # ------------------------------------------------------------
+    # One-step predictions using true context windows
+    # ------------------------------------------------------------
+    preds = []
+    targets = []
+    idxs = []
+    prefix = []
+
+    for idx, (context_seq, target, Dt) in enumerate(loader):
+        context_seq = context_seq.to(device)
+        if idx == 0:
+            context_means = context_seq.mean(dim=(2, 3, 4))[0].detach().cpu().numpy()
+            for context in context_means:
+                prefix.append(context.mean().item())
+        target = target.to(device)
+        Dt = Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_image = model(context_seq, in_vars, out_vars, Dt)
+
+        preds.append(pred_image.mean().item())
+        targets.append(target.mean().item())
+        idxs.append(idx)
+
+    plt.figure()
+    plt.scatter(idxs, preds, label="Predictions")
+    plt.scatter(idxs, targets, label="Truth")
+    plt.scatter(np.arange(len(prefix))-(len(prefix)), prefix, label='Initial Context Window')
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Sample index")
+    plt.ylabel("Mean magnitude/image value")
+    plt.tight_layout()
+    plt.savefig("pred_vs_truth_seq_context.png", dpi=200)
+
+    print("Saved pred_vs_truth_seq_context.png")
+
+
+    context_seq, target, Dt = next(iter(loader))
+
+    context_seq = context_seq.to(device)
+    Dt = Dt.to(torch.float32).to(device)
+
+    x = context_seq
+
+    preds_seq = []
+    truth_seq = []
+    idxs_seq = []
+
+    future_iter = iter(loader)
+
+    for step in range(args.n_future_steps):
+        try:
+            _, future_target, future_Dt = next(future_iter)
+        except StopIteration:
+            break
+
+        future_target = future_target.to(device)
+        future_Dt = future_Dt.to(torch.float32).to(device)
+
+        with torch.no_grad():
+            pred_image = model(x, in_vars, out_vars, future_Dt)
+
+        preds_seq.append(pred_image.mean().item())
+        truth_seq.append(future_target.mean().item())
+        idxs_seq.append(step)
+
+        # autoregressive update: append prediction
+        x = torch.cat([x[:, 1:], pred_image.unsqueeze(1)], dim=1)
+
+    plt.figure()
+
+    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
+    plt.scatter(idxs_seq, truth_seq, label="Truth")
+    plt.scatter(
+        np.arange(len(prefix)) - len(prefix),
+        prefix,
+        label="Initial context window",
+    )
+
+    plt.legend()
+    plt.gca().invert_yaxis()
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("Mean magnitude/image value")
+    plt.tight_layout()
+    plt.savefig("pred_vs_truth_seq_context_autoreg.png", dpi=200)
+
+
+    # ------------------------------------------------------------
+    # Image comparison for the final one-step batch above
+    # ------------------------------------------------------------
+    pred_plot = pred_image.squeeze().mean(dim=0).detach().cpu().numpy()
+    true_plot = target.squeeze().mean(dim=0).detach().cpu().numpy()
+    error_plot = pred_plot - true_plot
+
+    vmin = min(pred_plot.min(), true_plot.min())
+    vmax = max(pred_plot.max(), true_plot.max())
+    err_max = np.max(np.abs(error_plot))
+
+    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
+
+    im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
+    ax1.set_title("Prediction")
+
+    ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
+    ax2.set_title("Truth")
+
+    im3 = ax3.imshow(error_plot, origin="lower", vmin=-err_max, vmax=err_max)
+    ax3.set_title("Error (Pred - Truth)")
+
+    cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
+    cbar.set_label("Field value")
+
+    cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
+    cbar_err.set_label("Error")
+
+    for ax in (ax1, ax2, ax3):
+        ax.axis("off")
+
+    plt.tight_layout()
+    plt.savefig("img_comp_test_new.png", bbox_inches="tight", dpi=200)
+
+    print("Saved img_comp_test_new.png")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
new file mode 100644
index 00000000..2cf5bd5e
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -0,0 +1,1086 @@
+import os
+import time
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim.lr_scheduler import LambdaLR
+
+from yoke.models.vit.swin.bomberman import LodeRunner
+from yoke.datasets.lsc_dataset import LSC_rho2rho_temporal_DataSet
+from yoke.utils.training.epoch.loderunner import train_DDP_scalar_temporal_loderunner_epoch_gri
+from yoke.utils.restart import continuation_setup
+from yoke.utils.dataload import make_distributed_dataloader
+from yoke.utils.checkpointing import load_model_and_optimizer
+from yoke.utils.checkpointing import save_model_and_optimizer
+from yoke.lr_schedulers import CosineWithWarmupScheduler
+from yoke.helpers import cli
+
+# FIXME remove if restructure
+from torch.utils.data import Dataset, DataLoader, random_split
+import glob
+import random
+
+#MEAN = 24.694652705328807
+#STD = 4.67030961432848
+
+GLOBAL_GMAG_MEAN = 24.694652705328807
+GLOBAL_GMAG_STD  = 4.67030961432848
+EPS = 1e-6
+
+#############################################
+# Inputs
+#############################################
+descr_str = (
+    "Uses DDP to train LodeRunner architecture on single-timstep input and output "
+    "of the lsc240420 per-material density fields."
+)
+parser = argparse.ArgumentParser(
+    prog="DDP LodeRunner Training", description=descr_str, fromfile_prefix_chars="@"
+)
+parser = cli.add_default_args(parser=parser)
+parser = cli.add_filepath_args(parser=parser)
+parser = cli.add_computing_args(parser=parser)
+parser = cli.add_model_args(parser=parser)
+parser = cli.add_training_args(parser=parser)
+parser = cli.add_cosine_lr_scheduler_args(parser=parser)
+
+# DPOT‐style noise parameter
+parser.add_argument(
+    "--noise_scale",
+    type=float,
+    default=0.0,
+    help="Relative magnitude ε for Gaussian noise injection (e.g. 5e-5).",
+)
+
+# Change some default filepaths.
+parser.set_defaults(
+    train_filelist="lsc240420_prefixes_train_80pct.txt",
+    validation_filelist="lsc240420_prefixes_validation_10pct.txt",
+    test_filelist="lsc240420_prefixes_test_10pct.txt",
+)
+
+
+def compute_band_normalization(
+    file_prefix_list,
+    band_keys=("arr_ztfg", "arr_ztfr", "arr_ztfi"),
+    value_col=1,
+    stats_path="kilonova_gri_norm_stats.npz",
+):
+    """
+    Compute global per-band mean/std over the training files only.
+
+    Saves:
+        means: shape [3]
+        stds:  shape [3]
+    """
+
+    sums = np.zeros(len(band_keys), dtype=np.float64)
+    sums_sq = np.zeros(len(band_keys), dtype=np.float64)
+    counts = np.zeros(len(band_keys), dtype=np.float64)
+
+    for fn in file_prefix_list:
+        data = np.load(fn, allow_pickle=True)
+
+        for b, key in enumerate(band_keys):
+            vals = data[key][:, value_col].astype(np.float64)
+
+            finite = np.isfinite(vals)
+            vals = vals[finite]
+
+            sums[b] += vals.sum()
+            sums_sq[b] += np.square(vals).sum()
+            counts[b] += vals.size
+
+        data.close()
+
+    means = sums / counts
+    variances = sums_sq / counts - means**2
+    variances = np.maximum(variances, 1e-12)
+    stds = np.sqrt(variances)
+
+    means = means.astype(np.float32)
+    stds = stds.astype(np.float32)
+
+    np.savez(
+        stats_path,
+        means=means,
+        stds=stds,
+        band_keys=np.array(band_keys),
+        value_col=value_col,
+    )
+
+    print("Saved normalization stats:", stats_path)
+    print("means:", means)
+    print("stds:", stds)
+
+    return means, stds
+
+
+def load_or_compute_band_normalization(
+    stats_path="kilonova_gri_norm_stats.npz",
+    band_keys=("arr_ztfg", "arr_ztfr", "arr_ztfi"),
+    value_col=1,
+):
+    file_prefix_list = sorted(
+        glob.glob(
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz"
+        )
+    )
+
+    if os.path.exists(stats_path):
+        stats = np.load(stats_path, allow_pickle=True)
+        means = stats["means"].astype(np.float32)
+        stds = stats["stds"].astype(np.float32)
+        stats.close()
+
+        print("Loaded normalization stats:", stats_path)
+        print("means:", means)
+        print("stds:", stds)
+
+        return means, stds
+
+    return compute_band_normalization(
+        file_prefix_list=file_prefix_list,
+        band_keys=band_keys,
+        value_col=value_col,
+        stats_path=stats_path,
+    )
+
+
+class Kilonova_lc_scalar_context_DataSet_gri(Dataset):
+    def __init__(
+        self,
+        N_imgs=0,
+        context_len=5,
+        band_keys=("arr_ztfg", "arr_ztfr", "arr_ztfi"),
+        value_col=1,
+        means=None,
+        stds=None,
+        predicts_delta=True,
+    ):
+        file_prefix_list = sorted(
+            glob.glob(
+                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+                "uniform_dataset_20000/lc_*.npz"
+            )
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(
+                np.random.choice(file_prefix_list, N_imgs, replace=False)
+            )
+
+        random.shuffle(self.file_prefix_list)
+
+        self.context_len = context_len
+        self.band_keys = tuple(band_keys)
+        self.value_col = value_col
+        self.n_channels = len(self.band_keys)
+        self.predicts_delta = predicts_delta
+
+        if means is None:
+            raise ValueError(
+                "means must be provided for per-band normalization. "
+                "Expected shape [n_channels], e.g. [g_mean, r_mean, i_mean]."
+            )
+
+        if stds is None:
+            raise ValueError(
+                "stds must be provided for per-band normalization. "
+                "Expected shape [n_channels], e.g. [g_std, r_std, i_std]."
+            )
+
+        self.means = np.asarray(means, dtype=np.float32)
+        self.stds = np.asarray(stds, dtype=np.float32)
+
+        if self.means.shape[0] != self.n_channels:
+            raise ValueError(
+                f"means has length {self.means.shape[0]}, "
+                f"but n_channels={self.n_channels}"
+            )
+
+        if self.stds.shape[0] != self.n_channels:
+            raise ValueError(
+                f"stds has length {self.stds.shape[0]}, "
+                f"but n_channels={self.n_channels}"
+            )
+
+        if np.any(self.stds <= 0):
+            raise ValueError(f"All stds must be positive. Got stds={self.stds}")
+
+        self.samples = []
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+
+            # Assume all bands share the same time grid.
+            mjd = data[self.band_keys[0]][:, 0]
+            n_times = len(mjd)
+
+            data.close()
+
+            max_start = n_times - context_len - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((file_idx, startIDX))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        file_idx, startIDX = self.samples[index]
+        fn = self.file_prefix_list[file_idx]
+
+        data = np.load(fn, allow_pickle=True)
+
+        # Time grid from the first band.
+        # Assumes arr_ztfg, arr_ztfr, arr_ztfi have matching MJD columns.
+        mjd = data[self.band_keys[0]][:, 0].astype(np.float32)
+
+        # Stack one scalar value column from each band.
+        # vals shape: [T, 3] for g/r/i.
+        vals = np.stack(
+            [
+                data[key][:, self.value_col].astype(np.float32)
+                for key in self.band_keys
+            ],
+            axis=1,
+        )
+
+        data.close()
+
+        # Per-band normalization.
+        # vals[:, 0] = normalized g
+        # vals[:, 1] = normalized r
+        # vals[:, 2] = normalized i
+        vals = (vals - self.means[None, :]) / (self.stds[None, :] + EPS)
+
+        t0 = mjd.min()
+        t_obs = mjd - t0
+
+        target_idx = startIDX + self.context_len
+
+        # Context values shape: [context_len, 3]
+        context_vals = vals[startIDX:target_idx]
+
+        # Relative observation times shape: [context_len]
+        rel_t = t_obs[startIDX:target_idx] - t_obs[startIDX]
+        rel_t = rel_t.astype(np.float32)
+
+        # Flatten context as:
+        # [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK]
+        context_flat = context_vals.reshape(-1).astype(np.float32)
+
+        # Final input shape:
+        # [(3 * context_len) + context_len]
+        #
+        # For context_len=5:
+        # x.shape == [20]
+        x = np.concatenate([context_flat, rel_t], axis=0)
+        x = torch.tensor(x, dtype=torch.float32)
+
+        if self.predicts_delta:
+            # Predict normalized delta for each band:
+            # [delta_g, delta_r, delta_i]
+            target_vals = vals[target_idx] - vals[target_idx - 1]
+        else:
+            # Predict normalized absolute next value:
+            # [g_next, r_next, i_next]
+            target_vals = vals[target_idx]
+
+        target = torch.tensor(target_vals, dtype=torch.float32)
+
+        Dt = torch.tensor(
+            t_obs[target_idx] - t_obs[target_idx - 1],
+            dtype=torch.float32,
+        )
+
+        return x, target, Dt
+
+
+class Kilonova_lc_scalar_context_DataSet(Dataset):
+    def __init__(
+        self,
+        N_imgs=0,
+        context_len=5,
+    ):
+        file_prefix_list = sorted(
+            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
+
+        random.shuffle(self.file_prefix_list)
+
+        self.context_len = context_len
+        self.samples = []
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+            mjd = data["arr_ztfg"][:, 0]
+            n_times = len(mjd)
+            data.close()
+
+            max_start = n_times - context_len - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((file_idx, startIDX))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        file_idx, startIDX = self.samples[index]
+        fn = self.file_prefix_list[file_idx]
+
+        data = np.load(fn, allow_pickle=True)
+        arr = data["arr_ztfg"]
+
+        mjd = arr[:, 0]
+        g_mag = arr[:, 1].astype(np.float32)
+
+        g_mag = (g_mag - GLOBAL_GMAG_MEAN) / (GLOBAL_GMAG_STD + EPS)
+
+        t0 = mjd.min()
+        t_obs = mjd - t0
+
+        target_idx = startIDX + self.context_len
+
+        mags = g_mag[startIDX:target_idx].astype(np.float32)
+
+        # context_len values; first dt is 0, remaining are relative times
+        rel_t = t_obs[startIDX:target_idx] - t_obs[startIDX]
+        rel_t = rel_t.astype(np.float32)
+
+        x = np.concatenate([mags, rel_t], axis=0)
+        x = torch.tensor(x, dtype=torch.float32)
+
+        delta_mag = g_mag[target_idx] - g_mag[target_idx - 1]
+        target = torch.tensor(delta_mag, dtype=torch.float32)
+
+        Dt = torch.tensor(
+            t_obs[target_idx] - t_obs[target_idx - 1],
+            dtype=torch.float32,
+        )
+
+        # Keep target image-shaped only if your datastep still expects image targets.
+        # Better is to update datastep to use scalar target directly.
+        data.close()
+        return x, target, Dt
+
+
+class ScalarTemporalConditionedLodeRunner_gri(nn.Module):
+    def __init__(
+        self,
+        backbone: nn.Module,
+        context_len: int = 5,
+        n_input_channels: int = 3,
+        n_output_channels: int = 3,
+        image_size=(1120, 400),
+        backbone_channels: int = 8,
+        hidden: int = 64,
+    ):
+        super().__init__()
+
+        self.backbone = backbone
+        self.context_len = context_len
+        self.n_input_channels = n_input_channels
+        self.n_output_channels = n_output_channels
+        self.image_size = image_size
+        self.backbone_channels = backbone_channels
+
+        # Dataset x layout:
+        #   [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK, t0, t1, ..., tK]
+        #
+        # input_dim = context_len * n_input_channels + context_len
+        input_dim = context_len * n_input_channels + context_len
+
+        # Maps scalar temporal context into the 8 pseudo-channels expected by
+        # the pretrained LodeRunner backbone.
+        self.conditioner = nn.Sequential(
+            nn.Linear(input_dim, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, backbone_channels),
+        )
+
+        # Maps the 8-channel LodeRunner output back to 3 scalar predictions:
+        #   [delta_g, delta_r, delta_i]
+        self.output_head = nn.Sequential(
+            nn.Linear(backbone_channels, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, n_output_channels),
+        )
+
+    def forward(self, x, in_vars, out_vars, Dt):
+        """
+        x: [B, context_len * n_input_channels + context_len]
+
+        For 3-band, context_len=5:
+            x.shape == [B, 20]
+
+        Returns:
+            pred: [B, 3]
+        """
+        B = x.shape[0]
+        H, W = self.image_size
+
+        channel_vals = self.conditioner(x)  # [B, 8]
+
+        pseudo_img = channel_vals.view(
+            B,
+            self.backbone_channels,
+            1,
+            1,
+        ).expand(
+            B,
+            self.backbone_channels,
+            H,
+            W,
+        )
+
+        backbone_in_vars = torch.arange(self.backbone_channels, device=x.device)
+        backbone_out_vars = torch.arange(self.backbone_channels, device=x.device)
+
+        pred_img = self.backbone(
+            pseudo_img,
+            backbone_in_vars,
+            backbone_out_vars,
+            Dt,
+        )  # [B, 8, H, W]
+
+        # Collapse spatial dimensions to 8 backbone-channel summaries.
+        pred_channel_vals = pred_img.mean(dim=(2, 3))  # [B, 8]
+
+        # Convert 8 backbone channels to 3 output bands.
+        pred = self.output_head(pred_channel_vals)  # [B, 3]
+
+        return pred
+
+
+class ScalarTemporalConditionedLodeRunner(nn.Module):
+    def __init__(
+        self,
+        backbone: nn.Module,
+        context_len: int = 5,
+        image_size=(1120, 400),
+        n_channels: int = 8,
+        hidden: int = 64,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.context_len = context_len
+        self.image_size = image_size
+        self.n_channels = n_channels
+
+        # mags + dts
+        self.conditioner = nn.Sequential(
+            nn.Linear(2 * context_len, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, n_channels),
+        )
+
+    def forward(self, x, in_vars, out_vars, Dt):
+        """
+        x: [B, 2 * context_len]
+           first context_len entries are magnitudes
+           second context_len entries are temporal deltas
+        """
+        B = x.shape[0]
+        H, W = self.image_size
+
+        channel_vals = self.conditioner(x)  # [B, 8]
+
+        pseudo_img = channel_vals.view(B, self.n_channels, 1, 1).expand(
+            B,
+            self.n_channels,
+            H,
+            W,
+        )
+
+        pred_img = self.backbone(pseudo_img, in_vars, out_vars, Dt)
+
+        # Convert LodeRunner image output back to scalar delta prediction
+        pred_scalar = pred_img.mean(dim=(1, 2, 3))
+
+        return pred_scalar
+
+
+def load_direct_loderunner_checkpoint(
+    checkpoint_path,
+    model_args,
+    optimizer_kwargs,
+    device,
+):
+    checkpoint_data = torch.load(
+        checkpoint_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    saved_model_args = checkpoint_data.get("model_args", model_args)
+    context_len = checkpoint_data.get("context_len", 5)
+
+    backbone = LodeRunner(**saved_model_args).to(device)
+
+    model = ScalarTemporalConditionedLodeRunner_gri(
+        backbone=backbone,
+        context_len=context_len,
+        n_input_channels=checkpoint_data.get("n_input_channels", 3),
+        n_output_channels=checkpoint_data.get("n_output_channels", 3),
+        image_size=saved_model_args["image_size"],
+        backbone_channels=checkpoint_data.get("backbone_channels", 8),
+        hidden=checkpoint_data.get("hidden", 64),
+    ).to(device)
+
+    state_dict = checkpoint_data["model_state_dict"]
+
+    # Remove DDP prefix if present
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+
+    # Detect checkpoint type
+    is_wrapper_checkpoint = any(
+        k.startswith("backbone.") for k in state_dict.keys()
+    )
+
+    # -------------------------------------------------
+    # OLD plain LodeRunner checkpoint
+    # -------------------------------------------------
+    if not is_wrapper_checkpoint:
+
+        missing_keys, unexpected_keys = model.backbone.load_state_dict(
+            state_dict,
+            strict=False,
+        )
+
+        print("Loaded old LodeRunner checkpoint into model.backbone")
+        print("Missing backbone keys:", missing_keys)
+        print("Unexpected backbone keys:", unexpected_keys)
+
+        # This is NOT a true continuation.
+        # Conditioner is newly initialized.
+        starting_epoch = 0
+
+    # -------------------------------------------------
+    # NEW ScalarTemporalConditionedLodeRunner checkpoint
+    # -------------------------------------------------
+    else:
+
+        model.load_state_dict(state_dict, strict=True)
+
+        print("Loaded ScalarTemporalConditionedLodeRunner checkpoint")
+
+        starting_epoch = checkpoint_data.get("epoch", 0)
+
+    noise_scale = checkpoint_data.get("noise_scale", 0.0)
+    model.backbone.noise_scale = noise_scale
+
+    # Freeze pretrained backbone
+    for p in model.backbone.parameters():
+        p.requires_grad = False
+
+    # Train conditioner
+    for p in model.conditioner.parameters():
+        p.requires_grad = True
+
+    #optimizer = torch.optim.AdamW(
+    #    model.conditioner.parameters(),
+    #    **optimizer_kwargs,
+    #)
+    
+    optimizer = torch.optim.AdamW(
+        list(model.conditioner.parameters()) +
+        list(model.output_head.parameters()),
+        **optimizer_kwargs,
+    )
+
+    # Only restore optimizer for TRUE continuation checkpoints
+    if (
+        is_wrapper_checkpoint
+        and "optimizer_state_dict" in checkpoint_data
+    ):
+
+        optimizer.load_state_dict(
+            checkpoint_data["optimizer_state_dict"]
+        )
+
+        for state in optimizer.state.values():
+            for key, value in state.items():
+                if isinstance(value, torch.Tensor):
+                    state[key] = value.to(device)
+
+    return model, optimizer, starting_epoch
+
+
+def setup_distributed():
+    # ----- 1) Basic setup & environment variables -----
+    # Rely on Slurm variables: SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, etc.
+    rank = int(os.environ["SLURM_PROCID"])  # global rank
+    world_size = int(os.environ["SLURM_NTASKS"])  # total number of processes
+    local_rank = int(os.environ["SLURM_LOCALID"])  # local rank (GPU index on this node)
+
+    master_addr = os.environ["MASTER_ADDR"]
+    master_port = os.environ["MASTER_PORT"]
+
+    # ----- 2) Set the current GPU device for this process -----
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+
+    # ----- 3) Initialize the process group -----
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://{master_addr}:{master_port}",
+        world_size=world_size,
+        rank=rank,
+    )
+
+    return rank, world_size, local_rank, device
+
+
+def cleanup_distributed():
+    # ----- 8) Clean up (optional) -----
+    dist.destroy_process_group()
+
+
+def main(args, rank, world_size, local_rank, device):
+    #############################################
+    # Process Inputs
+    #############################################
+    # Study ID
+    studyIDX = args.studyIDX
+
+    # Resources
+    Ngpus = args.Ngpus
+    Knodes = args.Knodes
+
+    # Data Paths
+    train_filelist = args.FILELIST_DIR + args.train_filelist
+    validation_filelist = args.FILELIST_DIR + args.validation_filelist
+
+    # Model Parameters
+    embed_dim = args.embed_dim
+    block_structure = tuple(args.block_structure)
+
+    # Training Parameters
+    anchor_lr = args.anchor_lr
+    num_cycles = args.num_cycles
+    min_fraction = args.min_fraction
+    terminal_steps = args.terminal_steps
+    warmup_steps = args.warmup_steps
+    noise_scale = args.noise_scale
+
+    # Number of workers controls how batches of data are prefetched and,
+    # possibly, pre-loaded onto GPUs. If the number of workers is large they
+    # will swamp memory and jobs will fail.
+    num_workers = args.num_workers
+
+    # Epoch Parameters
+    batch_size = args.batch_size
+    total_epochs = args.total_epochs
+    cycle_epochs = args.cycle_epochs
+    train_batches = args.train_batches
+    val_batches = args.val_batches
+    train_per_val = args.TRAIN_PER_VAL
+    trn_rcrd_filename = args.trn_rcrd_filename
+    val_rcrd_filename = args.val_rcrd_filename
+    CONTINUATION = args.continuation
+    checkpoint = args.checkpoint
+
+    #############################################
+    # Model Arguments for Dynamic Reconstruction
+    #############################################
+    # Dictionary of available models.
+    available_models = {
+        "LodeRunner": LodeRunner
+    }
+
+    # Model arguments for LodeRunner.
+    model_args = {
+        "default_vars": [
+            "density_case",
+            "density_cushion",
+            "density_maincharge",
+            "density_outside_air",
+            "density_striker",
+            "density_throw",
+            "Uvelocity",
+            "Wvelocity",
+        ],
+        "image_size": (1120, 400),
+        "patch_size": (10, 5),
+        "embed_dim": embed_dim,
+        "emb_factor": 2,
+        "num_heads": 8,
+        "block_structure": block_structure,
+        "window_sizes": [(8, 8), (8, 8), (4, 4), (2, 2)],
+        "patch_merge_scales": [(2, 2), (2, 2), (2, 2)],
+        #"noise_scale": noise_scale,
+    }
+
+
+    CONTEXT_LEN = 5 #3
+    HIDDEN_CHANNELS = 64
+
+    optimizer_kwargs = {
+        "lr": 1e-4,# 1e-4, #1e-5
+        "betas": (0.9, 0.999),
+        "eps": 1e-08,
+        "weight_decay": 0.01,
+    }
+
+
+    if CONTINUATION:
+        model, optimizer, starting_epoch = load_direct_loderunner_checkpoint(
+            checkpoint_path=checkpoint,
+            model_args=model_args,
+            optimizer_kwargs=optimizer_kwargs,
+            device=device,
+        )
+
+        if rank == 0:
+            print(f"Loaded direct checkpoint from {checkpoint}")
+            print(f"Continuing from epoch {starting_epoch}")
+
+        '''  # FIXME block should be unindented if uncommented 
+        if CONTINUATION:
+            model, optimizer, starting_epoch = load_model_and_optimizer(
+                checkpoint,
+                optimizer_class=torch.optim.AdamW,
+                optimizer_kwargs=optimizer_kwargs,
+                available_models=available_models,
+                device=device,
+            )
+
+            if rank == 0:
+                print(f"Loaded temporal checkpoint from {checkpoint}")
+                print(f"Continuing from epoch {starting_epoch}")
+        '''
+
+    else:
+        starting_epoch = 0
+
+        model = LodeRunner(**model_args)
+        model.to(device)
+
+        manual_checkpoint = "/usr/projects/artimis/mpmm/pretrained_models/ddp_ldr_prod_250721/study005_modelState_epoch0100.pth"
+
+        checkpoint_data = torch.load(
+            manual_checkpoint,
+            map_location=device,
+            weights_only=False,
+        )
+
+        state_dict = checkpoint_data["model_state_dict"]
+
+        if all(k.startswith("module.") for k in state_dict.keys()):
+            state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+        missing_keys, unexpected_keys = model.load_state_dict(
+            state_dict,
+            strict=False,
+        )
+
+        if rank == 0:
+            print("Loaded pretrained backbone weights.")
+            print("Missing keys:", missing_keys)
+            print("Unexpected keys:", unexpected_keys)
+
+            model.noise_scale = noise_scale
+
+            backbone = model
+
+            model = ScalarTemporalConditionedLodeRunner_gri(
+                backbone=backbone,
+                context_len=CONTEXT_LEN,
+                n_input_channels=3,
+                n_output_channels=3,
+                image_size=model_args["image_size"],
+                backbone_channels=8,
+                hidden=HIDDEN_CHANNELS,
+            ).to(device)
+
+            # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
+            for p in model.backbone.parameters():
+                p.requires_grad = False
+
+            for p in model.conditioner.parameters():
+                p.requires_grad = True
+
+            for p in model.output_head.parameters():
+                p.requires_grad = True
+
+            optimizer = torch.optim.AdamW(
+                list(model.conditioner.parameters()) +
+                list(model.output_head.parameters()),
+                **optimizer_kwargs,
+            )
+
+    #loss_fn = nn.MSELoss(reduction="none")
+    loss_fn = nn.HuberLoss(delta=0.1, reduction="none")
+    model = DDP(model, device_ids=[local_rank], output_device=local_rank)
+
+    #############################################
+    # Learning Rate Scheduler
+    #############################################
+    print("Starting epoch: ", starting_epoch)
+    if starting_epoch == 0:
+        last_epoch = -1
+    else:
+        last_epoch = train_batches * (starting_epoch - 1)
+
+    # Scale the anchor LR by global batchsize
+    #
+    # # For multi-node
+    lr_scale = np.sqrt(float(Ngpus) * float(Knodes) * float(batch_size))
+    original_batchsize = 40.0  # 1 node, 4 gpus, 10 samples/gpu
+    ddp_anchor_lr = anchor_lr * lr_scale / original_batchsize
+    #
+    # For single node
+    # ddp_anchor_lr = anchor_lr
+
+    LRsched = LambdaLR(
+        optimizer,
+        lr_lambda=lambda step: 1.0,
+        last_epoch=last_epoch,
+    )
+
+
+    LRsched = CosineWithWarmupScheduler(
+        optimizer,
+        anchor_lr=ddp_anchor_lr,
+        terminal_steps=terminal_steps,
+        warmup_steps=warmup_steps,
+        num_cycles=num_cycles,
+        min_fraction=min_fraction,
+        last_epoch=last_epoch,
+    )
+
+    #############################################
+    # Data Initialization (Distributed Dataloader)
+    #############################################
+    #train_dataset = LSC_rho2rho_temporal_DataSet(
+    #    args.LSC_NPZ_DIR,
+    #    file_prefix_list=train_filelist,
+    #    max_timeIDX_offset=2,
+    #    max_file_checks=10,
+    #    half_image=True,
+    #)
+    #val_dataset = LSC_rho2rho_temporal_DataSet(
+    #    args.LSC_NPZ_DIR,
+    #    file_prefix_list=validation_filelist,
+    #    max_timeIDX_offset=2,
+    #    max_file_checks=10,
+    #    half_image=True,
+    #)
+
+    '''
+
+    train_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        context_len=CONTEXT_LEN,
+        #N_imgs=100,
+    )
+
+    val_dataset = Kilonova_lc_img_DataSet_channels_context(
+        half_image=False,
+        context_len=CONTEXT_LEN,
+        #N_imgs=20, #100,
+    )
+    '''
+    '''
+    train_dataset = Kilonova_lc_scalar_context_DataSet(
+        context_len=CONTEXT_LEN,
+    )
+
+    val_dataset = Kilonova_lc_scalar_context_DataSet(
+        context_len=CONTEXT_LEN,
+    )
+    '''
+
+    BAND_KEYS = ("arr_ztfg", "arr_ztfr", "arr_ztfi")
+    VALUE_COL = 1
+    N_BANDS = len(BAND_KEYS)
+
+    norm_stats_path = "kilonova_gri_norm_stats.npz"
+
+    if rank == 0:
+        band_means, band_stds = load_or_compute_band_normalization(
+            stats_path=norm_stats_path,
+            band_keys=BAND_KEYS,
+            value_col=VALUE_COL,
+        )
+
+    dist.barrier()
+
+    if rank != 0:
+        stats = np.load(norm_stats_path, allow_pickle=True)
+        band_means = stats["means"].astype(np.float32)
+        band_stds = stats["stds"].astype(np.float32)
+        stats.close()
+
+    if rank == 0:
+        print("Using band normalization:")
+        print("band_means:", band_means)
+        print("band_stds:", band_stds)
+
+    train_dataset = Kilonova_lc_scalar_context_DataSet_gri(
+        context_len=CONTEXT_LEN,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+        means=band_means,
+        stds=band_stds,
+    )
+
+    val_dataset = Kilonova_lc_scalar_context_DataSet_gri(
+        context_len=CONTEXT_LEN,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+        means=band_means,
+        stds=band_stds,
+    )
+
+
+    # NOTE: For DDP the batch_size is the per-GPU batch_size!!!
+    train_dataloader = make_distributed_dataloader(
+        train_dataset,
+        batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        rank=rank,
+        world_size=world_size,
+    )
+    val_dataloader = make_distributed_dataloader(
+        val_dataset,
+        batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    #############################################
+    # Training Loop (Modified for DDP)
+    #############################################
+    # Train Model
+    print("Training Model . . .")
+    starting_epoch += 1
+    ending_epoch = min(starting_epoch + cycle_epochs, total_epochs + 1)
+
+    TIME_EPOCH = True
+    for epochIDX in range(starting_epoch, ending_epoch):
+        print('%%%%%%%%%%%%%')
+        print(epochIDX)
+        print('%%%%%%%%%%%%%')
+        train_sampler = train_dataloader.sampler
+        train_sampler.set_epoch(epochIDX)
+
+        # For timing epochs
+        if TIME_EPOCH:
+            # Synchronize before starting the timer
+            #dist.barrier()  # Ensure that all nodes sync
+            torch.cuda.synchronize(device)  # Ensure GPUs on each node sync
+            # Time each epoch and print to stdout
+            startTime = time.time()
+
+
+        #train_DDP_loderunner_epoch(
+        train_DDP_scalar_temporal_loderunner_epoch_gri(
+            training_data=train_dataloader,
+            validation_data=val_dataloader,
+            num_train_batches=train_batches,
+            num_val_batches=val_batches,
+            model=model,
+            optimizer=optimizer,
+            loss_fn=loss_fn,
+            LRsched=LRsched,
+            epochIDX=epochIDX,
+            train_per_val=train_per_val,
+            train_rcrd_filename=trn_rcrd_filename,
+            val_rcrd_filename=val_rcrd_filename,
+            device=device,
+            rank=rank,
+            world_size=world_size,
+        )
+
+        print(f"[rank {rank}] finished epoch", flush=True)
+
+
+        if TIME_EPOCH:
+            # Synchronize before stopping the timer
+            torch.cuda.synchronize(device)  # Ensure GPUs on each node sync
+            #dist.barrier()  # Ensure that all nodes sync
+            # Time each epoch and print to stdout
+            endTime = time.time()
+
+        epoch_time = (endTime - startTime) / 60
+
+        # Print Summary Results
+        if rank == 0:
+            print(f"Completed epoch {epochIDX}...", flush=True)
+            print(f"Epoch time (minutes): {epoch_time:.2f}", flush=True)
+
+        # Save model and optimizer
+        #chkpt_name_str = f'study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth'
+        #new_chkpt_path = os.path.join("./", chkpt_name_str)
+
+        if rank == 0:
+            chkpt_name_str = f"study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth"
+            new_chkpt_path = os.path.join("./", chkpt_name_str)
+
+            print(f"Saving checkpoint: {new_chkpt_path}", flush=True)
+
+            torch.save(
+                {
+                    "epoch": epochIDX,
+                    "model_class": "ScalarTemporalConditionedLodeRunner_gri",
+                    "backbone_class": "LodeRunner",
+                    "model_args": model_args,
+                    "model_state_dict": model.module.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "noise_scale": noise_scale,
+                    "predicts_delta": True,
+                    "target_type": "delta_gri",
+                    "context_len": CONTEXT_LEN,
+                    "n_input_channels": 3,
+                    "n_output_channels": 3,
+                    "backbone_channels": 8,
+                    "hidden": 64,
+                },
+                new_chkpt_path,
+            )
+
+            print(f"Saved checkpoint: {new_chkpt_path}", flush=True)
+
+    if rank == 0:
+        #############################################
+        # Continue if Necessary
+        #############################################
+        FINISHED_TRAINING = epochIDX + 1 > total_epochs
+        if not FINISHED_TRAINING:
+            new_slurm_file = continuation_setup(
+                new_chkpt_path, studyIDX, last_epoch=epochIDX
+            )
+            os.system(f"sbatch {new_slurm_file}")
+
+if __name__ == "__main__":
+    print('running main')
+    args = parser.parse_args()
+
+    rank, world_size, local_rank, device = setup_distributed()
+
+    main(args, rank, world_size, local_rank, device)
+
+    cleanup_distributed()
diff --git a/applications/harnesses/KN_loderunner/training_START.input b/applications/harnesses/KN_loderunner/training_START.input
new file mode 100644
index 00000000..81e87bb2
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/training_START.input
@@ -0,0 +1,53 @@
+--studyIDX
+<studyIDX>
+--FILELIST_DIR
+/users/atoivonen/forks/Yoke/applications/filelists/
+--LSC_NPZ_DIR
+/lustre/scratch5/exempt/artimis/data/lsc240420/
+--train_filelist
+lsc240420_prefixes_train_80pct.txt
+--validation_filelist
+lsc240420_prefixes_validation_10pct.txt
+--block_structure
+<B0>
+<B1>
+<B2>
+<B3>
+--embed_dim
+<EMBED_DIM>
+--anchor_lr
+<ANCHOR_LR>
+--num_cycles
+<NUM_CYCLES>
+--min_fraction
+<MIN_FRACTION>
+--terminal_steps
+<TERMINAL_STEPS>
+--warmup_steps
+<WARMUP_STEPS>
+--noise_scale
+<NOISE_SCALE>
+--trn_rcrd_filename
+./training_study<studyIDX>_epoch<epochIDX>.csv
+--val_rcrd_filename
+./validation_study<studyIDX>_epoch<epochIDX>.csv
+--batch_size
+<BATCH_SIZE>
+--num_workers
+<NUM_WORKERS>
+--Ngpus
+<NGPUS>
+--Knodes
+<KNODES>
+--total_epochs
+500
+--cycle_epochs
+1
+--train_batches
+<NTRN_BATCH>
+--val_batches
+<NVAL_BATCH>
+--TRAIN_PER_VAL
+5
+--pretrained_model
+/usr/projects/artimis/mpmm/pretrained_models/ddp_ldr_prod_250721/study005_modelState_epoch0100.pth
diff --git a/applications/harnesses/KN_loderunner/training_START.slurm b/applications/harnesses/KN_loderunner/training_START.slurm
new file mode 100644
index 00000000..524e4f3f
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/training_START.slurm
@@ -0,0 +1,73 @@
+#!/bin/bash
+
+# This is a setup for GPU training on Venado. Find out how much
+# memory per node, number of CPUs/node.
+
+# NOTE: Number of CPUs per GPU must be an even number since there are
+# 2 threads per core. If an odd number is requested the next higher
+# even number gets used. 
+
+# The following are one set of SBATCH options for the Venado GPU
+# partition. There are optional other constraints.
+
+#SBATCH --job-name=ddp_s<studyIDX>_e0001
+#SBATCH --account=y26_artimis_ddc_g
+#SBATCH --time=1:00:00
+#SBATCH --partition=standard
+#SBATCH --nodes=<KNODES>
+#SBATCH --ntasks-per-node=<NGPUS>
+#SBATCH --gpus-per-node=<NGPUS>
+#SBATCH --mem-per-gpu=50G
+#SBATCH --output=study<studyIDX>_epoch0001.out
+#SBATCH --error=study<studyIDX>_epoch0001.err
+
+# Set the master node's address and port
+MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
+MASTER_PORT=$(shuf -i 1024-65535 -n 1)  # Choose a random port
+export MASTER_ADDR
+export MASTER_PORT
+
+# Enable shell debugging
+set -xv
+
+# Check available GPUs
+sinfo  -o "%P %.24G %N"
+srun -vv --cpu-bind=verbose /usr/bin/echo $CUDA_AVAILABLE_DEVICES
+nvidia-smi
+
+# Specify NCCL communication
+export NCCL_SOCKET_IFNAME=ib0  # Check possible interfaces with `ip link show`
+
+# for multi-node training
+export NCCL_IB_HCA=mlx5_0
+export NCCL_IB_GID_INDEX=3  # sometimes 0 or 3 depending on subnet manager config
+
+# Debugging distributed data parallel
+# export NCCL_DEBUG=INFO
+# export NCCL_DEBUG_SUBSYS=INIT
+
+# Debugging selene slurm
+export SLURM_CPU_BIND=verbose
+
+# Load correct conda environment
+module load anaconda/3.12
+source activate
+conda activate <YOKE_TORCH_ENV>
+
+# Set number of threads per GPU
+export OMP_NUM_THREADS=10
+
+# Get start time
+export date00=`date`
+
+# Start the Code
+# Explicitly set TCP environment for the following...
+srun -vv --cpu-bind=verbose python -u <train_script> @study<studyIDX>_START.input
+
+# Get end time and print to stdout
+export date01=`date`
+
+echo "===================TIME STARTED==================="
+echo $date00
+echo "===================TIME FINISHED==================="
+echo $date01
diff --git a/applications/harnesses/KN_loderunner/training_input.tmpl b/applications/harnesses/KN_loderunner/training_input.tmpl
new file mode 100644
index 00000000..03793f7f
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/training_input.tmpl
@@ -0,0 +1,56 @@
+--studyIDX
+<studyIDX>
+--FILELIST_DIR
+/users/atoivonen/forks/Yoke/applications/filelists/
+--LSC_NPZ_DIR
+/lustre/scratch5/exempt/artimis/data/lsc240420/
+--train_filelist
+lsc240420_prefixes_train_80pct.txt
+--validation_filelist
+lsc240420_prefixes_validation_10pct.txt
+--block_structure
+<B0>
+<B1>
+<B2>
+<B3>
+--embed_dim
+<EMBED_DIM>
+--anchor_lr
+<ANCHOR_LR>
+--num_cycles
+<NUM_CYCLES>
+--min_fraction
+<MIN_FRACTION>
+--terminal_steps
+<TERMINAL_STEPS>
+--warmup_steps
+<WARMUP_STEPS>
+--noise_scale
+<NOISE_SCALE>
+--trn_rcrd_filename
+./training_study<studyIDX>_epoch<epochIDX>.csv
+--val_rcrd_filename
+./validation_study<studyIDX>_epoch<epochIDX>.csv
+--batch_size
+<BATCH_SIZE>
+--num_workers
+<NUM_WORKERS>
+--Ngpus
+<NGPUS>
+--Knodes
+<KNODES>
+--total_epochs
+500
+--cycle_epochs
+1
+--train_batches
+<NTRN_BATCH>
+--val_batches
+<NVAL_BATCH>
+--TRAIN_PER_VAL
+5
+--continuation
+--checkpoint
+<CHECKPOINT>
+--pretrained_model
+/usr/projects/artimis/mpmm/pretrained_models/ddp_ldr_prod_250721/study005_modelState_epoch0100.pth
diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
new file mode 100644
index 00000000..3ba64beb
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -0,0 +1,73 @@
+#!/bin/bash
+
+# This is a setup for GPU training on Venado. Find out how much
+# memory per node, number of CPUs/node.
+
+# NOTE: Number of CPUs per GPU must be an even number since there are
+# 2 threads per core. If an odd number is requested the next higher
+# even number gets used. 
+
+# The following are one set of SBATCH options for the Venado GPU
+# partition. There are optional other constraints.
+
+#SBATCH --job-name=ddp_s<studyIDX>_e<epochIDX>
+#SBATCH --account=y26_artimis_ddc_g
+#SBATCH --time=1:00:00
+#SBATCH --partition=standard
+#SBATCH --nodes=<KNODES>
+#SBATCH --ntasks-per-node=<NGPUS>
+#SBATCH --gpus-per-node=<NGPUS>
+#SBATCH --mem-per-gpu=50G
+#SBATCH --output=study<studyIDX>_epoch<epochIDX>.out
+#SBATCH --error=study<studyIDX>_epoch<epochIDX>.err
+
+# Set the master node's address and port
+MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
+MASTER_PORT=$(shuf -i 1024-65535 -n 1)  # Choose a random port
+export MASTER_ADDR
+export MASTER_PORT
+
+# Enable shell debugging
+set -xv
+
+# Check available GPUs
+sinfo  -o "%P %.24G %N"
+srun -vv --cpu-bind=verbose /usr/bin/echo $CUDA_AVAILABLE_DEVICES
+nvidia-smi
+
+# Specify NCCL communication
+export NCCL_SOCKET_IFNAME=ib0  # Check possible interfaces with `ip link show`
+
+# for multi-node training
+export NCCL_IB_HCA=mlx5_0
+export NCCL_IB_GID_INDEX=3  # sometimes 0 or 3 depending on subnet manager config
+
+# Debugging distributed data parallel
+# export NCCL_DEBUG=INFO
+# export NCCL_DEBUG_SUBSYS=INIT
+
+# Debugging selene slurm
+export SLURM_CPU_BIND=verbose
+
+# Load correct conda environment
+module load anaconda/3.12
+source activate
+conda activate <YOKE_TORCH_ENV>
+
+# Set number of threads per GPU
+export OMP_NUM_THREADS=10
+
+# Get start time
+export date00=`date`
+
+# Start the Code
+# Explicitly set TCP environment for the following...
+srun -vv --cpu-bind=verbose python -u <train_script> @<INPUTFILE>
+
+# Get end time and print to stdout
+export date01=`date`
+
+echo "===================TIME STARTED==================="
+echo $date00
+echo "===================TIME FINISHED==================="
+echo $date01

From 72953b4aea572feb4172e0244a4f3e441871acf1 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 30 Jul 2026 15:50:22 -0600
Subject: [PATCH 10/66] Remove old and unused functions

---
 .../utils/KN_tmp/plot_loss_curves_channel.py  | 123 ---
 src/yoke/utils/KN_tmp/plot_pred_channel.py    | 214 -----
 .../utils/KN_tmp/plot_pred_channel_delta.py   | 254 ------
 src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py | 861 ------------------
 .../utils/training/datastep/loderunner.py     | 461 ----------
 src/yoke/utils/training/epoch/loderunner.py   |   6 -
 6 files changed, 1919 deletions(-)
 delete mode 100644 src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
 delete mode 100644 src/yoke/utils/KN_tmp/plot_pred_channel.py
 delete mode 100644 src/yoke/utils/KN_tmp/plot_pred_channel_delta.py
 delete mode 100644 src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py

diff --git a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py b/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
deleted file mode 100644
index 008d95b1..00000000
--- a/src/yoke/utils/KN_tmp/plot_loss_curves_channel.py
+++ /dev/null
@@ -1,123 +0,0 @@
-import argparse
-import glob
-import os
-import numpy as np
-import matplotlib.pyplot as plt
-
-
-def load_records(pattern):
-    files = sorted(glob.glob(pattern))
-
-    if len(files) == 0:
-        raise FileNotFoundError(f"No files matched pattern: {pattern}")
-
-    arrays = []
-
-    for fn in files:
-        try:
-            arr = np.loadtxt(fn, delimiter=",")
-        except Exception as e:
-            print(f"Skipping {fn}: {e}")
-            continue
-
-        if arr.size == 0:
-            continue
-
-        if arr.ndim == 1:
-            arr = arr[None, :]
-
-        arrays.append(arr)
-
-    if len(arrays) == 0:
-        raise RuntimeError(f"No valid data found for pattern: {pattern}")
-
-    data = np.vstack(arrays)
-
-    # columns: epoch, batch, loss
-    epochs = data[:, 0].astype(int)
-    batches = data[:, 1].astype(int)
-    losses = data[:, 2]
-
-    return epochs, batches, losses, files
-
-
-def epoch_means(epochs, losses):
-    unique_epochs = np.array(sorted(set(epochs)))
-    mean_losses = np.array([losses[epochs == e].mean() for e in unique_epochs])
-    std_losses = np.array([losses[epochs == e].std() for e in unique_epochs])
-    return unique_epochs, mean_losses, std_losses
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-    "--train_pattern",
-    type=str,
-    #default="runs/study_010/training_study010_epoch*.csv",
-    default="runs/study_013/training_study013_epoch*.csv",
-    )
-
-    parser.add_argument(
-        "--val_pattern",
-        type=str,
-        #default="runs/study_010/validation_study010_epoch*.csv",
-        default="runs/study_013/validation_study013_epoch*.csv",
-    )
-
-    parser.add_argument(
-        "--out",
-        type=str,
-        default="loss_curves_study013.png",
-    )
-
-    parser.add_argument(
-        "--logy",
-        action="store_true",
-        default=True,
-        help="Use log scale on y-axis.",
-    )
-
-    args = parser.parse_args()
-
-    train_epochs, train_batches, train_losses, train_files = load_records(args.train_pattern)
-
-    print("Loaded training files:")
-    for f in train_files:
-        print("  ", f)
-
-    train_ep, train_mean, train_std = epoch_means(train_epochs, train_losses)
-
-    plt.figure(figsize=(8, 5))
-    plt.plot(train_ep, train_mean, marker="o", label="Train")
-
-    # Try validation, but do not fail if absent
-    try:
-        val_epochs, val_batches, val_losses, val_files = load_records(args.val_pattern)
-
-        print("Loaded validation files:")
-        for f in val_files:
-            print("  ", f)
-
-        val_ep, val_mean, val_std = epoch_means(val_epochs, val_losses)
-        plt.plot(val_ep, val_mean, marker="s", label="Validation")
-
-    except Exception as e:
-        print(f"No validation curve plotted: {e}")
-
-    plt.xlabel("Epoch")
-    plt.ylabel("Mean loss")
-    plt.title("Loss curves")
-    plt.grid(True, alpha=0.3)
-    plt.legend()
-
-    if args.logy:
-        plt.yscale("log")
-
-    plt.tight_layout()
-    plt.savefig(args.out, dpi=200)
-    print(f"Saved {args.out}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/yoke/utils/KN_tmp/plot_pred_channel.py b/src/yoke/utils/KN_tmp/plot_pred_channel.py
deleted file mode 100644
index d126bbcb..00000000
--- a/src/yoke/utils/KN_tmp/plot_pred_channel.py
+++ /dev/null
@@ -1,214 +0,0 @@
-import argparse
-import numpy as np
-import torch
-import matplotlib
-import matplotlib.pyplot as plt
-
-from yoke.models.vit.swin.bomberman import LodeRunner
-from torch.utils.data import DataLoader
-
-from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
-
-matplotlib.rcParams["pdf.fonttype"] = 42
-matplotlib.rcParams["ps.fonttype"] = 42
-plt.rc("font", family="serif")
-plt.rcParams["figure.figsize"] = (6, 6)
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--ckpt",
-        type=str,
-        #default="runs/study_007/study007_modelState_epoch0100.pth",
-        #default="runs/study_010/study010_modelState_epoch0100.pth",
-        default="runs/study_012/study012_modelState_epoch0100.pth",
-
-    )
-    parser.add_argument("--N_imgs", type=int, default=1)
-    parser.add_argument("--batch_size", type=int, default=1)
-    parser.add_argument("--n_future_steps", type=int, default=10)
-
-    return parser.parse_args()
-
-
-def load_channel_model(ckpt_path, device):
-    ckpt = torch.load(
-        ckpt_path,
-        map_location=device,
-        weights_only=False,
-    )
-
-    model_args = ckpt["model_args"]
-    noise_scale = ckpt.get("noise_scale", 0.0)
-
-    model = LodeRunner(**model_args)
-    model.to(device)
-
-    state_dict = ckpt["model_state_dict"]
-
-    if all(k.startswith("module.") for k in state_dict.keys()):
-        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
-
-    missing, unexpected = model.load_state_dict(state_dict, strict=True)
-
-    print("Loaded checkpoint:", ckpt_path)
-    print("Missing keys:", missing)
-    print("Unexpected keys:", unexpected)
-    print("Loaded model_args:", model_args)
-
-    model.noise_scale = noise_scale
-    model.eval()
-
-    return model
-
-
-def main():
-    args = get_args()
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print("Using device:", device)
-
-    context_len = 5
-    model = load_channel_model(args.ckpt, device)
-
-    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
-        half_image=False,
-        N_imgs=args.N_imgs,
-        context_len=context_len,
-    )
-
-    loader = DataLoader(
-        eval_dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-    )
-
-    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
-    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
-
-    # ------------------------------------------------------------
-    # One-step predictions using true context windows
-    # ------------------------------------------------------------
-    preds = []
-    targets = []
-    idxs = []
-    prefix = []
-
-    for idx, (context_img, target, Dt) in enumerate(loader):
-        context_img = context_img.to(device)
-        if idx == 0:
-            context_means = context_img.mean(dim=(2, 3))[0].detach().cpu().numpy()
-            for context in context_means:
-                prefix.append(context.mean().item())
-        target = target.to(device)
-        Dt = Dt.to(torch.float32).to(device)
-
-        with torch.no_grad():
-            pred_image = model(context_img, in_vars, out_vars, Dt)
-
-        preds.append(pred_image.mean().item())
-        targets.append(target.mean().item())
-        idxs.append(idx)
-
-    plt.figure()
-    plt.scatter(idxs, preds, label="Predictions")
-    plt.scatter(idxs, targets, label="Truth")
-    plt.scatter(np.arange(len(prefix))-(len(prefix)), prefix, label='Initial Context Window')
-    plt.legend()
-    plt.gca().invert_yaxis()
-    plt.xlabel("Sample index")
-    plt.ylabel("Mean magnitude/image value")
-    plt.tight_layout()
-    plt.savefig("pred_vs_truth_channel_norm.png", dpi=200)
-
-
-
-    context_seq, target, Dt = next(iter(loader))
-
-    context_seq = context_seq.to(device)
-    Dt = Dt.to(torch.float32).to(device)
-
-    x = context_seq
-
-    preds_seq = []
-    truth_seq = []
-    idxs_seq = []
-
-    future_iter = iter(loader)
-
-    for step in range(args.n_future_steps):
-        try:
-            _, future_target, future_Dt = next(future_iter)
-        except StopIteration:
-            break
-
-        future_target = future_target.to(device)
-        future_Dt = future_Dt.to(torch.float32).to(device)
-
-        with torch.no_grad():
-            pred_image = model(x, in_vars, out_vars, future_Dt)
-
-        preds_seq.append(pred_image.mean().item())
-        truth_seq.append(future_target.mean().item())
-        idxs_seq.append(step)
-
-        # autoregressive update: append prediction
-        x = torch.cat([x[:, 1:], pred_image[:, -1:].detach()], dim=1)
-
-    plt.figure()
-
-    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
-    plt.scatter(idxs_seq, truth_seq, label="Truth")
-    plt.scatter(
-        np.arange(len(prefix)) - len(prefix),
-        prefix,
-        label="Initial context window",
-    )
-
-    plt.legend()
-    plt.gca().invert_yaxis()
-    plt.xlabel("Autoregressive step")
-    plt.ylabel("Mean magnitude/image value")
-    plt.tight_layout()
-    plt.savefig("pred_vs_truth_channel_norm_autoreg.png", dpi=200)
-
-
-    # ------------------------------------------------------------
-    # Image comparison for the final one-step batch above
-    # ------------------------------------------------------------
-    pred_plot = pred_image.squeeze().mean(dim=0).detach().cpu().numpy()
-    true_plot = target.squeeze().mean(dim=0).detach().cpu().numpy()
-    error_plot = pred_plot - true_plot
-
-    vmin = min(pred_plot.min(), true_plot.min())
-    vmax = max(pred_plot.max(), true_plot.max())
-    err_max = np.max(np.abs(error_plot))
-
-    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
-
-    im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
-    ax1.set_title("Prediction")
-
-    ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
-    ax2.set_title("Truth")
-
-    im3 = ax3.imshow(error_plot, origin="lower", vmin=-err_max, vmax=err_max)
-    ax3.set_title("Error (Pred - Truth)")
-
-    cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
-    cbar.set_label("Field value")
-
-    cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
-    cbar_err.set_label("Error")
-
-    for ax in (ax1, ax2, ax3):
-        ax.axis("off")
-
-    plt.tight_layout()
-    plt.savefig("img_comp_channel_norm.png", bbox_inches="tight", dpi=200)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/yoke/utils/KN_tmp/plot_pred_channel_delta.py b/src/yoke/utils/KN_tmp/plot_pred_channel_delta.py
deleted file mode 100644
index bf1874a0..00000000
--- a/src/yoke/utils/KN_tmp/plot_pred_channel_delta.py
+++ /dev/null
@@ -1,254 +0,0 @@
-import argparse
-import numpy as np
-import torch
-import matplotlib
-import matplotlib.pyplot as plt
-
-from yoke.models.vit.swin.bomberman import LodeRunner
-from torch.utils.data import DataLoader
-
-from train_LodeRunner_ddp import Kilonova_lc_img_DataSet_channels_context
-
-matplotlib.rcParams["pdf.fonttype"] = 42
-matplotlib.rcParams["ps.fonttype"] = 42
-plt.rc("font", family="serif")
-plt.rcParams["figure.figsize"] = (6, 6)
-
-# ============================================================
-# RUN IDENTIFIER
-# ============================================================
-RUN_ID = "013"
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--ckpt",
-        type=str,
-        default=f"runs/study_{RUN_ID}/study{RUN_ID}_modelState_epoch0100.pth",
-    )
-    parser.add_argument("--N_imgs", type=int, default=1)
-    parser.add_argument("--batch_size", type=int, default=1)
-    parser.add_argument("--n_future_steps", type=int, default=15)
-
-    return parser.parse_args()
-
-
-def load_channel_model(ckpt_path, device):
-    ckpt = torch.load(
-        ckpt_path,
-        map_location=device,
-        weights_only=False,
-    )
-
-    model_args = ckpt["model_args"]
-    noise_scale = ckpt.get("noise_scale", 0.0)
-    context_len = ckpt.get("context_len", 5)
-
-    print("Loaded checkpoint:", ckpt_path)
-    print("predicts_delta:", ckpt.get("predicts_delta", False))
-    print("target_type:", ckpt.get("target_type", "absolute"))
-    print("context_len:", context_len)
-
-    model = LodeRunner(**model_args)
-    model.to(device)
-
-    state_dict = ckpt["model_state_dict"]
-
-    if all(k.startswith("module.") for k in state_dict.keys()):
-        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
-
-    missing, unexpected = model.load_state_dict(state_dict, strict=True)
-
-    print("Missing keys:", missing)
-    print("Unexpected keys:", unexpected)
-
-    model.noise_scale = noise_scale
-    model.eval()
-
-    return model, context_len
-
-
-def main():
-    args = get_args()
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print("Using device:", device)
-
-    model, context_len = load_channel_model(args.ckpt, device)
-
-    eval_dataset = Kilonova_lc_img_DataSet_channels_context(
-        half_image=False,
-        N_imgs=args.N_imgs,
-        context_len=context_len,
-    )
-
-    loader = DataLoader(
-        eval_dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-    )
-
-    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
-    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
-
-    # ------------------------------------------------------------
-    # One-step predictions
-    # ------------------------------------------------------------
-    preds = []
-    targets = []
-    idxs = []
-    prefix = []
-
-    for idx, (context_img, target_delta, Dt) in enumerate(loader):
-        context_img = context_img.to(device)
-        target_delta = target_delta.to(device)
-        Dt = Dt.to(torch.float32).to(device)
-
-        if idx == 0:
-            context_means = context_img.mean(dim=(2, 3))[0].detach().cpu().numpy()
-            for context in context_means:
-                prefix.append(context.mean().item())
-
-        with torch.no_grad():
-            pred_delta_img = model(context_img, in_vars, out_vars, Dt)
-
-        last_mag_img = context_img[:, -1:]
-
-        pred_next_img = last_mag_img + pred_delta_img[:, -1:]
-        true_next_img = last_mag_img + target_delta[:, -1:]
-
-        preds.append(pred_next_img.mean().item())
-        targets.append(true_next_img.mean().item())
-        idxs.append(idx)
-
-    plt.figure()
-    plt.scatter(idxs, preds, label="Predicted next magnitude")
-    plt.scatter(idxs, targets, label="True next magnitude")
-    plt.scatter(
-        np.arange(len(prefix)) - len(prefix),
-        prefix,
-        label="Initial context window",
-    )
-
-    plt.legend()
-    plt.gca().invert_yaxis()
-    plt.xlabel("Sample index")
-    plt.ylabel("Normalized magnitude")
-    plt.tight_layout()
-
-    plt.savefig(
-        f"study{RUN_ID}_pred_vs_truth_channel_delta_onestep.png",
-        dpi=200,
-    )
-
-    # ------------------------------------------------------------
-    # Autoregressive rollout
-    # ------------------------------------------------------------
-    context_seq, target_delta, Dt = next(iter(loader))
-
-    x = context_seq.to(device)
-
-    preds_seq = []
-    truth_seq = []
-    idxs_seq = []
-
-    future_iter = iter(loader)
-
-    pred_next_img = None
-    true_next_img = None
-
-    for step in range(args.n_future_steps):
-        try:
-            _, future_target_delta, future_Dt = next(future_iter)
-        except StopIteration:
-            break
-
-        future_target_delta = future_target_delta.to(device)
-        future_Dt = future_Dt.to(torch.float32).to(device)
-
-        with torch.no_grad():
-            pred_delta_img = model(x, in_vars, out_vars, future_Dt)
-
-        last_mag_img = x[:, -1:]
-
-        pred_next_img = last_mag_img + pred_delta_img[:, -1:]
-        true_next_img = last_mag_img + future_target_delta[:, -1:]
-
-        preds_seq.append(pred_next_img.mean().item())
-        truth_seq.append(true_next_img.mean().item())
-        idxs_seq.append(step)
-
-        # Append predicted absolute next magnitude
-        x = torch.cat([x[:, 1:], pred_next_img.detach()], dim=1)
-
-    plt.figure()
-
-    plt.scatter(idxs_seq, preds_seq, label="Autoregressive predictions")
-    plt.scatter(idxs_seq, truth_seq, label="Truth")
-    plt.scatter(
-        np.arange(len(prefix)) - len(prefix),
-        prefix,
-        label="Initial context window",
-    )
-
-    plt.legend()
-    plt.gca().invert_yaxis()
-    plt.xlabel("Autoregressive step")
-    plt.ylabel("Normalized magnitude")
-    plt.tight_layout()
-
-    plt.savefig(
-        f"study{RUN_ID}_pred_vs_truth_channel_delta_autoreg.png",
-        dpi=200,
-    )
-
-    # ------------------------------------------------------------
-    # Image comparison
-    # ------------------------------------------------------------
-    if pred_next_img is not None and true_next_img is not None:
-        pred_plot = pred_next_img[0, 0].detach().cpu().numpy()
-        true_plot = true_next_img[0, 0].detach().cpu().numpy()
-        error_plot = pred_plot - true_plot
-
-        vmin = min(pred_plot.min(), true_plot.min())
-        vmax = max(pred_plot.max(), true_plot.max())
-        err_max = np.max(np.abs(error_plot))
-
-        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(22, 6))
-
-        im1 = ax1.imshow(pred_plot, origin="lower", vmin=vmin, vmax=vmax)
-        ax1.set_title("Prediction")
-
-        ax2.imshow(true_plot, origin="lower", vmin=vmin, vmax=vmax)
-        ax2.set_title("Truth")
-
-        im3 = ax3.imshow(
-            error_plot,
-            origin="lower",
-            vmin=-err_max,
-            vmax=err_max,
-        )
-        ax3.set_title("Error (Pred - Truth)")
-
-        cbar = fig.colorbar(im1, ax=[ax1, ax2], shrink=0.8)
-        cbar.set_label("Normalized magnitude")
-
-        cbar_err = fig.colorbar(im3, ax=ax3, shrink=0.8)
-        cbar_err.set_label("Error")
-
-        for ax in (ax1, ax2, ax3):
-            ax.axis("off")
-
-        plt.tight_layout()
-
-        plt.savefig(
-            f"study{RUN_ID}_img_comp_channel_delta.png",
-            bbox_inches="tight",
-            dpi=200,
-        )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py b/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
deleted file mode 100644
index ffb79712..00000000
--- a/src/yoke/utils/KN_tmp/train_LodeRunner_ddp.py
+++ /dev/null
@@ -1,861 +0,0 @@
-import os
-import time
-import argparse
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-
-from yoke.models.vit.swin.bomberman import LodeRunner
-from yoke.datasets.lsc_dataset import LSC_rho2rho_temporal_DataSet
-from yoke.utils.training.epoch.loderunner import train_DDP_loderunner_epoch
-from yoke.utils.training.epoch.loderunner import train_DDP_loderunner_epoch_seq_context
-from yoke.utils.restart import continuation_setup
-from yoke.utils.dataload import make_distributed_dataloader
-from yoke.utils.checkpointing import load_model_and_optimizer
-from yoke.utils.checkpointing import save_model_and_optimizer
-from yoke.lr_schedulers import CosineWithWarmupScheduler
-from yoke.helpers import cli
-
-# FIXME remove if restructure
-from torch.utils.data import Dataset, DataLoader, random_split
-import glob
-import random
-
-#MEAN = 24.694652705328807
-#STD = 4.67030961432848
-
-GLOBAL_GMAG_MEAN = 24.694652705328807
-GLOBAL_GMAG_STD  = 4.67030961432848
-EPS = 1e-6
-
-#############################################
-# Inputs
-#############################################
-descr_str = (
-    "Uses DDP to train LodeRunner architecture on single-timstep input and output "
-    "of the lsc240420 per-material density fields."
-)
-parser = argparse.ArgumentParser(
-    prog="DDP LodeRunner Training", description=descr_str, fromfile_prefix_chars="@"
-)
-parser = cli.add_default_args(parser=parser)
-parser = cli.add_filepath_args(parser=parser)
-parser = cli.add_computing_args(parser=parser)
-parser = cli.add_model_args(parser=parser)
-parser = cli.add_training_args(parser=parser)
-parser = cli.add_cosine_lr_scheduler_args(parser=parser)
-
-# DPOT‐style noise parameter
-parser.add_argument(
-    "--noise_scale",
-    type=float,
-    default=0.0,
-    help="Relative magnitude ε for Gaussian noise injection (e.g. 5e-5).",
-)
-
-# Change some default filepaths.
-parser.set_defaults(
-    train_filelist="lsc240420_prefixes_train_80pct.txt",
-    validation_filelist="lsc240420_prefixes_validation_10pct.txt",
-    test_filelist="lsc240420_prefixes_test_10pct.txt",
-)
-
-
-class Kilonova_lc_img_DataSet(Dataset):
-    def __init__(self, half_image=False, N_imgs=0):
-        file_prefix_list = sorted(
-            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
-        )
-
-        if N_imgs == 0:
-            self.file_prefix_list = file_prefix_list
-        else:
-            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
-
-        random.shuffle(self.file_prefix_list)
-
-        #self.max_timeIDX_offset = max_timeIDX_offset
-        self.half_image = half_image
-
-        # Build a global index: one entry per usable (file, startIDX)
-        self.samples = []
-        seqLen = 1
-
-        for file_idx, fn in enumerate(self.file_prefix_list):
-            data = np.load(fn, allow_pickle=True)
-            mjd = data["arr_ztfg"][:, 0]
-            n_times = len(mjd)
-            data.close()
-
-            max_start = n_times - seqLen - 1
-            for startIDX in range(max_start + 1):
-                self.samples.append((file_idx, startIDX))
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, index):
-        file_idx, startIDX = self.samples[index]
-        fn = self.file_prefix_list[file_idx]
-
-        data = np.load(fn, allow_pickle=True)
-
-        mjd = data["arr_ztfg"][:, 0]
-        t0 = mjd.min()
-        t_obs = mjd - t0
-        g_mag = data["arr_ztfg"][:, 1]
-
-        seqLen = 1
-        endIDX = startIDX + seqLen
-
-        start_mag = g_mag[startIDX]
-        end_mag = g_mag[endIDX]
-        start_t = t_obs[startIDX]
-        end_t = t_obs[endIDX]
-
-        Dt = torch.tensor(end_t - start_t, dtype=torch.float32)
-
-        H, W = 1120, 400
-
-        s = torch.tensor(start_mag, dtype=torch.float32)
-        start_img = s.view(1, 1, 1).expand(8, H, W)
-
-        e = torch.tensor(end_mag, dtype=torch.float32)
-        end_img = e.view(1, 1, 1).expand(8, H, W)
-
-        data.close()
-        return start_img, end_img, Dt
-
-
-class Kilonova_lc_img_DataSet_seq(Dataset):
-    def __init__(self, half_image=False, N_imgs=0):
-        file_prefix_list = sorted(
-            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
-        )
-
-        if N_imgs == 0:
-            self.file_prefix_list = file_prefix_list
-        else:
-            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
-
-        random.shuffle(self.file_prefix_list)
-
-        #self.max_timeIDX_offset = max_timeIDX_offset
-        self.half_image = half_image
-
-        # Build a global index: one entry per usable (file, startIDX)
-        self.samples = []
-        seqLen = 3
-
-        for file_idx, fn in enumerate(self.file_prefix_list):
-            data = np.load(fn, allow_pickle=True)
-            mjd = data["arr_ztfg"][:, 0]
-            n_times = len(mjd)
-            data.close()
-
-            max_start = n_times - seqLen - 1
-            for startIDX in range(max_start + 1):
-                self.samples.append((file_idx, startIDX))
-
-    def __len__(self):
-        return len(self.samples)
-
-    
-    def __getitem__(self, index):
-        file_idx, startIDX = self.samples[index]
-        fn = self.file_prefix_list[file_idx]
-
-        frames = []
-        seqLen = 3
-        H, W = 1120, 400
-
-        data = np.load(fn, allow_pickle=True)
-
-        mjd = data["arr_ztfg"][:, 0]
-        t0 = mjd.min()
-        t_obs = mjd - t0
-        g_mag = data["arr_ztfg"][:, 1]
-
-        endIDX = startIDX + seqLen
-
-        for i in range(seqLen):
-            seq_mag = g_mag[startIDX + i]
-            s = torch.tensor(seq_mag, dtype=torch.float32)
-            seq_img = s.view(1, 1, 1).expand(8, H, W)
-            frames.append(seq_img)
-
-        end_mag = g_mag[endIDX]
-        e = torch.tensor(end_mag, dtype=torch.float32)
-        end_img = e.view(1, 1, 1).expand(8, H, W)
-        frames.append(end_img)
-
-        start_t = t_obs[startIDX]
-        end_t = t_obs[endIDX]
-        Dt = torch.tensor(end_t - start_t, dtype=torch.float32)
-
-        data.close()
-
-        img_seq = torch.stack(frames, dim=0)
-        return img_seq, Dt
-
-
-class Kilonova_lc_img_DataSet_channels_context(Dataset):
-    def __init__(
-        self,
-        half_image=False,
-        N_imgs=0,
-        context_len=3,
-        H=1120,
-        W=400,
-        n_channels=8,
-    ):
-        assert context_len <= n_channels
-
-        file_prefix_list = sorted(
-            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
-        )
-
-        if N_imgs == 0:
-            self.file_prefix_list = file_prefix_list
-        else:
-            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
-
-        random.shuffle(self.file_prefix_list)
-
-        self.context_len = context_len
-        self.H = H
-        self.W = W
-        self.n_channels = n_channels
-        self.samples = []
-
-        for file_idx, fn in enumerate(self.file_prefix_list):
-            data = np.load(fn, allow_pickle=True)
-            mjd = data["arr_ztfg"][:, 0]
-            n_times = len(mjd)
-            data.close()
-
-            max_start = n_times - context_len - 1
-            for startIDX in range(max_start + 1):
-                self.samples.append((file_idx, startIDX))
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, index):
-        file_idx, startIDX = self.samples[index]
-        fn = self.file_prefix_list[file_idx]
-
-        data = np.load(fn, allow_pickle=True)
-        arr = data["arr_ztfg"]
-
-        #mjd = arr[:, 0]
-        #g_mag = arr[:, 1]
-
-        mjd = arr[:, 0]
-        g_mag = arr[:, 1].astype(np.float32)
-
-        # GLOBAL NORMALIZATION
-        g_mag = (g_mag - GLOBAL_GMAG_MEAN) / (GLOBAL_GMAG_STD + EPS)
-
-        t0 = mjd.min()
-        t_obs = mjd - t0
-
-        target_idx = startIDX + self.context_len
-
-        '''
-        # Input: [8, H, W], where channels encode previous timesteps.
-        # Right-align context. Unused earlier channels repeat earliest value.
-        context_img = torch.empty(self.n_channels, self.H, self.W, dtype=torch.float32)
-
-        earliest_mag = float(g_mag[startIDX])
-        context_img[:] = earliest_mag
-
-        offset = self.n_channels - self.context_len
-        for i in range(self.context_len):
-            ch = offset + i
-            context_img[ch] = float(g_mag[startIDX + i])
-
-        # Target: next scalar copied across all 8 channels.
-        target_mag = float(g_mag[target_idx])
-        target_img = torch.empty(self.n_channels, self.H, self.W, dtype=torch.float32)
-        target_img[:] = target_mag
-        '''
-
-        context_vals = torch.empty(self.n_channels, dtype=torch.float32)
-
-        earliest_mag = float(g_mag[startIDX])
-        context_vals[:] = earliest_mag
-
-        offset = self.n_channels - self.context_len
-        for i in range(self.context_len):
-            ch = offset + i
-            context_vals[ch] = float(g_mag[startIDX + i])
-
-        # expand() better for memory
-        context_img = context_vals.view(self.n_channels, 1, 1).expand(
-            self.n_channels,
-            self.H,
-            self.W,
-        )
-
-        # Predict DELTA (next - previous) rather than absolute next value
-        prev_mag = float(g_mag[target_idx - 1])
-        delta_mag = float(g_mag[target_idx] - g_mag[target_idx - 1])
-        target_val = torch.tensor(delta_mag, dtype=torch.float32)
-
-        target_img = target_val.view(1, 1, 1).expand(
-            self.n_channels,
-            self.H,
-            self.W,
-        )
-
-        Dt = torch.tensor(
-            t_obs[target_idx] - t_obs[target_idx - 1],
-            dtype=torch.float32,
-        )
-
-        data.close()
-
-        return context_img, target_img, Dt
-
-
-class ChannelStackAdapter(nn.Module):
-    """
-    Converts a sequence [B, K, C, H, W] into a fused image [B, C, H, W].
-    """
-
-    def __init__(self, in_channels: int, context_len: int, hidden_channels: int = 64):
-        super().__init__()
-        self.in_channels = in_channels
-        self.context_len = context_len
-        stacked_channels = in_channels * context_len
-
-        self.adapter = nn.Sequential(
-            nn.Conv2d(stacked_channels, hidden_channels, kernel_size=3, padding=1),
-            nn.GELU(),
-            nn.Conv2d(hidden_channels, in_channels, kernel_size=1),
-        )
-
-    def forward(self, x_seq: torch.Tensor) -> torch.Tensor:
-        """
-        x_seq: [B, K, C, H, W]
-        returns: [B, C, H, W]
-        """
-        if x_seq.ndim != 5:
-            raise ValueError(f"Expected x_seq to have shape [B, K, C, H, W], got {x_seq.shape}")
-
-        B, K, C, H, W = x_seq.shape
-        if C != self.in_channels:
-            raise ValueError(f"Expected {self.in_channels} channels, got {C}")
-        if K != self.context_len:
-            raise ValueError(f"Expected context_len={self.context_len}, got K={K}")
-
-        x = x_seq.reshape(B, K * C, H, W)
-        return self.adapter(x)
-
-
-class TemporalLodeRunner(nn.Module):
-    """
-    Wraps a pretrained one-step LodeRunner with a temporal adapter.
-
-    Input:
-        x_seq: [B, K, C, H, W]
-        in_vars, out_vars, Dt: same as original LodeRunner API
-
-    Output:
-        pred: [B, C, H, W]
-    """
-
-    def __init__(
-        self,
-        backbone: nn.Module,
-        in_channels: int = 8,
-        context_len: int = 3,
-        hidden_channels: int = 64,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.temporal_adapter = ChannelStackAdapter(
-            in_channels=in_channels,
-            context_len=context_len,
-            hidden_channels=hidden_channels,
-        )
-
-    def forward(
-        self,
-        x_seq: torch.Tensor,
-        in_vars: torch.Tensor,
-        out_vars: torch.Tensor,
-        Dt: torch.Tensor,
-    ) -> torch.Tensor:
-        fused_x = self.temporal_adapter(x_seq)          # [B, C, H, W]
-        pred = self.backbone(fused_x, in_vars, out_vars, Dt)
-        return pred
-
-
-def load_direct_loderunner_checkpoint(
-    checkpoint_path,
-    model_args,
-    optimizer_kwargs,
-    device,
-):
-    checkpoint_data = torch.load(
-        checkpoint_path,
-        map_location=device,
-        weights_only=False,
-    )
-
-    saved_model_args = checkpoint_data.get("model_args", model_args)
-
-    model = LodeRunner(**saved_model_args)
-    model.to(device)
-
-    state_dict = checkpoint_data["model_state_dict"]
-
-    if all(k.startswith("module.") for k in state_dict.keys()):
-        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
-
-    model.load_state_dict(state_dict, strict=True)
-
-    noise_scale = checkpoint_data.get("noise_scale", 0.0)
-    model.noise_scale = noise_scale
-
-    optimizer = torch.optim.AdamW(
-        model.parameters(),
-        **optimizer_kwargs,
-    )
-
-    if "optimizer_state_dict" in checkpoint_data:
-        optimizer.load_state_dict(checkpoint_data["optimizer_state_dict"])
-
-        for state in optimizer.state.values():
-            for key, value in state.items():
-                if isinstance(value, torch.Tensor):
-                    state[key] = value.to(device)
-
-    starting_epoch = checkpoint_data["epoch"]
-
-    return model, optimizer, starting_epoch
-
-
-def setup_distributed():
-    # ----- 1) Basic setup & environment variables -----
-    # Rely on Slurm variables: SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, etc.
-    rank = int(os.environ["SLURM_PROCID"])  # global rank
-    world_size = int(os.environ["SLURM_NTASKS"])  # total number of processes
-    local_rank = int(os.environ["SLURM_LOCALID"])  # local rank (GPU index on this node)
-
-    master_addr = os.environ["MASTER_ADDR"]
-    master_port = os.environ["MASTER_PORT"]
-
-    # ----- 2) Set the current GPU device for this process -----
-    torch.cuda.set_device(local_rank)
-    device = torch.device(f"cuda:{local_rank}")
-
-    # ----- 3) Initialize the process group -----
-    dist.init_process_group(
-        backend="nccl",
-        init_method=f"tcp://{master_addr}:{master_port}",
-        world_size=world_size,
-        rank=rank,
-    )
-
-    return rank, world_size, local_rank, device
-
-
-def cleanup_distributed():
-    # ----- 8) Clean up (optional) -----
-    dist.destroy_process_group()
-
-
-def main(args, rank, world_size, local_rank, device):
-    #############################################
-    # Process Inputs
-    #############################################
-    # Study ID
-    studyIDX = args.studyIDX
-
-    # Resources
-    Ngpus = args.Ngpus
-    Knodes = args.Knodes
-
-    # Data Paths
-    train_filelist = args.FILELIST_DIR + args.train_filelist
-    validation_filelist = args.FILELIST_DIR + args.validation_filelist
-
-    # Model Parameters
-    embed_dim = args.embed_dim
-    block_structure = tuple(args.block_structure)
-
-    # Training Parameters
-    anchor_lr = args.anchor_lr
-    num_cycles = args.num_cycles
-    min_fraction = args.min_fraction
-    terminal_steps = args.terminal_steps
-    warmup_steps = args.warmup_steps
-    noise_scale = args.noise_scale
-
-    # Number of workers controls how batches of data are prefetched and,
-    # possibly, pre-loaded onto GPUs. If the number of workers is large they
-    # will swamp memory and jobs will fail.
-    num_workers = args.num_workers
-
-    # Epoch Parameters
-    batch_size = args.batch_size
-    total_epochs = args.total_epochs
-    cycle_epochs = args.cycle_epochs
-    train_batches = args.train_batches
-    val_batches = args.val_batches
-    train_per_val = args.TRAIN_PER_VAL
-    trn_rcrd_filename = args.trn_rcrd_filename
-    val_rcrd_filename = args.val_rcrd_filename
-    CONTINUATION = args.continuation
-    checkpoint = args.checkpoint
-
-    #############################################
-    # Model Arguments for Dynamic Reconstruction
-    #############################################
-    # Dictionary of available models.
-    available_models = {
-        "LodeRunner": LodeRunner
-    }
-
-    # Model arguments for LodeRunner.
-    model_args = {
-        "default_vars": [
-            "density_case",
-            "density_cushion",
-            "density_maincharge",
-            "density_outside_air",
-            "density_striker",
-            "density_throw",
-            "Uvelocity",
-            "Wvelocity",
-        ],
-        "image_size": (1120, 400),
-        "patch_size": (10, 5),
-        "embed_dim": embed_dim,
-        "emb_factor": 2,
-        "num_heads": 8,
-        "block_structure": block_structure,
-        "window_sizes": [(8, 8), (8, 8), (4, 4), (2, 2)],
-        "patch_merge_scales": [(2, 2), (2, 2), (2, 2)],
-        #"noise_scale": noise_scale,
-    }
-
-
-    CONTEXT_LEN = 5 #3
-    HIDDEN_CHANNELS = 64
-
-    optimizer_kwargs = {
-        "lr": 1e-5,
-        "betas": (0.9, 0.999),
-        "eps": 1e-08,
-        "weight_decay": 0.01,
-    }
-
-
-    if CONTINUATION:
-        model, optimizer, starting_epoch = load_direct_loderunner_checkpoint(
-            checkpoint_path=checkpoint,
-            model_args=model_args,
-            optimizer_kwargs=optimizer_kwargs,
-            device=device,
-        )
-
-        if rank == 0:
-            print(f"Loaded direct checkpoint from {checkpoint}")
-            print(f"Continuing from epoch {starting_epoch}")
-
-        '''  # FIXME block should be unindented if uncommented 
-        if CONTINUATION:
-            model, optimizer, starting_epoch = load_model_and_optimizer(
-                checkpoint,
-                optimizer_class=torch.optim.AdamW,
-                optimizer_kwargs=optimizer_kwargs,
-                available_models=available_models,
-                device=device,
-            )
-
-            if rank == 0:
-                print(f"Loaded temporal checkpoint from {checkpoint}")
-                print(f"Continuing from epoch {starting_epoch}")
-        '''
-
-    else:
-        starting_epoch = 0
-
-        model = LodeRunner(**model_args)
-        model.to(device)
-
-        manual_checkpoint = "/usr/projects/artimis/mpmm/pretrained_models/ddp_ldr_prod_250721/study005_modelState_epoch0100.pth"
-
-        checkpoint_data = torch.load(
-            manual_checkpoint,
-            map_location=device,
-            weights_only=False,
-        )
-
-        state_dict = checkpoint_data["model_state_dict"]
-
-        if all(k.startswith("module.") for k in state_dict.keys()):
-            state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
-
-        missing_keys, unexpected_keys = model.load_state_dict(
-            state_dict,
-            strict=False,
-        )
-
-        if rank == 0:
-            print("Loaded pretrained backbone weights.")
-            print("Missing keys:", missing_keys)
-            print("Unexpected keys:", unexpected_keys)
-
-        model.noise_scale = noise_scale
-
-        # End-to-end fine-tuning: train adapter + backbone
-        for p in model.parameters():
-            p.requires_grad = True
-
-        optimizer = torch.optim.AdamW(
-            model.parameters(),
-            **optimizer_kwargs,
-        )
-
-    loss_fn = nn.MSELoss(reduction="none")
-
-    model = DDP(model, device_ids=[local_rank], output_device=local_rank)
-
-    #############################################
-    # Learning Rate Scheduler
-    #############################################
-    print("Starting epoch: ", starting_epoch)
-    if starting_epoch == 0:
-        last_epoch = -1
-    else:
-        last_epoch = train_batches * (starting_epoch - 1)
-
-    # Scale the anchor LR by global batchsize
-    #
-    # # For multi-node
-    lr_scale = np.sqrt(float(Ngpus) * float(Knodes) * float(batch_size))
-    original_batchsize = 40.0  # 1 node, 4 gpus, 10 samples/gpu
-    ddp_anchor_lr = anchor_lr * lr_scale / original_batchsize
-    #
-    # For single node
-    # ddp_anchor_lr = anchor_lr
-
-    LRsched = CosineWithWarmupScheduler(
-        optimizer,
-        anchor_lr=ddp_anchor_lr,
-        terminal_steps=terminal_steps,
-        warmup_steps=warmup_steps,
-        num_cycles=num_cycles,
-        min_fraction=min_fraction,
-        last_epoch=last_epoch,
-    )
-
-    #############################################
-    # Data Initialization (Distributed Dataloader)
-    #############################################
-    #train_dataset = LSC_rho2rho_temporal_DataSet(
-    #    args.LSC_NPZ_DIR,
-    #    file_prefix_list=train_filelist,
-    #    max_timeIDX_offset=2,
-    #    max_file_checks=10,
-    #    half_image=True,
-    #)
-    #val_dataset = LSC_rho2rho_temporal_DataSet(
-    #    args.LSC_NPZ_DIR,
-    #    file_prefix_list=validation_filelist,
-    #    max_timeIDX_offset=2,
-    #    max_file_checks=10,
-    #    half_image=True,
-    #)
-
-    '''
-    train_dataset = Kilonova_lc_img_DataSet_seq(
-        half_image=False,
-    )
-    val_dataset = Kilonova_lc_img_DataSet_seq(
-        half_image=False,
-    )
-    '''
-
-    train_dataset = Kilonova_lc_img_DataSet_channels_context(
-        half_image=False,
-        context_len=CONTEXT_LEN,
-        #N_imgs=100,
-    )
-
-    val_dataset = Kilonova_lc_img_DataSet_channels_context(
-        half_image=False,
-        context_len=CONTEXT_LEN,
-        #N_imgs=20, #100,
-    )
-
-    # NOTE: For DDP the batch_size is the per-GPU batch_size!!!
-    train_dataloader = make_distributed_dataloader(
-        train_dataset,
-        batch_size,
-        shuffle=True,
-        num_workers=num_workers,
-        rank=rank,
-        world_size=world_size,
-    )
-    val_dataloader = make_distributed_dataloader(
-        val_dataset,
-        batch_size,
-        shuffle=False,
-        num_workers=num_workers,
-        rank=rank,
-        world_size=world_size,
-    )
-
-    #############################################
-    # Training Loop (Modified for DDP)
-    #############################################
-    # Train Model
-    print("Training Model . . .")
-    starting_epoch += 1
-    ending_epoch = min(starting_epoch + cycle_epochs, total_epochs + 1)
-
-    TIME_EPOCH = True
-    for epochIDX in range(starting_epoch, ending_epoch):
-        print('%%%%%%%%%%%%%')
-        print(epochIDX)
-        print('%%%%%%%%%%%%%')
-        train_sampler = train_dataloader.sampler
-        train_sampler.set_epoch(epochIDX)
-
-        # For timing epochs
-        if TIME_EPOCH:
-            # Synchronize before starting the timer
-            #dist.barrier()  # Ensure that all nodes sync
-            torch.cuda.synchronize(device)  # Ensure GPUs on each node sync
-            # Time each epoch and print to stdout
-            startTime = time.time()
-
-
-        train_DDP_loderunner_epoch(
-            training_data=train_dataloader,
-            validation_data=val_dataloader,
-            num_train_batches=train_batches,
-            num_val_batches=val_batches,
-            model=model,
-            optimizer=optimizer,
-            loss_fn=loss_fn,
-            LRsched=LRsched,
-            epochIDX=epochIDX,
-            train_per_val=train_per_val,
-            train_rcrd_filename=trn_rcrd_filename,
-            val_rcrd_filename=val_rcrd_filename,
-            device=device,
-            rank=rank,
-            world_size=world_size,
-            seq=False,
-        )
-
-        print(f"[rank {rank}] finished epoch", flush=True)
-
-
-        if TIME_EPOCH:
-            # Synchronize before stopping the timer
-            torch.cuda.synchronize(device)  # Ensure GPUs on each node sync
-            #dist.barrier()  # Ensure that all nodes sync
-            # Time each epoch and print to stdout
-            endTime = time.time()
-
-        epoch_time = (endTime - startTime) / 60
-
-        # Print Summary Results
-        if rank == 0:
-            print(f"Completed epoch {epochIDX}...", flush=True)
-            print(f"Epoch time (minutes): {epoch_time:.2f}", flush=True)
-
-        # Save model and optimizer
-        #chkpt_name_str = f'study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth'
-        #new_chkpt_path = os.path.join("./", chkpt_name_str)
-
-        if rank == 0:
-            chkpt_name_str = f"study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth"
-            new_chkpt_path = os.path.join("./", chkpt_name_str)
-
-            print(f"Saving checkpoint: {new_chkpt_path}", flush=True)
-
-            torch.save(
-                {
-                    "epoch": epochIDX,
-                    "model_class": "LodeRunner",
-                    "model_args": model_args,
-                    "model_state_dict": model.module.state_dict(),
-                    "optimizer_state_dict": optimizer.state_dict(),
-                    "noise_scale": noise_scale,
-                    "predicts_delta": True,
-                    "target_type": "delta",
-                    "context_len": CONTEXT_LEN,
-                },
-                new_chkpt_path,
-            )
-
-            '''
-            save_model_and_optimizer(
-                model.module,
-                optimizer,
-                epochIDX,
-                new_chkpt_path,
-                model_class=LodeRunner,
-                model_args=model_args,
-            )
-            '''
-
-            print(f"Saved checkpoint: {new_chkpt_path}", flush=True)
-
-        '''
-        if rank == 0:
-            chkpt_name_str = f"study{studyIDX:03d}_modelState_epoch{epochIDX:04d}.pth"
-            new_chkpt_path = os.path.join("./", chkpt_name_str)
-
-            #save_model_and_optimizer(
-            #    model,
-            #    optimizer,
-            #    epochIDX,
-            #    new_chkpt_path,
-            #    model_class=LodeRunner,
-            #    model_args=model_args,
-            #)
-
-            print(f"Saved checkpoint: {new_chkpt_path}", flush=True)
-        '''
-    '''
-        save_model_and_optimizer(
-            model,
-            optimizer,
-            epochIDX,
-            new_chkpt_path,
-            model_class=LodeRunner,
-            model_args=model_args,
-        )
-    '''
-    if rank == 0:
-        #############################################
-        # Continue if Necessary
-        #############################################
-        FINISHED_TRAINING = epochIDX + 1 > total_epochs
-        if not FINISHED_TRAINING:
-            new_slurm_file = continuation_setup(
-                new_chkpt_path, studyIDX, last_epoch=epochIDX
-            )
-            os.system(f"sbatch {new_slurm_file}")
-
-if __name__ == "__main__":
-    print('running main')
-    args = parser.parse_args()
-
-    rank, world_size, local_rank, device = setup_distributed()
-
-    main(args, rank, world_size, local_rank, device)
-
-    cleanup_distributed()
diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py
index 6077e4ff..41a89dea 100644
--- a/src/yoke/utils/training/datastep/loderunner.py
+++ b/src/yoke/utils/training/datastep/loderunner.py
@@ -136,467 +136,6 @@ def eval_DDP_scalar_temporal_loderunner_datastep_gri(
     return target, pred, per_sample_loss.detach()
 
 
-def eval_DDP_scalar_temporal_loderunner_datastep(
-    data: tuple,
-    model: torch.nn.Module,
-    loss_fn: torch.nn.Module,
-    device: torch.device,
-    rank: int,
-    world_size: int,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    DDP evaluation datastep for scalar temporal LodeRunner wrapper.
-
-    Expected data:
-        x:      [B, input_dim]
-        target: [B, 3]
-        Dt:     [B]
-
-    Expected model output:
-        pred:   [B, 3]
-    """
-
-    model.eval()
-
-    x, target, Dt = data
-
-    x = x.to(device, non_blocking=True)
-    target = target.to(device, non_blocking=True)
-    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
-
-    in_vars = torch.arange(8, device=device)
-    out_vars = torch.arange(8, device=device)
-
-    with torch.no_grad():
-        pred = model(x, in_vars, out_vars, Dt)
-
-    if pred.shape != target.shape:
-        raise RuntimeError(
-            f"Validation prediction and target shapes do not match: "
-            f"pred.shape={pred.shape}, target.shape={target.shape}"
-        )
-
-    loss = loss_fn(pred, target)
-
-    if loss.ndim == 1:
-        per_sample_loss = loss
-    else:
-        per_sample_loss = loss.mean(dim=tuple(range(1, loss.ndim)))
-
-    return target, pred, per_sample_loss.detach()
-
-
-def train_loderunner_datastep(
-    data: tuple,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    device: torch.device,
-    channel_map: list[int],
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """A training step for which the data is of multi-input, multi-output type.
-
-    This is currently a proto-type function to get the LodeRunner architecture
-    training on a non-variable set of channels.
-
-    Args:
-        data (tuple): tuple of model input, corresponding ground truth, and lead time
-        model (torch.nn.Module): model to train
-        optimizer (torch.optim.Optimizer): optimizer for training set
-        loss_fn (torch.nn.Module): loss function for training set
-        device (torch.device): device index to select
-        channel_map (list): list of channel indices to use
-
-    Returns:
-        end_img (torch.Tensor): Ground truth end image
-        pred_img (torch.Tensor): Predicted end image
-        per_sample_loss (torch.Tensor): Per-sample loss for the batch
-    """
-    # Set model to train
-    model.train()
-
-    # Extract data
-    #(start_img, end_img, Dt) = data
-    img_seq, Dt = data
-
-    #start_img = start_img.to(device, non_blocking=True)
-    #Dt = Dt.to(torch.float32).to(device, non_blocking=True)
-    #end_img = end_img.to(device, non_blocking=True)
-
-    img_seq = img_seq.to(device, non_blocking=True)
-    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
-
-    start_img = img_seq[:, 0]
-    end_img = img_seq[:, -1]
-
-    # For our first LodeRunner training on the lsc240420 dataset the input and
-    # output prediction variables are fixed.
-    #
-    # Both in_vars and out_vars correspond to indices for every variable in
-    # this training setup...
-    #
-    # in_vars = ['density_case',
-    #            'density_cushion',
-    #            'density_maincharge',
-    #            'density_outside_air',
-    #            'density_striker',
-    #            'density_throw',
-    #            'Uvelocity',
-    #            'Wvelocity']
-    in_vars = torch.tensor(channel_map).to(device, non_blocking=True)
-    out_vars = torch.tensor(channel_map).to(device, non_blocking=True)
-
-    # Perform a forward pass
-    # NOTE: If training on GPU model should have already been moved to GPU
-    # prior to initalizing optimizer.
-    print("start_img entering model:", start_img.shape)   # expect [B, 8, 1120, 800]
-    print("len(in_vars):", len(in_vars))
-    print("in_vars:", in_vars)
-    pred_img = model(start_img, in_vars, out_vars, Dt)
-
-    # Expecting to use a *reduction="none"* loss function so we can track loss
-    # between individual samples. However, this will make the loss be computed
-    # element-wise so we need to still average over the (channel, height,
-    # width) dimensions to get the per-sample loss.
-    loss = loss_fn(pred_img, end_img)
-    per_sample_loss = loss.mean(dim=[1, 2, 3])  # Shape: (batch_size,)
-
-    # Perform backpropagation and update the weights
-    optimizer.zero_grad(set_to_none=True)  # Possible speed-up
-    loss.mean().backward()
-    optimizer.step()
-
-    return end_img, pred_img, per_sample_loss
-
-
-def train_scheduled_loderunner_datastep(
-    data: tuple,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    device: torch.device,
-    scheduled_prob: float,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Trainning step for LodeRunner with scheduled sampling.
-
-    This training step implements scheduled sampling, where the model
-    can either use the ground truth image or its own prediction as input.
-
-    Args:
-        data (tuple): Sequence of images in (img_seq, Dt) tuple.
-        model (loaded pytorch model): model to train.
-        optimizer (torch.optim): optimizer for training set.
-        loss_fn (torch.nn Loss Function): loss function for training set.
-        device (torch.device): device index to select.
-        scheduled_prob (float): Probability of using the ground truth as input.
-
-    Returns:
-        img_seq (torch.Tensor): Ground truth image sequence.
-        pred_seq (torch.Tensor): Predicted image sequence.
-        per_sample_loss (torch.Tensor): Per-sample loss for the batch.
-    """
-    # Set model to train
-    model.train()
-
-    # Extract data
-    img_seq, Dt = data
-
-    # [B, S, C, H, W] where S=seq-length
-    img_seq = img_seq.to(device, non_blocking=True)
-    # [B, 1]
-    Dt = Dt.to(device, non_blocking=True)
-
-    # Input and output variable indices
-    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
-    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
-
-    # Storage for predictions at each timestep
-    pred_seq = []
-
-    # Unbind and iterate over slices in sequence-length dimension
-    # NOTE: we exclude img_seq[:, :-1] since we don't have the next
-    #   timestep to compare to.
-    for k, k_img in enumerate(torch.unbind(img_seq[:, :-1], dim=1)):
-        if k == 0:
-            # Forward pass for the initial step
-            pred_img = model(k_img, in_vars, out_vars, Dt)
-        else:
-            # Apply scheduled sampling
-            if random.random() < scheduled_prob:
-                current_input = k_img
-            else:
-                current_input = pred_img
-
-            pred_img = model(current_input, in_vars, out_vars, Dt)
-
-        # Store the prediction
-        pred_seq.append(pred_img)
-
-    # Combine predictions into a tensor of shape [B, SeqLength, C, H, W]
-    pred_seq = torch.stack(pred_seq, dim=1)
-
-    # Compute loss
-    loss = loss_fn(pred_seq, img_seq[:, 1:])
-    per_sample_loss = loss.mean(dim=[1, 2, 3, 4])  # Shape: (batch_size,)
-
-    # Perform backpropagation and update the weights
-    optimizer.zero_grad(set_to_none=True)  # Possible speed-up
-    loss.mean().backward()
-    optimizer.step()
-
-    return img_seq[:, 1:], pred_seq, per_sample_loss
-
-
-def train_DDP_loderunner_datastep(
-    data: tuple,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    device: torch.device,
-    rank: int,
-    world_size: int,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """A DDP-compatible training step for multi-input, multi-output data.
-
-    Args:
-        data (tuple): tuple of model input, corresponding ground truth, and lead time
-        model (loaded pytorch model): model to train
-        optimizer (torch.optim): optimizer for training set
-        loss_fn (torch.nn Loss Function): loss function for training set
-        device (torch.device): device index to select
-        rank (int): Rank of device
-        world_size (int): Number of total DDP processes
-
-    Returns:
-        end_img (torch.Tensor): Ground truth end image
-        pred_img (torch.Tensor): Predicted end image
-        all_losses (torch.Tensor): Concatenated per-sample losses from all processes
-    """
-    # Set model to train mode
-    model.train()
-
-    # Extract data
-    start_img, end_img, Dt = data
-    start_img = start_img.to(device, non_blocking=True)
-    Dt = Dt.to(device, non_blocking=True)
-    end_img = end_img.to(device, non_blocking=True)
-
-    # Fixed input and output variable indices
-    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
-    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
-
-    # Forward pass
-    pred_img = model(start_img, in_vars, out_vars, Dt)
-
-    # Compute loss
-    loss = loss_fn(pred_img, end_img)
-    per_sample_loss = loss.mean(dim=[1, 2, 3])  # Per-sample loss
-
-    # Backward pass and optimization
-    optimizer.zero_grad(set_to_none=True)
-    loss.mean().backward()
-    optimizer.step()
-
-    # Gather per-sample losses from all processes
-    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
-    dist.all_gather(gathered_losses, per_sample_loss)
-
-    # Rank 0 concatenates and saves or returns all losses
-    if rank == 0:
-        all_losses = torch.cat(gathered_losses, dim=0)  # Shape: (total_batch_size,)
-    else:
-        all_losses = None
-
-    return end_img, pred_img, all_losses
-
-
-def train_DDP_loderunner_seq_channel_datastep(
-    data,
-    model,
-    optimizer,
-    loss_fn,
-    device,
-    rank,
-    world_size,
-):
-    model.train()
-
-    start_img, end_img, Dt = data
-
-    start_img = start_img.to(device, non_blocking=True)
-    end_img = end_img.to(device, non_blocking=True)
-    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
-
-    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
-    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7], device=device)
-
-    pred_img = model(start_img, in_vars, out_vars, Dt)
-
-    loss = loss_fn(pred_img, end_img)
-    per_sample_loss = loss.mean(dim=[1, 2, 3])
-
-    optimizer.zero_grad(set_to_none=True)
-    per_sample_loss.mean().backward()
-    optimizer.step()
-
-    return end_img, pred_img, per_sample_loss.detach()
-
-
-def train_DDP_loderunner_seq_datastep(
-    data: tuple,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    device: torch.device,
-    rank: int,
-    world_size: int,
-    scheduled_prob: float = 1.0,
-    channel_map: list[int] | None = None,
-):
-    """
-    DDP training step for autoregressive sequence training.
-
-    Expected data:
-        img_seq, Dt = data
-
-    Shapes:
-        img_seq: [B, S, C, H, W]
-        Dt:
-            either [B] / [B, 1] for a single constant Dt reused at every step,
-            or [B, S-1] / [B, S-1, 1] for per-step lead times.
-
-    Returns:
-        gt_seq:      [B, S-1, C, H, W]
-        pred_seq:    [B, S-1, C, H, W]
-        all_losses:  concatenated per-sample losses on rank 0, else None
-    """
-    model.train()
-
-    img_seq, Dt = data
-    img_seq = img_seq.to(device, non_blocking=True)
-    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
-
-    if channel_map is None:
-        channel_map = [0, 1, 2, 3, 4, 5, 6, 7]
-
-    in_vars = torch.tensor(channel_map, device=device)
-    out_vars = torch.tensor(channel_map, device=device)
-
-    B, S, C, H, W = img_seq.shape
-    assert S >= 2, "Sequence length must be at least 2."
-
-    pred_seq = []
-
-    # initial input is first frame
-    current_input = img_seq[:, 0]
-
-    for k in range(S - 1):
-        # support either one Dt for all steps or one Dt per step
-        if Dt.ndim == 1 or (Dt.ndim == 2 and Dt.shape[-1] == 1):
-            Dt_k = Dt
-        elif Dt.ndim == 2:
-            Dt_k = Dt[:, k].unsqueeze(-1)
-        elif Dt.ndim == 3:
-            Dt_k = Dt[:, k]
-        else:
-            raise ValueError(f"Unsupported Dt shape: {Dt.shape}")
-
-        pred_img = model(current_input, in_vars, out_vars, Dt_k)
-        pred_seq.append(pred_img)
-
-        if k < S - 2:
-            if random.random() < scheduled_prob:
-                current_input = img_seq[:, k + 1]   # teacher forcing
-            else:
-                current_input = pred_img.detach()   # autoregressive rollout
-
-    pred_seq = torch.stack(pred_seq, dim=1)   # [B, S-1, C, H, W]
-    gt_seq = img_seq[:, 1:]                   # [B, S-1, C, H, W]
-
-    loss = loss_fn(pred_seq, gt_seq)
-    per_sample_loss = loss.mean(dim=[1, 2, 3, 4])
-
-    optimizer.zero_grad(set_to_none=True)
-    loss.mean().backward()
-    optimizer.step()
-
-    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
-    dist.all_gather(gathered_losses, per_sample_loss)
-
-    if rank == 0:
-        all_losses = torch.cat(gathered_losses, dim=0)
-    else:
-        all_losses = None
-
-    return gt_seq, pred_seq, all_losses
-
-
-def train_DDP_loderunner_datastep_seq_old(
-    data: tuple,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    device: torch.device,
-    rank: int,
-    world_size: int,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """A DDP-compatible training step for multi-input, multi-output data.
-
-    Args:
-        data (tuple): tuple of model input, corresponding ground truth, and lead time
-        model (loaded pytorch model): model to train
-        optimizer (torch.optim): optimizer for training set
-        loss_fn (torch.nn Loss Function): loss function for training set
-        device (torch.device): device index to select
-        rank (int): Rank of device
-        world_size (int): Number of total DDP processes
-
-    Returns:
-        end_img (torch.Tensor): Ground truth end image
-        pred_img (torch.Tensor): Predicted end image
-        all_losses (torch.Tensor): Concatenated per-sample losses from all processes
-    """
-    # Set model to train mode
-    model.train()
-
-    # Extract data
-    #start_img, end_img, Dt = data
-    img_seq, Dt = data
-    #for img in img_seq:
-    #    # ...
-    start_img = start_img.to(device, non_blocking=True)
-    Dt = Dt.to(device, non_blocking=True)
-    end_img = end_img.to(device, non_blocking=True)
-
-    # Fixed input and output variable indices
-    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
-    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
-
-    # Forward pass
-    pred_img = model(start_img, in_vars, out_vars, Dt)
-
-    # Compute loss
-    loss = loss_fn(pred_img, end_img)
-    per_sample_loss = loss.mean(dim=[1, 2, 3])  # Per-sample loss
-
-    # Backward pass and optimization
-    optimizer.zero_grad(set_to_none=True)
-    loss.mean().backward()
-    optimizer.step()
-
-    # Gather per-sample losses from all processes
-    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
-    dist.all_gather(gathered_losses, per_sample_loss)
-
-    # Rank 0 concatenates and saves or returns all losses
-    if rank == 0:
-        all_losses = torch.cat(gathered_losses, dim=0)  # Shape: (total_batch_size,)
-    else:
-        all_losses = None
-
-    return end_img, pred_img, all_losses
-
-
 ####################################
 # Evaluating on a Datastep
 ####################################
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 18c7ca27..f471916c 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -11,14 +11,8 @@
     train_scheduled_loderunner_datastep,
     eval_scheduled_loderunner_datastep,
     train_DDP_loderunner_datastep,
-    train_DDP_loderunner_seq_datastep,
-    train_DDP_loderunner_seq_channel_datastep,
-    train_DDP_temporal_loderunner_datastep,
     train_DDP_scalar_temporal_loderunner_datastep_gri,
     eval_DDP_loderunner_datastep,
-    eval_DDP_loderunner_seq_datastep,
-    eval_DDP_loderunner_seq_context_datastep,
-    eval_DDP_loderunner_seq_channel_datastep,
     eval_DDP_scalar_temporal_loderunner_datastep_gri,
 )
 

From 8b625c0e72671d0e93cbeef2550e771678a48c8d Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 30 Jul 2026 16:26:29 -0600
Subject: [PATCH 11/66] Refactoring

---
 src/yoke/utils/checkpointing.py | 131 ++++++++++++++++++++++++++++++++
 src/yoke/utils/parallel.py      |  45 +++++++++++
 2 files changed, 176 insertions(+)

diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index 5ca580f7..4f407e64 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -10,6 +10,11 @@
 import torch.distributed as dist
 import h5py
 
+from yoke.models.vit.swin.bomberman import (
+    LodeRunner,
+    ScalarTemporalConditionedLodeRunner_gri,
+)
+
 
 def save_model_and_optimizer_hdf5(
     model: torch.nn.Module,
@@ -303,3 +308,129 @@ def load_model_and_optimizer(
         dist.barrier()
 
     return model, optimizer, checkpoint["epoch"]
+
+
+def load_direct_loderunner_checkpoint(
+    checkpoint_path: str,
+    model_args: dict,
+    optimizer_kwargs: dict,
+    device: torch.device,
+) -> tuple[torch.nn.Module, torch.optim.Optimizer, int]:
+    """Load a ScalarTemporalConditionedLodeRunner_gri model from a checkpoint.
+
+    Handles two checkpoint types:
+      - An old plain LodeRunner checkpoint, whose weights are loaded into the
+        wrapper's backbone (conditioner/output-head are freshly initialized and
+        this is not treated as a true continuation).
+      - A ScalarTemporalConditionedLodeRunner_gri wrapper checkpoint, which is
+        loaded in full and treated as a continuation.
+
+    The backbone is frozen and only the conditioner and output-head parameters
+    are trainable.
+
+    Args:
+        checkpoint_path (str): Path to the checkpoint file.
+        model_args (dict): Fallback LodeRunner init args if the checkpoint has
+            none stored.
+        optimizer_kwargs (dict): Kwargs for the AdamW optimizer.
+        device (torch.device): Device to load the model/optimizer onto.
+
+    Returns:
+        model (torch.nn.Module): The wrapper model.
+        optimizer (torch.optim.Optimizer): Optimizer over trainable parameters.
+        starting_epoch (int): Epoch to continue training from.
+    """
+    checkpoint_data = torch.load(
+        checkpoint_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    saved_model_args = checkpoint_data.get("model_args", model_args)
+    context_len = checkpoint_data.get("context_len", 5)
+
+    backbone = LodeRunner(**saved_model_args).to(device)
+
+    model = ScalarTemporalConditionedLodeRunner_gri(
+        backbone=backbone,
+        context_len=context_len,
+        n_input_channels=checkpoint_data.get("n_input_channels", 3),
+        n_output_channels=checkpoint_data.get("n_output_channels", 3),
+        image_size=saved_model_args["image_size"],
+        backbone_channels=checkpoint_data.get("backbone_channels", 8),
+        hidden=checkpoint_data.get("hidden", 64),
+    ).to(device)
+
+    state_dict = checkpoint_data["model_state_dict"]
+
+    # Remove DDP prefix if present
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+
+    # Detect checkpoint type
+    is_wrapper_checkpoint = any(
+        k.startswith("backbone.") for k in state_dict.keys()
+    )
+
+    # -------------------------------------------------
+    # OLD plain LodeRunner checkpoint
+    # -------------------------------------------------
+    if not is_wrapper_checkpoint:
+        missing_keys, unexpected_keys = model.backbone.load_state_dict(
+            state_dict,
+            strict=False,
+        )
+
+        print("Loaded old LodeRunner checkpoint into model.backbone")
+        print("Missing backbone keys:", missing_keys)
+        print("Unexpected backbone keys:", unexpected_keys)
+
+        # This is NOT a true continuation.
+        # Conditioner is newly initialized.
+        starting_epoch = 0
+
+    # -------------------------------------------------
+    # NEW ScalarTemporalConditionedLodeRunner checkpoint
+    # -------------------------------------------------
+    else:
+        model.load_state_dict(state_dict, strict=True)
+
+        print("Loaded ScalarTemporalConditionedLodeRunner checkpoint")
+
+        starting_epoch = checkpoint_data.get("epoch", 0)
+
+    noise_scale = checkpoint_data.get("noise_scale", 0.0)
+    model.backbone.noise_scale = noise_scale
+
+    # Freeze pretrained backbone
+    for p in model.backbone.parameters():
+        p.requires_grad = False
+
+    # Train conditioner
+    for p in model.conditioner.parameters():
+        p.requires_grad = True
+
+    optimizer = torch.optim.AdamW(
+        list(model.conditioner.parameters())
+        + list(model.output_head.parameters()),
+        **optimizer_kwargs,
+    )
+
+    # Only restore optimizer for TRUE continuation checkpoints
+    if (
+        is_wrapper_checkpoint
+        and "optimizer_state_dict" in checkpoint_data
+    ):
+        optimizer.load_state_dict(
+            checkpoint_data["optimizer_state_dict"]
+        )
+
+        for state in optimizer.state.values():
+            for key, value in state.items():
+                if isinstance(value, torch.Tensor):
+                    state[key] = value.to(device)
+
+    return model, optimizer, starting_epoch
diff --git a/src/yoke/utils/parallel.py b/src/yoke/utils/parallel.py
index 66cbb375..30d524df 100644
--- a/src/yoke/utils/parallel.py
+++ b/src/yoke/utils/parallel.py
@@ -5,8 +5,53 @@
 
 """
 
+import os
+
 import torch
 import torch.nn as nn
+import torch.distributed as dist
+
+
+def setup_distributed() -> tuple[int, int, int, torch.device]:
+    """Initialize the DDP process group from Slurm environment variables.
+
+    Relies on the Slurm-provided variables SLURM_PROCID, SLURM_NTASKS, and
+    SLURM_LOCALID, along with MASTER_ADDR and MASTER_PORT, to set up an NCCL
+    process group and select this process's GPU.
+
+    Returns:
+        rank (int): Global rank of this process.
+        world_size (int): Total number of processes.
+        local_rank (int): Local rank (GPU index) on this node.
+        device (torch.device): CUDA device for this process.
+    """
+    # ----- 1) Basic setup & environment variables -----
+    # Rely on Slurm variables: SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, etc.
+    rank = int(os.environ["SLURM_PROCID"])  # global rank
+    world_size = int(os.environ["SLURM_NTASKS"])  # total number of processes
+    local_rank = int(os.environ["SLURM_LOCALID"])  # local rank (GPU index on node)
+
+    master_addr = os.environ["MASTER_ADDR"]
+    master_port = os.environ["MASTER_PORT"]
+
+    # ----- 2) Set the current GPU device for this process -----
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+
+    # ----- 3) Initialize the process group -----
+    dist.init_process_group(
+        backend="nccl",
+        init_method=f"tcp://{master_addr}:{master_port}",
+        world_size=world_size,
+        rank=rank,
+    )
+
+    return rank, world_size, local_rank, device
+
+
+def cleanup_distributed() -> None:
+    """Destroy the DDP process group."""
+    dist.destroy_process_group()
 
 
 # Custom nn.DataParallel class to handle input to LodeRunner that should not be

From 0801b75a8d951972cf90f767cbdd5cd1742eb3f3 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 30 Jul 2026 16:45:43 -0600
Subject: [PATCH 12/66] clean-up

---
 .../utils/training/datastep/loderunner.py     | 245 ++++++++++++-
 src/yoke/utils/training/epoch/loderunner.py   | 321 +-----------------
 2 files changed, 243 insertions(+), 323 deletions(-)

diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py
index 41a89dea..2ec7eb4d 100644
--- a/src/yoke/utils/training/datastep/loderunner.py
+++ b/src/yoke/utils/training/datastep/loderunner.py
@@ -19,8 +19,7 @@ def train_DDP_scalar_temporal_loderunner_datastep_gri(
     rank: int,
     world_size: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """
-    DDP training datastep for scalar temporal LodeRunner wrapper.
+    """DDP training datastep for scalar temporal LodeRunner wrapper.
 
     Expected data:
         x:      [B, input_dim]
@@ -34,7 +33,6 @@ def train_DDP_scalar_temporal_loderunner_datastep_gri(
     Expected model output:
         pred:   [B, 3]
     """
-
     model.train()
 
     x, target, Dt = data
@@ -84,8 +82,7 @@ def eval_DDP_scalar_temporal_loderunner_datastep_gri(
     rank,
     world_size,
 ):
-    """
-    Evaluation datastep for ScalarTemporalConditionedLodeRunner.
+    """Evaluation datastep for ScalarTemporalConditionedLodeRunner.
 
     Expected dataset output:
         x:      [B, input_dim]
@@ -99,7 +96,6 @@ def eval_DDP_scalar_temporal_loderunner_datastep_gri(
     Expected model output:
         pred:   [B, 3]
     """
-
     model.eval()
 
     x, target, Dt = data
@@ -214,6 +210,7 @@ def eval_DDP_loderunner_seq_channel_datastep(
     rank,
     world_size,
 ):
+    """DDP eval step for channel-stacked LodeRunner on a single next-step target."""
     model.eval()
 
     start_img, end_img, Dt = data
@@ -243,6 +240,7 @@ def eval_DDP_loderunner_seq_datastep(
     rank,
     world_size,
 ):
+    """DDP eval step rolling LodeRunner over a full image sequence."""
     model.eval()
 
     img_seq, Dt = data
@@ -423,6 +421,7 @@ def train_DDP_temporal_loderunner_datastep(
     rank,
     world_size,
 ):
+    """DDP training step for temporal LodeRunner on a channel-stacked context."""
     model.train()
 
     context_seq, target_img, Dt = data
@@ -464,6 +463,7 @@ def train_DDP_scalar_temporal_loderunner_datastep(
     rank,
     world_size,
 ):
+    """DDP training step for scalar temporal LodeRunner on flattened context."""
     model.train()
 
     x, target, Dt = data
@@ -506,6 +506,7 @@ def eval_DDP_temporal_loderunner_datastep(
     rank,
     world_size,
 ):
+    """DDP eval step for temporal LodeRunner on a channel-stacked context."""
     model.eval()
 
     context_seq, target_img, Dt = data
@@ -542,14 +543,15 @@ def eval_DDP_scalar_temporal_loderunner_datastep(
     rank,
     world_size,
 ):
+    """DDP eval step for scalar temporal LodeRunner on flattened context."""
     model.eval()
 
     with torch.no_grad():
         x, target, Dt = data
 
-        x = x.to(torch.float32).to(device, non_blocking=True)          # [B, 2 * context_len]
-        target = target.to(torch.float32).to(device, non_blocking=True) # [B]
-        Dt = Dt.to(torch.float32).to(device, non_blocking=True)         # [B]
+        x = x.to(torch.float32).to(device, non_blocking=True)  # [B, 2 * context_len]
+        target = target.to(torch.float32).to(device, non_blocking=True)  # [B]
+        Dt = Dt.to(torch.float32).to(device, non_blocking=True)  # [B]
 
         C = 8
         in_vars = torch.arange(C, device=device)
@@ -580,8 +582,7 @@ def eval_DDP_loderunner_seq_context_datastep(
     rank,
     world_size,
 ):
-    """
-    DDP eval step for TemporalLodeRunner / channel-stacked context model.
+    """DDP eval step for TemporalLodeRunner / channel-stacked context model.
 
     Expected data:
         context_seq, target_img, Dt = data
@@ -591,9 +592,6 @@ def eval_DDP_loderunner_seq_context_datastep(
         target_img:  [B, C, H, W]
         Dt:          [B] or [B, 1]
     """
-    import torch
-    import torch.distributed as dist
-
     model.eval()
 
     context_seq, target_img, Dt = data
@@ -624,3 +622,222 @@ def eval_DDP_loderunner_seq_context_datastep(
         all_losses = None
 
     return target_img, pred_img, all_losses
+
+
+def train_loderunner_datastep(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    channel_map: list[int],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """A training step for which the data is of multi-input, multi-output type.
+
+    This is currently a proto-type function to get the LodeRunner architecture
+    training on a non-variable set of channels.
+
+    Args:
+        data (tuple): tuple of model input, corresponding ground truth, and lead time
+        model (torch.nn.Module): model to train
+        optimizer (torch.optim.Optimizer): optimizer for training set
+        loss_fn (torch.nn.Module): loss function for training set
+        device (torch.device): device index to select
+        channel_map (list): list of channel indices to use
+
+    Returns:
+        end_img (torch.Tensor): Ground truth end image
+        pred_img (torch.Tensor): Predicted end image
+        per_sample_loss (torch.Tensor): Per-sample loss for the batch
+    """
+    # Set model to train
+    model.train()
+
+    # Extract data
+    (start_img, end_img, Dt) = data
+
+    start_img = start_img.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+    end_img = end_img.to(device, non_blocking=True)
+
+    # For our first LodeRunner training on the lsc240420 dataset the input and
+    # output prediction variables are fixed.
+    #
+    # Both in_vars and out_vars correspond to indices for every variable in
+    # this training setup...
+    #
+    # in_vars = ['density_case',
+    #            'density_cushion',
+    #            'density_maincharge',
+    #            'density_outside_air',
+    #            'density_striker',
+    #            'density_throw',
+    #            'Uvelocity',
+    #            'Wvelocity']
+    in_vars = torch.tensor(channel_map).to(device, non_blocking=True)
+    out_vars = torch.tensor(channel_map).to(device, non_blocking=True)
+
+    # Perform a forward pass
+    # NOTE: If training on GPU model should have already been moved to GPU
+    # prior to initalizing optimizer.
+    pred_img = model(start_img, in_vars, out_vars, Dt)
+
+    # Expecting to use a *reduction="none"* loss function so we can track loss
+    # between individual samples. However, this will make the loss be computed
+    # element-wise so we need to still average over the (channel, height,
+    # width) dimensions to get the per-sample loss.
+    loss = loss_fn(pred_img, end_img)
+    per_sample_loss = loss.mean(dim=[1, 2, 3])  # Shape: (batch_size,)
+
+    # Perform backpropagation and update the weights
+    optimizer.zero_grad(set_to_none=True)  # Possible speed-up
+    loss.mean().backward()
+    optimizer.step()
+
+    return end_img, pred_img, per_sample_loss
+
+
+def train_scheduled_loderunner_datastep(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    scheduled_prob: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Trainning step for LodeRunner with scheduled sampling.
+
+    This training step implements scheduled sampling, where the model
+    can either use the ground truth image or its own prediction as input.
+
+    Args:
+        data (tuple): Sequence of images in (img_seq, Dt) tuple.
+        model (loaded pytorch model): model to train.
+        optimizer (torch.optim): optimizer for training set.
+        loss_fn (torch.nn Loss Function): loss function for training set.
+        device (torch.device): device index to select.
+        scheduled_prob (float): Probability of using the ground truth as input.
+
+    Returns:
+        img_seq (torch.Tensor): Ground truth image sequence.
+        pred_seq (torch.Tensor): Predicted image sequence.
+        per_sample_loss (torch.Tensor): Per-sample loss for the batch.
+    """
+    # Set model to train
+    model.train()
+
+    # Extract data
+    img_seq, Dt = data
+
+    # [B, S, C, H, W] where S=seq-length
+    img_seq = img_seq.to(device, non_blocking=True)
+    # [B, 1]
+    Dt = Dt.to(device, non_blocking=True)
+
+    # Input and output variable indices
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
+
+    # Storage for predictions at each timestep
+    pred_seq = []
+
+    # Unbind and iterate over slices in sequence-length dimension
+    # NOTE: we exclude img_seq[:, :-1] since we don't have the next
+    #   timestep to compare to.
+    for k, k_img in enumerate(torch.unbind(img_seq[:, :-1], dim=1)):
+        if k == 0:
+            # Forward pass for the initial step
+            pred_img = model(k_img, in_vars, out_vars, Dt)
+        else:
+            # Apply scheduled sampling
+            if random.random() < scheduled_prob:
+                current_input = k_img
+            else:
+                current_input = pred_img
+
+            pred_img = model(current_input, in_vars, out_vars, Dt)
+
+        # Store the prediction
+        pred_seq.append(pred_img)
+
+    # Combine predictions into a tensor of shape [B, SeqLength, C, H, W]
+    pred_seq = torch.stack(pred_seq, dim=1)
+
+    # Compute loss
+    loss = loss_fn(pred_seq, img_seq[:, 1:])
+    per_sample_loss = loss.mean(dim=[1, 2, 3, 4])  # Shape: (batch_size,)
+
+    # Perform backpropagation and update the weights
+    optimizer.zero_grad(set_to_none=True)  # Possible speed-up
+    loss.mean().backward()
+    optimizer.step()
+
+    return img_seq[:, 1:], pred_seq, per_sample_loss
+
+
+def train_DDP_loderunner_datastep(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """A DDP-compatible training step for multi-input, multi-output data.
+
+    Args:
+        data (tuple): tuple of model input, corresponding ground truth, and lead time
+        model (loaded pytorch model): model to train
+        optimizer (torch.optim): optimizer for training set
+        loss_fn (torch.nn Loss Function): loss function for training set
+        device (torch.device): device index to select
+        rank (int): Rank of device
+        world_size (int): Number of total DDP processes
+
+    Returns:
+        end_img (torch.Tensor): Ground truth end image
+        pred_img (torch.Tensor): Predicted end image
+        all_losses (torch.Tensor): Concatenated per-sample losses from all processes
+    """
+    # Set model to train mode
+    model.train()
+
+    # Extract data
+    start_img, end_img, Dt = data
+    start_img = start_img.to(device, non_blocking=True)
+    Dt = Dt.to(device, non_blocking=True)
+    end_img = end_img.to(device, non_blocking=True)
+
+    # Fixed input and output variable indices
+    in_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
+    out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True)
+
+    # Forward pass
+    pred_img = model(start_img, in_vars, out_vars, Dt)
+
+    # Compute loss
+    loss = loss_fn(pred_img, end_img)
+    per_sample_loss = loss.mean(dim=[1, 2, 3])  # Per-sample loss
+
+    # Backward pass and optimization
+    optimizer.zero_grad(set_to_none=True)
+    loss.mean().backward()
+    optimizer.step()
+
+    # Gather per-sample losses from all processes
+    gathered_losses = [torch.zeros_like(per_sample_loss) for _ in range(world_size)]
+    dist.all_gather(gathered_losses, per_sample_loss)
+
+    # Rank 0 concatenates and saves or returns all losses
+    if rank == 0:
+        all_losses = torch.cat(gathered_losses, dim=0)  # Shape: (total_batch_size,)
+    else:
+        all_losses = None
+
+    return end_img, pred_img, all_losses
+
+
+####################################
+# Evaluating on a Datastep
+####################################
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index f471916c..153c30d6 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -11,9 +11,7 @@
     train_scheduled_loderunner_datastep,
     eval_scheduled_loderunner_datastep,
     train_DDP_loderunner_datastep,
-    train_DDP_scalar_temporal_loderunner_datastep_gri,
     eval_DDP_loderunner_datastep,
-    eval_DDP_scalar_temporal_loderunner_datastep_gri,
 )
 
 
@@ -349,8 +347,7 @@ def train_DDP_scalar_temporal_loderunner_epoch_gri(
     rank: int,
     world_size: int,
 ) -> None:
-    """
-    DDP epoch function for scalar temporal LodeRunner training.
+    """DDP epoch function for scalar temporal LodeRunner training.
 
     Expected dataset output:
         x:      [B, input_dim]
@@ -367,7 +364,6 @@ def train_DDP_scalar_temporal_loderunner_epoch_gri(
     Expected model output:
         pred:   [B, 3]
     """
-
     train_rcrd_filename = train_rcrd_filename.replace(
         "<epochIDX>",
         f"{epochIDX:04d}",
@@ -485,93 +481,6 @@ def train_DDP_scalar_temporal_loderunner_epoch_gri(
                         np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
 
 
-def train_DDP_scalar_temporal_loderunner_epoch(
-    training_data: torch.utils.data.DataLoader,
-    validation_data: torch.utils.data.DataLoader,
-    num_train_batches: int,
-    num_val_batches: int,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    LRsched: torch.optim.lr_scheduler._LRScheduler,
-    epochIDX: int,
-    train_per_val: int,
-    train_rcrd_filename: str,
-    val_rcrd_filename: str,
-    device: torch.device,
-    rank: int,
-    world_size: int,
-) -> None:
-    trainbatch_ID = 0
-    valbatch_ID = 0
-
-    model.train()
-
-    train_rcrd_filename = train_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
-
-    with (
-        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
-    ) as train_rcrd_file:
-        for trainbatch_ID, traindata in enumerate(training_data):
-            if trainbatch_ID >= num_train_batches:
-                break
-
-            truth, pred, train_losses = train_DDP_scalar_temporal_loderunner_datastep_gri(
-                traindata,
-                model,
-                optimizer,
-                loss_fn,
-                device,
-                rank,
-                world_size,
-            )
-
-            LRsched.step()
-
-            if rank == 0:
-                batch_records = np.column_stack(
-                    [
-                        np.full(len(train_losses), epochIDX),
-                        np.full(len(train_losses), trainbatch_ID),
-                        train_losses.cpu().numpy().flatten(),
-                    ]
-                )
-                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
-
-    if epochIDX % train_per_val == 0:
-        print("Validating...", epochIDX)
-
-        val_rcrd_filename = val_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
-        model.eval()
-
-        with (
-            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
-        ) as val_rcrd_file:
-            for valbatch_ID, valdata in enumerate(validation_data):
-                if valbatch_ID >= num_val_batches:
-                    break
-
-
-                truth, pred, val_losses = eval_DDP_scalar_temporal_loderunner_datastep_gri(
-                    data=data,
-                    model=model,
-                    loss_fn=loss_fn,
-                    device=device,
-                    rank=rank,
-                    world_size=world_size,
-                )
-
-                if rank == 0:
-                    batch_records = np.column_stack(
-                        [
-                            np.full(len(val_losses), epochIDX),
-                            np.full(len(val_losses), valbatch_ID),
-                            val_losses.cpu().numpy().flatten(),
-                        ]
-                    )
-                    np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
-
-
 def train_DDP_loderunner_epoch(
     training_data: torch.utils.data.DataLoader,
     validation_data: torch.utils.data.DataLoader,
@@ -588,7 +497,6 @@ def train_DDP_loderunner_epoch(
     device: torch.device,
     rank: int,
     world_size: int,
-    seq: bool = False
 ) -> None:
     """Distributed data-parallel LodeRunner Epoch.
 
@@ -630,155 +538,10 @@ def train_DDP_loderunner_epoch(
             if trainbatch_ID >= num_train_batches:
                 break
 
-            if seq:  #all_losses??
-                gt_seq, pred_seq, train_losses = train_DDP_loderunner_seq_datastep(
-                    data=traindata,
-                    model=model,
-                    optimizer=optimizer,
-                    loss_fn=loss_fn,
-                    device=device,
-                    rank=rank,
-                    world_size=world_size,
-                    scheduled_prob=1.0,   # start with pure teacher forcing
-                )
-            else:
-                # Perform a single training step
-                truth, pred, train_losses = train_DDP_loderunner_seq_channel_datastep(
-                    traindata, model, optimizer, loss_fn, device, rank, world_size
-                )
-
-            # Increment the learning-rate scheduler
-            LRsched.step()
-
-            # Save training record (rank 0 only)
-            if rank == 0:
-                batch_records = np.column_stack(
-                    [
-                        np.full(len(train_losses), epochIDX),
-                        np.full(len(train_losses), trainbatch_ID),
-                        train_losses.cpu().numpy().flatten(),
-                    ]
-                )
-                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
-
-    # Validation loop
-    if epochIDX % train_per_val == 0:
-        print("Validating...", epochIDX)
-        val_rcrd_filename = val_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
-        model.eval()
-        with (
-            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
-        ) as val_rcrd_file:
-            with torch.no_grad():
-                for valbatch_ID, valdata in enumerate(validation_data):
-                    # Stop when number of training batches is reached
-                    if valbatch_ID >= num_val_batches:
-                        break
-
-                    if seq:  #all_losses??
-                        gt_seq, pred_seq, val_losses = eval_DDP_loderunner_seq_datastep(
-                            data=valdata,
-                            model=model,
-                            loss_fn=loss_fn,
-                            device=device,
-                            rank=rank,
-                            world_size=world_size,
-                        )
-                    else:
-                        # Perform a single training step
-                        end_img, pred_img, val_losses = eval_DDP_loderunner_seq_channel_datastep(
-                            valdata, model, loss_fn, device, rank, world_size
-                        )
-
-
-                    # Save validation record (rank 0 only)
-                    if rank == 0:
-                        batch_records = np.column_stack(
-                            [
-                                np.full(len(val_losses), epochIDX),
-                                np.full(len(val_losses), valbatch_ID),
-                                val_losses.cpu().numpy().flatten(),
-                            ]
-                        )
-                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
-
-
-def train_DDP_loderunner_epoch_seq_context(
-    training_data: torch.utils.data.DataLoader,
-    validation_data: torch.utils.data.DataLoader,
-    num_train_batches: int,
-    num_val_batches: int,
-    model: torch.nn.Module,
-    optimizer: torch.optim.Optimizer,
-    loss_fn: torch.nn.Module,
-    LRsched: torch.optim.lr_scheduler._LRScheduler,
-    epochIDX: int,
-    train_per_val: int,
-    train_rcrd_filename: str,
-    val_rcrd_filename: str,
-    device: torch.device,
-    rank: int,
-    world_size: int,
-    seq: bool = False
-) -> None:
-    """Distributed data-parallel LodeRunner Epoch.
-
-    Function to complete a training epoch on the LodeRunner architecture with
-    fixed channels in the input and output. Training and validation information
-    is saved to successive CSV files.
-
-    Args:
-        training_data (torch.utils.data.DataLoader): training dataloader
-        validation_data (torch.utils.data.DataLoader): validation dataloader
-        num_train_batches (int): Number of batches in training epoch
-        num_val_batches (int): Number of batches in validation epoch
-        model (torch.nn.Module): model to train
-        optimizer (torch.optim.Optimizer): optimizer for training set
-        loss_fn (torch.nn.Module): loss function for training set
-        LRsched (torch.optim.lr_scheduler._LRScheduler): Learning-rate scheduler called
-                                                         every training step.
-        epochIDX (int): Index of current training epoch
-        train_per_val (int): Number of Training epochs between each validation
-        train_rcrd_filename (str): Name of CSV file to save training sample stats to
-        val_rcrd_filename (str): Name of CSV file to save validation sample stats to
-        device (torch.device): device index to select
-        rank (int): rank of process
-        world_size (int): number of total processes
-
-    """
-    # Initialize things to save
-    trainbatch_ID = 0
-    valbatch_ID = 0
-
-    # Training loop
-    model.train()
-    train_rcrd_filename = train_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
-    max_train_batches = min(num_train_batches, len(training_data))
-
-    with (
-        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
-    ) as train_rcrd_file:
-        for trainbatch_ID, traindata in enumerate(training_data):
-            # Stop when number of training batches is reached
-            if trainbatch_ID >= max_train_batches:
-                break
-
-            if seq:  #all_losses??
-                gt_seq, pred_seq, train_losses = train_DDP_temporal_loderunner_datastep(
-                    data=traindata,
-                    model=model,
-                    optimizer=optimizer,
-                    loss_fn=loss_fn,
-                    device=device,
-                    rank=rank,
-                    world_size=world_size,
-                    #scheduled_prob=1.0,   # start with pure teacher forcing
-                )
-            else:
-                # Perform a single training step
-                truth, pred, train_losses = train_DDP_loderunner_seq_channel_datastep(
-                    traindata, model, optimizer, loss_fn, device, rank, world_size
-                )
+            # Perform a single training step
+            truth, pred, train_losses = train_DDP_loderunner_datastep(
+                traindata, model, optimizer, loss_fn, device, rank, world_size
+            )
 
             # Increment the learning-rate scheduler
             LRsched.step()
@@ -794,49 +557,6 @@ def train_DDP_loderunner_epoch_seq_context(
                 )
                 np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
 
-    # Validation loop
-    if epochIDX % train_per_val == 0:
-        print("Validating...", epochIDX)
-        val_rcrd_filename = val_rcrd_filename.replace("<epochIDX>", f"{epochIDX:04d}")
-        model.eval()
-        max_val_batches = min(num_val_batches, len(val_data))
-
-        with (open(val_rcrd_filename, "a") if rank == 0 else nullcontext()) as val_rcrd_file:
-            with torch.no_grad():
-                for valbatch_ID, valdata in enumerate(validation_data):
-                    if valbatch_ID >= max_val_batches:
-                        break
-
-                    if seq:
-                        end_img, pred_img, val_losses = eval_DDP_loderunner_seq_context_datastep(
-                            data=valdata,
-                            model=model,
-                            loss_fn=loss_fn,
-                            device=device,
-                            rank=rank,
-                            world_size=world_size,
-                        )
-                    else:
-                        end_img, pred_img, val_losses = eval_DDP_loderunner_seq_channel_datastep(
-                            valdata,
-                            model,
-                            loss_fn,
-                            device,
-                            rank,
-                            world_size,
-                        )
-
-                    if rank == 0:
-                        batch_records = np.column_stack(
-                            [
-                                np.full(len(val_losses), epochIDX),
-                                np.full(len(val_losses), valbatch_ID),
-                                val_losses.cpu().numpy().flatten(),
-                            ]
-                        )
-                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
-
-    '''
     # Validation loop
     if epochIDX % train_per_val == 0:
         print("Validating...", epochIDX)
@@ -851,30 +571,14 @@ def train_DDP_loderunner_epoch_seq_context(
                     if valbatch_ID >= num_val_batches:
                         break
 
-                if seq:  #all_losses??
-                    gt_seq, pred_seq, train_losses = eval_DDP_loderunner_seq_datastep(
-                        data=valdata,
-                        model=model,
-                        loss_fn=loss_fn,
-                        device=device,
-                        rank=rank,
-                        world_size=world_size,
+                    end_img, pred_img, val_losses = eval_DDP_loderunner_datastep(
+                        valdata,
+                        model,
+                        loss_fn,
+                        device,
+                        rank,
+                        world_size,
                     )
-                else:
-                    # Perform a single training step
-                    emd_img, pred_img, val_losses = eval_DDP_loderunner_datastep(
-                        valdata, model, loss_fn, device, rank, world_size
-                    )
-
-
-                    #end_img, pred_img, val_losses = eval_DDP_loderunner_datastep(
-                    #    valdata,
-                    #    model,
-                    #    loss_fn,
-                    #    device,
-                    #    rank,
-                    #    world_size,
-                    #)
 
                     # Save validation record (rank 0 only)
                     if rank == 0:
@@ -886,4 +590,3 @@ def train_DDP_loderunner_epoch_seq_context(
                             ]
                         )
                         np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
-        '''

From 9a4bb9f1d6a52162c14cf43fefb19c61633c1d97 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 31 Jul 2026 12:49:34 -0600
Subject: [PATCH 13/66] fix import

---
 .../KN_loderunner/train_LodeRunner_ddp.py     | 620 +-----------------
 src/yoke/datasets/kilonova_dataset.py         | 311 +++++++++
 src/yoke/models/vit/swin/bomberman.py         | 119 ++++
 3 files changed, 443 insertions(+), 607 deletions(-)
 create mode 100644 src/yoke/datasets/kilonova_dataset.py

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 2cf5bd5e..4bc82982 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -8,28 +8,26 @@
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim.lr_scheduler import LambdaLR
 
-from yoke.models.vit.swin.bomberman import LodeRunner
-from yoke.datasets.lsc_dataset import LSC_rho2rho_temporal_DataSet
-from yoke.utils.training.epoch.loderunner import train_DDP_scalar_temporal_loderunner_epoch_gri
+from yoke.models.vit.swin.bomberman import (
+    LodeRunner,
+    ScalarTemporalConditionedLodeRunner_gri,
+)
+from yoke.datasets.kilonova_dataset import (
+    Kilonova_lc_scalar_context_DataSet_gri,
+    load_or_compute_band_normalization,
+)
+from yoke.utils.training.epoch.loderunner import (
+    train_DDP_scalar_temporal_loderunner_epoch_gri,
+)
 from yoke.utils.restart import continuation_setup
 from yoke.utils.dataload import make_distributed_dataloader
 from yoke.utils.checkpointing import load_model_and_optimizer
 from yoke.utils.checkpointing import save_model_and_optimizer
+from yoke.utils.checkpointing import load_direct_loderunner_checkpoint
+from yoke.utils.parallel import setup_distributed, cleanup_distributed
 from yoke.lr_schedulers import CosineWithWarmupScheduler
 from yoke.helpers import cli
 
-# FIXME remove if restructure
-from torch.utils.data import Dataset, DataLoader, random_split
-import glob
-import random
-
-#MEAN = 24.694652705328807
-#STD = 4.67030961432848
-
-GLOBAL_GMAG_MEAN = 24.694652705328807
-GLOBAL_GMAG_STD  = 4.67030961432848
-EPS = 1e-6
-
 #############################################
 # Inputs
 #############################################
@@ -63,598 +61,6 @@
 )
 
 
-def compute_band_normalization(
-    file_prefix_list,
-    band_keys=("arr_ztfg", "arr_ztfr", "arr_ztfi"),
-    value_col=1,
-    stats_path="kilonova_gri_norm_stats.npz",
-):
-    """
-    Compute global per-band mean/std over the training files only.
-
-    Saves:
-        means: shape [3]
-        stds:  shape [3]
-    """
-
-    sums = np.zeros(len(band_keys), dtype=np.float64)
-    sums_sq = np.zeros(len(band_keys), dtype=np.float64)
-    counts = np.zeros(len(band_keys), dtype=np.float64)
-
-    for fn in file_prefix_list:
-        data = np.load(fn, allow_pickle=True)
-
-        for b, key in enumerate(band_keys):
-            vals = data[key][:, value_col].astype(np.float64)
-
-            finite = np.isfinite(vals)
-            vals = vals[finite]
-
-            sums[b] += vals.sum()
-            sums_sq[b] += np.square(vals).sum()
-            counts[b] += vals.size
-
-        data.close()
-
-    means = sums / counts
-    variances = sums_sq / counts - means**2
-    variances = np.maximum(variances, 1e-12)
-    stds = np.sqrt(variances)
-
-    means = means.astype(np.float32)
-    stds = stds.astype(np.float32)
-
-    np.savez(
-        stats_path,
-        means=means,
-        stds=stds,
-        band_keys=np.array(band_keys),
-        value_col=value_col,
-    )
-
-    print("Saved normalization stats:", stats_path)
-    print("means:", means)
-    print("stds:", stds)
-
-    return means, stds
-
-
-def load_or_compute_band_normalization(
-    stats_path="kilonova_gri_norm_stats.npz",
-    band_keys=("arr_ztfg", "arr_ztfr", "arr_ztfi"),
-    value_col=1,
-):
-    file_prefix_list = sorted(
-        glob.glob(
-            "/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz"
-        )
-    )
-
-    if os.path.exists(stats_path):
-        stats = np.load(stats_path, allow_pickle=True)
-        means = stats["means"].astype(np.float32)
-        stds = stats["stds"].astype(np.float32)
-        stats.close()
-
-        print("Loaded normalization stats:", stats_path)
-        print("means:", means)
-        print("stds:", stds)
-
-        return means, stds
-
-    return compute_band_normalization(
-        file_prefix_list=file_prefix_list,
-        band_keys=band_keys,
-        value_col=value_col,
-        stats_path=stats_path,
-    )
-
-
-class Kilonova_lc_scalar_context_DataSet_gri(Dataset):
-    def __init__(
-        self,
-        N_imgs=0,
-        context_len=5,
-        band_keys=("arr_ztfg", "arr_ztfr", "arr_ztfi"),
-        value_col=1,
-        means=None,
-        stds=None,
-        predicts_delta=True,
-    ):
-        file_prefix_list = sorted(
-            glob.glob(
-                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-                "uniform_dataset_20000/lc_*.npz"
-            )
-        )
-
-        if N_imgs == 0:
-            self.file_prefix_list = file_prefix_list
-        else:
-            self.file_prefix_list = list(
-                np.random.choice(file_prefix_list, N_imgs, replace=False)
-            )
-
-        random.shuffle(self.file_prefix_list)
-
-        self.context_len = context_len
-        self.band_keys = tuple(band_keys)
-        self.value_col = value_col
-        self.n_channels = len(self.band_keys)
-        self.predicts_delta = predicts_delta
-
-        if means is None:
-            raise ValueError(
-                "means must be provided for per-band normalization. "
-                "Expected shape [n_channels], e.g. [g_mean, r_mean, i_mean]."
-            )
-
-        if stds is None:
-            raise ValueError(
-                "stds must be provided for per-band normalization. "
-                "Expected shape [n_channels], e.g. [g_std, r_std, i_std]."
-            )
-
-        self.means = np.asarray(means, dtype=np.float32)
-        self.stds = np.asarray(stds, dtype=np.float32)
-
-        if self.means.shape[0] != self.n_channels:
-            raise ValueError(
-                f"means has length {self.means.shape[0]}, "
-                f"but n_channels={self.n_channels}"
-            )
-
-        if self.stds.shape[0] != self.n_channels:
-            raise ValueError(
-                f"stds has length {self.stds.shape[0]}, "
-                f"but n_channels={self.n_channels}"
-            )
-
-        if np.any(self.stds <= 0):
-            raise ValueError(f"All stds must be positive. Got stds={self.stds}")
-
-        self.samples = []
-
-        for file_idx, fn in enumerate(self.file_prefix_list):
-            data = np.load(fn, allow_pickle=True)
-
-            # Assume all bands share the same time grid.
-            mjd = data[self.band_keys[0]][:, 0]
-            n_times = len(mjd)
-
-            data.close()
-
-            max_start = n_times - context_len - 1
-            for startIDX in range(max_start + 1):
-                self.samples.append((file_idx, startIDX))
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, index):
-        file_idx, startIDX = self.samples[index]
-        fn = self.file_prefix_list[file_idx]
-
-        data = np.load(fn, allow_pickle=True)
-
-        # Time grid from the first band.
-        # Assumes arr_ztfg, arr_ztfr, arr_ztfi have matching MJD columns.
-        mjd = data[self.band_keys[0]][:, 0].astype(np.float32)
-
-        # Stack one scalar value column from each band.
-        # vals shape: [T, 3] for g/r/i.
-        vals = np.stack(
-            [
-                data[key][:, self.value_col].astype(np.float32)
-                for key in self.band_keys
-            ],
-            axis=1,
-        )
-
-        data.close()
-
-        # Per-band normalization.
-        # vals[:, 0] = normalized g
-        # vals[:, 1] = normalized r
-        # vals[:, 2] = normalized i
-        vals = (vals - self.means[None, :]) / (self.stds[None, :] + EPS)
-
-        t0 = mjd.min()
-        t_obs = mjd - t0
-
-        target_idx = startIDX + self.context_len
-
-        # Context values shape: [context_len, 3]
-        context_vals = vals[startIDX:target_idx]
-
-        # Relative observation times shape: [context_len]
-        rel_t = t_obs[startIDX:target_idx] - t_obs[startIDX]
-        rel_t = rel_t.astype(np.float32)
-
-        # Flatten context as:
-        # [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK]
-        context_flat = context_vals.reshape(-1).astype(np.float32)
-
-        # Final input shape:
-        # [(3 * context_len) + context_len]
-        #
-        # For context_len=5:
-        # x.shape == [20]
-        x = np.concatenate([context_flat, rel_t], axis=0)
-        x = torch.tensor(x, dtype=torch.float32)
-
-        if self.predicts_delta:
-            # Predict normalized delta for each band:
-            # [delta_g, delta_r, delta_i]
-            target_vals = vals[target_idx] - vals[target_idx - 1]
-        else:
-            # Predict normalized absolute next value:
-            # [g_next, r_next, i_next]
-            target_vals = vals[target_idx]
-
-        target = torch.tensor(target_vals, dtype=torch.float32)
-
-        Dt = torch.tensor(
-            t_obs[target_idx] - t_obs[target_idx - 1],
-            dtype=torch.float32,
-        )
-
-        return x, target, Dt
-
-
-class Kilonova_lc_scalar_context_DataSet(Dataset):
-    def __init__(
-        self,
-        N_imgs=0,
-        context_len=5,
-    ):
-        file_prefix_list = sorted(
-            glob.glob("/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz")
-        )
-
-        if N_imgs == 0:
-            self.file_prefix_list = file_prefix_list
-        else:
-            self.file_prefix_list = list(np.random.choice(file_prefix_list, N_imgs, replace=False))
-
-        random.shuffle(self.file_prefix_list)
-
-        self.context_len = context_len
-        self.samples = []
-
-        for file_idx, fn in enumerate(self.file_prefix_list):
-            data = np.load(fn, allow_pickle=True)
-            mjd = data["arr_ztfg"][:, 0]
-            n_times = len(mjd)
-            data.close()
-
-            max_start = n_times - context_len - 1
-            for startIDX in range(max_start + 1):
-                self.samples.append((file_idx, startIDX))
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, index):
-        file_idx, startIDX = self.samples[index]
-        fn = self.file_prefix_list[file_idx]
-
-        data = np.load(fn, allow_pickle=True)
-        arr = data["arr_ztfg"]
-
-        mjd = arr[:, 0]
-        g_mag = arr[:, 1].astype(np.float32)
-
-        g_mag = (g_mag - GLOBAL_GMAG_MEAN) / (GLOBAL_GMAG_STD + EPS)
-
-        t0 = mjd.min()
-        t_obs = mjd - t0
-
-        target_idx = startIDX + self.context_len
-
-        mags = g_mag[startIDX:target_idx].astype(np.float32)
-
-        # context_len values; first dt is 0, remaining are relative times
-        rel_t = t_obs[startIDX:target_idx] - t_obs[startIDX]
-        rel_t = rel_t.astype(np.float32)
-
-        x = np.concatenate([mags, rel_t], axis=0)
-        x = torch.tensor(x, dtype=torch.float32)
-
-        delta_mag = g_mag[target_idx] - g_mag[target_idx - 1]
-        target = torch.tensor(delta_mag, dtype=torch.float32)
-
-        Dt = torch.tensor(
-            t_obs[target_idx] - t_obs[target_idx - 1],
-            dtype=torch.float32,
-        )
-
-        # Keep target image-shaped only if your datastep still expects image targets.
-        # Better is to update datastep to use scalar target directly.
-        data.close()
-        return x, target, Dt
-
-
-class ScalarTemporalConditionedLodeRunner_gri(nn.Module):
-    def __init__(
-        self,
-        backbone: nn.Module,
-        context_len: int = 5,
-        n_input_channels: int = 3,
-        n_output_channels: int = 3,
-        image_size=(1120, 400),
-        backbone_channels: int = 8,
-        hidden: int = 64,
-    ):
-        super().__init__()
-
-        self.backbone = backbone
-        self.context_len = context_len
-        self.n_input_channels = n_input_channels
-        self.n_output_channels = n_output_channels
-        self.image_size = image_size
-        self.backbone_channels = backbone_channels
-
-        # Dataset x layout:
-        #   [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK, t0, t1, ..., tK]
-        #
-        # input_dim = context_len * n_input_channels + context_len
-        input_dim = context_len * n_input_channels + context_len
-
-        # Maps scalar temporal context into the 8 pseudo-channels expected by
-        # the pretrained LodeRunner backbone.
-        self.conditioner = nn.Sequential(
-            nn.Linear(input_dim, hidden),
-            nn.GELU(),
-            nn.Linear(hidden, hidden),
-            nn.GELU(),
-            nn.Linear(hidden, backbone_channels),
-        )
-
-        # Maps the 8-channel LodeRunner output back to 3 scalar predictions:
-        #   [delta_g, delta_r, delta_i]
-        self.output_head = nn.Sequential(
-            nn.Linear(backbone_channels, hidden),
-            nn.GELU(),
-            nn.Linear(hidden, n_output_channels),
-        )
-
-    def forward(self, x, in_vars, out_vars, Dt):
-        """
-        x: [B, context_len * n_input_channels + context_len]
-
-        For 3-band, context_len=5:
-            x.shape == [B, 20]
-
-        Returns:
-            pred: [B, 3]
-        """
-        B = x.shape[0]
-        H, W = self.image_size
-
-        channel_vals = self.conditioner(x)  # [B, 8]
-
-        pseudo_img = channel_vals.view(
-            B,
-            self.backbone_channels,
-            1,
-            1,
-        ).expand(
-            B,
-            self.backbone_channels,
-            H,
-            W,
-        )
-
-        backbone_in_vars = torch.arange(self.backbone_channels, device=x.device)
-        backbone_out_vars = torch.arange(self.backbone_channels, device=x.device)
-
-        pred_img = self.backbone(
-            pseudo_img,
-            backbone_in_vars,
-            backbone_out_vars,
-            Dt,
-        )  # [B, 8, H, W]
-
-        # Collapse spatial dimensions to 8 backbone-channel summaries.
-        pred_channel_vals = pred_img.mean(dim=(2, 3))  # [B, 8]
-
-        # Convert 8 backbone channels to 3 output bands.
-        pred = self.output_head(pred_channel_vals)  # [B, 3]
-
-        return pred
-
-
-class ScalarTemporalConditionedLodeRunner(nn.Module):
-    def __init__(
-        self,
-        backbone: nn.Module,
-        context_len: int = 5,
-        image_size=(1120, 400),
-        n_channels: int = 8,
-        hidden: int = 64,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.context_len = context_len
-        self.image_size = image_size
-        self.n_channels = n_channels
-
-        # mags + dts
-        self.conditioner = nn.Sequential(
-            nn.Linear(2 * context_len, hidden),
-            nn.GELU(),
-            nn.Linear(hidden, hidden),
-            nn.GELU(),
-            nn.Linear(hidden, n_channels),
-        )
-
-    def forward(self, x, in_vars, out_vars, Dt):
-        """
-        x: [B, 2 * context_len]
-           first context_len entries are magnitudes
-           second context_len entries are temporal deltas
-        """
-        B = x.shape[0]
-        H, W = self.image_size
-
-        channel_vals = self.conditioner(x)  # [B, 8]
-
-        pseudo_img = channel_vals.view(B, self.n_channels, 1, 1).expand(
-            B,
-            self.n_channels,
-            H,
-            W,
-        )
-
-        pred_img = self.backbone(pseudo_img, in_vars, out_vars, Dt)
-
-        # Convert LodeRunner image output back to scalar delta prediction
-        pred_scalar = pred_img.mean(dim=(1, 2, 3))
-
-        return pred_scalar
-
-
-def load_direct_loderunner_checkpoint(
-    checkpoint_path,
-    model_args,
-    optimizer_kwargs,
-    device,
-):
-    checkpoint_data = torch.load(
-        checkpoint_path,
-        map_location=device,
-        weights_only=False,
-    )
-
-    saved_model_args = checkpoint_data.get("model_args", model_args)
-    context_len = checkpoint_data.get("context_len", 5)
-
-    backbone = LodeRunner(**saved_model_args).to(device)
-
-    model = ScalarTemporalConditionedLodeRunner_gri(
-        backbone=backbone,
-        context_len=context_len,
-        n_input_channels=checkpoint_data.get("n_input_channels", 3),
-        n_output_channels=checkpoint_data.get("n_output_channels", 3),
-        image_size=saved_model_args["image_size"],
-        backbone_channels=checkpoint_data.get("backbone_channels", 8),
-        hidden=checkpoint_data.get("hidden", 64),
-    ).to(device)
-
-    state_dict = checkpoint_data["model_state_dict"]
-
-    # Remove DDP prefix if present
-    if any(k.startswith("module.") for k in state_dict.keys()):
-        state_dict = {
-            k.replace("module.", "", 1): v
-            for k, v in state_dict.items()
-        }
-
-    # Detect checkpoint type
-    is_wrapper_checkpoint = any(
-        k.startswith("backbone.") for k in state_dict.keys()
-    )
-
-    # -------------------------------------------------
-    # OLD plain LodeRunner checkpoint
-    # -------------------------------------------------
-    if not is_wrapper_checkpoint:
-
-        missing_keys, unexpected_keys = model.backbone.load_state_dict(
-            state_dict,
-            strict=False,
-        )
-
-        print("Loaded old LodeRunner checkpoint into model.backbone")
-        print("Missing backbone keys:", missing_keys)
-        print("Unexpected backbone keys:", unexpected_keys)
-
-        # This is NOT a true continuation.
-        # Conditioner is newly initialized.
-        starting_epoch = 0
-
-    # -------------------------------------------------
-    # NEW ScalarTemporalConditionedLodeRunner checkpoint
-    # -------------------------------------------------
-    else:
-
-        model.load_state_dict(state_dict, strict=True)
-
-        print("Loaded ScalarTemporalConditionedLodeRunner checkpoint")
-
-        starting_epoch = checkpoint_data.get("epoch", 0)
-
-    noise_scale = checkpoint_data.get("noise_scale", 0.0)
-    model.backbone.noise_scale = noise_scale
-
-    # Freeze pretrained backbone
-    for p in model.backbone.parameters():
-        p.requires_grad = False
-
-    # Train conditioner
-    for p in model.conditioner.parameters():
-        p.requires_grad = True
-
-    #optimizer = torch.optim.AdamW(
-    #    model.conditioner.parameters(),
-    #    **optimizer_kwargs,
-    #)
-    
-    optimizer = torch.optim.AdamW(
-        list(model.conditioner.parameters()) +
-        list(model.output_head.parameters()),
-        **optimizer_kwargs,
-    )
-
-    # Only restore optimizer for TRUE continuation checkpoints
-    if (
-        is_wrapper_checkpoint
-        and "optimizer_state_dict" in checkpoint_data
-    ):
-
-        optimizer.load_state_dict(
-            checkpoint_data["optimizer_state_dict"]
-        )
-
-        for state in optimizer.state.values():
-            for key, value in state.items():
-                if isinstance(value, torch.Tensor):
-                    state[key] = value.to(device)
-
-    return model, optimizer, starting_epoch
-
-
-def setup_distributed():
-    # ----- 1) Basic setup & environment variables -----
-    # Rely on Slurm variables: SLURM_PROCID, SLURM_NTASKS, SLURM_LOCALID, etc.
-    rank = int(os.environ["SLURM_PROCID"])  # global rank
-    world_size = int(os.environ["SLURM_NTASKS"])  # total number of processes
-    local_rank = int(os.environ["SLURM_LOCALID"])  # local rank (GPU index on this node)
-
-    master_addr = os.environ["MASTER_ADDR"]
-    master_port = os.environ["MASTER_PORT"]
-
-    # ----- 2) Set the current GPU device for this process -----
-    torch.cuda.set_device(local_rank)
-    device = torch.device(f"cuda:{local_rank}")
-
-    # ----- 3) Initialize the process group -----
-    dist.init_process_group(
-        backend="nccl",
-        init_method=f"tcp://{master_addr}:{master_port}",
-        world_size=world_size,
-        rank=rank,
-    )
-
-    return rank, world_size, local_rank, device
-
-
-def cleanup_distributed():
-    # ----- 8) Clean up (optional) -----
-    dist.destroy_process_group()
-
-
 def main(args, rank, world_size, local_rank, device):
     #############################################
     # Process Inputs
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
new file mode 100644
index 00000000..9a28500a
--- /dev/null
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -0,0 +1,311 @@
+"""Relating to kilonova light-curve data.
+
+Functions and classes for torch DataSets which sample kilonova light-curve
+(g/r/i band) data, along with helpers for computing and caching per-band
+normalization statistics.
+
+"""
+
+####################################
+# Packages
+####################################
+import glob
+import os
+import random
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+EPS = 1e-6
+
+
+def compute_band_normalization(
+    file_prefix_list: list[str],
+    band_keys: tuple[str, ...] = ("arr_ztfg", "arr_ztfr", "arr_ztfi"),
+    value_col: int = 1,
+    stats_path: str = "kilonova_gri_norm_stats.npz",
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute global per-band mean/std over the training files only.
+
+    Args:
+        file_prefix_list (list[str]): List of npz files to accumulate stats over.
+        band_keys (tuple[str, ...]): Keys of the bands to normalize.
+        value_col (int): Column index of the value to accumulate per band.
+        stats_path (str): Path to save the computed statistics to.
+
+    Returns:
+        means (np.ndarray): Per-band means, shape [n_bands].
+        stds (np.ndarray): Per-band standard deviations, shape [n_bands].
+    """
+    sums = np.zeros(len(band_keys), dtype=np.float64)
+    sums_sq = np.zeros(len(band_keys), dtype=np.float64)
+    counts = np.zeros(len(band_keys), dtype=np.float64)
+
+    for fn in file_prefix_list:
+        data = np.load(fn, allow_pickle=True)
+
+        for b, key in enumerate(band_keys):
+            vals = data[key][:, value_col].astype(np.float64)
+
+            finite = np.isfinite(vals)
+            vals = vals[finite]
+
+            sums[b] += vals.sum()
+            sums_sq[b] += np.square(vals).sum()
+            counts[b] += vals.size
+
+        data.close()
+
+    means = sums / counts
+    variances = sums_sq / counts - means**2
+    variances = np.maximum(variances, 1e-12)
+    stds = np.sqrt(variances)
+
+    means = means.astype(np.float32)
+    stds = stds.astype(np.float32)
+
+    np.savez(
+        stats_path,
+        means=means,
+        stds=stds,
+        band_keys=np.array(band_keys),
+        value_col=value_col,
+    )
+
+    print("Saved normalization stats:", stats_path)
+    print("means:", means)
+    print("stds:", stds)
+
+    return means, stds
+
+
+def load_or_compute_band_normalization(
+    stats_path: str = "kilonova_gri_norm_stats.npz",
+    band_keys: tuple[str, ...] = ("arr_ztfg", "arr_ztfr", "arr_ztfi"),
+    value_col: int = 1,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Load cached per-band normalization stats or compute them if missing.
+
+    Args:
+        stats_path (str): Path to load/save the statistics.
+        band_keys (tuple[str, ...]): Keys of the bands to normalize.
+        value_col (int): Column index of the value to accumulate per band.
+
+    Returns:
+        means (np.ndarray): Per-band means, shape [n_bands].
+        stds (np.ndarray): Per-band standard deviations, shape [n_bands].
+    """
+    # FIXME: hardcoded scratch path. Should be passed in as an argument so this
+    # library function does not depend on a user-specific filesystem location.
+    file_prefix_list = sorted(
+        glob.glob(
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz"
+        )
+    )
+
+    if os.path.exists(stats_path):
+        stats = np.load(stats_path, allow_pickle=True)
+        means = stats["means"].astype(np.float32)
+        stds = stats["stds"].astype(np.float32)
+        stats.close()
+
+        print("Loaded normalization stats:", stats_path)
+        print("means:", means)
+        print("stds:", stds)
+
+        return means, stds
+
+    return compute_band_normalization(
+        file_prefix_list=file_prefix_list,
+        band_keys=band_keys,
+        value_col=value_col,
+        stats_path=stats_path,
+    )
+
+
+class Kilonova_lc_scalar_context_DataSet_gri(Dataset):
+    """Scalar-context kilonova light-curve dataset for g/r/i bands.
+
+    Each sample provides a flattened window of normalized per-band values plus
+    relative observation times as input, and either the normalized next value
+    or the normalized delta as the target.
+    """
+
+    def __init__(
+        self,
+        N_imgs: int = 0,
+        context_len: int = 5,
+        band_keys: tuple[str, ...] = ("arr_ztfg", "arr_ztfr", "arr_ztfi"),
+        value_col: int = 1,
+        means: np.ndarray = None,
+        stds: np.ndarray = None,
+        predicts_delta: bool = True,
+    ) -> None:
+        """Initialize the dataset and build the sample index.
+
+        Args:
+            N_imgs (int): Number of light-curve files to sample; 0 uses all.
+            context_len (int): Number of context timesteps per sample.
+            band_keys (tuple[str, ...]): Keys of the bands to load.
+            value_col (int): Column index of the value to load per band.
+            means (np.ndarray): Per-band means for normalization, shape [n_bands].
+            stds (np.ndarray): Per-band stds for normalization, shape [n_bands].
+            predicts_delta (bool): If True target is the normalized delta,
+                otherwise the normalized next value.
+        """
+        # FIXME: hardcoded scratch path. Should be passed in as an argument so
+        # this dataset does not depend on a user-specific filesystem location.
+        file_prefix_list = sorted(
+            glob.glob(
+                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+                "uniform_dataset_20000/lc_*.npz"
+            )
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(
+                np.random.choice(file_prefix_list, N_imgs, replace=False)
+            )
+
+        random.shuffle(self.file_prefix_list)
+
+        self.context_len = context_len
+        self.band_keys = tuple(band_keys)
+        self.value_col = value_col
+        self.n_channels = len(self.band_keys)
+        self.predicts_delta = predicts_delta
+
+        if means is None:
+            raise ValueError(
+                "means must be provided for per-band normalization. "
+                "Expected shape [n_channels], e.g. [g_mean, r_mean, i_mean]."
+            )
+
+        if stds is None:
+            raise ValueError(
+                "stds must be provided for per-band normalization. "
+                "Expected shape [n_channels], e.g. [g_std, r_std, i_std]."
+            )
+
+        self.means = np.asarray(means, dtype=np.float32)
+        self.stds = np.asarray(stds, dtype=np.float32)
+
+        if self.means.shape[0] != self.n_channels:
+            raise ValueError(
+                f"means has length {self.means.shape[0]}, "
+                f"but n_channels={self.n_channels}"
+            )
+
+        if self.stds.shape[0] != self.n_channels:
+            raise ValueError(
+                f"stds has length {self.stds.shape[0]}, "
+                f"but n_channels={self.n_channels}"
+            )
+
+        if np.any(self.stds <= 0):
+            raise ValueError(f"All stds must be positive. Got stds={self.stds}")
+
+        self.samples = []
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+
+            # Assume all bands share the same time grid.
+            mjd = data[self.band_keys[0]][:, 0]
+            n_times = len(mjd)
+
+            data.close()
+
+            max_start = n_times - context_len - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((file_idx, startIDX))
+
+    def __len__(self) -> int:
+        """Return the number of samples in the dataset."""
+        return len(self.samples)
+
+    def __getitem__(
+        self, index: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return the (input, target, Dt) triple for a given sample index.
+
+        Args:
+            index (int): Sample index.
+
+        Returns:
+            x (torch.Tensor): Flattened context values and relative times.
+            target (torch.Tensor): Normalized next value or delta per band.
+            Dt (torch.Tensor): Time delta to the target step.
+        """
+        file_idx, startIDX = self.samples[index]
+        fn = self.file_prefix_list[file_idx]
+
+        data = np.load(fn, allow_pickle=True)
+
+        # Time grid from the first band.
+        # Assumes arr_ztfg, arr_ztfr, arr_ztfi have matching MJD columns.
+        mjd = data[self.band_keys[0]][:, 0].astype(np.float32)
+
+        # Stack one scalar value column from each band.
+        # vals shape: [T, 3] for g/r/i.
+        vals = np.stack(
+            [
+                data[key][:, self.value_col].astype(np.float32)
+                for key in self.band_keys
+            ],
+            axis=1,
+        )
+
+        data.close()
+
+        # Per-band normalization.
+        # vals[:, 0] = normalized g
+        # vals[:, 1] = normalized r
+        # vals[:, 2] = normalized i
+        vals = (vals - self.means[None, :]) / (self.stds[None, :] + EPS)
+
+        t0 = mjd.min()
+        t_obs = mjd - t0
+
+        target_idx = startIDX + self.context_len
+
+        # Context values shape: [context_len, 3]
+        context_vals = vals[startIDX:target_idx]
+
+        # Relative observation times shape: [context_len]
+        rel_t = t_obs[startIDX:target_idx] - t_obs[startIDX]
+        rel_t = rel_t.astype(np.float32)
+
+        # Flatten context as:
+        # [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK]
+        context_flat = context_vals.reshape(-1).astype(np.float32)
+
+        # Final input shape:
+        # [(3 * context_len) + context_len]
+        #
+        # For context_len=5:
+        # x.shape == [20]
+        x = np.concatenate([context_flat, rel_t], axis=0)
+        x = torch.tensor(x, dtype=torch.float32)
+
+        if self.predicts_delta:
+            # Predict normalized delta for each band:
+            # [delta_g, delta_r, delta_i]
+            target_vals = vals[target_idx] - vals[target_idx - 1]
+        else:
+            # Predict normalized absolute next value:
+            # [g_next, r_next, i_next]
+            target_vals = vals[target_idx]
+
+        target = torch.tensor(target_vals, dtype=torch.float32)
+
+        Dt = torch.tensor(
+            t_obs[target_idx] - t_obs[target_idx - 1],
+            dtype=torch.float32,
+        )
+
+        return x, target, Dt
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index e3c6af63..507cefa4 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -332,6 +332,125 @@ def validation_step(self, batch: torch.Tensor, batch_idx: int) -> None:
             self.log("val_loss", batch_loss, sync_dist=True)
 
 
+class ScalarTemporalConditionedLodeRunner_gri(nn.Module):
+    """Scalar-temporal wrapper around a pretrained LodeRunner backbone.
+
+    Maps a scalar temporal context (flattened per-band values plus relative
+    observation times) into the pseudo-channel image expected by a pretrained
+    LodeRunner backbone, then collapses the backbone's spatial output back to a
+    small number of scalar band predictions. Used for the kilonova light-curve
+    (g/r/i) forecasting task.
+
+    Args:
+        backbone (nn.Module): Pretrained LodeRunner backbone.
+        context_len (int): Number of context timesteps.
+        n_input_channels (int): Number of input bands (e.g. 3 for g/r/i).
+        n_output_channels (int): Number of predicted bands.
+        image_size (tuple): Spatial size (H, W) fed to the backbone.
+        backbone_channels (int): Number of pseudo-channels the backbone expects.
+        hidden (int): Hidden width of the conditioner/output-head MLPs.
+    """
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        context_len: int = 5,
+        n_input_channels: int = 3,
+        n_output_channels: int = 3,
+        image_size: tuple[int, int] = (1120, 400),
+        backbone_channels: int = 8,
+        hidden: int = 64,
+    ) -> None:
+        """Initialize conditioner and output-head around the backbone."""
+        super().__init__()
+
+        self.backbone = backbone
+        self.context_len = context_len
+        self.n_input_channels = n_input_channels
+        self.n_output_channels = n_output_channels
+        self.image_size = image_size
+        self.backbone_channels = backbone_channels
+
+        # Dataset x layout:
+        #   [g0, r0, i0, g1, r1, i1, ..., gK, rK, iK, t0, t1, ..., tK]
+        #
+        # input_dim = context_len * n_input_channels + context_len
+        input_dim = context_len * n_input_channels + context_len
+
+        # Maps scalar temporal context into the 8 pseudo-channels expected by
+        # the pretrained LodeRunner backbone.
+        self.conditioner = nn.Sequential(
+            nn.Linear(input_dim, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, backbone_channels),
+        )
+
+        # Maps the 8-channel LodeRunner output back to 3 scalar predictions:
+        #   [delta_g, delta_r, delta_i]
+        self.output_head = nn.Sequential(
+            nn.Linear(backbone_channels, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, n_output_channels),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        in_vars: torch.Tensor,
+        out_vars: torch.Tensor,
+        Dt: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            x (torch.Tensor): Scalar temporal context of shape
+                [B, context_len * n_input_channels + context_len]. For 3-band,
+                context_len=5, x.shape == [B, 20].
+            in_vars (torch.Tensor): Kept for LodeRunner API compatibility.
+            out_vars (torch.Tensor): Kept for LodeRunner API compatibility.
+            Dt (torch.Tensor): Lead-time tensor passed to the backbone.
+
+        Returns:
+            pred (torch.Tensor): Predictions of shape [B, n_output_channels].
+        """
+        B = x.shape[0]
+        H, W = self.image_size
+
+        channel_vals = self.conditioner(x)  # [B, 8]
+
+        pseudo_img = channel_vals.view(
+            B,
+            self.backbone_channels,
+            1,
+            1,
+        ).expand(
+            B,
+            self.backbone_channels,
+            H,
+            W,
+        )
+
+        backbone_in_vars = torch.arange(self.backbone_channels, device=x.device)
+        backbone_out_vars = torch.arange(self.backbone_channels, device=x.device)
+
+        pred_img = self.backbone(
+            pseudo_img,
+            backbone_in_vars,
+            backbone_out_vars,
+            Dt,
+        )  # [B, 8, H, W]
+
+        # Collapse spatial dimensions to 8 backbone-channel summaries.
+        pred_channel_vals = pred_img.mean(dim=(2, 3))  # [B, 8]
+
+        # Convert 8 backbone channels to 3 output bands.
+        pred = self.output_head(pred_channel_vals)  # [B, 3]
+
+        return pred
+
+
 if __name__ == "__main__":
     from yoke.utils.parameters import count_torch_params
 

From 801d51439c1d09ab33b1c65216ebd44b1204b178 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 11 Aug 2026 14:03:59 -0600
Subject: [PATCH 14/66] 9 band Rubin+ZTF first attempt

---
 .../harnesses/KN_loderunner/infer_9band.py    | 452 ++++++++++++++++++
 .../plot_pred_diagnostics_9band.py            | 353 ++++++++++++++
 .../KN_loderunner/train_LodeRunner_ddp.py     |  41 +-
 src/yoke/datasets/kilonova_dataset.py         | 246 ++++++++++
 src/yoke/models/vit/swin/bomberman.py         | 118 +++++
 .../utils/training/datastep/loderunner.py     | 103 ++++
 src/yoke/utils/training/epoch/loderunner.py   | 142 ++++++
 7 files changed, 1435 insertions(+), 20 deletions(-)
 create mode 100644 applications/harnesses/KN_loderunner/infer_9band.py
 create mode 100644 applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
new file mode 100644
index 00000000..7a0962c6
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -0,0 +1,452 @@
+"""Forecast all 9 bands into the future with the scalar temporal LodeRunner.
+
+This is the intended production use of the 9-band model: given a light curve of
+sparse, irregular, multi-band observations (real or simulated), take the most
+recent ``context_len`` observations and predict the value in every band at a
+grid of future lead times, measured from the last observation.
+
+For each requested lead time Dt, the model consumes the same fixed context
+window and emits a prediction for all 9 bands at that lead time. This directly
+answers "what will each observatory see next?" without assuming any band is
+observed at that future time.
+
+Predictions are denormalized back to magnitudes using the per-band normalization
+statistics. Results are plotted over the observed data and written to npz/csv.
+"""
+
+import argparse
+import csv
+import glob
+import json
+import os
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from yoke.models.vit.swin.bomberman import (
+    LodeRunner,
+    ScalarTemporalConditionedLodeRunner_9band,
+)
+from yoke.datasets.kilonova_dataset import (
+    EPS,
+    NINE_BAND_KEYS,
+    load_or_compute_band_normalization,
+)
+
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+
+BAND_KEYS = NINE_BAND_KEYS
+BAND_NAMES = ("ztfg", "ztfr", "ztfi", "u", "g", "r", "i", "z", "y")
+BAND_COLORS = (
+    "#2A9D8F",  # ztfg
+    "#E63946",  # ztfr
+    "#F4A261",  # ztfi
+    "#457B9D",  # u
+    "#1B9E77",  # g
+    "#D62828",  # r
+    "#E9C46A",  # i
+    "#8338EC",  # z
+    "#264653",  # y
+)
+VALUE_COL = 1
+N_BANDS = len(BAND_KEYS)
+
+
+def study_tag(study):
+    return f"{int(study):03d}"
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Forecast all 9 bands into the future from a light curve's most "
+            "recent observations."
+        )
+    )
+
+    parser.add_argument("--study", type=int, default=24)
+    parser.add_argument("--epoch", type=int, default=500)
+    parser.add_argument("--ckpt", type=str, default=None)
+
+    parser.add_argument(
+        "--data_glob",
+        type=str,
+        default=None,
+        help="Glob for light-curve npz files to forecast. If omitted, uses the "
+        "training data glob path.",
+    )
+    parser.add_argument(
+        "--n_curves",
+        type=int,
+        default=5,
+        help="Number of light-curve files to forecast.",
+    )
+
+    parser.add_argument(
+        "--horizon",
+        type=float,
+        default=5.0,
+        help="Forecast horizon in the same time units as the data (days), "
+        "measured from the last observation.",
+    )
+    parser.add_argument(
+        "--n_lead_times",
+        type=int,
+        default=25,
+        help="Number of lead times sampled between 0 and --horizon.",
+    )
+
+    parser.add_argument("--outdir", type=str, default=None)
+    parser.add_argument(
+        "--norm_stats_path",
+        type=str,
+        default="kilonova_9band_norm_stats.npz",
+    )
+
+    return parser.parse_args()
+
+
+def resolve_paths(args):
+    tag = study_tag(args.study)
+
+    if args.ckpt is None:
+        args.ckpt = (
+            f"runs/study_{tag}/study{tag}_modelState_epoch{args.epoch:04d}.pth"
+        )
+
+    if args.outdir is None:
+        args.outdir = f"runs/study_{tag}/forecast_9band"
+
+    if args.data_glob is None:
+        args.data_glob = (
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+            "uniform_dataset_20000/lc_*.npz"
+        )
+
+    return tag
+
+
+def strip_ddp_prefix(state_dict):
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        return {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+    return state_dict
+
+
+def load_9band_model(ckpt_path, device):
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+
+    model_args = ckpt["model_args"]
+    context_len = ckpt.get("context_len", 5)
+    n_bands = ckpt.get("n_bands", N_BANDS)
+    backbone_channels = ckpt.get("backbone_channels", 8)
+    hidden = ckpt.get("hidden", 64)
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("model_class:", ckpt.get("model_class", "unknown"))
+    print("target_type:", ckpt.get("target_type", "unknown"))
+    print("context_len:", context_len)
+    print("n_bands:", n_bands)
+
+    backbone = LodeRunner(**model_args).to(device)
+    backbone.noise_scale = noise_scale
+
+    model = ScalarTemporalConditionedLodeRunner_9band(
+        backbone=backbone,
+        context_len=context_len,
+        n_bands=n_bands,
+        image_size=model_args["image_size"],
+        backbone_channels=backbone_channels,
+        hidden=hidden,
+    ).to(device)
+
+    state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.eval()
+
+    return model, context_len, n_bands
+
+
+def load_event_stream(fn, means, stds):
+    """Load one npz file into a merged, time-sorted, normalized event stream.
+
+    Mirrors the dataset's stream construction so inference matches training.
+
+    Returns:
+        times (np.ndarray): Relative observation times [N].
+        values_norm (np.ndarray): Normalized values [N].
+        bands (np.ndarray): Band index per event [N].
+        raw (dict): Per-band raw (mjd, mag) arrays for plotting the observations.
+        t0 (float): The earliest MJD, used to align forecast times.
+    """
+    data = np.load(fn, allow_pickle=True)
+
+    times = []
+    values = []
+    bands = []
+    raw = {}
+
+    for band_idx, key in enumerate(BAND_KEYS):
+        if key not in data.files:
+            continue
+
+        arr = data[key]
+        if arr.size == 0:
+            continue
+
+        t = arr[:, 0].astype(np.float32)
+        v = arr[:, VALUE_COL].astype(np.float32)
+
+        times.append(t)
+        values.append(v)
+        bands.append(np.full(arr.shape[0], band_idx, dtype=np.int64))
+        raw[band_idx] = (t, v)
+
+    data.close()
+
+    if not times:
+        return None
+
+    times = np.concatenate(times)
+    values = np.concatenate(values)
+    bands = np.concatenate(bands)
+
+    order = np.argsort(times, kind="stable")
+    times = times[order]
+    values = values[order]
+    bands = bands[order]
+
+    t0 = float(times.min())
+    times = times - t0
+
+    values_norm = (values - means[bands]) / (stds[bands] + EPS)
+
+    return times, values_norm.astype(np.float32), bands, raw, t0
+
+
+def build_context_input(ctx_t, ctx_v, ctx_b, n_bands, device):
+    """Build the flattened per-event context input for the model.
+
+    Layout per event: [value, rel_t, one_hot_band(n_bands)].
+    """
+    context_len = len(ctx_t)
+
+    rel_t = (ctx_t - ctx_t[0]).astype(np.float32)
+
+    band_onehot = np.zeros((context_len, n_bands), dtype=np.float32)
+    band_onehot[np.arange(context_len), ctx_b] = 1.0
+
+    per_event = np.concatenate(
+        [ctx_v[:, None], rel_t[:, None], band_onehot],
+        axis=1,
+    )
+
+    x = torch.tensor(
+        per_event.reshape(-1),
+        dtype=torch.float32,
+        device=device,
+    ).unsqueeze(0)
+
+    return x
+
+
+def forecast_curve(
+    stream, model, device, context_len, n_bands, means, stds, lead_times
+):
+    """Forecast all bands at a grid of future lead times from the last context.
+
+    Returns a [n_lead_times, n_bands] array of denormalized (magnitude)
+    predictions and the absolute forecast times (in the last-observation frame).
+    """
+    times, values_norm, bands, _, _ = stream
+
+    # Most recent context_len events.
+    ctx_t = times[-context_len:]
+    ctx_v = values_norm[-context_len:]
+    ctx_b = bands[-context_len:]
+
+    x = build_context_input(ctx_t, ctx_v, ctx_b, n_bands, device)
+
+    last_t = float(times[-1])
+
+    preds_norm = np.zeros((len(lead_times), n_bands), dtype=np.float32)
+
+    with torch.no_grad():
+        for k, dt in enumerate(lead_times):
+            Dt = torch.tensor([dt], dtype=torch.float32, device=device)
+            pred = model(x, in_vars=None, out_vars=None, Dt=Dt)
+            preds_norm[k] = pred.reshape(n_bands).detach().cpu().numpy()
+
+    # Denormalize per band back to magnitudes.
+    preds_mag = preds_norm * (stds[None, :] + EPS) + means[None, :]
+
+    forecast_times = last_t + np.asarray(lead_times, dtype=np.float32)
+
+    return preds_mag, forecast_times, last_t
+
+
+def plot_forecast(stream, preds_mag, forecast_times, last_t, title, outpath):
+    _, _, _, raw, t0 = stream
+
+    plt.figure(figsize=(9, 6))
+
+    for band_idx in range(N_BANDS):
+        color = BAND_COLORS[band_idx]
+        name = BAND_NAMES[band_idx]
+
+        # Observed points for this band, aligned to the same relative-time
+        # frame as the stream and forecast (raw MJDs shifted by global t0).
+        if band_idx in raw:
+            t_obs, v_obs = raw[band_idx]
+            plt.scatter(
+                t_obs - t0,
+                v_obs,
+                s=18,
+                color=color,
+                alpha=0.6,
+                label=f"{name} obs",
+            )
+
+        # Forecast curve for this band.
+        plt.plot(
+            forecast_times,
+            preds_mag[:, band_idx],
+            linestyle="--",
+            color=color,
+            linewidth=1.6,
+            label=f"{name} forecast",
+        )
+
+    plt.axvline(
+        last_t,
+        color="k",
+        linewidth=1,
+        linestyle=":",
+        alpha=0.7,
+        label="last observation",
+    )
+
+    plt.gca().invert_yaxis()
+    plt.xlabel("Relative time (days)")
+    plt.ylabel("Magnitude")
+    plt.title(title)
+    plt.legend(fontsize=6, ncol=3, loc="best")
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close()
+
+
+def save_forecast_csv(preds_mag, forecast_times, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["forecast_time"] + [f"mag_{n}" for n in BAND_NAMES])
+        for k in range(len(forecast_times)):
+            writer.writerow(
+                [forecast_times[k]] + [preds_mag[k, b] for b in range(N_BANDS)]
+            )
+
+
+def main():
+    args = get_args()
+    run_id = resolve_paths(args)
+
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len, n_bands = load_9band_model(args.ckpt, device)
+
+    means, stds = load_or_compute_band_normalization(
+        stats_path=args.norm_stats_path,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+    )
+    means = np.asarray(means, dtype=np.float32)
+    stds = np.asarray(stds, dtype=np.float32)
+
+    files = sorted(glob.glob(args.data_glob))
+    if not files:
+        raise RuntimeError(f"No files matched data_glob: {args.data_glob}")
+
+    files = files[: args.n_curves]
+    print(f"Forecasting {len(files)} light curves.")
+
+    lead_times = np.linspace(0.0, args.horizon, args.n_lead_times)
+
+    for i, fn in enumerate(files):
+        stream = load_event_stream(fn, means, stds)
+
+        if stream is None:
+            print(f"Skipping {fn}: no observations.")
+            continue
+
+        times = stream[0]
+        if len(times) < context_len:
+            print(
+                f"Skipping {fn}: only {len(times)} events, "
+                f"need at least context_len={context_len}."
+            )
+            continue
+
+        preds_mag, forecast_times, last_t = forecast_curve(
+            stream=stream,
+            model=model,
+            device=device,
+            context_len=context_len,
+            n_bands=n_bands,
+            means=means,
+            stds=stds,
+            lead_times=lead_times,
+        )
+
+        base = os.path.splitext(os.path.basename(fn))[0]
+
+        png_path = os.path.join(
+            args.outdir, f"study{run_id}_forecast_{base}.png"
+        )
+        csv_path = os.path.join(
+            args.outdir, f"study{run_id}_forecast_{base}.csv"
+        )
+        npz_path = os.path.join(
+            args.outdir, f"study{run_id}_forecast_{base}.npz"
+        )
+
+        plot_forecast(
+            stream=stream,
+            preds_mag=preds_mag,
+            forecast_times=forecast_times,
+            last_t=last_t,
+            title=f"9-band forecast: {base}",
+            outpath=png_path,
+        )
+        save_forecast_csv(preds_mag, forecast_times, csv_path)
+        np.savez(
+            npz_path,
+            forecast_times=forecast_times,
+            preds_mag=preds_mag,
+            band_names=np.array(BAND_NAMES),
+            last_observation_time=last_t,
+        )
+
+        print(f"[{i + 1}/{len(files)}] saved forecast for {base}")
+
+    print("Done. Output directory:", args.outdir)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
new file mode 100644
index 00000000..1e0e0690
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -0,0 +1,353 @@
+"""Next-event prediction diagnostics for the 9-band scalar temporal LodeRunner.
+
+The 9-band model is trained on a merged event stream: each sample is a window of
+``context_len`` consecutive observations (across all bands) plus a lead time, and
+the model predicts the value in every band at that lead time. Each training
+target only observes one band, so diagnostics here compare the model's prediction
+for the observed target band against the truth, aggregated per band.
+
+Unlike the g/r/i diagnostics this script does NOT roll out autoregressively.
+Autoregression is ill-defined for a mixed-band event stream (each future event
+belongs to a single band), so we evaluate the direct one-step-ahead prediction
+the model is actually trained on.
+"""
+
+import argparse
+import csv
+import os
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from yoke.models.vit.swin.bomberman import (
+    LodeRunner,
+    ScalarTemporalConditionedLodeRunner_9band,
+)
+from yoke.datasets.kilonova_dataset import (
+    NINE_BAND_KEYS,
+    Kilonova_lc_scalar_context_DataSet_9band,
+    load_or_compute_band_normalization,
+)
+
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+
+BAND_KEYS = NINE_BAND_KEYS
+BAND_NAMES = ("ztfg", "ztfr", "ztfi", "u", "g", "r", "i", "z", "y")
+VALUE_COL = 1
+N_BANDS = len(BAND_KEYS)
+
+
+def study_tag(study):
+    return f"{int(study):03d}"
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Next-event prediction diagnostics for scalar temporal LodeRunner "
+            "9-band runs."
+        )
+    )
+
+    parser.add_argument("--study", type=int, default=24)
+    parser.add_argument("--epoch", type=int, default=500)
+    parser.add_argument("--ckpt", type=str, default=None)
+
+    parser.add_argument("--N_imgs", type=int, default=50)
+
+    parser.add_argument("--outdir", type=str, default=None)
+    parser.add_argument(
+        "--norm_stats_path",
+        type=str,
+        default="kilonova_9band_norm_stats.npz",
+    )
+
+    return parser.parse_args()
+
+
+def resolve_paths(args):
+    tag = study_tag(args.study)
+
+    if args.ckpt is None:
+        args.ckpt = (
+            f"runs/study_{tag}/study{tag}_modelState_epoch{args.epoch:04d}.pth"
+        )
+
+    if args.outdir is None:
+        args.outdir = f"runs/study_{tag}/next_event_diagnostics_9band"
+
+    return tag
+
+
+def strip_ddp_prefix(state_dict):
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        return {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+    return state_dict
+
+
+def load_9band_model(ckpt_path, device):
+    ckpt = torch.load(
+        ckpt_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    model_args = ckpt["model_args"]
+    context_len = ckpt.get("context_len", 5)
+
+    n_bands = ckpt.get("n_bands", N_BANDS)
+    backbone_channels = ckpt.get("backbone_channels", 8)
+    hidden = ckpt.get("hidden", 64)
+    noise_scale = ckpt.get("noise_scale", 0.0)
+
+    print("Loaded checkpoint:", ckpt_path)
+    print("model_class:", ckpt.get("model_class", "unknown"))
+    print("backbone_class:", ckpt.get("backbone_class", "LodeRunner"))
+    print("target_type:", ckpt.get("target_type", "unknown"))
+    print("context_len:", context_len)
+    print("n_bands:", n_bands)
+    print("band_keys:", ckpt.get("band_keys", list(BAND_KEYS)))
+    print("backbone_channels:", backbone_channels)
+    print("hidden:", hidden)
+
+    backbone = LodeRunner(**model_args).to(device)
+    backbone.noise_scale = noise_scale
+
+    model = ScalarTemporalConditionedLodeRunner_9band(
+        backbone=backbone,
+        context_len=context_len,
+        n_bands=n_bands,
+        image_size=model_args["image_size"],
+        backbone_channels=backbone_channels,
+        hidden=hidden,
+    ).to(device)
+
+    state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=True)
+
+    print("Loaded ScalarTemporalConditionedLodeRunner_9band checkpoint")
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
+
+    model.eval()
+
+    return model, context_len, n_bands
+
+
+def make_eval_dataset(args, context_len):
+    band_means, band_stds = load_or_compute_band_normalization(
+        stats_path=args.norm_stats_path,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+    )
+
+    print("Using band normalization:")
+    print("band_means:", band_means)
+    print("band_stds:", band_stds)
+
+    dataset = Kilonova_lc_scalar_context_DataSet_9band(
+        N_imgs=args.N_imgs,
+        context_len=context_len,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+        means=band_means,
+        stds=band_stds,
+    )
+
+    return dataset
+
+
+def collect_next_event_predictions(dataset, model, device, n_bands):
+    """Run the direct next-event prediction over every sample.
+
+    Returns a dict of per-band arrays of (pred, truth) for the observed band of
+    each sample, plus the residuals.
+    """
+    preds_by_band = [[] for _ in range(n_bands)]
+    truths_by_band = [[] for _ in range(n_bands)]
+
+    with torch.no_grad():
+        for idx in range(len(dataset)):
+            x, target, mask, Dt = dataset[idx]
+
+            x = torch.as_tensor(x, dtype=torch.float32, device=device).unsqueeze(0)
+            target = torch.as_tensor(target, dtype=torch.float32)
+            mask = torch.as_tensor(mask, dtype=torch.float32)
+            Dt = torch.as_tensor(Dt, dtype=torch.float32, device=device)
+
+            if Dt.ndim == 0:
+                Dt = Dt.unsqueeze(0)
+
+            pred = model(x, in_vars=None, out_vars=None, Dt=Dt)
+            pred = pred.reshape(n_bands).detach().cpu()
+
+            band_idx = int(torch.argmax(mask).item())
+
+            preds_by_band[band_idx].append(float(pred[band_idx]))
+            truths_by_band[band_idx].append(float(target[band_idx]))
+
+    results = []
+    for band_idx in range(n_bands):
+        p = np.asarray(preds_by_band[band_idx], dtype=np.float32)
+        t = np.asarray(truths_by_band[band_idx], dtype=np.float32)
+        r = p - t
+
+        if len(r) > 0:
+            mse = float(np.mean(r**2))
+        else:
+            mse = np.nan
+
+        results.append(
+            {
+                "band_idx": band_idx,
+                "band_name": BAND_NAMES[band_idx],
+                "pred": p,
+                "truth": t,
+                "residual": r,
+                "n": len(r),
+                "mse": mse,
+            }
+        )
+
+    return results
+
+
+def plot_pred_vs_truth(results, outpath):
+    fig, axes = plt.subplots(3, 3, figsize=(12, 12))
+    axes = axes.flatten()
+
+    for band_idx, res in enumerate(results):
+        ax = axes[band_idx]
+
+        if res["n"] == 0:
+            ax.set_title(f"{res['band_name']} (no samples)")
+            ax.axis("off")
+            continue
+
+        ax.scatter(res["truth"], res["pred"], s=8, alpha=0.5)
+
+        lo = float(min(res["truth"].min(), res["pred"].min()))
+        hi = float(max(res["truth"].max(), res["pred"].max()))
+        ax.plot([lo, hi], [lo, hi], linestyle="--", linewidth=1, color="k")
+
+        ax.set_xlabel("truth (norm)")
+        ax.set_ylabel("pred (norm)")
+        ax.set_title(f"{res['band_name']}  (n={res['n']}, MSE={res['mse']:.3g})")
+
+    fig.suptitle("Next-event prediction vs truth (normalized), per band", y=1.01)
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_residual_histograms(results, outpath):
+    fig, axes = plt.subplots(3, 3, figsize=(12, 12))
+    axes = axes.flatten()
+
+    for band_idx, res in enumerate(results):
+        ax = axes[band_idx]
+
+        if res["n"] == 0:
+            ax.set_title(f"{res['band_name']} (no samples)")
+            ax.axis("off")
+            continue
+
+        ax.hist(res["residual"], bins=min(30, max(1, res["n"])))
+        ax.axvline(0.0, linestyle="--", linewidth=1, color="k")
+        ax.set_xlabel("pred - truth (norm)")
+        ax.set_ylabel("count")
+        ax.set_title(f"{res['band_name']}  (n={res['n']})")
+
+    fig.suptitle("Next-event residual distributions, per band", y=1.01)
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_band_mse_bar(results, outpath):
+    names = [res["band_name"] for res in results]
+    mses = [res["mse"] if np.isfinite(res["mse"]) else 0.0 for res in results]
+
+    plt.figure(figsize=(9, 5))
+    plt.bar(names, mses)
+    plt.ylabel("Next-event MSE (normalized)")
+    plt.xlabel("Band")
+    plt.title("Per-band next-event MSE")
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200)
+    plt.close()
+
+
+def save_band_mse_csv(results, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["band_idx", "band_name", "n_samples", "mse"])
+        for res in results:
+            writer.writerow(
+                [res["band_idx"], res["band_name"], res["n"], res["mse"]]
+            )
+
+
+def main():
+    args = get_args()
+    run_id = resolve_paths(args)
+
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    model, context_len, n_bands = load_9band_model(args.ckpt, device)
+
+    eval_dataset = make_eval_dataset(args=args, context_len=context_len)
+
+    print("Dataset length:", len(eval_dataset))
+
+    if len(eval_dataset) == 0:
+        raise RuntimeError("Empty eval dataset. Check data path and N_imgs.")
+
+    results = collect_next_event_predictions(
+        dataset=eval_dataset,
+        model=model,
+        device=device,
+        n_bands=n_bands,
+    )
+
+    pred_truth_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_next_event_pred_vs_truth.png"
+    )
+    resid_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_next_event_residual_hist.png"
+    )
+    mse_bar_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_next_event_band_mse.png"
+    )
+    csv_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_next_event_band_mse.csv"
+    )
+
+    plot_pred_vs_truth(results, pred_truth_path)
+    plot_residual_histograms(results, resid_path)
+    plot_band_mse_bar(results, mse_bar_path)
+    save_band_mse_csv(results, csv_path)
+
+    print("Saved:")
+    print(" ", pred_truth_path)
+    print(" ", resid_path)
+    print(" ", mse_bar_path)
+    print(" ", csv_path)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 4bc82982..a0c9a869 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -10,14 +10,15 @@
 
 from yoke.models.vit.swin.bomberman import (
     LodeRunner,
-    ScalarTemporalConditionedLodeRunner_gri,
+    ScalarTemporalConditionedLodeRunner_9band,
 )
 from yoke.datasets.kilonova_dataset import (
-    Kilonova_lc_scalar_context_DataSet_gri,
+    Kilonova_lc_scalar_context_DataSet_9band,
+    NINE_BAND_KEYS,
     load_or_compute_band_normalization,
 )
 from yoke.utils.training.epoch.loderunner import (
-    train_DDP_scalar_temporal_loderunner_epoch_gri,
+    train_DDP_scalar_temporal_loderunner_epoch_9band,
 )
 from yoke.utils.restart import continuation_setup
 from yoke.utils.dataload import make_distributed_dataloader
@@ -140,6 +141,11 @@ def main(args, rank, world_size, local_rank, device):
     CONTEXT_LEN = 5 #3
     HIDDEN_CHANNELS = 64
 
+    # Nine-band merged event-stream setup (3 ZTF + 6 Rubin/LSST bands).
+    BAND_KEYS = NINE_BAND_KEYS
+    VALUE_COL = 1
+    N_BANDS = len(BAND_KEYS)
+
     optimizer_kwargs = {
         "lr": 1e-4,# 1e-4, #1e-5
         "betas": (0.9, 0.999),
@@ -208,11 +214,10 @@ def main(args, rank, world_size, local_rank, device):
 
             backbone = model
 
-            model = ScalarTemporalConditionedLodeRunner_gri(
+            model = ScalarTemporalConditionedLodeRunner_9band(
                 backbone=backbone,
                 context_len=CONTEXT_LEN,
-                n_input_channels=3,
-                n_output_channels=3,
+                n_bands=N_BANDS,
                 image_size=model_args["image_size"],
                 backbone_channels=8,
                 hidden=HIDDEN_CHANNELS,
@@ -316,11 +321,7 @@ def main(args, rank, world_size, local_rank, device):
     )
     '''
 
-    BAND_KEYS = ("arr_ztfg", "arr_ztfr", "arr_ztfi")
-    VALUE_COL = 1
-    N_BANDS = len(BAND_KEYS)
-
-    norm_stats_path = "kilonova_gri_norm_stats.npz"
+    norm_stats_path = "kilonova_9band_norm_stats.npz"
 
     if rank == 0:
         band_means, band_stds = load_or_compute_band_normalization(
@@ -342,7 +343,7 @@ def main(args, rank, world_size, local_rank, device):
         print("band_means:", band_means)
         print("band_stds:", band_stds)
 
-    train_dataset = Kilonova_lc_scalar_context_DataSet_gri(
+    train_dataset = Kilonova_lc_scalar_context_DataSet_9band(
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
@@ -350,7 +351,7 @@ def main(args, rank, world_size, local_rank, device):
         stds=band_stds,
     )
 
-    val_dataset = Kilonova_lc_scalar_context_DataSet_gri(
+    val_dataset = Kilonova_lc_scalar_context_DataSet_9band(
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
@@ -403,7 +404,7 @@ def main(args, rank, world_size, local_rank, device):
 
 
         #train_DDP_loderunner_epoch(
-        train_DDP_scalar_temporal_loderunner_epoch_gri(
+        train_DDP_scalar_temporal_loderunner_epoch_9band(
             training_data=train_dataloader,
             validation_data=val_dataloader,
             num_train_batches=train_batches,
@@ -451,19 +452,19 @@ def main(args, rank, world_size, local_rank, device):
             torch.save(
                 {
                     "epoch": epochIDX,
-                    "model_class": "ScalarTemporalConditionedLodeRunner_gri",
+                    "model_class": "ScalarTemporalConditionedLodeRunner_9band",
                     "backbone_class": "LodeRunner",
                     "model_args": model_args,
                     "model_state_dict": model.module.state_dict(),
                     "optimizer_state_dict": optimizer.state_dict(),
                     "noise_scale": noise_scale,
-                    "predicts_delta": True,
-                    "target_type": "delta_gri",
+                    "predicts_delta": False,
+                    "target_type": "value_9band",
                     "context_len": CONTEXT_LEN,
-                    "n_input_channels": 3,
-                    "n_output_channels": 3,
+                    "n_bands": N_BANDS,
+                    "band_keys": list(BAND_KEYS),
                     "backbone_channels": 8,
-                    "hidden": 64,
+                    "hidden": HIDDEN_CHANNELS,
                 },
                 new_chkpt_path,
             )
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 9a28500a..7a93554d 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -20,6 +20,22 @@
 
 EPS = 1e-6
 
+# Nine-band ordering used by the merged event-stream dataset: the three ZTF
+# bands followed by the six Rubin/LSST-style bands (u from SDSS, g/r/i/z/y from
+# PanSTARRS naming). The index of each key in this tuple is the band index used
+# throughout the 9-band pipeline.
+NINE_BAND_KEYS = (
+    "arr_ztfg",
+    "arr_ztfr",
+    "arr_ztfi",
+    "arr_sdssu",
+    "arr_ps1__g",
+    "arr_ps1__r",
+    "arr_ps1__i",
+    "arr_ps1__z",
+    "arr_ps1__y",
+)
+
 
 def compute_band_normalization(
     file_prefix_list: list[str],
@@ -309,3 +325,233 @@ def __getitem__(
         )
 
         return x, target, Dt
+
+
+class Kilonova_lc_scalar_context_DataSet_9band(Dataset):
+    """Merged event-stream kilonova dataset for sparse, irregular 9-band data.
+
+    Unlike the g/r/i dataset, this class makes no assumption that the bands
+    share a common time grid. Real and Rubin-simulated light curves record each
+    band on its own cadence, with differing numbers of observations per band and
+    entire bands sometimes missing. Every observation across all bands is read
+    as-is, tagged with its band index, and merged into a single time-sorted
+    stream of events.
+
+    Each sample is a window of ``context_len`` consecutive events used to predict
+    the next event in the stream. Because any single future observation belongs
+    to one band, the target is stored as a length-``n_bands`` vector with a mask
+    selecting the observed band; the loss is applied only there. Across the whole
+    dataset every band is supervised, and at inference the model emits a
+    prediction for all bands at a chosen lead time.
+
+    Each returned sample is a ``(x, target, mask, Dt)`` tuple where:
+        x:      flattened context of shape
+                [context_len * (2 + n_bands)], laid out per event as
+                [value, rel_t, one_hot_band(n_bands)].
+        target: normalized value per band, shape [n_bands]; only the observed
+                band is meaningful.
+        mask:   float mask, shape [n_bands]; 1.0 for the observed target band,
+                0.0 elsewhere.
+        Dt:     lead time from the last context event to the target event.
+    """
+
+    def __init__(
+        self,
+        N_imgs: int = 0,
+        context_len: int = 5,
+        band_keys: tuple[str, ...] = NINE_BAND_KEYS,
+        value_col: int = 1,
+        means: np.ndarray = None,
+        stds: np.ndarray = None,
+    ) -> None:
+        """Initialize the dataset and build the merged-event sample index.
+
+        Args:
+            N_imgs (int): Number of light-curve files to sample; 0 uses all.
+            context_len (int): Number of context events per sample.
+            band_keys (tuple[str, ...]): Keys of the bands to load. Their order
+                defines the band index used in the one-hot encoding and target.
+            value_col (int): Column index of the value to load per band.
+            means (np.ndarray): Per-band means for normalization, shape [n_bands].
+            stds (np.ndarray): Per-band stds for normalization, shape [n_bands].
+        """
+        # FIXME: hardcoded scratch path. Should be passed in as an argument so
+        # this dataset does not depend on a user-specific filesystem location.
+        file_prefix_list = sorted(
+            glob.glob(
+                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+                "uniform_dataset_20000/lc_*.npz"
+            )
+        )
+
+        if N_imgs == 0:
+            self.file_prefix_list = file_prefix_list
+        else:
+            self.file_prefix_list = list(
+                np.random.choice(file_prefix_list, N_imgs, replace=False)
+            )
+
+        random.shuffle(self.file_prefix_list)
+
+        self.context_len = context_len
+        self.band_keys = tuple(band_keys)
+        self.value_col = value_col
+        self.n_channels = len(self.band_keys)
+
+        if means is None:
+            raise ValueError(
+                "means must be provided for per-band normalization. "
+                "Expected shape [n_channels]."
+            )
+
+        if stds is None:
+            raise ValueError(
+                "stds must be provided for per-band normalization. "
+                "Expected shape [n_channels]."
+            )
+
+        self.means = np.asarray(means, dtype=np.float32)
+        self.stds = np.asarray(stds, dtype=np.float32)
+
+        if self.means.shape[0] != self.n_channels:
+            raise ValueError(
+                f"means has length {self.means.shape[0]}, "
+                f"but n_channels={self.n_channels}"
+            )
+
+        if self.stds.shape[0] != self.n_channels:
+            raise ValueError(
+                f"stds has length {self.stds.shape[0]}, "
+                f"but n_channels={self.n_channels}"
+            )
+
+        if np.any(self.stds <= 0):
+            raise ValueError(f"All stds must be positive. Got stds={self.stds}")
+
+        # Build the merged, time-sorted event stream for each file and index
+        # every context window into it. Events are (rel_time, value, band_idx).
+        # rel_time is relative to the earliest observation across all bands in
+        # the file so absolute MJD offsets do not leak into the model.
+        self.events_per_file = []
+        self.samples = []
+
+        for file_idx, fn in enumerate(self.file_prefix_list):
+            data = np.load(fn, allow_pickle=True)
+
+            times = []
+            values = []
+            bands = []
+
+            for band_idx, key in enumerate(self.band_keys):
+                # A band may be missing entirely in a given file.
+                if key not in data.files:
+                    continue
+
+                arr = data[key]
+                if arr.size == 0:
+                    continue
+
+                times.append(arr[:, 0].astype(np.float32))
+                values.append(arr[:, value_col].astype(np.float32))
+                bands.append(np.full(arr.shape[0], band_idx, dtype=np.int64))
+
+            data.close()
+
+            if not times:
+                continue
+
+            times = np.concatenate(times)
+            values = np.concatenate(values)
+            bands = np.concatenate(bands)
+
+            # Sort the merged stream by observation time.
+            order = np.argsort(times, kind="stable")
+            times = times[order]
+            values = values[order]
+            bands = bands[order]
+
+            # Relative times within the file.
+            times = times - times.min()
+
+            # Per-band normalization of the values.
+            values = (values - self.means[bands]) / (self.stds[bands] + EPS)
+
+            self.events_per_file.append(
+                (times, values.astype(np.float32), bands)
+            )
+
+            n_events = times.shape[0]
+            max_start = n_events - context_len - 1
+            for startIDX in range(max_start + 1):
+                self.samples.append((len(self.events_per_file) - 1, startIDX))
+
+    def __len__(self) -> int:
+        """Return the number of samples in the dataset."""
+        return len(self.samples)
+
+    def __getitem__(
+        self, index: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return the (input, target, mask, Dt) tuple for a given sample index.
+
+        Args:
+            index (int): Sample index.
+
+        Returns:
+            x (torch.Tensor): Flattened per-event context of shape
+                [context_len * (2 + n_bands)].
+            target (torch.Tensor): Normalized value per band, shape [n_bands];
+                only the observed band is meaningful.
+            mask (torch.Tensor): Float mask, shape [n_bands]; 1.0 for the
+                observed target band, 0.0 elsewhere.
+            Dt (torch.Tensor): Lead time to the target event.
+        """
+        file_idx, startIDX = self.samples[index]
+        times, values, bands = self.events_per_file[file_idx]
+
+        target_idx = startIDX + self.context_len
+
+        # Context events.
+        ctx_t = times[startIDX:target_idx]
+        ctx_v = values[startIDX:target_idx]
+        ctx_b = bands[startIDX:target_idx]
+
+        # Relative times within the context window.
+        rel_t = (ctx_t - ctx_t[0]).astype(np.float32)
+
+        # One-hot encode the band of each context event.
+        band_onehot = np.zeros(
+            (self.context_len, self.n_channels), dtype=np.float32
+        )
+        band_onehot[np.arange(self.context_len), ctx_b] = 1.0
+
+        # Per-event feature: [value, rel_t, one_hot_band(n_bands)].
+        per_event = np.concatenate(
+            [
+                ctx_v[:, None],
+                rel_t[:, None],
+                band_onehot,
+            ],
+            axis=1,
+        )
+
+        # Flatten to [context_len * (2 + n_bands)].
+        x = torch.tensor(per_event.reshape(-1), dtype=torch.float32)
+
+        # Target is the next event, placed into a per-band vector with a mask
+        # marking which band was actually observed.
+        target_band = int(bands[target_idx])
+        target = np.zeros(self.n_channels, dtype=np.float32)
+        mask = np.zeros(self.n_channels, dtype=np.float32)
+        target[target_band] = values[target_idx]
+        mask[target_band] = 1.0
+
+        target = torch.tensor(target, dtype=torch.float32)
+        mask = torch.tensor(mask, dtype=torch.float32)
+
+        Dt = torch.tensor(
+            times[target_idx] - times[target_idx - 1],
+            dtype=torch.float32,
+        )
+
+        return x, target, mask, Dt
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index 507cefa4..8f416597 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -451,6 +451,124 @@ def forward(
         return pred
 
 
+class ScalarTemporalConditionedLodeRunner_9band(nn.Module):
+    """Scalar-temporal wrapper for sparse, irregular 9-band light curves.
+
+    Like ``ScalarTemporalConditionedLodeRunner_gri`` this maps a scalar temporal
+    context into the pseudo-channel image expected by a pretrained LodeRunner
+    backbone and collapses the backbone output back to per-band predictions.
+    Unlike the g/r/i wrapper, the context is a merged event stream where each
+    event carries its own band identity (via a one-hot encoding), so the model
+    handles sparse and irregular sampling with missing bands. The output head
+    emits a prediction for every band at the requested lead time, which supports
+    forecasting all observatories' future observations from real data.
+
+    Args:
+        backbone (nn.Module): Pretrained LodeRunner backbone.
+        context_len (int): Number of context events.
+        n_bands (int): Number of bands (input band identities and output
+            predictions), e.g. 9.
+        image_size (tuple): Spatial size (H, W) fed to the backbone.
+        backbone_channels (int): Number of pseudo-channels the backbone expects.
+        hidden (int): Hidden width of the conditioner/output-head MLPs.
+    """
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        context_len: int = 5,
+        n_bands: int = 9,
+        image_size: tuple[int, int] = (1120, 400),
+        backbone_channels: int = 8,
+        hidden: int = 64,
+    ) -> None:
+        """Initialize conditioner and output-head around the backbone."""
+        super().__init__()
+
+        self.backbone = backbone
+        self.context_len = context_len
+        self.n_bands = n_bands
+        self.image_size = image_size
+        self.backbone_channels = backbone_channels
+
+        # Dataset x layout, flattened per event:
+        #   [value, rel_t, one_hot_band(n_bands)] * context_len
+        #
+        # input_dim = context_len * (2 + n_bands)
+        input_dim = context_len * (2 + n_bands)
+
+        # Maps the scalar temporal event stream into the pseudo-channels
+        # expected by the pretrained LodeRunner backbone.
+        self.conditioner = nn.Sequential(
+            nn.Linear(input_dim, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, backbone_channels),
+        )
+
+        # Maps the backbone-channel summary back to one prediction per band.
+        self.output_head = nn.Sequential(
+            nn.Linear(backbone_channels, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, n_bands),
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        in_vars: torch.Tensor,
+        out_vars: torch.Tensor,
+        Dt: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            x (torch.Tensor): Merged event-stream context of shape
+                [B, context_len * (2 + n_bands)].
+            in_vars (torch.Tensor): Kept for LodeRunner API compatibility.
+            out_vars (torch.Tensor): Kept for LodeRunner API compatibility.
+            Dt (torch.Tensor): Lead-time tensor passed to the backbone.
+
+        Returns:
+            pred (torch.Tensor): Predictions of shape [B, n_bands].
+        """
+        B = x.shape[0]
+        H, W = self.image_size
+
+        channel_vals = self.conditioner(x)  # [B, backbone_channels]
+
+        pseudo_img = channel_vals.view(
+            B,
+            self.backbone_channels,
+            1,
+            1,
+        ).expand(
+            B,
+            self.backbone_channels,
+            H,
+            W,
+        )
+
+        backbone_in_vars = torch.arange(self.backbone_channels, device=x.device)
+        backbone_out_vars = torch.arange(self.backbone_channels, device=x.device)
+
+        pred_img = self.backbone(
+            pseudo_img,
+            backbone_in_vars,
+            backbone_out_vars,
+            Dt,
+        )  # [B, backbone_channels, H, W]
+
+        # Collapse spatial dimensions to backbone-channel summaries.
+        pred_channel_vals = pred_img.mean(dim=(2, 3))  # [B, backbone_channels]
+
+        # Convert backbone channels to per-band predictions.
+        pred = self.output_head(pred_channel_vals)  # [B, n_bands]
+
+        return pred
+
+
 if __name__ == "__main__":
     from yoke.utils.parameters import count_torch_params
 
diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py
index 2ec7eb4d..470b71ba 100644
--- a/src/yoke/utils/training/datastep/loderunner.py
+++ b/src/yoke/utils/training/datastep/loderunner.py
@@ -132,6 +132,109 @@ def eval_DDP_scalar_temporal_loderunner_datastep_gri(
     return target, pred, per_sample_loss.detach()
 
 
+def train_DDP_scalar_temporal_loderunner_datastep_9band(
+    data: tuple,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """DDP training datastep for the masked 9-band scalar temporal wrapper.
+
+    The model predicts all bands, but each sample only observes one future band,
+    so the loss is masked to the observed band. Averaging over the mask gives one
+    scalar loss per sample.
+
+    Expected data:
+        x:      [B, context_len * (2 + n_bands)]
+        target: [B, n_bands]   normalized value, only observed band meaningful
+        mask:   [B, n_bands]   1.0 for observed band, 0.0 elsewhere
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, n_bands]
+    """
+    model.train()
+
+    x, target, mask, Dt = data
+
+    x = x.to(device, non_blocking=True)
+    target = target.to(device, non_blocking=True)
+    mask = mask.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    # Kept for LodeRunner-style API compatibility.
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    pred = model(x, in_vars, out_vars, Dt)
+
+    if pred.shape != target.shape:
+        raise RuntimeError(
+            f"Prediction and target shapes do not match: "
+            f"pred.shape={pred.shape}, target.shape={target.shape}"
+        )
+
+    # loss_fn uses reduction='none' -> [B, n_bands]. Mask to the observed band
+    # and average per sample over the observed entries.
+    loss = loss_fn(pred, target) * mask
+    per_sample_loss = loss.sum(dim=1) / (mask.sum(dim=1) + 1e-8)
+
+    optimizer.zero_grad(set_to_none=True)
+    per_sample_loss.mean().backward()
+    optimizer.step()
+
+    return target, pred, per_sample_loss.detach()
+
+
+def eval_DDP_scalar_temporal_loderunner_datastep_9band(
+    data,
+    model,
+    loss_fn,
+    device,
+    rank,
+    world_size,
+):
+    """Evaluation datastep for the masked 9-band scalar temporal wrapper.
+
+    Expected data:
+        x:      [B, context_len * (2 + n_bands)]
+        target: [B, n_bands]   normalized value, only observed band meaningful
+        mask:   [B, n_bands]   1.0 for observed band, 0.0 elsewhere
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, n_bands]
+    """
+    model.eval()
+
+    x, target, mask, Dt = data
+
+    x = x.to(device, non_blocking=True)
+    target = target.to(device, non_blocking=True)
+    mask = mask.to(device, non_blocking=True)
+    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    with torch.no_grad():
+        pred = model(x, in_vars, out_vars, Dt)
+
+    if pred.shape != target.shape:
+        raise RuntimeError(
+            "Prediction and target shapes do not match in eval datastep: "
+            f"pred.shape={pred.shape}, target.shape={target.shape}"
+        )
+
+    loss = loss_fn(pred, target) * mask
+    per_sample_loss = loss.sum(dim=1) / (mask.sum(dim=1) + 1e-8)
+
+    return target, pred, per_sample_loss.detach()
+
+
 ####################################
 # Evaluating on a Datastep
 ####################################
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 153c30d6..d6a19dff 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -481,6 +481,148 @@ def train_DDP_scalar_temporal_loderunner_epoch_gri(
                         np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
 
 
+def train_DDP_scalar_temporal_loderunner_epoch_9band(
+    training_data: torch.utils.data.DataLoader,
+    validation_data: torch.utils.data.DataLoader,
+    num_train_batches: int,
+    num_val_batches: int,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    LRsched: torch.optim.lr_scheduler._LRScheduler,
+    epochIDX: int,
+    train_per_val: int,
+    train_rcrd_filename: str,
+    val_rcrd_filename: str,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> None:
+    """DDP epoch function for the masked 9-band scalar temporal LodeRunner.
+
+    The dataloader yields a merged event-stream context plus a masked, per-band
+    target. The model predicts all bands; the loss is masked to the single band
+    observed at the target event.
+
+    Expected dataset output:
+        x:      [B, context_len * (2 + n_bands)]
+        target: [B, n_bands]   normalized value, only observed band meaningful
+        mask:   [B, n_bands]   1.0 for observed band, 0.0 elsewhere
+        Dt:     [B]
+
+    Expected model output:
+        pred:   [B, n_bands]
+    """
+    train_rcrd_filename = train_rcrd_filename.replace(
+        "<epochIDX>",
+        f"{epochIDX:04d}",
+    )
+
+    model.train()
+
+    with (
+        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
+    ) as train_rcrd_file:
+
+        for trainbatch_ID, data in enumerate(training_data):
+            if trainbatch_ID >= num_train_batches:
+                break
+
+            x, target, mask, Dt = data
+
+            x = x.to(device, non_blocking=True)
+            target = target.to(device, non_blocking=True)
+            mask = mask.to(device, non_blocking=True)
+            Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+            optimizer.zero_grad(set_to_none=True)
+
+            # Kept for API compatibility with LodeRunner-style wrappers.
+            in_vars = torch.arange(8, device=device)
+            out_vars = torch.arange(8, device=device)
+
+            pred = model(x, in_vars, out_vars, Dt)
+
+            if pred.shape != target.shape:
+                raise RuntimeError(
+                    f"Prediction and target shapes do not match: "
+                    f"pred.shape={pred.shape}, target.shape={target.shape}"
+                )
+
+            # loss_fn uses reduction='none' -> [B, n_bands]. Mask to the
+            # observed band and average per sample over observed entries.
+            loss = loss_fn(pred, target) * mask
+            per_sample_loss = loss.sum(dim=1) / (mask.sum(dim=1) + 1e-8)
+
+            batch_loss = per_sample_loss.mean()
+
+            batch_loss.backward()
+            optimizer.step()
+            LRsched.step()
+
+            if rank == 0:
+                batch_records = np.column_stack(
+                    [
+                        np.full(len(per_sample_loss), epochIDX),
+                        np.full(len(per_sample_loss), trainbatch_ID),
+                        per_sample_loss.detach().cpu().numpy().flatten(),
+                    ]
+                )
+                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+    if epochIDX % train_per_val == 0:
+        if rank == 0:
+            print("Validating...", epochIDX, flush=True)
+
+        val_rcrd_filename = val_rcrd_filename.replace(
+            "<epochIDX>",
+            f"{epochIDX:04d}",
+        )
+
+        model.eval()
+
+        with (
+            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
+        ) as val_rcrd_file:
+
+            with torch.no_grad():
+                for valbatch_ID, data in enumerate(validation_data):
+                    if valbatch_ID >= num_val_batches:
+                        break
+
+                    x, target, mask, Dt = data
+
+                    x = x.to(device, non_blocking=True)
+                    target = target.to(device, non_blocking=True)
+                    mask = mask.to(device, non_blocking=True)
+                    Dt = Dt.to(torch.float32).to(device, non_blocking=True)
+
+                    in_vars = torch.arange(8, device=device)
+                    out_vars = torch.arange(8, device=device)
+
+                    pred = model(x, in_vars, out_vars, Dt)
+
+                    if pred.shape != target.shape:
+                        raise RuntimeError(
+                            f"Validation prediction and target shapes do not "
+                            f"match: pred.shape={pred.shape}, "
+                            f"target.shape={target.shape}"
+                        )
+
+                    loss = loss_fn(pred, target) * mask
+                    per_sample_loss = loss.sum(dim=1) / (mask.sum(dim=1) + 1e-8)
+
+                    if rank == 0:
+                        batch_records = np.column_stack(
+                            [
+                                np.full(len(per_sample_loss), epochIDX),
+                                np.full(len(per_sample_loss), valbatch_ID),
+                                per_sample_loss.detach().cpu().numpy().flatten(),
+                            ]
+                        )
+                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+
 def train_DDP_loderunner_epoch(
     training_data: torch.utils.data.DataLoader,
     validation_data: torch.utils.data.DataLoader,

From 876c4d822dbf5dc7434cb31cbb95eaf882ec0cce Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 10:56:02 -0600
Subject: [PATCH 15/66] continuation fix

---
 .../KN_loderunner/plot_loss_curves_9band.py   |  96 +++++++++++++
 .../KN_loderunner/train_LodeRunner_ddp.py     |   4 +-
 src/yoke/utils/checkpointing.py               | 131 ++++++++++++++++++
 3 files changed, 229 insertions(+), 2 deletions(-)
 create mode 100644 applications/harnesses/KN_loderunner/plot_loss_curves_9band.py

diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
new file mode 100644
index 00000000..5a9cf361
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
@@ -0,0 +1,96 @@
+"""Plot training/validation loss curves for the 9-band scalar temporal LodeRunner.
+
+The 9-band training loop writes the same per-batch CSV records as the g/r/i runs
+(``training_study<studyIDX>_epoch<epochIDX>.csv`` and the matching validation
+files), with the format ``epoch, batch, loss`` where ``loss`` is the single
+masked scalar loss per sample. The loss-curve plotting logic is therefore
+identical; this module just reuses ``plot_loss_curves_gri`` with a 9-band title.
+
+Run directly, e.g.:
+    python plot_loss_curves_9band.py --study 24
+"""
+
+import plot_loss_curves_gri as base
+
+
+def main():
+    parser = base.argparse.ArgumentParser(
+        description=(
+            "Plot training/validation loss curves for scalar temporal "
+            "LodeRunner 9-band runs."
+        )
+    )
+
+    parser.add_argument("--study", type=int, default=base.DEFAULT_STUDY)
+    parser.add_argument("--runs_root", type=str, default=base.DEFAULT_RUNS_ROOT)
+    parser.add_argument("--train_pattern", type=str, default=None)
+    parser.add_argument("--val_pattern", type=str, default=None)
+    parser.add_argument("--out", type=str, default=None)
+
+    parser.add_argument(
+        "--columns",
+        type=str,
+        default=base.DEFAULT_COLUMNS,
+        help=(
+            "Comma-separated CSV column names. The first two must be "
+            "epoch,batch. For 9-band runs the loss is a single scalar column, "
+            "so the default epoch,batch,loss is correct."
+        ),
+    )
+
+    parser.add_argument("--title", type=str, default="9-band loss curves")
+    parser.add_argument("--dpi", type=int, default=200)
+
+    parser.add_argument(
+        "--logy",
+        dest="logy",
+        action="store_true",
+        default=True,
+        help="Use log scale on y-axis. This is the default.",
+    )
+    parser.add_argument(
+        "--linear",
+        dest="logy",
+        action="store_false",
+        help="Use linear y-axis.",
+    )
+
+    parser.add_argument(
+        "--show_std",
+        action="store_true",
+        help="Shade +/- one epoch standard deviation.",
+    )
+    parser.add_argument(
+        "--require_val",
+        action="store_true",
+        help=(
+            "Fail instead of continuing when validation records are missing "
+            "or invalid."
+        ),
+    )
+
+    args = parser.parse_args()
+
+    defaults = base.default_patterns(args.study, args.runs_root)
+
+    args.train_pattern = args.train_pattern or defaults["train"]
+    args.val_pattern = args.val_pattern or defaults["val"]
+    args.out = args.out or defaults["out"]
+
+    train = base.load_records(args.train_pattern, args.columns)
+    base.print_loaded("training", train)
+
+    try:
+        val = base.load_records(args.val_pattern, args.columns)
+        base.print_loaded("validation", val)
+    except Exception as exc:
+        if args.require_val:
+            raise
+        print(f"No validation curve plotted: {exc}")
+        val = None
+
+    base.plot_epoch_curves(train, val, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index a0c9a869..9e83b05e 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -24,7 +24,7 @@
 from yoke.utils.dataload import make_distributed_dataloader
 from yoke.utils.checkpointing import load_model_and_optimizer
 from yoke.utils.checkpointing import save_model_and_optimizer
-from yoke.utils.checkpointing import load_direct_loderunner_checkpoint
+from yoke.utils.checkpointing import load_direct_loderunner_checkpoint_9band
 from yoke.utils.parallel import setup_distributed, cleanup_distributed
 from yoke.lr_schedulers import CosineWithWarmupScheduler
 from yoke.helpers import cli
@@ -155,7 +155,7 @@ def main(args, rank, world_size, local_rank, device):
 
 
     if CONTINUATION:
-        model, optimizer, starting_epoch = load_direct_loderunner_checkpoint(
+        model, optimizer, starting_epoch = load_direct_loderunner_checkpoint_9band(
             checkpoint_path=checkpoint,
             model_args=model_args,
             optimizer_kwargs=optimizer_kwargs,
diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index 4f407e64..8c04b7c2 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -13,6 +13,7 @@
 from yoke.models.vit.swin.bomberman import (
     LodeRunner,
     ScalarTemporalConditionedLodeRunner_gri,
+    ScalarTemporalConditionedLodeRunner_9band,
 )
 
 
@@ -434,3 +435,133 @@ def load_direct_loderunner_checkpoint(
                     state[key] = value.to(device)
 
     return model, optimizer, starting_epoch
+
+
+def load_direct_loderunner_checkpoint_9band(
+    checkpoint_path: str,
+    model_args: dict,
+    optimizer_kwargs: dict,
+    device: torch.device,
+) -> tuple[torch.nn.Module, torch.optim.Optimizer, int]:
+    """Load a ScalarTemporalConditionedLodeRunner_9band model from a checkpoint.
+
+    The 9-band analogue of ``load_direct_loderunner_checkpoint``. Handles two
+    checkpoint types:
+      - An old plain LodeRunner checkpoint, whose weights are loaded into the
+        wrapper's backbone (conditioner/output-head are freshly initialized and
+        this is not treated as a true continuation).
+      - A ScalarTemporalConditionedLodeRunner_9band wrapper checkpoint, which is
+        loaded in full and treated as a continuation.
+
+    The backbone is frozen and only the conditioner and output-head parameters
+    are trainable.
+
+    Args:
+        checkpoint_path (str): Path to the checkpoint file.
+        model_args (dict): Fallback LodeRunner init args if the checkpoint has
+            none stored.
+        optimizer_kwargs (dict): Kwargs for the AdamW optimizer.
+        device (torch.device): Device to load the model/optimizer onto.
+
+    Returns:
+        model (torch.nn.Module): The wrapper model.
+        optimizer (torch.optim.Optimizer): Optimizer over trainable parameters.
+        starting_epoch (int): Epoch to continue training from.
+    """
+    checkpoint_data = torch.load(
+        checkpoint_path,
+        map_location=device,
+        weights_only=False,
+    )
+
+    saved_model_args = checkpoint_data.get("model_args", model_args)
+    context_len = checkpoint_data.get("context_len", 5)
+
+    backbone = LodeRunner(**saved_model_args).to(device)
+
+    model = ScalarTemporalConditionedLodeRunner_9band(
+        backbone=backbone,
+        context_len=context_len,
+        n_bands=checkpoint_data.get("n_bands", 9),
+        image_size=saved_model_args["image_size"],
+        backbone_channels=checkpoint_data.get("backbone_channels", 8),
+        hidden=checkpoint_data.get("hidden", 64),
+    ).to(device)
+
+    state_dict = checkpoint_data["model_state_dict"]
+
+    # Remove DDP prefix if present
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        state_dict = {
+            k.replace("module.", "", 1): v
+            for k, v in state_dict.items()
+        }
+
+    # Detect checkpoint type
+    is_wrapper_checkpoint = any(
+        k.startswith("backbone.") for k in state_dict.keys()
+    )
+
+    # -------------------------------------------------
+    # OLD plain LodeRunner checkpoint
+    # -------------------------------------------------
+    if not is_wrapper_checkpoint:
+        missing_keys, unexpected_keys = model.backbone.load_state_dict(
+            state_dict,
+            strict=False,
+        )
+
+        print("Loaded old LodeRunner checkpoint into model.backbone")
+        print("Missing backbone keys:", missing_keys)
+        print("Unexpected backbone keys:", unexpected_keys)
+
+        # This is NOT a true continuation.
+        # Conditioner is newly initialized.
+        starting_epoch = 0
+
+    # -------------------------------------------------
+    # NEW ScalarTemporalConditionedLodeRunner_9band checkpoint
+    # -------------------------------------------------
+    else:
+        model.load_state_dict(state_dict, strict=True)
+
+        print("Loaded ScalarTemporalConditionedLodeRunner_9band checkpoint")
+
+        starting_epoch = checkpoint_data.get("epoch", 0)
+
+    noise_scale = checkpoint_data.get("noise_scale", 0.0)
+    model.backbone.noise_scale = noise_scale
+
+    # Freeze pretrained backbone
+    for p in model.backbone.parameters():
+        p.requires_grad = False
+
+    # Train conditioner
+    for p in model.conditioner.parameters():
+        p.requires_grad = True
+
+    # Train output head
+    for p in model.output_head.parameters():
+        p.requires_grad = True
+
+    optimizer = torch.optim.AdamW(
+        list(model.conditioner.parameters())
+        + list(model.output_head.parameters()),
+        **optimizer_kwargs,
+    )
+
+    # Only restore optimizer for TRUE continuation checkpoints
+    if (
+        is_wrapper_checkpoint
+        and "optimizer_state_dict" in checkpoint_data
+    ):
+        optimizer.load_state_dict(
+            checkpoint_data["optimizer_state_dict"]
+        )
+
+        for state in optimizer.state.values():
+            for key, value in state.items():
+                if isinstance(value, torch.Tensor):
+                    state[key] = value.to(device)
+
+    return model, optimizer, starting_epoch

From 8d69591fb7da4884f123aecad47b3a10612dc448 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 11:26:35 -0600
Subject: [PATCH 16/66] ignore upper lims

---
 .../harnesses/KN_loderunner/infer_9band.py    | 15 +++++++
 .../plot_pred_diagnostics_9band.py            |  9 ++++
 .../KN_loderunner/train_LodeRunner_ddp.py     | 12 ++++++
 src/yoke/datasets/kilonova_dataset.py         | 41 +++++++++++++++++++
 4 files changed, 77 insertions(+)

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 7a0962c6..548443d7 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -56,8 +56,13 @@
     "#264653",  # y
 )
 VALUE_COL = 1
+ERROR_COL = 2
 N_BANDS = len(BAND_KEYS)
 
+# Match training: drop upper-limit (non-detection) observations, flagged by a
+# non-finite uncertainty in ERROR_COL, from the context fed to the model.
+DROP_UPPER_LIMITS = True
+
 
 def study_tag(study):
     return f"{int(study):03d}"
@@ -208,6 +213,14 @@ def load_event_stream(fn, means, stds):
         if arr.size == 0:
             continue
 
+        # Drop upper-limit (non-detection) rows, flagged by a non-finite
+        # uncertainty in ERROR_COL, matching how the model was trained.
+        if DROP_UPPER_LIMITS:
+            detected = np.isfinite(arr[:, ERROR_COL])
+            arr = arr[detected]
+            if arr.shape[0] == 0:
+                continue
+
         t = arr[:, 0].astype(np.float32)
         v = arr[:, VALUE_COL].astype(np.float32)
 
@@ -375,6 +388,8 @@ def main():
         stats_path=args.norm_stats_path,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
+        error_col=ERROR_COL,
+        drop_upper_limits=DROP_UPPER_LIMITS,
     )
     means = np.asarray(means, dtype=np.float32)
     stds = np.asarray(stds, dtype=np.float32)
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 1e0e0690..5d371c19 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -41,8 +41,13 @@
 BAND_KEYS = NINE_BAND_KEYS
 BAND_NAMES = ("ztfg", "ztfr", "ztfi", "u", "g", "r", "i", "z", "y")
 VALUE_COL = 1
+ERROR_COL = 2
 N_BANDS = len(BAND_KEYS)
 
+# Match training: drop upper-limit (non-detection) observations, flagged by a
+# non-finite uncertainty in ERROR_COL.
+DROP_UPPER_LIMITS = True
+
 
 def study_tag(study):
     return f"{int(study):03d}"
@@ -150,6 +155,8 @@ def make_eval_dataset(args, context_len):
         stats_path=args.norm_stats_path,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
+        error_col=ERROR_COL,
+        drop_upper_limits=DROP_UPPER_LIMITS,
     )
 
     print("Using band normalization:")
@@ -161,6 +168,8 @@ def make_eval_dataset(args, context_len):
         context_len=context_len,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
+        error_col=ERROR_COL,
+        drop_upper_limits=DROP_UPPER_LIMITS,
         means=band_means,
         stds=band_stds,
     )
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 9e83b05e..32180a12 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -144,8 +144,14 @@ def main(args, rank, world_size, local_rank, device):
     # Nine-band merged event-stream setup (3 ZTF + 6 Rubin/LSST bands).
     BAND_KEYS = NINE_BAND_KEYS
     VALUE_COL = 1
+    ERROR_COL = 2
     N_BANDS = len(BAND_KEYS)
 
+    # Upper-limit (non-detection) observations are flagged by a non-finite
+    # uncertainty in ERROR_COL. Drop them so the model trains only on real
+    # detections; normalization statistics are computed the same way.
+    DROP_UPPER_LIMITS = True
+
     optimizer_kwargs = {
         "lr": 1e-4,# 1e-4, #1e-5
         "betas": (0.9, 0.999),
@@ -328,6 +334,8 @@ def main(args, rank, world_size, local_rank, device):
             stats_path=norm_stats_path,
             band_keys=BAND_KEYS,
             value_col=VALUE_COL,
+            error_col=ERROR_COL,
+            drop_upper_limits=DROP_UPPER_LIMITS,
         )
 
     dist.barrier()
@@ -347,6 +355,8 @@ def main(args, rank, world_size, local_rank, device):
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
+        error_col=ERROR_COL,
+        drop_upper_limits=DROP_UPPER_LIMITS,
         means=band_means,
         stds=band_stds,
     )
@@ -355,6 +365,8 @@ def main(args, rank, world_size, local_rank, device):
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
+        error_col=ERROR_COL,
+        drop_upper_limits=DROP_UPPER_LIMITS,
         means=band_means,
         stds=band_stds,
     )
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 7a93554d..f84d48dc 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -41,6 +41,8 @@ def compute_band_normalization(
     file_prefix_list: list[str],
     band_keys: tuple[str, ...] = ("arr_ztfg", "arr_ztfr", "arr_ztfi"),
     value_col: int = 1,
+    error_col: int = 2,
+    drop_upper_limits: bool = False,
     stats_path: str = "kilonova_gri_norm_stats.npz",
 ) -> tuple[np.ndarray, np.ndarray]:
     """Compute global per-band mean/std over the training files only.
@@ -49,6 +51,11 @@ def compute_band_normalization(
         file_prefix_list (list[str]): List of npz files to accumulate stats over.
         band_keys (tuple[str, ...]): Keys of the bands to normalize.
         value_col (int): Column index of the value to accumulate per band.
+        error_col (int): Column index of the per-observation uncertainty. Only
+            used when drop_upper_limits is True.
+        drop_upper_limits (bool): If True, observations with a non-finite
+            uncertainty in error_col (upper limits / non-detections) are excluded
+            from the statistics, matching a dataset that drops them.
         stats_path (str): Path to save the computed statistics to.
 
     Returns:
@@ -66,6 +73,13 @@ def compute_band_normalization(
             vals = data[key][:, value_col].astype(np.float64)
 
             finite = np.isfinite(vals)
+
+            # Optionally exclude upper limits (non-finite uncertainty) so the
+            # normalization statistics match a dataset that drops them.
+            if drop_upper_limits:
+                errs = data[key][:, error_col].astype(np.float64)
+                finite = finite & np.isfinite(errs)
+
             vals = vals[finite]
 
             sums[b] += vals.sum()
@@ -101,6 +115,8 @@ def load_or_compute_band_normalization(
     stats_path: str = "kilonova_gri_norm_stats.npz",
     band_keys: tuple[str, ...] = ("arr_ztfg", "arr_ztfr", "arr_ztfi"),
     value_col: int = 1,
+    error_col: int = 2,
+    drop_upper_limits: bool = False,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Load cached per-band normalization stats or compute them if missing.
 
@@ -108,6 +124,10 @@ def load_or_compute_band_normalization(
         stats_path (str): Path to load/save the statistics.
         band_keys (tuple[str, ...]): Keys of the bands to normalize.
         value_col (int): Column index of the value to accumulate per band.
+        error_col (int): Column index of the per-observation uncertainty. Only
+            used when drop_upper_limits is True.
+        drop_upper_limits (bool): If True, exclude upper limits (non-finite
+            uncertainty) from the statistics, matching a dataset that drops them.
 
     Returns:
         means (np.ndarray): Per-band means, shape [n_bands].
@@ -137,6 +157,8 @@ def load_or_compute_band_normalization(
         file_prefix_list=file_prefix_list,
         band_keys=band_keys,
         value_col=value_col,
+        error_col=error_col,
+        drop_upper_limits=drop_upper_limits,
         stats_path=stats_path,
     )
 
@@ -361,6 +383,8 @@ def __init__(
         context_len: int = 5,
         band_keys: tuple[str, ...] = NINE_BAND_KEYS,
         value_col: int = 1,
+        error_col: int = 2,
+        drop_upper_limits: bool = True,
         means: np.ndarray = None,
         stds: np.ndarray = None,
     ) -> None:
@@ -372,6 +396,12 @@ def __init__(
             band_keys (tuple[str, ...]): Keys of the bands to load. Their order
                 defines the band index used in the one-hot encoding and target.
             value_col (int): Column index of the value to load per band.
+            error_col (int): Column index of the per-observation uncertainty.
+                Upper-limit (non-detection) rows are flagged by a non-finite
+                (e.g. inf) uncertainty in this column.
+            drop_upper_limits (bool): If True, observations with a non-finite
+                uncertainty in ``error_col`` are dropped from the event stream so
+                the model only sees real detections.
             means (np.ndarray): Per-band means for normalization, shape [n_bands].
             stds (np.ndarray): Per-band stds for normalization, shape [n_bands].
         """
@@ -396,6 +426,8 @@ def __init__(
         self.context_len = context_len
         self.band_keys = tuple(band_keys)
         self.value_col = value_col
+        self.error_col = error_col
+        self.drop_upper_limits = drop_upper_limits
         self.n_channels = len(self.band_keys)
 
         if means is None:
@@ -451,6 +483,15 @@ def __init__(
                 if arr.size == 0:
                     continue
 
+                # Drop upper-limit (non-detection) rows, which are flagged by a
+                # non-finite uncertainty (e.g. inf) in error_col. This keeps only
+                # real detections in the merged event stream.
+                if self.drop_upper_limits:
+                    detected = np.isfinite(arr[:, self.error_col])
+                    arr = arr[detected]
+                    if arr.shape[0] == 0:
+                        continue
+
                 times.append(arr[:, 0].astype(np.float32))
                 values.append(arr[:, value_col].astype(np.float32))
                 bands.append(np.full(arr.shape[0], band_idx, dtype=np.int64))

From 6d9cce9836294bfa1c4f4fa152e439eef324d5f1 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen13@github.com>
Date: Wed, 12 Aug 2026 11:28:33 -0600
Subject: [PATCH 17/66] update KN data path

---
 src/yoke/datasets/kilonova_dataset.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index f84d48dc..623735a1 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -137,8 +137,7 @@ def load_or_compute_band_normalization(
     # library function does not depend on a user-specific filesystem location.
     file_prefix_list = sorted(
         glob.glob(
-            "/net/sescratch1/atoivonen/data/KN_lightcurves/uniform_dataset_20000/lc_*.npz"
-        )
+        "/net/sescratch1/atoivonen/data/KN_lightcurves/rubin_ztf_10000_dataset/lc_*.npz")
     )
 
     if os.path.exists(stats_path):

From 5a1bb0d39e6b608b20982bcc88221c9a564d7486 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 13:38:53 -0600
Subject: [PATCH 18/66] update plot script

---
 .../plot_pred_diagnostics_9band.py            | 88 +++++++++++++++----
 1 file changed, 73 insertions(+), 15 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 5d371c19..b512b7da 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -15,11 +15,13 @@
 import argparse
 import csv
 import os
+import time
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+from torch.utils.data import DataLoader, Subset
 
 from yoke.models.vit.swin.bomberman import (
     LodeRunner,
@@ -67,6 +69,22 @@ def get_args():
 
     parser.add_argument("--N_imgs", type=int, default=50)
 
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=32,
+        help="Batch size for running the backbone over eval samples. Larger is "
+        "faster; the backbone forward dominates runtime.",
+    )
+    parser.add_argument(
+        "--max_samples",
+        type=int,
+        default=0,
+        help="Cap on the number of eval samples actually run through the model "
+        "(0 = all). The full 1120x400 backbone runs once per batch, so on CPU "
+        "this bounds runtime.",
+    )
+
     parser.add_argument("--outdir", type=str, default=None)
     parser.add_argument(
         "--norm_stats_path",
@@ -177,8 +195,15 @@ def make_eval_dataset(args, context_len):
     return dataset
 
 
-def collect_next_event_predictions(dataset, model, device, n_bands):
-    """Run the direct next-event prediction over every sample.
+def collect_next_event_predictions(
+    dataset, model, device, n_bands, batch_size=32, max_samples=0
+):
+    """Run the direct next-event prediction over the eval samples.
+
+    Samples are batched through the model so the (expensive) backbone forward
+    runs once per batch rather than once per sample. This is the difference
+    between the script appearing to hang and finishing quickly, especially on
+    CPU where a single 1120x400 backbone pass is not cheap.
 
     Returns a dict of per-band arrays of (pred, truth) for the observed band of
     each sample, plus the residuals.
@@ -186,25 +211,56 @@ def collect_next_event_predictions(dataset, model, device, n_bands):
     preds_by_band = [[] for _ in range(n_bands)]
     truths_by_band = [[] for _ in range(n_bands)]
 
-    with torch.no_grad():
-        for idx in range(len(dataset)):
-            x, target, mask, Dt = dataset[idx]
+    # Optionally cap the number of samples so runtime is bounded and visible.
+    if max_samples and max_samples < len(dataset):
+        eval_dataset = Subset(dataset, list(range(max_samples)))
+    else:
+        eval_dataset = dataset
+
+    loader = DataLoader(
+        eval_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=0,
+    )
 
-            x = torch.as_tensor(x, dtype=torch.float32, device=device).unsqueeze(0)
-            target = torch.as_tensor(target, dtype=torch.float32)
-            mask = torch.as_tensor(mask, dtype=torch.float32)
-            Dt = torch.as_tensor(Dt, dtype=torch.float32, device=device)
+    n_samples = len(eval_dataset)
+    n_batches = len(loader)
+    print(
+        f"Running {n_samples} samples through the model in {n_batches} "
+        f"batches of up to {batch_size}...",
+        flush=True,
+    )
+
+    start = time.time()
+    seen = 0
+
+    with torch.no_grad():
+        for batch_idx, (x, target, mask, Dt) in enumerate(loader):
+            x = x.to(device=device, dtype=torch.float32)
+            Dt = Dt.to(device=device, dtype=torch.float32)
 
             if Dt.ndim == 0:
                 Dt = Dt.unsqueeze(0)
 
             pred = model(x, in_vars=None, out_vars=None, Dt=Dt)
-            pred = pred.reshape(n_bands).detach().cpu()
-
-            band_idx = int(torch.argmax(mask).item())
-
-            preds_by_band[band_idx].append(float(pred[band_idx]))
-            truths_by_band[band_idx].append(float(target[band_idx]))
+            pred = pred.detach().cpu()  # [B, n_bands]
+
+            # Observed band per sample (mask is one-hot over bands).
+            band_idx = torch.argmax(mask, dim=1)  # [B]
+
+            for i in range(pred.shape[0]):
+                b = int(band_idx[i].item())
+                preds_by_band[b].append(float(pred[i, b]))
+                truths_by_band[b].append(float(target[i, b]))
+
+            seen += x.shape[0]
+            elapsed = time.time() - start
+            print(
+                f"  batch {batch_idx + 1}/{n_batches}  "
+                f"({seen}/{n_samples} samples, {elapsed:.1f}s)",
+                flush=True,
+            )
 
     results = []
     for band_idx in range(n_bands):
@@ -331,6 +387,8 @@ def main():
         model=model,
         device=device,
         n_bands=n_bands,
+        batch_size=args.batch_size,
+        max_samples=args.max_samples,
     )
 
     pred_truth_path = os.path.join(

From b76cd8f20cf00b145bc14de57f9d14080425cf4a Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 14:00:12 -0600
Subject: [PATCH 19/66] fix plotting script

---
 .../plot_pred_diagnostics_9band.py            | 596 ++++++++++++------
 1 file changed, 417 insertions(+), 179 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index b512b7da..b7fbf967 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -1,33 +1,42 @@
-"""Next-event prediction diagnostics for the 9-band scalar temporal LodeRunner.
-
-The 9-band model is trained on a merged event stream: each sample is a window of
-``context_len`` consecutive observations (across all bands) plus a lead time, and
-the model predicts the value in every band at that lead time. Each training
-target only observes one band, so diagnostics here compare the model's prediction
-for the observed target band against the truth, aggregated per band.
-
-Unlike the g/r/i diagnostics this script does NOT roll out autoregressively.
-Autoregression is ill-defined for a mixed-band event stream (each future event
-belongs to a single band), so we evaluate the direct one-step-ahead prediction
-the model is actually trained on.
+"""Autoregressive rollout diagnostics for the 9-band scalar temporal LodeRunner.
+
+This is the 9-band analogue of ``plot_pred_diagnostics_gri.py``. Rather than
+scoring a single next-event prediction, it produces an autoregressive forecast of
+the next several observations and feeds the model's own predictions back into the
+context, exactly like the g/r/i rollout.
+
+The 9-band data is a merged, time-sorted event stream where each observation
+belongs to a single filter. The model, however, emits a prediction for ALL nine
+bands at any requested lead time. So a rollout proceeds as:
+
+  1. Start from a context window of ``context_len`` true events.
+  2. Predict all 9 bands at the next event's lead time Dt.
+  3. The next true event is in one filter; take the model's prediction for that
+     filter as the forecast value, append it (with the event's true time and
+     band) back into the context, and drop the oldest event.
+  4. Repeat for ``n_future_steps``, so later predictions are conditioned on
+     earlier predictions.
+
+Because the observation schedule (times + which filter is seen next) is taken
+from the truth while the values are fed back from the model, this measures how
+well the model forecasts future observations in each filter over a rollout.
 """
 
 import argparse
 import csv
 import os
-import time
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
-from torch.utils.data import DataLoader, Subset
 
 from yoke.models.vit.swin.bomberman import (
     LodeRunner,
     ScalarTemporalConditionedLodeRunner_9band,
 )
 from yoke.datasets.kilonova_dataset import (
+    EPS,
     NINE_BAND_KEYS,
     Kilonova_lc_scalar_context_DataSet_9band,
     load_or_compute_band_normalization,
@@ -42,6 +51,17 @@
 
 BAND_KEYS = NINE_BAND_KEYS
 BAND_NAMES = ("ztfg", "ztfr", "ztfi", "u", "g", "r", "i", "z", "y")
+BAND_COLORS = (
+    "#2A9D8F",  # ztfg
+    "#E63946",  # ztfr
+    "#F4A261",  # ztfi
+    "#457B9D",  # u
+    "#1B9E77",  # g
+    "#D62828",  # r
+    "#E9C46A",  # i
+    "#8338EC",  # z
+    "#264653",  # y
+)
 VALUE_COL = 1
 ERROR_COL = 2
 N_BANDS = len(BAND_KEYS)
@@ -58,8 +78,8 @@ def study_tag(study):
 def get_args():
     parser = argparse.ArgumentParser(
         description=(
-            "Next-event prediction diagnostics for scalar temporal LodeRunner "
-            "9-band runs."
+            "Autoregressive rollout diagnostics for the scalar temporal "
+            "LodeRunner 9-band model."
         )
     )
 
@@ -67,22 +87,23 @@ def get_args():
     parser.add_argument("--epoch", type=int, default=500)
     parser.add_argument("--ckpt", type=str, default=None)
 
-    parser.add_argument("--N_imgs", type=int, default=50)
-
     parser.add_argument(
-        "--batch_size",
+        "--N_imgs",
+        type=int,
+        default=50,
+        help="Number of light-curve files to load into the eval dataset.",
+    )
+    parser.add_argument(
+        "--n_future_steps",
         type=int,
-        default=32,
-        help="Batch size for running the backbone over eval samples. Larger is "
-        "faster; the backbone forward dominates runtime.",
+        default=15,
+        help="Number of future events to forecast autoregressively per series.",
     )
     parser.add_argument(
-        "--max_samples",
+        "--n_series",
         type=int,
-        default=0,
-        help="Cap on the number of eval samples actually run through the model "
-        "(0 = all). The full 1120x400 backbone runs once per batch, so on CPU "
-        "this bounds runtime.",
+        default=10,
+        help="Number of light curves to roll out.",
     )
 
     parser.add_argument("--outdir", type=str, default=None)
@@ -104,7 +125,7 @@ def resolve_paths(args):
         )
 
     if args.outdir is None:
-        args.outdir = f"runs/study_{tag}/next_event_diagnostics_9band"
+        args.outdir = f"runs/study_{tag}/autoreg_diagnostics_9band"
 
     return tag
 
@@ -192,178 +213,357 @@ def make_eval_dataset(args, context_len):
         stds=band_stds,
     )
 
-    return dataset
+    return dataset, np.asarray(band_means), np.asarray(band_stds)
 
 
-def collect_next_event_predictions(
-    dataset, model, device, n_bands, batch_size=32, max_samples=0
-):
-    """Run the direct next-event prediction over the eval samples.
-
-    Samples are batched through the model so the (expensive) backbone forward
-    runs once per batch rather than once per sample. This is the difference
-    between the script appearing to hang and finishing quickly, especially on
-    CPU where a single 1120x400 backbone pass is not cheap.
+def build_context_input(win_v, win_t, win_b, context_len, n_bands, device):
+    """Build the flattened per-event context input for the model.
 
-    Returns a dict of per-band arrays of (pred, truth) for the observed band of
-    each sample, plus the residuals.
+    Layout per event: [value, rel_t, one_hot_band(n_bands)], relative time
+    measured from the first event in the window, matching the dataset.
     """
-    preds_by_band = [[] for _ in range(n_bands)]
-    truths_by_band = [[] for _ in range(n_bands)]
-
-    # Optionally cap the number of samples so runtime is bounded and visible.
-    if max_samples and max_samples < len(dataset):
-        eval_dataset = Subset(dataset, list(range(max_samples)))
-    else:
-        eval_dataset = dataset
-
-    loader = DataLoader(
-        eval_dataset,
-        batch_size=batch_size,
-        shuffle=False,
-        num_workers=0,
-    )
+    win_v = np.asarray(win_v, dtype=np.float32)
+    win_t = np.asarray(win_t, dtype=np.float32)
+    win_b = np.asarray(win_b, dtype=np.int64)
+
+    rel_t = (win_t - win_t[0]).astype(np.float32)
 
-    n_samples = len(eval_dataset)
-    n_batches = len(loader)
-    print(
-        f"Running {n_samples} samples through the model in {n_batches} "
-        f"batches of up to {batch_size}...",
-        flush=True,
+    band_onehot = np.zeros((context_len, n_bands), dtype=np.float32)
+    band_onehot[np.arange(context_len), win_b] = 1.0
+
+    per_event = np.concatenate(
+        [win_v[:, None], rel_t[:, None], band_onehot],
+        axis=1,
     )
 
-    start = time.time()
-    seen = 0
+    return torch.tensor(
+        per_event.reshape(-1),
+        dtype=torch.float32,
+        device=device,
+    ).unsqueeze(0)
+
+
+def get_rollout_from_stream(
+    times,
+    values,
+    bands,
+    model,
+    device,
+    start_idx,
+    n_future_steps,
+    context_len,
+    n_bands,
+    means,
+    stds,
+):
+    """Autoregressively forecast the next events of one merged event stream.
+
+    Args:
+        times (np.ndarray): Relative event times for the file [N].
+        values (np.ndarray): Normalized event values [N].
+        bands (np.ndarray): Band index per event [N].
+        start_idx (int): Index of the first context event.
+        n_future_steps (int): Number of future events to forecast.
+        context_len (int): Context window length.
+        n_bands (int): Number of bands.
+        means, stds (np.ndarray): Per-band normalization for denormalizing.
+
+    Returns:
+        dict describing the rollout (context, per-step forecasts, per-band MSE).
+    """
+    t_ref = float(times[start_idx])
+
+    # Running context window; values are fed back from predictions as we roll
+    # out, while times and band identities follow the true observation schedule.
+    ctx_t = list(times[start_idx : start_idx + context_len].astype(np.float32))
+    ctx_v = list(values[start_idx : start_idx + context_len].astype(np.float32))
+    ctx_b = list(bands[start_idx : start_idx + context_len].astype(np.int64))
+
+    context = {
+        "t_rel": np.asarray(ctx_t, dtype=np.float32) - t_ref,
+        "v_norm": np.asarray(ctx_v, dtype=np.float32),
+        "band": np.asarray(ctx_b, dtype=np.int64),
+    }
+
+    steps = []
+    band_sq_err = [[] for _ in range(n_bands)]
 
     with torch.no_grad():
-        for batch_idx, (x, target, mask, Dt) in enumerate(loader):
-            x = x.to(device=device, dtype=torch.float32)
-            Dt = Dt.to(device=device, dtype=torch.float32)
-
-            if Dt.ndim == 0:
-                Dt = Dt.unsqueeze(0)
-
-            pred = model(x, in_vars=None, out_vars=None, Dt=Dt)
-            pred = pred.detach().cpu()  # [B, n_bands]
-
-            # Observed band per sample (mask is one-hot over bands).
-            band_idx = torch.argmax(mask, dim=1)  # [B]
-
-            for i in range(pred.shape[0]):
-                b = int(band_idx[i].item())
-                preds_by_band[b].append(float(pred[i, b]))
-                truths_by_band[b].append(float(target[i, b]))
-
-            seen += x.shape[0]
-            elapsed = time.time() - start
-            print(
-                f"  batch {batch_idx + 1}/{n_batches}  "
-                f"({seen}/{n_samples} samples, {elapsed:.1f}s)",
-                flush=True,
+        for step in range(n_future_steps):
+            target_idx = start_idx + context_len + step
+
+            if target_idx >= len(times):
+                break
+
+            win_v = ctx_v[-context_len:]
+            win_t = ctx_t[-context_len:]
+            win_b = ctx_b[-context_len:]
+
+            x = build_context_input(
+                win_v=win_v,
+                win_t=win_t,
+                win_b=win_b,
+                context_len=context_len,
+                n_bands=n_bands,
+                device=device,
             )
 
-    results = []
-    for band_idx in range(n_bands):
-        p = np.asarray(preds_by_band[band_idx], dtype=np.float32)
-        t = np.asarray(truths_by_band[band_idx], dtype=np.float32)
-        r = p - t
-
-        if len(r) > 0:
-            mse = float(np.mean(r**2))
-        else:
-            mse = np.nan
-
-        results.append(
-            {
-                "band_idx": band_idx,
-                "band_name": BAND_NAMES[band_idx],
-                "pred": p,
-                "truth": t,
-                "residual": r,
-                "n": len(r),
-                "mse": mse,
-            }
-        )
+            # Lead time from the last context event to the next true event.
+            Dt = torch.tensor(
+                [float(times[target_idx]) - win_t[-1]],
+                dtype=torch.float32,
+                device=device,
+            )
 
-    return results
+            pred_all = model(x, in_vars=None, out_vars=None, Dt=Dt)
+            pred_all = pred_all.reshape(n_bands).detach().cpu().numpy()
 
+            target_band = int(bands[target_idx])
 
-def plot_pred_vs_truth(results, outpath):
-    fig, axes = plt.subplots(3, 3, figsize=(12, 12))
-    axes = axes.flatten()
+            pred_norm = float(pred_all[target_band])
+            true_norm = float(values[target_idx])
+            residual = pred_norm - true_norm
 
-    for band_idx, res in enumerate(results):
-        ax = axes[band_idx]
+            pred_mag = pred_norm * (stds[target_band] + EPS) + means[target_band]
+            true_mag = true_norm * (stds[target_band] + EPS) + means[target_band]
 
-        if res["n"] == 0:
-            ax.set_title(f"{res['band_name']} (no samples)")
-            ax.axis("off")
-            continue
+            band_sq_err[target_band].append(residual**2)
 
-        ax.scatter(res["truth"], res["pred"], s=8, alpha=0.5)
+            steps.append(
+                {
+                    "step": step,
+                    "t_rel": float(times[target_idx]) - t_ref,
+                    "band": target_band,
+                    "pred_norm": pred_norm,
+                    "true_norm": true_norm,
+                    "pred_mag": float(pred_mag),
+                    "true_mag": float(true_mag),
+                    "residual": residual,
+                }
+            )
 
-        lo = float(min(res["truth"].min(), res["pred"].min()))
-        hi = float(max(res["truth"].max(), res["pred"].max()))
-        ax.plot([lo, hi], [lo, hi], linestyle="--", linewidth=1, color="k")
+            # Feed the prediction back in as the newest context event, following
+            # the true schedule (time + band) for the observation just forecast.
+            ctx_t.append(float(times[target_idx]))
+            ctx_v.append(pred_norm)
+            ctx_b.append(target_band)
 
-        ax.set_xlabel("truth (norm)")
-        ax.set_ylabel("pred (norm)")
-        ax.set_title(f"{res['band_name']}  (n={res['n']}, MSE={res['mse']:.3g})")
+    residuals = np.asarray([s["residual"] for s in steps], dtype=np.float32)
+    total_mse = float(np.mean(residuals**2)) if len(residuals) else np.nan
 
-    fig.suptitle("Next-event prediction vs truth (normalized), per band", y=1.01)
-    fig.tight_layout()
-    fig.savefig(outpath, dpi=200, bbox_inches="tight")
-    plt.close(fig)
+    band_mse = np.full(n_bands, np.nan, dtype=np.float32)
+    for b in range(n_bands):
+        if band_sq_err[b]:
+            band_mse[b] = float(np.mean(band_sq_err[b]))
+
+    return {
+        "start_idx": start_idx,
+        "context": context,
+        "steps": steps,
+        "mse": total_mse,
+        "band_mse": band_mse,
+        "n_steps": len(steps),
+    }
 
 
-def plot_residual_histograms(results, outpath):
-    fig, axes = plt.subplots(3, 3, figsize=(12, 12))
+def select_series(dataset, context_len, n_future_steps, n_series):
+    """Pick files with enough events for a rollout, longest first.
+
+    Returns a list of (times, values, bands, start_idx) tuples.
+    """
+    min_events = context_len + 1  # need at least one future step
+
+    eligible = []
+    for times, values, bands in dataset.events_per_file:
+        if len(times) >= min_events:
+            eligible.append((times, values, bands))
+
+    # Prefer the longest streams so rollouts have the most future steps.
+    eligible.sort(key=lambda tvb: len(tvb[0]), reverse=True)
+
+    selected = []
+    for times, values, bands in eligible[:n_series]:
+        # Start at the beginning; the rollout naturally stops at the end of the
+        # stream if fewer than n_future_steps events remain.
+        selected.append((times, values, bands, 0))
+
+    return selected
+
+
+def plot_series_lightcurves(rollout, means, stds, outpath):
+    """Plot one rollout's context, truth, and forecast per band (3x3 grid)."""
+    fig, axes = plt.subplots(3, 3, figsize=(14, 12))
     axes = axes.flatten()
 
-    for band_idx, res in enumerate(results):
+    context = rollout["context"]
+    steps = rollout["steps"]
+
+    for band_idx in range(N_BANDS):
         ax = axes[band_idx]
+        color = BAND_COLORS[band_idx]
+
+        def denorm(v):
+            return v * (stds[band_idx] + EPS) + means[band_idx]
+
+        # Context observations that belong to this band.
+        ctx_mask = context["band"] == band_idx
+        if np.any(ctx_mask):
+            ax.scatter(
+                context["t_rel"][ctx_mask],
+                denorm(context["v_norm"][ctx_mask]),
+                s=30,
+                color=color,
+                marker="o",
+                label="context",
+            )
 
-        if res["n"] == 0:
-            ax.set_title(f"{res['band_name']} (no samples)")
-            ax.axis("off")
-            continue
+        # Future truth and forecast for this band.
+        b_steps = [s for s in steps if s["band"] == band_idx]
+        if b_steps:
+            t = np.asarray([s["t_rel"] for s in b_steps])
+            true_mag = np.asarray([s["true_mag"] for s in b_steps])
+            pred_mag = np.asarray([s["pred_mag"] for s in b_steps])
 
-        ax.hist(res["residual"], bins=min(30, max(1, res["n"])))
-        ax.axvline(0.0, linestyle="--", linewidth=1, color="k")
-        ax.set_xlabel("pred - truth (norm)")
-        ax.set_ylabel("count")
-        ax.set_title(f"{res['band_name']}  (n={res['n']})")
+            order = np.argsort(t)
+            t = t[order]
+            true_mag = true_mag[order]
+            pred_mag = pred_mag[order]
 
-    fig.suptitle("Next-event residual distributions, per band", y=1.01)
+            ax.plot(
+                t, true_mag, "-o", color=color, alpha=0.8, label="truth"
+            )
+            ax.plot(
+                t,
+                pred_mag,
+                "--s",
+                color="k",
+                alpha=0.8,
+                markerfacecolor="none",
+                label="forecast",
+            )
+
+        ax.axvline(0.0, color="gray", linewidth=1, linestyle=":", alpha=0.7)
+        ax.invert_yaxis()
+        ax.set_xlabel("Relative time (days)")
+        ax.set_ylabel("Magnitude")
+        ax.set_title(BAND_NAMES[band_idx])
+
+        # Only add a legend if this band actually plotted labeled artists.
+        if ax.get_legend_handles_labels()[1]:
+            ax.legend(fontsize=7, loc="best")
+
+    fig.suptitle(
+        f"Autoregressive forecast (start_idx={rollout['start_idx']}, "
+        f"{rollout['n_steps']} steps)",
+        y=1.01,
+    )
     fig.tight_layout()
     fig.savefig(outpath, dpi=200, bbox_inches="tight")
     plt.close(fig)
 
 
-def plot_band_mse_bar(results, outpath):
-    names = [res["band_name"] for res in results]
-    mses = [res["mse"] if np.isfinite(res["mse"]) else 0.0 for res in results]
+def plot_residuals_vs_step(rollouts, outpath):
+    """Residual vs autoregressive step, colored by band, across all series."""
+    plt.figure(figsize=(10, 6))
+
+    for band_idx in range(N_BANDS):
+        xs = []
+        ys = []
+        for rollout in rollouts:
+            for s in rollout["steps"]:
+                if s["band"] == band_idx:
+                    xs.append(s["step"])
+                    ys.append(s["residual"])
+        if xs:
+            plt.scatter(
+                xs,
+                ys,
+                s=20,
+                alpha=0.6,
+                color=BAND_COLORS[band_idx],
+                label=BAND_NAMES[band_idx],
+            )
+
+    plt.axhline(0.0, linestyle="--", linewidth=1, color="k")
+    plt.xlabel("Autoregressive step")
+    plt.ylabel("pred - truth (normalized)")
+    plt.title("Autoregressive residuals vs step, by band")
+    plt.legend(fontsize=7, ncol=3, loc="best")
+    plt.tight_layout()
+    plt.savefig(outpath, dpi=200, bbox_inches="tight")
+    plt.close()
+
+
+def plot_band_mse_bar(rollouts, outpath):
+    """Per-band MSE aggregated over all rollout steps and series."""
+    band_sq = [[] for _ in range(N_BANDS)]
+    for rollout in rollouts:
+        for s in rollout["steps"]:
+            band_sq[s["band"]].append(s["residual"] ** 2)
+
+    mses = [float(np.mean(band_sq[b])) if band_sq[b] else 0.0 for b in range(N_BANDS)]
 
     plt.figure(figsize=(9, 5))
-    plt.bar(names, mses)
-    plt.ylabel("Next-event MSE (normalized)")
+    plt.bar(list(BAND_NAMES), mses, color=list(BAND_COLORS))
+    plt.ylabel("Autoregressive MSE (normalized)")
     plt.xlabel("Band")
-    plt.title("Per-band next-event MSE")
+    plt.title("Per-band autoregressive rollout MSE")
     plt.tight_layout()
     plt.savefig(outpath, dpi=200)
     plt.close()
 
 
-def save_band_mse_csv(results, outpath):
+def save_mse_csv(rollouts, outpath):
     with open(outpath, "w", newline="") as f:
         writer = csv.writer(f)
-        writer.writerow(["band_idx", "band_name", "n_samples", "mse"])
-        for res in results:
+        writer.writerow(
+            ["series", "start_idx", "n_steps", "mse_total"]
+            + [f"mse_{n}" for n in BAND_NAMES]
+        )
+        for i, rollout in enumerate(rollouts):
             writer.writerow(
-                [res["band_idx"], res["band_name"], res["n"], res["mse"]]
+                [i, rollout["start_idx"], rollout["n_steps"], rollout["mse"]]
+                + [rollout["band_mse"][b] for b in range(N_BANDS)]
             )
 
 
+def save_step_csv(rollouts, outpath):
+    with open(outpath, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(
+            [
+                "series",
+                "start_idx",
+                "step",
+                "t_rel",
+                "band_idx",
+                "band_name",
+                "pred_norm",
+                "true_norm",
+                "pred_mag",
+                "true_mag",
+                "residual",
+            ]
+        )
+        for i, rollout in enumerate(rollouts):
+            for s in rollout["steps"]:
+                writer.writerow(
+                    [
+                        i,
+                        rollout["start_idx"],
+                        s["step"],
+                        s["t_rel"],
+                        s["band"],
+                        BAND_NAMES[s["band"]],
+                        s["pred_norm"],
+                        s["true_norm"],
+                        s["pred_mag"],
+                        s["true_mag"],
+                        s["residual"],
+                    ]
+                )
+
+
 def main():
     args = get_args()
     run_id = resolve_paths(args)
@@ -375,45 +575,83 @@ def main():
 
     model, context_len, n_bands = load_9band_model(args.ckpt, device)
 
-    eval_dataset = make_eval_dataset(args=args, context_len=context_len)
-
-    print("Dataset length:", len(eval_dataset))
+    eval_dataset, means, stds = make_eval_dataset(
+        args=args,
+        context_len=context_len,
+    )
 
-    if len(eval_dataset) == 0:
-        raise RuntimeError("Empty eval dataset. Check data path and N_imgs.")
+    print("Dataset files with events:", len(eval_dataset.events_per_file))
 
-    results = collect_next_event_predictions(
+    series = select_series(
         dataset=eval_dataset,
-        model=model,
-        device=device,
-        n_bands=n_bands,
-        batch_size=args.batch_size,
-        max_samples=args.max_samples,
+        context_len=context_len,
+        n_future_steps=args.n_future_steps,
+        n_series=args.n_series,
     )
 
-    pred_truth_path = os.path.join(
-        args.outdir, f"study{run_id}_9band_next_event_pred_vs_truth.png"
-    )
-    resid_path = os.path.join(
-        args.outdir, f"study{run_id}_9band_next_event_residual_hist.png"
+    if not series:
+        raise RuntimeError(
+            "No eval series with enough events for a rollout. Check data path, "
+            "N_imgs, and context_len."
+        )
+
+    print(f"Rolling out {len(series)} series, up to {args.n_future_steps} "
+          f"steps each.")
+
+    rollouts = []
+    for i, (times, values, bands, start_idx) in enumerate(series):
+        rollout = get_rollout_from_stream(
+            times=times,
+            values=values,
+            bands=bands,
+            model=model,
+            device=device,
+            start_idx=start_idx,
+            n_future_steps=args.n_future_steps,
+            context_len=context_len,
+            n_bands=n_bands,
+            means=means,
+            stds=stds,
+        )
+        rollouts.append(rollout)
+        print(
+            f"  series {i + 1}/{len(series)}: {rollout['n_steps']} steps, "
+            f"mse={rollout['mse']:.4g}",
+            flush=True,
+        )
+
+    # Per-series forecast light curves.
+    for i, rollout in enumerate(rollouts):
+        series_path = os.path.join(
+            args.outdir,
+            f"study{run_id}_9band_autoreg_series{i:02d}.png",
+        )
+        plot_series_lightcurves(rollout, means, stds, series_path)
+
+    residual_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_autoreg_residuals_vs_step.png"
     )
     mse_bar_path = os.path.join(
-        args.outdir, f"study{run_id}_9band_next_event_band_mse.png"
+        args.outdir, f"study{run_id}_9band_autoreg_band_mse.png"
+    )
+    mse_csv_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_autoreg_mse_by_series.csv"
     )
-    csv_path = os.path.join(
-        args.outdir, f"study{run_id}_9band_next_event_band_mse.csv"
+    step_csv_path = os.path.join(
+        args.outdir, f"study{run_id}_9band_autoreg_step_predictions.csv"
     )
 
-    plot_pred_vs_truth(results, pred_truth_path)
-    plot_residual_histograms(results, resid_path)
-    plot_band_mse_bar(results, mse_bar_path)
-    save_band_mse_csv(results, csv_path)
+    plot_residuals_vs_step(rollouts, residual_path)
+    plot_band_mse_bar(rollouts, mse_bar_path)
+    save_mse_csv(rollouts, mse_csv_path)
+    save_step_csv(rollouts, step_csv_path)
 
     print("Saved:")
-    print(" ", pred_truth_path)
-    print(" ", resid_path)
+    print("  per-series light curves in", args.outdir)
+    print(" ", residual_path)
     print(" ", mse_bar_path)
-    print(" ", csv_path)
+    print(" ", mse_csv_path)
+    print(" ", step_csv_path)
 
 
 if __name__ == "__main__":

From 2074692210e5b067e73e7c9ebf0ad59a51a5ff3b Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 15:11:26 -0600
Subject: [PATCH 20/66] update plot script to use all bands

---
 .../plot_pred_diagnostics_9band.py            | 56 ++++++++++++++-----
 1 file changed, 43 insertions(+), 13 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index b7fbf967..6836a14c 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -20,6 +20,11 @@
 Because the observation schedule (times + which filter is seen next) is taken
 from the truth while the values are fed back from the model, this measures how
 well the model forecasts future observations in each filter over a rollout.
+
+At every step the model emits a prediction for ALL nine bands, including bands
+with no context and bands not observed at that step. The per-series plots show
+this full all-band forecast, overlaid with truth wherever an observation exists.
+Only the observed band is fed back into the context and scored against truth.
 """
 
 import argparse
@@ -318,6 +323,14 @@ def get_rollout_from_stream(
             pred_all = model(x, in_vars=None, out_vars=None, Dt=Dt)
             pred_all = pred_all.reshape(n_bands).detach().cpu().numpy()
 
+            # The model predicts every band at this lead time, including bands
+            # with no context and bands not observed at this step. Keep the full
+            # 9-band prediction (normalized and denormalized) for plotting.
+            pred_all_norm = pred_all.astype(np.float32)
+            pred_all_mag = (
+                pred_all_norm * (stds + EPS) + means
+            ).astype(np.float32)
+
             target_band = int(bands[target_idx])
 
             pred_norm = float(pred_all[target_band])
@@ -339,6 +352,9 @@ def get_rollout_from_stream(
                     "pred_mag": float(pred_mag),
                     "true_mag": float(true_mag),
                     "residual": residual,
+                    # Full all-band prediction at this step's lead time.
+                    "pred_all_norm": pred_all_norm,
+                    "pred_all_mag": pred_all_mag,
                 }
             )
 
@@ -417,29 +433,40 @@ def denorm(v):
                 label="context",
             )
 
-        # Future truth and forecast for this band.
+        # Full forecast for this band at EVERY rollout step, whether or not this
+        # band was observed at that step and whether or not it had any context.
+        if steps:
+            t_all = np.asarray([s["t_rel"] for s in steps])
+            pred_all = np.asarray(
+                [s["pred_all_mag"][band_idx] for s in steps]
+            )
+
+            order = np.argsort(t_all)
+            t_all = t_all[order]
+            pred_all = pred_all[order]
+
+            ax.plot(
+                t_all,
+                pred_all,
+                "--s",
+                color="k",
+                alpha=0.8,
+                markerfacecolor="none",
+                label="forecast (all steps)",
+            )
+
+        # Truth for this band, at the steps where it was actually observed.
         b_steps = [s for s in steps if s["band"] == band_idx]
         if b_steps:
             t = np.asarray([s["t_rel"] for s in b_steps])
             true_mag = np.asarray([s["true_mag"] for s in b_steps])
-            pred_mag = np.asarray([s["pred_mag"] for s in b_steps])
 
             order = np.argsort(t)
             t = t[order]
             true_mag = true_mag[order]
-            pred_mag = pred_mag[order]
 
             ax.plot(
-                t, true_mag, "-o", color=color, alpha=0.8, label="truth"
-            )
-            ax.plot(
-                t,
-                pred_mag,
-                "--s",
-                color="k",
-                alpha=0.8,
-                markerfacecolor="none",
-                label="forecast",
+                t, true_mag, "-o", color=color, alpha=0.8, label="truth (obs)"
             )
 
         ax.axvline(0.0, color="gray", linewidth=1, linestyle=":", alpha=0.7)
@@ -544,6 +571,8 @@ def save_step_csv(rollouts, outpath):
                 "true_mag",
                 "residual",
             ]
+            # Full all-band predicted magnitude at this step's lead time.
+            + [f"pred_mag_{n}" for n in BAND_NAMES]
         )
         for i, rollout in enumerate(rollouts):
             for s in rollout["steps"]:
@@ -561,6 +590,7 @@ def save_step_csv(rollouts, outpath):
                         s["true_mag"],
                         s["residual"],
                     ]
+                    + [float(s["pred_all_mag"][b]) for b in range(N_BANDS)]
                 )
 
 

From 8805402f52329469680f491fde2ffabb76ca0381 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 15:33:00 -0600
Subject: [PATCH 21/66] add teacher forcing for diagnostics

---
 .../plot_pred_diagnostics_9band.py            | 167 +++++++++++++++---
 1 file changed, 138 insertions(+), 29 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 6836a14c..64dc5467 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -110,6 +110,13 @@ def get_args():
         default=10,
         help="Number of light curves to roll out.",
     )
+    parser.add_argument(
+        "--teacher_forced",
+        action="store_true",
+        help="Also compute a teacher-forced rollout (true values fed back "
+        "instead of predictions) and overlay it on the free-running rollout to "
+        "diagnose compounding rollout error. Off by default.",
+    )
 
     parser.add_argument("--outdir", type=str, default=None)
     parser.add_argument(
@@ -260,6 +267,7 @@ def get_rollout_from_stream(
     n_bands,
     means,
     stds,
+    teacher_forced=False,
 ):
     """Autoregressively forecast the next events of one merged event stream.
 
@@ -272,6 +280,11 @@ def get_rollout_from_stream(
         context_len (int): Context window length.
         n_bands (int): Number of bands.
         means, stds (np.ndarray): Per-band normalization for denormalizing.
+        teacher_forced (bool): If True, feed the TRUE observed value back into
+            the context at each step instead of the model's own prediction. This
+            isolates one-step-ahead skill from compounding rollout error: a model
+            that tracks truth when teacher-forced but drifts when free-running is
+            suffering from exposure bias, not a failure to learn the dynamics.
 
     Returns:
         dict describing the rollout (context, per-step forecasts, per-band MSE).
@@ -358,10 +371,12 @@ def get_rollout_from_stream(
                 }
             )
 
-            # Feed the prediction back in as the newest context event, following
-            # the true schedule (time + band) for the observation just forecast.
+            # Feed the newest context event, following the true schedule (time +
+            # band) for the observation just forecast. Free-running rollout feeds
+            # the model's own prediction back in; teacher forcing feeds the true
+            # value instead, so errors do not compound down the rollout.
             ctx_t.append(float(times[target_idx]))
-            ctx_v.append(pred_norm)
+            ctx_v.append(true_norm if teacher_forced else pred_norm)
             ctx_b.append(target_band)
 
     residuals = np.asarray([s["residual"] for s in steps], dtype=np.float32)
@@ -406,13 +421,28 @@ def select_series(dataset, context_len, n_future_steps, n_series):
     return selected
 
 
-def plot_series_lightcurves(rollout, means, stds, outpath):
-    """Plot one rollout's context, truth, and forecast per band (3x3 grid)."""
+def _band_forecast_curve(steps, band_idx):
+    """Return (t_rel, pred_mag) for a band across all steps, sorted by time."""
+    t_all = np.asarray([s["t_rel"] for s in steps])
+    pred_all = np.asarray([s["pred_all_mag"][band_idx] for s in steps])
+    order = np.argsort(t_all)
+    return t_all[order], pred_all[order]
+
+
+def plot_series_lightcurves(rollout, means, stds, outpath, tf_rollout=None):
+    """Plot one rollout's context, truth, and forecast per band (3x3 grid).
+
+    If ``tf_rollout`` (the teacher-forced rollout for the same series) is given,
+    its forecast is overlaid so free-running vs teacher-forced can be compared
+    directly: divergence between the two indicates compounding rollout error
+    (exposure bias) rather than a failure to learn one-step dynamics.
+    """
     fig, axes = plt.subplots(3, 3, figsize=(14, 12))
     axes = axes.flatten()
 
     context = rollout["context"]
     steps = rollout["steps"]
+    tf_steps = tf_rollout["steps"] if tf_rollout is not None else None
 
     for band_idx in range(N_BANDS):
         ax = axes[band_idx]
@@ -433,18 +463,10 @@ def denorm(v):
                 label="context",
             )
 
-        # Full forecast for this band at EVERY rollout step, whether or not this
-        # band was observed at that step and whether or not it had any context.
+        # Free-running forecast for this band at EVERY rollout step, whether or
+        # not this band was observed and whether or not it had any context.
         if steps:
-            t_all = np.asarray([s["t_rel"] for s in steps])
-            pred_all = np.asarray(
-                [s["pred_all_mag"][band_idx] for s in steps]
-            )
-
-            order = np.argsort(t_all)
-            t_all = t_all[order]
-            pred_all = pred_all[order]
-
+            t_all, pred_all = _band_forecast_curve(steps, band_idx)
             ax.plot(
                 t_all,
                 pred_all,
@@ -452,7 +474,20 @@ def denorm(v):
                 color="k",
                 alpha=0.8,
                 markerfacecolor="none",
-                label="forecast (all steps)",
+                label="forecast (free-run)",
+            )
+
+        # Teacher-forced forecast for the same band and steps.
+        if tf_steps:
+            t_tf, pred_tf = _band_forecast_curve(tf_steps, band_idx)
+            ax.plot(
+                t_tf,
+                pred_tf,
+                "--^",
+                color="tab:purple",
+                alpha=0.8,
+                markerfacecolor="none",
+                label="forecast (teacher-forced)",
             )
 
         # Truth for this band, at the steps where it was actually observed.
@@ -521,17 +556,45 @@ def plot_residuals_vs_step(rollouts, outpath):
     plt.close()
 
 
-def plot_band_mse_bar(rollouts, outpath):
-    """Per-band MSE aggregated over all rollout steps and series."""
+def _per_band_mse(rollouts):
+    """Aggregate per-band MSE over all steps and series."""
     band_sq = [[] for _ in range(N_BANDS)]
     for rollout in rollouts:
         for s in rollout["steps"]:
             band_sq[s["band"]].append(s["residual"] ** 2)
+    return np.asarray(
+        [float(np.mean(band_sq[b])) if band_sq[b] else 0.0 for b in range(N_BANDS)]
+    )
 
-    mses = [float(np.mean(band_sq[b])) if band_sq[b] else 0.0 for b in range(N_BANDS)]
 
-    plt.figure(figsize=(9, 5))
-    plt.bar(list(BAND_NAMES), mses, color=list(BAND_COLORS))
+def plot_band_mse_bar(rollouts, outpath, tf_rollouts=None):
+    """Per-band MSE aggregated over all rollout steps and series.
+
+    If ``tf_rollouts`` is given, free-running and teacher-forced MSE are shown as
+    grouped bars per band. A large free-run bar next to a small teacher-forced
+    bar is the signature of compounding rollout error (exposure bias).
+    """
+    mses = _per_band_mse(rollouts)
+
+    plt.figure(figsize=(10, 5))
+    x = np.arange(N_BANDS)
+
+    if tf_rollouts is not None:
+        tf_mses = _per_band_mse(tf_rollouts)
+        width = 0.4
+        plt.bar(x - width / 2, mses, width, color="tab:gray", label="free-run")
+        plt.bar(
+            x + width / 2,
+            tf_mses,
+            width,
+            color="tab:purple",
+            label="teacher-forced",
+        )
+        plt.legend(fontsize=9)
+    else:
+        plt.bar(x, mses, color=list(BAND_COLORS))
+
+    plt.xticks(x, list(BAND_NAMES))
     plt.ylabel("Autoregressive MSE (normalized)")
     plt.xlabel("Band")
     plt.title("Per-band autoregressive rollout MSE")
@@ -628,7 +691,10 @@ def main():
     print(f"Rolling out {len(series)} series, up to {args.n_future_steps} "
           f"steps each.")
 
+    teacher_forced = args.teacher_forced
+
     rollouts = []
+    tf_rollouts = [] if teacher_forced else None
     for i, (times, values, bands, start_idx) in enumerate(series):
         rollout = get_rollout_from_stream(
             times=times,
@@ -642,13 +708,39 @@ def main():
             n_bands=n_bands,
             means=means,
             stds=stds,
+            teacher_forced=False,
         )
         rollouts.append(rollout)
-        print(
-            f"  series {i + 1}/{len(series)}: {rollout['n_steps']} steps, "
-            f"mse={rollout['mse']:.4g}",
-            flush=True,
-        )
+
+        if teacher_forced:
+            tf_rollout = get_rollout_from_stream(
+                times=times,
+                values=values,
+                bands=bands,
+                model=model,
+                device=device,
+                start_idx=start_idx,
+                n_future_steps=args.n_future_steps,
+                context_len=context_len,
+                n_bands=n_bands,
+                means=means,
+                stds=stds,
+                teacher_forced=True,
+            )
+            tf_rollouts.append(tf_rollout)
+
+            print(
+                f"  series {i + 1}/{len(series)}: {rollout['n_steps']} steps, "
+                f"free-run mse={rollout['mse']:.4g}, "
+                f"teacher-forced mse={tf_rollout['mse']:.4g}",
+                flush=True,
+            )
+        else:
+            print(
+                f"  series {i + 1}/{len(series)}: {rollout['n_steps']} steps, "
+                f"mse={rollout['mse']:.4g}",
+                flush=True,
+            )
 
     # Per-series forecast light curves.
     for i, rollout in enumerate(rollouts):
@@ -656,7 +748,10 @@ def main():
             args.outdir,
             f"study{run_id}_9band_autoreg_series{i:02d}.png",
         )
-        plot_series_lightcurves(rollout, means, stds, series_path)
+        tf_rollout = tf_rollouts[i] if teacher_forced else None
+        plot_series_lightcurves(
+            rollout, means, stds, series_path, tf_rollout=tf_rollout
+        )
 
     residual_path = os.path.join(
         args.outdir, f"study{run_id}_9band_autoreg_residuals_vs_step.png"
@@ -672,10 +767,24 @@ def main():
     )
 
     plot_residuals_vs_step(rollouts, residual_path)
-    plot_band_mse_bar(rollouts, mse_bar_path)
+    plot_band_mse_bar(rollouts, mse_bar_path, tf_rollouts=tf_rollouts)
     save_mse_csv(rollouts, mse_csv_path)
     save_step_csv(rollouts, step_csv_path)
 
+    if teacher_forced:
+        tf_mse_csv_path = os.path.join(
+            args.outdir, f"study{run_id}_9band_autoreg_mse_by_series_tf.csv"
+        )
+        save_mse_csv(tf_rollouts, tf_mse_csv_path)
+
+        overall_free = np.nanmean([r["mse"] for r in rollouts])
+        overall_tf = np.nanmean([r["mse"] for r in tf_rollouts])
+        print(
+            f"Overall free-run MSE: {overall_free:.4g}  |  "
+            f"teacher-forced MSE: {overall_tf:.4g}",
+            flush=True,
+        )
+
     print("Saved:")
     print("  per-series light curves in", args.outdir)
     print(" ", residual_path)

From 3a16a6697a26dd99d02afc939231d2981a699dac Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 12 Aug 2026 16:11:18 -0600
Subject: [PATCH 22/66] Improve training and add scheduled sampling

---
 .../KN_loderunner/train_LodeRunner_ddp.py     | 111 +++++--
 src/yoke/datasets/kilonova_dataset.py         | 112 ++++++-
 src/yoke/utils/training/epoch/loderunner.py   | 276 ++++++++++++++++++
 3 files changed, 480 insertions(+), 19 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 32180a12..c5750f4a 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -19,6 +19,7 @@
 )
 from yoke.utils.training.epoch.loderunner import (
     train_DDP_scalar_temporal_loderunner_epoch_9band,
+    train_DDP_scalar_temporal_loderunner_epoch_9band_rollout,
 )
 from yoke.utils.restart import continuation_setup
 from yoke.utils.dataload import make_distributed_dataloader
@@ -54,6 +55,37 @@
     help="Relative magnitude ε for Gaussian noise injection (e.g. 5e-5).",
 )
 
+# Multi-step rollout training (scheduled sampling) to address exposure bias.
+parser.add_argument(
+    "--n_rollout_steps",
+    type=int,
+    default=1,
+    help="Number of future events supervised per sample. 1 uses the standard "
+    "single-step teacher-forced training; >1 enables scheduled-sampling "
+    "rollout training.",
+)
+parser.add_argument(
+    "--tf_start",
+    type=float,
+    default=1.0,
+    help="Teacher-forcing ratio at epoch 0 (probability of feeding the true "
+    "value back at each rollout step). Only used when --n_rollout_steps > 1.",
+)
+parser.add_argument(
+    "--tf_end",
+    type=float,
+    default=0.0,
+    help="Teacher-forcing ratio the schedule anneals down to. Only used when "
+    "--n_rollout_steps > 1.",
+)
+parser.add_argument(
+    "--tf_ramp_epochs",
+    type=int,
+    default=50,
+    help="Number of epochs over which the teacher-forcing ratio decays linearly "
+    "from --tf_start to --tf_end. Only used when --n_rollout_steps > 1.",
+)
+
 # Change some default filepaths.
 parser.set_defaults(
     train_filelist="lsc240420_prefixes_train_80pct.txt",
@@ -141,6 +173,13 @@ def main(args, rank, world_size, local_rank, device):
     CONTEXT_LEN = 5 #3
     HIDDEN_CHANNELS = 64
 
+    # Multi-step rollout training config (scheduled sampling). n_rollout_steps=1
+    # falls back to the standard single-step teacher-forced training.
+    n_rollout_steps = args.n_rollout_steps
+    tf_start = args.tf_start
+    tf_end = args.tf_end
+    tf_ramp_epochs = max(1, args.tf_ramp_epochs)
+
     # Nine-band merged event-stream setup (3 ZTF + 6 Rubin/LSST bands).
     BAND_KEYS = NINE_BAND_KEYS
     VALUE_COL = 1
@@ -359,6 +398,7 @@ def main(args, rank, world_size, local_rank, device):
         drop_upper_limits=DROP_UPPER_LIMITS,
         means=band_means,
         stds=band_stds,
+        n_rollout_steps=n_rollout_steps,
     )
 
     val_dataset = Kilonova_lc_scalar_context_DataSet_9band(
@@ -369,6 +409,7 @@ def main(args, rank, world_size, local_rank, device):
         drop_upper_limits=DROP_UPPER_LIMITS,
         means=band_means,
         stds=band_stds,
+        n_rollout_steps=n_rollout_steps,
     )
 
 
@@ -415,24 +456,57 @@ def main(args, rank, world_size, local_rank, device):
             startTime = time.time()
 
 
-        #train_DDP_loderunner_epoch(
-        train_DDP_scalar_temporal_loderunner_epoch_9band(
-            training_data=train_dataloader,
-            validation_data=val_dataloader,
-            num_train_batches=train_batches,
-            num_val_batches=val_batches,
-            model=model,
-            optimizer=optimizer,
-            loss_fn=loss_fn,
-            LRsched=LRsched,
-            epochIDX=epochIDX,
-            train_per_val=train_per_val,
-            train_rcrd_filename=trn_rcrd_filename,
-            val_rcrd_filename=val_rcrd_filename,
-            device=device,
-            rank=rank,
-            world_size=world_size,
-        )
+        if n_rollout_steps > 1:
+            # Linearly anneal the teacher-forcing ratio from tf_start to tf_end
+            # over tf_ramp_epochs, then hold at tf_end.
+            frac = min(1.0, (epochIDX - 1) / tf_ramp_epochs)
+            teacher_forcing_ratio = tf_start + (tf_end - tf_start) * frac
+
+            if rank == 0:
+                print(
+                    f"Rollout training: n_rollout_steps={n_rollout_steps}, "
+                    f"teacher_forcing_ratio={teacher_forcing_ratio:.4f}",
+                    flush=True,
+                )
+
+            train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
+                training_data=train_dataloader,
+                validation_data=val_dataloader,
+                num_train_batches=train_batches,
+                num_val_batches=val_batches,
+                model=model,
+                optimizer=optimizer,
+                loss_fn=loss_fn,
+                LRsched=LRsched,
+                epochIDX=epochIDX,
+                train_per_val=train_per_val,
+                train_rcrd_filename=trn_rcrd_filename,
+                val_rcrd_filename=val_rcrd_filename,
+                device=device,
+                rank=rank,
+                world_size=world_size,
+                n_bands=N_BANDS,
+                teacher_forcing_ratio=teacher_forcing_ratio,
+            )
+        else:
+            #train_DDP_loderunner_epoch(
+            train_DDP_scalar_temporal_loderunner_epoch_9band(
+                training_data=train_dataloader,
+                validation_data=val_dataloader,
+                num_train_batches=train_batches,
+                num_val_batches=val_batches,
+                model=model,
+                optimizer=optimizer,
+                loss_fn=loss_fn,
+                LRsched=LRsched,
+                epochIDX=epochIDX,
+                train_per_val=train_per_val,
+                train_rcrd_filename=trn_rcrd_filename,
+                val_rcrd_filename=val_rcrd_filename,
+                device=device,
+                rank=rank,
+                world_size=world_size,
+            )
 
         print(f"[rank {rank}] finished epoch", flush=True)
 
@@ -477,6 +551,7 @@ def main(args, rank, world_size, local_rank, device):
                     "band_keys": list(BAND_KEYS),
                     "backbone_channels": 8,
                     "hidden": HIDDEN_CHANNELS,
+                    "n_rollout_steps": n_rollout_steps,
                 },
                 new_chkpt_path,
             )
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 623735a1..845f2b45 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -386,6 +386,7 @@ def __init__(
         drop_upper_limits: bool = True,
         means: np.ndarray = None,
         stds: np.ndarray = None,
+        n_rollout_steps: int = 1,
     ) -> None:
         """Initialize the dataset and build the merged-event sample index.
 
@@ -403,6 +404,12 @@ def __init__(
                 the model only sees real detections.
             means (np.ndarray): Per-band means for normalization, shape [n_bands].
             stds (np.ndarray): Per-band stds for normalization, shape [n_bands].
+            n_rollout_steps (int): Number of future events supervised per sample.
+                When 1 (default) ``__getitem__`` returns the single-step
+                ``(x, target, mask, Dt)`` tuple. When >1 it returns a rollout
+                tuple carrying the initial context window plus the next
+                ``n_rollout_steps`` true events, for scheduled-sampling /
+                multi-step rollout training (see ``__getitem__``).
         """
         # FIXME: hardcoded scratch path. Should be passed in as an argument so
         # this dataset does not depend on a user-specific filesystem location.
@@ -429,6 +436,12 @@ def __init__(
         self.drop_upper_limits = drop_upper_limits
         self.n_channels = len(self.band_keys)
 
+        if n_rollout_steps < 1:
+            raise ValueError(
+                f"n_rollout_steps must be >= 1, got {n_rollout_steps}"
+            )
+        self.n_rollout_steps = n_rollout_steps
+
         if means is None:
             raise ValueError(
                 "means must be provided for per-band normalization. "
@@ -529,7 +542,26 @@ def __len__(self) -> int:
         """Return the number of samples in the dataset."""
         return len(self.samples)
 
-    def __getitem__(
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
+        """Return the sample for ``index``.
+
+        When ``n_rollout_steps == 1`` this returns the single-step
+        ``(x, target, mask, Dt)`` tuple. When ``n_rollout_steps > 1`` it returns
+        the rollout tuple ``(ctx_v, ctx_t, ctx_b, future_v, future_b, future_dt,
+        future_valid)`` described in :meth:`_getitem_rollout`.
+
+        Args:
+            index (int): Sample index.
+
+        Returns:
+            tuple[torch.Tensor, ...]: Single-step or rollout sample.
+        """
+        if self.n_rollout_steps > 1:
+            return self._getitem_rollout(index)
+
+        return self._getitem_single(index)
+
+    def _getitem_single(
         self, index: int
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         """Return the (input, target, mask, Dt) tuple for a given sample index.
@@ -595,3 +627,81 @@ def __getitem__(
         )
 
         return x, target, mask, Dt
+
+    def _getitem_rollout(
+        self, index: int
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        """Return a multi-step rollout sample for scheduled-sampling training.
+
+        The initial context window is returned in a structured form so the
+        training loop can rebuild the model input after feeding predictions back
+        in, and the next ``n_rollout_steps`` true events are returned as
+        fixed-length arrays. When the stream ends before ``n_rollout_steps``
+        future events are available, the remaining steps are padded and flagged
+        invalid via ``future_valid`` so the default collate can still stack
+        samples and the loss can ignore the padding.
+
+        Args:
+            index (int): Sample index.
+
+        Returns:
+            ctx_v (torch.Tensor): Normalized context values, shape [context_len].
+            ctx_t (torch.Tensor): Relative context times (relative to the first
+                context event), shape [context_len].
+            ctx_b (torch.Tensor): Context band indices (long), shape [context_len].
+            future_v (torch.Tensor): Normalized true value of each future event,
+                shape [n_rollout_steps]; padded steps are 0.
+            future_b (torch.Tensor): Band index of each future event (long),
+                shape [n_rollout_steps]; padded steps are 0.
+            future_dt (torch.Tensor): Lead time from the previous event to each
+                future event, shape [n_rollout_steps]; padded steps are 0.
+            future_valid (torch.Tensor): 1.0 for real future events, 0.0 for
+                padded steps, shape [n_rollout_steps].
+        """
+        file_idx, startIDX = self.samples[index]
+        times, values, bands = self.events_per_file[file_idx]
+
+        target_start = startIDX + self.context_len
+
+        # Initial context window, relative to its own first event so no absolute
+        # MJD offset leaks in (matching the single-step encoding).
+        ctx_t = (
+            times[startIDX:target_start] - times[startIDX]
+        ).astype(np.float32)
+        ctx_v = values[startIDX:target_start].astype(np.float32)
+        ctx_b = bands[startIDX:target_start].astype(np.int64)
+
+        n = self.n_rollout_steps
+        future_v = np.zeros(n, dtype=np.float32)
+        future_b = np.zeros(n, dtype=np.int64)
+        future_dt = np.zeros(n, dtype=np.float32)
+        future_valid = np.zeros(n, dtype=np.float32)
+
+        n_events = times.shape[0]
+        for step in range(n):
+            target_idx = target_start + step
+            if target_idx >= n_events:
+                break
+
+            future_v[step] = values[target_idx]
+            future_b[step] = bands[target_idx]
+            future_dt[step] = times[target_idx] - times[target_idx - 1]
+            future_valid[step] = 1.0
+
+        return (
+            torch.tensor(ctx_v, dtype=torch.float32),
+            torch.tensor(ctx_t, dtype=torch.float32),
+            torch.tensor(ctx_b, dtype=torch.long),
+            torch.tensor(future_v, dtype=torch.float32),
+            torch.tensor(future_b, dtype=torch.long),
+            torch.tensor(future_dt, dtype=torch.float32),
+            torch.tensor(future_valid, dtype=torch.float32),
+        )
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index d6a19dff..7d843688 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -623,6 +623,282 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
                         np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
 
 
+def _rollout_pass_9band(
+    ctx_v: torch.Tensor,
+    ctx_t: torch.Tensor,
+    ctx_b: torch.Tensor,
+    future_v: torch.Tensor,
+    future_b: torch.Tensor,
+    future_dt: torch.Tensor,
+    future_valid: torch.Tensor,
+    model: torch.nn.Module,
+    loss_fn: torch.nn.Module,
+    n_bands: int,
+    teacher_forcing_ratio: float,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Unroll the 9-band model over a batch of rollouts with scheduled sampling.
+
+    Maintains a batched sliding context window seeded from the initial context.
+    At each step the model predicts all bands at the true lead time; the loss is
+    taken on the single observed band (padding steps ignored via
+    ``future_valid``). The value fed back into the context is the true value with
+    probability ``teacher_forcing_ratio`` and the model's own (detached)
+    prediction otherwise, so gradients never flow through the rollout.
+
+    Args:
+        ctx_v (torch.Tensor): Initial context values [B, context_len].
+        ctx_t (torch.Tensor): Initial context relative times [B, context_len].
+        ctx_b (torch.Tensor): Initial context band indices [B, context_len].
+        future_v (torch.Tensor): True future values [B, n_rollout_steps].
+        future_b (torch.Tensor): Future band indices [B, n_rollout_steps].
+        future_dt (torch.Tensor): Future lead times [B, n_rollout_steps].
+        future_valid (torch.Tensor): Valid-step mask [B, n_rollout_steps].
+        model (torch.nn.Module): The 9-band wrapper model.
+        loss_fn (torch.nn.Module): Elementwise loss (reduction='none').
+        n_bands (int): Number of bands.
+        teacher_forcing_ratio (float): Probability of feeding the true value back
+            at each step (1.0 = fully teacher-forced, 0.0 = fully free-running).
+        device (torch.device): Compute device.
+
+    Returns:
+        per_sample_loss (torch.Tensor): Mean rollout loss per sample [B].
+        total_loss (torch.Tensor): Scalar mean loss over all valid steps.
+    """
+    B, context_len = ctx_v.shape
+    n_steps = future_v.shape[1]
+
+    # Running window; cloned so we can slide in-place without touching inputs.
+    win_v = ctx_v.clone()
+    win_t = ctx_t.clone()
+    win_b = ctx_b.clone()
+
+    batch_arange = torch.arange(B, device=device)
+
+    # Kept for API compatibility with the LodeRunner-style wrapper.
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    step_losses = []  # [B] per valid step
+    step_valid = []  # [B] per step
+
+    for step in range(n_steps):
+        # Build the flattened per-event input from the current window, with
+        # times made relative to the window's first event (matching the dataset
+        # encoding and the inference rollout).
+        rel_t = win_t - win_t[:, :1]
+        band_onehot = torch.zeros(
+            B, context_len, n_bands, device=device, dtype=win_v.dtype
+        )
+        band_onehot.scatter_(2, win_b.unsqueeze(-1), 1.0)
+
+        per_event = torch.cat(
+            [win_v.unsqueeze(-1), rel_t.unsqueeze(-1), band_onehot],
+            dim=-1,
+        )  # [B, context_len, 2 + n_bands]
+        x_step = per_event.reshape(B, -1)
+
+        Dt = future_dt[:, step]
+        pred_all = model(x_step, in_vars, out_vars, Dt)  # [B, n_bands]
+
+        tgt_band = future_b[:, step]
+        pred_obs = pred_all[batch_arange, tgt_band]  # [B]
+        true_obs = future_v[:, step]  # [B]
+        valid = future_valid[:, step]  # [B]
+
+        step_loss = loss_fn(pred_obs, true_obs) * valid
+        step_losses.append(step_loss)
+        step_valid.append(valid)
+
+        # Scheduled sampling: choose true vs own (detached) prediction per sample.
+        use_true = (
+            torch.rand(B, device=device) < teacher_forcing_ratio
+        )
+        fed = torch.where(use_true, true_obs, pred_obs.detach())
+
+        # For padded steps there is no real event to advance to; feeding the true
+        # (zero) value with a zero dt is harmless since their loss is masked out
+        # and later steps are also padded/masked.
+        new_t = win_t[:, -1] + Dt
+
+        # Slide the window: drop the oldest event, append the new one.
+        win_v = torch.cat([win_v[:, 1:], fed.unsqueeze(1)], dim=1)
+        win_t = torch.cat([win_t[:, 1:], new_t.unsqueeze(1)], dim=1)
+        win_b = torch.cat([win_b[:, 1:], tgt_band.unsqueeze(1)], dim=1)
+
+    step_losses = torch.stack(step_losses, dim=1)  # [B, n_steps]
+    step_valid = torch.stack(step_valid, dim=1)  # [B, n_steps]
+
+    per_sample_loss = step_losses.sum(dim=1) / (step_valid.sum(dim=1) + 1e-8)
+    total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
+
+    return per_sample_loss, total_loss
+
+
+def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
+    training_data: torch.utils.data.DataLoader,
+    validation_data: torch.utils.data.DataLoader,
+    num_train_batches: int,
+    num_val_batches: int,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: torch.nn.Module,
+    LRsched: torch.optim.lr_scheduler._LRScheduler,
+    epochIDX: int,
+    train_per_val: int,
+    train_rcrd_filename: str,
+    val_rcrd_filename: str,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+    n_bands: int = 9,
+    teacher_forcing_ratio: float = 1.0,
+) -> None:
+    """Multi-step rollout DDP epoch for the masked 9-band scalar temporal model.
+
+    Trains the model on autoregressive rollouts with scheduled sampling to
+    address exposure bias: instead of a single teacher-forced next-event step,
+    the model unrolls ``n_rollout_steps`` events, feeding its own (detached)
+    predictions back into the context with probability
+    ``1 - teacher_forcing_ratio`` at each step. Validation always uses a fully
+    free-running rollout (ratio 0.0) so the recorded metric reflects real
+    rollout skill.
+
+    Expected dataset output (per sample, from
+    ``Kilonova_lc_scalar_context_DataSet_9band`` with ``n_rollout_steps > 1``):
+        ctx_v, ctx_t, ctx_b:  [B, context_len]
+        future_v, future_b, future_dt, future_valid:  [B, n_rollout_steps]
+
+    Args:
+        n_bands (int): Number of bands the model predicts.
+        teacher_forcing_ratio (float): Per-epoch probability of feeding the true
+            value back at each rollout step during training.
+    """
+    train_rcrd_filename = train_rcrd_filename.replace(
+        "<epochIDX>",
+        f"{epochIDX:04d}",
+    )
+
+    model.train()
+
+    with (
+        open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
+    ) as train_rcrd_file:
+
+        for trainbatch_ID, data in enumerate(training_data):
+            if trainbatch_ID >= num_train_batches:
+                break
+
+            ctx_v, ctx_t, ctx_b, future_v, future_b, future_dt, future_valid = (
+                data
+            )
+
+            ctx_v = ctx_v.to(device, non_blocking=True)
+            ctx_t = ctx_t.to(device, non_blocking=True)
+            ctx_b = ctx_b.to(device, non_blocking=True)
+            future_v = future_v.to(device, non_blocking=True)
+            future_b = future_b.to(device, non_blocking=True)
+            future_dt = future_dt.to(torch.float32).to(device, non_blocking=True)
+            future_valid = future_valid.to(device, non_blocking=True)
+
+            optimizer.zero_grad(set_to_none=True)
+
+            per_sample_loss, batch_loss = _rollout_pass_9band(
+                ctx_v=ctx_v,
+                ctx_t=ctx_t,
+                ctx_b=ctx_b,
+                future_v=future_v,
+                future_b=future_b,
+                future_dt=future_dt,
+                future_valid=future_valid,
+                model=model,
+                loss_fn=loss_fn,
+                n_bands=n_bands,
+                teacher_forcing_ratio=teacher_forcing_ratio,
+                device=device,
+            )
+
+            batch_loss.backward()
+            optimizer.step()
+            LRsched.step()
+
+            if rank == 0:
+                batch_records = np.column_stack(
+                    [
+                        np.full(len(per_sample_loss), epochIDX),
+                        np.full(len(per_sample_loss), trainbatch_ID),
+                        per_sample_loss.detach().cpu().numpy().flatten(),
+                    ]
+                )
+                np.savetxt(train_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+    if epochIDX % train_per_val == 0:
+        if rank == 0:
+            print("Validating...", epochIDX, flush=True)
+
+        val_rcrd_filename = val_rcrd_filename.replace(
+            "<epochIDX>",
+            f"{epochIDX:04d}",
+        )
+
+        model.eval()
+
+        with (
+            open(val_rcrd_filename, "a") if rank == 0 else nullcontext()
+        ) as val_rcrd_file:
+
+            with torch.no_grad():
+                for valbatch_ID, data in enumerate(validation_data):
+                    if valbatch_ID >= num_val_batches:
+                        break
+
+                    (
+                        ctx_v,
+                        ctx_t,
+                        ctx_b,
+                        future_v,
+                        future_b,
+                        future_dt,
+                        future_valid,
+                    ) = data
+
+                    ctx_v = ctx_v.to(device, non_blocking=True)
+                    ctx_t = ctx_t.to(device, non_blocking=True)
+                    ctx_b = ctx_b.to(device, non_blocking=True)
+                    future_v = future_v.to(device, non_blocking=True)
+                    future_b = future_b.to(device, non_blocking=True)
+                    future_dt = future_dt.to(torch.float32).to(
+                        device, non_blocking=True
+                    )
+                    future_valid = future_valid.to(device, non_blocking=True)
+
+                    # Validation is always a pure free-running rollout.
+                    per_sample_loss, _ = _rollout_pass_9band(
+                        ctx_v=ctx_v,
+                        ctx_t=ctx_t,
+                        ctx_b=ctx_b,
+                        future_v=future_v,
+                        future_b=future_b,
+                        future_dt=future_dt,
+                        future_valid=future_valid,
+                        model=model,
+                        loss_fn=loss_fn,
+                        n_bands=n_bands,
+                        teacher_forcing_ratio=0.0,
+                        device=device,
+                    )
+
+                    if rank == 0:
+                        batch_records = np.column_stack(
+                            [
+                                np.full(len(per_sample_loss), epochIDX),
+                                np.full(len(per_sample_loss), valbatch_ID),
+                                per_sample_loss.detach().cpu().numpy().flatten(),
+                            ]
+                        )
+                        np.savetxt(val_rcrd_file, batch_records, fmt="%d, %d, %.8f")
+
+
 def train_DDP_loderunner_epoch(
     training_data: torch.utils.data.DataLoader,
     validation_data: torch.utils.data.DataLoader,

From 5826e370a2814b25f5972f88a1952f937d296ec1 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 13 Aug 2026 07:36:35 -0600
Subject: [PATCH 23/66] minor continuation fix

---
 .../KN_loderunner/train_LodeRunner_ddp.py     | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index c5750f4a..86e3be28 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -85,6 +85,16 @@
     help="Number of epochs over which the teacher-forcing ratio decays linearly "
     "from --tf_start to --tf_end. Only used when --n_rollout_steps > 1.",
 )
+parser.add_argument(
+    "--tf_ramp_start_epoch",
+    type=int,
+    default=0,
+    help="Absolute epoch at which the teacher-forcing anneal begins. Because "
+    "epoch numbering continues across restarts, set this to the checkpoint "
+    "epoch when starting rollout training as a continuation (e.g. 50 when "
+    "continuing from epoch 50) so the ramp is measured from there rather than "
+    "from epoch 0. Only used when --n_rollout_steps > 1.",
+)
 
 # Change some default filepaths.
 parser.set_defaults(
@@ -179,6 +189,7 @@ def main(args, rank, world_size, local_rank, device):
     tf_start = args.tf_start
     tf_end = args.tf_end
     tf_ramp_epochs = max(1, args.tf_ramp_epochs)
+    tf_ramp_start_epoch = args.tf_ramp_start_epoch
 
     # Nine-band merged event-stream setup (3 ZTF + 6 Rubin/LSST bands).
     BAND_KEYS = NINE_BAND_KEYS
@@ -458,8 +469,16 @@ def main(args, rank, world_size, local_rank, device):
 
         if n_rollout_steps > 1:
             # Linearly anneal the teacher-forcing ratio from tf_start to tf_end
-            # over tf_ramp_epochs, then hold at tf_end.
-            frac = min(1.0, (epochIDX - 1) / tf_ramp_epochs)
+            # over tf_ramp_epochs, then hold at tf_end. The ramp is measured from
+            # tf_ramp_start_epoch so it works correctly on a continuation, where
+            # epoch numbering carries over from the previous run.
+            frac = max(
+                0.0,
+                min(
+                    1.0,
+                    (epochIDX - 1 - tf_ramp_start_epoch) / tf_ramp_epochs,
+                ),
+            )
             teacher_forcing_ratio = tf_start + (tf_end - tf_start) * frac
 
             if rank == 0:

From a23cabc5a8a6047d10d2a7801b978e0204c85ec7 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 13 Aug 2026 13:30:53 -0600
Subject: [PATCH 24/66] Scheduled sampling fixes

---
 .../plot_pred_diagnostics_9band.py            | 94 ++++++++++++++++---
 .../KN_loderunner/train_LodeRunner_ddp.py     | 60 ++++++------
 .../KN_loderunner/training_START.input        | 12 ++-
 .../KN_loderunner/training_input.tmpl         | 12 ++-
 4 files changed, 133 insertions(+), 45 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 64dc5467..0843a223 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -387,6 +387,49 @@ def get_rollout_from_stream(
         if band_sq_err[b]:
             band_mse[b] = float(np.mean(band_sq_err[b]))
 
+    # Smooth "forecast from now": hold the initial context window fixed and sweep
+    # a monotonic lead-time grid, predicting all bands at each lead time. This is
+    # the physically meaningful forecast (matches infer_9band.py) and, unlike the
+    # per-step autoregressive predictions, does not sawtooth, because the lead
+    # time increases monotonically instead of oscillating with the true event
+    # cadence. Computed only for the free-run rollout (teacher forcing has no
+    # meaning without feedback).
+    fixed_forecast = None
+    if not teacher_forced and steps:
+        win_t0 = times[start_idx : start_idx + context_len].astype(np.float32)
+        win_v0 = values[start_idx : start_idx + context_len].astype(np.float32)
+        win_b0 = bands[start_idx : start_idx + context_len].astype(np.int64)
+
+        x0 = build_context_input(
+            win_v=win_v0,
+            win_t=win_t0,
+            win_b=win_b0,
+            context_len=context_len,
+            n_bands=n_bands,
+            device=device,
+        )
+
+        last_ctx_t_rel = float(win_t0[-1]) - t_ref
+        max_step_t_rel = max(s["t_rel"] for s in steps)
+        horizon = max(1e-3, max_step_t_rel - last_ctx_t_rel)
+
+        n_lead = 60
+        lead_times = np.linspace(0.0, horizon, n_lead).astype(np.float32)
+        preds_norm = np.zeros((n_lead, n_bands), dtype=np.float32)
+
+        with torch.no_grad():
+            for k, dt in enumerate(lead_times):
+                Dt = torch.tensor([dt], dtype=torch.float32, device=device)
+                pred = model(x0, in_vars=None, out_vars=None, Dt=Dt)
+                preds_norm[k] = pred.reshape(n_bands).detach().cpu().numpy()
+
+        preds_mag = preds_norm * (stds[None, :] + EPS) + means[None, :]
+
+        fixed_forecast = {
+            "t_rel": last_ctx_t_rel + lead_times,
+            "preds_mag": preds_mag.astype(np.float32),
+        }
+
     return {
         "start_idx": start_idx,
         "context": context,
@@ -394,6 +437,7 @@ def get_rollout_from_stream(
         "mse": total_mse,
         "band_mse": band_mse,
         "n_steps": len(steps),
+        "fixed_forecast": fixed_forecast,
     }
 
 
@@ -443,6 +487,7 @@ def plot_series_lightcurves(rollout, means, stds, outpath, tf_rollout=None):
     context = rollout["context"]
     steps = rollout["steps"]
     tf_steps = tf_rollout["steps"] if tf_rollout is not None else None
+    fixed_forecast = rollout.get("fixed_forecast")
 
     for band_idx in range(N_BANDS):
         ax = axes[band_idx]
@@ -463,31 +508,50 @@ def denorm(v):
                 label="context",
             )
 
-        # Free-running forecast for this band at EVERY rollout step, whether or
-        # not this band was observed and whether or not it had any context.
+        # Headline forecast: smooth "forecast from now" over a monotonic lead
+        # time grid, holding the initial context fixed. This replaces the old
+        # per-step connected curve, whose sawtooth was an artifact of the lead
+        # time Dt oscillating with the true (nightly-clustered) event cadence
+        # rather than any model instability.
+        if fixed_forecast is not None:
+            ax.plot(
+                fixed_forecast["t_rel"],
+                fixed_forecast["preds_mag"][:, band_idx],
+                "-",
+                color="k",
+                linewidth=1.8,
+                alpha=0.85,
+                label="forecast (fixed context)",
+            )
+
+        # Per-step autoregressive predictions as UNCONNECTED markers. Each is
+        # made at that step's own lead time, so connecting them produces the
+        # sawtooth; shown as points they still reveal free-run bias vs truth
+        # without the misleading zigzag line.
         if steps:
             t_all, pred_all = _band_forecast_curve(steps, band_idx)
-            ax.plot(
+            ax.scatter(
                 t_all,
                 pred_all,
-                "--s",
-                color="k",
-                alpha=0.8,
-                markerfacecolor="none",
-                label="forecast (free-run)",
+                s=22,
+                facecolors="none",
+                edgecolors="k",
+                alpha=0.6,
+                label="per-step pred (free-run)",
             )
 
-        # Teacher-forced forecast for the same band and steps.
+        # Teacher-forced per-step predictions, also as unconnected markers.
         if tf_steps:
             t_tf, pred_tf = _band_forecast_curve(tf_steps, band_idx)
-            ax.plot(
+            ax.scatter(
                 t_tf,
                 pred_tf,
-                "--^",
-                color="tab:purple",
-                alpha=0.8,
-                markerfacecolor="none",
-                label="forecast (teacher-forced)",
+                s=22,
+                marker="^",
+                facecolors="none",
+                edgecolors="tab:purple",
+                alpha=0.6,
+                label="per-step pred (teacher-forced)",
             )
 
         # Truth for this band, at the steps where it was actually observed.
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 86e3be28..ec25c7ec 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -266,34 +266,38 @@ def main(args, rank, world_size, local_rank, device):
             print("Missing keys:", missing_keys)
             print("Unexpected keys:", unexpected_keys)
 
-            model.noise_scale = noise_scale
-
-            backbone = model
-
-            model = ScalarTemporalConditionedLodeRunner_9band(
-                backbone=backbone,
-                context_len=CONTEXT_LEN,
-                n_bands=N_BANDS,
-                image_size=model_args["image_size"],
-                backbone_channels=8,
-                hidden=HIDDEN_CHANNELS,
-            ).to(device)
-
-            # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
-            for p in model.backbone.parameters():
-                p.requires_grad = False
-
-            for p in model.conditioner.parameters():
-                p.requires_grad = True
-
-            for p in model.output_head.parameters():
-                p.requires_grad = True
-
-            optimizer = torch.optim.AdamW(
-                list(model.conditioner.parameters()) +
-                list(model.output_head.parameters()),
-                **optimizer_kwargs,
-            )
+        # NOTE: model reconstruction and optimizer creation must run on ALL
+        # ranks. If gated behind `if rank == 0:` the non-zero ranks keep the bare
+        # LodeRunner and never define `optimizer`, which crashes DDP wrapping /
+        # the LR scheduler on multi-GPU runs.
+        model.noise_scale = noise_scale
+
+        backbone = model
+
+        model = ScalarTemporalConditionedLodeRunner_9band(
+            backbone=backbone,
+            context_len=CONTEXT_LEN,
+            n_bands=N_BANDS,
+            image_size=model_args["image_size"],
+            backbone_channels=8,
+            hidden=HIDDEN_CHANNELS,
+        ).to(device)
+
+        # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
+        for p in model.backbone.parameters():
+            p.requires_grad = False
+
+        for p in model.conditioner.parameters():
+            p.requires_grad = True
+
+        for p in model.output_head.parameters():
+            p.requires_grad = True
+
+        optimizer = torch.optim.AdamW(
+            list(model.conditioner.parameters()) +
+            list(model.output_head.parameters()),
+            **optimizer_kwargs,
+        )
 
     #loss_fn = nn.MSELoss(reduction="none")
     loss_fn = nn.HuberLoss(delta=0.1, reduction="none")
diff --git a/applications/harnesses/KN_loderunner/training_START.input b/applications/harnesses/KN_loderunner/training_START.input
index 81e87bb2..c477f515 100644
--- a/applications/harnesses/KN_loderunner/training_START.input
+++ b/applications/harnesses/KN_loderunner/training_START.input
@@ -27,6 +27,16 @@ lsc240420_prefixes_validation_10pct.txt
 <WARMUP_STEPS>
 --noise_scale
 <NOISE_SCALE>
+--n_rollout_steps
+5
+--tf_start
+1.0
+--tf_end
+0.0
+--tf_ramp_start_epoch
+20
+--tf_ramp_epochs
+20
 --trn_rcrd_filename
 ./training_study<studyIDX>_epoch<epochIDX>.csv
 --val_rcrd_filename
@@ -40,7 +50,7 @@ lsc240420_prefixes_validation_10pct.txt
 --Knodes
 <KNODES>
 --total_epochs
-500
+90
 --cycle_epochs
 1
 --train_batches
diff --git a/applications/harnesses/KN_loderunner/training_input.tmpl b/applications/harnesses/KN_loderunner/training_input.tmpl
index 03793f7f..4446f897 100644
--- a/applications/harnesses/KN_loderunner/training_input.tmpl
+++ b/applications/harnesses/KN_loderunner/training_input.tmpl
@@ -27,6 +27,16 @@ lsc240420_prefixes_validation_10pct.txt
 <WARMUP_STEPS>
 --noise_scale
 <NOISE_SCALE>
+--n_rollout_steps
+5
+--tf_start
+1.0
+--tf_end
+0.0
+--tf_ramp_start_epoch
+20
+--tf_ramp_epochs
+20
 --trn_rcrd_filename
 ./training_study<studyIDX>_epoch<epochIDX>.csv
 --val_rcrd_filename
@@ -40,7 +50,7 @@ lsc240420_prefixes_validation_10pct.txt
 --Knodes
 <KNODES>
 --total_epochs
-500
+90
 --cycle_epochs
 1
 --train_batches

From 7726a9ac238945691ec6979b285e5fe7f1902688 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 13 Aug 2026 13:35:17 -0600
Subject: [PATCH 25/66] loss plot show sheduled sampling regimes

---
 .../KN_loderunner/plot_loss_curves_9band.py   | 31 +++++++++
 .../KN_loderunner/plot_loss_curves_gri.py     | 64 +++++++++++++++++++
 2 files changed, 95 insertions(+)

diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
index 5a9cf361..0aa3e3b6 100644
--- a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
@@ -69,6 +69,37 @@ def main():
         ),
     )
 
+    # Scheduled-sampling regime shading. 9-band runs anneal the teacher-forcing
+    # ratio, so shade the warmup / anneal / free-run regimes by default. These
+    # must match the schedule in training_input.tmpl / training_START.input.
+    parser.add_argument(
+        "--shade_regimes",
+        dest="shade_regimes",
+        action="store_true",
+        default=True,
+        help="Shade the teacher-forcing warmup/anneal/free-run regimes. Default.",
+    )
+    parser.add_argument(
+        "--no_shade_regimes",
+        dest="shade_regimes",
+        action="store_false",
+        help="Disable teacher-forcing regime shading.",
+    )
+    parser.add_argument(
+        "--tf_ramp_start_epoch",
+        type=int,
+        default=20,
+        help="Absolute epoch at which the teacher-forcing anneal begins. Must "
+        "match the training schedule. Default 20.",
+    )
+    parser.add_argument(
+        "--tf_ramp_epochs",
+        type=int,
+        default=20,
+        help="Number of epochs the teacher-forcing ratio anneals over. Must "
+        "match the training schedule. Default 20.",
+    )
+
     args = parser.parse_args()
 
     defaults = base.default_patterns(args.study, args.runs_root)
diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
index 4d97aabb..568ead70 100644
--- a/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
@@ -124,6 +124,60 @@ def infer_loss_labels(loss_names, n_loss_cols):
     return loss_names
 
 
+def shade_tf_regimes(ax, ramp_start, ramp_epochs):
+    """Shade the scheduled-sampling teacher-forcing regimes on a loss axis.
+
+    The teacher-forcing ratio p anneals from 1.0 (fully teacher-forced) down to
+    0.0 (fully free-running) over training. This splits the run into three
+    regimes, defined entirely by ``ramp_start`` and ``ramp_epochs``:
+
+      - Warmup   [xmin, ramp_start]:              p = 1.0, equivalent to
+                                                  single-step training.
+      - Anneal   [ramp_start, ramp_start+epochs]: p ramps 1.0 -> 0.0.
+      - Free-run [ramp_start+epochs, xmax]:       p = 0.0, matches inference.
+
+    Only these training regimes are shaded; the validation curve is always pure
+    free-run regardless of regime, so a val bump at the anneal onset is expected.
+    """
+    xmin, xmax = ax.get_xlim()
+    anneal_end = ramp_start + ramp_epochs
+
+    # Clip regime edges to the visible epoch range so a partial run still shades
+    # sensibly (e.g. a plot that only reaches into the anneal phase).
+    warmup_lo, warmup_hi = xmin, min(ramp_start, xmax)
+    anneal_lo, anneal_hi = max(ramp_start, xmin), min(anneal_end, xmax)
+    free_lo, free_hi = max(anneal_end, xmin), xmax
+
+    regimes = [
+        (warmup_lo, warmup_hi, "tab:blue", "warmup (p=1)"),
+        (anneal_lo, anneal_hi, "tab:orange", "anneal (p:1→0)"),
+        (free_lo, free_hi, "tab:green", "free-run (p=0)"),
+    ]
+
+    for lo, hi, color, label in regimes:
+        if hi <= lo:
+            continue
+        ax.axvspan(lo, hi, color=color, alpha=0.08, zorder=0)
+        # Place the label near the top of the axis, centered in the band.
+        ax.text(
+            0.5 * (lo + hi),
+            0.97,
+            label,
+            transform=ax.get_xaxis_transform(),
+            ha="center",
+            va="top",
+            fontsize=8,
+            color=color,
+            alpha=0.9,
+        )
+
+    for boundary in (ramp_start, anneal_end):
+        if xmin < boundary < xmax:
+            ax.axvline(
+                boundary, color="gray", linewidth=1, linestyle=":", alpha=0.6
+            )
+
+
 def plot_epoch_curves(train, val, args):
     train_ep, train_mean, train_std = epoch_stats(
         train["epochs"],
@@ -194,6 +248,16 @@ def plot_epoch_curves(train, val, args):
     if args.logy:
         plt.yscale("log")
 
+    # Optionally shade the scheduled-sampling teacher-forcing regimes. Gated on
+    # attributes that only the 9-band wrapper sets, so g/r/i runs (and any caller
+    # that doesn't opt in) get the plain plot unchanged.
+    if getattr(args, "shade_regimes", False):
+        shade_tf_regimes(
+            plt.gca(),
+            ramp_start=getattr(args, "tf_ramp_start_epoch", 20),
+            ramp_epochs=getattr(args, "tf_ramp_epochs", 20),
+        )
+
     plt.tight_layout()
     plt.savefig(args.out, dpi=args.dpi)
     print(f"Saved {args.out}")

From 175c960d5ee3e36ad07c195d073765a21d9be391 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 13 Aug 2026 16:35:58 -0600
Subject: [PATCH 26/66] dataset fix

---
 applications/harnesses/KN_loderunner/ddp_production.csv | 5 +++++
 src/yoke/datasets/kilonova_dataset.py                   | 7 ++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
index a6695405..ffbbd96a 100644
--- a/applications/harnesses/KN_loderunner/ddp_production.csv
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -25,5 +25,10 @@ studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZ
 #22,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
 #23,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
 24,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+# study 25: 4 GPUs for ~4x faster wall-clock. NTRN_BATCH 1000->250 and
+# NVAL_BATCH 500->125 keep total samples/epoch the same as study 24
+# (4 ranks x 250 = 1000). TERMINAL/WARMUP steps scaled 4x down (are counted
+# in per-rank batches) so the LR schedule spans the same number of epochs.
+25,yoke311,1,4,128,1,1,9,1,2,10,250,125,2.0e-3,0.5,0.5,250,125,0.0,train_LodeRunner_ddp.py
 
 
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 845f2b45..d7b350e9 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -413,10 +413,15 @@ def __init__(
         """
         # FIXME: hardcoded scratch path. Should be passed in as an argument so
         # this dataset does not depend on a user-specific filesystem location.
+        # NOTE: must point at the SAME dataset as
+        # load_or_compute_band_normalization (the rubin_ztf_10000 set). The old
+        # uniform_dataset_20000 set is ZTF-only, so training on it left the six
+        # Rubin output heads without any targets (never trained) while the norm
+        # stats were computed over Rubin+ZTF -- a silent train/stats mismatch.
         file_prefix_list = sorted(
             glob.glob(
                 "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-                "uniform_dataset_20000/lc_*.npz"
+                "rubin_ztf_10000_dataset/lc_*.npz"
             )
         )
 

From 4f39c71a50244be8a7f57c9e8dec520c82463880 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 14 Aug 2026 11:58:29 -0600
Subject: [PATCH 27/66] add histogram plot

---
 .../KN_loderunner/ddp_production.csv          |  12 +-
 .../plot_observation_histograms.py            | 284 ++++++++++++++++++
 .../KN_loderunner/training_slurm.tmpl         |  11 +-
 3 files changed, 300 insertions(+), 7 deletions(-)
 create mode 100644 applications/harnesses/KN_loderunner/plot_observation_histograms.py

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
index ffbbd96a..8196bfb2 100644
--- a/applications/harnesses/KN_loderunner/ddp_production.csv
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -25,10 +25,12 @@ studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZ
 #22,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
 #23,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
 24,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
-# study 25: 4 GPUs for ~4x faster wall-clock. NTRN_BATCH 1000->250 and
-# NVAL_BATCH 500->125 keep total samples/epoch the same as study 24
-# (4 ranks x 250 = 1000). TERMINAL/WARMUP steps scaled 4x down (are counted
-# in per-rank batches) so the LR schedule spans the same number of epochs.
-25,yoke311,1,4,128,1,1,9,1,2,10,250,125,2.0e-3,0.5,0.5,250,125,0.0,train_LodeRunner_ddp.py
+# study 25: full 8-GPU node for ~8x faster wall-clock than 1 GPU. NTRN_BATCH
+# 1000->125 and NVAL_BATCH 500->63 keep total samples/epoch the same as study 24
+# (8 ranks x 125 = 1000). TERMINAL/WARMUP steps scaled ~8x down (counted in
+# per-rank batches) so the LR schedule spans the same number of epochs. These
+# are 8-GPU H100 nodes (se*); cpus-per-task=24 in the slurm template fits
+# 24*8=192 cores.
+25,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
 
 
diff --git a/applications/harnesses/KN_loderunner/plot_observation_histograms.py b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
new file mode 100644
index 00000000..0f8289c6
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
@@ -0,0 +1,284 @@
+"""Plot observation-count histograms for the 9-band kilonova light-curve data.
+
+Each ``lc_*.npz`` file holds one light curve as a set of per-band arrays keyed by
+``NINE_BAND_KEYS`` (``arr_ztfg`` ... ``arr_ps1__y``). Every row of a band array is
+one observation with columns ``[MJD, value, error, ...]``. Matching the 9-band
+dataset (``Kilonova_lc_scalar_context_DataSet_9band`` with
+``drop_upper_limits=True``), a row is counted as a real **detection** when its
+error column (col 2) is finite; a non-finite error flags an upper limit /
+non-detection. This script summarises the data set as histograms:
+
+  1. Total detections per band (bar chart) -- how much supervision each of the
+     nine output heads actually gets.
+  2. Total observations per band split into detections vs upper limits, when
+     ``--include_upper_limits`` is set.
+  3. Distribution of detections-per-light-curve, per band and summed over all
+     bands -- how long/rich a typical object is.
+
+Run directly, e.g.:
+    python plot_observation_histograms.py
+    python plot_observation_histograms.py --include_upper_limits
+    python plot_observation_histograms.py \
+        --data_glob '/path/to/lc_*.npz' --out obs_hist.png
+"""
+
+import argparse
+import glob
+import os
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Reuse the canonical band ordering / keys and column conventions from the
+# dataset so "observation" here means exactly what the model trains on.
+from yoke.datasets.kilonova_dataset import NINE_BAND_KEYS
+
+
+# Default to the same data the 9-band pipeline trains and computes norm stats on
+# (the Rubin+ZTF set), so the histogram reflects the real training distribution.
+DEFAULT_DATA_GLOB = (
+    "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+    "rubin_ztf_10000_dataset/lc_*.npz"
+)
+DEFAULT_ERROR_COL = 2
+
+# Short display labels for the bands, in NINE_BAND_KEYS order.
+BAND_LABELS = tuple(k.replace("arr_", "") for k in NINE_BAND_KEYS)
+
+
+def collect_counts(files, band_keys, error_col):
+    """Count detections and upper limits per band across all files.
+
+    Args:
+        files (list[str]): npz light-curve files to read.
+        band_keys (tuple[str, ...]): Band keys to count, in display order.
+        error_col (int): Column whose finiteness distinguishes a detection
+            (finite) from an upper limit / non-detection (non-finite).
+
+    Returns:
+        dict with:
+            det_totals (np.ndarray): Total detections per band, shape [n_bands].
+            lim_totals (np.ndarray): Total upper limits per band, shape [n_bands].
+            det_per_curve (list[np.ndarray]): For each band, an array holding the
+                detection count in each file that contains that band.
+            total_det_per_curve (np.ndarray): Total detections (all bands) per
+                file, one entry per file.
+            n_files (int): Number of files successfully read.
+    """
+    n_bands = len(band_keys)
+    det_totals = np.zeros(n_bands, dtype=np.int64)
+    lim_totals = np.zeros(n_bands, dtype=np.int64)
+    det_per_curve = [[] for _ in range(n_bands)]
+    total_det_per_curve = []
+
+    n_files = 0
+    for fn in files:
+        try:
+            data = np.load(fn, allow_pickle=True)
+        except Exception as exc:
+            print(f"  skipped {fn}: {exc}")
+            continue
+
+        file_total_det = 0
+        for b, key in enumerate(band_keys):
+            if key not in data.files:
+                continue
+
+            arr = data[key]
+            if arr.size == 0:
+                continue
+
+            errs = arr[:, error_col].astype(np.float64)
+            detected = np.isfinite(errs)
+            n_det = int(detected.sum())
+            n_lim = int(detected.size - n_det)
+
+            det_totals[b] += n_det
+            lim_totals[b] += n_lim
+            det_per_curve[b].append(n_det)
+            file_total_det += n_det
+
+        data.close()
+        total_det_per_curve.append(file_total_det)
+        n_files += 1
+
+    return {
+        "det_totals": det_totals,
+        "lim_totals": lim_totals,
+        "det_per_curve": [np.asarray(c, dtype=np.int64) for c in det_per_curve],
+        "total_det_per_curve": np.asarray(total_det_per_curve, dtype=np.int64),
+        "n_files": n_files,
+    }
+
+
+def plot_band_totals(ax, counts, include_upper_limits):
+    """Bar chart of total observations per band."""
+    n_bands = len(BAND_LABELS)
+    x = np.arange(n_bands)
+    det = counts["det_totals"]
+
+    if include_upper_limits:
+        lim = counts["lim_totals"]
+        ax.bar(x, det, color="tab:blue", label="detections")
+        ax.bar(x, lim, bottom=det, color="tab:gray", alpha=0.6, label="upper limits")
+        ax.legend()
+        title_extra = " (detections + upper limits)"
+    else:
+        ax.bar(x, det, color="tab:blue")
+        title_extra = " (detections only)"
+
+    # Annotate each bar with its detection count.
+    for xi, di in zip(x, det):
+        ax.text(xi, di, f"{int(di)}", ha="center", va="bottom", fontsize=8)
+
+    ax.set_xticks(x)
+    ax.set_xticklabels(BAND_LABELS, rotation=45, ha="right")
+    ax.set_ylabel("Number of observations")
+    ax.set_title(f"Total observations per band{title_extra}")
+    ax.grid(True, axis="y", alpha=0.3)
+
+
+def plot_total_hist(ax, counts):
+    """Histogram of total detections per light curve (summed over all bands)."""
+    totals = counts["total_det_per_curve"]
+    if totals.size == 0:
+        ax.set_visible(False)
+        return
+
+    hi = int(totals.max())
+    # One bin per integer count up to the max, capped so very long tails stay
+    # readable.
+    bins = np.arange(0, hi + 2) - 0.5 if hi <= 60 else 50
+    ax.hist(totals, bins=bins, color="tab:green", alpha=0.8)
+    ax.axvline(
+        totals.mean(),
+        color="k",
+        linestyle="--",
+        linewidth=1,
+        label=f"mean {totals.mean():.1f}",
+    )
+    ax.set_xlabel("Detections per light curve (all bands)")
+    ax.set_ylabel("Number of light curves")
+    ax.set_title("Total detections per light curve")
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+
+
+def plot_per_band_hist(ax, counts):
+    """Overlaid step histograms of detections-per-curve for each band."""
+    per_curve = counts["det_per_curve"]
+
+    # Common integer bins across bands so the overlays are comparable.
+    hi = max((c.max() if c.size else 0) for c in per_curve)
+    hi = int(hi)
+    bins = np.arange(0, max(hi, 1) + 2) - 0.5
+
+    cmap = plt.get_cmap("tab10")
+    for b, label in enumerate(BAND_LABELS):
+        c = per_curve[b]
+        if c.size == 0:
+            continue
+        ax.hist(
+            c,
+            bins=bins,
+            histtype="step",
+            linewidth=1.5,
+            color=cmap(b % 10),
+            label=label,
+        )
+
+    ax.set_xlabel("Detections per light curve")
+    ax.set_ylabel("Number of light curves")
+    ax.set_title("Detections per light curve, by band")
+    ax.legend(fontsize=8, ncol=2)
+    ax.grid(True, alpha=0.3)
+
+
+def print_summary(counts, include_upper_limits):
+    """Print the per-band and overall counts to stdout."""
+    det = counts["det_totals"]
+    lim = counts["lim_totals"]
+
+    print(f"Read {counts['n_files']} light-curve files.")
+    print(f"{'band':<10} {'detections':>12} {'upper_limits':>14}")
+    for b, label in enumerate(BAND_LABELS):
+        print(f"{label:<10} {int(det[b]):>12} {int(lim[b]):>14}")
+
+    print(f"{'TOTAL':<10} {int(det.sum()):>12} {int(lim.sum()):>14}")
+
+    totals = counts["total_det_per_curve"]
+    if totals.size:
+        print(
+            f"Detections per light curve: mean {totals.mean():.2f}, "
+            f"median {np.median(totals):.0f}, "
+            f"min {int(totals.min())}, max {int(totals.max())}"
+        )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Plot observation-count histograms (per band and total) for the "
+            "9-band kilonova light-curve data."
+        )
+    )
+
+    parser.add_argument(
+        "--data_glob",
+        type=str,
+        default=DEFAULT_DATA_GLOB,
+        help="Glob for the light-curve npz files. Default: the Rubin+ZTF set.",
+    )
+    parser.add_argument(
+        "--error_col",
+        type=int,
+        default=DEFAULT_ERROR_COL,
+        help=(
+            "Column whose finiteness marks a detection (finite) vs an upper "
+            "limit / non-detection (non-finite). Default 2."
+        ),
+    )
+    parser.add_argument(
+        "--include_upper_limits",
+        action="store_true",
+        help=(
+            "Also count and stack upper limits (non-detections) in the per-band "
+            "totals. By default only real detections are counted, matching the "
+            "9-band dataset with drop_upper_limits=True."
+        ),
+    )
+    parser.add_argument(
+        "--out",
+        type=str,
+        default="observation_histograms.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument("--dpi", type=int, default=200)
+
+    args = parser.parse_args()
+
+    files = sorted(glob.glob(args.data_glob))
+    if len(files) == 0:
+        raise FileNotFoundError(f"No files matched glob: {args.data_glob}")
+
+    print(f"Matched {len(files)} files for glob: {args.data_glob}")
+
+    counts = collect_counts(files, NINE_BAND_KEYS, args.error_col)
+    print_summary(counts, args.include_upper_limits)
+
+    fig, axes = plt.subplots(1, 3, figsize=(18, 5.5))
+    plot_band_totals(axes[0], counts, args.include_upper_limits)
+    plot_total_hist(axes[1], counts)
+    plot_per_band_hist(axes[2], counts)
+
+    fig.suptitle(
+        f"KN light-curve observations ({counts['n_files']} light curves)",
+        fontsize=13,
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.96))
+    fig.savefig(args.out, dpi=args.dpi)
+    print(f"Saved {os.path.abspath(args.out)}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index 3ba64beb..b354e89c 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -17,6 +17,12 @@
 #SBATCH --nodes=<KNODES>
 #SBATCH --ntasks-per-node=<NGPUS>
 #SBATCH --gpus-per-node=<NGPUS>
+# nodes have 8 H100s and 192-208 CPUs, i.e. ~24 cores/GPU. One task
+# per GPU (ntasks-per-node=NGPUS), so cpus-per-task=24 gives each rank enough
+# cores for its dataloader workers + OMP threads without over-requesting on the
+# 192-core nodes (24*NGPUS stays within budget for NGPUS up to 8). Without this
+# each task defaults to ~1 CPU and starves data loading, throttling the GPUs.
+#SBATCH --cpus-per-task=24
 #SBATCH --mem-per-gpu=50G
 #SBATCH --output=study<studyIDX>_epoch<epochIDX>.out
 #SBATCH --error=study<studyIDX>_epoch<epochIDX>.err
@@ -54,8 +60,9 @@ module load anaconda/3.12
 source activate
 conda activate <YOKE_TORCH_ENV>
 
-# Set number of threads per GPU
-export OMP_NUM_THREADS=10
+# Set number of threads per GPU. Must be <= cpus-per-task (24). Leaves cores for
+# the dataloader workers (NUM_WORKERS per rank) alongside the OMP compute threads.
+export OMP_NUM_THREADS=8
 
 # Get start time
 export date00=`date`

From f70cb3f7f4a981554c6ad850dc7731a01ec9ab7a Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 14 Aug 2026 13:31:24 -0600
Subject: [PATCH 28/66] add sweep plot

---
 .../plot_observation_histograms.py            | 115 +++++++++++++++++-
 1 file changed, 110 insertions(+), 5 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_observation_histograms.py b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
index 0f8289c6..963d4353 100644
--- a/applications/harnesses/KN_loderunner/plot_observation_histograms.py
+++ b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
@@ -14,10 +14,16 @@
      ``--include_upper_limits`` is set.
   3. Distribution of detections-per-light-curve, per band and summed over all
      bands -- how long/rich a typical object is.
+  4. A context-length sweep: how many training samples (context windows) and how
+     many light curves survive as ``context_len`` grows. This matters because the
+     9-band model bakes ``context_len`` into its first layer, so choosing a
+     longer context means a retrain -- this panel shows the data cost before you
+     pay for it.
 
 Run directly, e.g.:
     python plot_observation_histograms.py
     python plot_observation_histograms.py --include_upper_limits
+    python plot_observation_histograms.py --sweep_max 20
     python plot_observation_histograms.py \
         --data_glob '/path/to/lc_*.npz' --out obs_hist.png
 """
@@ -194,6 +200,78 @@ def plot_per_band_hist(ax, counts):
     ax.grid(True, alpha=0.3)
 
 
+def context_len_sweep(events_per_file, sweep_max):
+    """Compute surviving samples/curves as a function of ``context_len``.
+
+    The 9-band dataset turns each light curve of ``n_events`` detections into
+    ``max(0, n_events - context_len)`` training windows (one per start index,
+    ``max_start = n_events - context_len - 1``, inclusive). A curve contributes
+    at all only when ``n_events > context_len``. This mirrors
+    ``Kilonova_lc_scalar_context_DataSet_9band`` exactly.
+
+    Args:
+        events_per_file (np.ndarray): Detections (== merged event count) per
+            light curve, one entry per file.
+        sweep_max (int): Largest context length to evaluate.
+
+    Returns:
+        context_lens (np.ndarray): Candidate context lengths, shape [sweep_max].
+        n_samples (np.ndarray): Total training windows at each context length.
+        n_curves (np.ndarray): Number of light curves that yield >=1 window.
+    """
+    context_lens = np.arange(1, sweep_max + 1)
+    ev = events_per_file.astype(np.int64)
+
+    n_samples = np.array(
+        [np.maximum(ev - c, 0).sum() for c in context_lens], dtype=np.int64
+    )
+    n_curves = np.array(
+        [int((ev > c).sum()) for c in context_lens], dtype=np.int64
+    )
+    return context_lens, n_samples, n_curves
+
+
+def plot_context_sweep(ax, events_per_file, sweep_max):
+    """Plot surviving training samples and light curves vs context length."""
+    if events_per_file.size == 0:
+        ax.set_visible(False)
+        return
+
+    context_lens, n_samples, n_curves = context_len_sweep(
+        events_per_file, sweep_max
+    )
+    n_total = events_per_file.size
+
+    # Samples on the left axis (can be large); fraction of curves kept on the
+    # right axis so both trends are readable together.
+    ax.plot(context_lens, n_samples, marker="o", color="tab:purple",
+            label="training windows")
+    ax.set_xlabel("context_len")
+    ax.set_ylabel("Total training windows", color="tab:purple")
+    ax.tick_params(axis="y", labelcolor="tab:purple")
+    ax.set_title("Data cost of context length")
+    ax.grid(True, alpha=0.3)
+
+    ax2 = ax.twinx()
+    frac_curves = n_curves / n_total
+    ax2.plot(context_lens, frac_curves, marker="s", color="tab:red",
+             linestyle="--", label="fraction of curves kept")
+    ax2.set_ylabel("Fraction of light curves kept", color="tab:red")
+    ax2.tick_params(axis="y", labelcolor="tab:red")
+    ax2.set_ylim(0, 1.02)
+
+    # Mark the current default context_len=5 for reference.
+    if context_lens[0] <= 5 <= context_lens[-1]:
+        ax.axvline(5, color="gray", linestyle=":", linewidth=1, alpha=0.7)
+        ax.text(5, ax.get_ylim()[1], " default 5", ha="left", va="top",
+                fontsize=8, color="gray")
+
+    # Combined legend from both axes.
+    lines1, labels1 = ax.get_legend_handles_labels()
+    lines2, labels2 = ax2.get_legend_handles_labels()
+    ax.legend(lines1 + lines2, labels1 + labels2, fontsize=8, loc="center right")
+
+
 def print_summary(counts, include_upper_limits):
     """Print the per-band and overall counts to stdout."""
     det = counts["det_totals"]
@@ -215,6 +293,22 @@ def print_summary(counts, include_upper_limits):
         )
 
 
+def print_sweep(events_per_file, sweep_max):
+    """Print the context-length sweep table to stdout."""
+    if events_per_file.size == 0:
+        return
+
+    context_lens, n_samples, n_curves = context_len_sweep(
+        events_per_file, sweep_max
+    )
+    n_total = events_per_file.size
+
+    print("\nContext-length sweep (samples = training windows):")
+    print(f"{'context_len':>12} {'windows':>12} {'curves_kept':>12} {'frac':>7}")
+    for c, s, k in zip(context_lens, n_samples, n_curves):
+        print(f"{int(c):>12} {int(s):>12} {int(k):>12} {k / n_total:>7.2f}")
+
+
 def main():
     parser = argparse.ArgumentParser(
         description=(
@@ -247,6 +341,15 @@ def main():
             "9-band dataset with drop_upper_limits=True."
         ),
     )
+    parser.add_argument(
+        "--sweep_max",
+        type=int,
+        default=20,
+        help=(
+            "Largest context_len to evaluate in the data-cost sweep panel. "
+            "Default 20."
+        ),
+    )
     parser.add_argument(
         "--out",
         type=str,
@@ -265,17 +368,19 @@ def main():
 
     counts = collect_counts(files, NINE_BAND_KEYS, args.error_col)
     print_summary(counts, args.include_upper_limits)
+    print_sweep(counts["total_det_per_curve"], args.sweep_max)
 
-    fig, axes = plt.subplots(1, 3, figsize=(18, 5.5))
-    plot_band_totals(axes[0], counts, args.include_upper_limits)
-    plot_total_hist(axes[1], counts)
-    plot_per_band_hist(axes[2], counts)
+    fig, axes = plt.subplots(2, 2, figsize=(15, 11))
+    plot_band_totals(axes[0, 0], counts, args.include_upper_limits)
+    plot_total_hist(axes[0, 1], counts)
+    plot_per_band_hist(axes[1, 0], counts)
+    plot_context_sweep(axes[1, 1], counts["total_det_per_curve"], args.sweep_max)
 
     fig.suptitle(
         f"KN light-curve observations ({counts['n_files']} light curves)",
         fontsize=13,
     )
-    fig.tight_layout(rect=(0, 0, 1, 0.96))
+    fig.tight_layout(rect=(0, 0, 1, 0.97))
     fig.savefig(args.out, dpi=args.dpi)
     print(f"Saved {os.path.abspath(args.out)}")
 

From 8cfcf1b3307174b490ef575b068e706ea2d9a5d6 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 14 Aug 2026 15:01:22 -0600
Subject: [PATCH 29/66] improving context window

---
 .../harnesses/KN_loderunner/infer_9band.py    | 121 +++++++++--
 .../plot_observation_histograms.py            | 175 +++++++++++++++-
 .../plot_pred_diagnostics_9band.py            | 195 +++++++++++++++---
 .../KN_loderunner/train_LodeRunner_ddp.py     |  25 ++-
 src/yoke/datasets/kilonova_dataset.py         | 150 +++++++++++++-
 src/yoke/models/vit/swin/bomberman.py         |  27 ++-
 src/yoke/utils/checkpointing.py               |  10 +
 7 files changed, 637 insertions(+), 66 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 548443d7..2a2353ea 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -157,10 +157,21 @@ def load_9band_model(ckpt_path, device):
     hidden = ckpt.get("hidden", 64)
     noise_scale = ckpt.get("noise_scale", 0.0)
 
+    # Time-window context mode. When set, the model's first layer is sized by the
+    # padded width (max_context_len) and each event carries an extra validity
+    # flag. Legacy fixed-count checkpoints fall through to None.
+    context_window_days = ckpt.get("context_window_days", None)
+    if context_window_days is not None:
+        max_context_len = ckpt.get("max_context_len", context_len)
+    else:
+        max_context_len = context_len
+
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
     print("target_type:", ckpt.get("target_type", "unknown"))
     print("context_len:", context_len)
+    print("context_window_days:", context_window_days)
+    print("max_context_len:", max_context_len)
     print("n_bands:", n_bands)
 
     backbone = LodeRunner(**model_args).to(device)
@@ -168,11 +179,12 @@ def load_9band_model(ckpt_path, device):
 
     model = ScalarTemporalConditionedLodeRunner_9band(
         backbone=backbone,
-        context_len=context_len,
+        context_len=max_context_len,
         n_bands=n_bands,
         image_size=model_args["image_size"],
         backbone_channels=backbone_channels,
         hidden=hidden,
+        context_window_days=context_window_days,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
@@ -183,7 +195,7 @@ def load_9band_model(ckpt_path, device):
 
     model.eval()
 
-    return model, context_len, n_bands
+    return model, context_len, n_bands, context_window_days, max_context_len
 
 
 def load_event_stream(fn, means, stds):
@@ -251,22 +263,43 @@ def load_event_stream(fn, means, stds):
     return times, values_norm.astype(np.float32), bands, raw, t0
 
 
-def build_context_input(ctx_t, ctx_v, ctx_b, n_bands, device):
+def build_context_input(
+    ctx_t, ctx_v, ctx_b, n_bands, device, window_mode=False, max_context_len=None
+):
     """Build the flattened per-event context input for the model.
 
-    Layout per event: [value, rel_t, one_hot_band(n_bands)].
+    Fixed-count mode (``window_mode=False``): layout per event is
+    ``[value, rel_t, one_hot_band(n_bands)]`` (width ``2 + n_bands``).
+
+    Time-window mode (``window_mode=True``): the real events are padded to
+    ``max_context_len`` rows and each event gains a validity flag, giving
+    ``[value, rel_t, valid, one_hot_band(n_bands)]`` (width ``3 + n_bands``).
+    Real events fill the leading rows in time order with ``rel_t`` relative to
+    the first real event; padded rows are all-zero with ``valid = 0``. Matches
+    ``_getitem_window`` in the dataset.
     """
-    context_len = len(ctx_t)
+    ctx_t = np.asarray(ctx_t, dtype=np.float32)
+    ctx_v = np.asarray(ctx_v, dtype=np.float32)
+    ctx_b = np.asarray(ctx_b, dtype=np.int64)
 
     rel_t = (ctx_t - ctx_t[0]).astype(np.float32)
 
-    band_onehot = np.zeros((context_len, n_bands), dtype=np.float32)
-    band_onehot[np.arange(context_len), ctx_b] = 1.0
-
-    per_event = np.concatenate(
-        [ctx_v[:, None], rel_t[:, None], band_onehot],
-        axis=1,
-    )
+    if window_mode:
+        n_real = ctx_t.shape[0]
+        per_event = np.zeros((max_context_len, 3 + n_bands), dtype=np.float32)
+        per_event[:n_real, 0] = ctx_v
+        per_event[:n_real, 1] = rel_t
+        per_event[:n_real, 2] = 1.0  # validity flag for real events
+        per_event[np.arange(n_real), 3 + ctx_b] = 1.0
+    else:
+        context_len = ctx_t.shape[0]
+        band_onehot = np.zeros((context_len, n_bands), dtype=np.float32)
+        band_onehot[np.arange(context_len), ctx_b] = 1.0
+
+        per_event = np.concatenate(
+            [ctx_v[:, None], rel_t[:, None], band_onehot],
+            axis=1,
+        )
 
     x = torch.tensor(
         per_event.reshape(-1),
@@ -278,7 +311,17 @@ def build_context_input(ctx_t, ctx_v, ctx_b, n_bands, device):
 
 
 def forecast_curve(
-    stream, model, device, context_len, n_bands, means, stds, lead_times
+    stream,
+    model,
+    device,
+    context_len,
+    n_bands,
+    means,
+    stds,
+    lead_times,
+    window_mode=False,
+    context_window_days=None,
+    max_context_len=None,
 ):
     """Forecast all bands at a grid of future lead times from the last context.
 
@@ -287,12 +330,32 @@ def forecast_curve(
     """
     times, values_norm, bands, _, _ = stream
 
-    # Most recent context_len events.
-    ctx_t = times[-context_len:]
-    ctx_v = values_norm[-context_len:]
-    ctx_b = bands[-context_len:]
-
-    x = build_context_input(ctx_t, ctx_v, ctx_b, n_bands, device)
+    if window_mode:
+        # Time-window context: all events within context_window_days of the last
+        # observation, capped at max_context_len (matches _getitem_window).
+        anchor_t = times[-1]
+        lo = anchor_t - context_window_days
+        sel_idx = np.nonzero(times >= lo)[0]
+        if sel_idx.shape[0] > max_context_len:
+            sel_idx = sel_idx[-max_context_len:]
+        ctx_t = times[sel_idx]
+        ctx_v = values_norm[sel_idx]
+        ctx_b = bands[sel_idx]
+    else:
+        # Most recent context_len events.
+        ctx_t = times[-context_len:]
+        ctx_v = values_norm[-context_len:]
+        ctx_b = bands[-context_len:]
+
+    x = build_context_input(
+        ctx_t,
+        ctx_v,
+        ctx_b,
+        n_bands,
+        device,
+        window_mode=window_mode,
+        max_context_len=max_context_len,
+    )
 
     last_t = float(times[-1])
 
@@ -382,7 +445,15 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print("Using device:", device)
 
-    model, context_len, n_bands = load_9band_model(args.ckpt, device)
+    (
+        model,
+        context_len,
+        n_bands,
+        context_window_days,
+        max_context_len,
+    ) = load_9band_model(args.ckpt, device)
+
+    window_mode = context_window_days is not None
 
     means, stds = load_or_compute_band_normalization(
         stats_path=args.norm_stats_path,
@@ -411,10 +482,13 @@ def main():
             continue
 
         times = stream[0]
-        if len(times) < context_len:
+        # In window mode any curve with at least one detection can be forecast
+        # (the trailing window is padded); otherwise we need a full context.
+        min_events = 1 if window_mode else context_len
+        if len(times) < min_events:
             print(
                 f"Skipping {fn}: only {len(times)} events, "
-                f"need at least context_len={context_len}."
+                f"need at least {min_events}."
             )
             continue
 
@@ -427,6 +501,9 @@ def main():
             means=means,
             stds=stds,
             lead_times=lead_times,
+            window_mode=window_mode,
+            context_window_days=context_window_days,
+            max_context_len=max_context_len,
         )
 
         base = os.path.splitext(os.path.basename(fn))[0]
diff --git a/applications/harnesses/KN_loderunner/plot_observation_histograms.py b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
index 963d4353..5f06ff85 100644
--- a/applications/harnesses/KN_loderunner/plot_observation_histograms.py
+++ b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
@@ -69,6 +69,10 @@ def collect_counts(files, band_keys, error_col):
                 detection count in each file that contains that band.
             total_det_per_curve (np.ndarray): Total detections (all bands) per
                 file, one entry per file.
+            event_times_per_file (list[np.ndarray]): For each file, the merged,
+                time-sorted, file-relative detection times (all bands), matching
+                the event stream the 9-band dataset builds. Used by the
+                time-window sweep.
             n_files (int): Number of files successfully read.
     """
     n_bands = len(band_keys)
@@ -76,6 +80,7 @@ def collect_counts(files, band_keys, error_col):
     lim_totals = np.zeros(n_bands, dtype=np.int64)
     det_per_curve = [[] for _ in range(n_bands)]
     total_det_per_curve = []
+    event_times_per_file = []
 
     n_files = 0
     for fn in files:
@@ -86,6 +91,7 @@ def collect_counts(files, band_keys, error_col):
             continue
 
         file_total_det = 0
+        file_det_times = []
         for b, key in enumerate(band_keys):
             if key not in data.files:
                 continue
@@ -104,8 +110,21 @@ def collect_counts(files, band_keys, error_col):
             det_per_curve[b].append(n_det)
             file_total_det += n_det
 
+            # Collect detection times (col 0 = MJD) for the merged event stream.
+            if n_det:
+                file_det_times.append(arr[detected, 0].astype(np.float64))
+
         data.close()
         total_det_per_curve.append(file_total_det)
+
+        # Merge all bands' detections into one time-sorted, file-relative stream,
+        # exactly as Kilonova_lc_scalar_context_DataSet_9band does.
+        if file_det_times:
+            merged = np.concatenate(file_det_times)
+            merged.sort(kind="stable")
+            merged -= merged.min()
+            event_times_per_file.append(merged)
+
         n_files += 1
 
     return {
@@ -113,6 +132,7 @@ def collect_counts(files, band_keys, error_col):
         "lim_totals": lim_totals,
         "det_per_curve": [np.asarray(c, dtype=np.int64) for c in det_per_curve],
         "total_det_per_curve": np.asarray(total_det_per_curve, dtype=np.int64),
+        "event_times_per_file": event_times_per_file,
         "n_files": n_files,
     }
 
@@ -272,6 +292,101 @@ def plot_context_sweep(ax, events_per_file, sweep_max):
     ax.legend(lines1 + lines2, labels1 + labels2, fontsize=8, loc="center right")
 
 
+def time_window_counts(event_times_per_file, window_days):
+    """Real detections per context window for a fixed lookback in days.
+
+    Mirrors the planned window-mode selection in
+    ``Kilonova_lc_scalar_context_DataSet_9band``: for each target event (every
+    event after the first in a file's merged, time-sorted stream), the context
+    is every earlier event ``j`` with
+    ``times[target - 1] - times[j] <= window_days``. The window is anchored on
+    the event immediately preceding the target (the most recent observation),
+    so it always contains at least that one event.
+
+    Args:
+        event_times_per_file (list[np.ndarray]): Per-file merged, sorted,
+            file-relative detection times.
+        window_days (float): Trailing lookback length in days.
+
+    Returns:
+        np.ndarray: One entry per (file, target-event) sample: the number of
+            real detections that fall inside the trailing window.
+    """
+    counts = []
+    for times in event_times_per_file:
+        n = times.shape[0]
+        if n < 2:
+            continue
+        # target_idx runs over every event after the first; the window is
+        # anchored at times[target_idx - 1] and looks back window_days.
+        for target_idx in range(1, n):
+            anchor = times[target_idx - 1]
+            lo = anchor - window_days
+            # Events strictly before the target that are within the window.
+            in_window = times[:target_idx]
+            counts.append(int((in_window >= lo).sum()))
+    return np.asarray(counts, dtype=np.int64)
+
+
+def time_window_sweep(event_times_per_file, window_grid):
+    """Distribution of context size vs trailing window length.
+
+    Args:
+        event_times_per_file (list[np.ndarray]): Per-file merged detection times.
+        window_grid (np.ndarray): Candidate window lengths in days.
+
+    Returns:
+        window_grid (np.ndarray): The evaluated window lengths.
+        stats (dict): Percentile arrays keyed by label ("median", "p90", "p95",
+            "p99", "max", "mean"), each shape [len(window_grid)].
+        raw (list[np.ndarray]): Per-window arrays of per-sample context counts,
+            for histogramming.
+    """
+    pct_labels = [("median", 50), ("p90", 90), ("p95", 95), ("p99", 99)]
+    stats = {label: [] for label, _ in pct_labels}
+    stats["max"] = []
+    stats["mean"] = []
+    raw = []
+
+    for w in window_grid:
+        c = time_window_counts(event_times_per_file, float(w))
+        raw.append(c)
+        if c.size == 0:
+            for label, _ in pct_labels:
+                stats[label].append(0.0)
+            stats["max"].append(0.0)
+            stats["mean"].append(0.0)
+            continue
+        for label, q in pct_labels:
+            stats[label].append(float(np.percentile(c, q)))
+        stats["max"].append(float(c.max()))
+        stats["mean"].append(float(c.mean()))
+
+    stats = {k: np.asarray(v) for k, v in stats.items()}
+    return window_grid, stats, raw
+
+
+def plot_time_window_sweep(ax, event_times_per_file, window_grid):
+    """Plot context-size percentiles vs trailing window length."""
+    if len(event_times_per_file) == 0:
+        ax.set_visible(False)
+        return
+
+    window_grid, stats, _ = time_window_sweep(event_times_per_file, window_grid)
+
+    ax.plot(window_grid, stats["median"], marker="o", label="median")
+    ax.plot(window_grid, stats["p90"], marker="^", label="90th pct")
+    ax.plot(window_grid, stats["p95"], marker="s", label="95th pct")
+    ax.plot(window_grid, stats["p99"], marker="d", label="99th pct")
+    ax.plot(window_grid, stats["max"], linestyle=":", color="gray", label="max")
+
+    ax.set_xlabel("Context window length W (days)")
+    ax.set_ylabel("Real detections in window")
+    ax.set_title("Context size vs time window\n(pick max_context_len from a high pct)")
+    ax.legend(fontsize=8)
+    ax.grid(True, alpha=0.3)
+
+
 def print_summary(counts, include_upper_limits):
     """Print the per-band and overall counts to stdout."""
     det = counts["det_totals"]
@@ -309,6 +424,37 @@ def print_sweep(events_per_file, sweep_max):
         print(f"{int(c):>12} {int(s):>12} {int(k):>12} {k / n_total:>7.2f}")
 
 
+def print_time_window_sweep(event_times_per_file, window_grid):
+    """Print the time-window context-size sweep table to stdout.
+
+    This is the table to read when choosing the two window-mode hyperparameters:
+    ``context_window_days`` (a W with enough baseline for real evolution) and
+    ``max_context_len`` (a high percentile so padding rarely clips real events).
+    """
+    if len(event_times_per_file) == 0:
+        return
+
+    window_grid, stats, raw = time_window_sweep(event_times_per_file, window_grid)
+
+    print("\nTime-window context sweep (real detections inside trailing W days):")
+    print(
+        f"{'W_days':>8} {'n_samples':>10} {'median':>8} {'p90':>6} "
+        f"{'p95':>6} {'p99':>6} {'max':>6} {'mean':>7}"
+    )
+    for i, w in enumerate(window_grid):
+        n_s = raw[i].size
+        print(
+            f"{float(w):>8.2f} {int(n_s):>10} {stats['median'][i]:>8.0f} "
+            f"{stats['p90'][i]:>6.0f} {stats['p95'][i]:>6.0f} "
+            f"{stats['p99'][i]:>6.0f} {stats['max'][i]:>6.0f} "
+            f"{stats['mean'][i]:>7.1f}"
+        )
+    print(
+        "Pick context_window_days from W with enough baseline; set "
+        "max_context_len ~ the p95/p99 column at that W."
+    )
+
+
 def main():
     parser = argparse.ArgumentParser(
         description=(
@@ -350,6 +496,21 @@ def main():
             "Default 20."
         ),
     )
+    parser.add_argument(
+        "--window_max",
+        type=float,
+        default=10.0,
+        help=(
+            "Largest trailing window length (days) in the time-window context "
+            "sweep. Default 10."
+        ),
+    )
+    parser.add_argument(
+        "--window_step",
+        type=float,
+        default=1.0,
+        help="Step (days) between evaluated window lengths. Default 1.0.",
+    )
     parser.add_argument(
         "--out",
         type=str,
@@ -370,11 +531,19 @@ def main():
     print_summary(counts, args.include_upper_limits)
     print_sweep(counts["total_det_per_curve"], args.sweep_max)
 
-    fig, axes = plt.subplots(2, 2, figsize=(15, 11))
+    # Time-window sweep grid: window_step .. window_max (days).
+    window_grid = np.arange(
+        args.window_step, args.window_max + 0.5 * args.window_step, args.window_step
+    )
+    print_time_window_sweep(counts["event_times_per_file"], window_grid)
+
+    fig, axes = plt.subplots(2, 3, figsize=(21, 11))
     plot_band_totals(axes[0, 0], counts, args.include_upper_limits)
     plot_total_hist(axes[0, 1], counts)
-    plot_per_band_hist(axes[1, 0], counts)
-    plot_context_sweep(axes[1, 1], counts["total_det_per_curve"], args.sweep_max)
+    plot_per_band_hist(axes[0, 2], counts)
+    plot_context_sweep(axes[1, 0], counts["total_det_per_curve"], args.sweep_max)
+    plot_time_window_sweep(axes[1, 1], counts["event_times_per_file"], window_grid)
+    axes[1, 2].set_visible(False)
 
     fig.suptitle(
         f"KN light-curve observations ({counts['n_files']} light curves)",
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 0843a223..d53d7723 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -166,11 +166,22 @@ def load_9band_model(ckpt_path, device):
     hidden = ckpt.get("hidden", 64)
     noise_scale = ckpt.get("noise_scale", 0.0)
 
+    # Time-window context mode. When set, the model's first layer is sized by the
+    # padded width (max_context_len) and each event carries an extra validity
+    # flag. Legacy fixed-count checkpoints fall through to None.
+    context_window_days = ckpt.get("context_window_days", None)
+    if context_window_days is not None:
+        max_context_len = ckpt.get("max_context_len", context_len)
+    else:
+        max_context_len = context_len
+
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
     print("backbone_class:", ckpt.get("backbone_class", "LodeRunner"))
     print("target_type:", ckpt.get("target_type", "unknown"))
     print("context_len:", context_len)
+    print("context_window_days:", context_window_days)
+    print("max_context_len:", max_context_len)
     print("n_bands:", n_bands)
     print("band_keys:", ckpt.get("band_keys", list(BAND_KEYS)))
     print("backbone_channels:", backbone_channels)
@@ -181,11 +192,12 @@ def load_9band_model(ckpt_path, device):
 
     model = ScalarTemporalConditionedLodeRunner_9band(
         backbone=backbone,
-        context_len=context_len,
+        context_len=max_context_len,
         n_bands=n_bands,
         image_size=model_args["image_size"],
         backbone_channels=backbone_channels,
         hidden=hidden,
+        context_window_days=context_window_days,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
@@ -198,10 +210,15 @@ def load_9band_model(ckpt_path, device):
 
     model.eval()
 
-    return model, context_len, n_bands
+    return model, context_len, n_bands, context_window_days, max_context_len
 
 
-def make_eval_dataset(args, context_len):
+def make_eval_dataset(
+    args,
+    context_len,
+    context_window_days=None,
+    max_context_len=None,
+):
     band_means, band_stds = load_or_compute_band_normalization(
         stats_path=args.norm_stats_path,
         band_keys=BAND_KEYS,
@@ -223,30 +240,59 @@ def make_eval_dataset(args, context_len):
         drop_upper_limits=DROP_UPPER_LIMITS,
         means=band_means,
         stds=band_stds,
+        context_window_days=context_window_days,
+        max_context_len=max_context_len,
     )
 
     return dataset, np.asarray(band_means), np.asarray(band_stds)
 
 
-def build_context_input(win_v, win_t, win_b, context_len, n_bands, device):
+def build_context_input(
+    win_v,
+    win_t,
+    win_b,
+    context_len,
+    n_bands,
+    device,
+    window_mode=False,
+):
     """Build the flattened per-event context input for the model.
 
-    Layout per event: [value, rel_t, one_hot_band(n_bands)], relative time
-    measured from the first event in the window, matching the dataset.
+    Fixed-count mode (``window_mode=False``): layout per event is
+    ``[value, rel_t, one_hot_band(n_bands)]`` (width ``2 + n_bands``), with
+    ``rel_t`` measured from the first event in the window and exactly
+    ``context_len`` real events, matching the fixed-count dataset path.
+
+    Time-window mode (``window_mode=True``): the (real) events are padded to
+    ``context_len`` (= ``max_context_len``) rows and each event gains a validity
+    flag, giving ``[value, rel_t, valid, one_hot_band(n_bands)]`` (width
+    ``3 + n_bands``). Real events fill the leading rows in time order with
+    ``rel_t`` relative to the first real event; padded rows are all-zero with
+    ``valid = 0``. This matches ``_getitem_window`` in the dataset.
     """
     win_v = np.asarray(win_v, dtype=np.float32)
     win_t = np.asarray(win_t, dtype=np.float32)
     win_b = np.asarray(win_b, dtype=np.int64)
 
-    rel_t = (win_t - win_t[0]).astype(np.float32)
+    if window_mode:
+        n_real = win_v.shape[0]
+        rel_t = (win_t - win_t[0]).astype(np.float32)
 
-    band_onehot = np.zeros((context_len, n_bands), dtype=np.float32)
-    band_onehot[np.arange(context_len), win_b] = 1.0
+        per_event = np.zeros((context_len, 3 + n_bands), dtype=np.float32)
+        per_event[:n_real, 0] = win_v
+        per_event[:n_real, 1] = rel_t
+        per_event[:n_real, 2] = 1.0  # validity flag for real events
+        per_event[np.arange(n_real), 3 + win_b] = 1.0
+    else:
+        rel_t = (win_t - win_t[0]).astype(np.float32)
 
-    per_event = np.concatenate(
-        [win_v[:, None], rel_t[:, None], band_onehot],
-        axis=1,
-    )
+        band_onehot = np.zeros((context_len, n_bands), dtype=np.float32)
+        band_onehot[np.arange(context_len), win_b] = 1.0
+
+        per_event = np.concatenate(
+            [win_v[:, None], rel_t[:, None], band_onehot],
+            axis=1,
+        )
 
     return torch.tensor(
         per_event.reshape(-1),
@@ -255,6 +301,31 @@ def build_context_input(win_v, win_t, win_b, context_len, n_bands, device):
     ).unsqueeze(0)
 
 
+def _select_window(ctx_t, ctx_v, ctx_b, context_window_days, max_context_len):
+    """Select the trailing time-window subset of a growing context.
+
+    Mirrors ``_getitem_window`` in the dataset: keep every event within
+    ``context_window_days`` of the most recent context event, then keep the most
+    recent ``max_context_len`` if more qualify. Returns (win_v, win_t, win_b) as
+    lists in time order (oldest first).
+    """
+    ct = np.asarray(ctx_t, dtype=np.float32)
+    cv = np.asarray(ctx_v, dtype=np.float32)
+    cb = np.asarray(ctx_b, dtype=np.int64)
+
+    anchor_t = ct[-1]
+    lo = anchor_t - context_window_days
+    sel_idx = np.nonzero(ct >= lo)[0]
+    if sel_idx.shape[0] > max_context_len:
+        sel_idx = sel_idx[-max_context_len:]
+
+    return (
+        list(cv[sel_idx]),
+        list(ct[sel_idx]),
+        list(cb[sel_idx]),
+    )
+
+
 def get_rollout_from_stream(
     times,
     values,
@@ -268,6 +339,9 @@ def get_rollout_from_stream(
     means,
     stds,
     teacher_forced=False,
+    window_mode=False,
+    context_window_days=None,
+    max_context_len=None,
 ):
     """Autoregressively forecast the next events of one merged event stream.
 
@@ -291,11 +365,20 @@ def get_rollout_from_stream(
     """
     t_ref = float(times[start_idx])
 
+    # Number of true events used to warm-start the running context. In window
+    # mode we seed up to max_context_len so the trailing-W-days selection has
+    # enough events to draw from; otherwise the fixed count. Clamp so at least
+    # one true event remains past the seed for a future step (lets short curves,
+    # now trainable in window mode, still roll out).
+    requested_seed = max_context_len if window_mode else context_len
+    seed_len = min(requested_seed, len(times) - start_idx - 1)
+    seed_len = max(1, seed_len)
+
     # Running context window; values are fed back from predictions as we roll
     # out, while times and band identities follow the true observation schedule.
-    ctx_t = list(times[start_idx : start_idx + context_len].astype(np.float32))
-    ctx_v = list(values[start_idx : start_idx + context_len].astype(np.float32))
-    ctx_b = list(bands[start_idx : start_idx + context_len].astype(np.int64))
+    ctx_t = list(times[start_idx : start_idx + seed_len].astype(np.float32))
+    ctx_v = list(values[start_idx : start_idx + seed_len].astype(np.float32))
+    ctx_b = list(bands[start_idx : start_idx + seed_len].astype(np.int64))
 
     context = {
         "t_rel": np.asarray(ctx_t, dtype=np.float32) - t_ref,
@@ -308,22 +391,36 @@ def get_rollout_from_stream(
 
     with torch.no_grad():
         for step in range(n_future_steps):
-            target_idx = start_idx + context_len + step
+            target_idx = start_idx + seed_len + step
 
             if target_idx >= len(times):
                 break
 
-            win_v = ctx_v[-context_len:]
-            win_t = ctx_t[-context_len:]
-            win_b = ctx_b[-context_len:]
+            if window_mode:
+                # Time-window context: trailing W days of all events observed
+                # so far, padded to max_context_len (matches _getitem_window).
+                win_v, win_t, win_b = _select_window(
+                    ctx_t=ctx_t,
+                    ctx_v=ctx_v,
+                    ctx_b=ctx_b,
+                    context_window_days=context_window_days,
+                    max_context_len=max_context_len,
+                )
+                build_width = max_context_len
+            else:
+                win_v = ctx_v[-context_len:]
+                win_t = ctx_t[-context_len:]
+                win_b = ctx_b[-context_len:]
+                build_width = context_len
 
             x = build_context_input(
                 win_v=win_v,
                 win_t=win_t,
                 win_b=win_b,
-                context_len=context_len,
+                context_len=build_width,
                 n_bands=n_bands,
                 device=device,
+                window_mode=window_mode,
             )
 
             # Lead time from the last context event to the next true event.
@@ -396,17 +493,31 @@ def get_rollout_from_stream(
     # meaning without feedback).
     fixed_forecast = None
     if not teacher_forced and steps:
-        win_t0 = times[start_idx : start_idx + context_len].astype(np.float32)
-        win_v0 = values[start_idx : start_idx + context_len].astype(np.float32)
-        win_b0 = bands[start_idx : start_idx + context_len].astype(np.int64)
+        win_t0 = times[start_idx : start_idx + seed_len].astype(np.float32)
+        win_v0 = values[start_idx : start_idx + seed_len].astype(np.float32)
+        win_b0 = bands[start_idx : start_idx + seed_len].astype(np.int64)
+
+        if window_mode:
+            # Trailing W days of the seeded context, padded to max_context_len.
+            win_v0, win_t0, win_b0 = _select_window(
+                ctx_t=win_t0,
+                ctx_v=win_v0,
+                ctx_b=win_b0,
+                context_window_days=context_window_days,
+                max_context_len=max_context_len,
+            )
+            build_width0 = max_context_len
+        else:
+            build_width0 = context_len
 
         x0 = build_context_input(
             win_v=win_v0,
             win_t=win_t0,
             win_b=win_b0,
-            context_len=context_len,
+            context_len=build_width0,
             n_bands=n_bands,
             device=device,
+            window_mode=window_mode,
         )
 
         last_ctx_t_rel = float(win_t0[-1]) - t_ref
@@ -441,12 +552,22 @@ def get_rollout_from_stream(
     }
 
 
-def select_series(dataset, context_len, n_future_steps, n_series):
+def select_series(
+    dataset,
+    context_len,
+    n_future_steps,
+    n_series,
+    seed_len=None,
+):
     """Pick files with enough events for a rollout, longest first.
 
     Returns a list of (times, values, bands, start_idx) tuples.
     """
-    min_events = context_len + 1  # need at least one future step
+    # In window mode the context is seeded with up to max_context_len events
+    # (passed as seed_len); otherwise the fixed count. Either way we need at
+    # least one event past the seed for a future step.
+    warm = seed_len if seed_len is not None else context_len
+    min_events = warm + 1  # need at least one future step
 
     eligible = []
     for times, values, bands in dataset.events_per_file:
@@ -730,11 +851,22 @@ def main():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print("Using device:", device)
 
-    model, context_len, n_bands = load_9band_model(args.ckpt, device)
+    (
+        model,
+        context_len,
+        n_bands,
+        context_window_days,
+        max_context_len,
+    ) = load_9band_model(args.ckpt, device)
+
+    window_mode = context_window_days is not None
+    seed_len = max_context_len if window_mode else context_len
 
     eval_dataset, means, stds = make_eval_dataset(
         args=args,
         context_len=context_len,
+        context_window_days=context_window_days,
+        max_context_len=max_context_len if window_mode else None,
     )
 
     print("Dataset files with events:", len(eval_dataset.events_per_file))
@@ -744,6 +876,7 @@ def main():
         context_len=context_len,
         n_future_steps=args.n_future_steps,
         n_series=args.n_series,
+        seed_len=seed_len,
     )
 
     if not series:
@@ -773,6 +906,9 @@ def main():
             means=means,
             stds=stds,
             teacher_forced=False,
+            window_mode=window_mode,
+            context_window_days=context_window_days,
+            max_context_len=max_context_len,
         )
         rollouts.append(rollout)
 
@@ -790,6 +926,9 @@ def main():
                 means=means,
                 stds=stds,
                 teacher_forced=True,
+                window_mode=window_mode,
+                context_window_days=context_window_days,
+                max_context_len=max_context_len,
             )
             tf_rollouts.append(tf_rollout)
 
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index ec25c7ec..e8754f4e 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -183,6 +183,20 @@ def main(args, rank, world_size, local_rank, device):
     CONTEXT_LEN = 5 #3
     HIDDEN_CHANNELS = 64
 
+    # Time-window context mode. When CONTEXT_WINDOW_DAYS is not None, the dataset
+    # selects context by a trailing lookback in days (all detections within the
+    # last CONTEXT_WINDOW_DAYS), padded to MAX_CONTEXT_LEN with a per-event
+    # validity flag, instead of a fixed count of CONTEXT_LEN events. This gives
+    # the model real time evolution instead of a single dense night. Set these
+    # from the plot_observation_histograms.py time-window sweep. Leave
+    # CONTEXT_WINDOW_DAYS = None to use the legacy fixed-count context.
+    CONTEXT_WINDOW_DAYS = 2.0
+    MAX_CONTEXT_LEN = 12
+    # In window mode the model's first layer is sized by the padded width.
+    WRAPPER_CONTEXT_LEN = (
+        MAX_CONTEXT_LEN if CONTEXT_WINDOW_DAYS is not None else CONTEXT_LEN
+    )
+
     # Multi-step rollout training config (scheduled sampling). n_rollout_steps=1
     # falls back to the standard single-step teacher-forced training.
     n_rollout_steps = args.n_rollout_steps
@@ -276,11 +290,12 @@ def main(args, rank, world_size, local_rank, device):
 
         model = ScalarTemporalConditionedLodeRunner_9band(
             backbone=backbone,
-            context_len=CONTEXT_LEN,
+            context_len=WRAPPER_CONTEXT_LEN,
             n_bands=N_BANDS,
             image_size=model_args["image_size"],
             backbone_channels=8,
             hidden=HIDDEN_CHANNELS,
+            context_window_days=CONTEXT_WINDOW_DAYS,
         ).to(device)
 
         # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
@@ -406,6 +421,7 @@ def main(args, rank, world_size, local_rank, device):
         print("band_stds:", band_stds)
 
     train_dataset = Kilonova_lc_scalar_context_DataSet_9band(
+        N_imgs=100,
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
@@ -414,9 +430,12 @@ def main(args, rank, world_size, local_rank, device):
         means=band_means,
         stds=band_stds,
         n_rollout_steps=n_rollout_steps,
+        context_window_days=CONTEXT_WINDOW_DAYS,
+        max_context_len=MAX_CONTEXT_LEN,
     )
 
     val_dataset = Kilonova_lc_scalar_context_DataSet_9band(
+        N_imgs=100,
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
@@ -425,6 +444,8 @@ def main(args, rank, world_size, local_rank, device):
         means=band_means,
         stds=band_stds,
         n_rollout_steps=n_rollout_steps,
+        context_window_days=CONTEXT_WINDOW_DAYS,
+        max_context_len=MAX_CONTEXT_LEN,
     )
 
 
@@ -575,6 +596,8 @@ def main(args, rank, world_size, local_rank, device):
                     "backbone_channels": 8,
                     "hidden": HIDDEN_CHANNELS,
                     "n_rollout_steps": n_rollout_steps,
+                    "context_window_days": CONTEXT_WINDOW_DAYS,
+                    "max_context_len": MAX_CONTEXT_LEN,
                 },
                 new_chkpt_path,
             )
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index d7b350e9..5dae2748 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -387,12 +387,17 @@ def __init__(
         means: np.ndarray = None,
         stds: np.ndarray = None,
         n_rollout_steps: int = 1,
+        context_window_days: float = None,
+        max_context_len: int = None,
     ) -> None:
         """Initialize the dataset and build the merged-event sample index.
 
         Args:
             N_imgs (int): Number of light-curve files to sample; 0 uses all.
-            context_len (int): Number of context events per sample.
+            context_len (int): Number of context events per sample. In the
+                default fixed-count mode this is the exact context length. It is
+                ignored when ``context_window_days`` is set (window mode uses
+                ``max_context_len`` for the padded width instead).
             band_keys (tuple[str, ...]): Keys of the bands to load. Their order
                 defines the band index used in the one-hot encoding and target.
             value_col (int): Column index of the value to load per band.
@@ -409,7 +414,18 @@ def __init__(
                 ``(x, target, mask, Dt)`` tuple. When >1 it returns a rollout
                 tuple carrying the initial context window plus the next
                 ``n_rollout_steps`` true events, for scheduled-sampling /
-                multi-step rollout training (see ``__getitem__``).
+                multi-step rollout training (see ``__getitem__``). Only supported
+                in fixed-count mode (``context_window_days`` is None).
+            context_window_days (float): If set, switches to **time-window
+                mode**: the context of each sample is every detection within this
+                many days before the target event, padded to ``max_context_len``
+                with a per-event validity flag (see ``_getitem_window``). If None
+                (default) the dataset uses the legacy fixed-count context of
+                exactly ``context_len`` events.
+            max_context_len (int): Padded context width used in time-window mode.
+                Windows with more real events than this keep only the most recent
+                ``max_context_len``; windows with fewer are zero-padded. Defaults
+                to ``context_len`` when not given. Unused in fixed-count mode.
         """
         # FIXME: hardcoded scratch path. Should be passed in as an argument so
         # this dataset does not depend on a user-specific filesystem location.
@@ -447,6 +463,36 @@ def __init__(
             )
         self.n_rollout_steps = n_rollout_steps
 
+        # Time-window context mode. When context_window_days is set, the context
+        # is selected by a trailing lookback in days and padded to
+        # max_context_len with a validity flag, instead of a fixed event count.
+        self.context_window_days = context_window_days
+        self.window_mode = context_window_days is not None
+
+        if self.window_mode:
+            self.max_context_len = (
+                max_context_len if max_context_len is not None else context_len
+            )
+            if self.max_context_len < 1:
+                raise ValueError(
+                    f"max_context_len must be >= 1, got {self.max_context_len}"
+                )
+            if context_window_days <= 0:
+                raise ValueError(
+                    "context_window_days must be positive, got "
+                    f"{context_window_days}"
+                )
+            if n_rollout_steps > 1:
+                # Multi-step rollout window-sliding is not yet wired for the
+                # time-window selection rule (single-step first); fail loudly
+                # rather than silently mixing the two.
+                raise NotImplementedError(
+                    "time-window context (context_window_days) is currently "
+                    "only supported with n_rollout_steps=1."
+                )
+        else:
+            self.max_context_len = context_len
+
         if means is None:
             raise ValueError(
                 "means must be provided for per-band normalization. "
@@ -539,9 +585,20 @@ def __init__(
             )
 
             n_events = times.shape[0]
-            max_start = n_events - context_len - 1
-            for startIDX in range(max_start + 1):
-                self.samples.append((len(self.events_per_file) - 1, startIDX))
+            file_idx = len(self.events_per_file) - 1
+
+            if self.window_mode:
+                # One sample per event after the first: the target is event
+                # target_idx and the context is the trailing window ending at
+                # target_idx - 1 (always non-empty, so short curves contribute).
+                for target_idx in range(1, n_events):
+                    self.samples.append((file_idx, target_idx))
+            else:
+                # Legacy fixed-count windows: startIDX indexes the window start,
+                # target is startIDX + context_len.
+                max_start = n_events - context_len - 1
+                for startIDX in range(max_start + 1):
+                    self.samples.append((file_idx, startIDX))
 
     def __len__(self) -> int:
         """Return the number of samples in the dataset."""
@@ -561,6 +618,9 @@ def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
         Returns:
             tuple[torch.Tensor, ...]: Single-step or rollout sample.
         """
+        if self.window_mode:
+            return self._getitem_window(index)
+
         if self.n_rollout_steps > 1:
             return self._getitem_rollout(index)
 
@@ -633,6 +693,86 @@ def _getitem_single(
 
         return x, target, mask, Dt
 
+    def _getitem_window(
+        self, index: int
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return the (input, target, mask, Dt) tuple in time-window mode.
+
+        The context is every detection within ``context_window_days`` before the
+        target event, padded to ``max_context_len`` with a per-event validity
+        flag. The per-event feature gains a ``valid`` channel so the flattened
+        input carries the padding mask itself (the model does no masking):
+        ``[value, rel_t, valid, one_hot_band(n_bands)]``, width ``3 + n_bands``.
+
+        Real events fill the leading rows in time order (oldest first, matching
+        the fixed-count layout), with ``rel_t`` relative to the first *real*
+        event in the window; padded rows are all-zero with ``valid = 0``.
+
+        Args:
+            index (int): Sample index (maps to a (file_idx, target_idx) pair).
+
+        Returns:
+            x (torch.Tensor): Flattened padded context, shape
+                [max_context_len * (3 + n_bands)].
+            target (torch.Tensor): Normalized value per band, shape [n_bands];
+                only the observed band is meaningful.
+            mask (torch.Tensor): Float mask, shape [n_bands]; 1.0 for the
+                observed target band, 0.0 elsewhere.
+            Dt (torch.Tensor): Lead time from the most recent context event to
+                the target event.
+        """
+        file_idx, target_idx = self.samples[index]
+        times, values, bands = self.events_per_file[file_idx]
+
+        # Anchor the trailing window on the event immediately before the target
+        # (the most recent observation). Select all earlier events within the
+        # window, then keep the most recent max_context_len if there are more.
+        anchor_t = times[target_idx - 1]
+        lo = anchor_t - self.context_window_days
+
+        prior_t = times[:target_idx]
+        in_window = prior_t >= lo
+        sel_idx = np.nonzero(in_window)[0]
+        if sel_idx.shape[0] > self.max_context_len:
+            sel_idx = sel_idx[-self.max_context_len:]
+
+        ctx_t = times[sel_idx]
+        ctx_v = values[sel_idx]
+        ctx_b = bands[sel_idx]
+        n_real = sel_idx.shape[0]
+
+        # rel_t relative to the first real event in the window (same convention
+        # as the fixed-count path, which uses the window's first event).
+        rel_t = (ctx_t - ctx_t[0]).astype(np.float32)
+
+        # Padded per-event array: [value, rel_t, valid, one_hot_band].
+        per_event = np.zeros(
+            (self.max_context_len, 3 + self.n_channels), dtype=np.float32
+        )
+        per_event[:n_real, 0] = ctx_v
+        per_event[:n_real, 1] = rel_t
+        per_event[:n_real, 2] = 1.0  # validity flag for real events
+        per_event[np.arange(n_real), 3 + ctx_b] = 1.0
+
+        x = torch.tensor(per_event.reshape(-1), dtype=torch.float32)
+
+        # Target is the event at target_idx, in a per-band vector + mask.
+        target_band = int(bands[target_idx])
+        target = np.zeros(self.n_channels, dtype=np.float32)
+        mask = np.zeros(self.n_channels, dtype=np.float32)
+        target[target_band] = values[target_idx]
+        mask[target_band] = 1.0
+
+        target = torch.tensor(target, dtype=torch.float32)
+        mask = torch.tensor(mask, dtype=torch.float32)
+
+        Dt = torch.tensor(
+            times[target_idx] - times[target_idx - 1],
+            dtype=torch.float32,
+        )
+
+        return x, target, mask, Dt
+
     def _getitem_rollout(
         self, index: int
     ) -> tuple[
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index 8f416597..c163ecb8 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -481,8 +481,19 @@ def __init__(
         image_size: tuple[int, int] = (1120, 400),
         backbone_channels: int = 8,
         hidden: int = 64,
+        context_window_days: float = None,
     ) -> None:
-        """Initialize conditioner and output-head around the backbone."""
+        """Initialize conditioner and output-head around the backbone.
+
+        Args:
+            context_len (int): Number of context events per sample. In
+                time-window mode this is the padded width (``max_context_len``).
+            context_window_days (float): When set, the dataset selects context
+                by a trailing time window and pads it with a per-event validity
+                flag, so each event carries an extra ``valid`` feature and the
+                per-event width is ``3 + n_bands`` instead of ``2 + n_bands``.
+                When None (default), the legacy fixed-count layout is used.
+        """
         super().__init__()
 
         self.backbone = backbone
@@ -490,12 +501,14 @@ def __init__(
         self.n_bands = n_bands
         self.image_size = image_size
         self.backbone_channels = backbone_channels
-
-        # Dataset x layout, flattened per event:
-        #   [value, rel_t, one_hot_band(n_bands)] * context_len
-        #
-        # input_dim = context_len * (2 + n_bands)
-        input_dim = context_len * (2 + n_bands)
+        self.context_window_days = context_window_days
+
+        # Dataset x layout, flattened per event. Fixed-count mode:
+        #   [value, rel_t, one_hot_band(n_bands)] * context_len   -> 2 + n_bands
+        # Time-window mode adds a validity flag so padding is carried in x:
+        #   [value, rel_t, valid, one_hot_band(n_bands)] * context_len -> 3 + n_bands
+        per_event_width = 3 + n_bands if context_window_days is not None else 2 + n_bands
+        input_dim = context_len * per_event_width
 
         # Maps the scalar temporal event stream into the pseudo-channels
         # expected by the pretrained LodeRunner backbone.
diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index 8c04b7c2..1af2c726 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -477,6 +477,15 @@ def load_direct_loderunner_checkpoint_9band(
     saved_model_args = checkpoint_data.get("model_args", model_args)
     context_len = checkpoint_data.get("context_len", 5)
 
+    # Time-window context mode: the model's first layer is sized by the padded
+    # width (max_context_len) and each event carries an extra validity flag.
+    # Falls through to None for legacy fixed-count checkpoints, preserving the
+    # original sizing.
+    context_window_days = checkpoint_data.get("context_window_days", None)
+    if context_window_days is not None:
+        # In window mode the padded context width drives input_dim.
+        context_len = checkpoint_data.get("max_context_len", context_len)
+
     backbone = LodeRunner(**saved_model_args).to(device)
 
     model = ScalarTemporalConditionedLodeRunner_9band(
@@ -486,6 +495,7 @@ def load_direct_loderunner_checkpoint_9band(
         image_size=saved_model_args["image_size"],
         backbone_channels=checkpoint_data.get("backbone_channels", 8),
         hidden=checkpoint_data.get("hidden", 64),
+        context_window_days=context_window_days,
     ).to(device)
 
     state_dict = checkpoint_data["model_state_dict"]

From 47d913bd7394a80e34d0beb10cd663c435759b72 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 14 Aug 2026 15:31:52 -0600
Subject: [PATCH 30/66] bug fix

---
 .../harnesses/KN_loderunner/train_LodeRunner_ddp.py  | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index e8754f4e..816c5460 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -1,5 +1,6 @@
 import os
 import time
+import random
 import argparse
 import numpy as np
 import torch
@@ -420,6 +421,17 @@ def main(args, rank, world_size, local_rank, device):
         print("band_means:", band_means)
         print("band_stds:", band_stds)
 
+    # The 9-band dataset selects its N_imgs files with an unseeded RNG and each
+    # DDP rank builds its own dataset. With N_imgs>0 that would give every rank a
+    # DIFFERENT random file subset, hence different sample counts and different
+    # per-rank batch counts, so ranks desync and hang at gradient all-reduce
+    # (NCCL watchdog timeout). Seed both RNGs to the SAME value on every rank so
+    # all ranks pick the identical file subset. (With N_imgs=0 all files are used
+    # and this is moot, but seeding is harmless.)
+    DATA_SEED = 42
+    np.random.seed(DATA_SEED)
+    random.seed(DATA_SEED)
+
     train_dataset = Kilonova_lc_scalar_context_DataSet_9band(
         N_imgs=100,
         context_len=CONTEXT_LEN,

From 45664f1909d921dd8cf6dea9c20dc5dcf897d638 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 14 Aug 2026 16:30:36 -0600
Subject: [PATCH 31/66] scheduled sampling fix

---
 .../KN_loderunner/train_LodeRunner_ddp.py     |   3 +
 src/yoke/datasets/kilonova_dataset.py         | 119 ++++++-
 src/yoke/utils/training/epoch/loderunner.py   | 326 ++++++++++++++----
 3 files changed, 379 insertions(+), 69 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 816c5460..753fec10 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -543,6 +543,9 @@ def main(args, rank, world_size, local_rank, device):
                 world_size=world_size,
                 n_bands=N_BANDS,
                 teacher_forcing_ratio=teacher_forcing_ratio,
+                window_mode=CONTEXT_WINDOW_DAYS is not None,
+                context_window_days=CONTEXT_WINDOW_DAYS,
+                max_context_len=MAX_CONTEXT_LEN,
             )
         else:
             #train_DDP_loderunner_epoch(
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 5dae2748..3a2b5fdc 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -482,14 +482,6 @@ def __init__(
                     "context_window_days must be positive, got "
                     f"{context_window_days}"
                 )
-            if n_rollout_steps > 1:
-                # Multi-step rollout window-sliding is not yet wired for the
-                # time-window selection rule (single-step first); fail loudly
-                # rather than silently mixing the two.
-                raise NotImplementedError(
-                    "time-window context (context_window_days) is currently "
-                    "only supported with n_rollout_steps=1."
-                )
         else:
             self.max_context_len = context_len
 
@@ -619,6 +611,8 @@ def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
             tuple[torch.Tensor, ...]: Single-step or rollout sample.
         """
         if self.window_mode:
+            if self.n_rollout_steps > 1:
+                return self._getitem_window_rollout(index)
             return self._getitem_window(index)
 
         if self.n_rollout_steps > 1:
@@ -850,3 +844,112 @@ def _getitem_rollout(
             torch.tensor(future_dt, dtype=torch.float32),
             torch.tensor(future_valid, dtype=torch.float32),
         )
+
+    def _getitem_window_rollout(
+        self, index: int
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        """Return a multi-step rollout sample in time-window context mode.
+
+        Combines the time-window context selection of :meth:`_getitem_window`
+        with the multi-step future supervision of :meth:`_getitem_rollout`, so
+        scheduled-sampling rollout training uses the SAME trailing-W-day context
+        the model sees at inference (``get_rollout_from_stream`` /
+        ``_select_window`` in ``plot_pred_diagnostics_9band.py``).
+
+        The seed context is the trailing ``context_window_days`` window ending at
+        the anchor event (``target_idx - 1``), capped to the most recent
+        ``max_context_len`` events, and returned **padded** to ``max_context_len``
+        with a per-event validity flag. Unlike :meth:`_getitem_rollout`, the
+        context times are returned as **absolute stream times** (not made
+        relative here) so the training loop can append the true future event
+        times and re-select the W-day window each step, then make times relative
+        per step exactly as inference does.
+
+        Args:
+            index (int): Sample index (maps to a (file_idx, target_idx) pair).
+
+        Returns:
+            ctx_v (torch.Tensor): Padded context values, shape [max_context_len];
+                padded rows are 0.
+            ctx_t (torch.Tensor): Padded ABSOLUTE context times, shape
+                [max_context_len]; padded rows are 0.
+            ctx_b (torch.Tensor): Padded context band indices (long), shape
+                [max_context_len]; padded rows are 0.
+            ctx_valid (torch.Tensor): 1.0 for real seed events, 0.0 for padding,
+                shape [max_context_len].
+            future_v (torch.Tensor): Normalized true value of each future event,
+                shape [n_rollout_steps]; padded steps are 0.
+            future_b (torch.Tensor): Band index of each future event (long),
+                shape [n_rollout_steps]; padded steps are 0.
+            future_dt (torch.Tensor): Lead time from the previous event to each
+                future event, shape [n_rollout_steps]; padded steps are 0.
+            future_valid (torch.Tensor): 1.0 for real future events, 0.0 for
+                padded steps, shape [n_rollout_steps].
+        """
+        file_idx, target_idx = self.samples[index]
+        times, values, bands = self.events_per_file[file_idx]
+
+        # Seed context: trailing W-day window ending at the anchor event
+        # (target_idx - 1), capped to the most recent max_context_len. Identical
+        # selection to _getitem_window.
+        anchor_t = times[target_idx - 1]
+        lo = anchor_t - self.context_window_days
+
+        prior_t = times[:target_idx]
+        in_window = prior_t >= lo
+        sel_idx = np.nonzero(in_window)[0]
+        if sel_idx.shape[0] > self.max_context_len:
+            sel_idx = sel_idx[-self.max_context_len:]
+
+        n_real = sel_idx.shape[0]
+
+        # Padded seed context; times kept ABSOLUTE so the loop can append true
+        # future times and re-window (the loop makes them relative per step).
+        ctx_v = np.zeros(self.max_context_len, dtype=np.float32)
+        ctx_t = np.zeros(self.max_context_len, dtype=np.float32)
+        ctx_b = np.zeros(self.max_context_len, dtype=np.int64)
+        ctx_valid = np.zeros(self.max_context_len, dtype=np.float32)
+
+        ctx_v[:n_real] = values[sel_idx]
+        ctx_t[:n_real] = times[sel_idx]
+        ctx_b[:n_real] = bands[sel_idx]
+        ctx_valid[:n_real] = 1.0
+
+        # Future events target_idx … target_idx + n - 1, padded past the stream
+        # end and flagged invalid (same tail handling as _getitem_rollout).
+        n = self.n_rollout_steps
+        future_v = np.zeros(n, dtype=np.float32)
+        future_b = np.zeros(n, dtype=np.int64)
+        future_dt = np.zeros(n, dtype=np.float32)
+        future_valid = np.zeros(n, dtype=np.float32)
+
+        n_events = times.shape[0]
+        for step in range(n):
+            t_idx = target_idx + step
+            if t_idx >= n_events:
+                break
+
+            future_v[step] = values[t_idx]
+            future_b[step] = bands[t_idx]
+            future_dt[step] = times[t_idx] - times[t_idx - 1]
+            future_valid[step] = 1.0
+
+        return (
+            torch.tensor(ctx_v, dtype=torch.float32),
+            torch.tensor(ctx_t, dtype=torch.float32),
+            torch.tensor(ctx_b, dtype=torch.long),
+            torch.tensor(ctx_valid, dtype=torch.float32),
+            torch.tensor(future_v, dtype=torch.float32),
+            torch.tensor(future_b, dtype=torch.long),
+            torch.tensor(future_dt, dtype=torch.float32),
+            torch.tensor(future_valid, dtype=torch.float32),
+        )
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 7d843688..36ac303e 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -735,6 +735,182 @@ def _rollout_pass_9band(
     return per_sample_loss, total_loss
 
 
+def _rollout_pass_9band_window(
+    ctx_v: torch.Tensor,
+    ctx_t: torch.Tensor,
+    ctx_b: torch.Tensor,
+    ctx_valid: torch.Tensor,
+    future_v: torch.Tensor,
+    future_b: torch.Tensor,
+    future_dt: torch.Tensor,
+    future_valid: torch.Tensor,
+    model: torch.nn.Module,
+    loss_fn: torch.nn.Module,
+    n_bands: int,
+    context_window_days: float,
+    max_context_len: int,
+    teacher_forcing_ratio: float,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Unroll the 9-band model over a batch of rollouts in time-window mode.
+
+    The time-window analogue of :func:`_rollout_pass_9band`. Instead of a
+    fixed-count window slid by drop-oldest, this keeps a **growing** buffer of
+    absolute-time events and, at each step, re-selects the trailing
+    ``context_window_days`` window (capped to the most recent
+    ``max_context_len``), padded with a per-event validity flag. This mirrors the
+    inference rollout (``get_rollout_from_stream`` + ``_select_window`` in
+    ``plot_pred_diagnostics_9band.py``) so the model is trained on exactly the
+    context layout it is evaluated on: per-event width ``3 + n_bands``
+    (``[value, rel_t, valid, one_hot]``), ``rel_t`` relative to the window's
+    first real event.
+
+    Because the buffer is time-sorted and the window is "events within W days of
+    the most recent event, capped to the most recent ``max_context_len``", the
+    selected window is always a contiguous suffix, so re-selection is a batched
+    start-index + gather (no per-row Python loop, no scatter).
+
+    Args:
+        ctx_v (torch.Tensor): Padded seed context values [B, seed_width].
+        ctx_t (torch.Tensor): Padded seed ABSOLUTE context times [B, seed_width].
+        ctx_b (torch.Tensor): Padded seed context band indices [B, seed_width].
+        ctx_valid (torch.Tensor): Seed validity mask [B, seed_width].
+        future_v (torch.Tensor): True future values [B, n_rollout_steps].
+        future_b (torch.Tensor): Future band indices [B, n_rollout_steps].
+        future_dt (torch.Tensor): Future lead times [B, n_rollout_steps].
+        future_valid (torch.Tensor): Valid-step mask [B, n_rollout_steps].
+        model (torch.nn.Module): The 9-band wrapper model.
+        loss_fn (torch.nn.Module): Elementwise loss (reduction='none').
+        n_bands (int): Number of bands.
+        context_window_days (float): Trailing lookback W in days.
+        max_context_len (int): Padded context width M.
+        teacher_forcing_ratio (float): Probability of feeding the true value back
+            at each step (1.0 = fully teacher-forced, 0.0 = fully free-running).
+        device (torch.device): Compute device.
+
+    Returns:
+        per_sample_loss (torch.Tensor): Mean rollout loss per sample [B].
+        total_loss (torch.Tensor): Scalar mean loss over all valid steps.
+    """
+    B = ctx_v.shape[0]
+    seed_width = ctx_v.shape[1]
+    n_steps = future_v.shape[1]
+    W = float(context_window_days)
+    M = int(max_context_len)
+
+    # Growing buffer: the seed (<= M real events, left-packed) plus at most one
+    # appended event per rollout step. Left-packed and time-sorted throughout.
+    C = M + n_steps
+
+    buf_v = torch.zeros(B, C, device=device, dtype=ctx_v.dtype)
+    buf_t = torch.zeros(B, C, device=device, dtype=ctx_t.dtype)
+    buf_b = torch.zeros(B, C, device=device, dtype=torch.long)
+
+    buf_v[:, :seed_width] = ctx_v
+    buf_t[:, :seed_width] = ctx_t
+    buf_b[:, :seed_width] = ctx_b
+
+    # Number of real events currently in each row's buffer (>= 1 in window mode,
+    # since the anchor event is always inside its own trailing window).
+    count = ctx_valid.sum(dim=1).long()  # [B]
+
+    batch_arange = torch.arange(B, device=device)
+    pos = torch.arange(C, device=device).unsqueeze(0)  # [1, C]
+    out_pos = torch.arange(M, device=device).unsqueeze(0)  # [1, M]
+
+    # Kept for API compatibility with the LodeRunner-style wrapper.
+    in_vars = torch.arange(8, device=device)
+    out_vars = torch.arange(8, device=device)
+
+    step_losses = []  # [B] per step
+    step_valid = []  # [B] per step
+
+    for step in range(n_steps):
+        # Most recent real event time per row (left-packed => index count - 1).
+        last_idx = (count - 1).clamp(min=0)
+        last_t = buf_t.gather(1, last_idx.unsqueeze(1)).squeeze(1)  # [B]
+        lo = last_t - W
+
+        # Trailing W-day window is a contiguous suffix of the sorted buffer.
+        real_mask = pos < count.unsqueeze(1)  # [B, C]
+        ge_mask = buf_t >= lo.unsqueeze(1)  # [B, C]
+        n_in_window = (real_mask & ge_mask).sum(dim=1)  # [B]
+        win_len = torch.clamp(n_in_window, max=M)  # [B]
+        start = count - win_len  # [B]
+
+        # Source buffer index for each padded output position p in [0, M).
+        src_idx = start.unsqueeze(1) + out_pos  # [B, M]
+        valid_out = out_pos < win_len.unsqueeze(1)  # [B, M] bool
+        src_idx_c = src_idx.clamp(max=C - 1)
+
+        gathered_v = buf_v.gather(1, src_idx_c)  # [B, M]
+        gathered_t = buf_t.gather(1, src_idx_c)  # [B, M]
+        gathered_b = buf_b.gather(1, src_idx_c)  # [B, M]
+
+        # rel_t relative to the window's first real event (buf_t[start]).
+        first_t = buf_t.gather(1, start.clamp(max=C - 1).unsqueeze(1))  # [B, 1]
+        rel_t = gathered_t - first_t  # [B, M]
+
+        valid_f = valid_out.to(buf_v.dtype)  # [B, M]
+        val_col = gathered_v * valid_f
+        rel_col = rel_t * valid_f
+
+        band_onehot = torch.zeros(
+            B, M, n_bands, device=device, dtype=buf_v.dtype
+        )
+        band_onehot.scatter_(2, gathered_b.unsqueeze(-1), 1.0)
+        band_onehot = band_onehot * valid_f.unsqueeze(-1)
+
+        per_event = torch.cat(
+            [
+                val_col.unsqueeze(-1),
+                rel_col.unsqueeze(-1),
+                valid_f.unsqueeze(-1),
+                band_onehot,
+            ],
+            dim=-1,
+        )  # [B, M, 3 + n_bands]
+        x_step = per_event.reshape(B, -1)
+
+        Dt = future_dt[:, step]
+        pred_all = model(x_step, in_vars, out_vars, Dt)  # [B, n_bands]
+
+        tgt_band = future_b[:, step]
+        pred_obs = pred_all[batch_arange, tgt_band]  # [B]
+        true_obs = future_v[:, step]  # [B]
+        valid = future_valid[:, step]  # [B]
+
+        step_loss = loss_fn(pred_obs, true_obs) * valid
+        step_losses.append(step_loss)
+        step_valid.append(valid)
+
+        # Scheduled sampling: choose true vs own (detached) prediction per sample.
+        use_true = torch.rand(B, device=device) < teacher_forcing_ratio
+        fed = torch.where(use_true, true_obs, pred_obs.detach())
+
+        # Append the new event (following the true time/band schedule) at the
+        # left-packed write position and grow the count. For padded steps the
+        # appended event is spurious but harmless: those steps' losses are masked
+        # and all later steps for that row are padded too. The W-day re-selection
+        # at the next step handles dropping stale events (never drop-oldest here),
+        # matching the inference rollout which appends to its running context.
+        new_t = last_t + Dt
+
+        write_pos = count.clamp(max=C - 1).unsqueeze(1)  # [B, 1]
+        buf_v.scatter_(1, write_pos, fed.unsqueeze(1))
+        buf_t.scatter_(1, write_pos, new_t.unsqueeze(1))
+        buf_b.scatter_(1, write_pos, tgt_band.unsqueeze(1))
+        count = torch.clamp(count + 1, max=C)
+
+    step_losses = torch.stack(step_losses, dim=1)  # [B, n_steps]
+    step_valid = torch.stack(step_valid, dim=1)  # [B, n_steps]
+
+    per_sample_loss = step_losses.sum(dim=1) / (step_valid.sum(dim=1) + 1e-8)
+    total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
+
+    return per_sample_loss, total_loss
+
+
 def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
     training_data: torch.utils.data.DataLoader,
     validation_data: torch.utils.data.DataLoader,
@@ -753,6 +929,9 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
     world_size: int,
     n_bands: int = 9,
     teacher_forcing_ratio: float = 1.0,
+    window_mode: bool = False,
+    context_window_days: float = None,
+    max_context_len: int = None,
 ) -> None:
     """Multi-step rollout DDP epoch for the masked 9-band scalar temporal model.
 
@@ -766,14 +945,95 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
 
     Expected dataset output (per sample, from
     ``Kilonova_lc_scalar_context_DataSet_9band`` with ``n_rollout_steps > 1``):
-        ctx_v, ctx_t, ctx_b:  [B, context_len]
-        future_v, future_b, future_dt, future_valid:  [B, n_rollout_steps]
+        fixed-count mode (``window_mode=False``), 7-tuple:
+            ctx_v, ctx_t, ctx_b:  [B, context_len]
+            future_v, future_b, future_dt, future_valid:  [B, n_rollout_steps]
+        time-window mode (``window_mode=True``), 8-tuple (adds ``ctx_valid``):
+            ctx_v, ctx_t, ctx_b, ctx_valid:  [B, max_context_len]
+            future_v, future_b, future_dt, future_valid:  [B, n_rollout_steps]
 
     Args:
         n_bands (int): Number of bands the model predicts.
         teacher_forcing_ratio (float): Per-epoch probability of feeding the true
             value back at each rollout step during training.
+        window_mode (bool): If True, consume the 8-tuple time-window sample and
+            unroll with the trailing-W-day context re-selection
+            (:func:`_rollout_pass_9band_window`). If False (default), the legacy
+            fixed-count 7-tuple path (:func:`_rollout_pass_9band`).
+        context_window_days (float): Trailing lookback W in days. Required when
+            ``window_mode`` is True.
+        max_context_len (int): Padded context width M. Required when
+            ``window_mode`` is True.
     """
+    def _run_pass(
+        data: tuple[torch.Tensor, ...], ratio: float
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Unpack a batch (7- or 8-tuple), move to device, run the rollout."""
+        if window_mode:
+            (
+                ctx_v,
+                ctx_t,
+                ctx_b,
+                ctx_valid,
+                future_v,
+                future_b,
+                future_dt,
+                future_valid,
+            ) = data
+            ctx_valid = ctx_valid.to(device, non_blocking=True)
+        else:
+            (
+                ctx_v,
+                ctx_t,
+                ctx_b,
+                future_v,
+                future_b,
+                future_dt,
+                future_valid,
+            ) = data
+
+        ctx_v = ctx_v.to(device, non_blocking=True)
+        ctx_t = ctx_t.to(device, non_blocking=True)
+        ctx_b = ctx_b.to(device, non_blocking=True)
+        future_v = future_v.to(device, non_blocking=True)
+        future_b = future_b.to(device, non_blocking=True)
+        future_dt = future_dt.to(torch.float32).to(device, non_blocking=True)
+        future_valid = future_valid.to(device, non_blocking=True)
+
+        if window_mode:
+            return _rollout_pass_9band_window(
+                ctx_v=ctx_v,
+                ctx_t=ctx_t,
+                ctx_b=ctx_b,
+                ctx_valid=ctx_valid,
+                future_v=future_v,
+                future_b=future_b,
+                future_dt=future_dt,
+                future_valid=future_valid,
+                model=model,
+                loss_fn=loss_fn,
+                n_bands=n_bands,
+                context_window_days=context_window_days,
+                max_context_len=max_context_len,
+                teacher_forcing_ratio=ratio,
+                device=device,
+            )
+
+        return _rollout_pass_9band(
+            ctx_v=ctx_v,
+            ctx_t=ctx_t,
+            ctx_b=ctx_b,
+            future_v=future_v,
+            future_b=future_b,
+            future_dt=future_dt,
+            future_valid=future_valid,
+            model=model,
+            loss_fn=loss_fn,
+            n_bands=n_bands,
+            teacher_forcing_ratio=ratio,
+            device=device,
+        )
+
     train_rcrd_filename = train_rcrd_filename.replace(
         "<epochIDX>",
         f"{epochIDX:04d}",
@@ -789,33 +1049,10 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
             if trainbatch_ID >= num_train_batches:
                 break
 
-            ctx_v, ctx_t, ctx_b, future_v, future_b, future_dt, future_valid = (
-                data
-            )
-
-            ctx_v = ctx_v.to(device, non_blocking=True)
-            ctx_t = ctx_t.to(device, non_blocking=True)
-            ctx_b = ctx_b.to(device, non_blocking=True)
-            future_v = future_v.to(device, non_blocking=True)
-            future_b = future_b.to(device, non_blocking=True)
-            future_dt = future_dt.to(torch.float32).to(device, non_blocking=True)
-            future_valid = future_valid.to(device, non_blocking=True)
-
             optimizer.zero_grad(set_to_none=True)
 
-            per_sample_loss, batch_loss = _rollout_pass_9band(
-                ctx_v=ctx_v,
-                ctx_t=ctx_t,
-                ctx_b=ctx_b,
-                future_v=future_v,
-                future_b=future_b,
-                future_dt=future_dt,
-                future_valid=future_valid,
-                model=model,
-                loss_fn=loss_fn,
-                n_bands=n_bands,
-                teacher_forcing_ratio=teacher_forcing_ratio,
-                device=device,
+            per_sample_loss, batch_loss = _run_pass(
+                data, teacher_forcing_ratio
             )
 
             batch_loss.backward()
@@ -852,41 +1089,8 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
                     if valbatch_ID >= num_val_batches:
                         break
 
-                    (
-                        ctx_v,
-                        ctx_t,
-                        ctx_b,
-                        future_v,
-                        future_b,
-                        future_dt,
-                        future_valid,
-                    ) = data
-
-                    ctx_v = ctx_v.to(device, non_blocking=True)
-                    ctx_t = ctx_t.to(device, non_blocking=True)
-                    ctx_b = ctx_b.to(device, non_blocking=True)
-                    future_v = future_v.to(device, non_blocking=True)
-                    future_b = future_b.to(device, non_blocking=True)
-                    future_dt = future_dt.to(torch.float32).to(
-                        device, non_blocking=True
-                    )
-                    future_valid = future_valid.to(device, non_blocking=True)
-
                     # Validation is always a pure free-running rollout.
-                    per_sample_loss, _ = _rollout_pass_9band(
-                        ctx_v=ctx_v,
-                        ctx_t=ctx_t,
-                        ctx_b=ctx_b,
-                        future_v=future_v,
-                        future_b=future_b,
-                        future_dt=future_dt,
-                        future_valid=future_valid,
-                        model=model,
-                        loss_fn=loss_fn,
-                        n_bands=n_bands,
-                        teacher_forcing_ratio=0.0,
-                        device=device,
-                    )
+                    per_sample_loss, _ = _run_pass(data, 0.0)
 
                     if rank == 0:
                         batch_records = np.column_stack(

From 5804d3967b8d3aee8c09d47f21cc9f910590a072 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 14 Aug 2026 16:59:29 -0600
Subject: [PATCH 32/66] use entire dataset

---
 applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 753fec10..292f52cb 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -433,7 +433,7 @@ def main(args, rank, world_size, local_rank, device):
     random.seed(DATA_SEED)
 
     train_dataset = Kilonova_lc_scalar_context_DataSet_9band(
-        N_imgs=100,
+        N_imgs=0,
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,
@@ -447,7 +447,7 @@ def main(args, rank, world_size, local_rank, device):
     )
 
     val_dataset = Kilonova_lc_scalar_context_DataSet_9band(
-        N_imgs=100,
+        N_imgs=0,
         context_len=CONTEXT_LEN,
         band_keys=BAND_KEYS,
         value_col=VALUE_COL,

From 9a11281707e47d126949936542340eec32421f0d Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Mon, 17 Aug 2026 12:25:45 -0600
Subject: [PATCH 33/66] update sweep script

---
 .../plot_observation_histograms.py            | 145 +++++++++++++++++-
 1 file changed, 144 insertions(+), 1 deletion(-)

diff --git a/applications/harnesses/KN_loderunner/plot_observation_histograms.py b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
index 5f06ff85..d5f9abd1 100644
--- a/applications/harnesses/KN_loderunner/plot_observation_histograms.py
+++ b/applications/harnesses/KN_loderunner/plot_observation_histograms.py
@@ -19,6 +19,14 @@
      9-band model bakes ``context_len`` into its first layer, so choosing a
      longer context means a retrain -- this panel shows the data cost before you
      pay for it.
+  5. A time-window context sweep: context size (real detections) vs the trailing
+     lookback ``W`` in days, for choosing ``context_window_days`` /
+     ``max_context_len`` in window mode.
+  6. Supervised lead-time ``Delta_t`` vs forecast horizon: the CDF of the
+     gap-to-next-event (``h=1``, exactly the ``Delta_t`` training supervises)
+     against the lead times reachable at larger event offsets. Its right tail is
+     the coverage ceiling -- forecasts asked for a longer ``Delta_t`` than the
+     ``h=1`` p99/max extrapolate beyond anything the model was trained on.
 
 Run directly, e.g.:
     python plot_observation_histograms.py
@@ -387,6 +395,129 @@ def plot_time_window_sweep(ax, event_times_per_file, window_grid):
     ax.grid(True, alpha=0.3)
 
 
+def dt_horizon_stats(event_times_per_file, horizons):
+    """Distribution of lead time Delta_t as a function of event horizon.
+
+    Training supervises, for each context, the jump to the target event. When
+    the target is the *immediate next* event (``horizon=1``) the supervised
+    Delta_t is exactly the consecutive-event gap. This function collects, for
+    each horizon ``h``, every reachable lead time ``times[i + h] - times[i]``
+    across all light curves. Comparing ``h=1`` (what training actually sees)
+    against the forecast horizons used at inference reveals whether the model is
+    ever supervised at the Delta_t it is later asked to extrapolate to.
+
+    Args:
+        event_times_per_file (list[np.ndarray]): Per-file merged, sorted,
+            file-relative detection times.
+        horizons (iterable[int]): Event offsets ``h`` to evaluate. ``h=1`` is
+            the training Delta_t (gap to next event).
+
+    Returns:
+        horizons (list[int]): The evaluated horizons.
+        raw (dict[int, np.ndarray]): For each horizon, all reachable lead times
+            (days) pooled over every light curve.
+    """
+    horizons = list(horizons)
+    raw = {h: [] for h in horizons}
+    for times in event_times_per_file:
+        n = times.shape[0]
+        for h in horizons:
+            if n > h:
+                raw[h].append(times[h:] - times[:-h])
+    raw = {
+        h: (np.concatenate(v) if v else np.zeros(0, dtype=np.float64))
+        for h, v in raw.items()
+    }
+    return horizons, raw
+
+
+def plot_dt_horizon(ax, event_times_per_file, horizons):
+    """CDF of supervised lead time Delta_t vs event horizon.
+
+    The ``h=1`` curve is the distribution of training Delta_t (gap to the next
+    event); its right tail is where the model stops being supervised. Larger-``h``
+    curves show how far ahead (in events) a target must be drawn to reach a given
+    lead time -- the basis for horizon-covering target sampling. Percentile lines
+    for ``h=1`` make the coverage ceiling explicit.
+    """
+    if len(event_times_per_file) == 0:
+        ax.set_visible(False)
+        return
+
+    horizons, raw = dt_horizon_stats(event_times_per_file, horizons)
+    cmap = plt.get_cmap("viridis")
+
+    for j, h in enumerate(horizons):
+        c = raw[h]
+        if c.size == 0:
+            continue
+        xs = np.sort(c)
+        ys = np.arange(1, xs.size + 1) / xs.size
+        label = f"h={h}" + (" (train Δt)" if h == 1 else "")
+        ax.plot(
+            xs,
+            ys,
+            color=cmap(j / max(1, len(horizons) - 1)),
+            linewidth=2.0 if h == 1 else 1.3,
+            label=label,
+        )
+
+    # Percentile markers for the training Delta_t (h=1): the coverage ceiling.
+    base = raw[horizons[0]]
+    if base.size:
+        for q in (95, 99):
+            pv = float(np.percentile(base, q))
+            ax.axvline(pv, color="tab:red", linestyle=":", linewidth=1, alpha=0.7)
+            ax.text(
+                pv,
+                0.02,
+                f" p{q}={pv:.1f}d",
+                rotation=90,
+                va="bottom",
+                ha="right",
+                fontsize=7,
+                color="tab:red",
+            )
+
+    ax.set_xlabel("Lead time Δt (days)")
+    ax.set_ylabel("Cumulative fraction of samples")
+    ax.set_ylim(0, 1.02)
+    ax.set_title(
+        "Supervised Δt vs forecast horizon\n"
+        "(h=1 is training Δt; right tail = uncovered)"
+    )
+    ax.legend(fontsize=8, loc="lower right")
+    ax.grid(True, alpha=0.3)
+
+
+def print_dt_horizon(event_times_per_file, horizons):
+    """Print the lead-time-vs-horizon percentile table to stdout."""
+    if len(event_times_per_file) == 0:
+        return
+
+    horizons, raw = dt_horizon_stats(event_times_per_file, horizons)
+
+    print("\nSupervised lead time Δt by event horizon (days):")
+    print(
+        f"{'horizon':>8} {'n':>10} {'median':>8} {'p90':>7} {'p95':>7} "
+        f"{'p99':>7} {'max':>7}"
+    )
+    for h in horizons:
+        c = raw[h]
+        if c.size == 0:
+            print(f"{h:>8} {0:>10}")
+            continue
+        print(
+            f"{h:>8} {c.size:>10} {np.median(c):>8.2f} "
+            f"{np.percentile(c, 90):>7.2f} {np.percentile(c, 95):>7.2f} "
+            f"{np.percentile(c, 99):>7.2f} {c.max():>7.2f}"
+        )
+    print(
+        "h=1 is the Delta_t training actually supervises (target = next event). "
+        "Its p99/max is the coverage ceiling; forecasts beyond it extrapolate."
+    )
+
+
 def print_summary(counts, include_upper_limits):
     """Print the per-band and overall counts to stdout."""
     det = counts["det_totals"]
@@ -511,6 +642,17 @@ def main():
         default=1.0,
         help="Step (days) between evaluated window lengths. Default 1.0.",
     )
+    parser.add_argument(
+        "--dt_horizons",
+        type=int,
+        nargs="+",
+        default=[1, 2, 3, 5, 8],
+        help=(
+            "Event offsets h for the supervised-Delta_t panel. h=1 is the "
+            "training Delta_t (gap to next event); larger h show how far ahead "
+            "a target must be drawn to reach a given lead time. Default 1 2 3 5 8."
+        ),
+    )
     parser.add_argument(
         "--out",
         type=str,
@@ -536,6 +678,7 @@ def main():
         args.window_step, args.window_max + 0.5 * args.window_step, args.window_step
     )
     print_time_window_sweep(counts["event_times_per_file"], window_grid)
+    print_dt_horizon(counts["event_times_per_file"], args.dt_horizons)
 
     fig, axes = plt.subplots(2, 3, figsize=(21, 11))
     plot_band_totals(axes[0, 0], counts, args.include_upper_limits)
@@ -543,7 +686,7 @@ def main():
     plot_per_band_hist(axes[0, 2], counts)
     plot_context_sweep(axes[1, 0], counts["total_det_per_curve"], args.sweep_max)
     plot_time_window_sweep(axes[1, 1], counts["event_times_per_file"], window_grid)
-    axes[1, 2].set_visible(False)
+    plot_dt_horizon(axes[1, 2], counts["event_times_per_file"], args.dt_horizons)
 
     fig.suptitle(
         f"KN light-curve observations ({counts['n_files']} light curves)",

From 6528524a02c40ad0e8b073afe5643c0b523de732 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 18 Aug 2026 10:45:01 -0600
Subject: [PATCH 34/66] Training with dense and realistic curves

---
 .../eval_dense_latetime_9band.py              | 488 ++++++++++++++++++
 .../KN_loderunner/make_kn_object_lists.py     | 170 ++++++
 .../KN_loderunner/train_LodeRunner_ddp.py     | 175 +++++--
 src/yoke/datasets/kilonova_dataset.py         | 175 ++++++-
 4 files changed, 946 insertions(+), 62 deletions(-)
 create mode 100644 applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
 create mode 100644 applications/harnesses/KN_loderunner/make_kn_object_lists.py

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
new file mode 100644
index 00000000..9e2a5c72
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -0,0 +1,488 @@
+"""Dense late-time evaluation for the 9-band scalar temporal LodeRunner.
+
+The model is trained on REALISTIC light curves (sparse, upper limits dropped),
+which contain almost no late-time detections because kilonovae fade below the
+detection limit. This eval measures how well the model, given a REALISTIC
+observing context, forecasts the LATE-TIME behavior -- scored against a DENSE
+companion set (the same objects, denser cadence, no limiting-mag cut, so all
+late-time points are real detections). Realistic and dense views of an object
+are paired by filename stem.
+
+For each held-out (test-split) object:
+  1. Build the model input from the realistic stream (trailing time-window
+     context ending at the last realistic detection), exactly as in training.
+  2. For every dense point in the late-time region (phase from the first
+     realistic detection greater than ``--late_time_cutoff_days``), ask the model
+     to predict all nine bands at that point's lead time and score the predicted
+     magnitude of the dense point's band against the dense truth.
+  3. Also sweep a smooth lead-time grid for a per-object forecast plot.
+
+Only detections are used: realistic upper limits are dropped (matching
+training); the dense set is all detections by construction.
+
+IMPORTANT (time frames): the training dataset relativizes each stream to its own
+first event, so the realistic and dense views of one object live in different
+relative frames. Lead times (durations) are frame-independent and are what the
+model's ``Dt`` consumes, so this script reads ABSOLUTE MJD (column 0) from the
+raw npz files and works entirely in absolute-time differences.
+
+Normalization uses the TRAIN-ONLY realistic stats the model was trained with
+(``kilonova_9band_norm_stats_trainonly.npz`` by default) -- the exact encoding
+the model saw. This is a read-only diagnostic: it writes plots and a CSV and
+never trains.
+"""
+
+import argparse
+import csv
+import os
+import sys
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from yoke.datasets.kilonova_dataset import (
+    EPS,
+    NINE_BAND_KEYS,
+    load_or_compute_band_normalization,
+)
+
+# Reuse the model loader and window/input helpers from the rollout diagnostics
+# script that lives alongside this one. These harness scripts are run directly
+# (not as an installed package), so make the script directory importable.
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from plot_pred_diagnostics_9band import (  # noqa: E402
+    _select_window,
+    build_context_input,
+    load_9band_model,
+)
+
+
+matplotlib.rcParams["pdf.fonttype"] = 42
+matplotlib.rcParams["ps.fonttype"] = 42
+plt.rc("font", family="serif")
+plt.rcParams["figure.figsize"] = (7, 5)
+
+
+BAND_KEYS = NINE_BAND_KEYS
+BAND_NAMES = ("ztfg", "ztfr", "ztfi", "u", "g", "r", "i", "z", "y")
+BAND_COLORS = (
+    "#2A9D8F", "#E63946", "#F4A261", "#457B9D", "#1B9E77",
+    "#D62828", "#E9C46A", "#8338EC", "#264653",
+)
+VALUE_COL = 1
+ERROR_COL = 2
+N_BANDS = len(BAND_KEYS)
+DROP_UPPER_LIMITS = True  # matches training for the realistic (context) stream
+
+
+def study_tag(study: int) -> str:
+    """Zero-padded study id used in default paths."""
+    return f"{int(study):03d}"
+
+
+def _stem(path: str) -> str:
+    """Return the object identifier: filename without directory or extension."""
+    return os.path.splitext(os.path.basename(path))[0]
+
+
+def read_merged_stream(
+    npz_path: str, drop_upper_limits: bool
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Read one file's merged, time-sorted event stream in ABSOLUTE MJD.
+
+    Unlike the training dataset, times are NOT relativized here, so streams from
+    two directories (realistic and dense) remain on a common absolute clock.
+
+    Args:
+        npz_path (str): Path to the light-curve npz.
+        drop_upper_limits (bool): Drop non-detections (non-finite error) so the
+            realistic context stream matches training. The dense set is all
+            detections, so this is a no-op there.
+
+    Returns:
+        (times, values, bands): absolute MJD, raw magnitude, band index; each
+        [N] and sorted by time. Empty arrays if the file has no usable events.
+    """
+    data = np.load(npz_path, allow_pickle=True)
+    times, values, bands = [], [], []
+    for band_idx, key in enumerate(BAND_KEYS):
+        if key not in data.files:
+            continue
+        arr = data[key]
+        if arr.size == 0:
+            continue
+        if drop_upper_limits:
+            arr = arr[np.isfinite(arr[:, ERROR_COL])]
+            if arr.shape[0] == 0:
+                continue
+        times.append(arr[:, 0].astype(np.float64))
+        values.append(arr[:, VALUE_COL].astype(np.float32))
+        bands.append(np.full(arr.shape[0], band_idx, dtype=np.int64))
+    data.close()
+
+    if not times:
+        empty_f = np.empty(0, dtype=np.float64)
+        return empty_f, empty_f.astype(np.float32), np.empty(0, dtype=np.int64)
+
+    times = np.concatenate(times)
+    values = np.concatenate(values)
+    bands = np.concatenate(bands)
+    order = np.argsort(times, kind="stable")
+    return times[order], values[order], bands[order]
+
+
+def _stem_to_path(data_glob: str) -> dict:
+    """Map object stem -> file path for all files matched by a glob."""
+    import glob
+
+    return {_stem(f): f for f in glob.glob(data_glob)}
+
+
+def eval_object(
+    real_stream,
+    dense_stream,
+    model,
+    device,
+    means,
+    stds,
+    context_window_days,
+    max_context_len,
+    late_time_cutoff_days,
+):
+    """Score one object's late-time dense truth against a realistic-context forecast.
+
+    Returns a dict with the scored late-time points and a smooth forecast curve,
+    or None if the object cannot be evaluated (no realistic context, or no dense
+    points in the late-time region).
+    """
+    r_t, r_v, r_b = real_stream
+    d_t, d_v, d_b = dense_stream
+
+    if r_t.shape[0] < 1 or d_t.shape[0] < 1:
+        return None
+
+    # Phase zero = the first realistic detection (the observed trigger). The
+    # late-time region is dense points more than the cutoff past it.
+    t0 = float(r_t[0])
+    last_real_t = float(r_t[-1])
+
+    late_mask = (d_t - t0) > late_time_cutoff_days
+    if not np.any(late_mask):
+        return None
+
+    # Seed context from the realistic stream: trailing window ending at the last
+    # realistic detection, normalized as in training. build_context_input
+    # subtracts win_t[0], so absolute times are fine here.
+    r_v_norm = (r_v - means[r_b]) / (stds[r_b] + EPS)
+    win_v, win_t, win_b = _select_window(
+        ctx_t=list(r_t.astype(np.float32)),
+        ctx_v=list(r_v_norm.astype(np.float32)),
+        ctx_b=list(r_b),
+        context_window_days=context_window_days,
+        max_context_len=max_context_len,
+    )
+    x = build_context_input(
+        win_v=win_v,
+        win_t=win_t,
+        win_b=win_b,
+        context_len=max_context_len,
+        n_bands=N_BANDS,
+        device=device,
+        window_mode=True,
+    )
+
+    # Score each late-time dense point at its true lead time from the last
+    # realistic detection.
+    scored = []
+    with torch.no_grad():
+        for idx in np.nonzero(late_mask)[0]:
+            dt = float(d_t[idx]) - last_real_t
+            if dt <= 0:
+                # Dense point precedes the last realistic detection; not a
+                # forecast into the future. Skip.
+                continue
+            Dt = torch.tensor([dt], dtype=torch.float32, device=device)
+            pred_all = model(x, in_vars=None, out_vars=None, Dt=Dt)
+            pred_all = pred_all.reshape(N_BANDS).detach().cpu().numpy()
+
+            band = int(d_b[idx])
+            pred_mag = float(pred_all[band] * (stds[band] + EPS) + means[band])
+            true_mag = float(d_v[idx])
+            scored.append(
+                {
+                    "phase": float(d_t[idx]) - t0,
+                    "lead_time": dt,
+                    "band": band,
+                    "pred_mag": pred_mag,
+                    "true_mag": true_mag,
+                    "residual_mag": pred_mag - true_mag,
+                }
+            )
+
+    if not scored:
+        return None
+
+    # Smooth forecast curve for plotting: sweep lead time from 0 to the farthest
+    # scored late-time point, predicting all bands at each lead time.
+    max_dt = max(s["lead_time"] for s in scored)
+    lead_grid = np.linspace(0.0, max_dt, 60).astype(np.float32)
+    curve = np.zeros((lead_grid.shape[0], N_BANDS), dtype=np.float32)
+    with torch.no_grad():
+        for k, dt in enumerate(lead_grid):
+            Dt = torch.tensor([dt], dtype=torch.float32, device=device)
+            pred = model(x, in_vars=None, out_vars=None, Dt=Dt)
+            curve[k] = pred.reshape(N_BANDS).detach().cpu().numpy()
+    curve_mag = curve * (stds[None, :] + EPS) + means[None, :]
+
+    return {
+        "scored": scored,
+        "t0": t0,
+        "last_real_t": last_real_t,
+        "curve_phase": (last_real_t - t0) + lead_grid,
+        "curve_mag": curve_mag.astype(np.float32),
+        "real": (r_t - t0, r_v, r_b),
+        "dense": (d_t - t0, d_v, d_b),
+    }
+
+
+def plot_object(result, stem, outpath):
+    """Plot realistic context, dense truth, and the late-time forecast per band."""
+    fig, axes = plt.subplots(3, 3, figsize=(13, 10), sharex=True)
+    axes = axes.ravel()
+    r_ph, r_v, r_b = result["real"]
+    d_ph, d_v, d_b = result["dense"]
+
+    for b in range(N_BANDS):
+        ax = axes[b]
+        rm = r_b == b
+        dm = d_b == b
+        if np.any(dm):
+            ax.scatter(d_ph[dm], d_v[dm], s=14, c="0.6", label="dense truth")
+        if np.any(rm):
+            ax.scatter(
+                r_ph[rm], r_v[rm], s=26, c=BAND_COLORS[b],
+                edgecolor="k", linewidth=0.4, label="realistic ctx",
+            )
+        ax.plot(
+            result["curve_phase"], result["curve_mag"][:, b],
+            c=BAND_COLORS[b], lw=1.6, label="forecast",
+        )
+        ax.invert_yaxis()  # magnitudes: brighter is smaller
+        ax.set_title(BAND_NAMES[b], fontsize=9)
+        if b == 0:
+            ax.legend(fontsize=7, loc="best")
+
+    fig.suptitle(f"Dense late-time forecast: {stem}")
+    fig.supxlabel("Phase from first realistic detection [days]")
+    fig.supylabel("Magnitude")
+    fig.tight_layout()
+    fig.savefig(outpath, dpi=130)
+    plt.close(fig)
+
+
+def get_args():
+    """Parse command-line arguments."""
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--study", type=int, default=24)
+    p.add_argument("--epoch", type=int, default=500)
+    p.add_argument("--ckpt", type=str, default=None)
+    p.add_argument(
+        "--realistic_glob",
+        type=str,
+        default=(
+            "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+            "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
+        ),
+        help="Glob for the realistic light-curve files (observing context).",
+    )
+    p.add_argument(
+        "--dense_glob",
+        type=str,
+        required=True,
+        help="Glob for the dense light-curve files (late-time truth).",
+    )
+    p.add_argument(
+        "--test_filelist",
+        type=str,
+        default=None,
+        help="Path to the test-split stem list (one object stem per line). If "
+        "omitted, all objects present in BOTH globs are evaluated.",
+    )
+    p.add_argument(
+        "--norm_stats_path",
+        type=str,
+        default="kilonova_9band_norm_stats_trainonly.npz",
+        help="Train-only normalization stats the model was trained with.",
+    )
+    p.add_argument(
+        "--late_time_cutoff_days",
+        type=float,
+        default=8.0,
+        help="Dense points with phase (from first realistic detection) greater "
+        "than this are the late-time region scored here.",
+    )
+    p.add_argument("--outdir", type=str, default=None)
+    p.add_argument(
+        "--max_objects",
+        type=int,
+        default=0,
+        help="Cap the number of objects evaluated (0 = all).",
+    )
+    p.add_argument(
+        "--n_plots",
+        type=int,
+        default=12,
+        help="Number of per-object forecast plots to write.",
+    )
+    return p.parse_args()
+
+
+def main():
+    """Run the dense late-time evaluation."""
+    args = get_args()
+    tag = study_tag(args.study)
+    if args.ckpt is None:
+        args.ckpt = (
+            f"runs/study_{tag}/study{tag}_modelState_epoch{args.epoch:04d}.pth"
+        )
+    if args.outdir is None:
+        args.outdir = f"runs/study_{tag}/dense_latetime_eval_9band"
+    os.makedirs(args.outdir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    (
+        model,
+        context_len,
+        n_bands,
+        context_window_days,
+        max_context_len,
+    ) = load_9band_model(args.ckpt, device)
+
+    if context_window_days is None:
+        raise ValueError(
+            "This eval requires a time-window checkpoint (context_window_days "
+            "set); the loaded checkpoint is fixed-count."
+        )
+
+    # Load the TRAIN-ONLY stats the model was trained with (loaded if present;
+    # no eval-set recomputation).
+    means, stds = load_or_compute_band_normalization(
+        stats_path=args.norm_stats_path,
+        band_keys=BAND_KEYS,
+        value_col=VALUE_COL,
+        error_col=ERROR_COL,
+        drop_upper_limits=DROP_UPPER_LIMITS,
+    )
+    means = np.asarray(means, dtype=np.float32)
+    stds = np.asarray(stds, dtype=np.float32)
+
+    # Pair realistic and dense objects by stem, restricted to the test split.
+    real_map = _stem_to_path(args.realistic_glob)
+    dense_map = _stem_to_path(args.dense_glob)
+    stems = sorted(set(real_map) & set(dense_map))
+
+    if args.test_filelist is not None:
+        with open(args.test_filelist) as fh:
+            test_stems = {line.strip() for line in fh if line.strip()}
+        stems = [s for s in stems if s in test_stems]
+        print(f"Restricted to {len(stems)} test-split objects.")
+
+    print(
+        f"Realistic files: {len(real_map)}; dense files: {len(dense_map)}; "
+        f"paired & in-split: {len(stems)}"
+    )
+    if args.max_objects > 0:
+        stems = stems[: args.max_objects]
+
+    all_scored = []
+    plotted = 0
+    n_eval = 0
+    for stem in stems:
+        real_stream = read_merged_stream(real_map[stem], DROP_UPPER_LIMITS)
+        dense_stream = read_merged_stream(dense_map[stem], drop_upper_limits=False)
+        result = eval_object(
+            real_stream,
+            dense_stream,
+            model,
+            device,
+            means,
+            stds,
+            context_window_days,
+            max_context_len,
+            args.late_time_cutoff_days,
+        )
+        if result is None:
+            continue
+        n_eval += 1
+        for s in result["scored"]:
+            s["stem"] = stem
+            all_scored.append(s)
+        if plotted < args.n_plots:
+            plot_object(
+                result, stem,
+                os.path.join(args.outdir, f"latetime_{stem}.png"),
+            )
+            plotted += 1
+
+    if not all_scored:
+        print("No late-time points scored (check the cutoff and globs).")
+        return
+
+    # Per-band late-time error summary.
+    resid = np.asarray([s["residual_mag"] for s in all_scored])
+    bands = np.asarray([s["band"] for s in all_scored])
+    print(f"\nEvaluated {n_eval} objects; {len(all_scored)} late-time points "
+          f"(cutoff {args.late_time_cutoff_days} d).")
+    print(f"Overall late-time RMSE (mag): {np.sqrt(np.mean(resid**2)):.4f}  "
+          f"MAE: {np.mean(np.abs(resid)):.4f}")
+    print("Per-band late-time error (mag):")
+    for b in range(N_BANDS):
+        m = bands == b
+        if np.any(m):
+            print(f"  {BAND_NAMES[b]:>5}: n={m.sum():5d}  "
+                  f"RMSE={np.sqrt(np.mean(resid[m]**2)):.4f}  "
+                  f"MAE={np.mean(np.abs(resid[m])):.4f}  "
+                  f"bias={np.mean(resid[m]):+.4f}")
+
+    # Error vs lead time (binned) plot.
+    lead = np.asarray([s["lead_time"] for s in all_scored])
+    fig, ax = plt.subplots()
+    edges = np.linspace(0, lead.max(), 11)
+    centers = 0.5 * (edges[:-1] + edges[1:])
+    rmse_bin = np.full(centers.shape[0], np.nan)
+    for i in range(centers.shape[0]):
+        m = (lead >= edges[i]) & (lead < edges[i + 1])
+        if np.any(m):
+            rmse_bin[i] = np.sqrt(np.mean(resid[m] ** 2))
+    ax.plot(centers, rmse_bin, "o-")
+    ax.set_xlabel("Lead time from last realistic detection [days]")
+    ax.set_ylabel("Late-time forecast RMSE [mag]")
+    ax.set_title("Dense late-time forecast error vs lead time")
+    fig.tight_layout()
+    fig.savefig(os.path.join(args.outdir, "latetime_rmse_vs_lead.png"), dpi=130)
+    plt.close(fig)
+
+    # Full per-point CSV.
+    csv_path = os.path.join(args.outdir, "latetime_scored_points.csv")
+    with open(csv_path, "w", newline="") as fh:
+        w = csv.writer(fh)
+        w.writerow(
+            ["stem", "band", "phase_days", "lead_time_days",
+             "pred_mag", "true_mag", "residual_mag"]
+        )
+        for s in all_scored:
+            w.writerow([
+                s["stem"], BAND_NAMES[s["band"]], f"{s['phase']:.4f}",
+                f"{s['lead_time']:.4f}", f"{s['pred_mag']:.4f}",
+                f"{s['true_mag']:.4f}", f"{s['residual_mag']:.4f}",
+            ])
+
+    print(f"\nWrote {plotted} per-object plots, the RMSE-vs-lead plot, and "
+          f"{csv_path} in {args.outdir}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/make_kn_object_lists.py b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
new file mode 100644
index 00000000..eb404449
--- /dev/null
+++ b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
@@ -0,0 +1,170 @@
+"""Create a seeded, object-level train/val/test split for the 9-band KN data.
+
+The kilonova pipeline can train on the SAME objects viewed two ways: a
+"realistic" light-curve set (sparse, upper limits dropped) and a "dense" set
+(same objects, denser cadence, no limiting-mag cut). The two views live in
+separate directories but share filename stems (``lc_<id>``). To avoid a single
+object leaking across the split (its realistic view in train, its dense view in
+test, or vice versa), the split is computed ONCE at the object level and applied
+to BOTH directories by stem.
+
+This script writes three stem lists (one object identifier per line) --
+
+    kn_rubin_ztf_train.txt
+    kn_rubin_ztf_val.txt
+    kn_rubin_ztf_test.txt
+
+-- to ``applications/filelists/`` by default. The training harness reads each
+into a set and passes it as ``object_ids`` to
+``Kilonova_lc_scalar_context_DataSet_9band`` for both the realistic and dense
+directories.
+
+Determinism: the split uses a seeded ``numpy.random.default_rng`` and sorts the
+stems BEFORE shuffling, so the result is reproducible across machines regardless
+of filesystem ``glob`` ordering. Run it once and commit the three lists.
+
+Example:
+    python make_kn_object_lists.py \
+        --realistic_glob "/path/to/rubin_ztf_10000_dataset/lc_*.npz" \
+        --dense_glob "/path/to/rubin_ztf_10000_dense/lc_*.npz" \
+        --seed 20240817
+"""
+
+import argparse
+import glob
+import os
+
+
+def _stem(path: str) -> str:
+    """Return the object identifier: filename without directory or extension."""
+    return os.path.splitext(os.path.basename(path))[0]
+
+
+def make_object_split(
+    realistic_glob: str,
+    seed: int,
+    train_frac: float = 0.8,
+    val_frac: float = 0.1,
+    dense_glob: str = None,
+) -> tuple[list[str], list[str], list[str]]:
+    """Split object stems into train/val/test deterministically.
+
+    The split is computed over the REALISTIC universe of objects (always
+    present). The test fraction is the remainder, so no object is dropped.
+
+    Args:
+        realistic_glob (str): Glob for the realistic light-curve files.
+        seed (int): Seed for the shuffle RNG.
+        train_frac (float): Fraction of objects for training.
+        val_frac (float): Fraction of objects for validation. Test = remainder.
+        dense_glob (str): Optional glob for the dense files. Only used to report
+            how many realistic stems are (not) covered by the dense set.
+
+    Returns:
+        (train, val, test): three sorted lists of object stems.
+    """
+    import numpy as np
+
+    realistic_stems = sorted({_stem(f) for f in glob.glob(realistic_glob)})
+    if not realistic_stems:
+        raise ValueError(f"No files matched realistic_glob: {realistic_glob!r}")
+
+    if dense_glob is not None:
+        dense_stems = {_stem(f) for f in glob.glob(dense_glob)}
+        missing = set(realistic_stems) - dense_stems
+        extra = dense_stems - set(realistic_stems)
+        print(
+            f"Dense coverage: {len(dense_stems)} dense stems; "
+            f"{len(missing)} realistic objects have NO dense counterpart; "
+            f"{len(extra)} dense-only stems (ignored)."
+        )
+
+    # Sort-then-shuffle: deterministic regardless of glob/filesystem ordering.
+    stems = sorted(realistic_stems)
+    rng = np.random.default_rng(seed)
+    rng.shuffle(stems)
+
+    n = len(stems)
+    n_train = int(np.floor(train_frac * n))
+    n_val = int(np.floor(val_frac * n))
+
+    train = sorted(stems[:n_train])
+    val = sorted(stems[n_train:n_train + n_val])
+    test = sorted(stems[n_train + n_val:])  # remainder, no dropped objects
+
+    return train, val, test
+
+
+def _write_list(path: str, stems: list[str]) -> None:
+    """Write one stem per line to path."""
+    with open(path, "w") as fh:
+        for s in stems:
+            fh.write(s + "\n")
+    print(f"Wrote {len(stems):6d} stems -> {path}")
+
+
+def main() -> None:
+    """Parse args, build the split, and write the three stem lists."""
+    here = os.path.dirname(os.path.abspath(__file__))
+    default_out = os.path.abspath(
+        os.path.join(here, "..", "..", "filelists")
+    )
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--realistic_glob",
+        type=str,
+        default=(
+            "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+            "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
+        ),
+        help="Glob for the realistic light-curve files.",
+    )
+    parser.add_argument(
+        "--dense_glob",
+        type=str,
+        default=(
+            "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+            "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
+        ),
+        help="Optional glob for the dense files (for coverage reporting only).",
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        default=default_out,
+        help="Directory to write the stem lists into.",
+    )
+    parser.add_argument(
+        "--prefix",
+        type=str,
+        default="kn_rubin_ztf",
+        help="Filename prefix for the three output lists.",
+    )
+    parser.add_argument("--seed", type=int, default=20240817)
+    parser.add_argument("--train_frac", type=float, default=0.8)
+    parser.add_argument("--val_frac", type=float, default=0.1)
+    args = parser.parse_args()
+
+    train, val, test = make_object_split(
+        realistic_glob=args.realistic_glob,
+        seed=args.seed,
+        train_frac=args.train_frac,
+        val_frac=args.val_frac,
+        dense_glob=args.dense_glob,
+    )
+
+    total = len(train) + len(val) + len(test)
+    print(
+        f"Split {total} objects (seed={args.seed}): "
+        f"{len(train)} train / {len(val)} val / {len(test)} test"
+    )
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    _write_list(os.path.join(args.out_dir, f"{args.prefix}_train.txt"), train)
+    _write_list(os.path.join(args.out_dir, f"{args.prefix}_val.txt"), val)
+    _write_list(os.path.join(args.out_dir, f"{args.prefix}_test.txt"), test)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 292f52cb..38ba96e9 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -1,4 +1,5 @@
 import os
+import glob
 import time
 import random
 import argparse
@@ -8,6 +9,7 @@
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim.lr_scheduler import LambdaLR
+from torch.utils.data import ConcatDataset
 
 from yoke.models.vit.swin.bomberman import (
     LodeRunner,
@@ -31,6 +33,15 @@
 from yoke.lr_schedulers import CosineWithWarmupScheduler
 from yoke.helpers import cli
 
+def _read_stem_list(path: str) -> set:
+    """Read an object-stem split file (one stem per line) into a set.
+
+    Deterministic and RNG-free, so every DDP rank builds the identical set.
+    """
+    with open(path) as fh:
+        return {line.strip() for line in fh if line.strip()}
+
+
 #############################################
 # Inputs
 #############################################
@@ -97,11 +108,39 @@
     "from epoch 0. Only used when --n_rollout_steps > 1.",
 )
 
-# Change some default filepaths.
+# Paired-dataset globs. The realistic set is always used; the dense set (same
+# objects, denser cadence, no limiting-mag cut) is optional and concatenated onto
+# the realistic training data when present. Both are filtered to the object-level
+# split (see --train_filelist / --validation_filelist below).
+parser.add_argument(
+    "--kn_realistic_glob",
+    type=str,
+    default=(
+        "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+        "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
+    ),
+    help="Glob for the realistic light-curve files (primary training data).",
+)
+parser.add_argument(
+    "--kn_dense_glob",
+    type=str,
+    default=(
+        "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+        "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
+    ),
+    help="Optional glob for the dense light-curve files. When set (and it "
+    "matches files), the dense TRAIN objects are concatenated onto the "
+    "realistic TRAIN objects to supervise late-time behavior. Validation and "
+    "the primary metric stay realistic-only.",
+)
+
+# Change some default filepaths. The KN split lists hold object stems (one per
+# line), shared across the realistic and dense directories; see
+# make_kn_object_lists.py.
 parser.set_defaults(
-    train_filelist="lsc240420_prefixes_train_80pct.txt",
-    validation_filelist="lsc240420_prefixes_validation_10pct.txt",
-    test_filelist="lsc240420_prefixes_test_10pct.txt",
+    train_filelist="kn_rubin_ztf_train.txt",
+    validation_filelist="kn_rubin_ztf_val.txt",
+    test_filelist="kn_rubin_ztf_test.txt",
 )
 
 
@@ -116,9 +155,13 @@ def main(args, rank, world_size, local_rank, device):
     Ngpus = args.Ngpus
     Knodes = args.Knodes
 
-    # Data Paths
+    # Data Paths. The KN train/val/test lists hold object stems (one per line);
+    # the same stems select objects in both the realistic and dense directories,
+    # so an object is entirely in train or entirely in val/test in both views.
     train_filelist = args.FILELIST_DIR + args.train_filelist
     validation_filelist = args.FILELIST_DIR + args.validation_filelist
+    train_stems = _read_stem_list(train_filelist)
+    val_stems = _read_stem_list(validation_filelist)
 
     # Model Parameters
     embed_dim = args.embed_dim
@@ -193,6 +236,17 @@ def main(args, rank, world_size, local_rank, device):
     # CONTEXT_WINDOW_DAYS = None to use the legacy fixed-count context.
     CONTEXT_WINDOW_DAYS = 2.0
     MAX_CONTEXT_LEN = 12
+
+    # Horizon-covering target sampling (window mode only). When set, each sample
+    # draws its target lead time ~uniform in days over (0, TARGET_HORIZON_DAYS]
+    # and supervises the event nearest that lead time, instead of always the
+    # immediate next event. This flattens the supervised-Dt distribution so the
+    # model is trained at the multi-day lead times it is asked to forecast,
+    # rather than only at the short gap-to-next-event (p99 ~4d on this set) while
+    # rollouts forecast out to ~12d. Set from the plot_observation_histograms.py
+    # "Supervised Δt vs forecast horizon" panel (the h=1 tail is the uncovered
+    # region). Leave None to supervise the immediate next event as before.
+    TARGET_HORIZON_DAYS = 8.0
     # In window mode the model's first layer is sized by the padded width.
     WRAPPER_CONTEXT_LEN = (
         MAX_CONTEXT_LEN if CONTEXT_WINDOW_DAYS is not None else CONTEXT_LEN
@@ -397,7 +451,15 @@ def main(args, rank, world_size, local_rank, device):
     )
     '''
 
-    norm_stats_path = "kilonova_9band_norm_stats.npz"
+    # Normalization statistics are computed over the TRAIN objects only (of the
+    # realistic set) to avoid val/test leakage. The stats path is distinct from
+    # the old all-files cache so a stale/leaky cache can't be silently reused.
+    norm_stats_path = "kilonova_9band_norm_stats_trainonly.npz"
+    train_norm_files = sorted(
+        f
+        for f in glob.glob(args.kn_realistic_glob)
+        if os.path.splitext(os.path.basename(f))[0] in train_stems
+    )
 
     if rank == 0:
         band_means, band_stds = load_or_compute_band_normalization(
@@ -406,6 +468,7 @@ def main(args, rank, world_size, local_rank, device):
             value_col=VALUE_COL,
             error_col=ERROR_COL,
             drop_upper_limits=DROP_UPPER_LIMITS,
+            file_prefix_list=train_norm_files,
         )
 
     dist.barrier()
@@ -421,45 +484,75 @@ def main(args, rank, world_size, local_rank, device):
         print("band_means:", band_means)
         print("band_stds:", band_stds)
 
-    # The 9-band dataset selects its N_imgs files with an unseeded RNG and each
-    # DDP rank builds its own dataset. With N_imgs>0 that would give every rank a
-    # DIFFERENT random file subset, hence different sample counts and different
-    # per-rank batch counts, so ranks desync and hang at gradient all-reduce
-    # (NCCL watchdog timeout). Seed both RNGs to the SAME value on every rank so
-    # all ranks pick the identical file subset. (With N_imgs=0 all files are used
-    # and this is moot, but seeding is harmless.)
+    # Object-level split makes every DDP rank build an identical-length dataset:
+    # the train/val stem sets come from static files (no RNG), the file list is
+    # sorted(glob(...)) then filtered by stem, and N_imgs=0 uses all matched
+    # files (no np.random.choice). So no per-rank subset desync is possible.
+    # (The remaining random.shuffle inside the dataset only reorders files; it
+    # does not change the sample count.) Seed once for reproducibility only --
+    # this is no longer load-bearing for rank sync. Keep N_imgs=0; N_imgs>0 would
+    # reintroduce the unseeded-choice desync.
     DATA_SEED = 42
     np.random.seed(DATA_SEED)
     random.seed(DATA_SEED)
 
-    train_dataset = Kilonova_lc_scalar_context_DataSet_9band(
-        N_imgs=0,
-        context_len=CONTEXT_LEN,
-        band_keys=BAND_KEYS,
-        value_col=VALUE_COL,
-        error_col=ERROR_COL,
-        drop_upper_limits=DROP_UPPER_LIMITS,
-        means=band_means,
-        stds=band_stds,
-        n_rollout_steps=n_rollout_steps,
-        context_window_days=CONTEXT_WINDOW_DAYS,
-        max_context_len=MAX_CONTEXT_LEN,
-    )
+    def _make_9band(
+        data_glob: str, object_ids: set
+    ) -> Kilonova_lc_scalar_context_DataSet_9band:
+        """Build a 9-band dataset over one directory restricted to object_ids."""
+        return Kilonova_lc_scalar_context_DataSet_9band(
+            N_imgs=0,
+            context_len=CONTEXT_LEN,
+            band_keys=BAND_KEYS,
+            value_col=VALUE_COL,
+            error_col=ERROR_COL,
+            drop_upper_limits=DROP_UPPER_LIMITS,
+            means=band_means,
+            stds=band_stds,
+            n_rollout_steps=n_rollout_steps,
+            context_window_days=CONTEXT_WINDOW_DAYS,
+            max_context_len=MAX_CONTEXT_LEN,
+            target_horizon_days=TARGET_HORIZON_DAYS,
+            data_glob=data_glob,
+            object_ids=object_ids,
+        )
 
-    val_dataset = Kilonova_lc_scalar_context_DataSet_9band(
-        N_imgs=0,
-        context_len=CONTEXT_LEN,
-        band_keys=BAND_KEYS,
-        value_col=VALUE_COL,
-        error_col=ERROR_COL,
-        drop_upper_limits=DROP_UPPER_LIMITS,
-        means=band_means,
-        stds=band_stds,
-        n_rollout_steps=n_rollout_steps,
-        context_window_days=CONTEXT_WINDOW_DAYS,
-        max_context_len=MAX_CONTEXT_LEN,
+    # Realistic TRAIN objects (always present) plus, if a dense set is provided
+    # and matches files, the SAME train objects viewed densely -- concatenated to
+    # supervise late-time behavior. Validation stays realistic-only (matches the
+    # deployment metric).
+    train_real = _make_9band(args.kn_realistic_glob, train_stems)
+    train_parts = [train_real]
+
+    if args.kn_dense_glob and glob.glob(args.kn_dense_glob):
+        train_dense = _make_9band(args.kn_dense_glob, train_stems)
+        if len(train_dense) > 0:
+            train_parts.append(train_dense)
+            if rank == 0:
+                print(
+                    f"Dense training set added: {len(train_dense)} samples "
+                    f"(realistic: {len(train_real)} samples).",
+                    flush=True,
+                )
+        elif rank == 0:
+            print(
+                "Dense glob matched files but yielded 0 samples for the train "
+                "split; training on realistic set only.",
+                flush=True,
+            )
+    elif rank == 0 and args.kn_dense_glob:
+        print(
+            "Dense glob set but matched no files; training on realistic set "
+            "only.",
+            flush=True,
+        )
+
+    train_dataset = (
+        ConcatDataset(train_parts) if len(train_parts) > 1 else train_parts[0]
     )
 
+    val_dataset = _make_9band(args.kn_realistic_glob, val_stems)
+
 
     # NOTE: For DDP the batch_size is the per-GPU batch_size!!!
     train_dataloader = make_distributed_dataloader(
@@ -613,6 +706,12 @@ def main(args, rank, world_size, local_rank, device):
                     "n_rollout_steps": n_rollout_steps,
                     "context_window_days": CONTEXT_WINDOW_DAYS,
                     "max_context_len": MAX_CONTEXT_LEN,
+                    "target_horizon_days": TARGET_HORIZON_DAYS,
+                    "train_filelist": args.train_filelist,
+                    "validation_filelist": args.validation_filelist,
+                    "kn_realistic_glob": args.kn_realistic_glob,
+                    "kn_dense_glob": args.kn_dense_glob,
+                    "norm_stats_path": norm_stats_path,
                 },
                 new_chkpt_path,
             )
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 3a2b5fdc..67c5a042 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -37,6 +37,17 @@
 )
 
 
+def _stem(path: str) -> str:
+    """Return the object identifier (filename without directory or extension).
+
+    Light-curve files are named ``lc_<id>.npz``; the stem ``lc_<id>`` is the
+    object identity used to pair the same object across the realistic and dense
+    directories and to build object-level train/val/test splits. ``.npz`` is a
+    single extension so ``splitext`` is exact.
+    """
+    return os.path.splitext(os.path.basename(path))[0]
+
+
 def compute_band_normalization(
     file_prefix_list: list[str],
     band_keys: tuple[str, ...] = ("arr_ztfg", "arr_ztfr", "arr_ztfi"),
@@ -117,6 +128,7 @@ def load_or_compute_band_normalization(
     value_col: int = 1,
     error_col: int = 2,
     drop_upper_limits: bool = False,
+    file_prefix_list: list[str] = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Load cached per-band normalization stats or compute them if missing.
 
@@ -128,17 +140,24 @@ def load_or_compute_band_normalization(
             used when drop_upper_limits is True.
         drop_upper_limits (bool): If True, exclude upper limits (non-finite
             uncertainty) from the statistics, matching a dataset that drops them.
+        file_prefix_list (list[str]): Files to accumulate stats over. Pass the
+            TRAIN files only to avoid val/test leakage. If None, falls back to
+            the legacy hardcoded scratch glob (kept for backward compatibility).
 
     Returns:
         means (np.ndarray): Per-band means, shape [n_bands].
         stds (np.ndarray): Per-band standard deviations, shape [n_bands].
     """
-    # FIXME: hardcoded scratch path. Should be passed in as an argument so this
-    # library function does not depend on a user-specific filesystem location.
-    file_prefix_list = sorted(
-        glob.glob(
-        "/net/sescratch1/atoivonen/data/KN_lightcurves/rubin_ztf_10000_dataset/lc_*.npz")
-    )
+    if file_prefix_list is None:
+        # FIXME: hardcoded scratch path fallback. Prefer passing an explicit
+        # (train-only) file_prefix_list so this library function does not depend
+        # on a user-specific filesystem location and does not leak val/test data.
+        file_prefix_list = sorted(
+            glob.glob(
+                "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+                "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
+            )
+        )
 
     if os.path.exists(stats_path):
         stats = np.load(stats_path, allow_pickle=True)
@@ -389,6 +408,9 @@ def __init__(
         n_rollout_steps: int = 1,
         context_window_days: float = None,
         max_context_len: int = None,
+        target_horizon_days: float = None,
+        data_glob: str = None,
+        object_ids: set = None,
     ) -> None:
         """Initialize the dataset and build the merged-event sample index.
 
@@ -426,20 +448,52 @@ def __init__(
                 Windows with more real events than this keep only the most recent
                 ``max_context_len``; windows with fewer are zero-padded. Defaults
                 to ``context_len`` when not given. Unused in fixed-count mode.
+            target_horizon_days (float): If set (window mode only), enables
+                **horizon-covering target sampling**: instead of always
+                supervising the immediate next event, each sample draws a target
+                lead time ~uniform in days over ``(0, target_horizon_days]`` and
+                supervises the event whose gap from the anchor is nearest that
+                lead time (clamped to the last event in the curve). This flattens
+                the supervised-``Delta_t`` distribution so the model is trained at
+                the lead times it is later asked to forecast, rather than only at
+                the short gap-to-next-event (median ~0.3d, p99 ~4d) while
+                diagnostics forecast out to ~12d. The context selection (trailing
+                ``context_window_days`` window ending at the anchor) is unchanged,
+                preserving train/inference parity; only the target moves. In
+                rollout mode each of the ``n_rollout_steps`` steps draws its own
+                farther target from its current anchor, spreading coverage across
+                the whole rollout. None (default) keeps the immediate-next-event
+                target. Ignored in fixed-count mode.
+            data_glob (str): Glob pattern selecting the light-curve files (i.e.
+                which dataset directory). If None (default) the legacy hardcoded
+                rubin_ztf_10000 scratch glob is used, preserving prior behavior.
+                Used to point at either the realistic or the dense directory.
+            object_ids (set): If given, only files whose stem (filename without
+                directory or ``.npz``) is in this set are loaded. Used to apply a
+                shared object-level train/val/test split across the realistic and
+                dense directories (the same stems appear in both). None (default)
+                loads all files matched by ``data_glob``.
         """
-        # FIXME: hardcoded scratch path. Should be passed in as an argument so
-        # this dataset does not depend on a user-specific filesystem location.
-        # NOTE: must point at the SAME dataset as
-        # load_or_compute_band_normalization (the rubin_ztf_10000 set). The old
+        # Select the dataset directory. NOTE: the chosen set must be consistent
+        # with the normalization stats (both Rubin+ZTF). The old
         # uniform_dataset_20000 set is ZTF-only, so training on it left the six
         # Rubin output heads without any targets (never trained) while the norm
         # stats were computed over Rubin+ZTF -- a silent train/stats mismatch.
-        file_prefix_list = sorted(
-            glob.glob(
-                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-                "rubin_ztf_10000_dataset/lc_*.npz"
+        if data_glob is None:
+            # Legacy hardcoded scratch fallback (backward compatibility).
+            data_glob = (
+                "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+                "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
             )
-        )
+        file_prefix_list = sorted(glob.glob(data_glob))
+
+        # Restrict to an object-level split (shared across the realistic and
+        # dense directories via matching filename stems) when object_ids is set.
+        if object_ids is not None:
+            object_ids = set(object_ids)
+            file_prefix_list = [
+                f for f in file_prefix_list if _stem(f) in object_ids
+            ]
 
         if N_imgs == 0:
             self.file_prefix_list = file_prefix_list
@@ -485,6 +539,22 @@ def __init__(
         else:
             self.max_context_len = context_len
 
+        # Horizon-covering target sampling (window mode only). None keeps the
+        # immediate-next-event target; a positive value draws a target lead time
+        # ~uniform in days over (0, target_horizon_days].
+        self.target_horizon_days = target_horizon_days
+        if target_horizon_days is not None:
+            if not self.window_mode:
+                raise ValueError(
+                    "target_horizon_days is only supported in time-window mode "
+                    "(set context_window_days)."
+                )
+            if target_horizon_days <= 0:
+                raise ValueError(
+                    "target_horizon_days must be positive, got "
+                    f"{target_horizon_days}"
+                )
+
         if means is None:
             raise ValueError(
                 "means must be provided for per-band normalization. "
@@ -520,6 +590,12 @@ def __init__(
         # rel_time is relative to the earliest observation across all bands in
         # the file so absolute MJD offsets do not leak into the model.
         self.events_per_file = []
+        # Object stem (identity) parallel to events_per_file. Needed because a
+        # file's index in events_per_file is NOT its index in file_prefix_list
+        # (files are shuffled above, and empty curves are skipped below), so
+        # pairing the same object across the realistic and dense datasets must go
+        # through the stem, not a positional index.
+        self.stems_per_file = []
         self.samples = []
 
         for file_idx, fn in enumerate(self.file_prefix_list):
@@ -575,6 +651,7 @@ def __init__(
             self.events_per_file.append(
                 (times, values.astype(np.float32), bands)
             )
+            self.stems_per_file.append(_stem(fn))
 
             n_events = times.shape[0]
             file_idx = len(self.events_per_file) - 1
@@ -687,6 +764,32 @@ def _getitem_single(
 
         return x, target, mask, Dt
 
+    def _draw_target_idx(
+        self, times: np.ndarray, anchor_idx: int, n_events: int
+    ) -> int:
+        """Draw a horizon-covering target event index ahead of ``anchor_idx``.
+
+        Samples a target lead time ~uniform in days over
+        ``(0, target_horizon_days]`` and returns the future event whose gap from
+        the anchor is nearest that lead time. Because reachable lead times are
+        discrete (the actual future observation times), this approximates a
+        uniform-in-days target distribution up to data availability, and clamps
+        to the last event when the drawn lead time exceeds the curve. Assumes at
+        least one event follows the anchor (``anchor_idx + 1 < n_events``).
+
+        Args:
+            times (np.ndarray): Merged, sorted, file-relative event times.
+            anchor_idx (int): Index of the most recent context event.
+            n_events (int): Number of events in the curve.
+
+        Returns:
+            int: The drawn target event index, in ``(anchor_idx, n_events)``.
+        """
+        lead_time = np.random.uniform(0.0, self.target_horizon_days)
+        cand = np.arange(anchor_idx + 1, n_events)
+        gaps = times[cand] - times[anchor_idx]
+        return int(cand[np.argmin(np.abs(gaps - lead_time))])
+
     def _getitem_window(
         self, index: int
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -718,13 +821,23 @@ def _getitem_window(
         file_idx, target_idx = self.samples[index]
         times, values, bands = self.events_per_file[file_idx]
 
-        # Anchor the trailing window on the event immediately before the target
-        # (the most recent observation). Select all earlier events within the
-        # window, then keep the most recent max_context_len if there are more.
-        anchor_t = times[target_idx - 1]
+        # Anchor the trailing window on the event immediately before the enumerated
+        # target (the most recent observation). With horizon-covering target
+        # sampling the supervised target is redrawn to a farther event so the
+        # lead time Dt is ~uniform in days; the anchor (hence the context) is
+        # unchanged, preserving train/inference parity.
+        anchor_idx = target_idx - 1
+        if self.target_horizon_days is not None:
+            target_idx = self._draw_target_idx(
+                times, anchor_idx, times.shape[0]
+            )
+
+        anchor_t = times[anchor_idx]
         lo = anchor_t - self.context_window_days
 
-        prior_t = times[:target_idx]
+        # Context is events up to and including the anchor (never the target,
+        # which may now be several events ahead) within the trailing window.
+        prior_t = times[: anchor_idx + 1]
         in_window = prior_t >= lo
         sel_idx = np.nonzero(in_window)[0]
         if sel_idx.shape[0] > self.max_context_len:
@@ -760,8 +873,10 @@ def _getitem_window(
         target = torch.tensor(target, dtype=torch.float32)
         mask = torch.tensor(mask, dtype=torch.float32)
 
+        # Lead time from the anchor (most recent context event) to the target,
+        # which may be several events ahead under horizon-covering sampling.
         Dt = torch.tensor(
-            times[target_idx] - times[target_idx - 1],
+            times[target_idx] - times[anchor_idx],
             dtype=torch.float32,
         )
 
@@ -932,16 +1047,28 @@ def _getitem_window_rollout(
         future_dt = np.zeros(n, dtype=np.float32)
         future_valid = np.zeros(n, dtype=np.float32)
 
+        # Future targets. Without horizon sampling these are the consecutive
+        # events target_idx … target_idx + n - 1 (immediate-next chain). With
+        # horizon sampling each step draws a farther target from its own anchor
+        # (the previous step's target), so the rollout is supervised at
+        # ~uniform-in-days lead times at every step; future_dt is the anchor→
+        # target gap the training loop uses to advance the growing buffer.
         n_events = times.shape[0]
+        prev_idx = target_idx - 1
         for step in range(n):
-            t_idx = target_idx + step
-            if t_idx >= n_events:
+            if prev_idx + 1 >= n_events:
                 break
 
+            if self.target_horizon_days is not None:
+                t_idx = self._draw_target_idx(times, prev_idx, n_events)
+            else:
+                t_idx = prev_idx + 1
+
             future_v[step] = values[t_idx]
             future_b[step] = bands[t_idx]
-            future_dt[step] = times[t_idx] - times[t_idx - 1]
+            future_dt[step] = times[t_idx] - times[prev_idx]
             future_valid[step] = 1.0
+            prev_idx = t_idx
 
         return (
             torch.tensor(ctx_v, dtype=torch.float32),

From 58b36c99b5d0e0f0f844bdb0fed41bf489ffecf9 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 18 Aug 2026 10:49:01 -0600
Subject: [PATCH 35/66] fix paths

---
 .../harnesses/KN_loderunner/eval_dense_latetime_9band.py      | 2 +-
 applications/harnesses/KN_loderunner/make_kn_object_lists.py  | 4 ++--
 applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py  | 4 ++--
 src/yoke/datasets/kilonova_dataset.py                         | 4 ++--
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index 9e2a5c72..8afa24fd 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -292,7 +292,7 @@ def get_args():
         "--realistic_glob",
         type=str,
         default=(
-            "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
             "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Glob for the realistic light-curve files (observing context).",
diff --git a/applications/harnesses/KN_loderunner/make_kn_object_lists.py b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
index eb404449..42b1a4f0 100644
--- a/applications/harnesses/KN_loderunner/make_kn_object_lists.py
+++ b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
@@ -115,7 +115,7 @@ def main() -> None:
         "--realistic_glob",
         type=str,
         default=(
-            "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
             "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Glob for the realistic light-curve files.",
@@ -124,7 +124,7 @@ def main() -> None:
         "--dense_glob",
         type=str,
         default=(
-            "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
             "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Optional glob for the dense files (for coverage reporting only).",
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 38ba96e9..585ad5fe 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -116,7 +116,7 @@ def _read_stem_list(path: str) -> set:
     "--kn_realistic_glob",
     type=str,
     default=(
-        "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+        "/net/sescratch1/atoivonen/data/KN_lightcurves/"
         "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
     ),
     help="Glob for the realistic light-curve files (primary training data).",
@@ -125,7 +125,7 @@ def _read_stem_list(path: str) -> set:
     "--kn_dense_glob",
     type=str,
     default=(
-        "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+        "/net/sescratch1/atoivonen/data/KN_lightcurves/"
         "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
     ),
     help="Optional glob for the dense light-curve files. When set (and it "
diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py
index 67c5a042..74ce5d3c 100644
--- a/src/yoke/datasets/kilonova_dataset.py
+++ b/src/yoke/datasets/kilonova_dataset.py
@@ -154,7 +154,7 @@ def load_or_compute_band_normalization(
         # on a user-specific filesystem location and does not leak val/test data.
         file_prefix_list = sorted(
             glob.glob(
-                "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
                 "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
             )
         )
@@ -482,7 +482,7 @@ def __init__(
         if data_glob is None:
             # Legacy hardcoded scratch fallback (backward compatibility).
             data_glob = (
-                "/Users/atoivonen/Documents/repos/fake_kilonovae/"
+                "/net/sescratch1/atoivonen/data/KN_lightcurves/"
                 "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
             )
         file_prefix_list = sorted(glob.glob(data_glob))

From 371cedace5e01849ce210026b9b13ab8578b74b8 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 18 Aug 2026 12:28:30 -0600
Subject: [PATCH 36/66] fix filepath

---
 applications/harnesses/KN_loderunner/training_START.input | 2 +-
 applications/harnesses/KN_loderunner/training_input.tmpl  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_START.input b/applications/harnesses/KN_loderunner/training_START.input
index c477f515..ded5f8a9 100644
--- a/applications/harnesses/KN_loderunner/training_START.input
+++ b/applications/harnesses/KN_loderunner/training_START.input
@@ -1,7 +1,7 @@
 --studyIDX
 <studyIDX>
 --FILELIST_DIR
-/users/atoivonen/forks/Yoke/applications/filelists/
+/net/sescratch1/atoivonen/projects/filelists/
 --LSC_NPZ_DIR
 /lustre/scratch5/exempt/artimis/data/lsc240420/
 --train_filelist
diff --git a/applications/harnesses/KN_loderunner/training_input.tmpl b/applications/harnesses/KN_loderunner/training_input.tmpl
index 4446f897..37420a29 100644
--- a/applications/harnesses/KN_loderunner/training_input.tmpl
+++ b/applications/harnesses/KN_loderunner/training_input.tmpl
@@ -1,7 +1,7 @@
 --studyIDX
 <studyIDX>
 --FILELIST_DIR
-/users/atoivonen/forks/Yoke/applications/filelists/
+/net/sescratch1/atoivonen/projects/filelists/
 --LSC_NPZ_DIR
 /lustre/scratch5/exempt/artimis/data/lsc240420/
 --train_filelist

From e6ad38bd7679662e5e4ae3cd368b60d809fba88b Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 18 Aug 2026 12:41:56 -0600
Subject: [PATCH 37/66] fix more paths

---
 applications/harnesses/KN_loderunner/training_START.input | 4 ++--
 applications/harnesses/KN_loderunner/training_input.tmpl  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_START.input b/applications/harnesses/KN_loderunner/training_START.input
index ded5f8a9..cde36996 100644
--- a/applications/harnesses/KN_loderunner/training_START.input
+++ b/applications/harnesses/KN_loderunner/training_START.input
@@ -5,9 +5,9 @@
 --LSC_NPZ_DIR
 /lustre/scratch5/exempt/artimis/data/lsc240420/
 --train_filelist
-lsc240420_prefixes_train_80pct.txt
+kn_rubin_ztf_train.txt
 --validation_filelist
-lsc240420_prefixes_validation_10pct.txt
+kn_rubin_ztf_val.txt
 --block_structure
 <B0>
 <B1>
diff --git a/applications/harnesses/KN_loderunner/training_input.tmpl b/applications/harnesses/KN_loderunner/training_input.tmpl
index 37420a29..42e4e5fe 100644
--- a/applications/harnesses/KN_loderunner/training_input.tmpl
+++ b/applications/harnesses/KN_loderunner/training_input.tmpl
@@ -5,9 +5,9 @@
 --LSC_NPZ_DIR
 /lustre/scratch5/exempt/artimis/data/lsc240420/
 --train_filelist
-lsc240420_prefixes_train_80pct.txt
+kn_rubin_ztf_train.txt
 --validation_filelist
-lsc240420_prefixes_validation_10pct.txt
+kn_rubin_ztf_val.txt
 --block_structure
 <B0>
 <B1>

From ee9d62f7d35856433e92de1319ed3b424f0bcfff Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 18 Aug 2026 20:11:31 -0600
Subject: [PATCH 38/66] plotting fixes

---
 .../harnesses/KN_loderunner/infer_9band.py    | 31 +++++++++++++++++--
 .../plot_pred_diagnostics_9band.py            | 31 ++++++++++++++++++-
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 2a2353ea..d8969b04 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -112,7 +112,17 @@ def get_args():
     parser.add_argument(
         "--norm_stats_path",
         type=str,
-        default="kilonova_9band_norm_stats.npz",
+        default="kilonova_9band_norm_stats_trainonly.npz",
+        help="Train-only normalization stats the model was trained with. Must "
+        "match training so forecasts use the exact encoding the model saw.",
+    )
+    parser.add_argument(
+        "--test_filelist",
+        type=str,
+        default=None,
+        help="Path to the test-split stem list (one object stem per line). When "
+        "set, only held-out test objects are forecast. Omit to use every object "
+        "in --data_glob.",
     )
 
     return parser.parse_args()
@@ -132,7 +142,7 @@ def resolve_paths(args):
     if args.data_glob is None:
         args.data_glob = (
             "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-            "uniform_dataset_20000/lc_*.npz"
+            "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
         )
 
     return tag
@@ -469,6 +479,23 @@ def main():
     if not files:
         raise RuntimeError(f"No files matched data_glob: {args.data_glob}")
 
+    # Restrict to held-out test objects when a split list is given, so forecasts
+    # are not shown on training data.
+    if args.test_filelist:
+        with open(args.test_filelist) as fh:
+            test_stems = {line.strip() for line in fh if line.strip()}
+        files = [
+            f
+            for f in files
+            if os.path.splitext(os.path.basename(f))[0] in test_stems
+        ]
+        if not files:
+            raise RuntimeError(
+                "No files in --data_glob matched the test split "
+                f"({args.test_filelist})."
+            )
+        print(f"Restricting forecasts to {len(files)} test-split objects.")
+
     files = files[: args.n_curves]
     print(f"Forecasting {len(files)} light curves.")
 
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index d53d7723..089f5bd6 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -122,7 +122,26 @@ def get_args():
     parser.add_argument(
         "--norm_stats_path",
         type=str,
-        default="kilonova_9band_norm_stats.npz",
+        default="kilonova_9band_norm_stats_trainonly.npz",
+        help="Train-only normalization stats the model was trained with. Must "
+        "match training so plots use the exact encoding the model saw.",
+    )
+    parser.add_argument(
+        "--data_glob",
+        type=str,
+        default=(
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+            "rubin_ztf_10000_dataset_same_seed/lc_*.npz"
+        ),
+        help="Glob for the realistic light-curve files to diagnose.",
+    )
+    parser.add_argument(
+        "--test_filelist",
+        type=str,
+        default=None,
+        help="Path to the test-split stem list (one object stem per line). When "
+        "set, diagnostics run ONLY on held-out test objects, so plots are not "
+        "leaked by training data. Omit to run over every object in --data_glob.",
     )
 
     return parser.parse_args()
@@ -231,6 +250,14 @@ def make_eval_dataset(
     print("band_means:", band_means)
     print("band_stds:", band_stds)
 
+    # Restrict to held-out test objects when a split list is given, so the
+    # diagnostics are not leaked by training data.
+    object_ids = None
+    if getattr(args, "test_filelist", None):
+        with open(args.test_filelist) as fh:
+            object_ids = {line.strip() for line in fh if line.strip()}
+        print(f"Restricting diagnostics to {len(object_ids)} test-split objects.")
+
     dataset = Kilonova_lc_scalar_context_DataSet_9band(
         N_imgs=args.N_imgs,
         context_len=context_len,
@@ -242,6 +269,8 @@ def make_eval_dataset(
         stds=band_stds,
         context_window_days=context_window_days,
         max_context_len=max_context_len,
+        data_glob=getattr(args, "data_glob", None),
+        object_ids=object_ids,
     )
 
     return dataset, np.asarray(band_means), np.asarray(band_stds)

From 620619b0986f70c966e74eec72d0b4872f6da6b9 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 19 Aug 2026 10:36:58 -0600
Subject: [PATCH 39/66] fix context window

---
 .../eval_dense_latetime_9band.py              | 152 ++++++++++++------
 1 file changed, 103 insertions(+), 49 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index 8afa24fd..cf706147 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -140,6 +140,45 @@ def _stem_to_path(data_glob: str) -> dict:
     return {_stem(f): f for f in glob.glob(data_glob)}
 
 
+def _batched_forward(
+    model: torch.nn.Module,
+    x: torch.Tensor,
+    lead_times: np.ndarray,
+    device: torch.device,
+    max_batch: int = 256,
+) -> np.ndarray:
+    """Predict all bands for many lead times in one (chunked) forward pass.
+
+    The context ``x`` (shape [1, D]) is fixed; only the lead time varies. Tiling
+    ``x`` to the batch dimension and passing a Dt vector runs every lead time
+    together instead of one-at-a-time, which is dramatically faster on GPU and
+    numerically identical to the per-point loop. Chunked at ``max_batch`` so a
+    long lead-time sweep cannot exhaust GPU memory.
+
+    Args:
+        model: The 9-band scalar-temporal LodeRunner.
+        x (torch.Tensor): Context input of shape [1, D].
+        lead_times (np.ndarray): 1-D array of lead times (days).
+        device (torch.device): Device to run on.
+        max_batch (int): Maximum lead times evaluated per forward pass.
+
+    Returns:
+        np.ndarray: Predictions of shape [len(lead_times), N_BANDS] (normalized).
+    """
+    lead_times = np.asarray(lead_times, dtype=np.float32)
+    out = np.zeros((lead_times.shape[0], N_BANDS), dtype=np.float32)
+    with torch.no_grad():
+        for start in range(0, lead_times.shape[0], max_batch):
+            chunk = lead_times[start : start + max_batch]
+            x_batch = x.expand(chunk.shape[0], -1)
+            Dt = torch.tensor(chunk, dtype=torch.float32, device=device)
+            pred = model(x_batch, in_vars=None, out_vars=None, Dt=Dt)
+            out[start : start + chunk.shape[0]] = (
+                pred.reshape(chunk.shape[0], N_BANDS).detach().cpu().numpy()
+            )
+    return out
+
+
 def eval_object(
     real_stream,
     dense_stream,
@@ -163,23 +202,37 @@ def eval_object(
     if r_t.shape[0] < 1 or d_t.shape[0] < 1:
         return None
 
-    # Phase zero = the first realistic detection (the observed trigger). The
-    # late-time region is dense points more than the cutoff past it.
+    # Phase zero = the first realistic detection (the observed trigger).
     t0 = float(r_t[0])
-    last_real_t = float(r_t[-1])
+
+    # The cutoff splits context from forecast: the model may only see realistic
+    # detections up to the cutoff phase, and must FORECAST everything after it
+    # (scored against the dense truth). Truncating the context here -- rather
+    # than feeding the whole realistic stream and only scoring late points --
+    # makes every object forecast from the same phase boundary, instead of from
+    # wherever its realistic coverage happens to end. (Without this, a
+    # bright/well-covered object whose realistic detections run to ~14 d has an
+    # almost-zero forecast horizon and the curve collapses to a stub.)
+    ctx_mask = (r_t - t0) <= late_time_cutoff_days
+    if not np.any(ctx_mask):
+        return None
+    r_t_ctx = r_t[ctx_mask]
+    r_v_ctx = r_v[ctx_mask]
+    r_b_ctx = r_b[ctx_mask]
+    last_real_t = float(r_t_ctx[-1])
 
     late_mask = (d_t - t0) > late_time_cutoff_days
     if not np.any(late_mask):
         return None
 
-    # Seed context from the realistic stream: trailing window ending at the last
-    # realistic detection, normalized as in training. build_context_input
-    # subtracts win_t[0], so absolute times are fine here.
-    r_v_norm = (r_v - means[r_b]) / (stds[r_b] + EPS)
+    # Seed context from the truncated realistic stream: trailing window ending
+    # at the last pre-cutoff realistic detection, normalized as in training.
+    # build_context_input subtracts win_t[0], so absolute times are fine here.
+    r_v_norm = (r_v_ctx - means[r_b_ctx]) / (stds[r_b_ctx] + EPS)
     win_v, win_t, win_b = _select_window(
-        ctx_t=list(r_t.astype(np.float32)),
+        ctx_t=list(r_t_ctx.astype(np.float32)),
         ctx_v=list(r_v_norm.astype(np.float32)),
-        ctx_b=list(r_b),
+        ctx_b=list(r_b_ctx),
         context_window_days=context_window_days,
         max_context_len=max_context_len,
     )
@@ -194,46 +247,44 @@ def eval_object(
     )
 
     # Score each late-time dense point at its true lead time from the last
-    # realistic detection.
-    scored = []
-    with torch.no_grad():
-        for idx in np.nonzero(late_mask)[0]:
-            dt = float(d_t[idx]) - last_real_t
-            if dt <= 0:
-                # Dense point precedes the last realistic detection; not a
-                # forecast into the future. Skip.
-                continue
-            Dt = torch.tensor([dt], dtype=torch.float32, device=device)
-            pred_all = model(x, in_vars=None, out_vars=None, Dt=Dt)
-            pred_all = pred_all.reshape(N_BANDS).detach().cpu().numpy()
-
-            band = int(d_b[idx])
-            pred_mag = float(pred_all[band] * (stds[band] + EPS) + means[band])
-            true_mag = float(d_v[idx])
-            scored.append(
-                {
-                    "phase": float(d_t[idx]) - t0,
-                    "lead_time": dt,
-                    "band": band,
-                    "pred_mag": pred_mag,
-                    "true_mag": true_mag,
-                    "residual_mag": pred_mag - true_mag,
-                }
-            )
-
-    if not scored:
+    # realistic detection. The context ``x`` is fixed for this object, so all
+    # lead times are evaluated in a SINGLE batched forward pass (tile x to the
+    # batch dimension, pass a Dt vector) instead of one forward per point --
+    # numerically identical, but far faster on GPU.
+    late_idx = np.nonzero(late_mask)[0]
+    lead_times = (d_t[late_idx] - last_real_t).astype(np.float32)
+    # Only points strictly after the last realistic detection are forecasts.
+    keep = lead_times > 0
+    late_idx = late_idx[keep]
+    lead_times = lead_times[keep]
+
+    if late_idx.shape[0] == 0:
         return None
 
+    pred_scored = _batched_forward(model, x, lead_times, device)  # [P, N_BANDS]
+
+    scored = []
+    for j, idx in enumerate(late_idx):
+        band = int(d_b[idx])
+        pred_mag = float(pred_scored[j, band] * (stds[band] + EPS) + means[band])
+        true_mag = float(d_v[idx])
+        scored.append(
+            {
+                "phase": float(d_t[idx]) - t0,
+                "lead_time": float(lead_times[j]),
+                "band": band,
+                "pred_mag": pred_mag,
+                "true_mag": true_mag,
+                "residual_mag": pred_mag - true_mag,
+            }
+        )
+
     # Smooth forecast curve for plotting: sweep lead time from 0 to the farthest
-    # scored late-time point, predicting all bands at each lead time.
-    max_dt = max(s["lead_time"] for s in scored)
+    # scored late-time point, predicting all bands at each lead time -- also a
+    # single batched forward pass.
+    max_dt = float(lead_times.max())
     lead_grid = np.linspace(0.0, max_dt, 60).astype(np.float32)
-    curve = np.zeros((lead_grid.shape[0], N_BANDS), dtype=np.float32)
-    with torch.no_grad():
-        for k, dt in enumerate(lead_grid):
-            Dt = torch.tensor([dt], dtype=torch.float32, device=device)
-            pred = model(x, in_vars=None, out_vars=None, Dt=Dt)
-            curve[k] = pred.reshape(N_BANDS).detach().cpu().numpy()
+    curve = _batched_forward(model, x, lead_grid, device)  # [60, N_BANDS]
     curve_mag = curve * (stds[None, :] + EPS) + means[None, :]
 
     return {
@@ -242,7 +293,9 @@ def eval_object(
         "last_real_t": last_real_t,
         "curve_phase": (last_real_t - t0) + lead_grid,
         "curve_mag": curve_mag.astype(np.float32),
-        "real": (r_t - t0, r_v, r_b),
+        # Only the pre-cutoff realistic detections were shown to the model, so
+        # plot those as the context (not the full realistic stream).
+        "real": (r_t_ctx - t0, r_v_ctx, r_b_ctx),
         "dense": (d_t - t0, d_v, d_b),
     }
 
@@ -319,9 +372,10 @@ def get_args():
     p.add_argument(
         "--late_time_cutoff_days",
         type=float,
-        default=8.0,
-        help="Dense points with phase (from first realistic detection) greater "
-        "than this are the late-time region scored here.",
+        default=3.0,
+        help="Splits context from forecast. The model sees realistic detections "
+        "with phase (from first realistic detection) up to this value, and "
+        "forecasts all dense points after it -- the late-time region scored here.",
     )
     p.add_argument("--outdir", type=str, default=None)
     p.add_argument(

From 77321479292ad51c17651ecb3fc0f5a2aaead62c Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 19 Aug 2026 11:14:40 -0600
Subject: [PATCH 40/66] Dt encoding change

---
 .../harnesses/KN_loderunner/infer_9band.py    |  6 ++
 .../plot_pred_diagnostics_9band.py            |  6 ++
 .../KN_loderunner/train_LodeRunner_ddp.py     | 10 +++
 src/yoke/models/vit/swin/bomberman.py         | 79 ++++++++++++++++---
 src/yoke/utils/checkpointing.py               |  3 +
 5 files changed, 95 insertions(+), 9 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index d8969b04..3bd88116 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -176,6 +176,10 @@ def load_9band_model(ckpt_path, device):
     else:
         max_context_len = context_len
 
+    # 0 for legacy checkpoints (no key) -> Fourier Dt disabled -> matches saved
+    # weights so strict load succeeds.
+    dt_fourier_bands = ckpt.get("dt_fourier_bands", 0)
+
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
     print("target_type:", ckpt.get("target_type", "unknown"))
@@ -183,6 +187,7 @@ def load_9band_model(ckpt_path, device):
     print("context_window_days:", context_window_days)
     print("max_context_len:", max_context_len)
     print("n_bands:", n_bands)
+    print("dt_fourier_bands:", dt_fourier_bands)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -195,6 +200,7 @@ def load_9band_model(ckpt_path, device):
         backbone_channels=backbone_channels,
         hidden=hidden,
         context_window_days=context_window_days,
+        dt_fourier_bands=dt_fourier_bands,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 089f5bd6..3e1446dd 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -194,6 +194,10 @@ def load_9band_model(ckpt_path, device):
     else:
         max_context_len = context_len
 
+    # 0 for legacy checkpoints (no key) -> Fourier Dt disabled -> matches saved
+    # weights so strict load succeeds.
+    dt_fourier_bands = ckpt.get("dt_fourier_bands", 0)
+
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
     print("backbone_class:", ckpt.get("backbone_class", "LodeRunner"))
@@ -205,6 +209,7 @@ def load_9band_model(ckpt_path, device):
     print("band_keys:", ckpt.get("band_keys", list(BAND_KEYS)))
     print("backbone_channels:", backbone_channels)
     print("hidden:", hidden)
+    print("dt_fourier_bands:", dt_fourier_bands)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -217,6 +222,7 @@ def load_9band_model(ckpt_path, device):
         backbone_channels=backbone_channels,
         hidden=hidden,
         context_window_days=context_window_days,
+        dt_fourier_bands=dt_fourier_bands,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 585ad5fe..7fec9dd7 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -227,6 +227,14 @@ def main(args, rank, world_size, local_rank, device):
     CONTEXT_LEN = 5 #3
     HIDDEN_CHANNELS = 64
 
+    # Fourier lead-time conditioning. When > 0, the trainable conditioner and
+    # output head receive a 2*DT_FOURIER_BANDS sinusoidal encoding of the lead
+    # time Dt, so they can learn a real per-band decay curve instead of a flat
+    # persistence value. (Without this, Dt reaches the output only through the
+    # frozen backbone, which cannot adapt, so late-time forecasts plateau.) Set
+    # to 0 for the legacy architecture (byte-identical; old checkpoints load).
+    DT_FOURIER_BANDS = 8
+
     # Time-window context mode. When CONTEXT_WINDOW_DAYS is not None, the dataset
     # selects context by a trailing lookback in days (all detections within the
     # last CONTEXT_WINDOW_DAYS), padded to MAX_CONTEXT_LEN with a per-event
@@ -351,6 +359,7 @@ def main(args, rank, world_size, local_rank, device):
             backbone_channels=8,
             hidden=HIDDEN_CHANNELS,
             context_window_days=CONTEXT_WINDOW_DAYS,
+            dt_fourier_bands=DT_FOURIER_BANDS,
         ).to(device)
 
         # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
@@ -703,6 +712,7 @@ def _make_9band(
                     "band_keys": list(BAND_KEYS),
                     "backbone_channels": 8,
                     "hidden": HIDDEN_CHANNELS,
+                    "dt_fourier_bands": DT_FOURIER_BANDS,
                     "n_rollout_steps": n_rollout_steps,
                     "context_window_days": CONTEXT_WINDOW_DAYS,
                     "max_context_len": MAX_CONTEXT_LEN,
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index c163ecb8..3ec5e205 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -9,6 +9,7 @@
 """
 
 from collections.abc import Callable, Iterable
+import math
 import random
 
 import numpy as np
@@ -482,6 +483,7 @@ def __init__(
         backbone_channels: int = 8,
         hidden: int = 64,
         context_window_days: float = None,
+        dt_fourier_bands: int = 0,
     ) -> None:
         """Initialize conditioner and output-head around the backbone.
 
@@ -493,6 +495,18 @@ def __init__(
                 flag, so each event carries an extra ``valid`` feature and the
                 per-event width is ``3 + n_bands`` instead of ``2 + n_bands``.
                 When None (default), the legacy fixed-count layout is used.
+            dt_fourier_bands (int): Number of Fourier (sinusoidal) frequency
+                bands used to encode the lead time ``Dt`` and inject it into the
+                trainable conditioner and output head. When ``0`` (default) the
+                feature is DISABLED and the architecture is byte-identical to the
+                legacy model: neither MLP sees ``Dt`` directly (lead time reaches
+                the output only through the frozen backbone). When ``> 0``, a
+                ``2 * dt_fourier_bands``-wide encoding ``[sin(Dt·f), cos(Dt·f)]``
+                over a fixed log-spaced frequency bank is concatenated onto BOTH
+                MLP inputs, so the trainable path can learn a smooth, nonlinear,
+                per-band dependence on lead time (e.g. late-time decay) rather
+                than a lead-time-independent persistence value. The frequency
+                bank is a non-trainable buffer.
         """
         super().__init__()
 
@@ -502,6 +516,7 @@ def __init__(
         self.image_size = image_size
         self.backbone_channels = backbone_channels
         self.context_window_days = context_window_days
+        self.dt_fourier_bands = dt_fourier_bands
 
         # Dataset x layout, flattened per event. Fixed-count mode:
         #   [value, rel_t, one_hot_band(n_bands)] * context_len   -> 2 + n_bands
@@ -510,23 +525,57 @@ def __init__(
         per_event_width = 3 + n_bands if context_window_days is not None else 2 + n_bands
         input_dim = context_len * per_event_width
 
-        # Maps the scalar temporal event stream into the pseudo-channels
-        # expected by the pretrained LodeRunner backbone.
+        # Fourier lead-time encoding. When enabled, a fixed log-spaced frequency
+        # bank (periods ~0.5 -> 15 days, matching the forecast horizon) turns the
+        # scalar Dt into a 2*dt_fourier_bands feature vector that the trainable
+        # MLPs consume directly. Registered as a buffer so it saves/loads and
+        # moves with .to(device) but is excluded from the optimizer's param list.
+        if dt_fourier_bands > 0:
+            periods = torch.logspace(
+                math.log10(0.5), math.log10(15.0), dt_fourier_bands
+            )
+            self.register_buffer("dt_freqs", 2.0 * math.pi / periods)
+            dt_extra = 2 * dt_fourier_bands
+        else:
+            dt_extra = 0
+
+        # Maps the scalar temporal event stream (plus the Fourier Dt encoding,
+        # when enabled) into the pseudo-channels expected by the backbone.
         self.conditioner = nn.Sequential(
-            nn.Linear(input_dim, hidden),
+            nn.Linear(input_dim + dt_extra, hidden),
             nn.GELU(),
             nn.Linear(hidden, hidden),
             nn.GELU(),
             nn.Linear(hidden, backbone_channels),
         )
 
-        # Maps the backbone-channel summary back to one prediction per band.
+        # Maps the backbone-channel summary (plus the Fourier Dt encoding, when
+        # enabled) back to one prediction per band.
         self.output_head = nn.Sequential(
-            nn.Linear(backbone_channels, hidden),
+            nn.Linear(backbone_channels + dt_extra, hidden),
             nn.GELU(),
             nn.Linear(hidden, n_bands),
         )
 
+    def _encode_dt(self, Dt: torch.Tensor, batch_size: int) -> torch.Tensor:
+        """Fourier-encode the lead time for the trainable path.
+
+        Args:
+            Dt (torch.Tensor): Lead-time tensor of shape [B] (or broadcastable).
+            batch_size (int): Batch size B, used to size the disabled-path output.
+
+        Returns:
+            torch.Tensor: [B, 2 * dt_fourier_bands] of ``[sin(Dt·f), cos(Dt·f)]``
+            when enabled, else an empty [B, 0] tensor (so the concat is a no-op
+            and the disabled path is identical to the legacy model).
+        """
+        if self.dt_fourier_bands == 0:
+            return Dt.new_zeros((batch_size, 0))
+
+        # [B, 1] * [1, bands] -> [B, bands]
+        angles = Dt.reshape(batch_size, 1) * self.dt_freqs.reshape(1, -1)
+        return torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)
+
     def forward(
         self,
         x: torch.Tensor,
@@ -541,7 +590,10 @@ def forward(
                 [B, context_len * (2 + n_bands)].
             in_vars (torch.Tensor): Kept for LodeRunner API compatibility.
             out_vars (torch.Tensor): Kept for LodeRunner API compatibility.
-            Dt (torch.Tensor): Lead-time tensor passed to the backbone.
+            Dt (torch.Tensor): Lead-time tensor. Passed to the backbone and, when
+                ``dt_fourier_bands > 0``, Fourier-encoded and concatenated onto
+                both the conditioner and output-head inputs so the trainable path
+                can condition directly on lead time.
 
         Returns:
             pred (torch.Tensor): Predictions of shape [B, n_bands].
@@ -549,7 +601,13 @@ def forward(
         B = x.shape[0]
         H, W = self.image_size
 
-        channel_vals = self.conditioner(x)  # [B, backbone_channels]
+        # Fourier Dt encoding for the trainable path ([B, 0] when disabled, so
+        # both concats below are no-ops and match the legacy architecture).
+        dt_feat = self._encode_dt(Dt, B)
+
+        channel_vals = self.conditioner(
+            torch.cat([x, dt_feat], dim=1)
+        )  # [B, backbone_channels]
 
         pseudo_img = channel_vals.view(
             B,
@@ -576,8 +634,11 @@ def forward(
         # Collapse spatial dimensions to backbone-channel summaries.
         pred_channel_vals = pred_img.mean(dim=(2, 3))  # [B, backbone_channels]
 
-        # Convert backbone channels to per-band predictions.
-        pred = self.output_head(pred_channel_vals)  # [B, n_bands]
+        # Convert backbone channels to per-band predictions, conditioning the
+        # head on lead time via the same Fourier encoding ([B, 0] when disabled).
+        pred = self.output_head(
+            torch.cat([pred_channel_vals, dt_feat], dim=1)
+        )  # [B, n_bands]
 
         return pred
 
diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index 1af2c726..44297c9f 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -496,6 +496,9 @@ def load_direct_loderunner_checkpoint_9band(
         backbone_channels=checkpoint_data.get("backbone_channels", 8),
         hidden=checkpoint_data.get("hidden", 64),
         context_window_days=context_window_days,
+        # 0 for legacy checkpoints (no key) -> Fourier Dt disabled -> the
+        # architecture matches the saved weights and strict load succeeds.
+        dt_fourier_bands=checkpoint_data.get("dt_fourier_bands", 0),
     ).to(device)
 
     state_dict = checkpoint_data["model_state_dict"]

From b21352be1942dad1f98acd86af7f776140616647 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 19 Aug 2026 13:20:24 -0600
Subject: [PATCH 41/66] plotting speed up

---
 .../plot_pred_diagnostics_9band.py            | 45 ++++++++++++++++---
 1 file changed, 39 insertions(+), 6 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 3e1446dd..3a558628 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -336,6 +336,40 @@ def build_context_input(
     ).unsqueeze(0)
 
 
+def _batched_forward(model, x, lead_times, n_bands, device, max_batch=256):
+    """Predict all bands for many lead times in one (chunked) forward pass.
+
+    The context ``x`` (shape [1, D]) is fixed; only the lead time varies. Tiling
+    ``x`` to the batch dimension and passing a Dt vector runs every lead time
+    together instead of one-at-a-time, which is dramatically faster on GPU and
+    numerically identical to the per-lead-time loop. Chunked at ``max_batch`` so
+    a long lead-time sweep cannot exhaust GPU memory.
+
+    Args:
+        model: The 9-band scalar-temporal LodeRunner.
+        x (torch.Tensor): Context input of shape [1, D].
+        lead_times (np.ndarray): 1-D array of lead times (days).
+        n_bands (int): Number of bands the model emits.
+        device (torch.device): Device to run on.
+        max_batch (int): Maximum lead times evaluated per forward pass.
+
+    Returns:
+        np.ndarray: Predictions of shape [len(lead_times), n_bands] (normalized).
+    """
+    lead_times = np.asarray(lead_times, dtype=np.float32)
+    out = np.zeros((lead_times.shape[0], n_bands), dtype=np.float32)
+    with torch.no_grad():
+        for start in range(0, lead_times.shape[0], max_batch):
+            chunk = lead_times[start : start + max_batch]
+            x_batch = x.expand(chunk.shape[0], -1)
+            Dt = torch.tensor(chunk, dtype=torch.float32, device=device)
+            pred = model(x_batch, in_vars=None, out_vars=None, Dt=Dt)
+            out[start : start + chunk.shape[0]] = (
+                pred.reshape(chunk.shape[0], n_bands).detach().cpu().numpy()
+            )
+    return out
+
+
 def _select_window(ctx_t, ctx_v, ctx_b, context_window_days, max_context_len):
     """Select the trailing time-window subset of a growing context.
 
@@ -561,13 +595,12 @@ def get_rollout_from_stream(
 
         n_lead = 60
         lead_times = np.linspace(0.0, horizon, n_lead).astype(np.float32)
-        preds_norm = np.zeros((n_lead, n_bands), dtype=np.float32)
 
-        with torch.no_grad():
-            for k, dt in enumerate(lead_times):
-                Dt = torch.tensor([dt], dtype=torch.float32, device=device)
-                pred = model(x0, in_vars=None, out_vars=None, Dt=Dt)
-                preds_norm[k] = pred.reshape(n_bands).detach().cpu().numpy()
+        # Context x0 is fixed; only the lead time varies. Evaluate the whole
+        # sweep in a single batched forward pass (tile x0 to the batch dim, pass
+        # a Dt vector) instead of one forward per lead time -- numerically
+        # identical, but far faster on GPU.
+        preds_norm = _batched_forward(model, x0, lead_times, n_bands, device)
 
         preds_mag = preds_norm * (stds[None, :] + EPS) + means[None, :]
 

From d2ee98b7b29f52688fd386a63a005bfb47fbd269 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 19 Aug 2026 15:17:36 -0600
Subject: [PATCH 42/66] trying to improve training + late time rollouts

---
 .../eval_dense_latetime_9band.py              | 19 ++++++-
 src/yoke/models/vit/swin/bomberman.py         | 56 +++++++++++++------
 2 files changed, 57 insertions(+), 18 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index cf706147..cbd58b3d 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -189,12 +189,18 @@ def eval_object(
     context_window_days,
     max_context_len,
     late_time_cutoff_days,
+    late_time_max_days,
 ):
     """Score one object's late-time dense truth against a realistic-context forecast.
 
     Returns a dict with the scored late-time points and a smooth forecast curve,
     or None if the object cannot be evaluated (no realistic context, or no dense
     points in the late-time region).
+
+    The scored/forecast region is the phase band
+    ``late_time_cutoff_days < phase <= late_time_max_days`` (phase measured from
+    the first realistic detection). Points beyond ``late_time_max_days`` are
+    ignored so the forecast is only judged over a horizon we care about.
     """
     r_t, r_v, r_b = real_stream
     d_t, d_v, d_b = dense_stream
@@ -221,7 +227,9 @@ def eval_object(
     r_b_ctx = r_b[ctx_mask]
     last_real_t = float(r_t_ctx[-1])
 
-    late_mask = (d_t - t0) > late_time_cutoff_days
+    # Score the forecast only within the phase band cutoff < phase <= max_days.
+    d_phase = d_t - t0
+    late_mask = (d_phase > late_time_cutoff_days) & (d_phase <= late_time_max_days)
     if not np.any(late_mask):
         return None
 
@@ -377,6 +385,14 @@ def get_args():
         "with phase (from first realistic detection) up to this value, and "
         "forecasts all dense points after it -- the late-time region scored here.",
     )
+    p.add_argument(
+        "--late_time_max_days",
+        type=float,
+        default=10.0,
+        help="Upper bound (phase from first realistic detection) on the scored "
+        "forecast region. Dense points beyond this are ignored, so the forecast "
+        "is judged only over cutoff < phase <= this horizon.",
+    )
     p.add_argument("--outdir", type=str, default=None)
     p.add_argument(
         "--max_objects",
@@ -467,6 +483,7 @@ def main():
             context_window_days,
             max_context_len,
             args.late_time_cutoff_days,
+            args.late_time_max_days,
         )
         if result is None:
             continue
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index 3ec5e205..54a38185 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -501,12 +501,16 @@ def __init__(
                 feature is DISABLED and the architecture is byte-identical to the
                 legacy model: neither MLP sees ``Dt`` directly (lead time reaches
                 the output only through the frozen backbone). When ``> 0``, a
-                ``2 * dt_fourier_bands``-wide encoding ``[sin(Dt·f), cos(Dt·f)]``
-                over a fixed log-spaced frequency bank is concatenated onto BOTH
-                MLP inputs, so the trainable path can learn a smooth, nonlinear,
-                per-band dependence on lead time (e.g. late-time decay) rather
-                than a lead-time-independent persistence value. The frequency
-                bank is a non-trainable buffer.
+                ``2 * dt_fourier_bands + 1``-wide encoding
+                ``[sin(Dt·f), cos(Dt·f), log1p(Dt)]`` -- a fixed log-spaced
+                frequency bank (periods ~0.5 -> 60 days) plus one non-periodic
+                monotone channel -- is concatenated onto BOTH MLP inputs, so the
+                trainable path can learn a smooth, nonlinear, per-band dependence
+                on lead time (e.g. late-time decay) rather than a
+                lead-time-independent persistence value. The frequency bank is a
+                non-trainable buffer. The monotone ``log1p(Dt)`` channel gives an
+                explicit long-term trend so the forecast does not curve back up
+                at long lead times (which a purely periodic basis would).
         """
         super().__init__()
 
@@ -526,16 +530,28 @@ def __init__(
         input_dim = context_len * per_event_width
 
         # Fourier lead-time encoding. When enabled, a fixed log-spaced frequency
-        # bank (periods ~0.5 -> 15 days, matching the forecast horizon) turns the
-        # scalar Dt into a 2*dt_fourier_bands feature vector that the trainable
-        # MLPs consume directly. Registered as a buffer so it saves/loads and
-        # moves with .to(device) but is excluded from the optimizer's param list.
+        # bank turns the scalar Dt into a 2*dt_fourier_bands feature vector that
+        # the trainable MLPs consume directly. Registered as a buffer so it
+        # saves/loads and moves with .to(device) but is excluded from the
+        # optimizer's param list.
+        #
+        # The longest period is 60 days, well beyond any plausible forecast
+        # horizon: a purely sinusoidal basis MUST turn around (each component
+        # bottoms at half its period), which produced an unphysical "smile" in
+        # the late-time forecast (fade, then rise) as the lowest-frequency band
+        # -- previously 15 d, ~one full cycle over the window -- swung back up.
+        # With a 60 d max period the slowest component covers only a fraction of
+        # a cycle across a <=15 d horizon, so it stays monotone. In addition, a
+        # single non-periodic log1p(Dt) channel is appended (see _encode_dt),
+        # giving the MLPs an explicit monotone ramp for the long-term trend while
+        # the Fourier bank handles short-timescale structure.
         if dt_fourier_bands > 0:
             periods = torch.logspace(
-                math.log10(0.5), math.log10(15.0), dt_fourier_bands
+                math.log10(0.5), math.log10(60.0), dt_fourier_bands
             )
             self.register_buffer("dt_freqs", 2.0 * math.pi / periods)
-            dt_extra = 2 * dt_fourier_bands
+            # 2 * bands (sin + cos) + 1 monotone log1p(Dt) channel.
+            dt_extra = 2 * dt_fourier_bands + 1
         else:
             dt_extra = 0
 
@@ -565,16 +581,22 @@ def _encode_dt(self, Dt: torch.Tensor, batch_size: int) -> torch.Tensor:
             batch_size (int): Batch size B, used to size the disabled-path output.
 
         Returns:
-            torch.Tensor: [B, 2 * dt_fourier_bands] of ``[sin(Dt·f), cos(Dt·f)]``
-            when enabled, else an empty [B, 0] tensor (so the concat is a no-op
-            and the disabled path is identical to the legacy model).
+            torch.Tensor: [B, 2 * dt_fourier_bands + 1] of
+            ``[sin(Dt·f), cos(Dt·f), log1p(Dt)]`` when enabled, else an empty
+            [B, 0] tensor (so the concat is a no-op and the disabled path is
+            identical to the legacy model). The trailing ``log1p(Dt)`` is a
+            non-periodic, monotone channel so the trainable path always has an
+            explicit "further ahead -> keep fading" ramp, independent of the
+            periodic bands (which alone would eventually curve back up).
         """
         if self.dt_fourier_bands == 0:
             return Dt.new_zeros((batch_size, 0))
 
+        Dt = Dt.reshape(batch_size, 1)
         # [B, 1] * [1, bands] -> [B, bands]
-        angles = Dt.reshape(batch_size, 1) * self.dt_freqs.reshape(1, -1)
-        return torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)
+        angles = Dt * self.dt_freqs.reshape(1, -1)
+        mono = torch.log1p(Dt.clamp_min(0.0))  # [B, 1], monotone in lead time
+        return torch.cat([torch.sin(angles), torch.cos(angles), mono], dim=1)
 
     def forward(
         self,

From 893b2775ff1592799ea848a5a632381bc9ad2351 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 20 Aug 2026 11:32:06 -0600
Subject: [PATCH 43/66] Loss re-weighting

---
 .../eval_dense_latetime_9band.py              | 11 +++
 .../KN_loderunner/train_LodeRunner_ddp.py     | 31 +++++++
 src/yoke/utils/training/epoch/loderunner.py   | 85 +++++++++++++++++--
 3 files changed, 120 insertions(+), 7 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index cbd58b3d..320b963f 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -305,6 +305,9 @@ def eval_object(
         # plot those as the context (not the full realistic stream).
         "real": (r_t_ctx - t0, r_v_ctx, r_b_ctx),
         "dense": (d_t - t0, d_v, d_b),
+        # Right edge for plotting: the scored horizon. Beyond this the forecast
+        # is unsupervised extrapolation, so it is not shown.
+        "plot_max_phase": late_time_max_days,
     }
 
 
@@ -314,11 +317,17 @@ def plot_object(result, stem, outpath):
     axes = axes.ravel()
     r_ph, r_v, r_b = result["real"]
     d_ph, d_v, d_b = result["dense"]
+    # Show only the scored horizon; the forecast beyond it is unsupervised
+    # extrapolation (where the late-time upturn artifact lives).
+    plot_max = result.get("plot_max_phase")
 
     for b in range(N_BANDS):
         ax = axes[b]
         rm = r_b == b
         dm = d_b == b
+        # Clip the dense-truth scatter to the plotted horizon as well.
+        if plot_max is not None:
+            dm = dm & (d_ph <= plot_max)
         if np.any(dm):
             ax.scatter(d_ph[dm], d_v[dm], s=14, c="0.6", label="dense truth")
         if np.any(rm):
@@ -331,6 +340,8 @@ def plot_object(result, stem, outpath):
             c=BAND_COLORS[b], lw=1.6, label="forecast",
         )
         ax.invert_yaxis()  # magnitudes: brighter is smaller
+        if plot_max is not None:
+            ax.set_xlim(right=plot_max)
         ax.set_title(BAND_NAMES[b], fontsize=9)
         if b == 0:
             ax.legend(fontsize=7, loc="best")
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 7fec9dd7..6f6dfb65 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -279,6 +279,30 @@ def main(args, rank, world_size, local_rank, device):
     # detections; normalization statistics are computed the same way.
     DROP_UPPER_LIMITS = True
 
+    # Per-band loss weighting. Targets are per-band z-scored, so an equal-weight
+    # loss lets the large-dynamic-range blue bands (u, g fade to mag ~28-30) be
+    # under-fit -- the dense-eval showed a strong under-fade bias there (u
+    # ~ -5 mag, g ~ -2 mag) while ZTF/red bands fit well. Up-weighting u and g
+    # in the TRAINING backward (the recorded per-sample loss stays unweighted so
+    # the val CSV remains a comparable yardstick) pushes gradient toward the
+    # blue fade the model currently ignores. Order matches BAND_KEYS =
+    # (ztfg, ztfr, ztfi, sdssu, ps1_g, ps1_r, ps1_i, ps1_z, ps1_y). Set to None
+    # to recover the exact equal-weight objective.
+    BAND_WEIGHTS = torch.tensor(
+        [
+            1.0,  # ztfg
+            1.0,  # ztfr
+            1.0,  # ztfi
+            3.0,  # sdssu (u) -- worst under-fade, largest up-weight
+            2.0,  # ps1__g (g)
+            1.0,  # ps1__r (r)
+            1.0,  # ps1__i (i)
+            1.0,  # ps1__z (z)
+            1.0,  # ps1__y (y)
+        ],
+        dtype=torch.float32,
+    )
+
     optimizer_kwargs = {
         "lr": 1e-4,# 1e-4, #1e-5
         "betas": (0.9, 0.999),
@@ -648,6 +672,7 @@ def _make_9band(
                 window_mode=CONTEXT_WINDOW_DAYS is not None,
                 context_window_days=CONTEXT_WINDOW_DAYS,
                 max_context_len=MAX_CONTEXT_LEN,
+                band_weights=BAND_WEIGHTS,
             )
         else:
             #train_DDP_loderunner_epoch(
@@ -667,6 +692,7 @@ def _make_9band(
                 device=device,
                 rank=rank,
                 world_size=world_size,
+                band_weights=BAND_WEIGHTS,
             )
 
         print(f"[rank {rank}] finished epoch", flush=True)
@@ -713,6 +739,11 @@ def _make_9band(
                     "backbone_channels": 8,
                     "hidden": HIDDEN_CHANNELS,
                     "dt_fourier_bands": DT_FOURIER_BANDS,
+                    "band_weights": (
+                        BAND_WEIGHTS.tolist()
+                        if BAND_WEIGHTS is not None
+                        else None
+                    ),
                     "n_rollout_steps": n_rollout_steps,
                     "context_window_days": CONTEXT_WINDOW_DAYS,
                     "max_context_len": MAX_CONTEXT_LEN,
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 36ac303e..cc271f10 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -497,6 +497,7 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
     device: torch.device,
     rank: int,
     world_size: int,
+    band_weights: torch.Tensor = None,
 ) -> None:
     """DDP epoch function for the masked 9-band scalar temporal LodeRunner.
 
@@ -512,6 +513,13 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
 
     Expected model output:
         pred:   [B, n_bands]
+
+    ``band_weights`` (optional [n_bands] tensor) up-weights the backward pass
+    per observed band. Because targets are per-band z-scored, equal weighting
+    lets large-dynamic-range bands (u, g) contribute little to the loss and be
+    under-fit; weighting scales each sample by its observed band's weight. The
+    RECORDED per-sample loss stays unweighted so the CSV metric is comparable
+    across runs. ``None`` (default) reproduces the plain equal-weight behavior.
     """
     train_rcrd_filename = train_rcrd_filename.replace(
         "<epochIDX>",
@@ -520,6 +528,9 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
 
     model.train()
 
+    if band_weights is not None:
+        band_weights = band_weights.to(device)
+
     with (
         open(train_rcrd_filename, "a") if rank == 0 else nullcontext()
     ) as train_rcrd_file:
@@ -554,7 +565,18 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
             loss = loss_fn(pred, target) * mask
             per_sample_loss = loss.sum(dim=1) / (mask.sum(dim=1) + 1e-8)
 
-            batch_loss = per_sample_loss.mean()
+            if band_weights is None:
+                # Recorded metric == training objective: plain equal weight.
+                batch_loss = per_sample_loss.mean()
+            else:
+                # Weight the backward by the observed band's weight (mask is
+                # one-hot, so this picks each sample's band weight). The
+                # recorded per_sample_loss above stays unweighted so the CSV
+                # metric is comparable across runs and weightings.
+                sample_w = (mask * band_weights.reshape(1, -1)).sum(dim=1)
+                batch_loss = (per_sample_loss * sample_w).sum() / (
+                    sample_w.sum() + 1e-8
+                )
 
             batch_loss.backward()
             optimizer.step()
@@ -636,6 +658,7 @@ def _rollout_pass_9band(
     n_bands: int,
     teacher_forcing_ratio: float,
     device: torch.device,
+    band_weights: torch.Tensor = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Unroll the 9-band model over a batch of rollouts with scheduled sampling.
 
@@ -660,10 +683,16 @@ def _rollout_pass_9band(
         teacher_forcing_ratio (float): Probability of feeding the true value back
             at each step (1.0 = fully teacher-forced, 0.0 = fully free-running).
         device (torch.device): Compute device.
+        band_weights (torch.Tensor): Optional per-band gradient weights
+            [n_bands]. When given, ``total_loss`` (the backward objective) is a
+            per-band weighted mean over valid steps; ``per_sample_loss`` (the
+            recorded metric) stays unweighted. ``None`` reproduces the plain
+            equal-weight behavior exactly.
 
     Returns:
         per_sample_loss (torch.Tensor): Mean rollout loss per sample [B].
-        total_loss (torch.Tensor): Scalar mean loss over all valid steps.
+        total_loss (torch.Tensor): Scalar mean loss over all valid steps
+            (per-band weighted when ``band_weights`` is given).
     """
     B, context_len = ctx_v.shape
     n_steps = future_v.shape[1]
@@ -681,6 +710,7 @@ def _rollout_pass_9band(
 
     step_losses = []  # [B] per valid step
     step_valid = []  # [B] per step
+    step_bands = []  # [B] observed band index per step (for band weighting)
 
     for step in range(n_steps):
         # Build the flattened per-event input from the current window, with
@@ -709,6 +739,7 @@ def _rollout_pass_9band(
         step_loss = loss_fn(pred_obs, true_obs) * valid
         step_losses.append(step_loss)
         step_valid.append(valid)
+        step_bands.append(tgt_band)
 
         # Scheduled sampling: choose true vs own (detached) prediction per sample.
         use_true = (
@@ -730,7 +761,18 @@ def _rollout_pass_9band(
     step_valid = torch.stack(step_valid, dim=1)  # [B, n_steps]
 
     per_sample_loss = step_losses.sum(dim=1) / (step_valid.sum(dim=1) + 1e-8)
-    total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
+
+    if band_weights is None:
+        total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
+    else:
+        # Per-band-weighted objective: scale each valid step by its observed
+        # band's weight. per_sample_loss (recorded) stays unweighted above.
+        step_bands = torch.stack(step_bands, dim=1)  # [B, n_steps]
+        bw = band_weights.to(device)
+        step_w = bw[step_bands] * step_valid  # [B, n_steps]
+        total_loss = (step_losses * bw[step_bands]).sum() / (
+            step_w.sum() + 1e-8
+        )
 
     return per_sample_loss, total_loss
 
@@ -751,6 +793,7 @@ def _rollout_pass_9band_window(
     max_context_len: int,
     teacher_forcing_ratio: float,
     device: torch.device,
+    band_weights: torch.Tensor = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Unroll the 9-band model over a batch of rollouts in time-window mode.
 
@@ -787,10 +830,16 @@ def _rollout_pass_9band_window(
         teacher_forcing_ratio (float): Probability of feeding the true value back
             at each step (1.0 = fully teacher-forced, 0.0 = fully free-running).
         device (torch.device): Compute device.
+        band_weights (torch.Tensor): Optional per-band gradient weights
+            [n_bands]. When given, ``total_loss`` (the backward objective) is a
+            per-band weighted mean over valid steps; ``per_sample_loss`` (the
+            recorded metric) stays unweighted. ``None`` reproduces the plain
+            equal-weight behavior exactly.
 
     Returns:
         per_sample_loss (torch.Tensor): Mean rollout loss per sample [B].
-        total_loss (torch.Tensor): Scalar mean loss over all valid steps.
+        total_loss (torch.Tensor): Scalar mean loss over all valid steps
+            (per-band weighted when ``band_weights`` is given).
     """
     B = ctx_v.shape[0]
     seed_width = ctx_v.shape[1]
@@ -824,6 +873,7 @@ def _rollout_pass_9band_window(
 
     step_losses = []  # [B] per step
     step_valid = []  # [B] per step
+    step_bands = []  # [B] observed band index per step (for band weighting)
 
     for step in range(n_steps):
         # Most recent real event time per row (left-packed => index count - 1).
@@ -883,6 +933,7 @@ def _rollout_pass_9band_window(
         step_loss = loss_fn(pred_obs, true_obs) * valid
         step_losses.append(step_loss)
         step_valid.append(valid)
+        step_bands.append(tgt_band)
 
         # Scheduled sampling: choose true vs own (detached) prediction per sample.
         use_true = torch.rand(B, device=device) < teacher_forcing_ratio
@@ -906,7 +957,18 @@ def _rollout_pass_9band_window(
     step_valid = torch.stack(step_valid, dim=1)  # [B, n_steps]
 
     per_sample_loss = step_losses.sum(dim=1) / (step_valid.sum(dim=1) + 1e-8)
-    total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
+
+    if band_weights is None:
+        total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
+    else:
+        # Per-band-weighted objective: scale each valid step by its observed
+        # band's weight. per_sample_loss (recorded) stays unweighted above.
+        step_bands = torch.stack(step_bands, dim=1)  # [B, n_steps]
+        bw = band_weights.to(device)
+        step_w = bw[step_bands] * step_valid  # [B, n_steps]
+        total_loss = (step_losses * bw[step_bands]).sum() / (
+            step_w.sum() + 1e-8
+        )
 
     return per_sample_loss, total_loss
 
@@ -932,6 +994,7 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
     window_mode: bool = False,
     context_window_days: float = None,
     max_context_len: int = None,
+    band_weights: torch.Tensor = None,
 ) -> None:
     """Multi-step rollout DDP epoch for the masked 9-band scalar temporal model.
 
@@ -964,9 +1027,15 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
             ``window_mode`` is True.
         max_context_len (int): Padded context width M. Required when
             ``window_mode`` is True.
+        band_weights (torch.Tensor): Optional per-band gradient weights
+            [n_bands], forwarded to the rollout pass to up-weight
+            large-dynamic-range bands (u, g) in the backward objective. The
+            recorded per-sample loss stays unweighted. ``None`` (default)
+            reproduces the plain equal-weight behavior. Validation always runs
+            unweighted so the recorded metric is comparable.
     """
     def _run_pass(
-        data: tuple[torch.Tensor, ...], ratio: float
+        data: tuple[torch.Tensor, ...], ratio: float, weights: torch.Tensor = None
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Unpack a batch (7- or 8-tuple), move to device, run the rollout."""
         if window_mode:
@@ -1017,6 +1086,7 @@ def _run_pass(
                 max_context_len=max_context_len,
                 teacher_forcing_ratio=ratio,
                 device=device,
+                band_weights=weights,
             )
 
         return _rollout_pass_9band(
@@ -1032,6 +1102,7 @@ def _run_pass(
             n_bands=n_bands,
             teacher_forcing_ratio=ratio,
             device=device,
+            band_weights=weights,
         )
 
     train_rcrd_filename = train_rcrd_filename.replace(
@@ -1052,7 +1123,7 @@ def _run_pass(
             optimizer.zero_grad(set_to_none=True)
 
             per_sample_loss, batch_loss = _run_pass(
-                data, teacher_forcing_ratio
+                data, teacher_forcing_ratio, weights=band_weights
             )
 
             batch_loss.backward()

From cceef2bb8ebea43c9688ed2422019e921a50bd21 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 20 Aug 2026 11:43:22 -0600
Subject: [PATCH 44/66] plotting fix

---
 .../KN_loderunner/eval_dense_latetime_9band.py | 17 +++++++++++++----
 .../plot_pred_diagnostics_9band.py             | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index 320b963f..cd66e514 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -372,8 +372,14 @@ def get_args():
     p.add_argument(
         "--dense_glob",
         type=str,
-        required=True,
-        help="Glob for the dense light-curve files (late-time truth).",
+        default=(
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+            "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
+        ),
+        help="Glob for the dense light-curve files (late-time truth). Defaults to "
+        "the same dense set the model was trained on "
+        "(rubin_ztf_dense_10000_dataset_same_seed), whose Rubin bands reach "
+        "~11-12 d median so the 2->10 d scored region is well covered.",
     )
     p.add_argument(
         "--test_filelist",
@@ -391,10 +397,13 @@ def get_args():
     p.add_argument(
         "--late_time_cutoff_days",
         type=float,
-        default=3.0,
+        default=2.0,
         help="Splits context from forecast. The model sees realistic detections "
         "with phase (from first realistic detection) up to this value, and "
-        "forecasts all dense points after it -- the late-time region scored here.",
+        "forecasts all dense points after it -- the late-time region scored here. "
+        "Defaults to 2.0 to match the model's trained context_window_days (a "
+        "2-day trailing lookback), so the eval feeds the model the same context "
+        "span it saw in training rather than a wider ->3-day slice.",
     )
     p.add_argument(
         "--late_time_max_days",
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 3a558628..56bcff79 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -119,6 +119,16 @@ def get_args():
     )
 
     parser.add_argument("--outdir", type=str, default=None)
+    parser.add_argument(
+        "--fixed_forecast_max_days",
+        type=float,
+        default=10.0,
+        help="Cap (days past the last context event) on the smooth fixed-context "
+        "forecast sweep. The sweep would otherwise extend to the last true event "
+        "of each curve, which can run well past the region we care about and into "
+        "the unsupervised tail. Set to match the eval's --late_time_max_days so "
+        "both scripts show the same forecast horizon.",
+    )
     parser.add_argument(
         "--norm_stats_path",
         type=str,
@@ -411,6 +421,7 @@ def get_rollout_from_stream(
     window_mode=False,
     context_window_days=None,
     max_context_len=None,
+    fixed_forecast_max_days=None,
 ):
     """Autoregressively forecast the next events of one merged event stream.
 
@@ -593,6 +604,12 @@ def get_rollout_from_stream(
         max_step_t_rel = max(s["t_rel"] for s in steps)
         horizon = max(1e-3, max_step_t_rel - last_ctx_t_rel)
 
+        # Cap the smooth sweep at the horizon we care about (measured as lead
+        # time past the last context event), so it doesn't run into the
+        # unsupervised late tail. Matches the eval's --late_time_max_days.
+        if fixed_forecast_max_days is not None:
+            horizon = min(horizon, float(fixed_forecast_max_days))
+
         n_lead = 60
         lead_times = np.linspace(0.0, horizon, n_lead).astype(np.float32)
 
@@ -977,6 +994,7 @@ def main():
             window_mode=window_mode,
             context_window_days=context_window_days,
             max_context_len=max_context_len,
+            fixed_forecast_max_days=args.fixed_forecast_max_days,
         )
         rollouts.append(rollout)
 

From 697d40de413d4f8cdf4efce7c525624f17915366 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 20 Aug 2026 15:00:46 -0600
Subject: [PATCH 45/66] compare to uniform curves

---
 .../eval_dense_latetime_9band.py              | 60 ++++++++++++++++++-
 .../KN_loderunner/make_kn_object_lists.py     |  2 +-
 .../KN_loderunner/train_LodeRunner_ddp.py     |  2 +-
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index cd66e514..2df70bea 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -190,6 +190,7 @@ def eval_object(
     max_context_len,
     late_time_cutoff_days,
     late_time_max_days,
+    uniform_stream: tuple | None = None,
 ):
     """Score one object's late-time dense truth against a realistic-context forecast.
 
@@ -295,6 +296,17 @@ def eval_object(
     curve = _batched_forward(model, x, lead_grid, device)  # [60, N_BANDS]
     curve_mag = curve * (stds[None, :] + EPS) + means[None, :]
 
+    # Optional uniform-grid "true curve" for plotting, phase-aligned to the same
+    # t0 (first realistic detection) so it overlays in the same frame. Never
+    # scored or fed to the model -- it is the noise-free, no-limiting-mag target
+    # curve, shown so the forecast is readable even where the survey-limited
+    # dense truth goes dark.
+    uniform = None
+    if uniform_stream is not None:
+        u_t, u_v, u_b = uniform_stream
+        if u_t.shape[0] > 0:
+            uniform = (u_t - t0, u_v, u_b)
+
     return {
         "scored": scored,
         "t0": t0,
@@ -305,6 +317,7 @@ def eval_object(
         # plot those as the context (not the full realistic stream).
         "real": (r_t_ctx - t0, r_v_ctx, r_b_ctx),
         "dense": (d_t - t0, d_v, d_b),
+        "uniform": uniform,
         # Right edge for plotting: the scored horizon. Beyond this the forecast
         # is unsupervised extrapolation, so it is not shown.
         "plot_max_phase": late_time_max_days,
@@ -317,6 +330,7 @@ def plot_object(result, stem, outpath):
     axes = axes.ravel()
     r_ph, r_v, r_b = result["real"]
     d_ph, d_v, d_b = result["dense"]
+    uniform = result.get("uniform")
     # Show only the scored horizon; the forecast beyond it is unsupervised
     # extrapolation (where the late-time upturn artifact lives).
     plot_max = result.get("plot_max_phase")
@@ -330,6 +344,20 @@ def plot_object(result, stem, outpath):
             dm = dm & (d_ph <= plot_max)
         if np.any(dm):
             ax.scatter(d_ph[dm], d_v[dm], s=14, c="0.6", label="dense truth")
+        # Uniform-grid true curve: continuous line so forecast-vs-truth reads as
+        # curve-vs-curve, including where the survey-limited dense truth has no
+        # points. Clip to the plotted horizon and sort by phase for a clean line.
+        if uniform is not None:
+            u_ph, u_v, u_b = uniform
+            um = u_b == b
+            if plot_max is not None:
+                um = um & (u_ph <= plot_max)
+            if np.any(um):
+                order = np.argsort(u_ph[um])
+                ax.plot(
+                    u_ph[um][order], u_v[um][order],
+                    c="0.4", lw=1.0, ls="--", alpha=0.9, label="uniform truth",
+                )
         if np.any(rm):
             ax.scatter(
                 r_ph[rm], r_v[rm], s=26, c=BAND_COLORS[b],
@@ -374,13 +402,29 @@ def get_args():
         type=str,
         default=(
             "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-            "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
+            "rubin_ztf_moredense_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Glob for the dense light-curve files (late-time truth). Defaults to "
         "the same dense set the model was trained on "
-        "(rubin_ztf_dense_10000_dataset_same_seed), whose Rubin bands reach "
+        "(rubin_ztf_moredense_10000_dataset_same_seed), whose Rubin bands reach "
         "~11-12 d median so the 2->10 d scored region is well covered.",
     )
+    p.add_argument(
+        "--uniform_glob",
+        type=str,
+        default=(
+            "/net/sescratch1/atoivonen/data/KN_lightcurves/"
+            "rubin_ztf_uniform_10000_dataset_same_seed/lc_*.npz"
+        ),
+        help="Optional glob for a UNIFORM-grid, noise-free, no-limiting-mag "
+        "companion set (same objects/seed, sampled on a dense regular phase "
+        "grid). When it matches files, each per-object plot overlays this as a "
+        "continuous 'true curve' line so the forecast can be read curve-vs-curve "
+        "even where the survey-limited dense truth has no detections (e.g. deep "
+        "late-time u/g). Plotting only -- never scored, never fed to the model. "
+        "Silently skipped if no files match, so it auto-activates once the set "
+        "is generated.",
+    )
     p.add_argument(
         "--test_filelist",
         type=str,
@@ -474,6 +518,12 @@ def main():
     dense_map = _stem_to_path(args.dense_glob)
     stems = sorted(set(real_map) & set(dense_map))
 
+    # Optional uniform-grid plotting companion (same stems). Absent files are
+    # fine: the overlay just doesn't appear.
+    uniform_map = _stem_to_path(args.uniform_glob) if args.uniform_glob else {}
+    if uniform_map:
+        print(f"Uniform-grid overlay files: {len(uniform_map)}")
+
     if args.test_filelist is not None:
         with open(args.test_filelist) as fh:
             test_stems = {line.strip() for line in fh if line.strip()}
@@ -493,6 +543,11 @@ def main():
     for stem in stems:
         real_stream = read_merged_stream(real_map[stem], DROP_UPPER_LIMITS)
         dense_stream = read_merged_stream(dense_map[stem], drop_upper_limits=False)
+        uniform_stream = None
+        if stem in uniform_map:
+            uniform_stream = read_merged_stream(
+                uniform_map[stem], drop_upper_limits=False
+            )
         result = eval_object(
             real_stream,
             dense_stream,
@@ -504,6 +559,7 @@ def main():
             max_context_len,
             args.late_time_cutoff_days,
             args.late_time_max_days,
+            uniform_stream=uniform_stream,
         )
         if result is None:
             continue
diff --git a/applications/harnesses/KN_loderunner/make_kn_object_lists.py b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
index 42b1a4f0..7eb73f38 100644
--- a/applications/harnesses/KN_loderunner/make_kn_object_lists.py
+++ b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
@@ -125,7 +125,7 @@ def main() -> None:
         type=str,
         default=(
             "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-            "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
+            "rubin_ztf_moredense_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Optional glob for the dense files (for coverage reporting only).",
     )
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 6f6dfb65..0c20bfb2 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -126,7 +126,7 @@ def _read_stem_list(path: str) -> set:
     type=str,
     default=(
         "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-        "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
+        "rubin_ztf_moredense_10000_dataset_same_seed/lc_*.npz"
     ),
     help="Optional glob for the dense light-curve files. When set (and it "
     "matches files), the dense TRAIN objects are concatenated onto the "

From 94790d540b4b788d4d087049890f63611bb24ab7 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Mon, 24 Aug 2026 12:09:06 -0600
Subject: [PATCH 46/66] use original dense training set + config adjustments

---
 .../harnesses/KN_loderunner/eval_dense_latetime_9band.py    | 4 ++--
 .../harnesses/KN_loderunner/make_kn_object_lists.py         | 2 +-
 .../harnesses/KN_loderunner/train_LodeRunner_ddp.py         | 2 +-
 applications/harnesses/KN_loderunner/training_START.input   | 6 +++---
 applications/harnesses/KN_loderunner/training_input.tmpl    | 6 +++---
 5 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index 2df70bea..f61d0405 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -402,11 +402,11 @@ def get_args():
         type=str,
         default=(
             "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-            "rubin_ztf_moredense_10000_dataset_same_seed/lc_*.npz"
+            "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Glob for the dense light-curve files (late-time truth). Defaults to "
         "the same dense set the model was trained on "
-        "(rubin_ztf_moredense_10000_dataset_same_seed), whose Rubin bands reach "
+        "(rubin_ztf_dense_10000_dataset_same_seed), whose Rubin bands reach "
         "~11-12 d median so the 2->10 d scored region is well covered.",
     )
     p.add_argument(
diff --git a/applications/harnesses/KN_loderunner/make_kn_object_lists.py b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
index 7eb73f38..42b1a4f0 100644
--- a/applications/harnesses/KN_loderunner/make_kn_object_lists.py
+++ b/applications/harnesses/KN_loderunner/make_kn_object_lists.py
@@ -125,7 +125,7 @@ def main() -> None:
         type=str,
         default=(
             "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-            "rubin_ztf_moredense_10000_dataset_same_seed/lc_*.npz"
+            "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
         ),
         help="Optional glob for the dense files (for coverage reporting only).",
     )
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 0c20bfb2..6f6dfb65 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -126,7 +126,7 @@ def _read_stem_list(path: str) -> set:
     type=str,
     default=(
         "/net/sescratch1/atoivonen/data/KN_lightcurves/"
-        "rubin_ztf_moredense_10000_dataset_same_seed/lc_*.npz"
+        "rubin_ztf_dense_10000_dataset_same_seed/lc_*.npz"
     ),
     help="Optional glob for the dense light-curve files. When set (and it "
     "matches files), the dense TRAIN objects are concatenated onto the "
diff --git a/applications/harnesses/KN_loderunner/training_START.input b/applications/harnesses/KN_loderunner/training_START.input
index cde36996..f019984d 100644
--- a/applications/harnesses/KN_loderunner/training_START.input
+++ b/applications/harnesses/KN_loderunner/training_START.input
@@ -28,15 +28,15 @@ kn_rubin_ztf_val.txt
 --noise_scale
 <NOISE_SCALE>
 --n_rollout_steps
-5
+12
 --tf_start
 1.0
 --tf_end
 0.0
 --tf_ramp_start_epoch
-20
+8
 --tf_ramp_epochs
-20
+12
 --trn_rcrd_filename
 ./training_study<studyIDX>_epoch<epochIDX>.csv
 --val_rcrd_filename
diff --git a/applications/harnesses/KN_loderunner/training_input.tmpl b/applications/harnesses/KN_loderunner/training_input.tmpl
index 42e4e5fe..f552b660 100644
--- a/applications/harnesses/KN_loderunner/training_input.tmpl
+++ b/applications/harnesses/KN_loderunner/training_input.tmpl
@@ -28,15 +28,15 @@ kn_rubin_ztf_val.txt
 --noise_scale
 <NOISE_SCALE>
 --n_rollout_steps
-5
+12
 --tf_start
 1.0
 --tf_end
 0.0
 --tf_ramp_start_epoch
-20
+8
 --tf_ramp_epochs
-20
+12
 --trn_rcrd_filename
 ./training_study<studyIDX>_epoch<epochIDX>.csv
 --val_rcrd_filename

From a71bbd6a8aca491463528157923ee8216fdb658e Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Mon, 24 Aug 2026 15:24:13 -0600
Subject: [PATCH 47/66] parallel fix

---
 src/yoke/utils/parallel.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/yoke/utils/parallel.py b/src/yoke/utils/parallel.py
index 30d524df..bc694981 100644
--- a/src/yoke/utils/parallel.py
+++ b/src/yoke/utils/parallel.py
@@ -35,6 +35,14 @@ def setup_distributed() -> tuple[int, int, int, torch.device]:
     master_port = os.environ["MASTER_PORT"]
 
     # ----- 2) Set the current GPU device for this process -----
+    # Map local rank onto the GPUs this process can actually see. When Slurm
+    # binds one GPU per task (cgroup isolation), each rank sees a single device
+    # renumbered to 0, so device_count() == 1 and local_rank must fold to 0.
+    # When every rank sees all node GPUs, device_count() == NGPUS and the modulo
+    # is a no-op. Guards against "invalid device ordinal" under per-task binding.
+    n_visible = torch.cuda.device_count()
+    if n_visible > 0:
+        local_rank = local_rank % n_visible
     torch.cuda.set_device(local_rank)
     device = torch.device(f"cuda:{local_rank}")
 

From 8f3101a261bc0a63ee1b5de03732b95fb2f4fea4 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Mon, 24 Aug 2026 15:31:18 -0600
Subject: [PATCH 48/66] input template fix

---
 applications/harnesses/KN_loderunner/training_START.slurm | 2 +-
 applications/harnesses/KN_loderunner/training_slurm.tmpl  | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_START.slurm b/applications/harnesses/KN_loderunner/training_START.slurm
index 524e4f3f..d8a7d44f 100644
--- a/applications/harnesses/KN_loderunner/training_START.slurm
+++ b/applications/harnesses/KN_loderunner/training_START.slurm
@@ -62,7 +62,7 @@ export date00=`date`
 
 # Start the Code
 # Explicitly set TCP environment for the following...
-srun -vv --cpu-bind=verbose python -u <train_script> @study<studyIDX>_START.input
+srun -vv --cpu-bind=verbose --gpu-bind=none python -u <train_script> @study<studyIDX>_START.input
 
 # Get end time and print to stdout
 export date01=`date`
diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index b354e89c..fa60773f 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -68,8 +68,11 @@ export OMP_NUM_THREADS=8
 export date00=`date`
 
 # Start the Code
-# Explicitly set TCP environment for the following...
-srun -vv --cpu-bind=verbose python -u <train_script> @<INPUTFILE>
+# --gpu-bind=none exposes ALL node GPUs to every task so setup_distributed()'s
+# local_rank (0..NGPUS-1) maps 1:1 onto distinct devices. Without it Slurm's
+# default binding hands each task only a subset, so ranks collide on the same
+# CUDA device ("Duplicate GPU detected" in NCCL).
+srun -vv --cpu-bind=verbose --gpu-bind=none python -u <train_script> @<INPUTFILE>
 
 # Get end time and print to stdout
 export date01=`date`

From 3af6627a45caf281094578484e7d8ecb02cf28f6 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 25 Aug 2026 08:24:12 -0600
Subject: [PATCH 49/66] dubugging

---
 src/yoke/utils/parallel.py | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/yoke/utils/parallel.py b/src/yoke/utils/parallel.py
index bc694981..973a9174 100644
--- a/src/yoke/utils/parallel.py
+++ b/src/yoke/utils/parallel.py
@@ -35,14 +35,25 @@ def setup_distributed() -> tuple[int, int, int, torch.device]:
     master_port = os.environ["MASTER_PORT"]
 
     # ----- 2) Set the current GPU device for this process -----
-    # Map local rank onto the GPUs this process can actually see. When Slurm
-    # binds one GPU per task (cgroup isolation), each rank sees a single device
-    # renumbered to 0, so device_count() == 1 and local_rank must fold to 0.
-    # When every rank sees all node GPUs, device_count() == NGPUS and the modulo
-    # is a no-op. Guards against "invalid device ordinal" under per-task binding.
+    # DDP requires each rank to own a DISTINCT physical GPU. That must be
+    # arranged by the launcher's GPU visibility, not by index arithmetic here:
+    # folding local_rank onto a partial visible set (e.g. local_rank % count)
+    # silently maps two ranks to the same device ("Duplicate GPU detected").
+    # Log what this rank actually sees so a bad Slurm binding is obvious.
     n_visible = torch.cuda.device_count()
-    if n_visible > 0:
-        local_rank = local_rank % n_visible
+    print(
+        f"[setup_distributed] rank={rank} local_rank={local_rank} "
+        f"world_size={world_size} visible_gpus={n_visible} "
+        f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '<unset>')}",
+        flush=True,
+    )
+    if local_rank >= n_visible:
+        raise RuntimeError(
+            f"local_rank {local_rank} >= visible GPUs {n_visible}. Each task sees "
+            "only a subset of the node's GPUs -- the Slurm launcher is binding "
+            "GPUs per task. Launch with `srun --gpu-bind=none` so every rank sees "
+            "all node GPUs and local_rank maps 1:1 to a distinct device."
+        )
     torch.cuda.set_device(local_rank)
     device = torch.device(f"cuda:{local_rank}")
 

From 77cfff9176839463ac639d2195a3632fa90b74c6 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Tue, 25 Aug 2026 11:28:30 -0600
Subject: [PATCH 50/66] fixing parallel submission

---
 .../KN_loderunner/training_START.slurm        |  2 +-
 .../KN_loderunner/training_slurm.tmpl         | 10 ++---
 src/yoke/utils/parallel.py                    | 42 ++++++++++++-------
 3 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_START.slurm b/applications/harnesses/KN_loderunner/training_START.slurm
index d8a7d44f..524e4f3f 100644
--- a/applications/harnesses/KN_loderunner/training_START.slurm
+++ b/applications/harnesses/KN_loderunner/training_START.slurm
@@ -62,7 +62,7 @@ export date00=`date`
 
 # Start the Code
 # Explicitly set TCP environment for the following...
-srun -vv --cpu-bind=verbose --gpu-bind=none python -u <train_script> @study<studyIDX>_START.input
+srun -vv --cpu-bind=verbose python -u <train_script> @study<studyIDX>_START.input
 
 # Get end time and print to stdout
 export date01=`date`
diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index fa60773f..149df61f 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -68,11 +68,11 @@ export OMP_NUM_THREADS=8
 export date00=`date`
 
 # Start the Code
-# --gpu-bind=none exposes ALL node GPUs to every task so setup_distributed()'s
-# local_rank (0..NGPUS-1) maps 1:1 onto distinct devices. Without it Slurm's
-# default binding hands each task only a subset, so ranks collide on the same
-# CUDA device ("Duplicate GPU detected" in NCCL).
-srun -vv --cpu-bind=verbose --gpu-bind=none python -u <train_script> @<INPUTFILE>
+# Default Slurm per-task GPU binding is correct here: each task sees exactly one
+# GPU (as cuda:0), and setup_distributed() selects that device. Do NOT add
+# --gpu-bind=none -- it exposes partial GPU subsets per task and causes ranks to
+# collide on the same device ("Duplicate GPU detected" in NCCL).
+srun -vv --cpu-bind=verbose python -u <train_script> @<INPUTFILE>
 
 # Get end time and print to stdout
 export date01=`date`
diff --git a/src/yoke/utils/parallel.py b/src/yoke/utils/parallel.py
index 973a9174..ed3217c1 100644
--- a/src/yoke/utils/parallel.py
+++ b/src/yoke/utils/parallel.py
@@ -22,7 +22,11 @@ def setup_distributed() -> tuple[int, int, int, torch.device]:
     Returns:
         rank (int): Global rank of this process.
         world_size (int): Total number of processes.
-        local_rank (int): Local rank (GPU index) on this node.
+        gpu_index (int): The CUDA device index this process was assigned. This is
+            the actual device ordinal to use for DDP ``device_ids``/
+            ``output_device`` -- 0 under per-task GPU binding (one visible GPU),
+            or the local rank when all node GPUs are visible. NOT necessarily
+            equal to SLURM_LOCALID.
         device (torch.device): CUDA device for this process.
     """
     # ----- 1) Basic setup & environment variables -----
@@ -35,27 +39,33 @@ def setup_distributed() -> tuple[int, int, int, torch.device]:
     master_port = os.environ["MASTER_PORT"]
 
     # ----- 2) Set the current GPU device for this process -----
-    # DDP requires each rank to own a DISTINCT physical GPU. That must be
-    # arranged by the launcher's GPU visibility, not by index arithmetic here:
-    # folding local_rank onto a partial visible set (e.g. local_rank % count)
-    # silently maps two ranks to the same device ("Duplicate GPU detected").
-    # Log what this rank actually sees so a bad Slurm binding is obvious.
+    # Two Slurm GPU-visibility models must both work, so pick the device index
+    # from what THIS task actually sees rather than assuming local_rank is it:
+    #   * Per-task binding (default cgroup isolation): each task sees exactly one
+    #     GPU, renumbered to cuda:0. device_count() == 1, so the device index is
+    #     0 for every rank (local_rank 1,2,3 would be invalid ordinals here).
+    #   * All-visible (e.g. srun --gpu-bind=none): each task sees all node GPUs,
+    #     device_count() == NGPUS, and local_rank is the correct distinct index.
+    # Choosing 0 when only one GPU is visible, else local_rank, gives each rank a
+    # DISTINCT physical GPU under both models and avoids "invalid device ordinal"
+    # and "Duplicate GPU detected".
     n_visible = torch.cuda.device_count()
+    gpu_index = 0 if n_visible <= 1 else local_rank
     print(
         f"[setup_distributed] rank={rank} local_rank={local_rank} "
-        f"world_size={world_size} visible_gpus={n_visible} "
+        f"world_size={world_size} visible_gpus={n_visible} gpu_index={gpu_index} "
         f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '<unset>')}",
         flush=True,
     )
-    if local_rank >= n_visible:
+    if gpu_index >= n_visible:
         raise RuntimeError(
-            f"local_rank {local_rank} >= visible GPUs {n_visible}. Each task sees "
-            "only a subset of the node's GPUs -- the Slurm launcher is binding "
-            "GPUs per task. Launch with `srun --gpu-bind=none` so every rank sees "
-            "all node GPUs and local_rank maps 1:1 to a distinct device."
+            f"Chosen GPU index {gpu_index} >= visible GPUs {n_visible} for "
+            f"local_rank {local_rank}. Each task sees only a partial subset of the "
+            "node's GPUs (neither clean per-task binding nor full visibility). "
+            "Check the Slurm GPU request/binding for this job."
         )
-    torch.cuda.set_device(local_rank)
-    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(gpu_index)
+    device = torch.device(f"cuda:{gpu_index}")
 
     # ----- 3) Initialize the process group -----
     dist.init_process_group(
@@ -65,7 +75,9 @@ def setup_distributed() -> tuple[int, int, int, torch.device]:
         rank=rank,
     )
 
-    return rank, world_size, local_rank, device
+    # Return the ACTUAL device index (gpu_index), not SLURM_LOCALID, so callers
+    # place DDP on the device this process really owns under both binding models.
+    return rank, world_size, gpu_index, device
 
 
 def cleanup_distributed() -> None:

From 5ca6087c25474a695fbd87be81d4966f9841a9c7 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 07:40:00 -0600
Subject: [PATCH 51/66] anchor to last point

---
 .../KN_loderunner/ddp_production.csv          |  11 ++
 .../eval_dense_latetime_9band.py              | 181 ++++++++++++++++--
 .../harnesses/KN_loderunner/infer_9band.py    |   5 +
 .../plot_pred_diagnostics_9band.py            |   5 +
 .../KN_loderunner/train_LodeRunner_ddp.py     |  42 +++-
 .../KN_loderunner/training_START.slurm        |  12 +-
 src/yoke/models/vit/swin/bomberman.py         |  91 +++++++++
 src/yoke/utils/checkpointing.py               |   3 +
 src/yoke/utils/training/epoch/loderunner.py   |  47 +++--
 9 files changed, 360 insertions(+), 37 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
index 8196bfb2..a0460437 100644
--- a/applications/harnesses/KN_loderunner/ddp_production.csv
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -32,5 +32,16 @@ studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZ
 # are 8-GPU H100 nodes (se*); cpus-per-task=24 in the slurm template fits
 # 24*8=192 cores.
 25,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+# study 40: Run A (n_rollout_steps=12). 1 GPU, BATCH_SIZE=5 keeps the 12-step
+# rollout under 80GB (5x12=60 activation units vs the OOM'd 10x12=120). Single
+# GPU avoids the multi-GPU Slurm GPU-binding issue on these H100 nodes.
+40,yoke311,1,1,128,1,1,9,1,2,5,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+# study 41: 12 GPUs = 2 nodes x 6 GPUs/node (8-H100 se* nodes; under the 32-GPU
+# QOS cap). BATCH_SIZE=5 per GPU keeps the 12-step rollout under 80GB. NTRN/NVAL
+# 84/42 per rank hold ~1000/500 samples per epoch across 12 ranks (12*84~=1000).
+# TERMINAL/WARMUP scaled to per-rank batches (84/42) so the LR schedule spans the
+# same epochs. Uses the current train script: PREDICT_DELTA=True + DT_WEIGHT_TAU=3
+# + ZTF band weights bumped to 2 (A+B+ZTF combined, not isolated).
+41,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
 
 
diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index f61d0405..42690830 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -179,6 +179,109 @@ def _batched_forward(
     return out
 
 
+def _rollout_scored(
+    model: torch.nn.Module,
+    device: torch.device,
+    means: np.ndarray,
+    stds: np.ndarray,
+    ctx_t0: list,
+    ctx_v0: list,
+    ctx_b0: list,
+    target_t: np.ndarray,
+    target_v: np.ndarray,
+    target_b: np.ndarray,
+    t0: float,
+    context_window_days: float,
+    max_context_len: int,
+) -> list:
+    """Autoregressive late-time forecast: feed each prediction back as context.
+
+    Steps through the chronologically-ordered late-time dense targets. At each
+    step the trailing-window context is rebuilt from the growing lists, the model
+    predicts all bands at the lead time ``target_t[k] - ctx_t[-1]`` (from the last
+    FED event, not the fixed last realistic detection), the target band's residual
+    is recorded, then the model's own NORMALIZED prediction for that band is
+    appended to the context -- matching the training rollout (``pred_obs.detach()``)
+    and ``get_rollout_from_stream`` (``pred_norm``). Produces the same per-point
+    ``scored`` dicts as the direct path, so plotting/CSV/aggregation are unchanged.
+
+    Args:
+        model: The 9-band scalar-temporal LodeRunner.
+        device (torch.device): Device to run on.
+        means (np.ndarray): Per-band normalization means.
+        stds (np.ndarray): Per-band normalization standard deviations.
+        ctx_t0 (list): Seed context absolute times (pre-cutoff realistic).
+        ctx_v0 (list): Seed context NORMALIZED values (pre-cutoff realistic).
+        ctx_b0 (list): Seed context band indices (pre-cutoff realistic).
+        target_t (np.ndarray): Late-time dense target absolute times, chronological.
+        target_v (np.ndarray): Late-time dense target magnitudes.
+        target_b (np.ndarray): Late-time dense target band indices.
+        t0 (float): Phase-zero time (first realistic detection).
+        context_window_days (float): Trailing lookback W.
+        max_context_len (int): Padded context width M.
+
+    Returns:
+        list: One scored dict per target (phase/lead_time/band/pred_mag/true_mag/
+        residual_mag).
+    """
+    ctx_t = list(ctx_t0)
+    ctx_v = list(ctx_v0)
+    ctx_b = list(ctx_b0)
+
+    scored = []
+    with torch.no_grad():
+        for k in range(target_t.shape[0]):
+            win_v, win_t, win_b = _select_window(
+                ctx_t=ctx_t,
+                ctx_v=ctx_v,
+                ctx_b=ctx_b,
+                context_window_days=context_window_days,
+                max_context_len=max_context_len,
+            )
+            x = build_context_input(
+                win_v=win_v,
+                win_t=win_t,
+                win_b=win_b,
+                context_len=max_context_len,
+                n_bands=N_BANDS,
+                device=device,
+                window_mode=True,
+            )
+            # Lead time from the last FED event (the running context tip).
+            dt = float(target_t[k]) - float(ctx_t[-1])
+            if dt <= 0:
+                # Non-increasing time; skip feeding but still score at a tiny dt.
+                dt = max(dt, 1e-3)
+            Dt = torch.tensor([dt], dtype=torch.float32, device=device)
+            pred_all = (
+                model(x, in_vars=None, out_vars=None, Dt=Dt)
+                .reshape(N_BANDS)
+                .detach()
+                .cpu()
+                .numpy()
+            )
+            band = int(target_b[k])
+            pred_norm = float(pred_all[band])
+            pred_mag = pred_norm * (stds[band] + EPS) + means[band]
+            true_mag = float(target_v[k])
+            scored.append(
+                {
+                    "phase": float(target_t[k]) - t0,
+                    "lead_time": dt,
+                    "band": band,
+                    "pred_mag": float(pred_mag),
+                    "true_mag": true_mag,
+                    "residual_mag": float(pred_mag) - true_mag,
+                }
+            )
+            # Feed the NORMALIZED prediction back as the next context event.
+            ctx_t.append(float(target_t[k]))
+            ctx_v.append(pred_norm)
+            ctx_b.append(band)
+
+    return scored
+
+
 def eval_object(
     real_stream,
     dense_stream,
@@ -191,6 +294,7 @@ def eval_object(
     late_time_cutoff_days,
     late_time_max_days,
     uniform_stream: tuple | None = None,
+    rollout: bool = False,
 ):
     """Score one object's late-time dense truth against a realistic-context forecast.
 
@@ -202,6 +306,19 @@ def eval_object(
     ``late_time_cutoff_days < phase <= late_time_max_days`` (phase measured from
     the first realistic detection). Points beyond ``late_time_max_days`` are
     ignored so the forecast is only judged over a horizon we care about.
+
+    Two forecast modes:
+
+    * DIRECT (``rollout=False``, default): the pre-cutoff realistic context is
+      fixed, and every late-time point is predicted in one batched pass at its
+      true lead time from the last realistic detection. No feedback -- this
+      measures the model's raw Dt-conditioned response.
+    * AUTOREGRESSIVE (``rollout=True``): the context grows -- at each late-time
+      point (chronological order) the model predicts, and its own (normalized)
+      prediction is appended to the context before the next step, exactly as the
+      training rollout and ``get_rollout_from_stream`` do. Each step's ``Dt`` is
+      measured from the last FED event, not the fixed last realistic detection.
+      This measures the true inference path (and exposes drift).
     """
     r_t, r_v, r_b = real_stream
     d_t, d_v, d_b = dense_stream
@@ -270,23 +387,42 @@ def eval_object(
     if late_idx.shape[0] == 0:
         return None
 
-    pred_scored = _batched_forward(model, x, lead_times, device)  # [P, N_BANDS]
-
-    scored = []
-    for j, idx in enumerate(late_idx):
-        band = int(d_b[idx])
-        pred_mag = float(pred_scored[j, band] * (stds[band] + EPS) + means[band])
-        true_mag = float(d_v[idx])
-        scored.append(
-            {
-                "phase": float(d_t[idx]) - t0,
-                "lead_time": float(lead_times[j]),
-                "band": band,
-                "pred_mag": pred_mag,
-                "true_mag": true_mag,
-                "residual_mag": pred_mag - true_mag,
-            }
+    if rollout:
+        # AUTOREGRESSIVE: grow the context, feeding each (normalized) prediction
+        # back before the next step. Mirrors get_rollout_from_stream and the
+        # training rollout. Each step's Dt is from the last FED event's time.
+        scored = _rollout_scored(
+            model=model,
+            device=device,
+            means=means,
+            stds=stds,
+            ctx_t0=list(r_t_ctx.astype(np.float32)),
+            ctx_v0=list(r_v_norm.astype(np.float32)),
+            ctx_b0=list(r_b_ctx),
+            target_t=d_t[late_idx].astype(np.float32),
+            target_v=d_v[late_idx].astype(np.float32),
+            target_b=d_b[late_idx].astype(np.int64),
+            t0=t0,
+            context_window_days=context_window_days,
+            max_context_len=max_context_len,
         )
+    else:
+        pred_scored = _batched_forward(model, x, lead_times, device)  # [P, N_BANDS]
+        scored = []
+        for j, idx in enumerate(late_idx):
+            band = int(d_b[idx])
+            pred_mag = float(pred_scored[j, band] * (stds[band] + EPS) + means[band])
+            true_mag = float(d_v[idx])
+            scored.append(
+                {
+                    "phase": float(d_t[idx]) - t0,
+                    "lead_time": float(lead_times[j]),
+                    "band": band,
+                    "pred_mag": pred_mag,
+                    "true_mag": true_mag,
+                    "residual_mag": pred_mag - true_mag,
+                }
+            )
 
     # Smooth forecast curve for plotting: sweep lead time from 0 to the farthest
     # scored late-time point, predicting all bands at each lead time -- also a
@@ -470,6 +606,14 @@ def get_args():
         default=12,
         help="Number of per-object forecast plots to write.",
     )
+    p.add_argument(
+        "--rollout",
+        action="store_true",
+        help="Score the AUTOREGRESSIVE forecast: feed each prediction back as "
+        "context before the next late-time point (the true inference path), "
+        "instead of the default DIRECT single-pass forecast from a fixed "
+        "pre-cutoff context. Comparing the two isolates rollout drift.",
+    )
     return p.parse_args()
 
 
@@ -560,6 +704,7 @@ def main():
             args.late_time_cutoff_days,
             args.late_time_max_days,
             uniform_stream=uniform_stream,
+            rollout=args.rollout,
         )
         if result is None:
             continue
@@ -581,7 +726,9 @@ def main():
     # Per-band late-time error summary.
     resid = np.asarray([s["residual_mag"] for s in all_scored])
     bands = np.asarray([s["band"] for s in all_scored])
-    print(f"\nEvaluated {n_eval} objects; {len(all_scored)} late-time points "
+    mode = "AUTOREGRESSIVE rollout" if args.rollout else "DIRECT single-pass"
+    print(f"\nForecast mode: {mode}")
+    print(f"Evaluated {n_eval} objects; {len(all_scored)} late-time points "
           f"(cutoff {args.late_time_cutoff_days} d).")
     print(f"Overall late-time RMSE (mag): {np.sqrt(np.mean(resid**2)):.4f}  "
           f"MAE: {np.mean(np.abs(resid)):.4f}")
diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 3bd88116..2171b9bc 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -179,6 +179,9 @@ def load_9band_model(ckpt_path, device):
     # 0 for legacy checkpoints (no key) -> Fourier Dt disabled -> matches saved
     # weights so strict load succeeds.
     dt_fourier_bands = ckpt.get("dt_fourier_bands", 0)
+    # False for legacy checkpoints (no key) -> absolute head. Adds no params, so
+    # it only changes forward() behavior, never the state_dict shape.
+    predict_delta = ckpt.get("predict_delta", False)
 
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
@@ -188,6 +191,7 @@ def load_9band_model(ckpt_path, device):
     print("max_context_len:", max_context_len)
     print("n_bands:", n_bands)
     print("dt_fourier_bands:", dt_fourier_bands)
+    print("predict_delta:", predict_delta)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -201,6 +205,7 @@ def load_9band_model(ckpt_path, device):
         hidden=hidden,
         context_window_days=context_window_days,
         dt_fourier_bands=dt_fourier_bands,
+        predict_delta=predict_delta,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 56bcff79..0d7acf51 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -207,6 +207,9 @@ def load_9band_model(ckpt_path, device):
     # 0 for legacy checkpoints (no key) -> Fourier Dt disabled -> matches saved
     # weights so strict load succeeds.
     dt_fourier_bands = ckpt.get("dt_fourier_bands", 0)
+    # False for legacy checkpoints (no key) -> absolute head. Adds no params, so
+    # it only changes forward() behavior, never the state_dict shape.
+    predict_delta = ckpt.get("predict_delta", False)
 
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
@@ -220,6 +223,7 @@ def load_9band_model(ckpt_path, device):
     print("backbone_channels:", backbone_channels)
     print("hidden:", hidden)
     print("dt_fourier_bands:", dt_fourier_bands)
+    print("predict_delta:", predict_delta)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -233,6 +237,7 @@ def load_9band_model(ckpt_path, device):
         hidden=hidden,
         context_window_days=context_window_days,
         dt_fourier_bands=dt_fourier_bands,
+        predict_delta=predict_delta,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 6f6dfb65..14799211 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -235,6 +235,17 @@ def main(args, rank, world_size, local_rank, device):
     # to 0 for the legacy architecture (byte-identical; old checkpoints load).
     DT_FOURIER_BANDS = 8
 
+    # Delta-anchored head. When True, the output head predicts a CHANGE relative
+    # to the per-band last observed magnitude (fallback: most-recent observation
+    # in any band) instead of an absolute magnitude, so the forecast starts AT the
+    # last observation at Dt=0 rather than reconstructing the zero-point from
+    # scratch. This kills the ~0.5-0.8 mag lead-0 offset seen in the dense eval.
+    # Adds NO parameters (the anchor is derived from x), so old checkpoints still
+    # load strict=True; the flag is saved and restored by the loaders. Requires
+    # window mode (needs the per-event validity flag). Set False for the absolute
+    # head (byte-identical numerics to the pre-delta model).
+    PREDICT_DELTA = True
+
     # Time-window context mode. When CONTEXT_WINDOW_DAYS is not None, the dataset
     # selects context by a trailing lookback in days (all detections within the
     # last CONTEXT_WINDOW_DAYS), padded to MAX_CONTEXT_LEN with a per-event
@@ -282,17 +293,20 @@ def main(args, rank, world_size, local_rank, device):
     # Per-band loss weighting. Targets are per-band z-scored, so an equal-weight
     # loss lets the large-dynamic-range blue bands (u, g fade to mag ~28-30) be
     # under-fit -- the dense-eval showed a strong under-fade bias there (u
-    # ~ -5 mag, g ~ -2 mag) while ZTF/red bands fit well. Up-weighting u and g
-    # in the TRAINING backward (the recorded per-sample loss stays unweighted so
-    # the val CSV remains a comparable yardstick) pushes gradient toward the
-    # blue fade the model currently ignores. Order matches BAND_KEYS =
+    # ~ -5 mag, g ~ -2 mag). Up-weighting these bands in the TRAINING backward
+    # (the recorded per-sample loss stays unweighted so the val CSV remains a
+    # comparable yardstick) pushes gradient toward the fades the model currently
+    # ignores. The three ZTF bands are also up-weighted (1->2): in autoregressive
+    # rollout they lag because ZTF is realistically sampled near peak/early with
+    # little late-time context to re-anchor, so they get the least gradient
+    # pressure exactly where they fail. Order matches BAND_KEYS =
     # (ztfg, ztfr, ztfi, sdssu, ps1_g, ps1_r, ps1_i, ps1_z, ps1_y). Set to None
     # to recover the exact equal-weight objective.
     BAND_WEIGHTS = torch.tensor(
         [
-            1.0,  # ztfg
-            1.0,  # ztfr
-            1.0,  # ztfi
+            2.0,  # ztfg -- ZTF bands lag in rollout; up-weight from 1->2
+            2.0,  # ztfr -- ZTF bands lag in rollout; up-weight from 1->2
+            2.0,  # ztfi -- ZTF bands lag in rollout; up-weight from 1->2
             3.0,  # sdssu (u) -- worst under-fade, largest up-weight
             2.0,  # ps1__g (g)
             1.0,  # ps1__r (r)
@@ -303,6 +317,16 @@ def main(args, rank, world_size, local_rank, device):
         dtype=torch.float32,
     )
 
+    # Lead-time loss weighting (window-mode rollout only). The dense truth has far
+    # more points in the slowly-fading late tail than in the 2-5 d rise, so an
+    # equal-per-step objective rewards nailing the flat tail and the forecast
+    # collapses to a plateau. Weighting each rollout step by 1 / (1 + Dt / tau)
+    # gives short-lead steps (the rise, where the model fails) relatively more
+    # gradient. Composes multiplicatively with BAND_WEIGHTS. The recorded
+    # per-sample loss stays unweighted so the val CSV remains comparable. Set to
+    # None to recover the equal-per-step objective. tau ~ a few days.
+    DT_WEIGHT_TAU = 3.0
+
     optimizer_kwargs = {
         "lr": 1e-4,# 1e-4, #1e-5
         "betas": (0.9, 0.999),
@@ -384,6 +408,7 @@ def main(args, rank, world_size, local_rank, device):
             hidden=HIDDEN_CHANNELS,
             context_window_days=CONTEXT_WINDOW_DAYS,
             dt_fourier_bands=DT_FOURIER_BANDS,
+            predict_delta=PREDICT_DELTA,
         ).to(device)
 
         # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
@@ -673,6 +698,7 @@ def _make_9band(
                 context_window_days=CONTEXT_WINDOW_DAYS,
                 max_context_len=MAX_CONTEXT_LEN,
                 band_weights=BAND_WEIGHTS,
+                dt_weight_tau=DT_WEIGHT_TAU,
             )
         else:
             #train_DDP_loderunner_epoch(
@@ -739,6 +765,8 @@ def _make_9band(
                     "backbone_channels": 8,
                     "hidden": HIDDEN_CHANNELS,
                     "dt_fourier_bands": DT_FOURIER_BANDS,
+                    "predict_delta": PREDICT_DELTA,
+                    "dt_weight_tau": DT_WEIGHT_TAU,
                     "band_weights": (
                         BAND_WEIGHTS.tolist()
                         if BAND_WEIGHTS is not None
diff --git a/applications/harnesses/KN_loderunner/training_START.slurm b/applications/harnesses/KN_loderunner/training_START.slurm
index 524e4f3f..3c845428 100644
--- a/applications/harnesses/KN_loderunner/training_START.slurm
+++ b/applications/harnesses/KN_loderunner/training_START.slurm
@@ -17,6 +17,13 @@
 #SBATCH --nodes=<KNODES>
 #SBATCH --ntasks-per-node=<NGPUS>
 #SBATCH --gpus-per-node=<NGPUS>
+# nodes have 8 H100s and 192-208 CPUs, i.e. ~24 cores/GPU. One task per GPU
+# (ntasks-per-node=NGPUS), so cpus-per-task=24 gives each rank enough cores for
+# its dataloader workers + OMP threads (24*NGPUS stays within budget for NGPUS
+# up to 8). Must match the continuation template so epoch 1 behaves like later
+# epochs; without it Slurm may hand each task a single core and starve the
+# dataloaders at multi-GPU scale.
+#SBATCH --cpus-per-task=24
 #SBATCH --mem-per-gpu=50G
 #SBATCH --output=study<studyIDX>_epoch0001.out
 #SBATCH --error=study<studyIDX>_epoch0001.err
@@ -54,8 +61,9 @@ module load anaconda/3.12
 source activate
 conda activate <YOKE_TORCH_ENV>
 
-# Set number of threads per GPU
-export OMP_NUM_THREADS=10
+# Set number of threads per GPU. Must be <= cpus-per-task (24). Leaves cores for
+# the dataloader workers (NUM_WORKERS per rank) alongside the OMP compute threads.
+export OMP_NUM_THREADS=8
 
 # Get start time
 export date00=`date`
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index 54a38185..e991a761 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -472,6 +472,10 @@ class ScalarTemporalConditionedLodeRunner_9band(nn.Module):
         image_size (tuple): Spatial size (H, W) fed to the backbone.
         backbone_channels (int): Number of pseudo-channels the backbone expects.
         hidden (int): Hidden width of the conditioner/output-head MLPs.
+        predict_delta (bool): When True, the output head predicts a CHANGE relative
+            to the last observed magnitude per band (anchored regression) instead of
+            an absolute magnitude. See ``__init__`` for the anchoring rule. Adds no
+            parameters, so it is backward-compatible with existing checkpoints.
     """
 
     def __init__(
@@ -484,6 +488,7 @@ def __init__(
         hidden: int = 64,
         context_window_days: float = None,
         dt_fourier_bands: int = 0,
+        predict_delta: bool = False,
     ) -> None:
         """Initialize conditioner and output-head around the backbone.
 
@@ -511,9 +516,29 @@ def __init__(
                 non-trainable buffer. The monotone ``log1p(Dt)`` channel gives an
                 explicit long-term trend so the forecast does not curve back up
                 at long lead times (which a purely periodic basis would).
+            predict_delta (bool): When True (default False), the output head predicts
+                a normalized-space DELTA that is added to a per-band anchor -- the
+                most-recent observed value in that band within the context window --
+                so the forecast starts AT the last observation at lead time 0 instead
+                of reconstructing the absolute magnitude from scratch. Bands with no
+                observation in the window fall back to the most-recent observation in
+                ANY band (global-last); if the whole window is empty the anchor is 0
+                (the normalized mean). The anchor is derived from ``x`` inside
+                ``forward`` and requires the ``valid`` column, so this mode is only
+                valid in time-window mode (``context_window_days`` set). It adds NO
+                parameters, so the ``state_dict`` is byte-identical to the absolute
+                model and existing checkpoints load unchanged; the flag is recorded
+                in the checkpoint and restored by the loaders.
         """
         super().__init__()
 
+        if predict_delta and context_window_days is None:
+            raise ValueError(
+                "predict_delta=True requires time-window mode "
+                "(context_window_days set); the per-band anchor needs the "
+                "'valid' column present only in the window layout."
+            )
+
         self.backbone = backbone
         self.context_len = context_len
         self.n_bands = n_bands
@@ -521,6 +546,7 @@ def __init__(
         self.backbone_channels = backbone_channels
         self.context_window_days = context_window_days
         self.dt_fourier_bands = dt_fourier_bands
+        self.predict_delta = predict_delta
 
         # Dataset x layout, flattened per event. Fixed-count mode:
         #   [value, rel_t, one_hot_band(n_bands)] * context_len   -> 2 + n_bands
@@ -598,6 +624,63 @@ def _encode_dt(self, Dt: torch.Tensor, batch_size: int) -> torch.Tensor:
         mono = torch.log1p(Dt.clamp_min(0.0))  # [B, 1], monotone in lead time
         return torch.cat([torch.sin(angles), torch.cos(angles), mono], dim=1)
 
+    def _band_anchor(self, x: torch.Tensor) -> torch.Tensor:
+        """Per-band last observed value from the windowed context, for delta mode.
+
+        Reconstructs, for each band, the most-recent observed (normalized) value in
+        the trailing context window. The window-mode ``x`` flattens per event as
+        ``[value, rel_t, valid, one_hot_band(n_bands)]`` (see the dataset's
+        ``_getitem_window``); "most recent" is the valid event of that band with the
+        largest ``rel_t``. Bands with no observation in the window fall back to the
+        most-recent observation in ANY band (global-last). A fully empty window (no
+        valid events) yields an all-zero anchor (the normalized mean).
+
+        Args:
+            x (torch.Tensor): Window-mode context, shape
+                [B, context_len * (3 + n_bands)].
+
+        Returns:
+            torch.Tensor: Per-band anchor of shape [B, n_bands] in normalized units.
+        """
+        B = x.shape[0]
+        nb = self.n_bands
+        ev = x.view(B, self.context_len, 3 + nb)  # [B, L, 3+nb]
+
+        value = ev[..., 0]  # [B, L]
+        rel_t = ev[..., 1]  # [B, L]
+        valid = ev[..., 2] > 0.5  # [B, L] bool
+        band_oh = ev[..., 3:]  # [B, L, nb]
+        band_idx = band_oh.argmax(dim=-1)  # [B, L]
+
+        neg_inf = torch.finfo(rel_t.dtype).min
+
+        # Global-last: value of the valid event with the largest rel_t (any band).
+        g_score = torch.where(valid, rel_t, torch.full_like(rel_t, neg_inf))
+        g_any = valid.any(dim=1)  # [B]
+        g_arg = g_score.argmax(dim=1)  # [B]
+        global_last = value.gather(1, g_arg.unsqueeze(1)).squeeze(1)  # [B]
+        global_last = torch.where(g_any, global_last, torch.zeros_like(global_last))
+
+        # Per-band last: for each band b, the valid event of band b with max rel_t.
+        # match[b] over events, scored by rel_t; [B, nb, L].
+        band_match = (
+            valid.unsqueeze(1)
+            & (band_idx.unsqueeze(1) == torch.arange(nb, device=x.device).view(1, nb, 1))
+        )  # [B, nb, L]
+        pb_score = torch.where(
+            band_match,
+            rel_t.unsqueeze(1),
+            torch.full_like(rel_t.unsqueeze(1), neg_inf),
+        )  # [B, nb, L]
+        pb_has = band_match.any(dim=2)  # [B, nb]
+        pb_arg = pb_score.argmax(dim=2)  # [B, nb]
+        per_band_last = torch.gather(
+            value.unsqueeze(1).expand(B, nb, self.context_len), 2, pb_arg.unsqueeze(2)
+        ).squeeze(2)  # [B, nb]
+
+        # Fall back to global-last where a band was never observed in the window.
+        return torch.where(pb_has, per_band_last, global_last.unsqueeze(1))
+
     def forward(
         self,
         x: torch.Tensor,
@@ -662,6 +745,14 @@ def forward(
             torch.cat([pred_channel_vals, dt_feat], dim=1)
         )  # [B, n_bands]
 
+        # Delta mode: the head predicts a change relative to the per-band last
+        # observed value (anchored regression), so the forecast starts at the last
+        # observation at Dt=0 instead of reconstructing the absolute magnitude. The
+        # anchor is derived from x (no parameters), so the disabled path is
+        # byte-identical to the absolute model.
+        if self.predict_delta:
+            pred = pred + self._band_anchor(x)
+
         return pred
 
 
diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index 44297c9f..dda20f05 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -499,6 +499,9 @@ def load_direct_loderunner_checkpoint_9band(
         # 0 for legacy checkpoints (no key) -> Fourier Dt disabled -> the
         # architecture matches the saved weights and strict load succeeds.
         dt_fourier_bands=checkpoint_data.get("dt_fourier_bands", 0),
+        # False for legacy checkpoints (no key) -> absolute head. Adds no params,
+        # so this only changes forward() behavior, never the state_dict.
+        predict_delta=checkpoint_data.get("predict_delta", False),
     ).to(device)
 
     state_dict = checkpoint_data["model_state_dict"]
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index cc271f10..7d634231 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -794,6 +794,7 @@ def _rollout_pass_9band_window(
     teacher_forcing_ratio: float,
     device: torch.device,
     band_weights: torch.Tensor = None,
+    dt_weight_tau: float = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Unroll the 9-band model over a batch of rollouts in time-window mode.
 
@@ -835,6 +836,14 @@ def _rollout_pass_9band_window(
             per-band weighted mean over valid steps; ``per_sample_loss`` (the
             recorded metric) stays unweighted. ``None`` reproduces the plain
             equal-weight behavior exactly.
+        dt_weight_tau (float): Optional lead-time weighting time-constant in days.
+            When set, each rollout step is weighted by ``1 / (1 + Dt / tau)`` in the
+            backward objective, so short-lead steps (the early rise, where the model
+            fails) get relatively more gradient than the many slowly-fading
+            late-tail steps (which otherwise dominate a per-point mean and pull the
+            forecast toward a flat plateau). Composes multiplicatively with
+            ``band_weights``. ``per_sample_loss`` (the recorded metric) stays
+            unweighted. ``None`` (default) disables it, reproducing prior behavior.
 
     Returns:
         per_sample_loss (torch.Tensor): Mean rollout loss per sample [B].
@@ -958,17 +967,23 @@ def _rollout_pass_9band_window(
 
     per_sample_loss = step_losses.sum(dim=1) / (step_valid.sum(dim=1) + 1e-8)
 
-    if band_weights is None:
+    if band_weights is None and dt_weight_tau is None:
         total_loss = step_losses.sum() / (step_valid.sum() + 1e-8)
     else:
-        # Per-band-weighted objective: scale each valid step by its observed
-        # band's weight. per_sample_loss (recorded) stays unweighted above.
-        step_bands = torch.stack(step_bands, dim=1)  # [B, n_steps]
-        bw = band_weights.to(device)
-        step_w = bw[step_bands] * step_valid  # [B, n_steps]
-        total_loss = (step_losses * bw[step_bands]).sum() / (
-            step_w.sum() + 1e-8
-        )
+        # Weighted objective: scale each valid step by its observed band's weight
+        # (band_weights) and/or a lead-time weight that down-weights long horizons
+        # (dt_weight_tau). Both compose multiplicatively; per_sample_loss
+        # (recorded) stays unweighted above.
+        step_w = step_valid.clone()  # [B, n_steps]
+        if band_weights is not None:
+            step_bands = torch.stack(step_bands, dim=1)  # [B, n_steps]
+            step_w = step_w * band_weights.to(device)[step_bands]
+        if dt_weight_tau is not None:
+            # 1 / (1 + Dt / tau): 1.0 at Dt=0, decaying with lead time. future_dt
+            # is [B, n_steps]; clamp Dt>=0 so padded/degenerate steps stay sane.
+            dt_w = 1.0 / (1.0 + future_dt.clamp_min(0.0) / float(dt_weight_tau))
+            step_w = step_w * dt_w
+        total_loss = (step_losses * step_w).sum() / (step_w.sum() + 1e-8)
 
     return per_sample_loss, total_loss
 
@@ -995,6 +1010,7 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
     context_window_days: float = None,
     max_context_len: int = None,
     band_weights: torch.Tensor = None,
+    dt_weight_tau: float = None,
 ) -> None:
     """Multi-step rollout DDP epoch for the masked 9-band scalar temporal model.
 
@@ -1033,9 +1049,17 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
             recorded per-sample loss stays unweighted. ``None`` (default)
             reproduces the plain equal-weight behavior. Validation always runs
             unweighted so the recorded metric is comparable.
+        dt_weight_tau (float): Optional lead-time weighting time-constant in days
+            (window mode only), forwarded to the rollout pass to down-weight
+            long-horizon steps so the early rise is not swamped by the many
+            late-tail points. The recorded per-sample loss stays unweighted, and
+            validation always runs unweighted. ``None`` (default) disables it.
     """
     def _run_pass(
-        data: tuple[torch.Tensor, ...], ratio: float, weights: torch.Tensor = None
+        data: tuple[torch.Tensor, ...],
+        ratio: float,
+        weights: torch.Tensor = None,
+        dt_tau: float = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Unpack a batch (7- or 8-tuple), move to device, run the rollout."""
         if window_mode:
@@ -1087,6 +1111,7 @@ def _run_pass(
                 teacher_forcing_ratio=ratio,
                 device=device,
                 band_weights=weights,
+                dt_weight_tau=dt_tau,
             )
 
         return _rollout_pass_9band(
@@ -1123,7 +1148,7 @@ def _run_pass(
             optimizer.zero_grad(set_to_none=True)
 
             per_sample_loss, batch_loss = _run_pass(
-                data, teacher_forcing_ratio, weights=band_weights
+                data, teacher_forcing_ratio, weights=band_weights, dt_tau=dt_weight_tau
             )
 
             batch_loss.backward()

From 2576ca7b1c215354c5e3ed67e8702aeeec6ded57 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 12:17:30 -0600
Subject: [PATCH 52/66] add 5th and 95th loss percentiles

---
 .../KN_loderunner/plot_loss_curves_9band.py   |  26 +++++
 .../KN_loderunner/plot_loss_curves_gri.py     | 102 +++++++++++++++++-
 2 files changed, 126 insertions(+), 2 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
index 0aa3e3b6..11371346 100644
--- a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
@@ -60,6 +60,32 @@ def main():
         action="store_true",
         help="Shade +/- one epoch standard deviation.",
     )
+    parser.add_argument(
+        "--show_percentiles",
+        dest="show_percentiles",
+        action="store_true",
+        default=True,
+        help="Shade the per-epoch percentile band (default 5th-95th) of the "
+        "per-batch loss. Default on for 9-band runs.",
+    )
+    parser.add_argument(
+        "--no_show_percentiles",
+        dest="show_percentiles",
+        action="store_false",
+        help="Disable the per-epoch percentile band shading.",
+    )
+    parser.add_argument(
+        "--pct_lo",
+        type=float,
+        default=5.0,
+        help="Lower percentile for --show_percentiles. Default 5.",
+    )
+    parser.add_argument(
+        "--pct_hi",
+        type=float,
+        default=95.0,
+        help="Upper percentile for --show_percentiles. Default 95.",
+    )
     parser.add_argument(
         "--require_val",
         action="store_true",
diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
index 568ead70..69bf32b4 100644
--- a/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py
@@ -110,6 +110,40 @@ def epoch_stats(epochs, losses):
     return unique_epochs, mean_losses, std_losses
 
 
+def epoch_percentiles(
+    epochs: np.ndarray,
+    losses: np.ndarray,
+    pct_lo: float = 5.0,
+    pct_hi: float = 95.0,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Per-epoch lower/upper percentiles of the per-batch loss.
+
+    Complements ``epoch_stats``: where the std band is symmetric and can dip
+    below zero on a log axis, percentile bands capture the actual spread of the
+    per-batch losses within each epoch (e.g. the 5th and 95th) and are robust to
+    the heavy-tailed batches that make the raw loss curve jagged.
+
+    Args:
+        epochs (np.ndarray): Per-row epoch index, shape [N].
+        losses (np.ndarray): Per-row loss columns, shape [N, n_loss_cols].
+        pct_lo (float): Lower percentile in [0, 100]. Default 5.0.
+        pct_hi (float): Upper percentile in [0, 100]. Default 95.0.
+
+    Returns:
+        unique_epochs (np.ndarray): Sorted unique epoch indices, shape [E].
+        lo_losses (np.ndarray): Lower-percentile loss per epoch, [E, n_loss_cols].
+        hi_losses (np.ndarray): Upper-percentile loss per epoch, [E, n_loss_cols].
+    """
+    unique_epochs = np.array(sorted(set(epochs)))
+    lo_losses = np.vstack(
+        [np.percentile(losses[epochs == e], pct_lo, axis=0) for e in unique_epochs]
+    )
+    hi_losses = np.vstack(
+        [np.percentile(losses[epochs == e], pct_hi, axis=0) for e in unique_epochs]
+    )
+    return unique_epochs, lo_losses, hi_losses
+
+
 def infer_loss_labels(loss_names, n_loss_cols):
     """Make nicer labels for common scalar/GRI cases."""
     if n_loss_cols == 1:
@@ -187,6 +221,17 @@ def plot_epoch_curves(train, val, args):
     n_loss_cols = train["losses"].shape[1]
     loss_labels = infer_loss_labels(train["loss_names"], n_loss_cols)
 
+    # Percentile bands (default 5th/95th) of the per-batch loss within each epoch.
+    show_percentiles = getattr(args, "show_percentiles", False)
+    pct_lo = getattr(args, "pct_lo", 5.0)
+    pct_hi = getattr(args, "pct_hi", 95.0)
+    if show_percentiles:
+        _, train_plo, train_phi = epoch_percentiles(
+            train["epochs"], train["losses"], pct_lo, pct_hi
+        )
+    else:
+        train_plo = train_phi = None
+
     if val is not None:
         val_ep, val_mean, val_std = epoch_stats(
             val["epochs"],
@@ -203,17 +248,27 @@ def plot_epoch_curves(train, val, args):
             val_ep = None
             val_mean = None
             val_std = None
+            val_plo = val_phi = None
+        elif show_percentiles:
+            _, val_plo, val_phi = epoch_percentiles(
+                val["epochs"], val["losses"], pct_lo, pct_hi
+            )
+        else:
+            val_plo = val_phi = None
     else:
         val_ep = None
         val_mean = None
         val_std = None
+        val_plo = val_phi = None
 
     plt.figure(figsize=(9, 5.5))
 
+    band_label = f"{pct_lo:g}-{pct_hi:g} pct"
+
     for idx, label in enumerate(loss_labels):
         suffix = "" if n_loss_cols == 1 else f" {label}"
 
-        plt.plot(
+        (train_line,) = plt.plot(
             train_ep,
             train_mean[:, idx],
             marker="o",
@@ -225,8 +280,20 @@ def plot_epoch_curves(train, val, args):
             hi = train_mean[:, idx] + train_std[:, idx]
             plt.fill_between(train_ep, lo, hi, alpha=0.15)
 
+        if show_percentiles:
+            lo = np.maximum(train_plo[:, idx], 1e-30)
+            hi = train_phi[:, idx]
+            plt.fill_between(
+                train_ep,
+                lo,
+                hi,
+                alpha=0.15,
+                color=train_line.get_color(),
+                label=f"Train {band_label}{suffix}",
+            )
+
         if val is not None:
-            plt.plot(
+            (val_line,) = plt.plot(
                 val_ep,
                 val_mean[:, idx],
                 marker="s",
@@ -239,6 +306,18 @@ def plot_epoch_curves(train, val, args):
                 hi = val_mean[:, idx] + val_std[:, idx]
                 plt.fill_between(val_ep, lo, hi, alpha=0.10)
 
+            if show_percentiles:
+                lo = np.maximum(val_plo[:, idx], 1e-30)
+                hi = val_phi[:, idx]
+                plt.fill_between(
+                    val_ep,
+                    lo,
+                    hi,
+                    alpha=0.10,
+                    color=val_line.get_color(),
+                    label=f"Validation {band_label}{suffix}",
+                )
+
     plt.xlabel("Epoch")
     plt.ylabel("Mean loss")
     plt.title(args.title)
@@ -321,6 +400,25 @@ def main():
         help="Shade +/- one epoch standard deviation.",
     )
 
+    parser.add_argument(
+        "--show_percentiles",
+        action="store_true",
+        help="Shade the per-epoch percentile band (default 5th-95th) of the "
+        "per-batch loss.",
+    )
+    parser.add_argument(
+        "--pct_lo",
+        type=float,
+        default=5.0,
+        help="Lower percentile for --show_percentiles. Default 5.",
+    )
+    parser.add_argument(
+        "--pct_hi",
+        type=float,
+        default=95.0,
+        help="Upper percentile for --show_percentiles. Default 95.",
+    )
+
     parser.add_argument(
         "--require_val",
         action="store_true",

From 6fb71e117d30cd81e6c5fca75a19751e01d1aeca Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 14:29:50 -0600
Subject: [PATCH 53/66] exponential moving average and adjustment to late obs
 pinning

---
 .../eval_dense_latetime_9band.py              |  9 +-
 .../harnesses/KN_loderunner/infer_9band.py    | 43 +++++++++-
 .../plot_pred_diagnostics_9band.py            | 45 +++++++++-
 .../KN_loderunner/train_LodeRunner_ddp.py     | 78 ++++++++++++++++++
 src/yoke/models/vit/swin/bomberman.py         | 82 ++++++++++++++++++-
 src/yoke/utils/checkpointing.py               |  5 ++
 src/yoke/utils/training/epoch/loderunner.py   | 24 ++++++
 7 files changed, 277 insertions(+), 9 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
index 42690830..f2a045bb 100644
--- a/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
+++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py
@@ -524,6 +524,13 @@ def get_args():
     p.add_argument("--study", type=int, default=24)
     p.add_argument("--epoch", type=int, default=500)
     p.add_argument("--ckpt", type=str, default=None)
+    p.add_argument(
+        "--use_ema",
+        action="store_true",
+        help="Overlay the EMA (Polyak) shadow of the trainable params instead "
+        "of the raw weights. Falls back to raw weights if the checkpoint has "
+        "no EMA shadow.",
+    )
     p.add_argument(
         "--realistic_glob",
         type=str,
@@ -637,7 +644,7 @@ def main():
         n_bands,
         context_window_days,
         max_context_len,
-    ) = load_9band_model(args.ckpt, device)
+    ) = load_9band_model(args.ckpt, device, use_ema=getattr(args, "use_ema", False))
 
     if context_window_days is None:
         raise ValueError(
diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 2171b9bc..4bcba32c 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -34,6 +34,7 @@
     NINE_BAND_KEYS,
     load_or_compute_band_normalization,
 )
+from yoke.utils.ema import ParamEMA
 
 
 matplotlib.rcParams["pdf.fonttype"] = 42
@@ -80,6 +81,14 @@ def get_args():
     parser.add_argument("--epoch", type=int, default=500)
     parser.add_argument("--ckpt", type=str, default=None)
 
+    parser.add_argument(
+        "--use_ema",
+        action="store_true",
+        help="Overlay the EMA (Polyak) shadow of the trainable params instead "
+        "of the raw weights. Falls back to raw weights if the checkpoint has "
+        "no EMA shadow.",
+    )
+
     parser.add_argument(
         "--data_glob",
         type=str,
@@ -157,7 +166,7 @@ def strip_ddp_prefix(state_dict):
     return state_dict
 
 
-def load_9band_model(ckpt_path, device):
+def load_9band_model(ckpt_path, device, use_ema: bool = False):
     ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
 
     model_args = ckpt["model_args"]
@@ -182,6 +191,10 @@ def load_9band_model(ckpt_path, device):
     # False for legacy checkpoints (no key) -> absolute head. Adds no params, so
     # it only changes forward() behavior, never the state_dict shape.
     predict_delta = ckpt.get("predict_delta", False)
+    # False/3 for legacy checkpoints (no key) -> flat-hold anchor. Derived from
+    # x/Dt, so it only changes forward() behavior, never the state_dict shape.
+    trend_decay_anchor = ckpt.get("trend_decay_anchor", False)
+    trend_slope_k = ckpt.get("trend_slope_k", 3)
 
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
@@ -192,6 +205,7 @@ def load_9band_model(ckpt_path, device):
     print("n_bands:", n_bands)
     print("dt_fourier_bands:", dt_fourier_bands)
     print("predict_delta:", predict_delta)
+    print("trend_decay_anchor:", trend_decay_anchor)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -206,6 +220,8 @@ def load_9band_model(ckpt_path, device):
         context_window_days=context_window_days,
         dt_fourier_bands=dt_fourier_bands,
         predict_delta=predict_delta,
+        trend_decay_anchor=trend_decay_anchor,
+        trend_slope_k=trend_slope_k,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
@@ -214,6 +230,29 @@ def load_9band_model(ckpt_path, device):
     print("Missing keys:", missing)
     print("Unexpected keys:", unexpected)
 
+    # Optionally overlay the EMA (Polyak) shadow (conditioner + output_head
+    # only, hence a PARTIAL overlay). Falls back to raw weights with a notice.
+    if use_ema:
+        ema_sd = ckpt.get("ema_state_dict")
+        if ema_sd is None:
+            print("--use_ema requested but checkpoint has no ema_state_dict; "
+                  "using raw weights.")
+        else:
+            ema = ParamEMA(
+                (
+                    (n, p)
+                    for n, p in model.named_parameters()
+                    if p.requires_grad
+                ),
+                decay=ckpt.get("ema_decay", 0.999),
+            )
+            ema.load_state_dict(ema_sd)
+            ema.copy_to(model.named_parameters())
+            print(f"Overlaid EMA weights ({len(ema_sd)} tensors, "
+                  f"decay={ckpt.get('ema_decay', 0.999)}).")
+    else:
+        print("Using raw (non-EMA) weights.")
+
     model.eval()
 
     return model, context_len, n_bands, context_window_days, max_context_len
@@ -472,7 +511,7 @@ def main():
         n_bands,
         context_window_days,
         max_context_len,
-    ) = load_9band_model(args.ckpt, device)
+    ) = load_9band_model(args.ckpt, device, use_ema=getattr(args, "use_ema", False))
 
     window_mode = context_window_days is not None
 
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 0d7acf51..36323740 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -46,6 +46,7 @@
     Kilonova_lc_scalar_context_DataSet_9band,
     load_or_compute_band_normalization,
 )
+from yoke.utils.ema import ParamEMA
 
 
 matplotlib.rcParams["pdf.fonttype"] = 42
@@ -92,6 +93,14 @@ def get_args():
     parser.add_argument("--epoch", type=int, default=500)
     parser.add_argument("--ckpt", type=str, default=None)
 
+    parser.add_argument(
+        "--use_ema",
+        action="store_true",
+        help="Overlay the EMA (Polyak) shadow of the trainable params instead "
+        "of the raw weights. Falls back to raw weights if the checkpoint has "
+        "no EMA shadow.",
+    )
+
     parser.add_argument(
         "--N_imgs",
         type=int,
@@ -180,7 +189,7 @@ def strip_ddp_prefix(state_dict):
     return state_dict
 
 
-def load_9band_model(ckpt_path, device):
+def load_9band_model(ckpt_path, device, use_ema: bool = False):
     ckpt = torch.load(
         ckpt_path,
         map_location=device,
@@ -210,6 +219,10 @@ def load_9band_model(ckpt_path, device):
     # False for legacy checkpoints (no key) -> absolute head. Adds no params, so
     # it only changes forward() behavior, never the state_dict shape.
     predict_delta = ckpt.get("predict_delta", False)
+    # False/3 for legacy checkpoints (no key) -> flat-hold anchor. Derived from
+    # x/Dt, so it only changes forward() behavior, never the state_dict shape.
+    trend_decay_anchor = ckpt.get("trend_decay_anchor", False)
+    trend_slope_k = ckpt.get("trend_slope_k", 3)
 
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
@@ -224,6 +237,7 @@ def load_9band_model(ckpt_path, device):
     print("hidden:", hidden)
     print("dt_fourier_bands:", dt_fourier_bands)
     print("predict_delta:", predict_delta)
+    print("trend_decay_anchor:", trend_decay_anchor)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -238,6 +252,8 @@ def load_9band_model(ckpt_path, device):
         context_window_days=context_window_days,
         dt_fourier_bands=dt_fourier_bands,
         predict_delta=predict_delta,
+        trend_decay_anchor=trend_decay_anchor,
+        trend_slope_k=trend_slope_k,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
@@ -248,6 +264,31 @@ def load_9band_model(ckpt_path, device):
     print("Missing keys:", missing)
     print("Unexpected keys:", unexpected)
 
+    # Optionally overlay the EMA (Polyak) shadow of the trainable params. The
+    # shadow covers only conditioner + output_head, so this is a PARTIAL overlay
+    # (never a strict load). Falls back to the raw weights, with a notice, when
+    # the checkpoint predates EMA or has none.
+    if use_ema:
+        ema_sd = ckpt.get("ema_state_dict")
+        if ema_sd is None:
+            print("--use_ema requested but checkpoint has no ema_state_dict; "
+                  "using raw weights.")
+        else:
+            ema = ParamEMA(
+                (
+                    (n, p)
+                    for n, p in model.named_parameters()
+                    if p.requires_grad
+                ),
+                decay=ckpt.get("ema_decay", 0.999),
+            )
+            ema.load_state_dict(ema_sd)
+            ema.copy_to(model.named_parameters())
+            print(f"Overlaid EMA weights ({len(ema_sd)} tensors, "
+                  f"decay={ckpt.get('ema_decay', 0.999)}).")
+    else:
+        print("Using raw (non-EMA) weights.")
+
     model.eval()
 
     return model, context_len, n_bands, context_window_days, max_context_len
@@ -947,7 +988,7 @@ def main():
         n_bands,
         context_window_days,
         max_context_len,
-    ) = load_9band_model(args.ckpt, device)
+    ) = load_9band_model(args.ckpt, device, use_ema=getattr(args, "use_ema", False))
 
     window_mode = context_window_days is not None
     seed_len = max_context_len if window_mode else context_len
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 14799211..49c60bca 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -24,6 +24,7 @@
     train_DDP_scalar_temporal_loderunner_epoch_9band,
     train_DDP_scalar_temporal_loderunner_epoch_9band_rollout,
 )
+from yoke.utils.ema import ParamEMA
 from yoke.utils.restart import continuation_setup
 from yoke.utils.dataload import make_distributed_dataloader
 from yoke.utils.checkpointing import load_model_and_optimizer
@@ -108,6 +109,19 @@ def _read_stem_list(path: str) -> set:
     "from epoch 0. Only used when --n_rollout_steps > 1.",
 )
 
+# Per-step weight EMA (Polyak averaging) of the trainable params (conditioner +
+# output_head). A shadow average that survives the cycle_epochs=1 process restart
+# through the checkpoint. Smooths the jagged per-batch trajectory; eval can opt in
+# via --use_ema.
+parser.add_argument(
+    "--ema_decay",
+    type=float,
+    default=0.999,
+    help="EMA decay for the trainable-parameter shadow (Polyak averaging), "
+    "updated after each optimizer step. Higher = slower/smoother. Set 0 to "
+    "disable EMA entirely.",
+)
+
 # Paired-dataset globs. The realistic set is always used; the dense set (same
 # objects, denser cadence, no limiting-mag cut) is optional and concatenated onto
 # the realistic training data when present. Both are filtered to the object-level
@@ -246,6 +260,27 @@ def main(args, rank, world_size, local_rank, device):
     # head (byte-identical numerics to the pre-delta model).
     PREDICT_DELTA = True
 
+    # Trend/decay anchor (delta head only). When True, the per-band anchor the
+    # head predicts a residual on top of is no longer the flat last-observed value
+    # but a locally-extrapolated one: anchor[b] = v_last[b] + slope[b] * Dt, with
+    # slope from a least-squares fit over the band's most-recent TREND_SLOPE_K
+    # valid events. This lets the forecast LEAN INTO the fade instead of holding
+    # flat -- targeting the residual blue-band under-fade (u/ztfg/g) that the
+    # flat-hold anchor structurally cannot track. The slope is RAW (unclamped):
+    # bands sampled pre-peak can extrapolate continued brightening (accepted
+    # risk; the head learns a residual on top). To forbid brightening, clamp the
+    # slope non-negative in _band_anchor (one line, flagged there). Adds NO
+    # parameters (derived from x/Dt), so old checkpoints load strict=True; the
+    # flags round-trip via the loaders. Requires PREDICT_DELTA + window mode.
+    TREND_DECAY_ANCHOR = True
+    TREND_SLOPE_K = 3
+
+    # Per-step weight EMA (Polyak averaging) decay for the trainable params. Read
+    # from --ema_decay so it flows through the @input file and survives resubmits.
+    # 0 disables EMA. The shadow is saved into and restored from the .pth each
+    # epoch so it survives the cycle_epochs=1 process restart.
+    EMA_DECAY = args.ema_decay
+
     # Time-window context mode. When CONTEXT_WINDOW_DAYS is not None, the dataset
     # selects context by a trailing lookback in days (all detections within the
     # last CONTEXT_WINDOW_DAYS), padded to MAX_CONTEXT_LEN with a per-event
@@ -409,6 +444,8 @@ def main(args, rank, world_size, local_rank, device):
             context_window_days=CONTEXT_WINDOW_DAYS,
             dt_fourier_bands=DT_FOURIER_BANDS,
             predict_delta=PREDICT_DELTA,
+            trend_decay_anchor=TREND_DECAY_ANCHOR,
+            trend_slope_k=TREND_SLOPE_K,
         ).to(device)
 
         # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
@@ -431,6 +468,39 @@ def main(args, rank, world_size, local_rank, device):
     loss_fn = nn.HuberLoss(delta=0.1, reduction="none")
     model = DDP(model, device_ids=[local_rank], output_device=local_rank)
 
+    #############################################
+    # Per-step weight EMA (Polyak averaging)
+    #############################################
+    # Shadow the trainable params only (conditioner + output_head; the frozen
+    # backbone is excluded by the requires_grad filter). Keyed by name so it
+    # round-trips through the .pth and survives the cycle_epochs=1 restart. On a
+    # CONTINUATION, restore the shadow from the checkpoint BEFORE training so the
+    # average is not silently reset to the current weights every epoch.
+    ema = None
+    if EMA_DECAY > 0:
+        ema = ParamEMA(
+            (
+                (n, p)
+                for n, p in model.module.named_parameters()
+                if p.requires_grad
+            ),
+            decay=EMA_DECAY,
+        )
+        if CONTINUATION:
+            ema_ckpt = torch.load(checkpoint, map_location=device, weights_only=False)
+            ema_sd = ema_ckpt.get("ema_state_dict")
+            if ema_sd is not None:
+                ema.load_state_dict(ema_sd)
+                if rank == 0:
+                    print(
+                        f"Restored EMA shadow ({len(ema_sd)} tensors) from checkpoint."
+                    )
+            elif rank == 0:
+                print(
+                    "No ema_state_dict in checkpoint; EMA initialized from current "
+                    "weights (expected on the first EMA-enabled epoch)."
+                )
+
     #############################################
     # Learning Rate Scheduler
     #############################################
@@ -699,6 +769,7 @@ def _make_9band(
                 max_context_len=MAX_CONTEXT_LEN,
                 band_weights=BAND_WEIGHTS,
                 dt_weight_tau=DT_WEIGHT_TAU,
+                ema=ema,
             )
         else:
             #train_DDP_loderunner_epoch(
@@ -719,6 +790,7 @@ def _make_9band(
                 rank=rank,
                 world_size=world_size,
                 band_weights=BAND_WEIGHTS,
+                ema=ema,
             )
 
         print(f"[rank {rank}] finished epoch", flush=True)
@@ -766,6 +838,12 @@ def _make_9band(
                     "hidden": HIDDEN_CHANNELS,
                     "dt_fourier_bands": DT_FOURIER_BANDS,
                     "predict_delta": PREDICT_DELTA,
+                    "trend_decay_anchor": TREND_DECAY_ANCHOR,
+                    "trend_slope_k": TREND_SLOPE_K,
+                    "ema_decay": EMA_DECAY,
+                    "ema_state_dict": (
+                        ema.state_dict() if ema is not None else None
+                    ),
                     "dt_weight_tau": DT_WEIGHT_TAU,
                     "band_weights": (
                         BAND_WEIGHTS.tolist()
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index e991a761..a226de88 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -489,6 +489,8 @@ def __init__(
         context_window_days: float = None,
         dt_fourier_bands: int = 0,
         predict_delta: bool = False,
+        trend_decay_anchor: bool = False,
+        trend_slope_k: int = 3,
     ) -> None:
         """Initialize conditioner and output-head around the backbone.
 
@@ -529,6 +531,21 @@ def __init__(
                 parameters, so the ``state_dict`` is byte-identical to the absolute
                 model and existing checkpoints load unchanged; the flag is recorded
                 in the checkpoint and restored by the loaders.
+            trend_decay_anchor (bool): When True (default False), the per-band delta
+                anchor is EXTRAPOLATED along that band's recent local slope instead of
+                held flat: ``anchor[b] = v_last[b] + slope[b] * Dt``. The slope is a
+                least-squares fit of value vs. ``rel_t`` over that band's most-recent
+                ``trend_slope_k`` valid events (bands with < 2 valid events get slope 0,
+                i.e. the flat-hold behavior). This lets the forecast lean into a fade
+                rather than plateau at the last value. Requires ``predict_delta`` (it
+                augments the same anchor) and window mode. The slope is RAW (unclamped):
+                for bands whose recent points are pre-peak and rising this extrapolates
+                continued brightening, so the output head must learn a residual to
+                correct it. Adds NO parameters (derived from ``x``/``Dt``), so the
+                ``state_dict`` is unchanged and existing checkpoints load
+                ``strict=True``.
+            trend_slope_k (int): Number of most-recent valid events per band used to fit
+                the local slope when ``trend_decay_anchor`` is on. Default 3.
         """
         super().__init__()
 
@@ -539,6 +556,12 @@ def __init__(
                 "'valid' column present only in the window layout."
             )
 
+        if trend_decay_anchor and not predict_delta:
+            raise ValueError(
+                "trend_decay_anchor=True requires predict_delta=True; the trend "
+                "slope augments the per-band delta anchor."
+            )
+
         self.backbone = backbone
         self.context_len = context_len
         self.n_bands = n_bands
@@ -547,6 +570,8 @@ def __init__(
         self.context_window_days = context_window_days
         self.dt_fourier_bands = dt_fourier_bands
         self.predict_delta = predict_delta
+        self.trend_decay_anchor = trend_decay_anchor
+        self.trend_slope_k = trend_slope_k
 
         # Dataset x layout, flattened per event. Fixed-count mode:
         #   [value, rel_t, one_hot_band(n_bands)] * context_len   -> 2 + n_bands
@@ -624,8 +649,8 @@ def _encode_dt(self, Dt: torch.Tensor, batch_size: int) -> torch.Tensor:
         mono = torch.log1p(Dt.clamp_min(0.0))  # [B, 1], monotone in lead time
         return torch.cat([torch.sin(angles), torch.cos(angles), mono], dim=1)
 
-    def _band_anchor(self, x: torch.Tensor) -> torch.Tensor:
-        """Per-band last observed value from the windowed context, for delta mode.
+    def _band_anchor(self, x: torch.Tensor, Dt: torch.Tensor) -> torch.Tensor:
+        """Per-band anchor from the windowed context, for delta mode.
 
         Reconstructs, for each band, the most-recent observed (normalized) value in
         the trailing context window. The window-mode ``x`` flattens per event as
@@ -635,9 +660,21 @@ def _band_anchor(self, x: torch.Tensor) -> torch.Tensor:
         most-recent observation in ANY band (global-last). A fully empty window (no
         valid events) yields an all-zero anchor (the normalized mean).
 
+        When ``self.trend_decay_anchor`` is set, the per-band anchor is additionally
+        extrapolated along that band's recent local slope,
+        ``anchor[b] = v_last[b] + slope[b] * Dt``, where ``slope[b]`` is a
+        least-squares fit of value vs. ``rel_t`` over that band's most-recent
+        ``self.trend_slope_k`` valid events. Bands with < 2 valid events get slope 0
+        (flat hold). The slope is RAW (unclamped): a pre-peak/rising band extrapolates
+        continued brightening. ``Dt`` and ``rel_t`` share the same time units (days),
+        so the extrapolation is unit-consistent.
+
         Args:
             x (torch.Tensor): Window-mode context, shape
                 [B, context_len * (3 + n_bands)].
+            Dt (torch.Tensor): Lead time (days) from the anchor (most-recent event) to
+                the target, shape [B] or broadcastable. Only used when
+                ``self.trend_decay_anchor``; the flat-hold path ignores it.
 
         Returns:
             torch.Tensor: Per-band anchor of shape [B, n_bands] in normalized units.
@@ -679,7 +716,44 @@ def _band_anchor(self, x: torch.Tensor) -> torch.Tensor:
         ).squeeze(2)  # [B, nb]
 
         # Fall back to global-last where a band was never observed in the window.
-        return torch.where(pb_has, per_band_last, global_last.unsqueeze(1))
+        v_last = torch.where(pb_has, per_band_last, global_last.unsqueeze(1))
+
+        if not self.trend_decay_anchor:
+            return v_last
+
+        # Trend/decay anchor: extrapolate each band along the local slope of its
+        # most-recent trend_slope_k valid events, anchor = v_last + slope * Dt.
+        # Select those events per band as the top-k by rel_t (pb_score already masks
+        # non-matching/invalid events to -inf), then least-squares-fit value vs rel_t.
+        k = min(self.trend_slope_k, self.context_len)
+        topk_t, topk_idx = pb_score.topk(k, dim=2)  # [B, nb, k]; -inf where < k matches
+        # Robust validity of each selected slot: gather the band-match boolean at the
+        # chosen indices (top-k of an all -inf row lands on non-matching events).
+        topk_valid = torch.gather(
+            band_match.to(value.dtype), 2, topk_idx
+        ) > 0.5  # [B, nb, k]
+        topk_v = torch.gather(
+            value.unsqueeze(1).expand(B, nb, self.context_len), 2, topk_idx
+        )  # [B, nb, k]
+
+        # Weighted (valid-only) least-squares slope of v vs t per (batch, band).
+        w = topk_valid.to(value.dtype)  # [B, nb, k]
+        t = torch.where(topk_valid, topk_t, torch.zeros_like(topk_t))
+        v = torch.where(topk_valid, topk_v, torch.zeros_like(topk_v))
+        n = w.sum(dim=2)  # [B, nb]
+        denom = n.clamp_min(1.0)
+        t_bar = (w * t).sum(dim=2) / denom  # [B, nb]
+        v_bar = (w * v).sum(dim=2) / denom  # [B, nb]
+        dt_ = t - t_bar.unsqueeze(2)
+        dv_ = v - v_bar.unsqueeze(2)
+        cov = (w * dt_ * dv_).sum(dim=2)  # [B, nb]
+        var = (w * dt_ * dt_).sum(dim=2)  # [B, nb]
+        slope = torch.where(
+            (n >= 2.0) & (var > 1e-8), cov / var.clamp_min(1e-8), torch.zeros_like(cov)
+        )  # [B, nb]; bands with < 2 valid events -> 0 (flat hold)
+
+        # RAW slope (no clamp). To forbid brightening, add: slope = slope.clamp_min(0.0)
+        return v_last + slope * Dt.reshape(B, 1)
 
     def forward(
         self,
@@ -751,7 +825,7 @@ def forward(
         # anchor is derived from x (no parameters), so the disabled path is
         # byte-identical to the absolute model.
         if self.predict_delta:
-            pred = pred + self._band_anchor(x)
+            pred = pred + self._band_anchor(x, Dt)
 
         return pred
 
diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index dda20f05..1183b71f 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -502,6 +502,11 @@ def load_direct_loderunner_checkpoint_9band(
         # False for legacy checkpoints (no key) -> absolute head. Adds no params,
         # so this only changes forward() behavior, never the state_dict.
         predict_delta=checkpoint_data.get("predict_delta", False),
+        # False/3 for legacy checkpoints (no key) -> flat-hold anchor. Derived
+        # from x/Dt, so this only changes forward() behavior, never the
+        # state_dict, and strict load stays valid.
+        trend_decay_anchor=checkpoint_data.get("trend_decay_anchor", False),
+        trend_slope_k=checkpoint_data.get("trend_slope_k", 3),
     ).to(device)
 
     state_dict = checkpoint_data["model_state_dict"]
diff --git a/src/yoke/utils/training/epoch/loderunner.py b/src/yoke/utils/training/epoch/loderunner.py
index 7d634231..9d12c641 100644
--- a/src/yoke/utils/training/epoch/loderunner.py
+++ b/src/yoke/utils/training/epoch/loderunner.py
@@ -498,6 +498,7 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
     rank: int,
     world_size: int,
     band_weights: torch.Tensor = None,
+    ema: object = None,
 ) -> None:
     """DDP epoch function for the masked 9-band scalar temporal LodeRunner.
 
@@ -520,6 +521,10 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
     under-fit; weighting scales each sample by its observed band's weight. The
     RECORDED per-sample loss stays unweighted so the CSV metric is comparable
     across runs. ``None`` (default) reproduces the plain equal-weight behavior.
+
+    ``ema`` (optional :class:`yoke.utils.ema.ParamEMA`) shadows the trainable
+    params and is updated after each ``optimizer.step()`` (Polyak averaging);
+    ``None`` (default) is a no-op.
     """
     train_rcrd_filename = train_rcrd_filename.replace(
         "<epochIDX>",
@@ -582,6 +587,12 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band(
             optimizer.step()
             LRsched.step()
 
+            if ema is not None:
+                core = model.module if hasattr(model, "module") else model
+                ema.update(
+                    (n, p) for n, p in core.named_parameters() if p.requires_grad
+                )
+
             if rank == 0:
                 batch_records = np.column_stack(
                     [
@@ -1011,6 +1022,7 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
     max_context_len: int = None,
     band_weights: torch.Tensor = None,
     dt_weight_tau: float = None,
+    ema: object = None,
 ) -> None:
     """Multi-step rollout DDP epoch for the masked 9-band scalar temporal model.
 
@@ -1054,6 +1066,12 @@ def train_DDP_scalar_temporal_loderunner_epoch_9band_rollout(
             long-horizon steps so the early rise is not swamped by the many
             late-tail points. The recorded per-sample loss stays unweighted, and
             validation always runs unweighted. ``None`` (default) disables it.
+        ema (object): Optional :class:`yoke.utils.ema.ParamEMA` shadow of the
+            trainable params (conditioner + output_head). When supplied, its
+            ``update`` is called after every ``optimizer.step()`` so the average
+            tracks the per-batch trajectory (Polyak averaging). Only the
+            ``requires_grad`` params of the underlying (DDP-unwrapped) module are
+            shadowed. ``None`` (default) is a no-op.
     """
     def _run_pass(
         data: tuple[torch.Tensor, ...],
@@ -1155,6 +1173,12 @@ def _run_pass(
             optimizer.step()
             LRsched.step()
 
+            if ema is not None:
+                core = model.module if hasattr(model, "module") else model
+                ema.update(
+                    (n, p) for n, p in core.named_parameters() if p.requires_grad
+                )
+
             if rank == 0:
                 batch_records = np.column_stack(
                     [

From 8af7647b027173be11dfc06832f13b563dda9d62 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 14:39:53 -0600
Subject: [PATCH 54/66] reduce number of epochs

---
 applications/harnesses/KN_loderunner/training_START.input | 2 +-
 applications/harnesses/KN_loderunner/training_input.tmpl  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_START.input b/applications/harnesses/KN_loderunner/training_START.input
index f019984d..9ba19892 100644
--- a/applications/harnesses/KN_loderunner/training_START.input
+++ b/applications/harnesses/KN_loderunner/training_START.input
@@ -50,7 +50,7 @@ kn_rubin_ztf_val.txt
 --Knodes
 <KNODES>
 --total_epochs
-90
+40
 --cycle_epochs
 1
 --train_batches
diff --git a/applications/harnesses/KN_loderunner/training_input.tmpl b/applications/harnesses/KN_loderunner/training_input.tmpl
index f552b660..c68d90a3 100644
--- a/applications/harnesses/KN_loderunner/training_input.tmpl
+++ b/applications/harnesses/KN_loderunner/training_input.tmpl
@@ -50,7 +50,7 @@ kn_rubin_ztf_val.txt
 --Knodes
 <KNODES>
 --total_epochs
-90
+40
 --cycle_epochs
 1
 --train_batches

From b2a3fe14e813012a4cedd850392d6cafd8a45ce5 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 14:56:27 -0600
Subject: [PATCH 55/66] adding missing exponential moving average file

---
 src/yoke/utils/ema.py | 114 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 src/yoke/utils/ema.py

diff --git a/src/yoke/utils/ema.py b/src/yoke/utils/ema.py
new file mode 100644
index 00000000..4c069033
--- /dev/null
+++ b/src/yoke/utils/ema.py
@@ -0,0 +1,114 @@
+"""Exponential moving average (Polyak averaging) of selected model parameters.
+
+A small, framework-light EMA shadow for a *subset* of a model's parameters --
+here the trainable conditioner + output-head of the scalar-temporal LodeRunner
+wrapper, with the frozen backbone deliberately excluded. Keeping the shadow keyed
+by parameter NAME (rather than by position) makes it robust to reload: the
+state_dict round-trips through the training checkpoint so the average survives the
+per-epoch process restart used by the DDP harness (``cycle_epochs=1``), and the
+eval scripts can overlay it onto a freshly-constructed model without needing an
+optimizer.
+
+Typical use in training::
+
+    ema = ParamEMA(
+        ((n, p) for n, p in model.named_parameters() if p.requires_grad),
+        decay=0.999,
+    )
+    ...
+    optimizer.step()
+    ema.update((n, p) for n, p in model.named_parameters() if p.requires_grad)
+    ...
+    torch.save({..., "ema_state_dict": ema.state_dict()}, path)
+
+and in eval to compare against the raw weights::
+
+    ema = ParamEMA(named_trainable_params, decay=ckpt["ema_decay"])
+    ema.load_state_dict(ckpt["ema_state_dict"])
+    ema.copy_to(model.named_parameters())  # partial overlay of the shadowed subset
+"""
+
+from collections.abc import Iterable
+
+import torch
+
+
+class ParamEMA:
+    """Name-keyed exponential moving average over a subset of parameters.
+
+    The shadow tracks ``shadow <- decay * shadow + (1 - decay) * param`` for each
+    supplied (name, parameter) pair. Only the provided parameters are shadowed, so
+    passing just the trainable params keeps the frozen backbone out of the average
+    (and out of the checkpoint).
+    """
+
+    def __init__(
+        self,
+        named_params: Iterable[tuple[str, torch.Tensor]],
+        decay: float = 0.999,
+    ) -> None:
+        """Initialize the shadow from the current parameter values.
+
+        Args:
+            named_params (Iterable[tuple[str, torch.Tensor]]): (name, parameter)
+                pairs to shadow, e.g. the ``requires_grad`` subset of
+                ``model.named_parameters()``. Consumed once (may be a generator).
+            decay (float): EMA decay in [0, 1). Higher = slower/smoother. The shadow
+                is initialized to a detached clone of each parameter.
+        """
+        self.decay = float(decay)
+        self.shadow: dict[str, torch.Tensor] = {
+            name: p.detach().clone() for name, p in named_params
+        }
+
+    @torch.no_grad()
+    def update(self, named_params: Iterable[tuple[str, torch.Tensor]]) -> None:
+        """Update the shadow toward the current parameter values.
+
+        Args:
+            named_params (Iterable[tuple[str, torch.Tensor]]): The same (name,
+                parameter) pairs supplied at construction (order-independent; matched
+                by name). Names not present in the shadow are ignored; missing shadow
+                entries are skipped.
+        """
+        d = self.decay
+        for name, p in named_params:
+            s = self.shadow.get(name)
+            if s is None:
+                continue
+            # shadow <- decay * shadow + (1 - decay) * param
+            s.lerp_(p.detach().to(s.device, s.dtype), 1.0 - d)
+
+    def state_dict(self) -> dict[str, torch.Tensor]:
+        """Return the shadow as a CPU-resident name->tensor dict for checkpointing."""
+        return {name: s.detach().cpu().clone() for name, s in self.shadow.items()}
+
+    def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
+        """Load a previously saved shadow, copying into the existing tensors in place.
+
+        Args:
+            state_dict (dict[str, torch.Tensor]): Mapping produced by
+                :meth:`state_dict`. Only keys present in the current shadow are
+                loaded; the copy preserves each shadow tensor's device/dtype.
+        """
+        for name, s in self.shadow.items():
+            saved = state_dict.get(name)
+            if saved is not None:
+                s.copy_(saved.to(s.device, s.dtype))
+
+    @torch.no_grad()
+    def copy_to(self, named_params: Iterable[tuple[str, torch.Tensor]]) -> None:
+        """Overlay the shadow onto live parameters (a partial, in-place assignment).
+
+        Only parameters whose names are in the shadow are overwritten, so this is a
+        PARTIAL overlay -- the frozen/backbone params of the target model are left
+        untouched. Use this at eval time to swap in the averaged weights.
+
+        Args:
+            named_params (Iterable[tuple[str, torch.Tensor]]): (name, parameter) pairs
+                of the target model, e.g. ``model.named_parameters()``.
+        """
+        for name, p in named_params:
+            s = self.shadow.get(name)
+            if s is not None:
+                p.data.copy_(s.to(p.device, p.dtype))

From 51a33957ac3dac8fb3064a9b89cd93e0701ec764 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 15:34:50 -0600
Subject: [PATCH 56/66] EMA to .99

---
 applications/harnesses/KN_loderunner/infer_9band.py        | 4 ++--
 .../harnesses/KN_loderunner/plot_pred_diagnostics_9band.py | 4 ++--
 .../harnesses/KN_loderunner/train_LodeRunner_ddp.py        | 7 +++++--
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 4bcba32c..34cdc963 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -244,12 +244,12 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
                     for n, p in model.named_parameters()
                     if p.requires_grad
                 ),
-                decay=ckpt.get("ema_decay", 0.999),
+                decay=ckpt.get("ema_decay", 0.99),
             )
             ema.load_state_dict(ema_sd)
             ema.copy_to(model.named_parameters())
             print(f"Overlaid EMA weights ({len(ema_sd)} tensors, "
-                  f"decay={ckpt.get('ema_decay', 0.999)}).")
+                  f"decay={ckpt.get('ema_decay', 0.99)}).")
     else:
         print("Using raw (non-EMA) weights.")
 
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 36323740..380d07f1 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -280,12 +280,12 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
                     for n, p in model.named_parameters()
                     if p.requires_grad
                 ),
-                decay=ckpt.get("ema_decay", 0.999),
+                decay=ckpt.get("ema_decay", 0.99),
             )
             ema.load_state_dict(ema_sd)
             ema.copy_to(model.named_parameters())
             print(f"Overlaid EMA weights ({len(ema_sd)} tensors, "
-                  f"decay={ckpt.get('ema_decay', 0.999)}).")
+                  f"decay={ckpt.get('ema_decay', 0.99)}).")
     else:
         print("Using raw (non-EMA) weights.")
 
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 49c60bca..efd30c03 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -116,9 +116,12 @@ def _read_stem_list(path: str) -> set:
 parser.add_argument(
     "--ema_decay",
     type=float,
-    default=0.999,
+    default=0.99,
     help="EMA decay for the trainable-parameter shadow (Polyak averaging), "
-    "updated after each optimizer step. Higher = slower/smoother. Set 0 to "
+    "updated after each optimizer step. Higher = slower/smoother; the averaging "
+    "window is ~1/(1-decay) steps. 0.99 ~= 100 steps ~= 1.2 epochs at "
+    "NTRN_BATCH=84, matched to the 40-epoch run so the shadow tracks the "
+    "trajectory instead of lagging toward the initial weights. Set 0 to "
     "disable EMA entirely.",
 )
 

From ca9a78725c1e75879d4192b618b6978fdb484a46 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 16:44:19 -0600
Subject: [PATCH 57/66] fix plot regimes

---
 .../harnesses/KN_loderunner/plot_loss_curves_9band.py  | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
index 11371346..d896e08f 100644
--- a/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py
@@ -114,16 +114,18 @@ def main():
     parser.add_argument(
         "--tf_ramp_start_epoch",
         type=int,
-        default=20,
+        default=8,
         help="Absolute epoch at which the teacher-forcing anneal begins. Must "
-        "match the training schedule. Default 20.",
+        "match --tf_ramp_start_epoch in training_START.input / "
+        "training_input.tmpl. Default 8.",
     )
     parser.add_argument(
         "--tf_ramp_epochs",
         type=int,
-        default=20,
+        default=12,
         help="Number of epochs the teacher-forcing ratio anneals over. Must "
-        "match the training schedule. Default 20.",
+        "match --tf_ramp_epochs in training_START.input / training_input.tmpl. "
+        "Default 12.",
     )
 
     args = parser.parse_args()

From 604b64b4e2cf6c1087ce94c558feacca184cad03 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Wed, 26 Aug 2026 22:31:06 -0600
Subject: [PATCH 58/66] training bug fix

---
 .../harnesses/KN_loderunner/infer_9band.py    |  3 ++
 .../plot_pred_diagnostics_9band.py            |  3 ++
 .../KN_loderunner/train_LodeRunner_ddp.py     | 13 +++++++
 src/yoke/models/vit/swin/bomberman.py         | 35 ++++++++++++++++---
 src/yoke/utils/checkpointing.py               |  1 +
 5 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py
index 34cdc963..9bfdafdb 100644
--- a/applications/harnesses/KN_loderunner/infer_9band.py
+++ b/applications/harnesses/KN_loderunner/infer_9band.py
@@ -195,6 +195,7 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
     # x/Dt, so it only changes forward() behavior, never the state_dict shape.
     trend_decay_anchor = ckpt.get("trend_decay_anchor", False)
     trend_slope_k = ckpt.get("trend_slope_k", 3)
+    trend_max_offset = ckpt.get("trend_max_offset", None)
 
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
@@ -206,6 +207,7 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
     print("dt_fourier_bands:", dt_fourier_bands)
     print("predict_delta:", predict_delta)
     print("trend_decay_anchor:", trend_decay_anchor)
+    print("trend_max_offset:", trend_max_offset)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -222,6 +224,7 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
         predict_delta=predict_delta,
         trend_decay_anchor=trend_decay_anchor,
         trend_slope_k=trend_slope_k,
+        trend_max_offset=trend_max_offset,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
diff --git a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
index 380d07f1..cf3d7bfe 100644
--- a/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
+++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py
@@ -223,6 +223,7 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
     # x/Dt, so it only changes forward() behavior, never the state_dict shape.
     trend_decay_anchor = ckpt.get("trend_decay_anchor", False)
     trend_slope_k = ckpt.get("trend_slope_k", 3)
+    trend_max_offset = ckpt.get("trend_max_offset", None)
 
     print("Loaded checkpoint:", ckpt_path)
     print("model_class:", ckpt.get("model_class", "unknown"))
@@ -238,6 +239,7 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
     print("dt_fourier_bands:", dt_fourier_bands)
     print("predict_delta:", predict_delta)
     print("trend_decay_anchor:", trend_decay_anchor)
+    print("trend_max_offset:", trend_max_offset)
 
     backbone = LodeRunner(**model_args).to(device)
     backbone.noise_scale = noise_scale
@@ -254,6 +256,7 @@ def load_9band_model(ckpt_path, device, use_ema: bool = False):
         predict_delta=predict_delta,
         trend_decay_anchor=trend_decay_anchor,
         trend_slope_k=trend_slope_k,
+        trend_max_offset=trend_max_offset,
     ).to(device)
 
     state_dict = strip_ddp_prefix(ckpt["model_state_dict"])
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index efd30c03..89d6b060 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -278,6 +278,17 @@ def main(args, rank, world_size, local_rank, device):
     TREND_DECAY_ANCHOR = True
     TREND_SLOPE_K = 3
 
+    # Cap on the extrapolated anchor offset slope*Dt, in per-band z-score units
+    # (values are z-scored per band). A steep raw slope times a large Dt (horizon
+    # up to ~8 days, 12-step rollout) can push the anchor far past the last real
+    # point, so the head must learn a big corrective residual -- which floors the
+    # loss and blocks descent (observed in the first raw-slope run). Clamping the
+    # offset symmetrically to +/- TREND_MAX_OFFSET keeps the anchor near the data
+    # while still letting it lean into the fade. Sign-agnostic: bounds both an
+    # over-fast fade and the near-peak brightening the raw slope extrapolates. Set
+    # None to disable the cap (raw slope).
+    TREND_MAX_OFFSET = 1.0
+
     # Per-step weight EMA (Polyak averaging) decay for the trainable params. Read
     # from --ema_decay so it flows through the @input file and survives resubmits.
     # 0 disables EMA. The shadow is saved into and restored from the .pth each
@@ -449,6 +460,7 @@ def main(args, rank, world_size, local_rank, device):
             predict_delta=PREDICT_DELTA,
             trend_decay_anchor=TREND_DECAY_ANCHOR,
             trend_slope_k=TREND_SLOPE_K,
+            trend_max_offset=TREND_MAX_OFFSET,
         ).to(device)
 
         # Stage 1: freeze pretrained LodeRunner, train only conditioner + output head
@@ -843,6 +855,7 @@ def _make_9band(
                     "predict_delta": PREDICT_DELTA,
                     "trend_decay_anchor": TREND_DECAY_ANCHOR,
                     "trend_slope_k": TREND_SLOPE_K,
+                    "trend_max_offset": TREND_MAX_OFFSET,
                     "ema_decay": EMA_DECAY,
                     "ema_state_dict": (
                         ema.state_dict() if ema is not None else None
diff --git a/src/yoke/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py
index a226de88..040e56d2 100644
--- a/src/yoke/models/vit/swin/bomberman.py
+++ b/src/yoke/models/vit/swin/bomberman.py
@@ -491,6 +491,7 @@ def __init__(
         predict_delta: bool = False,
         trend_decay_anchor: bool = False,
         trend_slope_k: int = 3,
+        trend_max_offset: float = None,
     ) -> None:
         """Initialize conditioner and output-head around the backbone.
 
@@ -546,6 +547,18 @@ def __init__(
                 ``strict=True``.
             trend_slope_k (int): Number of most-recent valid events per band used to fit
                 the local slope when ``trend_decay_anchor`` is on. Default 3.
+            trend_max_offset (float): When set (and ``trend_decay_anchor`` is on), the
+                extrapolated anchor offset ``slope[b] * Dt`` is symmetrically clamped to
+                ``[-trend_max_offset, +trend_max_offset]`` (in per-band normalized /
+                z-score units) before being added to ``v_last``. Values are z-scored per
+                band, so this bounds the anchor displacement to a fixed number of
+                standard deviations regardless of fade direction or lead time -- capping
+                the overshoot that a steep raw slope produces at large ``Dt`` (the head
+                would otherwise have to learn a large corrective residual, flooring the
+                loss). It is deliberately SIGN-AGNOSTIC: it limits both an over-fast fade
+                and the near-peak pre-peak "brightening" extrapolation of the raw slope.
+                When None (default) the offset is RAW (unclamped), matching the original
+                trend-anchor behavior. Adds NO parameters; round-trips via checkpoint.
         """
         super().__init__()
 
@@ -572,6 +585,7 @@ def __init__(
         self.predict_delta = predict_delta
         self.trend_decay_anchor = trend_decay_anchor
         self.trend_slope_k = trend_slope_k
+        self.trend_max_offset = trend_max_offset
 
         # Dataset x layout, flattened per event. Fixed-count mode:
         #   [value, rel_t, one_hot_band(n_bands)] * context_len   -> 2 + n_bands
@@ -665,9 +679,12 @@ def _band_anchor(self, x: torch.Tensor, Dt: torch.Tensor) -> torch.Tensor:
         ``anchor[b] = v_last[b] + slope[b] * Dt``, where ``slope[b]`` is a
         least-squares fit of value vs. ``rel_t`` over that band's most-recent
         ``self.trend_slope_k`` valid events. Bands with < 2 valid events get slope 0
-        (flat hold). The slope is RAW (unclamped): a pre-peak/rising band extrapolates
-        continued brightening. ``Dt`` and ``rel_t`` share the same time units (days),
-        so the extrapolation is unit-consistent.
+        (flat hold). When ``self.trend_max_offset`` is set, the offset ``slope[b] * Dt``
+        is symmetrically clamped to ``[-trend_max_offset, +trend_max_offset]`` (z-score
+        units) so a steep slope cannot overshoot at large ``Dt``; when None the slope is
+        RAW (unclamped) and a pre-peak/rising band extrapolates continued brightening.
+        ``Dt`` and ``rel_t`` share the same time units (days), so the extrapolation is
+        unit-consistent.
 
         Args:
             x (torch.Tensor): Window-mode context, shape
@@ -752,8 +769,16 @@ def _band_anchor(self, x: torch.Tensor, Dt: torch.Tensor) -> torch.Tensor:
             (n >= 2.0) & (var > 1e-8), cov / var.clamp_min(1e-8), torch.zeros_like(cov)
         )  # [B, nb]; bands with < 2 valid events -> 0 (flat hold)
 
-        # RAW slope (no clamp). To forbid brightening, add: slope = slope.clamp_min(0.0)
-        return v_last + slope * Dt.reshape(B, 1)
+        # Extrapolated offset in z-score units. Symmetrically clamp its MAGNITUDE when
+        # trend_max_offset is set: a steep raw slope times a large Dt would otherwise
+        # push the anchor far past the data, forcing the head to learn a large
+        # corrective residual (which floors the loss). The cap is sign-agnostic, so it
+        # bounds both an over-fast fade and the near-peak pre-peak brightening the raw
+        # slope produces. When None, the offset is raw (unclamped).
+        offset = slope * Dt.reshape(B, 1)
+        if self.trend_max_offset is not None:
+            offset = offset.clamp(-self.trend_max_offset, self.trend_max_offset)
+        return v_last + offset
 
     def forward(
         self,
diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py
index 1183b71f..cd5305ad 100644
--- a/src/yoke/utils/checkpointing.py
+++ b/src/yoke/utils/checkpointing.py
@@ -507,6 +507,7 @@ def load_direct_loderunner_checkpoint_9band(
         # state_dict, and strict load stays valid.
         trend_decay_anchor=checkpoint_data.get("trend_decay_anchor", False),
         trend_slope_k=checkpoint_data.get("trend_slope_k", 3),
+        trend_max_offset=checkpoint_data.get("trend_max_offset", None),
     ).to(device)
 
     state_dict = checkpoint_data["model_state_dict"]

From e8a67183f4191b70e330319aef590c47efd00b7e Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 27 Aug 2026 08:13:44 -0600
Subject: [PATCH 59/66] config change

---
 .../KN_loderunner/training_slurm.tmpl          | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index 149df61f..bf209b61 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -17,12 +17,16 @@
 #SBATCH --nodes=<KNODES>
 #SBATCH --ntasks-per-node=<NGPUS>
 #SBATCH --gpus-per-node=<NGPUS>
-# nodes have 8 H100s and 192-208 CPUs, i.e. ~24 cores/GPU. One task
-# per GPU (ntasks-per-node=NGPUS), so cpus-per-task=24 gives each rank enough
-# cores for its dataloader workers + OMP threads without over-requesting on the
-# 192-core nodes (24*NGPUS stays within budget for NGPUS up to 8). Without this
-# each task defaults to ~1 CPU and starves data loading, throttling the GPUs.
-#SBATCH --cpus-per-task=24
+# nodes have 8 H100s and 144 CPUs, i.e. 18 cores/GPU. One task per GPU
+# (ntasks-per-node=NGPUS), so cpus-per-task=18 gives each rank enough cores for
+# its dataloader workers + OMP threads. Do NOT set this to 24: with the observed
+# 144-core nodes, 24*6=144 packs the node to 100%, and under the implicit --exact
+# (auto-enabled by -c) the step cannot be created until every core is free at
+# once -- so a lingering prior-epoch step (cycle_epochs=1 resubmit) or the
+# pre-flight diagnostic srun blocks step creation forever ("Requested nodes are
+# busy", retrying). 18*6=108 leaves 36 cores of slack so the step always creates;
+# 18 comfortably fits OMP_NUM_THREADS=8 + NUM_WORKERS dataloader procs per rank.
+#SBATCH --cpus-per-task=18
 #SBATCH --mem-per-gpu=50G
 #SBATCH --output=study<studyIDX>_epoch<epochIDX>.out
 #SBATCH --error=study<studyIDX>_epoch<epochIDX>.err
@@ -60,7 +64,7 @@ module load anaconda/3.12
 source activate
 conda activate <YOKE_TORCH_ENV>
 
-# Set number of threads per GPU. Must be <= cpus-per-task (24). Leaves cores for
+# Set number of threads per GPU. Must be <= cpus-per-task (18). Leaves cores for
 # the dataloader workers (NUM_WORKERS per rank) alongside the OMP compute threads.
 export OMP_NUM_THREADS=8
 

From 5aa17548e02fc607bbd4fc012e7a5130f62cdca8 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 27 Aug 2026 09:35:18 -0600
Subject: [PATCH 60/66] update config for 8 gpu run

---
 .../KN_loderunner/ddp_production.csv          | 12 +++++++++-
 .../KN_loderunner/training_slurm.tmpl         | 23 ++++++++++---------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
index a0460437..0b5eddbd 100644
--- a/applications/harnesses/KN_loderunner/ddp_production.csv
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -42,6 +42,16 @@ studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZ
 # TERMINAL/WARMUP scaled to per-rank batches (84/42) so the LR schedule spans the
 # same epochs. Uses the current train script: PREDICT_DELTA=True + DT_WEIGHT_TAU=3
 # + ZTF band weights bumped to 2 (A+B+ZTF combined, not isolated).
-41,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
+#41,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
+# study 42: single node, all 8 GPUs (the max on these 8-H100 se* nodes). Replaces
+# the 2-node study 41, whose cross-node job kept hanging in Slurm step creation
+# ("Requested nodes are busy") -- a single-node job needs only one node free and
+# has no cross-node NCCL/IB teardown between the cycle_epochs=1 per-epoch jobs.
+# BATCH_SIZE=5 keeps the 12-step rollout under 80GB (same as study 40/41). NTRN/NVAL
+# 125/63 per rank hold ~1000/500 samples per epoch across 8 ranks (8*125=1000),
+# identical epoch size to studies 24/25/40/41. TERMINAL/WARMUP=125/63 (per-rank
+# batches) so the LR schedule spans the same epochs. Same model settings as 41:
+# PREDICT_DELTA + trend/decay anchor (now offset-capped) + EMA + ZTF band weights.
+42,yoke311,1,8,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
 
 
diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index bf209b61..59a5dc37 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -17,16 +17,17 @@
 #SBATCH --nodes=<KNODES>
 #SBATCH --ntasks-per-node=<NGPUS>
 #SBATCH --gpus-per-node=<NGPUS>
-# nodes have 8 H100s and 144 CPUs, i.e. 18 cores/GPU. One task per GPU
-# (ntasks-per-node=NGPUS), so cpus-per-task=18 gives each rank enough cores for
-# its dataloader workers + OMP threads. Do NOT set this to 24: with the observed
-# 144-core nodes, 24*6=144 packs the node to 100%, and under the implicit --exact
-# (auto-enabled by -c) the step cannot be created until every core is free at
-# once -- so a lingering prior-epoch step (cycle_epochs=1 resubmit) or the
-# pre-flight diagnostic srun blocks step creation forever ("Requested nodes are
-# busy", retrying). 18*6=108 leaves 36 cores of slack so the step always creates;
-# 18 comfortably fits OMP_NUM_THREADS=8 + NUM_WORKERS dataloader procs per rank.
-#SBATCH --cpus-per-task=18
+# nodes have 8 H100s and 144 CPUs. One task per GPU (ntasks-per-node=NGPUS), so
+# cpus-per-task=16 gives each rank enough cores for its dataloader workers + OMP
+# threads while leaving the node UNDER-packed. Do NOT raise this to 18: at the
+# max 8 GPUs (study 42), 18*8=144 packs the 144-core node to 100%, and under the
+# implicit --exact (auto-enabled by -c) the step cannot be created until every
+# core is free at once -- so a lingering prior-epoch step (cycle_epochs=1
+# resubmit) or the pre-flight diagnostic srun blocks step creation forever
+# ("Requested nodes are busy", retrying). 16*8=128 leaves 16 cores of slack so
+# the step always creates; 16 comfortably fits OMP_NUM_THREADS=8 + NUM_WORKERS
+# dataloader procs per rank.
+#SBATCH --cpus-per-task=16
 #SBATCH --mem-per-gpu=50G
 #SBATCH --output=study<studyIDX>_epoch<epochIDX>.out
 #SBATCH --error=study<studyIDX>_epoch<epochIDX>.err
@@ -64,7 +65,7 @@ module load anaconda/3.12
 source activate
 conda activate <YOKE_TORCH_ENV>
 
-# Set number of threads per GPU. Must be <= cpus-per-task (18). Leaves cores for
+# Set number of threads per GPU. Must be <= cpus-per-task (16). Leaves cores for
 # the dataloader workers (NUM_WORKERS per rank) alongside the OMP compute threads.
 export OMP_NUM_THREADS=8
 

From 2e5cc95259770af6746334e311e267fb36e08e59 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Thu, 27 Aug 2026 21:39:16 -0600
Subject: [PATCH 61/66] exclude anchor feature

---
 .../harnesses/KN_loderunner/train_LodeRunner_ddp.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 89d6b060..45b3b511 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -275,7 +275,18 @@ def main(args, rank, world_size, local_rank, device):
     # slope non-negative in _band_anchor (one line, flagged there). Adds NO
     # parameters (derived from x/Dt), so old checkpoints load strict=True; the
     # flags round-trip via the loaders. Requires PREDICT_DELTA + window mode.
-    TREND_DECAY_ANCHOR = True
+    #
+    # ABLATION (set False): the trend anchor is disabled to isolate whether it is
+    # what floored the training loss. Studies 24/40 (flat-hold anchor) descended;
+    # studies 049 (raw slope) and its capped variant (TREND_MAX_OFFSET=1.0) both
+    # stayed flat, so the offset cap did NOT restore descent -- ruling out
+    # raw-slope overshoot as the cause. With this False the head predicts a
+    # residual on the flat last-observed anchor again (EMA stays on but is inert
+    # for the training loss). If the loss now descends -> the trend anchor is the
+    # cause and stays off; if it stays flat -> the cause is elsewhere (12-step
+    # rollout, DT_WEIGHT_TAU, ZTF weight bump, or Huber delta). TREND_MAX_OFFSET
+    # below is ignored while this is False.
+    TREND_DECAY_ANCHOR = False
     TREND_SLOPE_K = 3
 
     # Cap on the extrapolated anchor offset slope*Dt, in per-band z-score units

From a169ed5df8a17d0157a1606308473c5ad64c5cb6 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 28 Aug 2026 09:10:29 -0600
Subject: [PATCH 62/66] adjust Huber loss

---
 .../harnesses/KN_loderunner/train_LodeRunner_ddp.py    | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 45b3b511..131e006f 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -491,7 +491,15 @@ def main(args, rank, world_size, local_rank, device):
         )
 
     #loss_fn = nn.MSELoss(reduction="none")
-    loss_fn = nn.HuberLoss(delta=0.1, reduction="none")
+    # delta=0.1 was in place since the harness was created (studies 24/25 too), so
+    # it is NOT what newly flattened the loss -- but it does throttle descent: with
+    # z-scored targets (sigma~=1) and a mean loss ~0.09, typical errors are ~0.95
+    # sigma, deep in Huber's LINEAR regime where the per-point gradient is capped
+    # at delta=0.1. Raising delta to 1.0 gives an error-proportional gradient (~10x
+    # stronger here) so the model can actually push down the bulk error, and only
+    # reverts to the robust linear regime past 1 sigma. delta=1.0 ~= MSE for the
+    # in-sigma bulk while still clipping the heavy-tailed outliers.
+    loss_fn = nn.HuberLoss(delta=1.0, reduction="none")
     model = DDP(model, device_ids=[local_rank], output_device=local_rank)
 
     #############################################

From d94b7c8fbbc7108a65d2aa3a90972d9f3032c541 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 28 Aug 2026 10:55:56 -0600
Subject: [PATCH 63/66] trying to fix 802 error

---
 .../KN_loderunner/ddp_production.csv          | 60 +++++++++----------
 .../KN_loderunner/training_slurm.tmpl         | 38 ++++++++++++
 2 files changed, 67 insertions(+), 31 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
index 0b5eddbd..2b24434d 100644
--- a/applications/harnesses/KN_loderunner/ddp_production.csv
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -24,34 +24,32 @@ studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZ
 #21,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
 #22,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
 #23,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
-24,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
-# study 25: full 8-GPU node for ~8x faster wall-clock than 1 GPU. NTRN_BATCH
-# 1000->125 and NVAL_BATCH 500->63 keep total samples/epoch the same as study 24
-# (8 ranks x 125 = 1000). TERMINAL/WARMUP steps scaled ~8x down (counted in
-# per-rank batches) so the LR schedule spans the same number of epochs. These
-# are 8-GPU H100 nodes (se*); cpus-per-task=24 in the slurm template fits
-# 24*8=192 cores.
-25,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
-# study 40: Run A (n_rollout_steps=12). 1 GPU, BATCH_SIZE=5 keeps the 12-step
-# rollout under 80GB (5x12=60 activation units vs the OOM'd 10x12=120). Single
-# GPU avoids the multi-GPU Slurm GPU-binding issue on these H100 nodes.
-40,yoke311,1,1,128,1,1,9,1,2,5,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
-# study 41: 12 GPUs = 2 nodes x 6 GPUs/node (8-H100 se* nodes; under the 32-GPU
-# QOS cap). BATCH_SIZE=5 per GPU keeps the 12-step rollout under 80GB. NTRN/NVAL
-# 84/42 per rank hold ~1000/500 samples per epoch across 12 ranks (12*84~=1000).
-# TERMINAL/WARMUP scaled to per-rank batches (84/42) so the LR schedule spans the
-# same epochs. Uses the current train script: PREDICT_DELTA=True + DT_WEIGHT_TAU=3
-# + ZTF band weights bumped to 2 (A+B+ZTF combined, not isolated).
-#41,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
-# study 42: single node, all 8 GPUs (the max on these 8-H100 se* nodes). Replaces
-# the 2-node study 41, whose cross-node job kept hanging in Slurm step creation
-# ("Requested nodes are busy") -- a single-node job needs only one node free and
-# has no cross-node NCCL/IB teardown between the cycle_epochs=1 per-epoch jobs.
-# BATCH_SIZE=5 keeps the 12-step rollout under 80GB (same as study 40/41). NTRN/NVAL
-# 125/63 per rank hold ~1000/500 samples per epoch across 8 ranks (8*125=1000),
-# identical epoch size to studies 24/25/40/41. TERMINAL/WARMUP=125/63 (per-rank
-# batches) so the LR schedule spans the same epochs. Same model settings as 41:
-# PREDICT_DELTA + trend/decay anchor (now offset-capped) + EMA + ZTF band weights.
-42,yoke311,1,8,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
-
-
+#24,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#25,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#26,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#27,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#28,yoke311,1,1,128,1,1,9,1,2,10,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#29,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py 
+#30,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#31,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#32,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#33,yoke311,1,8,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#34,yoke311,1,1,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#35,yoke311,1,1,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#36,yoke311,1,4,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#37,yoke311,1,1,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#38,yoke311,1,1,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#39,yoke311,1,1,128,1,1,9,1,2,10,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#40,yoke311,1,2,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#41,yoke311,1,1,128,1,1,9,1,2,5,1000,500,2.0e-3,0.5,0.5,1000,500,0.0,train_LodeRunner_ddp.py
+#42,yoke311,1,2,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#43,yoke311,1,8,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#44,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
+#45,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
+#46,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
+#47,yoke311,2,6,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,84,42,0.0,train_LodeRunner_ddp.py
+#48,yoke311,1,8,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,125,63,0.0,train_LodeRunner_ddp.py
+#49,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
+#50,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
+#51,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
+52,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index 59a5dc37..6d854297 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -69,6 +69,44 @@ conda activate <YOKE_TORCH_ENV>
 # the dataloader workers (NUM_WORKERS per rank) alongside the OMP compute threads.
 export OMP_NUM_THREADS=8
 
+# ---------------------------------------------------------------------------
+# Fabric-ready gate. On these NVSwitch (multi-GPU NVLink) nodes, the NVIDIA
+# Fabric Manager must finish initializing the fabric before any process calls
+# CUDA on more than one GPU. nvidia-smi (above) only QUERIES the driver, so it
+# succeeds even when the fabric is not up yet -- and then the first
+# torch.cuda.set_device() in setup_distributed() fails with
+#   "Error 802: system not yet initialized"
+# on a perfectly healthy node. This is a launch-time RACE between FM readiness
+# and our CUDA init, seen only on multi-GPU jobs (single-GPU needs no fabric),
+# and it clears on resubmit. Rather than manually resubmit, wait here until each
+# allocated node reports the fabric ready (nvidia-smi -q "Fabric ... State:
+# Completed"), then launch. Falls through after a timeout so a node that never
+# reports (or a non-NVSwitch node with no Fabric section) can't hang the job --
+# the training srun then either succeeds or fails loudly as before.
+FABRIC_WAIT_TIMEOUT=${FABRIC_WAIT_TIMEOUT:-120}  # seconds, per node
+srun --ntasks=${SLURM_NNODES} --ntasks-per-node=1 --cpu-bind=none bash -c '
+  deadline=$(( SECONDS + '"${FABRIC_WAIT_TIMEOUT}"' ))
+  while true; do
+    # Grab the Fabric "State" lines from nvidia-smi. No Fabric section (older
+    # driver / non-NVSwitch) -> nothing to wait for, proceed immediately.
+    fabric=$(nvidia-smi -q 2>/dev/null | grep -A3 -i "^\s*Fabric" | grep -i "State")
+    if [ -z "$fabric" ]; then
+      echo "[$(hostname)] no Fabric section reported; proceeding."
+      exit 0
+    fi
+    # Ready when every reported State is Completed (no not-yet-Completed lines).
+    if ! echo "$fabric" | grep -viq "Completed"; then
+      echo "[$(hostname)] fabric ready: $fabric"
+      exit 0
+    fi
+    if [ $SECONDS -ge $deadline ]; then
+      echo "[$(hostname)] WARNING: fabric not ready after '"${FABRIC_WAIT_TIMEOUT}"'s; proceeding anyway. Last: $fabric"
+      exit 0
+    fi
+    sleep 3
+  done
+'
+
 # Get start time
 export date00=`date`
 

From 76f835c2896ee0982058c1ccf1ba501756b75110 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 28 Aug 2026 11:14:05 -0600
Subject: [PATCH 64/66] trying again

---
 .../KN_loderunner/training_slurm.tmpl         | 42 ++++++++++---------
 1 file changed, 22 insertions(+), 20 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index 6d854297..e0e9f394 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -72,35 +72,37 @@ export OMP_NUM_THREADS=8
 # ---------------------------------------------------------------------------
 # Fabric-ready gate. On these NVSwitch (multi-GPU NVLink) nodes, the NVIDIA
 # Fabric Manager must finish initializing the fabric before any process calls
-# CUDA on more than one GPU. nvidia-smi (above) only QUERIES the driver, so it
-# succeeds even when the fabric is not up yet -- and then the first
-# torch.cuda.set_device() in setup_distributed() fails with
+# CUDA on more than one GPU. If a rank calls torch.cuda.set_device() before FM
+# is up, it dies with
 #   "Error 802: system not yet initialized"
 # on a perfectly healthy node. This is a launch-time RACE between FM readiness
 # and our CUDA init, seen only on multi-GPU jobs (single-GPU needs no fabric),
-# and it clears on resubmit. Rather than manually resubmit, wait here until each
-# allocated node reports the fabric ready (nvidia-smi -q "Fabric ... State:
-# Completed"), then launch. Falls through after a timeout so a node that never
-# reports (or a non-NVSwitch node with no Fabric section) can't hang the job --
-# the training srun then either succeeds or fails loudly as before.
-FABRIC_WAIT_TIMEOUT=${FABRIC_WAIT_TIMEOUT:-120}  # seconds, per node
+# and it clears on resubmit.
+#
+# nvidia-smi is NOT a usable readiness probe here: it only queries the DRIVER,
+# which is up long before the fabric, so it "passes" while set_device() still
+# throws 802. So we probe the ACTUAL failing operation instead -- initialize a
+# real CUDA context on this node's assigned GPU -- and retry until it succeeds.
+# Each attempt runs in a FRESH `python -c` process: CUDA device state is cached
+# per-process, so a too-early attempt that saw 0 usable GPUs must be discarded
+# with the process or it would poison every later check. One probe per node
+# (fabric readiness is per-node), then launch. Falls through after a timeout so
+# the job can't hang forever -- the training srun then succeeds or fails loudly
+# exactly as before.
+FABRIC_WAIT_TIMEOUT=${FABRIC_WAIT_TIMEOUT:-180}  # seconds, per node
 srun --ntasks=${SLURM_NNODES} --ntasks-per-node=1 --cpu-bind=none bash -c '
   deadline=$(( SECONDS + '"${FABRIC_WAIT_TIMEOUT}"' ))
+  attempt=0
   while true; do
-    # Grab the Fabric "State" lines from nvidia-smi. No Fabric section (older
-    # driver / non-NVSwitch) -> nothing to wait for, proceed immediately.
-    fabric=$(nvidia-smi -q 2>/dev/null | grep -A3 -i "^\s*Fabric" | grep -i "State")
-    if [ -z "$fabric" ]; then
-      echo "[$(hostname)] no Fabric section reported; proceeding."
-      exit 0
-    fi
-    # Ready when every reported State is Completed (no not-yet-Completed lines).
-    if ! echo "$fabric" | grep -viq "Completed"; then
-      echo "[$(hostname)] fabric ready: $fabric"
+    attempt=$(( attempt + 1 ))
+    # Fresh process each time: force a real CUDA context (device 0 is the only
+    # GPU this task sees under per-task binding). Exit 0 only if init succeeds.
+    if python -c "import torch; torch.cuda.init(); torch.zeros(1, device=\"cuda:0\"); torch.cuda.synchronize()" >/dev/null 2>&1; then
+      echo "[$(hostname)] CUDA/fabric ready after ${attempt} attempt(s)."
       exit 0
     fi
     if [ $SECONDS -ge $deadline ]; then
-      echo "[$(hostname)] WARNING: fabric not ready after '"${FABRIC_WAIT_TIMEOUT}"'s; proceeding anyway. Last: $fabric"
+      echo "[$(hostname)] WARNING: CUDA not ready after '"${FABRIC_WAIT_TIMEOUT}"'s (${attempt} attempts); proceeding anyway."
       exit 0
     fi
     sleep 3

From 23a9ce38231e1a776b3b128143b1ace44e4507cc Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 28 Aug 2026 11:40:26 -0600
Subject: [PATCH 65/66] wip fix

---
 .../KN_loderunner/training_slurm.tmpl         | 95 ++++++++++---------
 1 file changed, 50 insertions(+), 45 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/training_slurm.tmpl b/applications/harnesses/KN_loderunner/training_slurm.tmpl
index e0e9f394..0a633cf4 100644
--- a/applications/harnesses/KN_loderunner/training_slurm.tmpl
+++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl
@@ -69,55 +69,60 @@ conda activate <YOKE_TORCH_ENV>
 # the dataloader workers (NUM_WORKERS per rank) alongside the OMP compute threads.
 export OMP_NUM_THREADS=8
 
+# Get start time
+export date00=`date`
+
 # ---------------------------------------------------------------------------
-# Fabric-ready gate. On these NVSwitch (multi-GPU NVLink) nodes, the NVIDIA
-# Fabric Manager must finish initializing the fabric before any process calls
-# CUDA on more than one GPU. If a rank calls torch.cuda.set_device() before FM
-# is up, it dies with
+# 802 retry loop. On these NVSwitch (multi-GPU NVLink) nodes the NVIDIA Fabric
+# Manager must finish initializing the fabric before a rank calls
+# torch.cuda.set_device(). If a rank gets there first it dies at CUDA init with
 #   "Error 802: system not yet initialized"
-# on a perfectly healthy node. This is a launch-time RACE between FM readiness
-# and our CUDA init, seen only on multi-GPU jobs (single-GPU needs no fabric),
-# and it clears on resubmit.
+# on a perfectly healthy node -- device_count() already succeeded, it's the
+# context-init-on-a-device step that loses the race. Seen only on multi-GPU
+# jobs (single-GPU needs no fabric); it clears on resubmit.
 #
-# nvidia-smi is NOT a usable readiness probe here: it only queries the DRIVER,
-# which is up long before the fabric, so it "passes" while set_device() still
-# throws 802. So we probe the ACTUAL failing operation instead -- initialize a
-# real CUDA context on this node's assigned GPU -- and retry until it succeeds.
-# Each attempt runs in a FRESH `python -c` process: CUDA device state is cached
-# per-process, so a too-early attempt that saw 0 usable GPUs must be discarded
-# with the process or it would poison every later check. One probe per node
-# (fabric readiness is per-node), then launch. Falls through after a timeout so
-# the job can't hang forever -- the training srun then succeeds or fails loudly
-# exactly as before.
-FABRIC_WAIT_TIMEOUT=${FABRIC_WAIT_TIMEOUT:-180}  # seconds, per node
-srun --ntasks=${SLURM_NNODES} --ntasks-per-node=1 --cpu-bind=none bash -c '
-  deadline=$(( SECONDS + '"${FABRIC_WAIT_TIMEOUT}"' ))
-  attempt=0
-  while true; do
-    attempt=$(( attempt + 1 ))
-    # Fresh process each time: force a real CUDA context (device 0 is the only
-    # GPU this task sees under per-task binding). Exit 0 only if init succeeds.
-    if python -c "import torch; torch.cuda.init(); torch.zeros(1, device=\"cuda:0\"); torch.cuda.synchronize()" >/dev/null 2>&1; then
-      echo "[$(hostname)] CUDA/fabric ready after ${attempt} attempt(s)."
-      exit 0
-    fi
-    if [ $SECONDS -ge $deadline ]; then
-      echo "[$(hostname)] WARNING: CUDA not ready after '"${FABRIC_WAIT_TIMEOUT}"'s (${attempt} attempts); proceeding anyway."
-      exit 0
+# We do NOT pre-probe with a separate srun: this job runs in all-visible mode
+# (each task sees all GPUs, ntasks-per-node=NGPUS), so a probe srun with a
+# different --ntasks shape can't get a step created under the implicit --exact
+# and silently no-ops (which is exactly why the earlier gate printed nothing).
+# Instead we retry the TRAINING srun itself. A relaunch reuses this job's exact
+# allocation (no step-shape mismatch) and starts fresh processes for every rank
+# (which also clears the per-process CUDA error state that 802 leaves behind).
+# We only retry when the run's output shows the 802 signature -- any other
+# non-zero exit is a real failure and surfaces immediately instead of looping.
+MAX_802_RETRIES=${MAX_802_RETRIES:-5}
+RUN_LOG=study<studyIDX>_epoch<epochIDX>.runlog
+attempt=0
+while true; do
+  attempt=$(( attempt + 1 ))
+
+  # Start the Code
+  # Default Slurm per-task GPU binding is correct here: each task sees exactly
+  # one GPU (as cuda:0), and setup_distributed() selects that device. Do NOT add
+  # --gpu-bind=none -- it exposes partial GPU subsets per task and causes ranks
+  # to collide on the same device ("Duplicate GPU detected" in NCCL).
+  # tee: keep the live .out/.err stream AND capture a copy to scan for 802.
+  srun -vv --cpu-bind=verbose python -u <train_script> @<INPUTFILE> 2>&1 | tee "$RUN_LOG"
+  rc=${PIPESTATUS[0]}
+
+  if [ $rc -eq 0 ]; then
+    echo "[802-retry] training srun succeeded on attempt ${attempt}."
+    break
+  fi
+
+  if grep -q "Error 802: system not yet initialized" "$RUN_LOG"; then
+    if [ $attempt -ge $MAX_802_RETRIES ]; then
+      echo "[802-retry] hit Error 802 on attempt ${attempt} (max ${MAX_802_RETRIES}); giving up."
+      exit $rc
     fi
-    sleep 3
-  done
-'
-
-# Get start time
-export date00=`date`
-
-# Start the Code
-# Default Slurm per-task GPU binding is correct here: each task sees exactly one
-# GPU (as cuda:0), and setup_distributed() selects that device. Do NOT add
-# --gpu-bind=none -- it exposes partial GPU subsets per task and causes ranks to
-# collide on the same device ("Duplicate GPU detected" in NCCL).
-srun -vv --cpu-bind=verbose python -u <train_script> @<INPUTFILE>
+    echo "[802-retry] Error 802 (fabric not ready) on attempt ${attempt}; relaunching in 15s."
+    sleep 15
+    continue
+  fi
+
+  echo "[802-retry] training srun failed (rc=${rc}) with no Error 802 signature; not a fabric race, exiting."
+  exit $rc
+done
 
 # Get end time and print to stdout
 export date01=`date`

From ca19dc94b5e6ce85ddd811973e050fbdf4b30604 Mon Sep 17 00:00:00 2001
From: Andrew Michael Toivonen <atoivonen@pn2405823.lanl.gov>
Date: Fri, 28 Aug 2026 15:03:07 -0600
Subject: [PATCH 66/66] trying to reproduce 44

---
 .../harnesses/KN_loderunner/ddp_production.csv |  9 ++++++++-
 .../KN_loderunner/train_LodeRunner_ddp.py      | 18 +++++++++---------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv
index 2b24434d..a86aaca6 100644
--- a/applications/harnesses/KN_loderunner/ddp_production.csv
+++ b/applications/harnesses/KN_loderunner/ddp_production.csv
@@ -52,4 +52,11 @@ studyIDX,YOKE_TORCH_ENV,KNODES,NGPUS,EMBED_DIM,B0,B1,B2,B3,NUM_WORKERS,BATCH_SIZ
 #49,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
 #50,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
 #51,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
-52,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
+#52,yoke311,1,1,128,1,1,9,1,2,5,125,63,2.0e-3,0.5,0.5,100,50,0.0,train_LodeRunner_ddp.py
+# Study 56: reproduce study 54's DESCENDING schedule on current KN code.
+# 1 node x 1 GPU (no NVSwitch fabric -> no Error 802), NTRN_BATCH=84, and the
+# long schedule TERMINAL_STEPS=500/WARMUP=100 that let study 54's validation
+# descend 1.0->0.3 (study 55 at 84/42 never left warmup and looked flat).
+# Model config matches study 44: PREDICT_DELTA=True, TREND_DECAY_ANCHOR=False,
+# Huber delta=0.1. Once this descends, layer features back one at a time.
+56,yoke311,1,1,128,1,1,9,1,2,5,84,42,2.0e-3,0.5,0.5,500,100,0.0,train_LodeRunner_ddp.py
diff --git a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
index 131e006f..0842bdae 100644
--- a/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
+++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py
@@ -491,15 +491,15 @@ def main(args, rank, world_size, local_rank, device):
         )
 
     #loss_fn = nn.MSELoss(reduction="none")
-    # delta=0.1 was in place since the harness was created (studies 24/25 too), so
-    # it is NOT what newly flattened the loss -- but it does throttle descent: with
-    # z-scored targets (sigma~=1) and a mean loss ~0.09, typical errors are ~0.95
-    # sigma, deep in Huber's LINEAR regime where the per-point gradient is capped
-    # at delta=0.1. Raising delta to 1.0 gives an error-proportional gradient (~10x
-    # stronger here) so the model can actually push down the bulk error, and only
-    # reverts to the robust linear regime past 1 sigma. delta=1.0 ~= MSE for the
-    # in-sigma bulk while still clipping the heavy-tailed outliers.
-    loss_fn = nn.HuberLoss(delta=1.0, reduction="none")
+    # delta=0.1 to MATCH study 44's config exactly (the RMSE-1.52 run we are
+    # reproducing on current code). The flat-loss investigation cleared the model
+    # code -- baseline == study-44 config still trained -- and showed the real
+    # descent driver is the SCHEDULE (study 54's TERMINAL_STEPS=500/WARMUP=100
+    # descends; study 55's 84/42 stays in warmup and looks flat). Once a clean
+    # study-44 reproduction is in hand, delta=1.0 is a deliberate next experiment
+    # (error-proportional gradient ~10x stronger for the in-sigma bulk); it also
+    # rescales the loss magnitude, so do NOT compare curve heights across the two.
+    loss_fn = nn.HuberLoss(delta=0.1, reduction="none")
     model = DDP(model, device_ids=[local_rank], output_device=local_rank)
 
     #############################################