Cross_day_lin.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Wiener baseline cross-day evaluation with two training methods:
  - GD  : PyTorch gradient descent
  - OLS : closed-form least squares (two-stage Wiener-cascade)

Outputs (same schema style as your crossday script):
  crossday_results_Wiener_<DIMRED>_<TRAINMETHOD>.pkl

Examples:
  python wiener_crossday_gd_vs_ols.py --train_method OLS --dimred PCA --n_folds 5
  python wiener_crossday_gd_vs_ols.py --train_method GD  --dimred PCA --n_folds 5
"""

import os
import argparse
import warnings
import numpy as np
import pandas as pd

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader

from scipy.signal import butter, filtfilt
from scipy.ndimage import gaussian_filter1d
from numpy.linalg import pinv
from sklearn.decomposition import PCA

# UMAP (optional)
try:
    import umap
except Exception:
    umap = None

warnings.filterwarnings("ignore", message="n_jobs value 1 overridden to 1 by setting random_state.")

# ============================== GLOBAL CONFIG ==============================
COMBINED_PICKLE_FILE = "combined.pkl"

SEED              = 42
BIN_FACTOR        = 20
BIN_SIZE          = 0.001        # 1 ms
SMOOTHING_LENGTH  = 0.05         # 50 ms
SAMPLING_RATE     = 1000
GAUSS_TRUNCATE    = 4.0

DEVICE            = torch.device("cuda" if torch.cuda.is_available() else "cpu")
BATCH_SIZE        = 256

# Hyperparameters are stored here (single dict), like your first script style.
# You can override any of these via CLI flags.
HYP = dict(
    n_pca=32,
    k_lag=20,
    stride_mode="K",          # "K" => stride = k_lag, "1" => stride = 1
    embargo_includes_klag=False,  # False => smoothing embargo only, True => max(k_lag, smoothing_embargo)

    # GD-only knobs
    epochs=150,
    lr=1e-3,
    weight_decay=0.0,
)

# ============================== UTILS ==============================
def set_seed(seed: int):
    import random
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def get_all_unit_names(combined_df):
    unit_set = set()
    for _, row in combined_df.iterrows():
        sc = row.get("spike_counts", None)
        if isinstance(sc, pd.DataFrame):
            unit_set.update(sc.columns)
    return sorted(list(unit_set))

def butter_lowpass(data, fs, order=4, cutoff_hz=5.0):
    nyq = 0.5 * fs
    norm = cutoff_hz / nyq
    b, a = butter(order, norm, btype='low', analog=False)
    return filtfilt(b, a, data, axis=0)

def downsample_spike_and_emg(spike_df, emg_data, bin_factor=10):
    if spike_df.empty or spike_df.shape[0] < bin_factor:
        return spike_df, emg_data

    T_old, n_units = spike_df.shape
    T_new = T_old // bin_factor
    spk_arr = spike_df.values[: T_new * bin_factor, :]
    spk_arr = spk_arr.reshape(T_new, bin_factor, n_units).sum(axis=1)
    ds_spike_df = pd.DataFrame(spk_arr, columns=spike_df.columns)

    if isinstance(emg_data, pd.DataFrame):
        e_arr = emg_data.values
        col_names = emg_data.columns
    else:
        e_arr = np.array(emg_data)
        col_names = None

    if e_arr.shape[0] < bin_factor:
        return ds_spike_df, emg_data

    e_arr = e_arr[: T_new * bin_factor, ...]
    if e_arr.ndim == 2:
        e_arr = e_arr.reshape(T_new, bin_factor, e_arr.shape[1]).mean(axis=1)
        ds_emg = pd.DataFrame(e_arr, columns=col_names) if col_names is not None else e_arr
    else:
        ds_emg = emg_data

    return ds_spike_df, ds_emg

def build_continuous_dataset_raw(df, bin_factor, all_units=None):
    """Concat trials (downsampled) without smoothing/filtering; return X_raw, Y_raw, cuts."""
    spikes_all, emg_all, lengths = [], [], []
    for _, row in df.iterrows():
        spike_df = row.get("spike_counts", None)
        emg_val  = row.get("EMG", None)
        if not isinstance(spike_df, pd.DataFrame) or spike_df.empty:
            continue
        if emg_val is None:
            continue

        if all_units is not None:
            spike_df = spike_df.reindex(columns=all_units, fill_value=0)

        ds_spike_df, ds_emg = downsample_spike_and_emg(spike_df, emg_val, bin_factor)
        if ds_spike_df.shape[0] == 0:
            continue

        Xr = ds_spike_df.values.astype(np.float32)
        Yr = ds_emg.values.astype(np.float32) if isinstance(ds_emg, pd.DataFrame) else np.asarray(ds_emg, dtype=np.float32)

        spikes_all.append(Xr)
        emg_all.append(Yr)
        lengths.append(len(Xr))

    if not spikes_all:
        return np.empty((0,), dtype=np.float32), np.empty((0,), dtype=np.float32), []

    cuts = np.cumsum(lengths)[:-1].tolist()
    return np.concatenate(spikes_all, axis=0), np.concatenate(emg_all, axis=0), cuts

def smooth_spike_data(x_2d, eff_bin, smoothing_length):
    sigma = (smoothing_length / eff_bin) / 2.0
    return gaussian_filter1d(x_2d.astype(np.float32), sigma=sigma, axis=0)

def preprocess_segment(Xseg, Yseg, bin_factor, bin_size=BIN_SIZE, smoothing_length=SMOOTHING_LENGTH):
    """Segment-wise smoothing/filtering (no leak across segments)."""
    eff_fs  = SAMPLING_RATE // bin_factor
    eff_bin = bin_factor * bin_size
    Xs = smooth_spike_data(Xseg, eff_bin, smoothing_length)
    Ys = butter_lowpass(np.abs(Yseg), eff_fs)
    return Xs, Ys

def sigma_bins(bin_factor, bin_size=BIN_SIZE, smoothing_length=SMOOTHING_LENGTH):
    eff_bin = bin_factor * bin_size
    return (smoothing_length / eff_bin) / 2.0

def smoothing_embargo_bins(bin_factor=BIN_FACTOR, bin_size=BIN_SIZE, smoothing_length=SMOOTHING_LENGTH, truncate=GAUSS_TRUNCATE):
    emb = int(np.ceil(truncate * sigma_bins(bin_factor, bin_size, smoothing_length)))
    return emb

def embargo_bins(k_lag, include_klag=True):
    emb = smoothing_embargo_bins()
    return max(k_lag, emb) if include_klag else emb

def time_kfold_splits(n_time, n_splits):
    block = n_time // n_splits
    splits = []
    for k in range(n_splits):
        v0 = k * block
        v1 = (k + 1) * block if k < n_splits - 1 else n_time
        splits.append((v0, v1))
    return splits

def adjust_cuts_for_segment(start, end, cuts_global, trim_left=0, trim_right=0, seg_len=None):
    local = [c - start for c in cuts_global if start < c < end]
    if seg_len is None:
        seg_len = end - start
    new_start = trim_left
    new_end   = seg_len - trim_right
    return [c - new_start for c in local if new_start < c < new_end]

def valid_window_indices(n_time, k, cuts, start=0, end=None, stride=1):
    end = n_time if end is None else end
    out = []
    for t in range(start + k, end, stride):
        if any(t - k < c < t for c in cuts):
            continue
        out.append(t)
    return out

def build_seq_with_cuts_flat(Z, Y, k_lag, cuts, stride=1):
    idx = valid_window_indices(Z.shape[0], k_lag, cuts, stride=stride)
    if not idx:
        return (np.empty((0, k_lag * Z.shape[1]), dtype=np.float32),
                np.empty((0, Y.shape[1]), dtype=np.float32))
    X = np.stack([Z[t-k_lag:t, :].reshape(-1) for t in idx], axis=0).astype(np.float32)
    Yb = np.stack([Y[t, :] for t in idx], axis=0).astype(np.float32)
    return X, Yb

# ========================= DIM REDUCTION =========================
def get_dimred_model(data, method, n_components, seed):
    if method.upper() == "PCA":
        model = PCA(n_components=n_components, random_state=seed)
    elif method.upper() == "UMAP":
        if umap is None:
            raise RuntimeError("umap-learn not installed: pip install umap-learn")
        model = umap.UMAP(n_components=n_components, random_state=seed)
    else:
        raise ValueError(f"Unknown dim. reduction: {method}")
    model.fit(data)
    return model

def transform_dimred(model, data):
    return model.transform(data)

# ============================ WIENER MODELS ============================
class WienerCascadeTorch(nn.Module):
    """Linear FIR stage + per-output quadratic."""
    def __init__(self, input_dim: int, n_out: int):
        super().__init__()
        self.lin = nn.Linear(input_dim, n_out, bias=True)
        self.c0  = nn.Parameter(torch.zeros(n_out))
        self.c1  = nn.Parameter(torch.ones(n_out))
        self.c2  = nn.Parameter(torch.ones(n_out) * 0.1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.lin(x)
        return self.c0 + self.c1 * z + self.c2 * (z ** 2)

def train_torch(model, X_train, Y_train, num_epochs, lr, weight_decay=0.0):
    ds = TensorDataset(torch.tensor(X_train, dtype=torch.float32),
                       torch.tensor(Y_train, dtype=torch.float32))
    dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True)
    opt  = optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    crit = nn.MSELoss()
    model.train()
    for _ in range(num_epochs):
        for xb, yb in dl:
            xb, yb = xb.to(DEVICE), yb.to(DEVICE)
            opt.zero_grad(set_to_none=True)
            loss = crit(model(xb), yb)
            loss.backward()
            opt.step()
    return model

def solve_wiener_cascade_ols(X_train: np.ndarray, Y_train: np.ndarray):
    """
    OLS teacher method:
      1) Linear stage W by pinv on [1, X]
      2) Per-output quadratic c by pinv on [1, z, z^2]
    Returns W_lin (bias+weights), C_poly (per-output [c0,c1,c2]).
    """
    X_bias = np.column_stack([np.ones(X_train.shape[0], dtype=np.float32), X_train.astype(np.float32)])
    W_lin = (np.linalg.pinv(X_bias) @ Y_train).astype(np.float32)      # (1+din) x n_out
    Z = (X_bias @ W_lin).astype(np.float32)                            # n x n_out

    C_poly = np.zeros((Y_train.shape[1], 3), dtype=np.float32)
    for m in range(Y_train.shape[1]):
        z_m = Z[:, m].astype(np.float32)
        M = np.column_stack([np.ones_like(z_m), z_m, z_m * z_m]).astype(np.float32)
        c = (np.linalg.pinv(M) @ Y_train[:, m]).astype(np.float32)
        C_poly[m] = c
    return W_lin, C_poly

def predict_wiener_ols(X: np.ndarray, W_lin: np.ndarray, C_poly: np.ndarray) -> np.ndarray:
    X_bias = np.column_stack([np.ones(X.shape[0], dtype=np.float32), X.astype(np.float32)])
    Z = (X_bias @ W_lin).astype(np.float32)  # n x n_out
    c0 = C_poly[:, 0][None, :]
    c1 = C_poly[:, 1][None, :]
    c2 = C_poly[:, 2][None, :]
    Yp = c0 + c1 * Z + c2 * (Z * Z)
    return Yp.astype(np.float32)

# ============================ METRICS ============================
def eval_vaf_full_numpy(yt: np.ndarray, yp: np.ndarray) -> np.ndarray:
    v = []
    for ch in range(yt.shape[1]):
        vt = np.var(yt[:, ch])
        v.append(np.nan if vt < 1e-12 else 1.0 - np.var(yt[:, ch] - yp[:, ch]) / vt)
    return np.asarray(v, dtype=np.float32)

# ============================ ALIGNMENT ============================
def align_linear_pinv(Zx: np.ndarray, Z0: np.ndarray, lam: float = 1e-6) -> np.ndarray:
    """Solve Zx A ≈ Z0 (centered) => A = (X^T X + lam I)^(-1) X^T Y."""
    if Zx.shape != Z0.shape:
        raise ValueError(f"pinv align requires same shape, got {Zx.shape} vs {Z0.shape}")
    X = Zx - Zx.mean(axis=0, keepdims=True)
    Y = Z0 - Z0.mean(axis=0, keepdims=True)
    d = X.shape[1]
    A = np.linalg.solve(X.T @ X + lam * np.eye(d), X.T @ Y)
    return (Zx - Zx.mean(axis=0, keepdims=True)) @ A + Z0.mean(axis=0, keepdims=True)

# ================================ MAIN ================================
def main():
    ap = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ap.add_argument("--train_method", type=str, required=True, choices=["GD", "OLS"])
    ap.add_argument("--dimred", type=str, default="PCA", choices=["PCA", "UMAP"])
    ap.add_argument("--n_folds", type=int, default=5)
    ap.add_argument("--seed", type=int, default=SEED)
    ap.add_argument("--save_dir", type=str, default=".")
    ap.add_argument("--combined_pickle", type=str, default=COMBINED_PICKLE_FILE)

    # Optional overrides of HYP
    ap.add_argument("--n_pca", type=int, default=None)
    ap.add_argument("--k_lag", type=int, default=None)
    ap.add_argument("--stride_mode", type=str, default=None, choices=["K", "1"])
    ap.add_argument("--embargo_includes_klag", action="store_true", help="If set: EMB=max(k_lag, smoothing_embargo).")
    ap.add_argument("--epochs", type=int, default=None)
    ap.add_argument("--lr", type=float, default=None)
    ap.add_argument("--weight_decay", type=float, default=None)

    args = ap.parse_args()
    set_seed(args.seed)

    # Merge HYP with CLI overrides
    hyp = dict(HYP)
    if args.n_pca is not None: hyp["n_pca"] = int(args.n_pca)
    if args.k_lag is not None: hyp["k_lag"] = int(args.k_lag)
    if args.stride_mode is not None: hyp["stride_mode"] = str(args.stride_mode)
    if args.epochs is not None: hyp["epochs"] = int(args.epochs)
    if args.lr is not None: hyp["lr"] = float(args.lr)
    if args.weight_decay is not None: hyp["weight_decay"] = float(args.weight_decay)

    # The CLI flag sets the boolean (kept in hyp for logging)
    hyp["embargo_includes_klag"] = bool(args.embargo_includes_klag)

    N_PCA = int(hyp["n_pca"])
    K_LAG = int(hyp["k_lag"])
    STRIDE = K_LAG if str(hyp["stride_mode"]).upper() == "K" else 1
    EMB = embargo_bins(K_LAG, include_klag=bool(hyp["embargo_includes_klag"]))

    # Load
    combined_df = pd.read_pickle(args.combined_pickle)
    if not np.issubdtype(combined_df["date"].dtype, np.datetime64):
        combined_df["date"] = pd.to_datetime(combined_df["date"], errors="coerce")

    unique_days = sorted(combined_df["date"].dropna().unique())
    if len(unique_days) == 0:
        raise RuntimeError("No days found in combined_df")

    day0 = unique_days[0]
    train_df = combined_df[combined_df["date"] == day0].reset_index(drop=True)

    ALL_UNITS = get_all_unit_names(combined_df)

    # Detect EMG channels
    n_emg_channels = 0
    for _, row in combined_df.iterrows():
        emg_val = row.get("EMG", None)
        if emg_val is not None:
            if isinstance(emg_val, pd.DataFrame) and not emg_val.empty:
                n_emg_channels = emg_val.shape[1]; break
            if isinstance(emg_val, np.ndarray) and emg_val.size > 0:
                n_emg_channels = emg_val.shape[1]; break
    if n_emg_channels == 0:
        raise RuntimeError("Could not detect EMG channels")

    # Raw Day0
    X0_raw, Y0_raw, cuts0 = build_continuous_dataset_raw(train_df, BIN_FACTOR, all_units=ALL_UNITS)
    if X0_raw.size == 0:
        raise RuntimeError("Empty day0 after downsampling")

    splits = time_kfold_splits(X0_raw.shape[0], args.n_folds)
    results = []

    for fold, (val_start, val_end) in enumerate(splits, 1):
        # raw segments
        X_left_raw  = X0_raw[:val_start];        Y_left_raw  = Y0_raw[:val_start]
        X_val_raw   = X0_raw[val_start:val_end]; Y_val_raw   = Y0_raw[val_start:val_end]
        X_right_raw = X0_raw[val_end:];          Y_right_raw = Y0_raw[val_end:]

        # preprocess each segment independently (no leak)
        Xl, Yl = preprocess_segment(X_left_raw,  Y_left_raw,  BIN_FACTOR, BIN_SIZE, SMOOTHING_LENGTH) if len(X_left_raw) else (np.empty((0, X0_raw.shape[1]), np.float32), np.empty((0, n_emg_channels), np.float32))
        Xv, Yv = preprocess_segment(X_val_raw,   Y_val_raw,   BIN_FACTOR, BIN_SIZE, SMOOTHING_LENGTH) if len(X_val_raw)  else (np.empty((0, X0_raw.shape[1]), np.float32), np.empty((0, n_emg_channels), np.float32))
        Xr, Yr = preprocess_segment(X_right_raw, Y_right_raw, BIN_FACTOR, BIN_SIZE, SMOOTHING_LENGTH) if len(X_right_raw) else (np.empty((0, X0_raw.shape[1]), np.float32), np.empty((0, n_emg_channels), np.float32))

        # trim by embargo and adjust cuts
        if len(Xl) > EMB:
            Xl = Xl[:len(Xl)-EMB]; Yl = Yl[:len(Yl)-EMB]
            cuts_left = adjust_cuts_for_segment(0, len(X_left_raw), cuts0, trim_left=0, trim_right=EMB, seg_len=len(X_left_raw))
        else:
            Xl = np.empty((0, X0_raw.shape[1]), np.float32); Yl = np.empty((0, n_emg_channels), np.float32); cuts_left = []

        if len(Xv) > 2*EMB:
            Xv = Xv[EMB:len(Xv)-EMB]; Yv = Yv[EMB:len(Yv)-EMB]
            cuts_val = adjust_cuts_for_segment(val_start, val_end, cuts0, trim_left=EMB, trim_right=EMB, seg_len=len(X_val_raw))
        else:
            Xv = np.empty((0, X0_raw.shape[1]), np.float32); Yv = np.empty((0, n_emg_channels), np.float32); cuts_val = []

        if len(Xr) > EMB:
            Xr = Xr[EMB:]; Yr = Yr[EMB:]
            cuts_right = adjust_cuts_for_segment(val_end, len(X0_raw), cuts0, trim_left=EMB, trim_right=0, seg_len=len(X_right_raw))
        else:
            Xr = np.empty((0, X0_raw.shape[1]), np.float32); Yr = np.empty((0, n_emg_channels), np.float32); cuts_right = []

        # concat train time
        if Xl.size and Xr.size:
            Xtr_time = np.vstack([Xl, Xr])
            Ytr_time = np.vstack([Yl, Yr])
            cuts_train = cuts_left + [len(Xl)] + [c + len(Xl) for c in cuts_right]
        elif Xl.size:
            Xtr_time, Ytr_time, cuts_train = Xl, Yl, cuts_left
        else:
            Xtr_time, Ytr_time, cuts_train = Xr, Yr, cuts_right

        if Xtr_time.shape[0] <= K_LAG or Xv.shape[0] <= K_LAG:
            print(f"[fold={fold}] not enough samples after embargo; skip")
            continue

        # Fit manifold on TRAIN only
        dimred_model_day0 = get_dimred_model(Xtr_time, args.dimred, max(N_PCA, 2), args.seed + fold)
        Z_tr = transform_dimred(dimred_model_day0, Xtr_time)[:, :N_PCA]
        Z_va = transform_dimred(dimred_model_day0, Xv)[:, :N_PCA]

        # Windowing with cuts
        X_tr, Y_tr = build_seq_with_cuts_flat(Z_tr, Ytr_time, K_LAG, cuts_train, stride=STRIDE)
        X_te, Y_te = build_seq_with_cuts_flat(Z_va, Yv,       K_LAG, cuts_val,   stride=STRIDE)
        if X_tr.shape[0] == 0 or X_te.shape[0] == 0:
            print(f"[fold={fold}] empty after windowing; skip")
            continue

        # Train Wiener (GD or OLS)
        if args.train_method == "GD":
            model = WienerCascadeTorch(K_LAG * N_PCA, n_emg_channels).to(DEVICE)
            train_torch(model, X_tr, Y_tr,
                        num_epochs=int(hyp["epochs"]),
                        lr=float(hyp["lr"]),
                        weight_decay=float(hyp["weight_decay"]))
            with torch.no_grad():
                Yp = model(torch.from_numpy(X_te).float().to(DEVICE)).cpu().numpy().astype(np.float32)
        else:
            W_lin, C_poly = solve_wiener_cascade_ols(X_tr, Y_tr)
            Yp = predict_wiener_ols(X_te, W_lin, C_poly)

        vaf_ch = eval_vaf_full_numpy(Y_te, Yp)
        mean_vaf = float(np.nanmean(vaf_ch))

        # Save crossval rows (day0)
        for ch_idx, v in enumerate(vaf_ch):
            results.append({
                "day": day0,
                "day_int": 0,
                "align": "crossval",
                "decoder": "Wiener",
                "train_method": args.train_method,
                "dim_red": args.dimred,
                "fold": fold - 1,
                "emg_channel": int(ch_idx),
                "vaf": float(v),
                "mean_vaf": mean_vaf,

                # Hyperparams compacted: store the dict as-is + a few computed values
                "hyp": dict(hyp),
                "stride": int(STRIDE),
                "embargo": int(EMB),
            })

        # -------- Cross-days (direct + aligned) using same fold-trained model/params --------
        for d_val in unique_days:
            if pd.to_datetime(d_val) == pd.to_datetime(day0):
                continue

            day_df = combined_df[combined_df["date"] == d_val].reset_index(drop=True)
            spikeX_raw, emgX_raw, cutsX = build_continuous_dataset_raw(day_df, BIN_FACTOR, all_units=ALL_UNITS)
            if spikeX_raw.shape[0] == 0:
                continue

            # preprocess entire dayX (test-time)
            spikeX, emgX = preprocess_segment(spikeX_raw, emgX_raw, BIN_FACTOR, BIN_SIZE, SMOOTHING_LENGTH)

            zx_direct = transform_dimred(dimred_model_day0, spikeX)[:, :N_PCA]

            # Build dayX manifold for alignment
            dimred_model_dayX = get_dimred_model(spikeX, args.dimred, max(N_PCA, 2), args.seed + fold)
            zx_dayX = transform_dimred(dimred_model_dayX, spikeX)[:, :N_PCA]

            for align_mode in ["direct", "aligned"]:
                if align_mode == "direct":
                    zx_test = zx_direct
                else:
                    if args.dimred.upper() == "PCA":
                        V0 = dimred_model_day0.components_[:N_PCA, :].T
                        Vx = dimred_model_dayX.components_[:N_PCA, :].T
                        try:
                            R = pinv(Vx) @ V0
                            zx_test = zx_dayX @ R
                        except Exception:
                            zx_test = zx_dayX
                    else:
                        try:
                            zx_test = align_linear_pinv(zx_dayX, zx_direct, lam=1e-6)
                        except Exception:
                            zx_test = zx_dayX

                X_seq, Y_seq = build_seq_with_cuts_flat(zx_test, emgX, K_LAG, cutsX, stride=STRIDE)
                if X_seq.shape[0] == 0:
                    continue

                if args.train_method == "GD":
                    with torch.no_grad():
                        Yp2 = model(torch.from_numpy(X_seq).float().to(DEVICE)).cpu().numpy().astype(np.float32)
                else:
                    Yp2 = predict_wiener_ols(X_seq, W_lin, C_poly)

                vaf_ch2 = eval_vaf_full_numpy(Y_seq, Yp2)
                mean_vaf2 = float(np.nanmean(vaf_ch2))

                for ch_idx, v in enumerate(vaf_ch2):
                    results.append({
                        "day": d_val,
                        "day_int": int((pd.to_datetime(d_val) - pd.to_datetime(day0)).days),
                        "align": align_mode,
                        "decoder": "Wiener",
                        "train_method": args.train_method,
                        "dim_red": args.dimred,
                        "fold": fold - 1,
                        "emg_channel": int(ch_idx),
                        "vaf": float(v),
                        "mean_vaf": mean_vaf2,
                        "hyp": dict(hyp),
                        "stride": int(STRIDE),
                        "embargo": int(EMB),
                    })

        if args.train_method == "GD":
            del model
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        print(f"[fold={fold}/{args.n_folds}] done.")

    # Save
    os.makedirs(args.save_dir, exist_ok=True)
    out_name = f"crossday_results_Wiener_{args.dimred}_{args.train_method}.pkl"
    save_path = os.path.join(args.save_dir, out_name)
    pd.to_pickle(pd.DataFrame(results), save_path)
    print(f"\n[INFO] Saved all results to {save_path}")

if __name__ == "__main__":
    main()