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/applications/harnesses/KN_loderunner/ddp_production.csv b/applications/harnesses/KN_loderunner/ddp_production.csv new file mode 100644 index 00000000..a86aaca6 --- /dev/null +++ b/applications/harnesses/KN_loderunner/ddp_production.csv @@ -0,0 +1,62 @@ +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 +#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 +# 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/eval_dense_latetime_9band.py b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py new file mode 100644 index 00000000..f2a045bb --- /dev/null +++ b/applications/harnesses/KN_loderunner/eval_dense_latetime_9band.py @@ -0,0 +1,789 @@ +"""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 _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 _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, + model, + device, + means, + stds, + context_window_days, + max_context_len, + 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. + + 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. + + 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 + + if r_t.shape[0] < 1 or d_t.shape[0] < 1: + return None + + # Phase zero = the first realistic detection (the observed trigger). + t0 = float(r_t[0]) + + # 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]) + + # 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 + + # 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_ctx.astype(np.float32)), + ctx_v=list(r_v_norm.astype(np.float32)), + ctx_b=list(r_b_ctx), + 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. 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 + + 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 + # single batched forward pass. + max_dt = float(lead_times.max()) + lead_grid = np.linspace(0.0, max_dt, 60).astype(np.float32) + 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, + "last_real_t": last_real_t, + "curve_phase": (last_real_t - t0) + lead_grid, + "curve_mag": curve_mag.astype(np.float32), + # 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), + "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, + } + + +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"] + 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") + + 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") + # 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], + 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 + 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") + + 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( + "--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, + default=( + "/net/sescratch1/atoivonen/data/KN_lightcurves/" + "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, + 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( + "--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, + 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=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. " + "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", + 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", + 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.", + ) + 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() + + +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, use_ema=getattr(args, "use_ema", False)) + + 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)) + + # 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()} + 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) + 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, + model, + device, + means, + stds, + context_window_days, + max_context_len, + args.late_time_cutoff_days, + args.late_time_max_days, + uniform_stream=uniform_stream, + rollout=args.rollout, + ) + 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]) + 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}") + 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/infer_9band.py b/applications/harnesses/KN_loderunner/infer_9band.py new file mode 100644 index 00000000..9bfdafdb --- /dev/null +++ b/applications/harnesses/KN_loderunner/infer_9band.py @@ -0,0 +1,624 @@ +"""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, +) +from yoke.utils.ema import ParamEMA + + +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 +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}" + + +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( + "--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, + 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_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() + + +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/" + "rubin_ztf_10000_dataset_same_seed/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, use_ema: bool = False): + 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) + + # 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 + + # 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) + # 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) + trend_max_offset = ckpt.get("trend_max_offset", None) + + 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) + 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 + + model = ScalarTemporalConditionedLodeRunner_9band( + backbone=backbone, + 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, + dt_fourier_bands=dt_fourier_bands, + 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"]) + missing, unexpected = model.load_state_dict(state_dict, strict=True) + + 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.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.99)}).") + else: + print("Using raw (non-EMA) weights.") + + model.eval() + + return model, context_len, n_bands, context_window_days, max_context_len + + +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 + + # 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) + + 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, window_mode=False, max_context_len=None +): + """Build the flattened per-event context input for the model. + + 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. + """ + 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) + + 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), + dtype=torch.float32, + device=device, + ).unsqueeze(0) + + return x + + +def forecast_curve( + 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. + + 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 + + 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]) + + 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, + context_window_days, + max_context_len, + ) = load_9band_model(args.ckpt, device, use_ema=getattr(args, "use_ema", False)) + + window_mode = context_window_days is not None + + 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) + + files = sorted(glob.glob(args.data_glob)) + 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.") + + 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] + # 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 {min_events}." + ) + 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, + 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] + + 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/make_kn_object_lists.py b/applications/harnesses/KN_loderunner/make_kn_object_lists.py new file mode 100644 index 00000000..42b1a4f0 --- /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_``). 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=( + "/net/sescratch1/atoivonen/data/KN_lightcurves/" + "rubin_ztf_10000_dataset_same_seed/lc_*.npz" + ), + help="Glob for the realistic light-curve files.", + ) + parser.add_argument( + "--dense_glob", + type=str, + default=( + "/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).", + ) + 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/plot_loss_curves_9band.py b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py new file mode 100644 index 00000000..d896e08f --- /dev/null +++ b/applications/harnesses/KN_loderunner/plot_loss_curves_9band.py @@ -0,0 +1,155 @@ +"""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_epoch.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( + "--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", + help=( + "Fail instead of continuing when validation records are missing " + "or invalid." + ), + ) + + # 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=8, + help="Absolute epoch at which the teacher-forcing anneal begins. Must " + "match --tf_ramp_start_epoch in training_START.input / " + "training_input.tmpl. Default 8.", + ) + parser.add_argument( + "--tf_ramp_epochs", + type=int, + default=12, + help="Number of epochs the teacher-forcing ratio anneals over. Must " + "match --tf_ramp_epochs in training_START.input / training_input.tmpl. " + "Default 12.", + ) + + 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/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..69bf32b4 --- /dev/null +++ b/applications/harnesses/KN_loderunner/plot_loss_curves_gri.py @@ -0,0 +1,452 @@ +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 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: + 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 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"], + train["losses"], + ) + + 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"], + 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 + 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}" + + (train_line,) = 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 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: + (val_line,) = 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) + + 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) + plt.grid(True, alpha=0.3) + plt.legend() + + 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}") + + +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( + "--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", + 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_observation_histograms.py b/applications/harnesses/KN_loderunner/plot_observation_histograms.py new file mode 100644 index 00000000..d5f9abd1 --- /dev/null +++ b/applications/harnesses/KN_loderunner/plot_observation_histograms.py @@ -0,0 +1,701 @@ +"""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. + 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. + 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 + 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 +""" + +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. + 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) + 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 = [] + event_times_per_file = [] + + 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 + file_det_times = [] + 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 + + # 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 { + "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), + "event_times_per_file": event_times_per_file, + "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 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 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 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"] + 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 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 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=( + "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( + "--sweep_max", + type=int, + default=20, + help=( + "Largest context_len to evaluate in the data-cost sweep panel. " + "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( + "--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, + 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) + print_sweep(counts["total_det_per_curve"], args.sweep_max) + + # 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) + 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) + plot_total_hist(axes[0, 1], counts) + 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) + 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)", + fontsize=13, + ) + fig.tight_layout(rect=(0, 0, 1, 0.97)) + 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/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_9band.py b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py new file mode 100644 index 00000000..cf3d7bfe --- /dev/null +++ b/applications/harnesses/KN_loderunner/plot_pred_diagnostics_9band.py @@ -0,0 +1,1135 @@ +"""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. + +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 +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 ( + EPS, + NINE_BAND_KEYS, + Kilonova_lc_scalar_context_DataSet_9band, + load_or_compute_band_normalization, +) +from yoke.utils.ema import ParamEMA + + +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 +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}" + + +def get_args(): + parser = argparse.ArgumentParser( + description=( + "Autoregressive rollout diagnostics for the scalar temporal " + "LodeRunner 9-band model." + ) + ) + + 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( + "--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, + default=50, + help="Number of light-curve files to load into the eval dataset.", + ) + parser.add_argument( + "--n_future_steps", + type=int, + default=15, + help="Number of future events to forecast autoregressively per series.", + ) + parser.add_argument( + "--n_series", + type=int, + 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( + "--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, + 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() + + +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_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, use_ema: bool = False): + 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) + + # 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 + + # 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) + # 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) + trend_max_offset = ckpt.get("trend_max_offset", None) + + 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) + print("hidden:", hidden) + 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 + + model = ScalarTemporalConditionedLodeRunner_9band( + backbone=backbone, + 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, + dt_fourier_bands=dt_fourier_bands, + 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"]) + + missing, unexpected = model.load_state_dict(state_dict, strict=True) + + print("Loaded ScalarTemporalConditionedLodeRunner_9band checkpoint") + 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.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.99)}).") + else: + print("Using raw (non-EMA) weights.") + + model.eval() + + return model, context_len, n_bands, context_window_days, max_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, + value_col=VALUE_COL, + error_col=ERROR_COL, + drop_upper_limits=DROP_UPPER_LIMITS, + ) + + print("Using band normalization:") + 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, + band_keys=BAND_KEYS, + value_col=VALUE_COL, + error_col=ERROR_COL, + drop_upper_limits=DROP_UPPER_LIMITS, + means=band_means, + 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) + + +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. + + 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) + + if window_mode: + n_real = win_v.shape[0] + rel_t = (win_t - win_t[0]).astype(np.float32) + + 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) + + 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), + dtype=torch.float32, + device=device, + ).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. + + 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, + bands, + model, + device, + start_idx, + n_future_steps, + context_len, + n_bands, + means, + stds, + teacher_forced=False, + 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. + + 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. + 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). + """ + 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 + 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, + "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 step in range(n_future_steps): + target_idx = start_idx + seed_len + step + + if target_idx >= len(times): + break + + 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=build_width, + n_bands=n_bands, + device=device, + window_mode=window_mode, + ) + + # 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, + ) + + 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]) + true_norm = float(values[target_idx]) + residual = pred_norm - true_norm + + pred_mag = pred_norm * (stds[target_band] + EPS) + means[target_band] + true_mag = true_norm * (stds[target_band] + EPS) + means[target_band] + + band_sq_err[target_band].append(residual**2) + + 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, + # Full all-band prediction at this step's lead time. + "pred_all_norm": pred_all_norm, + "pred_all_mag": pred_all_mag, + } + ) + + # 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(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) + total_mse = float(np.mean(residuals**2)) if len(residuals) else np.nan + + 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])) + + # 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 + 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=build_width0, + n_bands=n_bands, + device=device, + window_mode=window_mode, + ) + + 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) + + # 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) + + # 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, :] + + fixed_forecast = { + "t_rel": last_ctx_t_rel + lead_times, + "preds_mag": preds_mag.astype(np.float32), + } + + return { + "start_idx": start_idx, + "context": context, + "steps": steps, + "mse": total_mse, + "band_mse": band_mse, + "n_steps": len(steps), + "fixed_forecast": fixed_forecast, + } + + +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. + """ + # 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: + 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 _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 + fixed_forecast = rollout.get("fixed_forecast") + + 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", + ) + + # 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.scatter( + t_all, + pred_all, + s=22, + facecolors="none", + edgecolors="k", + alpha=0.6, + label="per-step pred (free-run)", + ) + + # Teacher-forced per-step predictions, also as unconnected markers. + if tf_steps: + t_tf, pred_tf = _band_forecast_curve(tf_steps, band_idx) + ax.scatter( + t_tf, + pred_tf, + 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. + 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]) + + order = np.argsort(t) + t = t[order] + true_mag = true_mag[order] + + ax.plot( + t, true_mag, "-o", color=color, alpha=0.8, label="truth (obs)" + ) + + 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_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 _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)] + ) + + +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") + 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( + ["series", "start_idx", "n_steps", "mse_total"] + + [f"mse_{n}" for n in BAND_NAMES] + ) + for i, rollout in enumerate(rollouts): + writer.writerow( + [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", + ] + # 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"]: + 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"], + ] + + [float(s["pred_all_mag"][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, + context_window_days, + max_context_len, + ) = 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 + + 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)) + + series = select_series( + dataset=eval_dataset, + context_len=context_len, + n_future_steps=args.n_future_steps, + n_series=args.n_series, + seed_len=seed_len, + ) + + 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.") + + 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, + 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=False, + 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) + + 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, + window_mode=window_mode, + context_window_days=context_window_days, + max_context_len=max_context_len, + ) + 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): + series_path = os.path.join( + args.outdir, + f"study{run_id}_9band_autoreg_series{i:02d}.png", + ) + 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" + ) + mse_bar_path = os.path.join( + 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" + ) + step_csv_path = os.path.join( + args.outdir, f"study{run_id}_9band_autoreg_step_predictions.csv" + ) + + plot_residuals_vs_step(rollouts, residual_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) + print(" ", mse_bar_path) + print(" ", mse_csv_path) + print(" ", step_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..0842bdae --- /dev/null +++ b/applications/harnesses/KN_loderunner/train_LodeRunner_ddp.py @@ -0,0 +1,922 @@ +import os +import glob +import time +import random +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 torch.utils.data import ConcatDataset + +from yoke.models.vit.swin.bomberman import ( + LodeRunner, + ScalarTemporalConditionedLodeRunner_9band, +) +from yoke.datasets.kilonova_dataset import ( + 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_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 +from yoke.utils.checkpointing import save_model_and_optimizer +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 + +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 +############################################# +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).", +) + +# 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.", +) +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.", +) + +# 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.99, + help="EMA decay for the trainable-parameter shadow (Polyak averaging), " + "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.", +) + +# 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=( + "/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).", +) +parser.add_argument( + "--kn_dense_glob", + type=str, + default=( + "/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 " + "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="kn_rubin_ztf_train.txt", + validation_filelist="kn_rubin_ztf_val.txt", + test_filelist="kn_rubin_ztf_test.txt", +) + + +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. 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 + 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 + + # 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 + + # 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 + + # 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. + # + # 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 + # (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 + # 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 + # 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 + + # 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 + ) + + # 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) + 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 + 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 + + # 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). 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( + [ + 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) + 1.0, # ps1__i (i) + 1.0, # ps1__z (z) + 1.0, # ps1__y (y) + ], + 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), + "eps": 1e-08, + "weight_decay": 0.01, + } + + + if CONTINUATION: + model, optimizer, starting_epoch = load_direct_loderunner_checkpoint_9band( + 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) + + # 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=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, + dt_fourier_bands=DT_FOURIER_BANDS, + 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 + 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") + # 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) + + ############################################# + # 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 + ############################################# + 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, + ) + ''' + + # 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( + stats_path=norm_stats_path, + band_keys=BAND_KEYS, + value_col=VALUE_COL, + error_col=ERROR_COL, + drop_upper_limits=DROP_UPPER_LIMITS, + file_prefix_list=train_norm_files, + ) + + 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) + + # 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) + + 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, + ) + + # 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( + 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() + + + 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. 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: + 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, + window_mode=CONTEXT_WINDOW_DAYS is not None, + context_window_days=CONTEXT_WINDOW_DAYS, + max_context_len=MAX_CONTEXT_LEN, + band_weights=BAND_WEIGHTS, + dt_weight_tau=DT_WEIGHT_TAU, + ema=ema, + ) + 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, + band_weights=BAND_WEIGHTS, + ema=ema, + ) + + 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_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": False, + "target_type": "value_9band", + "context_len": CONTEXT_LEN, + "n_bands": N_BANDS, + "band_keys": list(BAND_KEYS), + "backbone_channels": 8, + "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, + "trend_max_offset": TREND_MAX_OFFSET, + "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() + 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, + "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, + ) + + 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..9ba19892 --- /dev/null +++ b/applications/harnesses/KN_loderunner/training_START.input @@ -0,0 +1,63 @@ +--studyIDX + +--FILELIST_DIR +/net/sescratch1/atoivonen/projects/filelists/ +--LSC_NPZ_DIR +/lustre/scratch5/exempt/artimis/data/lsc240420/ +--train_filelist +kn_rubin_ztf_train.txt +--validation_filelist +kn_rubin_ztf_val.txt +--block_structure + + + + +--embed_dim + +--anchor_lr + +--num_cycles + +--min_fraction + +--terminal_steps + +--warmup_steps + +--noise_scale + +--n_rollout_steps +12 +--tf_start +1.0 +--tf_end +0.0 +--tf_ramp_start_epoch +8 +--tf_ramp_epochs +12 +--trn_rcrd_filename +./training_study_epoch.csv +--val_rcrd_filename +./validation_study_epoch.csv +--batch_size + +--num_workers + +--Ngpus + +--Knodes + +--total_epochs +40 +--cycle_epochs +1 +--train_batches + +--val_batches + +--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..3c845428 --- /dev/null +++ b/applications/harnesses/KN_loderunner/training_START.slurm @@ -0,0 +1,81 @@ +#!/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_e0001 +#SBATCH --account=y26_artimis_ddc_g +#SBATCH --time=1:00:00 +#SBATCH --partition=standard +#SBATCH --nodes= +#SBATCH --ntasks-per-node= +#SBATCH --gpus-per-node= +# 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_epoch0001.out +#SBATCH --error=study_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 + +# 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` + +# Start the Code +# Explicitly set TCP environment for the following... +srun -vv --cpu-bind=verbose python -u @study_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..c68d90a3 --- /dev/null +++ b/applications/harnesses/KN_loderunner/training_input.tmpl @@ -0,0 +1,66 @@ +--studyIDX + +--FILELIST_DIR +/net/sescratch1/atoivonen/projects/filelists/ +--LSC_NPZ_DIR +/lustre/scratch5/exempt/artimis/data/lsc240420/ +--train_filelist +kn_rubin_ztf_train.txt +--validation_filelist +kn_rubin_ztf_val.txt +--block_structure + + + + +--embed_dim + +--anchor_lr + +--num_cycles + +--min_fraction + +--terminal_steps + +--warmup_steps + +--noise_scale + +--n_rollout_steps +12 +--tf_start +1.0 +--tf_end +0.0 +--tf_ramp_start_epoch +8 +--tf_ramp_epochs +12 +--trn_rcrd_filename +./training_study_epoch.csv +--val_rcrd_filename +./validation_study_epoch.csv +--batch_size + +--num_workers + +--Ngpus + +--Knodes + +--total_epochs +40 +--cycle_epochs +1 +--train_batches + +--val_batches + +--TRAIN_PER_VAL +5 +--continuation +--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..0a633cf4 --- /dev/null +++ b/applications/harnesses/KN_loderunner/training_slurm.tmpl @@ -0,0 +1,133 @@ +#!/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_e +#SBATCH --account=y26_artimis_ddc_g +#SBATCH --time=1:00:00 +#SBATCH --partition=standard +#SBATCH --nodes= +#SBATCH --ntasks-per-node= +#SBATCH --gpus-per-node= +# 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_epoch.out +#SBATCH --error=study_epoch.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 + +# 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 + +# Get start time +export date00=`date` + +# --------------------------------------------------------------------------- +# 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 -- 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. +# +# 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_epoch.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 @ 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 + 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` + +echo "===================TIME STARTED===================" +echo $date00 +echo "===================TIME FINISHED===================" +echo $date01 diff --git a/src/yoke/datasets/kilonova_dataset.py b/src/yoke/datasets/kilonova_dataset.py new file mode 100644 index 00000000..74ce5d3c --- /dev/null +++ b/src/yoke/datasets/kilonova_dataset.py @@ -0,0 +1,1082 @@ +"""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 + +# 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 _stem(path: str) -> str: + """Return the object identifier (filename without directory or extension). + + Light-curve files are named ``lc_.npz``; the stem ``lc_`` 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"), + 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. + + 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. + 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: + 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) + + # 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() + 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, + 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. + + 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. + 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. + 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]. + """ + 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( + "/net/sescratch1/atoivonen/data/KN_lightcurves/" + "rubin_ztf_10000_dataset_same_seed/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, + error_col=error_col, + drop_upper_limits=drop_upper_limits, + 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 + + +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, + error_col: int = 2, + drop_upper_limits: bool = True, + means: np.ndarray = None, + stds: np.ndarray = None, + 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. + + Args: + N_imgs (int): Number of light-curve files to sample; 0 uses all. + 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. + 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]. + 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__``). 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. + 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``. + """ + # 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. + if data_glob is None: + # Legacy hardcoded scratch fallback (backward compatibility). + data_glob = ( + "/net/sescratch1/atoivonen/data/KN_lightcurves/" + "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 + 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.error_col = error_col + 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 + + # 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}" + ) + 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. " + "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 = [] + # 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): + 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 + + # 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)) + + 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) + ) + self.stems_per_file.append(_stem(fn)) + + n_events = times.shape[0] + 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.""" + return len(self.samples) + + 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.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: + 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. + + 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 + + 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]: + """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 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 + + # 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: + 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) + + # 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[anchor_idx], + dtype=torch.float32, + ) + + 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), + ) + + 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) + + # 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): + 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[prev_idx] + future_valid[step] = 1.0 + prev_idx = t_idx + + 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/models/vit/swin/bomberman.py b/src/yoke/models/vit/swin/bomberman.py index e3c6af63..040e56d2 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 @@ -332,6 +333,528 @@ 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 + + +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. + 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__( + 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, + context_window_days: float = None, + dt_fourier_bands: int = 0, + 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. + + 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. + 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 + 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). + 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. + 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. + 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__() + + 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." + ) + + 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 + self.image_size = image_size + self.backbone_channels = backbone_channels + 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 + 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 + # 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 + + # Fourier lead-time encoding. When enabled, a fixed log-spaced frequency + # 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(60.0), dt_fourier_bands + ) + self.register_buffer("dt_freqs", 2.0 * math.pi / periods) + # 2 * bands (sin + cos) + 1 monotone log1p(Dt) channel. + dt_extra = 2 * dt_fourier_bands + 1 + 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 + dt_extra, hidden), + nn.GELU(), + nn.Linear(hidden, hidden), + nn.GELU(), + nn.Linear(hidden, backbone_channels), + ) + + # 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 + 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 + 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 * 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 _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 + ``[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). + + 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). 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 + [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. + """ + 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. + 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) + + # 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, + 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 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]. + """ + B = x.shape[0] + H, W = self.image_size + + # 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, + 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, 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] + + # 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, Dt) + + return pred + + if __name__ == "__main__": from yoke.utils.parameters import count_torch_params diff --git a/src/yoke/utils/checkpointing.py b/src/yoke/utils/checkpointing.py index 5ca580f7..cd5305ad 100644 --- a/src/yoke/utils/checkpointing.py +++ b/src/yoke/utils/checkpointing.py @@ -10,6 +10,12 @@ import torch.distributed as dist import h5py +from yoke.models.vit.swin.bomberman import ( + LodeRunner, + ScalarTemporalConditionedLodeRunner_gri, + ScalarTemporalConditionedLodeRunner_9band, +) + def save_model_and_optimizer_hdf5( model: torch.nn.Module, @@ -303,3 +309,281 @@ 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 + + +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) + + # 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( + 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), + 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), + # 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), + trend_max_offset=checkpoint_data.get("trend_max_offset", None), + ).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 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)) diff --git a/src/yoke/utils/parallel.py b/src/yoke/utils/parallel.py index 66cbb375..ed3217c1 100644 --- a/src/yoke/utils/parallel.py +++ b/src/yoke/utils/parallel.py @@ -5,8 +5,84 @@ """ +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. + 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 ----- + # 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 ----- + # 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} gpu_index={gpu_index} " + f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}", + flush=True, + ) + if gpu_index >= n_visible: + raise RuntimeError( + 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(gpu_index) + device = torch.device(f"cuda:{gpu_index}") + + # ----- 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 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: + """Destroy the DDP process group.""" + dist.destroy_process_group() # Custom nn.DataParallel class to handle input to LodeRunner that should not be diff --git a/src/yoke/utils/training/datastep/loderunner.py b/src/yoke/utils/training/datastep/loderunner.py index 600996ab..470b71ba 100644 --- a/src/yoke/utils/training/datastep/loderunner.py +++ b/src/yoke/utils/training/datastep/loderunner.py @@ -10,40 +10,267 @@ import random -def train_loderunner_datastep( +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 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 +#################################### +def eval_loderunner_datastep( + data: tuple, + model: torch.nn.Module, + 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. + """An evaluation 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 + model (torch.nn.Module): model to evaluate + loss_fn (torch.nn.Module): loss function for evaluation device (torch.device): device index to select - channel_map (list): list of channel indices to use + channel_map (list[int]): 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() + model.eval() # 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) + Dt = Dt.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 @@ -75,15 +302,93 @@ def train_loderunner_datastep( 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( +def eval_DDP_loderunner_seq_channel_datastep( + data, + model, + loss_fn, + device, + rank, + world_size, +): + """DDP eval step for channel-stacked LodeRunner on a single next-step target.""" + 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, +): + """DDP eval step rolling LodeRunner over a full image sequence.""" + 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, optimizer: torch.optim.Optimizer, @@ -91,10 +396,7 @@ def train_scheduled_loderunner_datastep( 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. + """Evaluation step for LodeRunner with scheduled sampling. Args: data (tuple): Sequence of images in (img_seq, Dt) tuple. @@ -109,8 +411,8 @@ def train_scheduled_loderunner_datastep( pred_seq (torch.Tensor): Predicted image sequence. per_sample_loss (torch.Tensor): Per-sample loss for the batch. """ - # Set model to train - model.train() + # Set model to evaluation + model.eval() # Extract data img_seq, Dt = data @@ -153,41 +455,34 @@ def train_scheduled_loderunner_datastep( 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( +def eval_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. + """A DDP-compatible evaluation step. 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 + world_size (int): Total number of 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() + # Set model to evaluation mode + model.eval() # Extract data start_img, end_img, Dt = data @@ -200,17 +495,13 @@ def train_DDP_loderunner_datastep( 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) + with torch.no_grad(): + 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) @@ -224,42 +515,252 @@ def train_DDP_loderunner_datastep( return end_img, pred_img, all_losses -#################################### -# Evaluating on a Datastep -#################################### -def eval_loderunner_datastep( +def train_DDP_temporal_loderunner_datastep( + data, + model, + optimizer, + loss_fn, + device, + rank, + world_size, +): + """DDP training step for temporal LodeRunner on a channel-stacked context.""" + 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 train_DDP_scalar_temporal_loderunner_datastep( + data, + model, + optimizer, + loss_fn, + device, + rank, + world_size, +): + """DDP training step for scalar temporal LodeRunner on flattened context.""" + 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, + loss_fn, + device, + rank, + world_size, +): + """DDP eval step for temporal LodeRunner on a channel-stacked context.""" + 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_scalar_temporal_loderunner_datastep( + data, + model, + loss_fn, + device, + 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] + + 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, + 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] + """ + 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 + + +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]: - """An evaluation step for which the data is of multi-input, multi-output type. + """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 evaluate - loss_fn (torch.nn.Module): loss function for evaluation + 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[int]): list of channel indices to use + 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.eval() + 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) + 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 @@ -291,10 +792,15 @@ def eval_loderunner_datastep( 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 eval_scheduled_loderunner_datastep( +def train_scheduled_loderunner_datastep( data: tuple, model: torch.nn.Module, optimizer: torch.optim.Optimizer, @@ -302,7 +808,10 @@ def eval_scheduled_loderunner_datastep( device: torch.device, scheduled_prob: float, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Evaluation step for LodeRunner with scheduled sampling. + """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. @@ -317,8 +826,8 @@ def eval_scheduled_loderunner_datastep( pred_seq (torch.Tensor): Predicted image sequence. per_sample_loss (torch.Tensor): Per-sample loss for the batch. """ - # Set model to evaluation - model.eval() + # Set model to train + model.train() # Extract data img_seq, Dt = data @@ -361,34 +870,41 @@ def eval_scheduled_loderunner_datastep( 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 eval_DDP_loderunner_datastep( +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 evaluation step. + """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): Total number of DDP processes + 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 evaluation mode - model.eval() + # Set model to train mode + model.train() # Extract data start_img, end_img, Dt = data @@ -401,13 +917,17 @@ def eval_DDP_loderunner_datastep( out_vars = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]).to(device, non_blocking=True) # Forward pass - with torch.no_grad(): - pred_img = model(start_img, in_vars, out_vars, Dt) + 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) @@ -419,3 +939,8 @@ def eval_DDP_loderunner_datastep( 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 26629e4e..9d12c641 100644 --- a/src/yoke/utils/training/epoch/loderunner.py +++ b/src/yoke/utils/training/epoch/loderunner.py @@ -330,6 +330,899 @@ 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( + "", + 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( + "", + 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_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, + band_weights: torch.Tensor = None, + ema: object = None, +) -> 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] + + ``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. + + ``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( + "", + f"{epochIDX:04d}", + ) + + 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: + + 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) + + 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() + 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( + [ + 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( + "", + 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 _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, + band_weights: torch.Tensor = None, +) -> 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. + 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 + (per-band weighted when ``band_weights`` is given). + """ + 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 + 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 + # 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) + 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 + ) + 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) + + 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 + + +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, + 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. + + 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. + 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. + 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]. + 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] + 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 + 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). + 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) + 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 + 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) + + if band_weights is None and dt_weight_tau is None: + total_loss = step_losses.sum() / (step_valid.sum() + 1e-8) + else: + # 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 + + +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, + window_mode: bool = False, + context_window_days: float = None, + 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. + + 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``): + 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. + 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. + 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. + 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, ...], + 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: + ( + 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, + band_weights=weights, + dt_weight_tau=dt_tau, + ) + + 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, + band_weights=weights, + ) + + train_rcrd_filename = train_rcrd_filename.replace( + "", + 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 + + optimizer.zero_grad(set_to_none=True) + + per_sample_loss, batch_loss = _run_pass( + data, teacher_forcing_ratio, weights=band_weights, dt_tau=dt_weight_tau + ) + + batch_loss.backward() + 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( + [ + 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( + "", + 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 + + # Validation is always a pure free-running rollout. + per_sample_loss, _ = _run_pass(data, 0.0) + + 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,