{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%writefile /kaggle/working/20_train_pcsbn.py\n#!/usr/bin/env python\n# =============================================================================\n# ARIEL 2025 -- SCRIPT 20 : PC-SBN TRAINER\n#   contribution 1  Differentiable Transit Projection (DTP)\n#   contribution 2  Structure-aware output head\n# =============================================================================\n# THE MODEL\n#\n#   x(t,lambda) = flux/baseline - 1        (Cottaar step-5 v3 cache)\n#   physically   x ~= -s(t; theta) * d(lambda)\n#\n#   1. white-light curve      w(t) = mean_lambda x(t,lambda)\n#   2. GEOMETRY CNN           w -> theta = (t0, T14, tau, u)          [NEW 1]\n#   3. TRANSIT SHAPE          s(t; theta) in [0,1], analytic, 0 params\n#   4. CLOSED-FORM PROJECTION d_phys(lambda), var(lambda)             [NEW 1]\n#   5. RESIDUAL               r = x + s (x) d_phys        -> encoder\n#   6. ModernTCN ENCODER      r -> z in R^192             (frozen config)\n#   7. STRUCTURE-AWARE HEAD   z, d_phys, var -> mu, sigma             [NEW 2]\n#\n# THE PROJECTION, exactly. The instrument noise in the v3 cache is\n# per-wavelength and constant in time, so a weighted least squares over t at\n# fixed lambda has a constant weight -- which CANCELS in the depth estimate and\n# survives only in its variance (verified to 5e-18):\n#\n#     d_phys(lambda) = sum_t s_t * (-x(t,lambda))  /  sum_t s_t^2\n#     var(lambda)    = noise_rel(lambda)^2         /  sum_t s_t^2\n#\n# sum_t s_t^2 is one scalar per planet: the effective number of in-transit bins.\n# noise_rel is known to carry a uniform 4.81x scale error; it cancels in d_phys\n# and is absorbed by the learned scale c in the head, so no cache rebuild is\n# needed. Its RELATIVE structure across wavelength is correct (corr 0.976) and\n# its within-planet spread is p90/p10 = 3.0, so the weights do real work.\n#\n# THE HEAD, exactly.\n#     mu(lambda)     = d_phys(lambda) + mu_L + sum_k a_k B_k(lambda)\n#     sigma^2(lambda) = c^2 var(lambda) + sigma_L^2 + sum_k B_k(lambda)^2 exp(2 l_k)\n# B is a FIXED cosine (DCT-II) basis, k = 1..K, so it is smooth by construction\n# and fitted on no data at all -- nothing to leak. The correction terms are\n# zero-initialised, so at step 0 the model outputs exactly d_phys.\n#\n# ABLATION SWITCHES (environment variables, all read by _env below)\n#   ARIEL_USE_DTP      1 physics path on   0 off (encoder sees the raw tensor)\n#   ARIEL_USE_ENCODER  1 encoder on        0 off (pure physics, rung 2)\n#   ARIEL_HEAD_TYPE    plain | level | basis\n#   ARIEL_K_BASIS      number of cosine basis functions (default 24)\n#   ARIEL_COV_TYPE     diag | chol        (basis-space covariance form)\n#   ARIEL_FOLD         -1 historical split, 0..4 five-fold\n#\n# The ablation ladder (each rung differs from the previous in ONE switch):\n#   1 baseline      USE_DTP=0 USE_ENCODER=1 HEAD_TYPE=plain   (= script 04)\n#   2 physics only  USE_DTP=1 USE_ENCODER=0 HEAD_TYPE=level\n#   3 DTP + plain   USE_DTP=1 USE_ENCODER=1 HEAD_TYPE=plain\n#   4 struct only   USE_DTP=0 USE_ENCODER=1 HEAD_TYPE=basis\n#   5 full          USE_DTP=1 USE_ENCODER=1 HEAD_TYPE=basis\n#\n# -----------------------------------------------------------------------------\n# NOTEBOOK SETUP  (GPU T4 x2, Internet ON, HF_TOKEN secret)\n#   Add Input : Competitions -> ariel-data-challenge-2025\n#   Add Input : Datasets -> moderntcn-src      (for the ModernTCN blocks)\n# =============================================================================\n\nimport os\nimport gc\nimport glob\nimport json\nimport math\nimport random\nimport sys\nimport time\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom scipy.stats import norm as _norm\nfrom torch.utils.data import Dataset, DataLoader\nfrom huggingface_hub import HfApi, create_repo, hf_hub_download\n\n\n# =============================================================================\n# CONFIG\n# =============================================================================\n\ndef _env(name, default, cast=str):\n    v = os.environ.get(\"ARIEL_\" + name)\n    return default if v is None else cast(v)\n\n\ndef _envb(name, default):\n    return bool(int(_env(name, int(default), int)))\n\n\nUSE_DTP      = _envb(\"USE_DTP\", True)\nUSE_ENCODER  = _envb(\"USE_ENCODER\", True)\nHEAD_TYPE    = _env(\"HEAD_TYPE\", \"basis\")        # plain | level | basis\nK_BASIS      = _env(\"K_BASIS\", 24, int)\nCOV_TYPE     = _env(\"COV_TYPE\", \"diag\")          # diag | chol\nFOLD         = _env(\"FOLD\", -1, int)\nN_FOLDS      = _env(\"N_FOLDS\", 5, int)\nTRAIN_SEED   = _env(\"TRAIN_SEED\", 1, int)\nDROPOUT_P    = _env(\"DROPOUT_P\", 0.15, float)\nNUM_EPOCHS   = _env(\"NUM_EPOCHS\", 600, int)\nBETA_NLL     = _env(\"BETA_NLL\", 0.5, float)\n# theta is supervised against the cache's stored transit box for this many\n# epochs. Without it the trapezoid's gradient is sparse at initialisation (only\n# the ingress/egress bins carry signal), so the geometry CNN can sit in a bad\n# basin. The cache values come from a label-free box search, so this adds no\n# label information.\nTHETA_SUP_EPOCHS = _env(\"THETA_SUP_EPOCHS\", 40, int)\nTHETA_SUP_WEIGHT = _env(\"THETA_SUP_WEIGHT\", 1.0, float)\n\nif HEAD_TYPE not in (\"plain\", \"level\", \"basis\"):\n    raise ValueError(f\"HEAD_TYPE must be plain|level|basis, got {HEAD_TYPE!r}\")\nif COV_TYPE not in (\"diag\", \"chol\"):\n    raise ValueError(f\"COV_TYPE must be diag|chol, got {COV_TYPE!r}\")\nif not USE_DTP and not USE_ENCODER:\n    raise ValueError(\"USE_DTP=0 and USE_ENCODER=0 leaves no model at all\")\nif not USE_DTP and HEAD_TYPE == \"level\":\n    raise ValueError(\"HEAD_TYPE=level without DTP has nothing to correct; \"\n                     \"use plain or basis\")\n\nCACHE_REPO = \"maariaaa12/ariel-2025-cottaar-step5-cache-v3\"\nLOCAL_CACHE_DIR = \"/kaggle/working/cottaar_step5_cache_v3\"\nTRAIN_CSV = \"/kaggle/input/competitions/ariel-data-challenge-2025/train.csv\"\nSTAR_INFO_CSV = \"/kaggle/input/competitions/ariel-data-challenge-2025/train_star_info.csv\"\nMETADATA_COLS = [\"Rs\", \"Ms\", \"Ts\", \"Mp\", \"e\", \"P\", \"sma\", \"i\"]\nSRC_NAME = \"train_spectrum_moderntcn_v2_robust_cache.py\"\n\nN_WAVE = 283\nSEQ_LEN = 187\nPPM_SCALE = 1e6\nSIGMA_FLOOR = 1e-6          # 1 ppm, native (Rp/Rs)^2 units\nGLL_IDEAL_SIGMA_PPM = 10.0\nECE_N_BINS = 10\nVAL_FRACTION = 0.1\nTRAIN_EVAL_SUBSET_SIZE = 110\n\n# float16 safety: the encoder runs under autocast, and near 1.0 float16 resolves\n# only ~500-1000 ppm. The residual is already centred near 0, so it only needs\n# the x100 scale to sit clear of the subnormal range.\nENCODER_SCALE = 100.0\n\nBATCH_SIZE = 16\nLR = 2e-4\nWEIGHT_DECAY = 9e-4\nHIDDEN_DIM = 192\nKERNEL_SIZES = [15, 15, 15]\nN_BLOCKS = [1, 1, 1]\nGRAD_CLIP_NORM = 1.0\nMSE_WARMUP_EPOCHS = 20\nTRANSITION_LR_DECAY = 0.3\nEMA_DECAY = 0.995\nNUM_WORKERS = 2\nHF_UPLOAD_EVERY_N_EPOCHS = 10\nRESUME_IF_AVAILABLE = True\n\n# ---- tag ----\n_t = \"pcsbn\"\n_t += \"_dtp\" if USE_DTP else \"_nodtp\"\n_t += \"_enc\" if USE_ENCODER else \"_noenc\"\n_t += f\"_{HEAD_TYPE}\"\nif HEAD_TYPE == \"basis\":\n    _t += f\"_k{K_BASIS}\"\n    if COV_TYPE != \"diag\":\n        _t += f\"_{COV_TYPE}\"\nif BETA_NLL != 0.5:\n    _t += \"_b\" + str(BETA_NLL).replace(\".\", \"\")\n_t += f\"_s{TRAIN_SEED}\"\nif DROPOUT_P != 0.15:\n    _t += \"_d\" + str(DROPOUT_P).replace(\".\", \"\")\n_repo_tag = _t\nif FOLD >= 0:\n    if not (0 <= FOLD < N_FOLDS):\n        raise ValueError(f\"FOLD must be in [0,{N_FOLDS}), got {FOLD}\")\n    _t += f\"_f{FOLD}\"\n    _repo_tag += \"_5fold\"\nDATA_TAG = _t\nCHECKPOINT_REPO = f\"maariaaa12/ariel-2025-{_repo_tag.replace('_', '-')}\"\n\nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n_hf_api = HfApi()\n\n\ndef banner(t):\n    print(f\"\\n{'='*78}\\n{t}\\n{'='*78}\", flush=True)\n\n\n# =============================================================================\n# THE ModernTCN BLOCKS COME FROM THE EXISTING TRAINER -- same blocks, so the\n# encoder is literally the frozen winner, not a re-implementation of it.\n# =============================================================================\n\ndef _load_src():\n    hits = []\n    for pat in (f\"/kaggle/input/**/{SRC_NAME}\", f\"/kaggle/working/{SRC_NAME}\"):\n        hits += sorted(glob.glob(pat, recursive=True))\n    if not hits:\n        raise FileNotFoundError(f\"{SRC_NAME} not found -- attach the moderntcn-src dataset.\")\n    d = os.path.dirname(hits[0])\n    if d not in sys.path:\n        sys.path.insert(0, d)\n    import importlib\n    m = importlib.import_module(SRC_NAME[:-3])\n    print(f\"  ModernTCN blocks from {hits[0]}\")\n    return m\n\n\nM = _load_src()\n\n\n# =============================================================================\n# DATASET  --  returns the RAW tensor. Scaling happens inside the model, so the\n# physics path works in native relative-flux units and only the encoder sees the\n# x100 version.\n# =============================================================================\n\nclass PCSBNDataset(Dataset):\n    def __init__(self, planet_ids, targets_df, star_info_df, meta_mean, meta_std):\n        self.pids = [str(p) for p in planet_ids]\n        self.dir = Path(LOCAL_CACHE_DIR)\n        self.dir.mkdir(parents=True, exist_ok=True)\n        t = targets_df.copy()\n        t[\"planet_id\"] = t[\"planet_id\"].astype(str)\n        self.targets = t.set_index(\"planet_id\")\n        s = star_info_df.copy()\n        s[\"planet_id\"] = s[\"planet_id\"].astype(str)\n        self.star = s.set_index(\"planet_id\")\n        self.meta_mean, self.meta_std = meta_mean, meta_std\n\n    def __len__(self):\n        return len(self.pids)\n\n    def _path(self, pid):\n        p = self.dir / f\"{pid}.npz\"\n        if p.exists():\n            return p\n        return Path(hf_hub_download(repo_id=CACHE_REPO, repo_type=\"dataset\",\n                                    filename=f\"{pid}.npz\", local_dir=str(self.dir)))\n\n    def __getitem__(self, i):\n        pid = self.pids[i]\n        z = np.load(self._path(pid))\n        x = z[\"spectrum_clean\"].astype(np.float32)             # (187, 283)\n        nr = z[\"noise_rel\"].astype(np.float32)                  # (283,)\n        # a zero or non-finite noise entry would make var(lambda) explode\n        nr = np.nan_to_num(nr, nan=0.0, posinf=0.0, neginf=0.0)\n        med = np.median(nr[nr > 0]) if (nr > 0).any() else 1e-4\n        nr = np.where(nr > 0, nr, med).astype(np.float32)\n        y = self.targets.loc[pid].values.astype(np.float32)      # (283,) native\n        meta = ((self.star.loc[pid, METADATA_COLS].values.astype(np.float32)\n                 - self.meta_mean) / self.meta_std)\n        c = float(z[\"transit_centre_bin\"]) if \"transit_centre_bin\" in z else SEQ_LEN / 2\n        w = float(z[\"transit_width_bin\"]) if \"transit_width_bin\" in z else SEQ_LEN / 3\n        return {\"planet_id\": pid,\n                \"x\": torch.from_numpy(x),\n                \"noise_rel\": torch.from_numpy(nr),\n                \"target\": torch.from_numpy(y),\n                \"metadata\": torch.from_numpy(meta),\n                \"theta_ref\": torch.tensor([c, w], dtype=torch.float32)}\n\n\n# =============================================================================\n# CONTRIBUTION 1 -- DIFFERENTIABLE TRANSIT PROJECTION\n# =============================================================================\n\nclass GeometryCNN(nn.Module):\n    \"\"\"White-light curve -> theta = (t0, T14, tau, u). ~20k parameters.\n\n    Only the white-light curve is used, never the labels. The outputs are\n    squashed into physically valid ranges so no epoch can produce a nonsensical\n    transit that would make the projection blow up.\n    \"\"\"\n\n    def __init__(self, seq_len=SEQ_LEN, ch=32):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Conv1d(1, ch, 7, padding=3), nn.BatchNorm1d(ch), nn.GELU(),\n            nn.Conv1d(ch, ch, 7, stride=2, padding=3), nn.BatchNorm1d(ch), nn.GELU(),\n            nn.Conv1d(ch, ch, 5, stride=2, padding=2), nn.BatchNorm1d(ch), nn.GELU(),\n            nn.Conv1d(ch, ch, 5, padding=2), nn.BatchNorm1d(ch), nn.GELU(),\n        )\n        self.head = nn.Linear(2 * ch, 4)\n        nn.init.zeros_(self.head.weight)\n        nn.init.zeros_(self.head.bias)      # starts at the centre of every range\n        self.seq_len = seq_len\n\n    def forward(self, w):\n        # w: (B, T). Normalise per planet so the CNN sees shape, not depth.\n        wc = w - w.mean(1, keepdim=True)\n        sd = wc.std(1, keepdim=True).clamp_min(1e-8)\n        h = self.net((wc / sd).unsqueeze(1))                       # (B, ch, T')\n        h = torch.cat([h.mean(-1), h.amax(-1)], dim=-1)            # (B, 2ch)\n        raw = self.head(h)\n        T = self.seq_len\n        # t0 inside the middle half of the window\n        t0 = 0.5 * T + 0.25 * T * torch.tanh(raw[:, 0])\n        # full duration between 10% and 80% of the window\n        T14 = T * (0.10 + 0.70 * torch.sigmoid(raw[:, 1]))\n        # ingress as a fraction of the half-duration, 5%..50%\n        tau = 0.05 + 0.45 * torch.sigmoid(raw[:, 2])\n        # linear limb-darkening coefficient, 0..0.9\n        u = 0.9 * torch.sigmoid(raw[:, 3])\n        return t0, T14, tau, u\n\n\ndef transit_shape(t0, T14, tau, u, seq_len=SEQ_LEN, device=None):\n    \"\"\"s(t; theta) in [0,1]: a trapezoid with a linear limb-darkening profile.\n\n    Differentiable in every parameter. At u=0 it is the plain trapezoid, which\n    is the correct first-order shape for a transit. The limb-darkening factor is\n    normalised so the disc-average intensity is 1, which keeps d_phys on the\n    same scale as the label.\n    \"\"\"\n    t = torch.arange(seq_len, device=device, dtype=t0.dtype).unsqueeze(0)  # (1,T)\n    half = (0.5 * T14).unsqueeze(1)                                        # (B,1)\n    d = (t - t0.unsqueeze(1)).abs()                                        # (B,T)\n    ramp = (tau.unsqueeze(1) * half).clamp_min(1.0)   # >= 1 bin, so the ingress\n                                                       # always spans real bins\n                                                       # and carries a gradient\n    g = ((half - d) / ramp).clamp(0.0, 1.0)                                # (B,T)\n    # radial position on the stellar disc while in transit\n    rho = (d / half.clamp_min(1e-6)).clamp(0.0, 1.0)\n    mu_ld = torch.sqrt((1.0 - rho ** 2).clamp_min(0.0))\n    uu = u.unsqueeze(1)\n    I = (1.0 - uu * (1.0 - mu_ld)) / (1.0 - uu / 3.0)\n    return g * I\n\n\ndef project(x, s, noise_rel):\n    \"\"\"Closed-form per-wavelength depth and its variance.\n\n    x: (B,T,W) = flux/baseline - 1, so the transit is a NEGATIVE dip and\n       x ~= -s (x) d.  noise_rel: (B,W).  s: (B,T).\n    Returns d_phys (B,W), var (B,W), s2 (B,1) = effective in-transit bins.\n    \"\"\"\n    s2 = (s ** 2).sum(1, keepdim=True).clamp_min(1e-3)          # (B,1)\n    d = torch.einsum(\"bt,btw->bw\", s, -x) / s2                   # (B,W)\n    var = (noise_rel ** 2) / s2                                  # (B,W)\n    return d, var, s2\n\n\n# =============================================================================\n# ENCODER -- the frozen winner's own blocks, stopped before its head so the\n# pooled feature vector z is available to the new head.\n# =============================================================================\n\nclass Encoder(nn.Module):\n    def __init__(self, hidden_dim=HIDDEN_DIM, kernels=KERNEL_SIZES,\n                 blocks=N_BLOCKS, dropout_p=DROPOUT_P, n_wave=N_WAVE):\n        super().__init__()\n        self.input_proj = nn.Conv1d(n_wave, hidden_dim, kernel_size=1)\n        self.stages = nn.ModuleList([\n            nn.Sequential(*[M.ModernTCNBlock(hidden_dim, k, dropout_p,\n                                             pointwise_type=\"standard\")\n                            for _ in range(b)])\n            for k, b in zip(kernels, blocks)])\n        self.downsample = nn.MaxPool1d(kernel_size=2, stride=2)\n        self.attn_score = nn.Conv1d(hidden_dim, 1, kernel_size=1)\n        self.out_dim = hidden_dim\n\n    def forward(self, x):\n        # x: (B, T, W) -> channels are wavelength, sequence is time\n        h = self.input_proj(x.transpose(1, 2))\n        n = len(self.stages)\n        for i, st in enumerate(self.stages):\n            h = st(h)\n            if i < n - 1:\n                h = self.downsample(h)\n        wts = F.softmax(self.attn_score(h), dim=-1)\n        return (h * wts).sum(-1)                                # (B, hidden_dim)\n\n\n# =============================================================================\n# CONTRIBUTION 2 -- STRUCTURE-AWARE OUTPUT HEAD\n# =============================================================================\n\ndef cosine_basis(n_wave=N_WAVE, k=K_BASIS):\n    \"\"\"DCT-II basis, k = 1..K (the constant term is left out: mu_L already is\n    the level). Fixed, analytic, fitted on nothing -- so nothing can leak.\"\"\"\n    i = np.arange(n_wave)[:, None]\n    kk = np.arange(1, k + 1)[None, :]\n    B = np.cos(np.pi * kk * (i + 0.5) / n_wave)\n    B = B / np.linalg.norm(B, axis=0, keepdims=True)     # unit-norm columns\n    return torch.tensor(B, dtype=torch.float32)          # (W, K)\n\n\nclass StructureHead(nn.Module):\n    \"\"\"mu   = d_phys + mu_L + sum_k a_k B_k\n       var  = c^2 var_phys + sigma_L^2 + sum_k B_k^2 exp(2 l_k)      [diag]\n              c^2 var_phys + sigma_L^2 + diag(B L L^T B^T)           [chol]\n\n    Every correction is zero-initialised, so at step 0 mu is exactly d_phys and\n    the model degrades gracefully to pure least squares. With HEAD_TYPE=plain it\n    predicts 283 free means and 283 free sigmas instead -- the baseline's head,\n    kept here so the ablation changes one thing at a time.\n    \"\"\"\n\n    def __init__(self, in_dim, n_wave=N_WAVE, head_type=HEAD_TYPE,\n                 k=K_BASIS, cov_type=COV_TYPE, use_phys=USE_DTP,\n                 dropout_p=DROPOUT_P):\n        super().__init__()\n        self.head_type, self.cov_type, self.use_phys = head_type, cov_type, use_phys\n        self.n_wave, self.k = n_wave, k\n        self.dropout = nn.Dropout(dropout_p)\n\n        if head_type == \"basis\":\n            self.register_buffer(\"B\", cosine_basis(n_wave, k))          # (W,K)\n            n_cov = k if cov_type == \"diag\" else k * (k + 1) // 2\n            n_out = 1 + k + n_cov                    # mu_L, a, covariance\n        elif head_type == \"level\":\n            n_out = 1                                 # mu_L only\n            n_cov = 0\n        else:                                         # plain\n            n_out = 2 * n_wave\n            n_cov = 0\n        self.n_cov = n_cov\n        self.fc = nn.Linear(in_dim, n_out)\n        if head_type in (\"basis\", \"level\"):\n            nn.init.zeros_(self.fc.weight)\n            nn.init.zeros_(self.fc.bias)\n        else:\n            nn.init.zeros_(self.fc.bias)\n            self.fc.weight.data.mul_(0.01)\n\n        # log sigma_L and log c, learned globally (not per planet)\n        self.log_sigma_L = nn.Parameter(torch.tensor(math.log(3e-4)))\n        self.log_c = nn.Parameter(torch.zeros(()))     # c = 1 at init\n\n    @torch.no_grad()\n    def warm_init(self, mu_native, sigma_native):\n        \"\"\"The plain head predicts sigma as SIGMA_FLOOR + softplus(raw). With a\n        zero bias that is softplus(0) = 0.693 in native units -- 693,000 ppm,\n        absurd, and it takes many epochs to walk down from there. The baseline\n        trainer solves this with warm_init_head; the same is needed here or the\n        plain-head rungs would be handicapped and the ablation unfair.\n        The structured heads need no warm start: their sigma comes from\n        c^2 var_phys + sigma_L^2, already initialised at a physical scale.\"\"\"\n        if self.head_type != \"plain\":\n            return False\n        s = np.clip(np.asarray(sigma_native) - SIGMA_FLOOR, 1e-8, None)\n        self.fc.bias[:self.n_wave] = torch.tensor(\n            np.asarray(mu_native), dtype=torch.float32)\n        self.fc.bias[self.n_wave:] = torch.tensor(\n            np.log(np.expm1(s)), dtype=torch.float32)\n        return True\n\n    def forward(self, z, d_phys, var_phys):\n        B_ = d_phys.shape[0] if d_phys is not None else z.shape[0]\n        out = self.fc(self.dropout(z)) if z is not None else None\n\n        if self.head_type == \"plain\":\n            mu_corr, log_s = out.chunk(2, dim=-1)\n            sigma = SIGMA_FLOOR + F.softplus(log_s)\n            mu = mu_corr + (d_phys if self.use_phys else 0.0)\n            return mu, sigma, {}\n\n        sigma_L = torch.exp(self.log_sigma_L)\n        c = torch.exp(self.log_c)\n        mu = d_phys.clone() if self.use_phys else torch.zeros(\n            B_, self.n_wave, device=z.device, dtype=z.dtype)\n        v = (c ** 2) * var_phys if self.use_phys else torch.zeros_like(mu)\n        v = v + sigma_L ** 2\n\n        if self.head_type == \"level\":\n            mu = mu + out[:, :1]\n            return mu, torch.sqrt(v.clamp_min(SIGMA_FLOOR ** 2)), {\"c\": c}\n\n        mu_L = out[:, :1]\n        a = out[:, 1:1 + self.k]                                   # (B,K)\n        raw_cov = out[:, 1 + self.k:]\n        mu = mu + mu_L + a @ self.B.T                              # (B,W)\n\n        if self.cov_type == \"diag\":\n            # Sigma_a = diag(exp(2 l_k)), l_k = raw - 4 so the basis variance\n            # starts small. diag(B Sigma_a B^T) = sum_k B_k(lambda)^2 exp(2 l_k),\n            # which is (B,K) @ (K,W) -> (B,W). Diagonal in BASIS space is still\n            # wavelength-CORRELATED, which is the claim.\n            lam = torch.exp(2.0 * (raw_cov - 4.0).clamp(-12.0, 6.0))    # (B,K)\n            v = v + lam @ (self.B ** 2).T                               # (B,W)\n        else:\n            # full Cholesky in basis space: Sigma_a = L L^T\n            L = torch.zeros(B_, self.k, self.k, device=out.device, dtype=out.dtype)\n            ii, jj = torch.tril_indices(self.k, self.k, device=out.device)\n            L[:, ii, jj] = raw_cov\n            dg = torch.arange(self.k, device=out.device)\n            L[:, dg, dg] = torch.exp((L[:, dg, dg] - 4.0).clamp(-12.0, 6.0))\n            BL = torch.einsum(\"wk,bkj->bwj\", self.B, L)                 # (B,W,K)\n            v = v + (BL ** 2).sum(-1)                                   # (B,W)\n\n        return mu, torch.sqrt(v.clamp_min(SIGMA_FLOOR ** 2)), {\"c\": c, \"a\": a}\n\n\n# =============================================================================\n# THE FULL MODEL\n# =============================================================================\n\nclass PCSBN(nn.Module):\n    \"\"\"Every switch is an explicit argument, defaulting to the environment.\n\n    That matters for two things the globals could not do: the EMA shadow copy\n    must be built with the SAME switches (clone_empty), and the out-of-fold\n    assembler must rebuild any checkpoint's architecture from the config stored\n    inside it (build_from_config) rather than from whatever the environment\n    happens to say at that moment.\n    \"\"\"\n\n    def __init__(self, use_dtp=None, use_encoder=None, head_type=None,\n                 k_basis=None, cov_type=None, dropout_p=None, n_wave=N_WAVE):\n        super().__init__()\n        self.kw = dict(\n            use_dtp=USE_DTP if use_dtp is None else bool(use_dtp),\n            use_encoder=USE_ENCODER if use_encoder is None else bool(use_encoder),\n            head_type=HEAD_TYPE if head_type is None else head_type,\n            k_basis=K_BASIS if k_basis is None else int(k_basis),\n            cov_type=COV_TYPE if cov_type is None else cov_type,\n            dropout_p=DROPOUT_P if dropout_p is None else float(dropout_p),\n            n_wave=n_wave)\n        kw = self.kw\n        self.use_dtp, self.use_encoder = kw[\"use_dtp\"], kw[\"use_encoder\"]\n        self.n_wave = kw[\"n_wave\"]\n        if self.use_dtp:\n            self.geom = GeometryCNN()\n        if self.use_encoder:\n            self.encoder = Encoder(dropout_p=kw[\"dropout_p\"], n_wave=kw[\"n_wave\"])\n            in_dim = self.encoder.out_dim\n        else:\n            # pure physics: the head still needs an input, so give it a single\n            # learned vector -- the corrections become per-dataset constants\n            self.z_const = nn.Parameter(torch.zeros(1, 16))\n            in_dim = 16\n        self.head = StructureHead(in_dim, n_wave=kw[\"n_wave\"],\n                                  head_type=kw[\"head_type\"], k=kw[\"k_basis\"],\n                                  cov_type=kw[\"cov_type\"], use_phys=self.use_dtp,\n                                  dropout_p=kw[\"dropout_p\"])\n\n    def clone_empty(self):\n        return PCSBN(**self.kw)\n\n    def forward(self, x, noise_rel, metadata=None):\n        \"\"\"PRECISION SPLIT -- this is not optional.\n\n        var(lambda) = noise_rel^2 / sum_t s^2 is around 8e-10, and float16's\n        smallest normal is 6.1e-5, so under autocast it would flush to EXACTLY\n        zero -- as would SIGMA_FLOOR^2 = 1e-12 -- giving sigma = 0 and a NaN\n        loss. d_phys ~ 0.012 also only resolves ~12 ppm steps in float16, and the\n        signal is ~10 ppm; this is the same failure that cost +0.048 earlier.\n\n        So the physics path and the head run in float32. Only the ENCODER runs in\n        float16, where it is safe: the residual is already centred at 0, so\n        float16's constant relative precision works in its favour there.\n        \"\"\"\n        B = x.shape[0]\n        theta = None\n        with torch.autocast(\"cuda\", enabled=False):\n            x32 = x.float()\n            if self.use_dtp:\n                w = x32.mean(-1)                                       # (B,T)\n                t0, T14, tau, u = self.geom(w)\n                s = transit_shape(t0, T14, tau, u, x32.shape[1], x32.device)\n                d_phys, var_phys, s2 = project(x32, s, noise_rel.float())\n                theta = dict(t0=t0, T14=T14, tau=tau, u=u, s2=s2.squeeze(1))\n                resid = x32 + s.unsqueeze(-1) * d_phys.unsqueeze(1)    # (B,T,W)\n            else:\n                d_phys = torch.zeros(B, self.n_wave,\n                                     device=x32.device, dtype=x32.dtype)\n                var_phys = torch.zeros_like(d_phys)\n                resid = x32\n\n        # the one part that DOES run in float16, under the caller's autocast\n        if self.use_encoder:\n            z = self.encoder(resid * ENCODER_SCALE)\n        else:\n            z = self.z_const.expand(B, -1)\n\n        with torch.autocast(\"cuda\", enabled=False):\n            mu, sigma, extra = self.head(z.float(), d_phys, var_phys)\n        if theta is not None:\n            extra.update(theta)\n            extra[\"d_phys\"] = d_phys\n        return mu, sigma, extra\n\n\n# =============================================================================\n# LOSSES / METRICS  (identical definitions to every other script here)\n# =============================================================================\n\ndef beta_nll(mu, sigma, target, beta=BETA_NLL):\n    nll = 0.5 * (math.log(2 * math.pi) + 2 * torch.log(sigma)\n                 + ((target - mu) / sigma) ** 2)\n    return ((sigma.detach() ** (2 * beta)) * nll).mean()\n\n\ndef theta_supervision(extra, theta_ref):\n    \"\"\"Pull t0 and T14 toward the cache's label-free box search. Scale-free, so\n    it cannot dominate the likelihood.\"\"\"\n    if \"t0\" not in extra:\n        return torch.zeros((), device=theta_ref.device)\n    t0_r, w_r = theta_ref[:, 0], theta_ref[:, 1]\n    return (((extra[\"t0\"] - t0_r) / SEQ_LEN) ** 2).mean() \\\n        + (((extra[\"T14\"] - w_r) / SEQ_LEN) ** 2).mean()\n\n\ndef gll_ppm(y, mu, sg):\n    sg = np.clip(sg, 1e-8, None)\n    return -0.5 * (np.log(2 * np.pi) + np.log(sg ** 2) + ((y - mu) ** 2) / sg ** 2)\n\n\ndef ece(y, mu, sg, n_bins=ECE_N_BINS):\n    pit = _norm.cdf((y - mu) / np.clip(sg, 1e-8, None))\n    lv = np.linspace(0, 1, n_bins + 1)[1:]\n    return float(np.mean(np.abs(np.array([(pit <= l).mean() for l in lv]) - lv)))\n\n\ndef build_from_config(cfg):\n    \"\"\"Rebuild a model exactly as a checkpoint was trained. Used by the\n    out-of-fold assembler, so it never has to guess an architecture.\"\"\"\n    return PCSBN(use_dtp=cfg[\"USE_DTP\"], use_encoder=cfg[\"USE_ENCODER\"],\n                 head_type=cfg[\"HEAD_TYPE\"], k_basis=cfg.get(\"K_BASIS\", K_BASIS),\n                 cov_type=cfg.get(\"COV_TYPE\", \"diag\"),\n                 dropout_p=cfg.get(\"DROPOUT_P\", DROPOUT_P))\n\n\nclass EMA:\n    def __init__(self, model, decay=EMA_DECAY):\n        self.decay = decay\n        self.shadow = model.clone_empty().to(DEVICE)\n        self.shadow.load_state_dict(model.state_dict())\n        for p in self.shadow.parameters():\n            p.requires_grad_(False)\n\n    @torch.no_grad()\n    def update(self, model):\n        msd = model.state_dict()\n        for k, v in self.shadow.state_dict().items():\n            if v.dtype.is_floating_point:\n                v.mul_(self.decay).add_(msd[k].detach(), alpha=1 - self.decay)\n            else:\n                v.copy_(msd[k])\n\n\n@torch.no_grad()\ndef evaluate(model, loader, ref_mean_ppm, ref_std_ppm):\n    model.eval()\n    ys, mus, sgs, th = [], [], [], []\n    for b in loader:\n        x = b[\"x\"].to(DEVICE)\n        nr = b[\"noise_rel\"].to(DEVICE)\n        if DEVICE.type == \"cuda\":\n            with torch.autocast(\"cuda\", dtype=torch.float16):\n                mu, sg, ex = model(x, nr, b[\"metadata\"].to(DEVICE))\n        else:\n            mu, sg, ex = model(x, nr, b[\"metadata\"].to(DEVICE))\n        ys.append(b[\"target\"].numpy())\n        mus.append(mu.float().cpu().numpy())\n        sgs.append(sg.float().cpu().numpy())\n        if \"t0\" in ex:\n            th.append(np.stack([ex[\"t0\"].float().cpu().numpy(),\n                                ex[\"T14\"].float().cpu().numpy(),\n                                ex[\"tau\"].float().cpu().numpy(),\n                                ex[\"u\"].float().cpu().numpy()], 1))\n    y = np.concatenate(ys) * PPM_SCALE\n    mu = np.concatenate(mus) * PPM_SCALE\n    sg = np.clip(np.concatenate(sgs) * PPM_SCALE, 1.0, None)\n    l_ideal = -0.5 * (np.log(2 * np.pi) + np.log(GLL_IDEAL_SIGMA_PPM ** 2)) * y.size\n    l_ref = float(np.sum(gll_ppm(y, np.broadcast_to(ref_mean_ppm, y.shape),\n                                 np.broadcast_to(ref_std_ppm, y.shape))))\n    score = (float(np.sum(gll_ppm(y, mu, sg))) - l_ref) / (l_ideal - l_ref)\n    out = {\"GLL_score\": float(score), \"ECE\": ece(y, mu, sg),\n           \"RMSE_ppm\": float(np.sqrt(np.mean((mu - y) ** 2))),\n           \"mean_sigma_ppm\": float(sg.mean()),\n           \"cov68\": float((np.abs((y - mu) / sg) <= 1).mean()),\n           \"level_err_ppm\": float((mu - y).mean(1).mean())}\n    if th:\n        T = np.concatenate(th)\n        out[\"theta_median\"] = {\"t0\": float(np.median(T[:, 0])),\n                               \"T14\": float(np.median(T[:, 1])),\n                               \"tau\": float(np.median(T[:, 2])),\n                               \"u\": float(np.median(T[:, 3]))}\n    model.train()\n    return out\n\n\n# =============================================================================\n# MAIN\n# =============================================================================\n\ndef main():\n    torch.manual_seed(TRAIN_SEED)\n    torch.cuda.manual_seed_all(TRAIN_SEED)\n    np.random.seed(TRAIN_SEED)\n    random.seed(TRAIN_SEED)\n\n    banner(\"PC-SBN\")\n    print(f\"  DTP {USE_DTP}   encoder {USE_ENCODER}   head {HEAD_TYPE}\"\n          + (f\" K={K_BASIS} cov={COV_TYPE}\" if HEAD_TYPE == \"basis\" else \"\"))\n    print(f\"  fold {FOLD}   seed {TRAIN_SEED}   dropout {DROPOUT_P}   \"\n          f\"epochs {NUM_EPOCHS}   beta {BETA_NLL}\")\n    print(f\"  tag  {DATA_TAG}\")\n    print(f\"  repo {CHECKPOINT_REPO}\")\n    print(f\"  cache {CACHE_REPO}\")\n\n    train_df = pd.read_csv(TRAIN_CSV)\n    all_ids = train_df[\"planet_id\"].astype(str).tolist()\n    wl_cols = [c for c in train_df.columns if c != \"planet_id\"]\n\n    rng = np.random.RandomState(42)\n    shuffled = all_ids.copy()\n    rng.shuffle(shuffled)\n    if FOLD < 0:\n        n_val = max(1, int(len(shuffled) * VAL_FRACTION))\n        val_ids, train_ids = shuffled[:n_val], shuffled[n_val:]\n        print(f\"  split: historical single split\")\n    else:\n        blocks = np.array_split(np.arange(len(shuffled)), N_FOLDS)\n        vp = set(int(i) for i in blocks[FOLD])\n        val_ids = [p for j, p in enumerate(shuffled) if j in vp]\n        train_ids = [p for j, p in enumerate(shuffled) if j not in vp]\n        print(f\"  split: {N_FOLDS}-fold CV, fold {FOLD}\")\n    print(f\"Train planets: {len(train_ids)}  Val planets: {len(val_ids)}\")\n\n    # ---- preflight: the v3 cache must be complete ----\n    local = {p.stem for p in Path(LOCAL_CACHE_DIR).glob(\"*.npz\")} \\\n        if Path(LOCAL_CACHE_DIR).exists() else set()\n    try:\n        remote = {f[:-4] for f in _hf_api.list_repo_files(CACHE_REPO, repo_type=\"dataset\")\n                  if f.endswith(\".npz\")}\n    except Exception as e:\n        remote = set()\n        print(f\"  [preflight] cannot list {CACHE_REPO}: {type(e).__name__}\")\n    need = set(train_ids) | set(val_ids)\n    miss = sorted(need - (local | remote))\n    print(f\"  [preflight] cache planets available: {len(need) - len(miss)} / {len(need)}\")\n    if miss:\n        raise RuntimeError(\n            f\"v3 cache incomplete: {len(miss)} planets missing (e.g. {miss[:5]}). \"\n            f\"Run 13_add_noise_to_cache.py until it reports 1100 / 1100.\")\n\n    ref_mean_ppm, ref_std_ppm = M.compute_target_reference_stats(\n        train_df[train_df[\"planet_id\"].astype(str).isin(train_ids)], wl_cols)\n\n    star = pd.read_csv(STAR_INFO_CSV)\n    star[\"planet_id\"] = star[\"planet_id\"].astype(str)\n    tsi = star[star[\"planet_id\"].isin(train_ids)]\n    meta_mean = tsi[METADATA_COLS].values.astype(np.float32).mean(0)\n    meta_std = tsi[METADATA_COLS].values.astype(np.float32).std(0) + 1e-8\n\n    tr_ds = PCSBNDataset(train_ids, train_df, star, meta_mean, meta_std)\n    va_ds = PCSBNDataset(val_ids, train_df, star, meta_mean, meta_std)\n    te_ds = PCSBNDataset(train_ids[:min(TRAIN_EVAL_SUBSET_SIZE, len(train_ids))],\n                         train_df, star, meta_mean, meta_std)\n    tl = DataLoader(tr_ds, batch_size=BATCH_SIZE, shuffle=True,\n                    num_workers=NUM_WORKERS, pin_memory=True, drop_last=True)\n    vl = DataLoader(va_ds, batch_size=BATCH_SIZE, shuffle=False,\n                    num_workers=NUM_WORKERS, pin_memory=True)\n    el = DataLoader(te_ds, batch_size=BATCH_SIZE, shuffle=False,\n                    num_workers=NUM_WORKERS, pin_memory=True)\n\n    model = PCSBN().to(DEVICE)\n    # the plain head needs a physical sigma starting point (see warm_init).\n    # mu starts at the train mean level only when there is no physics path to\n    # start from; with DTP the physics already supplies the level.\n    _mu0 = np.zeros(N_WAVE) if USE_DTP else (ref_mean_ppm / PPM_SCALE)\n    _sg0 = np.full(N_WAVE, 300.0 / PPM_SCALE)          # 300 ppm, a typical scale\n    if model.head.warm_init(_mu0, _sg0):\n        print(f\"  [warm init] plain head: sigma -> 300 ppm, \"\n              f\"mu bias -> {'0 (DTP supplies the level)' if USE_DTP else 'train mean'}\")\n    n_par = sum(p.numel() for p in model.parameters())\n    n_geom = sum(p.numel() for p in model.geom.parameters()) if USE_DTP else 0\n    n_enc = sum(p.numel() for p in model.encoder.parameters()) if USE_ENCODER else 0\n    n_head = sum(p.numel() for p in model.head.parameters())\n    print(f\"\\n  params total {n_par:,}   geometry {n_geom:,}   \"\n          f\"encoder {n_enc:,}   head {n_head:,}\")\n\n    # ---- SANITY: with a zero-initialised head, mu must equal d_phys ----\n    if USE_DTP and HEAD_TYPE in (\"basis\", \"level\"):\n        model.eval()\n        with torch.no_grad():\n            b = next(iter(vl))\n            mu0, sg0, ex0 = model(b[\"x\"].to(DEVICE), b[\"noise_rel\"].to(DEVICE))\n            dev = (mu0 - ex0[\"d_phys\"]).abs().max().item()\n            print(f\"  [sanity] max |mu - d_phys| at init = {dev:.3e}  \"\n                  f\"({'OK' if dev < 1e-6 else 'FAIL -- head is not zero-init'})\")\n            if dev >= 1e-6:\n                raise RuntimeError(\"the correction head is not zero-initialised\")\n            dp = ex0[\"d_phys\"].float().cpu().numpy() * PPM_SCALE\n            yt = b[\"target\"].numpy() * PPM_SCALE\n            print(f\"  [sanity] d_phys level median {np.median(dp.mean(1)):8.1f} ppm   \"\n                  f\"label median {np.median(yt.mean(1)):8.1f} ppm\")\n            print(f\"  [sanity] corr(d_phys level, label level) = \"\n                  f\"{np.corrcoef(dp.mean(1), yt.mean(1))[0,1]:+.4f}   \"\n                  f\"(a physics path that works starts well above +0.9)\")\n            print(f\"  [sanity] effective in-transit bins, median \"\n                  f\"{float(ex0['s2'].median()):.1f} of {SEQ_LEN}\")\n        model.train()\n\n    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=NUM_EPOCHS)\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=(DEVICE.type == \"cuda\"))\n    ema = EMA(model)\n\n    ckpt_local = f\"/kaggle/working/pcsbn_{DATA_TAG}_latest.pt\"\n    best_local = f\"/kaggle/working/pcsbn_{DATA_TAG}_best.pt\"\n    start_ep, best, best_ep = 0, -1e9, -1\n    if RESUME_IF_AVAILABLE and os.path.exists(ckpt_local):\n        try:\n            ck = torch.load(ckpt_local, map_location=DEVICE, weights_only=False)\n            model.load_state_dict(ck[\"model_state\"])\n            ema.shadow.load_state_dict(ck[\"ema_state\"])\n            opt.load_state_dict(ck[\"optimizer_state\"])\n            start_ep, best, best_ep = ck[\"epoch\"] + 1, ck[\"best\"], ck[\"best_epoch\"]\n            print(f\"  [resume] from epoch {start_ep} (best {best:.4f} @ {best_ep})\")\n        except Exception as e:\n            print(f\"  [resume] failed ({type(e).__name__}) -- starting fresh\")\n\n    banner(\"TRAINING\")\n    t_start = time.time()\n    for ep in range(start_ep, NUM_EPOCHS):\n        if ep == MSE_WARMUP_EPOCHS:\n            for g in opt.param_groups:\n                g[\"lr\"] = LR * TRANSITION_LR_DECAY\n            sched = torch.optim.lr_scheduler.CosineAnnealingLR(\n                opt, T_max=max(NUM_EPOCHS - MSE_WARMUP_EPOCHS, 1))\n        model.train()\n        tot = 0.0\n        for b in tl:\n            x = b[\"x\"].to(DEVICE, non_blocking=True)\n            nr = b[\"noise_rel\"].to(DEVICE, non_blocking=True)\n            yt = b[\"target\"].to(DEVICE, non_blocking=True)\n            tr = b[\"theta_ref\"].to(DEVICE, non_blocking=True)\n            opt.zero_grad(set_to_none=True)\n            with torch.autocast(\"cuda\", dtype=torch.float16,\n                                enabled=(DEVICE.type == \"cuda\")):\n                mu, sg, ex = model(x, nr, b[\"metadata\"].to(DEVICE))\n                loss = (F.mse_loss(mu, yt) if ep < MSE_WARMUP_EPOCHS\n                        else beta_nll(mu, sg, yt))\n                if ep < THETA_SUP_EPOCHS and USE_DTP:\n                    loss = loss + THETA_SUP_WEIGHT * theta_supervision(ex, tr)\n            scaler.scale(loss).backward()\n            scaler.unscale_(opt)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP_NORM)\n            scaler.step(opt)\n            scaler.update()\n            ema.update(model)\n            tot += float(loss.detach())\n        sched.step()\n\n        if ep % 5 == 0 or ep == NUM_EPOCHS - 1:\n            vm = evaluate(ema.shadow, vl, ref_mean_ppm, ref_std_ppm)\n            msg = (f\"  ep {ep:>3d}  loss {tot/max(len(tl),1):+.5f}  \"\n                   f\"val GLL {vm['GLL_score']:.4f}  ECE {vm['ECE']:.4f}  \"\n                   f\"RMSE {vm['RMSE_ppm']:.1f}  cov68 {vm['cov68']:.3f}\")\n            if \"theta_median\" in vm:\n                t = vm[\"theta_median\"]\n                msg += (f\"  theta t0={t['t0']:.0f} T14={t['T14']:.0f} \"\n                        f\"tau={t['tau']:.2f} u={t['u']:.2f}\")\n            print(msg, flush=True)\n            if vm[\"GLL_score\"] > best:\n                best, best_ep = vm[\"GLL_score\"], ep\n                torch.save({\"model_state\": model.state_dict(),\n                            \"ema_state\": ema.shadow.state_dict(),\n                            \"epoch\": ep, \"val_metrics\": vm,\n                            \"config\": _cfg_dict()}, best_local)\n                print(f\"  [best] {best:.4f} @ {ep} -- saved\")\n            torch.save({\"model_state\": model.state_dict(),\n                        \"ema_state\": ema.shadow.state_dict(),\n                        \"optimizer_state\": opt.state_dict(), \"epoch\": ep,\n                        \"best\": best, \"best_epoch\": best_ep,\n                        \"config\": _cfg_dict()}, ckpt_local)\n            if ep % HF_UPLOAD_EVERY_N_EPOCHS == 0:\n                _upload(ckpt_local, f\"pcsbn_{DATA_TAG}_latest.pt\")\n\n    banner(\"FINAL\")\n    vm = evaluate(ema.shadow, vl, ref_mean_ppm, ref_std_ppm)\n    tm = evaluate(ema.shadow, el, ref_mean_ppm, ref_std_ppm)\n    print(\"  val   (EMA, last epoch):\\n\" + json.dumps(vm, indent=2))\n    print(f\"  train subset GLL {tm['GLL_score']:.4f}  \"\n          f\"(generalisation gap {tm['GLL_score'] - vm['GLL_score']:+.4f})\")\n    fin = f\"/kaggle/working/pcsbn_{DATA_TAG}_final.pt\"\n    torch.save({\"model_state\": ema.shadow.state_dict(), \"val_metrics\": vm,\n                \"train_subset_metrics\": tm, \"config\": _cfg_dict()}, fin)\n    _upload(fin, f\"pcsbn_{DATA_TAG}_final.pt\")\n    if os.path.exists(best_local):\n        _upload(best_local, f\"pcsbn_{DATA_TAG}_best.pt\")\n    with open(f\"/kaggle/working/pcsbn_metrics_{DATA_TAG}.json\", \"w\") as f:\n        json.dump({\"val_last\": vm, \"train_subset_last\": tm,\n                   \"best_val_gll\": best, \"best_epoch\": best_ep,\n                   \"n_params\": n_par, \"config\": _cfg_dict()}, f, indent=2)\n    print(f\"\\n  BEST val GLL {best:.4f} @ epoch {best_ep}\")\n    print(f\"  wall time {(time.time()-t_start)/60:.1f} min\")\n    print(f\"\\n  Next: 21_assemble_oof_pcsbn.py, then 19_paired_compare.py\")\n    print(f\"  against the locked baseline (pooled OOF 0.4290 / 0.4890 excl 1).\")\n\n\ndef _cfg_dict():\n    return {\"USE_DTP\": USE_DTP, \"USE_ENCODER\": USE_ENCODER,\n            \"HEAD_TYPE\": HEAD_TYPE, \"K_BASIS\": K_BASIS, \"COV_TYPE\": COV_TYPE,\n            \"FOLD\": FOLD, \"N_FOLDS\": N_FOLDS, \"TRAIN_SEED\": TRAIN_SEED,\n            \"DROPOUT_P\": DROPOUT_P, \"NUM_EPOCHS\": NUM_EPOCHS,\n            \"BETA_NLL\": BETA_NLL, \"hidden_dim\": HIDDEN_DIM,\n            \"kernels\": KERNEL_SIZES, \"tag\": DATA_TAG, \"cache\": CACHE_REPO}\n\n\ndef _upload(path, name):\n    try:\n        create_repo(repo_id=CHECKPOINT_REPO, repo_type=\"model\",\n                    exist_ok=True, private=False)\n        _hf_api.upload_file(path_or_fileobj=path, path_in_repo=name,\n                            repo_id=CHECKPOINT_REPO, repo_type=\"model\")\n    except Exception as e:\n        print(f\"  [warn] HF upload failed for {name}: {type(e).__name__}\")\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-09-25T05:21:13.810251Z","iopub.execute_input":"2026-09-25T05:21:13.810528Z","iopub.status.idle":"2026-09-25T05:21:13.832321Z","shell.execute_reply.started":"2026-09-25T05:21:13.810509Z","shell.execute_reply":"2026-09-25T05:21:13.831506Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!/usr/bin/env python\n# =============================================================================\n# ARIEL 2025 -- SCRIPT 22 : RUN THE PC-SBN ABLATION LADDER\n# =============================================================================\n# Each rung differs from the baseline in exactly ONE thing, and every rung is\n# trained on the same 5 folds, so every comparison is paired over the same 1100\n# out-of-fold planets.\n#\n#   rung            DTP  encoder  head    what it isolates\n#   1 baseline      no   yes      plain   (already measured: 0.4290 / 0.4890)\n#   2 physics only  yes  no       level   what the physics alone can do\n#   3 DTP + plain   yes  yes      plain   CONTRIBUTION 1 on its own\n#   4 struct only   no   yes      basis   CONTRIBUTION 2 on its own\n#   5 full          yes  yes      basis   both, and whether they interact\n#\n# ORDER MATTERS FOR COMPUTE. Rung 5 runs FIRST: if the full model does not work\n# there is no point measuring its parts. Then 3 and 4, the two isolating\n# ablations that carry the contribution claims. Then 2, which is expected to be\n# worse than the baseline and is there to show what the physics supplies on its\n# own. Covariance-form and K sweeps come last, inside the winner only.\n#\n# Budget: ~20 min per fold, so ~1.7 h per rung and ~7 h for rungs 5, 3, 4, 2.\n# Resumable: a fold whose exact checkpoint is already on HuggingFace is skipped,\n# so re-running this cell after a session dies continues where it stopped.\n# -----------------------------------------------------------------------------\n# HOW TO USE  (GPU T4 x2, Internet ON, HF_TOKEN, competition data attached,\n#              moderntcn-src attached)\n#\n#   Cell 1:  %%writefile /kaggle/working/20_train_pcsbn.py\n#            <paste the whole of 20_train_pcsbn.py below that line>\n#            -- output should be only \"Writing /kaggle/working/20_train_pcsbn.py\"\n#\n#   Cell 2:  <paste this script>\n#\n# Then in a second notebook: script 21, then script 19.\n# =============================================================================\n\nimport os\nimport re\nimport sys\nimport glob\nimport time\nimport subprocess\n\n# =============================================================================\n# CONFIG\n# =============================================================================\n\nCOMMON = {\"TRAIN_SEED\": \"1\", \"DROPOUT_P\": \"0.15\", \"NUM_EPOCHS\": \"600\",\n          \"BETA_NLL\": \"0.5\", \"N_FOLDS\": \"5\", \"THETA_SUP_EPOCHS\": \"40\"}\n\nRUNGS = [\n    (\"5 full\",         {\"USE_DTP\": \"1\", \"USE_ENCODER\": \"1\", \"HEAD_TYPE\": \"basis\",\n                        \"K_BASIS\": \"24\", \"COV_TYPE\": \"diag\"}),\n    (\"3 DTP + plain\",  {\"USE_DTP\": \"1\", \"USE_ENCODER\": \"1\", \"HEAD_TYPE\": \"plain\"}),\n    (\"4 struct only\",  {\"USE_DTP\": \"0\", \"USE_ENCODER\": \"1\", \"HEAD_TYPE\": \"basis\",\n                        \"K_BASIS\": \"24\", \"COV_TYPE\": \"diag\"}),\n    (\"2 physics only\", {\"USE_DTP\": \"1\", \"USE_ENCODER\": \"0\", \"HEAD_TYPE\": \"level\"}),\n]\nFOLDS = [0, 1, 2, 3, 4]\n\nTRAIN_SCRIPT = \"/kaggle/working/20_train_pcsbn.py\"\nMAX_HOURS = 11.0\nSKIP_IF_ON_HF = True\nHF_USER = \"maariaaa12\"\n\n\ndef banner(t):\n    print(f\"\\n{'#'*78}\\n# {t}\\n{'#'*78}\", flush=True)\n\n\nT0 = time.time()\n\nif not os.path.isfile(TRAIN_SCRIPT):\n    found = sorted(glob.glob(\"/kaggle/working/*pcsbn*.py\")\n                   + glob.glob(\"/kaggle/working/*train*.py\"))\n    raise FileNotFoundError(\n        f\"{TRAIN_SCRIPT} not found.\\n\"\n        f\"  In the cell above, put this as the FIRST line:\\n\"\n        f\"      %%writefile {TRAIN_SCRIPT}\\n\"\n        f\"  then paste the whole trainer under it and run that cell once.\\n\"\n        f\"  python files currently around: {found}\")\n\n\n# -----------------------------------------------------------------------------\n# TAG LOGIC -- must reproduce the trainer's _t / _repo_tag EXACTLY.\n# A mismatch silently retrains everything, which is why it is verified below.\n# -----------------------------------------------------------------------------\ndef tag_for(cfg, fold=None):\n    t = \"pcsbn\"\n    t += \"_dtp\" if cfg[\"USE_DTP\"] == \"1\" else \"_nodtp\"\n    t += \"_enc\" if cfg[\"USE_ENCODER\"] == \"1\" else \"_noenc\"\n    t += f\"_{cfg['HEAD_TYPE']}\"\n    if cfg[\"HEAD_TYPE\"] == \"basis\":\n        t += f\"_k{int(cfg.get('K_BASIS', 24))}\"\n        if cfg.get(\"COV_TYPE\", \"diag\") != \"diag\":\n            t += f\"_{cfg['COV_TYPE']}\"\n    b = float(cfg.get(\"BETA_NLL\", 0.5))\n    if b != 0.5:\n        t += \"_b\" + str(b).replace(\".\", \"\")\n    t += f\"_s{cfg['TRAIN_SEED']}\"\n    d = float(cfg.get(\"DROPOUT_P\", 0.15))\n    if d != 0.15:\n        t += \"_d\" + str(d).replace(\".\", \"\")\n    base = t\n    if fold is not None:\n        t += f\"_f{fold}\"\n    return t, base\n\n\ndef repo_for(base):\n    return f\"{HF_USER}/ariel-2025-{base.replace('_', '-')}-5fold\"\n\n\n# -----------------------------------------------------------------------------\n# what is already trained\n# -----------------------------------------------------------------------------\ndone = set()\nif SKIP_IF_ON_HF:\n    try:\n        from huggingface_hub import HfApi, login\n        try:\n            from kaggle_secrets import UserSecretsClient\n            login(token=UserSecretsClient().get_secret(\"HF_TOKEN\"))\n        except Exception:\n            pass\n        api = HfApi()\n        cache = {}\n        for _, cfg in RUNGS:\n            full = {**COMMON, **cfg}\n            for f in FOLDS:\n                tag, base = tag_for(full, f)\n                r = repo_for(base)\n                if r not in cache:\n                    try:\n                        cache[r] = set(api.list_repo_files(r, repo_type=\"model\"))\n                    except Exception:\n                        cache[r] = set()\n                # the FINAL checkpoint is the protocol's selection rule, so that\n                # is what \"done\" means -- a _best.pt alone is not enough\n                if f\"pcsbn_{tag}_final.pt\" in cache[r]:\n                    done.add(tag)\n    except Exception as e:\n        print(f\"  [warn] could not check HuggingFace ({type(e).__name__}); \"\n              f\"running everything\")\n\nJOBS = []\nfor name, cfg in RUNGS:\n    for f in FOLDS:\n        tag, base = tag_for({**COMMON, **cfg}, f)\n        JOBS.append((name, cfg, f, tag, base))\n\nbanner(f\"PLAN -- {len(RUNGS)} rungs x {len(FOLDS)} folds = {len(JOBS)} runs, \"\n       f\"{len(done)} already trained\")\nfor name, cfg, f, tag, base in JOBS:\n    print(f\"  {name:<15s} fold {f}   \"\n          f\"{'SKIP (on HF)' if tag in done else 'run':<13s} {tag}\")\nprint(f\"\\n  budget {MAX_HOURS} h; ~20 min per run, so expect \"\n       f\"~{len([j for j in JOBS if j[3] not in done]) * 20 / 60:.1f} h of work\")\n\nresults = []\nfor i, (name, cfg, f, tag, base) in enumerate(JOBS, 1):\n    if tag in done:\n        continue\n    if (time.time() - T0) / 3600 > MAX_HOURS:\n        print(f\"\\n  time budget {MAX_HOURS} h reached -- stopping. Re-run this \"\n              f\"cell to continue; finished folds are skipped.\")\n        break\n\n    banner(f\"{name}   fold {f}   ({i}/{len(JOBS)})   {tag}\")\n    # a stale *_latest.pt from ANOTHER run would make the trainer resume into the\n    # wrong weights. This run's own latest is kept, so a killed session resumes.\n    for p in glob.glob(\"/kaggle/working/pcsbn_*_latest.pt\"):\n        if tag not in os.path.basename(p):\n            os.remove(p)\n\n    env = dict(os.environ)\n    env.update({f\"ARIEL_{k}\": v for k, v in {**COMMON, **cfg, \"FOLD\": str(f)}.items()})\n    t0 = time.time()\n    proc = subprocess.run([sys.executable, TRAIN_SCRIPT], env=env,\n                          stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n    mins = (time.time() - t0) / 60\n    out = proc.stdout or \"\"\n\n    # sanity lines worth surfacing even on success\n    sp = re.search(r\"Train planets: (\\d+)\\s+Val planets: (\\d+)\", out)\n    if sp:\n        print(f\"  split: train {sp.group(1)}  val {sp.group(2)}\")\n        if int(sp.group(2)) < 200:\n            print(f\"  [!!] validated on {sp.group(2)} planets -- the FOLD switch \"\n                  f\"did not take effect\")\n    for pat in (r\"\\[sanity\\] max \\|mu - d_phys\\|.*\",\n                r\"\\[sanity\\] corr\\(d_phys level, label level\\).*\",\n                r\"\\[sanity\\] effective in-transit bins.*\",\n                r\"\\[warm init\\].*\"):\n        m = re.search(pat, out)\n        if m:\n            print(f\"  {m.group(0).strip()}\")\n\n    best = re.search(r\"BEST val GLL ([\\-\\d.]+) @ epoch (\\d+)\", out)\n    if proc.returncode != 0 or not best:\n        print(out[-5000:])\n        print(f\"  [FAILED] returncode {proc.returncode} after {mins:.1f} min\")\n        results.append((name, f, tag, None, None, mins, \"FAILED\"))\n        continue\n\n    gll, ep = float(best.group(1)), int(best.group(2))\n    gap = re.search(r\"generalisation gap ([+\\-\\d.]+)\", out)\n    th = re.search(r'\"theta_median\".*?\"t0\": ([\\d.\\-]+).*?\"T14\": ([\\d.\\-]+)', out, re.S)\n    note = f\"gap {gap.group(1)}\" if gap else \"\"\n    if th:\n        note += f\"  theta t0={float(th.group(1)):.0f} T14={float(th.group(2)):.0f}\"\n    print(f\"  done in {mins:.1f} min   best val GLL {gll:.4f} @ epoch {ep}   {note}\")\n    results.append((name, f, tag, gll, ep, mins, note))\n    print(f\"  elapsed {(time.time() - T0)/60:.0f} min total\", flush=True)\n\n# =============================================================================\n# SUMMARY\n# =============================================================================\nbanner(\"SUMMARY\")\nprint(f\"  {'rung':<15s} {'fold':>4s} {'GLL':>8s} {'epoch':>6s} {'min':>6s}  note\")\nfor name, f, tag, gll, ep, mins, note in results:\n    g = f\"{gll:.4f}\" if gll is not None else \"FAILED\"\n    e = str(ep) if ep is not None else \"--\"\n    print(f\"  {name:<15s} {f:>4d} {g:>8s} {e:>6s} {mins:>6.1f}  {note}\")\n\nok = [r for r in results if r[3] is not None]\nif ok:\n    import statistics as st\n    print()\n    for name, _ in RUNGS:\n        v = [r[3] for r in ok if r[0] == name]\n        if len(v) > 1:\n            print(f\"  {name:<15s} per-fold mean {st.mean(v):.4f}  \"\n                  f\"sd {st.stdev(v):.4f}  n={len(v)}\")\n    print(\"\\n  These are TRAINING-REPORTED per-fold numbers and are NOT the result.\")\n    print(\"  The fold sd is fold difficulty, not an error bar.\")\n\nn_have = len(done) + len(ok)\nprint(f\"\\n  {n_have}/{len(JOBS)} runs complete.\")\nif n_have >= len(JOBS):\n    print(\"  All rungs trained. Next, in a second notebook:\")\n    print(\"    script 21  assemble out-of-fold predictions for every rung\")\n    print(\"    script 19  paired comparison against the locked baseline\")\n    print(\"               (0.4290 all-1100 / 0.4890 excl 1, MDE 0.0026)\")\nelse:\n    print(\"  Re-run this cell to continue.\")\nprint(f\"\\n  total wall time {(time.time() - T0)/60:.0f} min\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-09-25T05:22:02.275195Z","iopub.execute_input":"2026-09-25T05:22:02.275947Z","iopub.status.idle":"2026-09-25T05:24:23.93207Z","shell.execute_reply.started":"2026-09-25T05:22:02.275919Z","shell.execute_reply":"2026-09-25T05:24:23.931006Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# NOTEBOOK E — score karo","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/20_train_pcsbn.py\n#!/usr/bin/env python\n# =============================================================================\n# ARIEL 2025 -- SCRIPT 20 : PC-SBN TRAINER\n#   contribution 1  Differentiable Transit Projection (DTP)\n#   contribution 2  Structure-aware output head\n# =============================================================================\n# THE MODEL\n#\n#   x(t,lambda) = flux/baseline - 1        (Cottaar step-5 v3 cache)\n#   physically   x ~= -s(t; theta) * d(lambda)\n#\n#   1. white-light curve      w(t) = mean_lambda x(t,lambda)\n#   2. GEOMETRY CNN           w -> theta = (t0, T14, tau, u)          [NEW 1]\n#   3. TRANSIT SHAPE          s(t; theta) in [0,1], analytic, 0 params\n#   4. CLOSED-FORM PROJECTION d_phys(lambda), var(lambda)             [NEW 1]\n#   5. RESIDUAL               r = x + s (x) d_phys        -> encoder\n#   6. ModernTCN ENCODER      r -> z in R^192             (frozen config)\n#   7. STRUCTURE-AWARE HEAD   z, d_phys, var -> mu, sigma             [NEW 2]\n#\n# THE PROJECTION, exactly. The instrument noise in the v3 cache is\n# per-wavelength and constant in time, so a weighted least squares over t at\n# fixed lambda has a constant weight -- which CANCELS in the depth estimate and\n# survives only in its variance (verified to 5e-18):\n#\n#     d_phys(lambda) = sum_t s_t * (-x(t,lambda))  /  sum_t s_t^2\n#     var(lambda)    = noise_rel(lambda)^2         /  sum_t s_t^2\n#\n# sum_t s_t^2 is one scalar per planet: the effective number of in-transit bins.\n# noise_rel is known to carry a uniform 4.81x scale error; it cancels in d_phys\n# and is absorbed by the learned scale c in the head, so no cache rebuild is\n# needed. Its RELATIVE structure across wavelength is correct (corr 0.976) and\n# its within-planet spread is p90/p10 = 3.0, so the weights do real work.\n#\n# THE HEAD, exactly.\n#     mu(lambda)     = d_phys(lambda) + mu_L + sum_k a_k B_k(lambda)\n#     sigma^2(lambda) = c^2 var(lambda) + sigma_L^2 + sum_k B_k(lambda)^2 exp(2 l_k)\n# B is a FIXED cosine (DCT-II) basis, k = 1..K, so it is smooth by construction\n# and fitted on no data at all -- nothing to leak. The correction terms are\n# zero-initialised, so at step 0 the model outputs exactly d_phys.\n#\n# ABLATION SWITCHES (environment variables, all read by _env below)\n#   ARIEL_USE_DTP      1 physics path on   0 off (encoder sees the raw tensor)\n#   ARIEL_USE_ENCODER  1 encoder on        0 off (pure physics, rung 2)\n#   ARIEL_HEAD_TYPE    plain | level | basis\n#   ARIEL_K_BASIS      number of cosine basis functions (default 24)\n#   ARIEL_COV_TYPE     diag | chol        (basis-space covariance form)\n#   ARIEL_FOLD         -1 historical split, 0..4 five-fold\n#\n# The ablation ladder (each rung differs from the previous in ONE switch):\n#   1 baseline      USE_DTP=0 USE_ENCODER=1 HEAD_TYPE=plain   (= script 04)\n#   2 physics only  USE_DTP=1 USE_ENCODER=0 HEAD_TYPE=level\n#   3 DTP + plain   USE_DTP=1 USE_ENCODER=1 HEAD_TYPE=plain\n#   4 struct only   USE_DTP=0 USE_ENCODER=1 HEAD_TYPE=basis\n#   5 full          USE_DTP=1 USE_ENCODER=1 HEAD_TYPE=basis\n#\n# -----------------------------------------------------------------------------\n# NOTEBOOK SETUP  (GPU T4 x2, Internet ON, HF_TOKEN secret)\n#   Add Input : Competitions -> ariel-data-challenge-2025\n#   Add Input : Datasets -> moderntcn-src      (for the ModernTCN blocks)\n# =============================================================================\n\nimport os\nimport gc\nimport glob\nimport json\nimport math\nimport random\nimport sys\nimport time\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom scipy.stats import norm as _norm\nfrom torch.utils.data import Dataset, DataLoader\nfrom huggingface_hub import HfApi, create_repo, hf_hub_download\n\n\n# =============================================================================\n# CONFIG\n# =============================================================================\n\ndef _env(name, default, cast=str):\n    v = os.environ.get(\"ARIEL_\" + name)\n    return default if v is None else cast(v)\n\n\ndef _envb(name, default):\n    return bool(int(_env(name, int(default), int)))\n\n\nUSE_DTP      = _envb(\"USE_DTP\", True)\nUSE_ENCODER  = _envb(\"USE_ENCODER\", True)\nHEAD_TYPE    = _env(\"HEAD_TYPE\", \"basis\")        # plain | level | basis\nK_BASIS      = _env(\"K_BASIS\", 24, int)\nCOV_TYPE     = _env(\"COV_TYPE\", \"diag\")          # diag | chol\nFOLD         = _env(\"FOLD\", -1, int)\nN_FOLDS      = _env(\"N_FOLDS\", 5, int)\nTRAIN_SEED   = _env(\"TRAIN_SEED\", 1, int)\nDROPOUT_P    = _env(\"DROPOUT_P\", 0.15, float)\nNUM_EPOCHS   = _env(\"NUM_EPOCHS\", 600, int)\nBETA_NLL     = _env(\"BETA_NLL\", 0.5, float)\n# theta is supervised against the cache's stored transit box for this many\n# epochs. Without it the trapezoid's gradient is sparse at initialisation (only\n# the ingress/egress bins carry signal), so the geometry CNN can sit in a bad\n# basin. The cache values come from a label-free box search, so this adds no\n# label information.\nTHETA_SUP_EPOCHS = _env(\"THETA_SUP_EPOCHS\", 40, int)\nTHETA_SUP_WEIGHT = _env(\"THETA_SUP_WEIGHT\", 1.0, float)\n\nif HEAD_TYPE not in (\"plain\", \"level\", \"basis\"):\n    raise ValueError(f\"HEAD_TYPE must be plain|level|basis, got {HEAD_TYPE!r}\")\nif COV_TYPE not in (\"diag\", \"chol\"):\n    raise ValueError(f\"COV_TYPE must be diag|chol, got {COV_TYPE!r}\")\nif not USE_DTP and not USE_ENCODER:\n    raise ValueError(\"USE_DTP=0 and USE_ENCODER=0 leaves no model at all\")\nif not USE_DTP and HEAD_TYPE == \"level\":\n    raise ValueError(\"HEAD_TYPE=level without DTP has nothing to correct; \"\n                     \"use plain or basis\")\n\nCACHE_REPO = \"maariaaa12/ariel-2025-cottaar-step5-cache-v3\"\nLOCAL_CACHE_DIR = \"/kaggle/working/cottaar_step5_cache_v3\"\nTRAIN_CSV = \"/kaggle/input/competitions/ariel-data-challenge-2025/train.csv\"\nSTAR_INFO_CSV = \"/kaggle/input/competitions/ariel-data-challenge-2025/train_star_info.csv\"\nMETADATA_COLS = [\"Rs\", \"Ms\", \"Ts\", \"Mp\", \"e\", \"P\", \"sma\", \"i\"]\nSRC_NAME = \"train_spectrum_moderntcn_v2_robust_cache.py\"\n\nN_WAVE = 283\nSEQ_LEN = 187\nPPM_SCALE = 1e6\nSIGMA_FLOOR = 1e-6          # 1 ppm, native (Rp/Rs)^2 units\nGLL_IDEAL_SIGMA_PPM = 10.0\nECE_N_BINS = 10\nVAL_FRACTION = 0.1\nTRAIN_EVAL_SUBSET_SIZE = 110\n\n# float16 safety: the encoder runs under autocast, and near 1.0 float16 resolves\n# only ~500-1000 ppm. The residual is already centred near 0, so it only needs\n# the x100 scale to sit clear of the subnormal range.\nENCODER_SCALE = 100.0\n\nBATCH_SIZE = 16\nLR = 2e-4\nWEIGHT_DECAY = 9e-4\nHIDDEN_DIM = 192\nKERNEL_SIZES = [15, 15, 15]\nN_BLOCKS = [1, 1, 1]\nGRAD_CLIP_NORM = 1.0\nMSE_WARMUP_EPOCHS = 20\nTRANSITION_LR_DECAY = 0.3\nEMA_DECAY = 0.995\nNUM_WORKERS = 2\nHF_UPLOAD_EVERY_N_EPOCHS = 10\nRESUME_IF_AVAILABLE = True\n\n# ---- tag ----\n_t = \"pcsbn\"\n_t += \"_dtp\" if USE_DTP else \"_nodtp\"\n_t += \"_enc\" if USE_ENCODER else \"_noenc\"\n_t += f\"_{HEAD_TYPE}\"\nif HEAD_TYPE == \"basis\":\n    _t += f\"_k{K_BASIS}\"\n    if COV_TYPE != \"diag\":\n        _t += f\"_{COV_TYPE}\"\nif BETA_NLL != 0.5:\n    _t += \"_b\" + str(BETA_NLL).replace(\".\", \"\")\n_t += f\"_s{TRAIN_SEED}\"\nif DROPOUT_P != 0.15:\n    _t += \"_d\" + str(DROPOUT_P).replace(\".\", \"\")\n_repo_tag = _t\nif FOLD >= 0:\n    if not (0 <= FOLD < N_FOLDS):\n        raise ValueError(f\"FOLD must be in [0,{N_FOLDS}), got {FOLD}\")\n    _t += f\"_f{FOLD}\"\n    _repo_tag += \"_5fold\"\nDATA_TAG = _t\nCHECKPOINT_REPO = f\"maariaaa12/ariel-2025-{_repo_tag.replace('_', '-')}\"\n\nDEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n_hf_api = HfApi()\n\n\ndef banner(t):\n    print(f\"\\n{'='*78}\\n{t}\\n{'='*78}\", flush=True)\n\n\n# =============================================================================\n# THE ModernTCN BLOCKS COME FROM THE EXISTING TRAINER -- same blocks, so the\n# encoder is literally the frozen winner, not a re-implementation of it.\n# =============================================================================\n\ndef _load_src():\n    hits = []\n    for pat in (f\"/kaggle/input/**/{SRC_NAME}\", f\"/kaggle/working/{SRC_NAME}\"):\n        hits += sorted(glob.glob(pat, recursive=True))\n    if not hits:\n        raise FileNotFoundError(f\"{SRC_NAME} not found -- attach the moderntcn-src dataset.\")\n    d = os.path.dirname(hits[0])\n    if d not in sys.path:\n        sys.path.insert(0, d)\n    import importlib\n    m = importlib.import_module(SRC_NAME[:-3])\n    print(f\"  ModernTCN blocks from {hits[0]}\")\n    return m\n\n\nM = _load_src()\n\n\n# =============================================================================\n# DATASET  --  returns the RAW tensor. Scaling happens inside the model, so the\n# physics path works in native relative-flux units and only the encoder sees the\n# x100 version.\n# =============================================================================\n\nclass PCSBNDataset(Dataset):\n    def __init__(self, planet_ids, targets_df, star_info_df, meta_mean, meta_std):\n        self.pids = [str(p) for p in planet_ids]\n        self.dir = Path(LOCAL_CACHE_DIR)\n        self.dir.mkdir(parents=True, exist_ok=True)\n        t = targets_df.copy()\n        t[\"planet_id\"] = t[\"planet_id\"].astype(str)\n        self.targets = t.set_index(\"planet_id\")\n        s = star_info_df.copy()\n        s[\"planet_id\"] = s[\"planet_id\"].astype(str)\n        self.star = s.set_index(\"planet_id\")\n        self.meta_mean, self.meta_std = meta_mean, meta_std\n\n    def __len__(self):\n        return len(self.pids)\n\n    def _path(self, pid):\n        p = self.dir / f\"{pid}.npz\"\n        if p.exists():\n            return p\n        return Path(hf_hub_download(repo_id=CACHE_REPO, repo_type=\"dataset\",\n                                    filename=f\"{pid}.npz\", local_dir=str(self.dir)))\n\n    def __getitem__(self, i):\n        pid = self.pids[i]\n        z = np.load(self._path(pid))\n        x = z[\"spectrum_clean\"].astype(np.float32)             # (187, 283)\n        nr = z[\"noise_rel\"].astype(np.float32)                  # (283,)\n        # a zero or non-finite noise entry would make var(lambda) explode\n        nr = np.nan_to_num(nr, nan=0.0, posinf=0.0, neginf=0.0)\n        med = np.median(nr[nr > 0]) if (nr > 0).any() else 1e-4\n        nr = np.where(nr > 0, nr, med).astype(np.float32)\n        y = self.targets.loc[pid].values.astype(np.float32)      # (283,) native\n        meta = ((self.star.loc[pid, METADATA_COLS].values.astype(np.float32)\n                 - self.meta_mean) / self.meta_std)\n        c = float(z[\"transit_centre_bin\"]) if \"transit_centre_bin\" in z else SEQ_LEN / 2\n        w = float(z[\"transit_width_bin\"]) if \"transit_width_bin\" in z else SEQ_LEN / 3\n        return {\"planet_id\": pid,\n                \"x\": torch.from_numpy(x),\n                \"noise_rel\": torch.from_numpy(nr),\n                \"target\": torch.from_numpy(y),\n                \"metadata\": torch.from_numpy(meta),\n                \"theta_ref\": torch.tensor([c, w], dtype=torch.float32)}\n\n\n# =============================================================================\n# CONTRIBUTION 1 -- DIFFERENTIABLE TRANSIT PROJECTION\n# =============================================================================\n\nclass GeometryCNN(nn.Module):\n    \"\"\"White-light curve -> theta = (t0, T14, tau, u). ~20k parameters.\n\n    Only the white-light curve is used, never the labels. The outputs are\n    squashed into physically valid ranges so no epoch can produce a nonsensical\n    transit that would make the projection blow up.\n    \"\"\"\n\n    def __init__(self, seq_len=SEQ_LEN, ch=32):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Conv1d(1, ch, 7, padding=3), nn.BatchNorm1d(ch), nn.GELU(),\n            nn.Conv1d(ch, ch, 7, stride=2, padding=3), nn.BatchNorm1d(ch), nn.GELU(),\n            nn.Conv1d(ch, ch, 5, stride=2, padding=2), nn.BatchNorm1d(ch), nn.GELU(),\n            nn.Conv1d(ch, ch, 5, padding=2), nn.BatchNorm1d(ch), nn.GELU(),\n        )\n        self.head = nn.Linear(2 * ch, 4)\n        nn.init.zeros_(self.head.weight)\n        nn.init.zeros_(self.head.bias)      # starts at the centre of every range\n        self.seq_len = seq_len\n\n    def forward(self, w):\n        # w: (B, T). Normalise per planet so the CNN sees shape, not depth.\n        wc = w - w.mean(1, keepdim=True)\n        sd = wc.std(1, keepdim=True).clamp_min(1e-8)\n        h = self.net((wc / sd).unsqueeze(1))                       # (B, ch, T')\n        h = torch.cat([h.mean(-1), h.amax(-1)], dim=-1)            # (B, 2ch)\n        raw = self.head(h)\n        T = self.seq_len\n        # t0 inside the middle half of the window\n        t0 = 0.5 * T + 0.25 * T * torch.tanh(raw[:, 0])\n        # full duration between 10% and 80% of the window\n        T14 = T * (0.10 + 0.70 * torch.sigmoid(raw[:, 1]))\n        # ingress as a fraction of the half-duration, 5%..50%\n        tau = 0.05 + 0.45 * torch.sigmoid(raw[:, 2])\n        # linear limb-darkening coefficient, 0..0.9\n        u = 0.9 * torch.sigmoid(raw[:, 3])\n        return t0, T14, tau, u\n\n\ndef transit_shape(t0, T14, tau, u, seq_len=SEQ_LEN, device=None):\n    \"\"\"s(t; theta) in [0,1]: a trapezoid with a linear limb-darkening profile.\n\n    Differentiable in every parameter. At u=0 it is the plain trapezoid, which\n    is the correct first-order shape for a transit. The limb-darkening factor is\n    normalised so the disc-average intensity is 1, which keeps d_phys on the\n    same scale as the label.\n    \"\"\"\n    t = torch.arange(seq_len, device=device, dtype=t0.dtype).unsqueeze(0)  # (1,T)\n    half = (0.5 * T14).unsqueeze(1)                                        # (B,1)\n    d = (t - t0.unsqueeze(1)).abs()                                        # (B,T)\n    ramp = (tau.unsqueeze(1) * half).clamp_min(1.0)   # >= 1 bin, so the ingress\n                                                       # always spans real bins\n                                                       # and carries a gradient\n    g = ((half - d) / ramp).clamp(0.0, 1.0)                                # (B,T)\n    # radial position on the stellar disc while in transit\n    rho = (d / half.clamp_min(1e-6)).clamp(0.0, 1.0)\n    mu_ld = torch.sqrt((1.0 - rho ** 2).clamp_min(0.0))\n    uu = u.unsqueeze(1)\n    I = (1.0 - uu * (1.0 - mu_ld)) / (1.0 - uu / 3.0)\n    return g * I\n\n\ndef project(x, s, noise_rel):\n    \"\"\"Closed-form per-wavelength depth and its variance.\n\n    x: (B,T,W) = flux/baseline - 1, so the transit is a NEGATIVE dip and\n       x ~= -s (x) d.  noise_rel: (B,W).  s: (B,T).\n    Returns d_phys (B,W), var (B,W), s2 (B,1) = effective in-transit bins.\n    \"\"\"\n    s2 = (s ** 2).sum(1, keepdim=True).clamp_min(1e-3)          # (B,1)\n    d = torch.einsum(\"bt,btw->bw\", s, -x) / s2                   # (B,W)\n    var = (noise_rel ** 2) / s2                                  # (B,W)\n    return d, var, s2\n\n\n# =============================================================================\n# ENCODER -- the frozen winner's own blocks, stopped before its head so the\n# pooled feature vector z is available to the new head.\n# =============================================================================\n\nclass Encoder(nn.Module):\n    def __init__(self, hidden_dim=HIDDEN_DIM, kernels=KERNEL_SIZES,\n                 blocks=N_BLOCKS, dropout_p=DROPOUT_P, n_wave=N_WAVE):\n        super().__init__()\n        self.input_proj = nn.Conv1d(n_wave, hidden_dim, kernel_size=1)\n        self.stages = nn.ModuleList([\n            nn.Sequential(*[M.ModernTCNBlock(hidden_dim, k, dropout_p,\n                                             pointwise_type=\"standard\")\n                            for _ in range(b)])\n            for k, b in zip(kernels, blocks)])\n        self.downsample = nn.MaxPool1d(kernel_size=2, stride=2)\n        self.attn_score = nn.Conv1d(hidden_dim, 1, kernel_size=1)\n        self.out_dim = hidden_dim\n\n    def forward(self, x):\n        # x: (B, T, W) -> channels are wavelength, sequence is time\n        h = self.input_proj(x.transpose(1, 2))\n        n = len(self.stages)\n        for i, st in enumerate(self.stages):\n            h = st(h)\n            if i < n - 1:\n                h = self.downsample(h)\n        wts = F.softmax(self.attn_score(h), dim=-1)\n        return (h * wts).sum(-1)                                # (B, hidden_dim)\n\n\n# =============================================================================\n# CONTRIBUTION 2 -- STRUCTURE-AWARE OUTPUT HEAD\n# =============================================================================\n\ndef cosine_basis(n_wave=N_WAVE, k=K_BASIS):\n    \"\"\"DCT-II basis, k = 1..K (the constant term is left out: mu_L already is\n    the level). Fixed, analytic, fitted on nothing -- so nothing can leak.\"\"\"\n    i = np.arange(n_wave)[:, None]\n    kk = np.arange(1, k + 1)[None, :]\n    B = np.cos(np.pi * kk * (i + 0.5) / n_wave)\n    B = B / np.linalg.norm(B, axis=0, keepdims=True)     # unit-norm columns\n    return torch.tensor(B, dtype=torch.float32)          # (W, K)\n\n\nclass StructureHead(nn.Module):\n    \"\"\"mu   = d_phys + mu_L + sum_k a_k B_k\n       var  = c^2 var_phys + sigma_L^2 + sum_k B_k^2 exp(2 l_k)      [diag]\n              c^2 var_phys + sigma_L^2 + diag(B L L^T B^T)           [chol]\n\n    Every correction is zero-initialised, so at step 0 mu is exactly d_phys and\n    the model degrades gracefully to pure least squares. With HEAD_TYPE=plain it\n    predicts 283 free means and 283 free sigmas instead -- the baseline's head,\n    kept here so the ablation changes one thing at a time.\n    \"\"\"\n\n    def __init__(self, in_dim, n_wave=N_WAVE, head_type=HEAD_TYPE,\n                 k=K_BASIS, cov_type=COV_TYPE, use_phys=USE_DTP,\n                 dropout_p=DROPOUT_P):\n        super().__init__()\n        self.head_type, self.cov_type, self.use_phys = head_type, cov_type, use_phys\n        self.n_wave, self.k = n_wave, k\n        self.dropout = nn.Dropout(dropout_p)\n\n        if head_type == \"basis\":\n            self.register_buffer(\"B\", cosine_basis(n_wave, k))          # (W,K)\n            n_cov = k if cov_type == \"diag\" else k * (k + 1) // 2\n            n_out = 1 + k + n_cov                    # mu_L, a, covariance\n        elif head_type == \"level\":\n            n_out = 1                                 # mu_L only\n            n_cov = 0\n        else:                                         # plain\n            n_out = 2 * n_wave\n            n_cov = 0\n        self.n_cov = n_cov\n        self.fc = nn.Linear(in_dim, n_out)\n        if head_type in (\"basis\", \"level\"):\n            nn.init.zeros_(self.fc.weight)\n            nn.init.zeros_(self.fc.bias)\n        else:\n            nn.init.zeros_(self.fc.bias)\n            self.fc.weight.data.mul_(0.01)\n\n        # log sigma_L and log c, learned globally (not per planet)\n        self.log_sigma_L = nn.Parameter(torch.tensor(math.log(3e-4)))\n        self.log_c = nn.Parameter(torch.zeros(()))     # c = 1 at init\n\n    @torch.no_grad()\n    def warm_init(self, mu_native, sigma_native):\n        \"\"\"The plain head predicts sigma as SIGMA_FLOOR + softplus(raw). With a\n        zero bias that is softplus(0) = 0.693 in native units -- 693,000 ppm,\n        absurd, and it takes many epochs to walk down from there. The baseline\n        trainer solves this with warm_init_head; the same is needed here or the\n        plain-head rungs would be handicapped and the ablation unfair.\n        The structured heads need no warm start: their sigma comes from\n        c^2 var_phys + sigma_L^2, already initialised at a physical scale.\"\"\"\n        if self.head_type != \"plain\":\n            return False\n        s = np.clip(np.asarray(sigma_native) - SIGMA_FLOOR, 1e-8, None)\n        self.fc.bias[:self.n_wave] = torch.tensor(\n            np.asarray(mu_native), dtype=torch.float32)\n        self.fc.bias[self.n_wave:] = torch.tensor(\n            np.log(np.expm1(s)), dtype=torch.float32)\n        return True\n\n    def forward(self, z, d_phys, var_phys):\n        B_ = d_phys.shape[0] if d_phys is not None else z.shape[0]\n        out = self.fc(self.dropout(z)) if z is not None else None\n\n        if self.head_type == \"plain\":\n            mu_corr, log_s = out.chunk(2, dim=-1)\n            sigma = SIGMA_FLOOR + F.softplus(log_s)\n            mu = mu_corr + (d_phys if self.use_phys else 0.0)\n            return mu, sigma, {}\n\n        sigma_L = torch.exp(self.log_sigma_L)\n        c = torch.exp(self.log_c)\n        mu = d_phys.clone() if self.use_phys else torch.zeros(\n            B_, self.n_wave, device=z.device, dtype=z.dtype)\n        v = (c ** 2) * var_phys if self.use_phys else torch.zeros_like(mu)\n        v = v + sigma_L ** 2\n\n        if self.head_type == \"level\":\n            mu = mu + out[:, :1]\n            return mu, torch.sqrt(v.clamp_min(SIGMA_FLOOR ** 2)), {\"c\": c}\n\n        mu_L = out[:, :1]\n        a = out[:, 1:1 + self.k]                                   # (B,K)\n        raw_cov = out[:, 1 + self.k:]\n        mu = mu + mu_L + a @ self.B.T                              # (B,W)\n\n        if self.cov_type == \"diag\":\n            # Sigma_a = diag(exp(2 l_k)), l_k = raw - 4 so the basis variance\n            # starts small. diag(B Sigma_a B^T) = sum_k B_k(lambda)^2 exp(2 l_k),\n            # which is (B,K) @ (K,W) -> (B,W). Diagonal in BASIS space is still\n            # wavelength-CORRELATED, which is the claim.\n            lam = torch.exp(2.0 * (raw_cov - 4.0).clamp(-12.0, 6.0))    # (B,K)\n            v = v + lam @ (self.B ** 2).T                               # (B,W)\n        else:\n            # full Cholesky in basis space: Sigma_a = L L^T\n            L = torch.zeros(B_, self.k, self.k, device=out.device, dtype=out.dtype)\n            ii, jj = torch.tril_indices(self.k, self.k, device=out.device)\n            L[:, ii, jj] = raw_cov\n            dg = torch.arange(self.k, device=out.device)\n            L[:, dg, dg] = torch.exp((L[:, dg, dg] - 4.0).clamp(-12.0, 6.0))\n            BL = torch.einsum(\"wk,bkj->bwj\", self.B, L)                 # (B,W,K)\n            v = v + (BL ** 2).sum(-1)                                   # (B,W)\n\n        return mu, torch.sqrt(v.clamp_min(SIGMA_FLOOR ** 2)), {\"c\": c, \"a\": a}\n\n\n# =============================================================================\n# THE FULL MODEL\n# =============================================================================\n\nclass PCSBN(nn.Module):\n    \"\"\"Every switch is an explicit argument, defaulting to the environment.\n\n    That matters for two things the globals could not do: the EMA shadow copy\n    must be built with the SAME switches (clone_empty), and the out-of-fold\n    assembler must rebuild any checkpoint's architecture from the config stored\n    inside it (build_from_config) rather than from whatever the environment\n    happens to say at that moment.\n    \"\"\"\n\n    def __init__(self, use_dtp=None, use_encoder=None, head_type=None,\n                 k_basis=None, cov_type=None, dropout_p=None, n_wave=N_WAVE):\n        super().__init__()\n        self.kw = dict(\n            use_dtp=USE_DTP if use_dtp is None else bool(use_dtp),\n            use_encoder=USE_ENCODER if use_encoder is None else bool(use_encoder),\n            head_type=HEAD_TYPE if head_type is None else head_type,\n            k_basis=K_BASIS if k_basis is None else int(k_basis),\n            cov_type=COV_TYPE if cov_type is None else cov_type,\n            dropout_p=DROPOUT_P if dropout_p is None else float(dropout_p),\n            n_wave=n_wave)\n        kw = self.kw\n        self.use_dtp, self.use_encoder = kw[\"use_dtp\"], kw[\"use_encoder\"]\n        self.n_wave = kw[\"n_wave\"]\n        if self.use_dtp:\n            self.geom = GeometryCNN()\n        if self.use_encoder:\n            self.encoder = Encoder(dropout_p=kw[\"dropout_p\"], n_wave=kw[\"n_wave\"])\n            in_dim = self.encoder.out_dim\n        else:\n            # pure physics: the head still needs an input, so give it a single\n            # learned vector -- the corrections become per-dataset constants\n            self.z_const = nn.Parameter(torch.zeros(1, 16))\n            in_dim = 16\n        self.head = StructureHead(in_dim, n_wave=kw[\"n_wave\"],\n                                  head_type=kw[\"head_type\"], k=kw[\"k_basis\"],\n                                  cov_type=kw[\"cov_type\"], use_phys=self.use_dtp,\n                                  dropout_p=kw[\"dropout_p\"])\n\n    def clone_empty(self):\n        return PCSBN(**self.kw)\n\n    def forward(self, x, noise_rel, metadata=None):\n        \"\"\"PRECISION SPLIT -- this is not optional.\n\n        var(lambda) = noise_rel^2 / sum_t s^2 is around 8e-10, and float16's\n        smallest normal is 6.1e-5, so under autocast it would flush to EXACTLY\n        zero -- as would SIGMA_FLOOR^2 = 1e-12 -- giving sigma = 0 and a NaN\n        loss. d_phys ~ 0.012 also only resolves ~12 ppm steps in float16, and the\n        signal is ~10 ppm; this is the same failure that cost +0.048 earlier.\n\n        So the physics path and the head run in float32. Only the ENCODER runs in\n        float16, where it is safe: the residual is already centred at 0, so\n        float16's constant relative precision works in its favour there.\n        \"\"\"\n        B = x.shape[0]\n        theta = None\n        with torch.autocast(\"cuda\", enabled=False):\n            x32 = x.float()\n            if self.use_dtp:\n                w = x32.mean(-1)                                       # (B,T)\n                t0, T14, tau, u = self.geom(w)\n                s = transit_shape(t0, T14, tau, u, x32.shape[1], x32.device)\n                d_phys, var_phys, s2 = project(x32, s, noise_rel.float())\n                theta = dict(t0=t0, T14=T14, tau=tau, u=u, s2=s2.squeeze(1))\n                resid = x32 + s.unsqueeze(-1) * d_phys.unsqueeze(1)    # (B,T,W)\n            else:\n                d_phys = torch.zeros(B, self.n_wave,\n                                     device=x32.device, dtype=x32.dtype)\n                var_phys = torch.zeros_like(d_phys)\n                resid = x32\n\n        # the one part that DOES run in float16, under the caller's autocast\n        if self.use_encoder:\n            z = self.encoder(resid * ENCODER_SCALE)\n        else:\n            z = self.z_const.expand(B, -1)\n\n        with torch.autocast(\"cuda\", enabled=False):\n            mu, sigma, extra = self.head(z.float(), d_phys, var_phys)\n        if theta is not None:\n            extra.update(theta)\n            extra[\"d_phys\"] = d_phys\n        return mu, sigma, extra\n\n\n# =============================================================================\n# LOSSES / METRICS  (identical definitions to every other script here)\n# =============================================================================\n\ndef beta_nll(mu, sigma, target, beta=BETA_NLL):\n    nll = 0.5 * (math.log(2 * math.pi) + 2 * torch.log(sigma)\n                 + ((target - mu) / sigma) ** 2)\n    return ((sigma.detach() ** (2 * beta)) * nll).mean()\n\n\ndef theta_supervision(extra, theta_ref):\n    \"\"\"Pull t0 and T14 toward the cache's label-free box search. Scale-free, so\n    it cannot dominate the likelihood.\"\"\"\n    if \"t0\" not in extra:\n        return torch.zeros((), device=theta_ref.device)\n    t0_r, w_r = theta_ref[:, 0], theta_ref[:, 1]\n    return (((extra[\"t0\"] - t0_r) / SEQ_LEN) ** 2).mean() \\\n        + (((extra[\"T14\"] - w_r) / SEQ_LEN) ** 2).mean()\n\n\ndef gll_ppm(y, mu, sg):\n    sg = np.clip(sg, 1e-8, None)\n    return -0.5 * (np.log(2 * np.pi) + np.log(sg ** 2) + ((y - mu) ** 2) / sg ** 2)\n\n\ndef ece(y, mu, sg, n_bins=ECE_N_BINS):\n    pit = _norm.cdf((y - mu) / np.clip(sg, 1e-8, None))\n    lv = np.linspace(0, 1, n_bins + 1)[1:]\n    return float(np.mean(np.abs(np.array([(pit <= l).mean() for l in lv]) - lv)))\n\n\ndef build_from_config(cfg):\n    \"\"\"Rebuild a model exactly as a checkpoint was trained. Used by the\n    out-of-fold assembler, so it never has to guess an architecture.\"\"\"\n    return PCSBN(use_dtp=cfg[\"USE_DTP\"], use_encoder=cfg[\"USE_ENCODER\"],\n                 head_type=cfg[\"HEAD_TYPE\"], k_basis=cfg.get(\"K_BASIS\", K_BASIS),\n                 cov_type=cfg.get(\"COV_TYPE\", \"diag\"),\n                 dropout_p=cfg.get(\"DROPOUT_P\", DROPOUT_P))\n\n\nclass EMA:\n    def __init__(self, model, decay=EMA_DECAY):\n        self.decay = decay\n        self.shadow = model.clone_empty().to(DEVICE)\n        self.shadow.load_state_dict(model.state_dict())\n        for p in self.shadow.parameters():\n            p.requires_grad_(False)\n\n    @torch.no_grad()\n    def update(self, model):\n        msd = model.state_dict()\n        for k, v in self.shadow.state_dict().items():\n            if v.dtype.is_floating_point:\n                v.mul_(self.decay).add_(msd[k].detach(), alpha=1 - self.decay)\n            else:\n                v.copy_(msd[k])\n\n\n@torch.no_grad()\ndef evaluate(model, loader, ref_mean_ppm, ref_std_ppm):\n    model.eval()\n    ys, mus, sgs, th = [], [], [], []\n    for b in loader:\n        x = b[\"x\"].to(DEVICE)\n        nr = b[\"noise_rel\"].to(DEVICE)\n        if DEVICE.type == \"cuda\":\n            with torch.autocast(\"cuda\", dtype=torch.float16):\n                mu, sg, ex = model(x, nr, b[\"metadata\"].to(DEVICE))\n        else:\n            mu, sg, ex = model(x, nr, b[\"metadata\"].to(DEVICE))\n        ys.append(b[\"target\"].numpy())\n        mus.append(mu.float().cpu().numpy())\n        sgs.append(sg.float().cpu().numpy())\n        if \"t0\" in ex:\n            th.append(np.stack([ex[\"t0\"].float().cpu().numpy(),\n                                ex[\"T14\"].float().cpu().numpy(),\n                                ex[\"tau\"].float().cpu().numpy(),\n                                ex[\"u\"].float().cpu().numpy()], 1))\n    y = np.concatenate(ys) * PPM_SCALE\n    mu = np.concatenate(mus) * PPM_SCALE\n    sg = np.clip(np.concatenate(sgs) * PPM_SCALE, 1.0, None)\n    l_ideal = -0.5 * (np.log(2 * np.pi) + np.log(GLL_IDEAL_SIGMA_PPM ** 2)) * y.size\n    l_ref = float(np.sum(gll_ppm(y, np.broadcast_to(ref_mean_ppm, y.shape),\n                                 np.broadcast_to(ref_std_ppm, y.shape))))\n    score = (float(np.sum(gll_ppm(y, mu, sg))) - l_ref) / (l_ideal - l_ref)\n    out = {\"GLL_score\": float(score), \"ECE\": ece(y, mu, sg),\n           \"RMSE_ppm\": float(np.sqrt(np.mean((mu - y) ** 2))),\n           \"mean_sigma_ppm\": float(sg.mean()),\n           \"cov68\": float((np.abs((y - mu) / sg) <= 1).mean()),\n           \"level_err_ppm\": float((mu - y).mean(1).mean())}\n    if th:\n        T = np.concatenate(th)\n        out[\"theta_median\"] = {\"t0\": float(np.median(T[:, 0])),\n                               \"T14\": float(np.median(T[:, 1])),\n                               \"tau\": float(np.median(T[:, 2])),\n                               \"u\": float(np.median(T[:, 3]))}\n    model.train()\n    return out\n\n\n# =============================================================================\n# MAIN\n# =============================================================================\n\ndef main():\n    torch.manual_seed(TRAIN_SEED)\n    torch.cuda.manual_seed_all(TRAIN_SEED)\n    np.random.seed(TRAIN_SEED)\n    random.seed(TRAIN_SEED)\n\n    banner(\"PC-SBN\")\n    print(f\"  DTP {USE_DTP}   encoder {USE_ENCODER}   head {HEAD_TYPE}\"\n          + (f\" K={K_BASIS} cov={COV_TYPE}\" if HEAD_TYPE == \"basis\" else \"\"))\n    print(f\"  fold {FOLD}   seed {TRAIN_SEED}   dropout {DROPOUT_P}   \"\n          f\"epochs {NUM_EPOCHS}   beta {BETA_NLL}\")\n    print(f\"  tag  {DATA_TAG}\")\n    print(f\"  repo {CHECKPOINT_REPO}\")\n    print(f\"  cache {CACHE_REPO}\")\n\n    train_df = pd.read_csv(TRAIN_CSV)\n    all_ids = train_df[\"planet_id\"].astype(str).tolist()\n    wl_cols = [c for c in train_df.columns if c != \"planet_id\"]\n\n    rng = np.random.RandomState(42)\n    shuffled = all_ids.copy()\n    rng.shuffle(shuffled)\n    if FOLD < 0:\n        n_val = max(1, int(len(shuffled) * VAL_FRACTION))\n        val_ids, train_ids = shuffled[:n_val], shuffled[n_val:]\n        print(f\"  split: historical single split\")\n    else:\n        blocks = np.array_split(np.arange(len(shuffled)), N_FOLDS)\n        vp = set(int(i) for i in blocks[FOLD])\n        val_ids = [p for j, p in enumerate(shuffled) if j in vp]\n        train_ids = [p for j, p in enumerate(shuffled) if j not in vp]\n        print(f\"  split: {N_FOLDS}-fold CV, fold {FOLD}\")\n    print(f\"Train planets: {len(train_ids)}  Val planets: {len(val_ids)}\")\n\n    # ---- preflight: the v3 cache must be complete ----\n    local = {p.stem for p in Path(LOCAL_CACHE_DIR).glob(\"*.npz\")} \\\n        if Path(LOCAL_CACHE_DIR).exists() else set()\n    try:\n        remote = {f[:-4] for f in _hf_api.list_repo_files(CACHE_REPO, repo_type=\"dataset\")\n                  if f.endswith(\".npz\")}\n    except Exception as e:\n        remote = set()\n        print(f\"  [preflight] cannot list {CACHE_REPO}: {type(e).__name__}\")\n    need = set(train_ids) | set(val_ids)\n    miss = sorted(need - (local | remote))\n    print(f\"  [preflight] cache planets available: {len(need) - len(miss)} / {len(need)}\")\n    if miss:\n        raise RuntimeError(\n            f\"v3 cache incomplete: {len(miss)} planets missing (e.g. {miss[:5]}). \"\n            f\"Run 13_add_noise_to_cache.py until it reports 1100 / 1100.\")\n\n    ref_mean_ppm, ref_std_ppm = M.compute_target_reference_stats(\n        train_df[train_df[\"planet_id\"].astype(str).isin(train_ids)], wl_cols)\n\n    star = pd.read_csv(STAR_INFO_CSV)\n    star[\"planet_id\"] = star[\"planet_id\"].astype(str)\n    tsi = star[star[\"planet_id\"].isin(train_ids)]\n    meta_mean = tsi[METADATA_COLS].values.astype(np.float32).mean(0)\n    meta_std = tsi[METADATA_COLS].values.astype(np.float32).std(0) + 1e-8\n\n    tr_ds = PCSBNDataset(train_ids, train_df, star, meta_mean, meta_std)\n    va_ds = PCSBNDataset(val_ids, train_df, star, meta_mean, meta_std)\n    te_ds = PCSBNDataset(train_ids[:min(TRAIN_EVAL_SUBSET_SIZE, len(train_ids))],\n                         train_df, star, meta_mean, meta_std)\n    tl = DataLoader(tr_ds, batch_size=BATCH_SIZE, shuffle=True,\n                    num_workers=NUM_WORKERS, pin_memory=True, drop_last=True)\n    vl = DataLoader(va_ds, batch_size=BATCH_SIZE, shuffle=False,\n                    num_workers=NUM_WORKERS, pin_memory=True)\n    el = DataLoader(te_ds, batch_size=BATCH_SIZE, shuffle=False,\n                    num_workers=NUM_WORKERS, pin_memory=True)\n\n    model = PCSBN().to(DEVICE)\n    # the plain head needs a physical sigma starting point (see warm_init).\n    # mu starts at the train mean level only when there is no physics path to\n    # start from; with DTP the physics already supplies the level.\n    _mu0 = np.zeros(N_WAVE) if USE_DTP else (ref_mean_ppm / PPM_SCALE)\n    _sg0 = np.full(N_WAVE, 300.0 / PPM_SCALE)          # 300 ppm, a typical scale\n    if model.head.warm_init(_mu0, _sg0):\n        print(f\"  [warm init] plain head: sigma -> 300 ppm, \"\n              f\"mu bias -> {'0 (DTP supplies the level)' if USE_DTP else 'train mean'}\")\n    n_par = sum(p.numel() for p in model.parameters())\n    n_geom = sum(p.numel() for p in model.geom.parameters()) if USE_DTP else 0\n    n_enc = sum(p.numel() for p in model.encoder.parameters()) if USE_ENCODER else 0\n    n_head = sum(p.numel() for p in model.head.parameters())\n    print(f\"\\n  params total {n_par:,}   geometry {n_geom:,}   \"\n          f\"encoder {n_enc:,}   head {n_head:,}\")\n\n    # ---- SANITY: with a zero-initialised head, mu must equal d_phys ----\n    if USE_DTP and HEAD_TYPE in (\"basis\", \"level\"):\n        model.eval()\n        with torch.no_grad():\n            b = next(iter(vl))\n            mu0, sg0, ex0 = model(b[\"x\"].to(DEVICE), b[\"noise_rel\"].to(DEVICE))\n            dev = (mu0 - ex0[\"d_phys\"]).abs().max().item()\n            print(f\"  [sanity] max |mu - d_phys| at init = {dev:.3e}  \"\n                  f\"({'OK' if dev < 1e-6 else 'FAIL -- head is not zero-init'})\")\n            if dev >= 1e-6:\n                raise RuntimeError(\"the correction head is not zero-initialised\")\n            dp = ex0[\"d_phys\"].float().cpu().numpy() * PPM_SCALE\n            yt = b[\"target\"].numpy() * PPM_SCALE\n            print(f\"  [sanity] d_phys level median {np.median(dp.mean(1)):8.1f} ppm   \"\n                  f\"label median {np.median(yt.mean(1)):8.1f} ppm\")\n            print(f\"  [sanity] corr(d_phys level, label level) = \"\n                  f\"{np.corrcoef(dp.mean(1), yt.mean(1))[0,1]:+.4f}   \"\n                  f\"(a physics path that works starts well above +0.9)\")\n            print(f\"  [sanity] effective in-transit bins, median \"\n                  f\"{float(ex0['s2'].median()):.1f} of {SEQ_LEN}\")\n        model.train()\n\n    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=NUM_EPOCHS)\n    scaler = torch.amp.GradScaler(\"cuda\", enabled=(DEVICE.type == \"cuda\"))\n    ema = EMA(model)\n\n    ckpt_local = f\"/kaggle/working/pcsbn_{DATA_TAG}_latest.pt\"\n    best_local = f\"/kaggle/working/pcsbn_{DATA_TAG}_best.pt\"\n    start_ep, best, best_ep = 0, -1e9, -1\n    if RESUME_IF_AVAILABLE and os.path.exists(ckpt_local):\n        try:\n            ck = torch.load(ckpt_local, map_location=DEVICE, weights_only=False)\n            model.load_state_dict(ck[\"model_state\"])\n            ema.shadow.load_state_dict(ck[\"ema_state\"])\n            opt.load_state_dict(ck[\"optimizer_state\"])\n            start_ep, best, best_ep = ck[\"epoch\"] + 1, ck[\"best\"], ck[\"best_epoch\"]\n            print(f\"  [resume] from epoch {start_ep} (best {best:.4f} @ {best_ep})\")\n        except Exception as e:\n            print(f\"  [resume] failed ({type(e).__name__}) -- starting fresh\")\n\n    banner(\"TRAINING\")\n    t_start = time.time()\n    for ep in range(start_ep, NUM_EPOCHS):\n        if ep == MSE_WARMUP_EPOCHS:\n            for g in opt.param_groups:\n                g[\"lr\"] = LR * TRANSITION_LR_DECAY\n            sched = torch.optim.lr_scheduler.CosineAnnealingLR(\n                opt, T_max=max(NUM_EPOCHS - MSE_WARMUP_EPOCHS, 1))\n        model.train()\n        tot = 0.0\n        for b in tl:\n            x = b[\"x\"].to(DEVICE, non_blocking=True)\n            nr = b[\"noise_rel\"].to(DEVICE, non_blocking=True)\n            yt = b[\"target\"].to(DEVICE, non_blocking=True)\n            tr = b[\"theta_ref\"].to(DEVICE, non_blocking=True)\n            opt.zero_grad(set_to_none=True)\n            with torch.autocast(\"cuda\", dtype=torch.float16,\n                                enabled=(DEVICE.type == \"cuda\")):\n                mu, sg, ex = model(x, nr, b[\"metadata\"].to(DEVICE))\n                loss = (F.mse_loss(mu, yt) if ep < MSE_WARMUP_EPOCHS\n                        else beta_nll(mu, sg, yt))\n                if ep < THETA_SUP_EPOCHS and USE_DTP:\n                    loss = loss + THETA_SUP_WEIGHT * theta_supervision(ex, tr)\n            scaler.scale(loss).backward()\n            scaler.unscale_(opt)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP_NORM)\n            scaler.step(opt)\n            scaler.update()\n            ema.update(model)\n            tot += float(loss.detach())\n        sched.step()\n\n        if ep % 5 == 0 or ep == NUM_EPOCHS - 1:\n            vm = evaluate(ema.shadow, vl, ref_mean_ppm, ref_std_ppm)\n            msg = (f\"  ep {ep:>3d}  loss {tot/max(len(tl),1):+.5f}  \"\n                   f\"val GLL {vm['GLL_score']:.4f}  ECE {vm['ECE']:.4f}  \"\n                   f\"RMSE {vm['RMSE_ppm']:.1f}  cov68 {vm['cov68']:.3f}\")\n            if \"theta_median\" in vm:\n                t = vm[\"theta_median\"]\n                msg += (f\"  theta t0={t['t0']:.0f} T14={t['T14']:.0f} \"\n                        f\"tau={t['tau']:.2f} u={t['u']:.2f}\")\n            print(msg, flush=True)\n            if vm[\"GLL_score\"] > best:\n                best, best_ep = vm[\"GLL_score\"], ep\n                torch.save({\"model_state\": model.state_dict(),\n                            \"ema_state\": ema.shadow.state_dict(),\n                            \"epoch\": ep, \"val_metrics\": vm,\n                            \"config\": _cfg_dict()}, best_local)\n                print(f\"  [best] {best:.4f} @ {ep} -- saved\")\n            torch.save({\"model_state\": model.state_dict(),\n                        \"ema_state\": ema.shadow.state_dict(),\n                        \"optimizer_state\": opt.state_dict(), \"epoch\": ep,\n                        \"best\": best, \"best_epoch\": best_ep,\n                        \"config\": _cfg_dict()}, ckpt_local)\n            if ep % HF_UPLOAD_EVERY_N_EPOCHS == 0:\n                _upload(ckpt_local, f\"pcsbn_{DATA_TAG}_latest.pt\")\n\n    banner(\"FINAL\")\n    vm = evaluate(ema.shadow, vl, ref_mean_ppm, ref_std_ppm)\n    tm = evaluate(ema.shadow, el, ref_mean_ppm, ref_std_ppm)\n    print(\"  val   (EMA, last epoch):\\n\" + json.dumps(vm, indent=2))\n    print(f\"  train subset GLL {tm['GLL_score']:.4f}  \"\n          f\"(generalisation gap {tm['GLL_score'] - vm['GLL_score']:+.4f})\")\n    fin = f\"/kaggle/working/pcsbn_{DATA_TAG}_final.pt\"\n    torch.save({\"model_state\": ema.shadow.state_dict(), \"val_metrics\": vm,\n                \"train_subset_metrics\": tm, \"config\": _cfg_dict()}, fin)\n    _upload(fin, f\"pcsbn_{DATA_TAG}_final.pt\")\n    if os.path.exists(best_local):\n        _upload(best_local, f\"pcsbn_{DATA_TAG}_best.pt\")\n    with open(f\"/kaggle/working/pcsbn_metrics_{DATA_TAG}.json\", \"w\") as f:\n        json.dump({\"val_last\": vm, \"train_subset_last\": tm,\n                   \"best_val_gll\": best, \"best_epoch\": best_ep,\n                   \"n_params\": n_par, \"config\": _cfg_dict()}, f, indent=2)\n    print(f\"\\n  BEST val GLL {best:.4f} @ epoch {best_ep}\")\n    print(f\"  wall time {(time.time()-t_start)/60:.1f} min\")\n    print(f\"\\n  Next: 21_assemble_oof_pcsbn.py, then 19_paired_compare.py\")\n    print(f\"  against the locked baseline (pooled OOF 0.4290 / 0.4890 excl 1).\")\n\n\ndef _cfg_dict():\n    return {\"USE_DTP\": USE_DTP, \"USE_ENCODER\": USE_ENCODER,\n            \"HEAD_TYPE\": HEAD_TYPE, \"K_BASIS\": K_BASIS, \"COV_TYPE\": COV_TYPE,\n            \"FOLD\": FOLD, \"N_FOLDS\": N_FOLDS, \"TRAIN_SEED\": TRAIN_SEED,\n            \"DROPOUT_P\": DROPOUT_P, \"NUM_EPOCHS\": NUM_EPOCHS,\n            \"BETA_NLL\": BETA_NLL, \"hidden_dim\": HIDDEN_DIM,\n            \"kernels\": KERNEL_SIZES, \"tag\": DATA_TAG, \"cache\": CACHE_REPO}\n\n\ndef _upload(path, name):\n    try:\n        create_repo(repo_id=CHECKPOINT_REPO, repo_type=\"model\",\n                    exist_ok=True, private=False)\n        _hf_api.upload_file(path_or_fileobj=path, path_in_repo=name,\n                            repo_id=CHECKPOINT_REPO, repo_type=\"model\")\n    except Exception as e:\n        print(f\"  [warn] HF upload failed for {name}: {type(e).__name__}\")\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!/usr/bin/env python\n# =============================================================================\n# ARIEL 2025 -- SCRIPT 21 : OUT-OF-FOLD ASSEMBLER FOR PC-SBN\n# =============================================================================\n# The PC-SBN counterpart of script 15. Two differences, both deliberate:\n#\n#   * No architecture guessing. Every checkpoint stores its own `config`, and\n#     build_from_config rebuilds the model from it, then loads strictly. Script\n#     15 had to infer ModernTCN's shape from the state_dict because the old\n#     trainer saved no config; that guessing is gone here.\n#\n#   * The output file is byte-compatible with script 15's: the same 1100 planets\n#     in the same order (folds 0..4, each in its own val_ids order), with\n#     planet_ids and fold_id. So scripts 16 and 19 pick it up automatically and\n#     PAIR it against the locked baseline with no extra work.\n#\n# The comparison target, for reference:\n#     baseline, 5-fold, final checkpoints\n#       pooled GLL all 1100      0.4290\n#       pooled GLL excl 1        0.4890\n#       median per-planet        0.5095\n#       oracle sigma ceiling     0.6069\n#\n# -----------------------------------------------------------------------------\n# NOTEBOOK SETUP  (GPU T4 x2, Internet ON, HF_TOKEN)\n#   Add Input : Competitions -> ariel-data-challenge-2025\n#   Add Input : Datasets -> moderntcn-src\n#\n#   Cell 1:  %%writefile /kaggle/working/20_train_pcsbn.py\n#            <paste the whole of 20_train_pcsbn.py>\n#   Cell 2:  <paste this script>\n# =============================================================================\n\nimport os\nimport sys\nimport glob\nimport json\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import DataLoader\nfrom huggingface_hub import HfApi, hf_hub_download\n\n# =============================================================================\n# CONFIG\n# =============================================================================\n\n# PC-SBN tags WITHOUT the _f<fold> suffix, exactly as script 22 trained them.\n# Leave empty to auto-discover every pcsbn *-5fold repo on the account.\nBASES = []\n\nCKPT_KIND = \"final\"            # \"final\" (fixed epoch, the locked protocol) or \"best\"\nN_FOLDS = 5\nVERIFY_TOL = 0.002\nOUT_DIR = \"/kaggle/working\"\nTRAINER = \"20_train_pcsbn.py\"\nHF_USER = \"maariaaa12\"\n\n\ndef banner(t):\n    print(f\"\\n{'='*78}\\n{t}\\n{'='*78}\", flush=True)\n\n\n# =============================================================================\n# 1. IMPORT THE PC-SBN TRAINER (definitions only; main() is guarded)\n# =============================================================================\n\nbanner(\"1. LOCATING THE PC-SBN TRAINER\")\nhits = sorted(glob.glob(f\"/kaggle/working/{TRAINER}\")\n              + glob.glob(f\"/kaggle/input/**/{TRAINER}\", recursive=True))\nif not hits:\n    raise FileNotFoundError(\n        f\"{TRAINER} not found. In the cell above, put this as the FIRST line:\\n\"\n        f\"    %%writefile /kaggle/working/{TRAINER}\\n\"\n        f\"then paste the whole trainer under it and run that cell once.\")\nd = os.path.dirname(hits[0])\nif d not in sys.path:\n    sys.path.insert(0, d)\nimport importlib                                                    # noqa: E402\nP = importlib.import_module(TRAINER[:-3])\nprint(f\"  {hits[0]}\")\nprint(f\"  imported OK   build_from_config: {hasattr(P, 'build_from_config')}\")\n\n_api = HfApi()\n\n# =============================================================================\n# 2. DISCOVER WHAT HAS BEEN TRAINED\n# =============================================================================\n\nbanner(\"2. DISCOVERING PC-SBN 5-FOLD REPOS\")\nif BASES:\n    bases = list(BASES)\n    print(f\"  using the {len(bases)} base tag(s) in BASES\")\nelse:\n    try:\n        repos = [m.id for m in _api.list_models(author=HF_USER)]\n    except Exception as e:\n        raise RuntimeError(f\"cannot list {HF_USER}'s models ({type(e).__name__}); \"\n                           f\"set BASES explicitly.\")\n    bases = []\n    for r in repos:\n        name = r.split(\"/\", 1)[1]\n        if name.startswith(\"ariel-2025-pcsbn-\") and name.endswith(\"-5fold\"):\n            bases.append(name[len(\"ariel-2025-\"):-len(\"-5fold\")].replace(\"-\", \"_\"))\n    bases = sorted(set(bases))\n    print(f\"  auto-discovered {len(bases)} PC-SBN 5-fold configuration(s)\")\n    for b in bases:\n        print(f\"    {b}\")\nif not bases:\n    raise RuntimeError(\"No PC-SBN 5-fold checkpoints found. Run script 22 first.\")\n\n# =============================================================================\n# 3. REBUILD THE FOLD SPLITS  (identical to scripts 04 / 15 / 20)\n# =============================================================================\n\nbanner(\"3. REBUILDING THE 5-FOLD SPLITS\")\ntrain_df = pd.read_csv(P.TRAIN_CSV)\nall_ids = train_df[\"planet_id\"].astype(str).tolist()\nwl_cols = [c for c in train_df.columns if c != \"planet_id\"]\n\nrng = np.random.RandomState(42)\nshuffled = all_ids.copy()\nrng.shuffle(shuffled)\nblocks = np.array_split(np.arange(len(shuffled)), N_FOLDS)\n\nstar = pd.read_csv(P.STAR_INFO_CSV)\nstar[\"planet_id\"] = star[\"planet_id\"].astype(str)\n\nFOLDS = []\nfor f in range(N_FOLDS):\n    vp = set(int(i) for i in blocks[f])\n    va = [p for j, p in enumerate(shuffled) if j in vp]\n    tr = [p for j, p in enumerate(shuffled) if j not in vp]\n    tsi = star[star[\"planet_id\"].isin(tr)]\n    rm, rs = P.M.compute_target_reference_stats(\n        train_df[train_df[\"planet_id\"].astype(str).isin(tr)], wl_cols)\n    FOLDS.append(dict(val=va, train=tr,\n                      meta_mean=tsi[P.METADATA_COLS].values.astype(np.float32).mean(0),\n                      meta_std=tsi[P.METADATA_COLS].values.astype(np.float32).std(0) + 1e-8,\n                      ref_mean=rm, ref_std=rs))\n    print(f\"  fold {f}: train {len(tr)}  val {len(va)}\")\n\ncover = [p for f in FOLDS for p in f[\"val\"]]\nassert len(cover) == len(all_ids) == len(set(cover)), \"folds do not partition the planets\"\nprint(f\"  coverage: {len(cover)} planets, each held out exactly once  [OK]\")\n\n\n# =============================================================================\n# 4. HELPERS\n# =============================================================================\n\ndef gll_score(y, mu, sg, ref_mean, ref_std):\n    sg = np.clip(sg, 1.0, None)\n    tot = float(np.sum(P.gll_ppm(y, mu, sg)))\n    ideal = -0.5 * (np.log(2 * np.pi) + np.log(P.GLL_IDEAL_SIGMA_PPM ** 2)) * y.size\n    ref = float(np.sum(P.gll_ppm(y, np.broadcast_to(ref_mean, y.shape),\n                                 np.broadcast_to(ref_std, y.shape))))\n    return (tot - ref) / (ideal - ref)\n\n\n@torch.no_grad()\ndef infer(model, loader):\n    ys, mus, sgs, th = [], [], [], []\n    for b in loader:\n        x = b[\"x\"].to(P.DEVICE)\n        nr = b[\"noise_rel\"].to(P.DEVICE)\n        if P.DEVICE.type == \"cuda\":\n            with torch.autocast(\"cuda\", dtype=torch.float16):\n                mu, sg, ex = model(x, nr, b[\"metadata\"].to(P.DEVICE))\n        else:\n            mu, sg, ex = model(x, nr, b[\"metadata\"].to(P.DEVICE))\n        ys.append(b[\"target\"].numpy())\n        mus.append(mu.float().cpu().numpy())\n        sgs.append(sg.float().cpu().numpy())\n        if \"t0\" in ex:\n            th.append(np.stack([ex[k].float().cpu().numpy()\n                                for k in (\"t0\", \"T14\", \"tau\", \"u\")], 1))\n    return (np.concatenate(ys) * P.PPM_SCALE,\n            np.concatenate(mus) * P.PPM_SCALE,\n            np.concatenate(sgs) * P.PPM_SCALE,\n            np.concatenate(th) if th else None)\n\n\n# =============================================================================\n# 5. ASSEMBLE\n# =============================================================================\n\nsummary = []\nfor base in bases:\n    banner(f\"ASSEMBLING  {base}   [{CKPT_KIND}]\")\n    repo = f\"{HF_USER}/ariel-2025-{base.replace('_', '-')}-5fold\"\n    print(f\"  repo  {repo}\")\n    try:\n        files = set(_api.list_repo_files(repo, repo_type=\"model\"))\n    except Exception as e:\n        print(f\"  [FAILED] cannot reach it ({type(e).__name__})\")\n        continue\n    want = [f\"pcsbn_{base}_f{f}_{CKPT_KIND}.pt\" for f in range(N_FOLDS)]\n    miss = [w for w in want if w not in files]\n    if miss:\n        print(f\"  [SKIP] {len(miss)} fold checkpoint(s) missing: {miss}\")\n        continue\n\n    parts, ok, cfg0 = [], True, None\n    for f in range(N_FOLDS):\n        FD = FOLDS[f]\n        print(f\"\\n  --- fold {f} ---\")\n        path = hf_hub_download(repo_id=repo, repo_type=\"model\",\n                               filename=want[f], local_dir=\"/kaggle/working\")\n        ck = torch.load(path, map_location=P.DEVICE, weights_only=False)\n        cfg = ck[\"config\"]\n        cfg0 = cfg0 or cfg\n        for k in (\"USE_DTP\", \"USE_ENCODER\", \"HEAD_TYPE\", \"K_BASIS\", \"COV_TYPE\"):\n            if cfg.get(k) != cfg0.get(k):\n                raise ValueError(f\"fold {f} was trained with {k}={cfg.get(k)} but \"\n                                 f\"fold 0 with {cfg0.get(k)} -- not the same model\")\n        model = P.build_from_config(cfg).to(P.DEVICE)\n        W = ck[\"ema_state\"] if \"ema_state\" in ck else ck[\"model_state\"]\n        model.load_state_dict(W, strict=True)\n        model.eval()\n\n        ds = P.PCSBNDataset(FD[\"val\"], train_df, star, FD[\"meta_mean\"], FD[\"meta_std\"])\n        dl = DataLoader(ds, batch_size=P.BATCH_SIZE, shuffle=False,\n                        num_workers=2, pin_memory=True)\n        y, mu, sg, th = infer(model, dl)\n        g = gll_score(y, mu, sg, FD[\"ref_mean\"], FD[\"ref_std\"])\n        rec = (ck.get(\"val_metrics\") or {}).get(\"GLL_score\")\n        print(f\"    planets {len(y)}   GLL {g:.4f}   \"\n              f\"RMSE {np.sqrt(np.mean((mu-y)**2)):.1f} ppm\")\n        if rec is not None:\n            dd = g - float(rec)\n            print(f\"    recorded {float(rec):.4f}   difference {dd:+.5f}   \"\n                  f\"{'MATCH' if abs(dd) < VERIFY_TOL else 'MISMATCH'}\")\n            if abs(dd) >= VERIFY_TOL:\n                ok = False\n                print(f\"    [!!] fold {f} does not reproduce its recorded GLL\")\n        if th is not None:\n            print(f\"    theta median: t0 {np.median(th[:,0]):.0f}  \"\n                  f\"T14 {np.median(th[:,1]):.0f}  tau {np.median(th[:,2]):.2f}  \"\n                  f\"u {np.median(th[:,3]):.2f}\")\n        parts.append(dict(y=y, mu=mu, sg=sg, th=th, gll=g,\n                          ids=np.array(FD[\"val\"]), fold=f))\n        del model\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\n    if not ok:\n        print(f\"\\n  [ABORT] a fold failed verification. Nothing written for {base}.\")\n        continue\n\n    y = np.concatenate([p[\"y\"] for p in parts])\n    mu = np.concatenate([p[\"mu\"] for p in parts])\n    sg = np.concatenate([p[\"sg\"] for p in parts])\n    ids = np.concatenate([p[\"ids\"] for p in parts])\n    fid = np.concatenate([np.full(len(p[\"ids\"]), p[\"fold\"], np.int16) for p in parts])\n    fin = np.isfinite(mu).all(1) & np.isfinite(sg).all(1)\n    if (~fin).any():\n        print(f\"\\n  [warn] dropping {int((~fin).sum())} non-finite planets\")\n        y, mu, sg, ids, fid = y[fin], mu[fin], sg[fin], ids[fin], fid[fin]\n\n    ref_all_m, ref_all_s = P.M.compute_target_reference_stats(train_df, wl_cols)\n    g_oof = gll_score(y, mu, sg, ref_all_m, ref_all_s)\n    suffix = \"\" if CKPT_KIND == \"best\" else f\"_{CKPT_KIND}\"\n    out = os.path.join(OUT_DIR, f\"preds_{base}_oof{suffix}.npz\")\n    save = dict(y_ppm=y, mu_ppm=mu, sigma_ppm=sg, planet_ids=ids, fold_id=fid)\n    if all(p[\"th\"] is not None for p in parts):\n        save[\"theta\"] = np.concatenate([p[\"th\"] for p in parts])[fin]\n    np.savez_compressed(out, **save)\n\n    per = [p[\"gll\"] for p in parts]\n    print(f\"\\n  OUT-OF-FOLD RESULT  ({len(y)} planets)\")\n    print(f\"    per-fold GLL      {' '.join(f'{v:.4f}' for v in per)}\")\n    print(f\"    fold mean         {np.mean(per):.4f}   sd {np.std(per, ddof=1):.4f}\")\n    print(f\"    pooled OOF GLL    {g_oof:.4f}   ECE {P.ece(y, mu, sg):.4f}   \"\n          f\"RMSE {np.sqrt(np.mean((mu-y)**2)):.1f} ppm\")\n    print(f\"    vs baseline       {g_oof - 0.4290:+.4f} on the all-1100 statistic\")\n    print(f\"    saved {out}\")\n    summary.append((base + suffix, g_oof, np.mean(per), np.std(per, ddof=1), len(y)))\n\n# =============================================================================\n# 6. SUMMARY\n# =============================================================================\n\nbanner(\"SUMMARY\")\nif not summary:\n    print(\"  Nothing assembled. Fix the errors above.\")\nelse:\n    print(f\"  {'config':<46s} {'OOF GLL':>9s} {'fold mean':>10s} {'fold sd':>8s} {'n':>6s}\")\n    for b, g, fm, fs, n in summary:\n        print(f\"  {b:<46s} {g:>9.4f} {fm:>10.4f} {fs:>8.4f} {n:>6d}\")\n    print(f\"\\n  locked baseline: 0.4290 all-1100 / 0.4890 excl 1 planet\")\n    print(f\"\\n  DO NOT read the all-1100 column as the verdict: one planet moves it\")\n    print(f\"  by 0.06 and the paired MDE on it is 0.0913. Run script 19 next -- it\")\n    print(f\"  pairs every set against the baseline on four statistics at once and\")\n    print(f\"  is the only thing that can resolve a 0.01 gain.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!/usr/bin/env python\n# =============================================================================\n# ARIEL 2025 -- SCRIPT 19 : ROBUST PAIRED COMPARISON OF TWO PREDICTION SETS\n# =============================================================================\n# WHY THIS EXISTS\n#\n# The pooled GLL is the competition metric and must stay the headline. But it is\n# a SUM with a 1 ppm sigma floor, so a single planet can contribute without\n# bound. Measured on the 5-fold out-of-fold set:\n#\n#     planet 1843015807   |z| = 31.6   s_p = -70.2   (1 planet of 1100)\n#     paired MDE on pooled GLL          0.0913\n#     paired MDE the protocol should give   ~0.003-0.005\n#\n# One planet costs a factor of ~20 in statistical power. Every architectural\n# gain this thesis is chasing is smaller than 0.02, so on the pooled statistic\n# alone NOTHING would ever come out significant -- not because the models are\n# equal, but because one planet dominates the variance of the difference.\n#\n# THE FIX IS NOT TO DELETE THE PLANET. It is to report several statistics of\n# the SAME paired comparison, each declared in advance:\n#\n#   1. pooled GLL, all planets        the competition-comparable headline\n#   2. pooled GLL, rule-excluded      the same metric with a PRE-DECLARED\n#                                     exclusion rule (|z| above a threshold\n#                                     under the REFERENCE model only, so the\n#                                     rule never depends on the model being\n#                                     tested)\n#   3. median per-planet score        naturally immune to any single planet\n#   4. 5% trimmed mean per-planet     between the two\n#\n# All four are computed on ONE bootstrap resample of planets, so their CIs are\n# directly comparable. Report 1 as the result and 2-4 as the power-bearing\n# companions. If they disagree, that disagreement is itself a finding.\n#\n# The exclusion rule uses only the REFERENCE set. A new model is never allowed\n# to choose which planets it is scored on.\n#\n# CPU, seconds. Reads preds_*_oof*.npz from /kaggle/working.\n# =============================================================================\n\nimport os\nimport glob\nimport json\nfrom pathlib import Path\n\nimport numpy as np\n\nPRED_GLOB = \"/kaggle/working/preds_*_oof*.npz\"\nOUT_DIR = \"/kaggle/working/diagnostics\"\n\n# The reference every other set is compared against. \"auto\" picks the\n# _oof_final set if present (fixed-epoch selection is the locked protocol),\n# else the first file found.\nREFERENCE = \"auto\"\n\nZ_FLAG = 10.0            # pre-declared: exclude planets whose MEDIAN |z| under\n                          # the REFERENCE exceeds this. 10 sigma is far outside\n                          # anything a working uncertainty head produces.\nTRIM_FRAC = 0.05\nN_BOOT = 10000\nSEED = 42\nSIGMA_FLOOR_PPM = 1.0\nGLL_IDEAL_SIGMA_PPM = 10.0\n\n\ndef banner(t):\n    print(f\"\\n{'='*78}\\n{t}\\n{'='*78}\", flush=True)\n\n\n# ------------------------------------------------------------------ metric ---\n\ndef gll_elem(y, mu, s):\n    s = np.clip(s, 1e-8, None)\n    return -0.5 * (np.log(2 * np.pi) + np.log(s ** 2) + (y - mu) ** 2 / s ** 2)\n\n\ndef parts(y, mu, s, ref_mean, ref_std):\n    g = gll_elem(y, mu, s).sum(1)\n    ref = gll_elem(y, np.broadcast_to(ref_mean, y.shape),\n                   np.broadcast_to(ref_std, y.shape)).sum(1)\n    ideal = np.full(len(y), -0.5 * (np.log(2 * np.pi)\n                                    + np.log(GLL_IDEAL_SIGMA_PPM ** 2)) * y.shape[1])\n    return g, ref, ideal\n\n\n# ---------------------------------------------------------------- load -------\n\nbanner(\"LOADING\")\nfiles = sorted(glob.glob(PRED_GLOB))\nif not files:\n    raise FileNotFoundError(f\"No file matching {PRED_GLOB}. Run script 15 first.\")\n\nP = {}\nfor fp in files:\n    z = np.load(fp, allow_pickle=True)\n    lab = Path(fp).stem.replace(\"preds_\", \"\")\n    P[lab] = dict(y=z[\"y_ppm\"].astype(np.float64),\n                  mu=z[\"mu_ppm\"].astype(np.float64),\n                  sg=np.clip(z[\"sigma_ppm\"].astype(np.float64), SIGMA_FLOOR_PPM, None),\n                  pid=z[\"planet_ids\"].astype(str) if \"planet_ids\" in z.files else None)\n    print(f\"  {lab}   {len(P[lab]['y'])} planets\")\n\nref_lab = REFERENCE\nif ref_lab == \"auto\":\n    fin = [k for k in P if k.endswith(\"_final\")]\n    ref_lab = fin[0] if fin else sorted(P)[0]\nif ref_lab not in P:\n    raise KeyError(f\"REFERENCE {ref_lab!r} not loaded. Available: {sorted(P)}\")\nprint(f\"\\n  reference: {ref_lab}\")\n\nR = P[ref_lab]\nn = len(R[\"y\"])\n# every set must be the same planets in the same order\nfor lab, D in P.items():\n    if len(D[\"y\"]) != n or not np.allclose(D[\"y\"], R[\"y\"]):\n        raise ValueError(f\"{lab}: different planets or order than {ref_lab} -- \"\n                         f\"cannot pair. Re-assemble both with script 15.\")\n\nref_mean, ref_std = R[\"y\"].mean(0), R[\"y\"].std(0) + 1e-8\n\n# ------------------------------------------------- pre-declared exclusion ----\n\nbanner(f\"PRE-DECLARED EXCLUSION RULE   median |z| > {Z_FLAG} under {ref_lab}\")\nz_ref = np.median(np.abs((R[\"y\"] - R[\"mu\"]) / R[\"sg\"]), axis=1)\nflag = z_ref > Z_FLAG\nkeep = ~flag\nprint(f\"  {int(flag.sum())} of {n} planets excluded ({flag.sum()/n*100:.2f}%)\")\nif flag.any():\n    order = np.argsort(-z_ref[flag])\n    idx = np.where(flag)[0][order]\n    print(f\"\\n  {'planet_id':>12s} {'median |z|':>11s} {'RMSE ppm':>10s} \"\n          f\"{'med sigma':>10s}\")\n    print(f\"  {'-'*48}\")\n    for i in idx:\n        pid = R[\"pid\"][i] if R[\"pid\"] is not None else str(i)\n        print(f\"  {pid:>12s} {z_ref[i]:>11.2f} \"\n              f\"{np.sqrt(((R['mu'][i]-R['y'][i])**2).mean()):>10.1f} \"\n              f\"{np.median(R['sg'][i]):>10.1f}\")\n    print(f\"\\n  The rule uses ONLY the reference set, so no model being tested can\")\n    print(f\"  influence which planets it is scored on. State the rule and this\")\n    print(f\"  list in the thesis; never choose the threshold after seeing a result.\")\nelse:\n    print(\"  nothing excluded -- the pooled and excluded statistics coincide\")\n\n# ---------------------------------------------------------- statistics -------\n\nrng = np.random.default_rng(SEED)\nBOOT = rng.integers(0, n, size=(N_BOOT, n))\n\n\ndef four_stats(g, ref, ideal, sp, rows):\n    \"\"\"rows: (B, n) index matrix, or None for the point estimate.\"\"\"\n    if rows is None:\n        pooled = (g.sum() - ref.sum()) / (ideal.sum() - ref.sum())\n        gk, rk, ik = g[keep], ref[keep], ideal[keep]\n        pooled_ex = (gk.sum() - rk.sum()) / (ik.sum() - rk.sum())\n        med = np.median(sp)\n        k = int(round(len(sp) * TRIM_FRAC))\n        s = np.sort(sp)\n        trim = s[k:len(s) - k].mean() if len(s) - 2 * k > 0 else np.nan\n        return np.array([pooled, pooled_ex, med, trim])\n\n    gg, rr, ii = g[rows], ref[rows], ideal[rows]\n    pooled = (gg.sum(1) - rr.sum(1)) / (ii.sum(1) - rr.sum(1))\n    # excluded: mask the flagged planets INSIDE each resample\n    m = keep[rows]\n    pooled_ex = ((gg * m).sum(1) - (rr * m).sum(1)) / ((ii * m).sum(1) - (rr * m).sum(1))\n    S = sp[rows]\n    med = np.median(S, axis=1)\n    k = int(round(S.shape[1] * TRIM_FRAC))\n    Ss = np.sort(S, axis=1)\n    trim = Ss[:, k:S.shape[1] - k].mean(1)\n    return np.stack([pooled, pooled_ex, med, trim], axis=1)\n\n\nNAMES = [\"pooled GLL (all)\", f\"pooled GLL (excl {int(flag.sum())})\",\n         \"median per-planet\", f\"{TRIM_FRAC:.0%}-trimmed mean\"]\n\nCACHE = {}\nfor lab, D in P.items():\n    g, ref, ideal = parts(D[\"y\"], D[\"mu\"], D[\"sg\"], ref_mean, ref_std)\n    sp = (g - ref) / (ideal - ref)\n    CACHE[lab] = dict(point=four_stats(g, ref, ideal, sp, None),\n                      boot=four_stats(g, ref, ideal, sp, BOOT))\n\nbanner(\"POINT ESTIMATES\")\nprint(f\"  {'set':<44s}\" + \"\".join(f\"{nm:>22s}\" for nm in NAMES))\nprint(f\"  {'-'*44}\" + \"-\" * (22 * 4))\nfor lab in [ref_lab] + [k for k in sorted(P) if k != ref_lab]:\n    star = \" *\" if lab == ref_lab else \"  \"\n    print(f\"  {lab[:42]:<42s}{star}\"\n          + \"\".join(f\"{v:>22.4f}\" for v in CACHE[lab][\"point\"]))\nprint(f\"\\n  * = reference.  The FIRST column is the competition metric and the\")\nprint(f\"  number the thesis quotes. The other three exist to give the paired\")\nprint(f\"  comparison below enough power to resolve a real difference.\")\n\n# ------------------------------------------------------- paired results ------\n\nbanner(\"PAIRED DIFFERENCES vs THE REFERENCE\")\nprint(f\"  bootstrap: {N_BOOT} planet resamples, shared across all statistics\\n\")\n\nresults = {}\nothers = [k for k in sorted(P) if k != ref_lab]\nif not others:\n    print(\"  Only the reference is loaded -- nothing to compare yet.\")\n    print(\"  Assemble a second prediction set with script 15 and re-run.\")\nfor lab in others:\n    dboot = CACHE[lab][\"boot\"] - CACHE[ref_lab][\"boot\"]\n    dpoint = CACHE[lab][\"point\"] - CACHE[ref_lab][\"point\"]\n    print(f\"  {lab}   minus   {ref_lab}\")\n    print(f\"  {'statistic':<26s} {'observed':>10s} {'95% CI':>22s} \"\n          f\"{'SE':>8s} {'MDE':>8s} {'sig':>5s}\")\n    print(f\"  {'-'*84}\")\n    rows = []\n    for j, nm in enumerate(NAMES):\n        d = dboot[:, j]\n        lo, hi = np.percentile(d, [2.5, 97.5])\n        se = float(d.std(ddof=1))\n        sig = \"YES\" if (lo > 0 or hi < 0) else \"no\"\n        print(f\"  {nm:<26s} {dpoint[j]:>+10.4f} \"\n              f\"[{lo:>+9.4f},{hi:>+9.4f}] {se:>8.4f} {1.96*se:>8.4f} {sig:>5s}\")\n        rows.append(dict(statistic=nm, observed=float(dpoint[j]),\n                         ci95=[float(lo), float(hi)], se=se, mde=float(1.96 * se),\n                         significant=(sig == \"YES\")))\n    results[lab] = rows\n    print()\n\n# ------------------------------------------------------------- power note ----\n\nbanner(\"WHICH STATISTIC CAN RESOLVE WHAT\")\nprint(f\"  {'statistic':<26s} {'MDE':>9s}   what it can detect\")\nprint(f\"  {'-'*76}\")\nif others:\n    ref_rows = results[others[0]]\n    for r in ref_rows:\n        m = r[\"mde\"]\n        if m <= 0.006:\n            note = \"a 0.01 architectural gain, comfortably\"\n        elif m <= 0.012:\n            note = \"a 0.02 gain; 0.01 is marginal\"\n        elif m <= 0.03:\n            note = \"only gains above ~0.03\"\n        else:\n            note = \"nothing this thesis is chasing\"\n        print(f\"  {r['statistic']:<26s} {m:>9.4f}   {note}\")\n    print(f\"\\n  MDEs come from the first comparison listed; they are a property of\")\n    print(f\"  the statistic and the planet set, not of the models.\")\n\nprint(f\"\\n  HOW TO REPORT AN ABLATION\")\nprint(f\"    headline    pooled GLL (all planets), with its CI\")\nprint(f\"    claim       only if the power-bearing statistics agree in sign AND\")\nprint(f\"                at least one of them is significant\")\nprint(f\"    always      state the exclusion rule, the excluded planet(s), and\")\nprint(f\"                that the pooled CI is dominated by them\")\nprint(f\"    never       pick the statistic after seeing which one is significant\")\n\nPath(OUT_DIR).mkdir(parents=True, exist_ok=True)\nout = os.path.join(OUT_DIR, \"paired_compare.json\")\nwith open(out, \"w\") as fh:\n    json.dump({\"reference\": ref_lab, \"n_planets\": int(n),\n               \"z_flag\": Z_FLAG, \"n_excluded\": int(flag.sum()),\n               \"excluded_ids\": ([str(R[\"pid\"][i]) for i in np.where(flag)[0]]\n                                if R[\"pid\"] is not None else []),\n               \"statistics\": NAMES,\n               \"point_estimates\": {k: [float(v) for v in CACHE[k][\"point\"]]\n                                   for k in CACHE},\n               \"paired\": results}, fh, indent=2)\nprint(f\"\\n  saved {out}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}