{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"sourceType":"competition"},{"sourceId":9629432,"sourceType":"datasetVersion","datasetId":5846888},{"sourceId":13108789,"sourceType":"datasetVersion","datasetId":8303788},{"sourceId":564250,"sourceType":"modelInstanceVersion","modelInstanceId":426187,"modelId":443655},{"sourceId":588855,"sourceType":"modelInstanceVersion","modelInstanceId":440307,"modelId":456857},{"sourceId":589522,"sourceType":"modelInstanceVersion","modelInstanceId":440854,"modelId":456857}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"this is my attempt on training resnet model\n\nthe original ideal is from https://www.kaggle.com/code/qaedtgyh/resnet-for-airs-0-360-lb","metadata":{}},{"cell_type":"markdown","source":"you can view my write up [here](https://www.kaggle.com/competitions/ariel-data-challenge-2025/writeups/58th-place-solution)","metadata":{}},{"cell_type":"code","source":"# =========================================================\n# Ariel Data Challenge 2025 — 0.360 → 0.361狙いの軽パッチ\n# 変更点：\n#  - (#5) predictions1 を波長方向に極軽スムージングして 0.7:0.3 でブレンド\n#  - (#4) sigma_fgs / sigma_air のクリップ幅と最終係数を僅かに拡張\n# 他はオリジナルに準拠\n# =========================================================\n\n# install pqdm for parallel processing (Kaggleオフライン用ホイール)\n!pip install --no-index --find-links=/kaggle/input/ariel-2024-pqdm pqdm\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:47.203713Z","iopub.execute_input":"2025-09-24T17:36:47.203869Z","iopub.status.idle":"2025-09-24T17:36:50.202281Z","shell.execute_reply.started":"2025-09-24T17:36:47.203856Z","shell.execute_reply":"2025-09-24T17:36:50.201374Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nimport os\nimport time\nimport itertools\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom tqdm import tqdm\nfrom pqdm.threads import pqdm\nfrom astropy.stats import sigma_clip\nfrom scipy.optimize import minimize\nfrom torch.utils.data import DataLoader, TensorDataset, random_split\nfrom sklearn.preprocessing import StandardScaler\nfrom scipy.signal import savgol_filter\nfrom sklearn.metrics import mean_squared_error\nimport matplotlib.pyplot as plt\n\n# ======== パスとモード ========\nROOT_PATH = \"/kaggle/input/ariel-data-challenge-2025\"\nMODE = \"train\"\n__t0 = time.perf_counter()\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.203492Z","iopub.execute_input":"2025-09-24T17:36:50.203823Z","iopub.status.idle":"2025-09-24T17:36:50.210486Z","shell.execute_reply.started":"2025-09-24T17:36:50.203782Z","shell.execute_reply":"2025-09-24T17:36:50.209795Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"\nfgs_weight = 0.4 / 1.95 * 282\n# ======== Config ========\nclass Config:\n    DATA_PATH = '/kaggle/input/ariel-data-challenge-2025'\n    DATASET = \"train\"\n    PRETRAIN_PATH = \"/kaggle/input/ariel2025-5folds-resnetse/pytorch/full-sigma/1\"\n    SCALE = 0.952          # ※ #9 は今回は変更しない\n    SIGMA = 0.00055\n    \n    CUT_INF = 39\n    CUT_SUP = 321\n    \n    SENSOR_CONFIG = {\n        \"AIRS-CH0\": {\n            \"raw_shape\": [11250, 32, 356],\n            \"calibrated_shape\": [1, 32, CUT_SUP - CUT_INF],\n            \"linear_corr_shape\": (6, 32, 356),\n            \"dt_pattern\": (0.1, 4.5), \n            \"binning\": 30\n        },\n        \"FGS1\": {\n            \"raw_shape\": [135000, 32, 32],\n            \"calibrated_shape\": [1, 32, 32],\n            \"linear_corr_shape\": (6, 32, 32),\n            \"dt_pattern\": (0.1, 0.1),\n            \"binning\": 30 * 12\n        }\n    }\n    \n    MODEL_PHASE_DETECTION_SLICE = slice(30, 140)\n    MODEL_OPTIMIZATION_DELTA = 11\n    MODEL_POLYNOMIAL_DEGREE = 3\n    \n    N_JOBS = 3\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.2115Z","iopub.execute_input":"2025-09-24T17:36:50.211745Z","iopub.status.idle":"2025-09-24T17:36:50.236077Z","shell.execute_reply.started":"2025-09-24T17:36:50.211727Z","shell.execute_reply":"2025-09-24T17:36:50.235496Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Utils","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport pandas.api.types\nimport scipy.stats\n\n\nclass ParticipantVisibleError(Exception):\n    pass\n\n\ndef score(\n    solution: pd.DataFrame,\n    submission: pd.DataFrame,\n    row_id_column_name: str,\n    naive_mean: float,\n    naive_sigma: float,\n    fsg_sigma_true: float = 1e-6,\n    airs_sigma_true: float = 1e-5,\n    fgs_weight: float = 1,\n) -> float:\n    \"\"\"\n    This is a Gaussian Log Likelihood based metric. For a submission, which contains the predicted mean (x_hat) and variance (x_hat_std),\n    we calculate the Gaussian Log-likelihood (GLL) value to the provided ground truth (x). We treat each pair of x_hat,\n    x_hat_std as a 1D gaussian, meaning there will be 283 1D gaussian distributions, hence 283 values for each test spectrum,\n    the GLL value for one spectrum is the sum of all of them.\n\n    Inputs:\n        - solution: Ground Truth spectra (from test set)\n            - shape: (nsamples, n_wavelengths)\n        - submission: Predicted spectra and errors (from participants)\n            - shape: (nsamples, n_wavelengths*2)\n        naive_mean: (float) mean from the train set.\n        naive_sigma: (float) standard deviation from the train set.\n        fsg_sigma_true: (float) standard deviation from the FSG1 instrument for the test set.\n        airs_sigma_true: (float) standard deviation from the AIRS instrument for the test set.\n        fgs_weight: (float) relative weight of the fgs channel\n    \"\"\"\n\n    del solution[row_id_column_name]\n    del submission[row_id_column_name]\n\n    if submission.min().min() < 0:\n        raise ParticipantVisibleError('Negative values in the submission')\n    for col in submission.columns:\n        if not pandas.api.types.is_numeric_dtype(submission[col]):\n            raise ParticipantVisibleError(f'Submission column {col} must be a number')\n\n    n_wavelengths = len(solution.columns)\n    if len(submission.columns) != n_wavelengths * 2:\n        raise ParticipantVisibleError('Wrong number of columns in the submission')\n\n    y_pred = submission.iloc[:, :n_wavelengths].values\n    # Set a non-zero minimum sigma pred to prevent division by zero errors.\n    sigma_pred = np.clip(submission.iloc[:, n_wavelengths:].values, a_min=10**-15, a_max=None)\n    sigma_true = np.append(\n        np.array(\n            [\n                fsg_sigma_true,\n            ]\n        ),\n        np.ones(n_wavelengths - 1) * airs_sigma_true,\n    )\n    y_true = solution.values\n\n    GLL_pred = scipy.stats.norm.logpdf(y_true, loc=y_pred, scale=sigma_pred)\n    GLL_true = scipy.stats.norm.logpdf(y_true, loc=y_true, scale=sigma_true * np.ones_like(y_true))\n    GLL_mean = scipy.stats.norm.logpdf(y_true, loc=naive_mean * np.ones_like(y_true), scale=naive_sigma * np.ones_like(y_true))\n\n    # normalise the score, right now it becomes a matrix instead of a scalar.\n    ind_scores = (GLL_pred - GLL_mean) / (GLL_true - GLL_mean)\n\n    weights = np.append(np.array([fgs_weight]), np.ones(len(solution.columns) - 1))\n    weights = weights * np.ones_like(ind_scores)\n    submit_score = np.average(ind_scores, weights=weights)\n    return float(np.clip(submit_score, 0.0, 1.0))\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.237047Z","iopub.execute_input":"2025-09-24T17:36:50.237353Z","iopub.status.idle":"2025-09-24T17:36:50.255411Z","shell.execute_reply.started":"2025-09-24T17:36:50.237327Z","shell.execute_reply":"2025-09-24T17:36:50.254845Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n# ======== 位相補助 ========\ndef _phase_detector_signal(signal, cfg):\n    sl = cfg.MODEL_PHASE_DETECTION_SLICE\n    min_idx = int(np.argmin(signal[sl])) + sl.start\n    s1 = signal[:min_idx]; s2 = signal[min_idx:]\n    if s1.size < 3 or s2.size < 3:\n        return 0, len(signal) - 1\n    g1 = np.gradient(s1); g1_max = np.max(g1) if np.size(g1) else 0.0\n    g2 = np.gradient(s2); g2_max = np.max(g2) if np.size(g2) else 0.0\n    if g1_max != 0: g1 /= g1_max\n    if g2_max != 0: g2 /= g2_max\n    phase1 = int(np.argmin(g1)); phase2 = int(np.argmax(g2)) + min_idx\n    return phase1, phase2\n\n\n# ======== σ推定（#4: クリップ幅＆仕上げ係数を拡張） ========\ndef estimate_sigma_fgs(preprocessed_data, cfg):\n    \"\"\"FGS1用 σ のソフト推定（保守拡張版）\"\"\"\n    sig_rel = []\n    delta = cfg.MODEL_OPTIMIZATION_DELTA\n    eps = 1e-12\n    for single in preprocessed_data:\n        air_white = savgol_filter(single[:, 1:].mean(axis=1), 20, 2)\n        p1, p2 = _phase_detector_signal(air_white, cfg)\n        p1 = max(delta, p1)\n        p2 = min(len(air_white) - delta - 1, p2)\n\n        fgs = single[:, 0]\n        oot = (fgs[: p1 - delta] if p1 - delta > 0 else np.empty(0, fgs.dtype))\n        if p2 + delta < fgs.size:\n            oot = np.concatenate([oot, fgs[p2 + delta :]])\n        inn = fgs[p1 + delta : max(p1 + delta, p2 - delta)]\n\n        if oot.size == 0 or inn.size == 0:\n            sig_rel.append(np.nan); continue\n\n        n_oot, n_in = len(oot), len(inn)\n        var_oot = np.nanvar(oot, ddof=1)\n        var_in  = np.nanvar(inn, ddof=1)\n        oot_mean = float(np.nanmean(oot)) if np.isfinite(np.nanmean(oot)) else float(np.nanmean(fgs))\n        sigma_rel = np.sqrt(var_oot / max(n_oot,1) + var_in / max(n_in,1)) / max(oot_mean, eps)\n        sig_rel.append(sigma_rel)\n\n    s = np.asarray(sig_rel, dtype=float)\n    mask = np.isfinite(s) & (s > 0)\n    med = float(np.nanmedian(s[mask])) if mask.any() else 1.0\n\n    k = np.ones_like(s)\n    if med > 0 and np.isfinite(med):\n        k[mask] = np.sqrt(s[mask] / med)\n\n    # --- #4: clipをやや緩める（0.8–1.25 → 0.85–1.30）---\n    k = np.clip(k, 0.85, 1.30)\n\n    sigma_fgs = k * cfg.SIGMA\n\n    # --- #4: 仕上げの全体係数（わずかに拡張）---\n    sigma_fgs *= 1.04\n    return sigma_fgs\n\n\ndef estimate_sigma_air(preprocessed_data, cfg):\n    \"\"\"AIRS用 σ のソフト推定（保守拡張版）\"\"\"\n    sig_rel = []\n    delta = cfg.MODEL_OPTIMIZATION_DELTA\n    eps = 1e-12\n\n    for single in preprocessed_data:\n        white = np.nanmean(single[:, 1:], axis=1)\n        white_s = savgol_filter(white, 20, 2)\n\n        p1, p2 = _phase_detector_signal(white_s, cfg)\n        p1 = max(delta, p1)\n        p2 = min(len(white) - delta - 1, p2)\n\n        oot_left = white[: p1 - delta] if p1 - delta > 0 else np.empty(0, white.dtype)\n        oot_right = white[p2 + delta :] if (p2 + delta) < white.size else np.empty(0, white.dtype)\n        oot = np.concatenate([oot_left, oot_right]) if (oot_left.size + oot_right.size) else oot_left\n        inn = white[p1 + delta : max(p1 + delta, p2 - delta)]\n\n        if oot.size == 0 or inn.size == 0:\n            sig_rel.append(np.nan); continue\n\n        n_oot, n_in = len(oot), len(inn)\n        var_oot = np.nanvar(oot, ddof=1)\n        var_in  = np.nanvar(inn, ddof=1)\n        oot_mean = float(np.nanmean(oot)) if np.isfinite(np.nanmean(oot)) else float(np.nanmean(white))\n\n        sigma_rel = np.sqrt(var_oot / max(n_oot,1) + var_in / max(n_in,1)) / max(oot_mean, eps)\n        sig_rel.append(sigma_rel)\n\n    s = np.asarray(sig_rel, dtype=float)\n    mask = np.isfinite(s) & (s > 0)\n    med = float(np.nanmedian(s[mask])) if mask.any() else 1.0\n\n    k = np.ones_like(s)\n    if med > 0 and np.isfinite(med):\n        k[mask] = np.sqrt(s[mask] / med)\n\n    # --- #4: clipをやや緩める（0.90–1.20 → 0.92–1.22）---\n    k = np.clip(k, 0.92, 1.22)\n\n    sigma_air = k * cfg.SIGMA\n\n    # --- #4: 仕上げの全体係数（わずかに拡張）---\n    sigma_air *= 1.04\n    return sigma_air\n\n\n# ======== 前処理パイプライン ========\nclass SignalProcessor:\n    def __init__(self, config):\n        self.cfg = config\n        self.adc_info = pd.read_csv(f\"{self.cfg.DATA_PATH}/adc_info.csv\")\n        self.planet_ids = pd.read_csv(f'{self.cfg.DATA_PATH}/{self.cfg.DATASET}_star_info.csv', index_col='planet_id').index.astype(int)\n\n    def _apply_linear_corr(self, linear_corr, signal):\n        coeffs = np.flip(linear_corr, axis=0)\n        x = signal.astype(np.float64, copy=False)\n        out = np.empty_like(x, dtype=np.float64)\n        out[...] = coeffs[0]\n        for k in range(1, coeffs.shape[0]):\n            np.multiply(out, x, out=out)\n            out += coeffs[k]\n        return out.astype(signal.dtype, copy=False)\n\n    def _calibrate_single_signal(self, planet_id, sensor):\n        sensor_cfg = self.cfg.SENSOR_CONFIG[sensor]\n    \n        signal = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_signal_0.parquet\"\n        ).to_numpy()\n        dark = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/dark.parquet\"\n        ).to_numpy()\n        dead = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/dead.parquet\"\n        ).to_numpy()\n        flat = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/flat.parquet\"\n        ).to_numpy()\n        linear_corr = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/linear_corr.parquet\"\n        ).values.astype(np.float64).reshape(sensor_cfg[\"linear_corr_shape\"])\n    \n        signal = signal.reshape(sensor_cfg[\"raw_shape\"])\n        gain = self.adc_info[f\"{sensor}_adc_gain\"].iloc[0]\n        offset = self.adc_info[f\"{sensor}_adc_offset\"].iloc[0]\n        signal = signal / gain + offset\n    \n        hot = sigma_clip(dark, sigma=5, maxiters=5).mask\n    \n        if sensor == \"AIRS-CH0\":\n            signal = signal[:, :, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            linear_corr = linear_corr[:, :, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            dark = dark[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            dead = dead[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            flat = flat[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            hot = hot[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n    \n        if sensor == \"FGS1\":\n            y0, y1, x0, x1 = 10, 22, 10, 22\n            signal = signal[:, y0:y1, x0:x1]\n            dark   = dark[y0:y1, x0:x1]\n            dead   = dead[y0:y1, x0:x1]\n            flat   = flat[y0:y1, x0:x1]\n            linear_corr = linear_corr[:, y0:y1, x0:x1]\n            hot    = hot[y0:y1, x0:x1]\n    \n        np.maximum(signal, 0, out=signal)\n    \n        if sensor == \"FGS1\":\n            signal = self._apply_linear_corr(linear_corr, signal)\n        elif sensor == \"AIRS-CH0\":\n            sl = (slice(None), slice(10, 22), slice(None))\n            signal[sl] = self._apply_linear_corr(linear_corr[:, 10:22, :], signal[sl])\n        else:\n            signal = self._apply_linear_corr(linear_corr, signal)\n    \n        base_dt, increment = sensor_cfg[\"dt_pattern\"]\n        even_scale = base_dt\n        odd_scale  = base_dt + increment\n        signal[::2]  -= dark * even_scale\n        signal[1::2] -= dark * odd_scale\n    \n        if sensor == \"FGS1\":\n            flat_roi = flat.astype(signal.dtype, copy=False).copy()\n            bad = (dead) | ~np.isfinite(flat_roi) | (flat_roi == 0)\n            flat_roi[bad] = np.nan\n            signal /= flat_roi\n    \n        elif sensor == \"AIRS-CH0\":\n            y0, y1 = 10, 22\n            flat_roi = flat[y0:y1, :].astype(signal.dtype, copy=False).copy()\n            bad = (dead[y0:y1, :]) | ~np.isfinite(flat_roi) | (flat_roi == 0)\n            flat_roi[bad] = np.nan\n            signal[:, y0:y1, :] /= flat_roi\n    \n        else:\n            flat2 = flat.astype(signal.dtype, copy=False).copy()\n            bad2 = (dead) | ~np.isfinite(flat2) | (flat2 == 0)\n            flat2[bad2] = np.nan\n            signal /= flat2\n    \n        return signal\n\n    def _preprocess_calibrated_signal(self, calibrated_signal, sensor):\n        sensor_cfg = self.cfg.SENSOR_CONFIG[sensor]\n        binning = sensor_cfg[\"binning\"]\n\n        if sensor == \"AIRS-CH0\":\n            signal_roi = calibrated_signal[:, 10:22, :]\n        elif sensor == \"FGS1\":\n            signal_roi = calibrated_signal[:, 10:22, 10:22]\n            signal_roi = signal_roi.reshape(signal_roi.shape[0], -1)\n        \n        mean_signal = np.nanmean(signal_roi, axis=1)\n        cds_signal = mean_signal[1::2] - mean_signal[0::2]\n\n        n_bins = cds_signal.shape[0] // binning\n        binned = np.array([\n            cds_signal[j*binning : (j+1)*binning].mean(axis=0) \n            for j in range(n_bins)\n        ])\n\n        if sensor == \"AIRS-CH0\":\n            q_lo = np.nanpercentile(binned, 5.0, axis=1, keepdims=True)\n            q_hi = np.nanpercentile(binned, 95.0, axis=1, keepdims=True)\n            np.clip(binned, q_lo, q_hi, out=binned)\n\n        if sensor == \"FGS1\":\n            binned = binned.reshape((binned.shape[0], 1))\n\n        if sensor == \"AIRS-CH0\":\n            var = np.nanvar(binned, axis=0, ddof=1)\n            med = np.nanmedian(var)\n            safe_var = np.where(~np.isfinite(var) | (var <= 0), med if (np.isfinite(med) and med > 0) else 1.0, var)\n            w = 1.0 / safe_var\n\n            lo, hi = np.nanpercentile(w, 5.0), np.nanpercentile(w, 95.0)\n            if np.isfinite(lo) and np.isfinite(hi) and lo < hi:\n                w = np.clip(w, lo, hi)\n\n            M = binned.shape[1]\n            s = np.nansum(w)\n            if np.isfinite(s) and s > 0:\n                w = w * (M / s)\n            else:\n                w = np.ones_like(w)\n\n            binned *= w[None, :]\n\n        return binned\n\n    def _process_planet_sensor(self, args):\n        planet_id, sensor = args['planet_id'], args['sensor']\n        calibrated = self._calibrate_single_signal(planet_id, sensor)\n        preprocessed = self._preprocess_calibrated_signal(calibrated, sensor)\n        return preprocessed\n\n    def process_all_data(self):\n        args_fgs1 = [dict(planet_id=planet_id, sensor=\"FGS1\") for planet_id in self.planet_ids]\n        preprocessed_fgs1 = pqdm(args_fgs1, self._process_planet_sensor, n_jobs=self.cfg.N_JOBS)\n\n        args_airs_ch0 = [dict(planet_id=planet_id, sensor=\"AIRS-CH0\") for planet_id in self.planet_ids]\n        preprocessed_airs_ch0 = pqdm(args_airs_ch0, self._process_planet_sensor, n_jobs=self.cfg.N_JOBS)\n\n        preprocessed_signal = np.concatenate(\n            [np.stack(preprocessed_fgs1), np.stack(preprocessed_airs_ch0)], axis=2\n        )\n        return preprocessed_signal","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.256031Z","iopub.execute_input":"2025-09-24T17:36:50.256327Z","iopub.status.idle":"2025-09-24T17:36:50.290686Z","shell.execute_reply.started":"2025-09-24T17:36:50.256309Z","shell.execute_reply":"2025-09-24T17:36:50.290178Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Transit\n","metadata":{}},{"cell_type":"code","source":"class TransitModel:\n    def __init__(self, config):\n        self.cfg = config\n\n    def _phase_detector(self, signal):\n        search_slice = self.cfg.MODEL_PHASE_DETECTION_SLICE\n        min_index = np.argmin(signal[search_slice]) + search_slice.start\n        \n        signal1 = signal[:min_index]\n        signal2 = signal[min_index:]\n\n        grad1 = np.gradient(signal1)\n        grad1 /= grad1.max()\n        \n        grad2 = np.gradient(signal2)\n        grad2 /= grad2.max()\n\n        phase1 = np.argmin(grad1)\n        phase2 = np.argmax(grad2) + min_index\n\n        return phase1, phase2\n    \n    def _objective_function(self, s, signal, phase1, phase2):\n        delta = self.cfg.MODEL_OPTIMIZATION_DELTA\n        power = self.cfg.MODEL_POLYNOMIAL_DEGREE\n\n        if phase1 - delta <= 0 or phase2 + delta >= len(signal) or phase2 - delta - (phase1 + delta) < 5:\n            delta = 2\n\n        y = np.concatenate([\n            signal[: phase1 - delta],\n            signal[phase1 + delta : phase2 - delta] * (1 + s),\n            signal[phase2 + delta :]\n        ])\n        x = np.arange(len(y))\n\n        coeffs = np.polyfit(x, y, deg=power)\n        poly = np.poly1d(coeffs)\n        error = np.abs(poly(x) - y).mean()\n        \n        return error\n\n    def predict(self, single_preprocessed_signal):\n        signal_1d = single_preprocessed_signal[:, 1:].mean(axis=1)\n        signal_1d = savgol_filter(signal_1d, 23, 2)   # #2 は今回は触らない\n        \n        phase1, phase2 = self._phase_detector(signal_1d)\n\n        phase1 = max(self.cfg.MODEL_OPTIMIZATION_DELTA, phase1)\n        phase2 = min(len(signal_1d) - self.cfg.MODEL_OPTIMIZATION_DELTA - 1, phase2)    \n\n        result = minimize(\n            fun=self._objective_function,\n            x0=[0.0001],\n            args=(signal_1d, phase1, phase2),\n            method=\"Nelder-Mead\"\n        )\n        \n        return result.x[0]\n\n    def predict_all(self, preprocessed_signals):\n        predictions = [\n            self.predict(preprocessed_signal)\n            for preprocessed_signal in tqdm(preprocessed_signals)\n        ]\n        return np.array(predictions) * self.cfg.SCALE\n\n\n# ======== メタデータ ========\nStarInfo = pd.read_csv(ROOT_PATH + f\"/{MODE}_star_info.csv\")\nStarInfo[\"planet_id\"] = StarInfo[\"planet_id\"].astype(int)\nPlanetIds = StarInfo[\"planet_id\"].tolist()\nStarInfo = StarInfo.set_index(\"planet_id\")\n\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.291366Z","iopub.execute_input":"2025-09-24T17:36:50.2916Z","iopub.status.idle":"2025-09-24T17:36:50.322255Z","shell.execute_reply.started":"2025-09-24T17:36:50.291576Z","shell.execute_reply":"2025-09-24T17:36:50.321526Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# SE structure","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass SEBlockMLP(nn.Module):\n    \"\"\"Squeeze-and-Excitation block for MLP/tabular input (B, F)\"\"\"\n    def __init__(self, channels, reduction=8):\n        super().__init__()\n        self.fc1 = nn.Linear(channels, channels // reduction, bias=True)\n        self.fc2 = nn.Linear(channels // reduction, channels, bias=True)\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        # x: (B, F)\n        se = x.mean(dim=0, keepdim=True)           # (1, F)\n        se = F.relu(self.fc1(se), inplace=True)    # (1, F//r)\n        se = self.sigmoid(self.fc2(se))            # (1, F)\n        return x * se                              # (B, F)\n\nclass GaussianNoise(nn.Module):\n    \"\"\"Add Gaussian noise to input tensor (for tabular/MLP)\"\"\"\n    def __init__(self, stddev):\n        super().__init__()\n        self.stddev = stddev\n\n    def forward(self, x):\n        if self.training and self.stddev > 0:\n            noise = torch.randn_like(x, device=x.device) * self.stddev  # <-- FIXED\n            return x + noise\n        return x\n\nclass ResidualBlockMLP(nn.Module):\n    def __init__(self, dim, p=0.2, use_se=True, use_noise=False, noise_std=0.01):\n        super().__init__()\n        self.fc1 = nn.Linear(dim, dim)\n        self.bn1 = nn.BatchNorm1d(dim)\n        self.fc2 = nn.Linear(dim, dim)\n        self.bn2 = nn.BatchNorm1d(dim)\n        self.activation = nn.ReLU()\n        self.dropout = nn.Dropout(p)\n        self.se = SEBlockMLP(dim) if use_se else nn.Identity()\n        self.noise = GaussianNoise(noise_std) if use_noise else nn.Identity()\n\n    def forward(self, x):\n        identity = x\n        out = self.activation(self.bn1(self.fc1(x)))\n        out = self.dropout(out)\n        out = self.bn2(self.fc2(out))\n        out = self.se(out)\n        out = self.noise(out)\n        return self.activation(out + identity)\n\nclass ResNetMLP2(nn.Module):\n    def __init__(self, input_dim=9, hidden_dim=128, output_dim=283, num_blocks=3, dropout_rate=0.2, use_se=True, use_noise=False, noise_std=0.01):\n        super().__init__()\n        self.input_layer = nn.Linear(input_dim, hidden_dim)\n        self.blocks = nn.Sequential(*[\n            ResidualBlockMLP(hidden_dim, p=dropout_rate, use_se=use_se, use_noise=use_noise, noise_std=noise_std)\n            for _ in range(num_blocks)\n        ])\n        self.output_layer = nn.Linear(hidden_dim, output_dim)\n\n    def forward(self, x):\n        x = self.input_layer(x)\n        x = self.blocks(x)\n        x = self.output_layer(x)\n        return x\n    \n# class ResNetMLP2(nn.Module):\n#     def __init__(self, input_dim=9, hidden_dim=192, output_dim=282, num_blocks=12, dropout_rate=0.2,use_se=True, use_noise=False, noise_std=0.01):\n#         super().__init__()\n#         self.input_layer = nn.Sequential(\n#             nn.Linear(input_dim, hidden_dim),\n#             nn.LayerNorm(hidden_dim),\n#             nn.GELU()\n#         )\n#         self.blocks = nn.Sequential(*[ResidualBlockMLP(hidden_dim, p=dropout_rate) for _ in range(num_blocks)])\n#         self.output_mu = nn.Linear(hidden_dim, output_dim)\n#         self.output_log_sigma = nn.Linear(hidden_dim, output_dim)\n\n#     def forward(self, x):\n#         x = self.input_layer(x)\n#         x = self.blocks(x)\n#         mu = F.softplus(self.output_mu(x))  # smooth non-neg\n#         log_sigma = self.output_log_sigma(x)\n#         sigma = F.softplus(log_sigma) + 1e-6\n#         return mu, sigma\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.322916Z","iopub.execute_input":"2025-09-24T17:36:50.323138Z","iopub.status.idle":"2025-09-24T17:36:50.335108Z","shell.execute_reply.started":"2025-09-24T17:36:50.323105Z","shell.execute_reply":"2025-09-24T17:36:50.334567Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# # ======== AIRS 出力（282次元）モデル ========\n# class ResidualBlock2(nn.Module):\n#     def __init__(self, dim, p=0.2, activation=nn.ReLU):\n#         super().__init__()\n#         self.fc1 = nn.Linear(dim, dim)\n#         self.bn1 = nn.BatchNorm1d(dim)\n#         self.fc2 = nn.Linear(dim, dim)\n#         self.bn2 = nn.BatchNorm1d(dim)\n#         self.activation = activation()\n#         self.relu = nn.ReLU()\n#         self.dropout = nn.Dropout(p)\n\n#     def forward(self, x):\n#         identity = x\n#         out = self.relu(self.bn1(self.fc1(x)))\n#         out = self.dropout(out)\n#         out = self.bn2(self.fc2(out))\n#         return self.activation(out + identity)\n\n\n# class ResNetMLP2(nn.Module):\n#     def __init__(self, input_dim=3, hidden_dim=128, output_dim=282, num_blocks=3, dropout_rate=0.2):\n#         super().__init__()\n#         self.input_layer = nn.Linear(input_dim, hidden_dim)\n#         self.blocks = nn.Sequential(*[ResidualBlock2(hidden_dim, p=dropout_rate) for _ in range(num_blocks)])\n#         self.output_layer = nn.Linear(hidden_dim, output_dim)\n\n#     def forward(self, x):\n#         x = self.input_layer(x)\n#         x = self.blocks(x)\n#         x = self.output_layer(x)\n#         return x\n\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.335761Z","iopub.execute_input":"2025-09-24T17:36:50.335932Z","iopub.status.idle":"2025-09-24T17:36:50.35667Z","shell.execute_reply.started":"2025-09-24T17:36:50.335918Z","shell.execute_reply":"2025-09-24T17:36:50.356168Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Generator","metadata":{}},{"cell_type":"code","source":"\n# ======== 提出生成 ========\nclass SubmissionGenerator:\n    def __init__(self, config):\n        self.cfg = config\n        self.sample_submission = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/sample_submission.csv\", index_col=\"planet_id\")\n\n    def create(self, predictions1, predictions, sigma_fgs=None, sigma_air=None):\n        planet_ids = self.sample_submission.index\n        n_mu = self.sample_submission.shape[1] // 2  # 283\n\n        preds = np.asarray(predictions, dtype=float).reshape(-1)\n        mu = np.tile(preds.reshape(-1, 1), (1, n_mu))\n        mu = np.clip(mu, 0, None)\n\n        sigmas = np.full_like(mu, self.cfg.SIGMA, dtype=float)\n        if sigma_fgs is not None:\n            sigma_fgs = np.asarray(sigma_fgs, dtype=float).reshape(-1)\n            sigmas[:, 0] = np.clip(sigma_fgs, 1e-6, 0.1)\n        if sigma_air is not None:\n            sigma_air = np.asarray(sigma_air, dtype=float).reshape(-1, 1)\n            sigmas[:, 1:] = np.clip(sigma_air, 1e-6, 0.1)\n\n        submission_df = pd.DataFrame(\n            np.concatenate([mu, sigmas], axis=1),\n            columns=self.sample_submission.columns,\n            index=planet_ids\n        )\n        # μの1列目(白色FGS)と AIRS 1:282 を元の仕様どおり再配置\n        submission_df.iloc[:, 0] = predictions\n        submission_df.iloc[:, 1:283] = predictions1\n\n        submission_df.to_csv(\"submission.csv\")\n        return submission_df\n    def create_from_full(self, mu_full, sigma_fgs=None, sigma_air=None):\n        \"\"\"\n        mu_full: (n, 283) where [:,0]=FGS, [:,1:]=AIRS\n        \"\"\"\n        planet_ids = self.sample_submission.index\n        mu = np.asarray(mu_full, dtype=float)\n        assert mu.shape[1] == self.sample_submission.shape[1] // 2, f\"mu_full has wrong width {mu.shape[1]}\"\n        mu = np.clip(mu, 0, None)\n\n        sigmas = np.full_like(mu, self.cfg.SIGMA, dtype=float)\n        if sigma_fgs is not None:\n            sigmas[:, 0] = np.clip(np.asarray(sigma_fgs).reshape(-1), 1e-6, 0.1)\n        if sigma_air is not None:\n            if len(sigma_air[0]) == 282:\n                sigmas[:, 1:] = np.clip(np.asarray(sigma_air), 1e-6, 0.1)\n            else:\n                sigmas[:, 1:] = np.clip(np.asarray(sigma_air).reshape(-1, 1), 1e-6, 0.1)\n\n        submission_df = pd.DataFrame(\n            np.concatenate([mu, sigmas], axis=1),\n            columns=self.sample_submission.columns,\n            index=planet_ids\n        )\n        submission_df.to_csv(\"submission.csv\")\n        return submission_df\n    \n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.357416Z","iopub.execute_input":"2025-09-24T17:36:50.357652Z","iopub.status.idle":"2025-09-24T17:36:50.381786Z","shell.execute_reply.started":"2025-09-24T17:36:50.357627Z","shell.execute_reply":"2025-09-24T17:36:50.381304Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"markdown","source":"## Preprocess","metadata":{}},{"cell_type":"code","source":"config = Config()\nsignal_processor = SignalProcessor(config)\n\nif MODE == \"test\":\n    preprocessed_data = signal_processor.process_all_data()\n\nelse:\n    preprocessed_path = '/kaggle/input/train-preprocessed-ariel'\n    preprocessed_data = np.load(f'{preprocessed_path}/preprocessed_data.npy')\n# 1D深さ\nmodel = TransitModel(config)\npredictions = model.predict_all(preprocessed_data)\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:36:50.382379Z","iopub.execute_input":"2025-09-24T17:36:50.382634Z","iopub.status.idle":"2025-09-24T17:37:00.807111Z","shell.execute_reply.started":"2025-09-24T17:36:50.382608Z","shell.execute_reply":"2025-09-24T17:37:00.806402Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.preprocessing import StandardScaler\nimport joblib\n#load scaler\n\n# --- #4: 保守拡張版の σ 推定 ---\nsigma_fgs_vec = estimate_sigma_fgs(preprocessed_data, config)\nsigma_air_vec = estimate_sigma_air(preprocessed_data, config)\ntransit_df = pd.DataFrame({\n    \"planet_id\": PlanetIds,\n    \"transit_depth\": predictions\n})\ninput_df = pd.merge(transit_df, StarInfo, on=\"planet_id\", how=\"left\")\ninput_df[\"transit_depth\"] *= 10000\n# features = ['transit_depth','Rs','i']\nfeatures = [\n    'transit_depth', 'Rs', 'Ms', 'Ts', 'Mp', 'e', 'P', 'sma', 'i'\n]\n\n# if config.PRETRAIN_PATH:\n#     scaler = joblib.load(os.path.join(config.PRETRAIN_PATH, 'scaler.joblib'))\n#     X = scaler.transform(input_df[features].values.astype(np.float32))\n# else:\n#     scaler = StandardScaler()\n#     X = scaler.fit_transform(input_df[features].values.astype(np.float32))\n\n#save scaler\n\n# joblib.dump(scaler, 'scaler.joblib')\n\nX_tensor = torch.tensor(input_df[features].values.astype(np.float32), dtype=torch.float32)","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:00.80796Z","iopub.execute_input":"2025-09-24T17:37:00.808252Z","iopub.status.idle":"2025-09-24T17:37:02.926089Z","shell.execute_reply.started":"2025-09-24T17:37:00.808222Z","shell.execute_reply":"2025-09-24T17:37:02.925188Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## eval function","metadata":{}},{"cell_type":"code","source":"# ...existing code...\ndef calculate_competition_score_for_validation(predictions1, predictions_val, sigma_fgs_val, sigma_air_val, val_labels, naive_mean, naive_sigma):\n    \"\"\"\n    Build solution (284 incl. id) and submission (567 incl. id) and compute score.\n      - mu columns: wl_1..wl_283  (FGS is wl_1; AIRS are wl_2..wl_283)\n      - sigma columns: sigma_1..sigma_283\n    \"\"\"\n    try:\n        # Shapes\n        predictions1 = np.asarray(predictions1, dtype=float)\n        predictions_val = np.asarray(predictions_val, dtype=float).reshape(-1)\n        sigma_fgs_val = np.asarray(sigma_fgs_val, dtype=float).reshape(-1)\n        sigma_air_val = np.asarray(sigma_air_val, dtype=float).reshape(-1)\n\n        n = len(predictions_val)\n        assert predictions1.shape[0] == n, f\"predictions1 rows {predictions1.shape[0]} != {n}\"\n        # Ensure AIRS has 282 columns\n        if predictions1.ndim == 1:\n            predictions1 = predictions1.reshape(n, -1)\n        if predictions1.shape[1] == 283:\n            # If someone accidentally passed 283, drop one AIRS column (keep first 282)\n            predictions1 = predictions1[:, :282]\n        assert predictions1.shape[1] == 282, f\"AIRS width must be 282, got {predictions1.shape[1]}\"\n\n        # Column names (start at 1)\n        mu_cols = [f'wl_{i}' for i in range(1, 284)]        # 283\n        sigma_cols = [f'sigma_{i}' for i in range(1, 284)]   # 283\n        all_cols = mu_cols + sigma_cols                      # 566\n\n        # Submission matrix (n, 566)\n        sub = np.zeros((n, 566), dtype=float)\n\n        # Fill mu (clip non-negative)\n        sub[:, 0] = np.clip(predictions_val, 0.0, None)        # wl_1 (FGS)\n        sub[:, 1:283] = np.clip(predictions1, 0.0, None)       # wl_2..wl_283 (AIRS)\n\n        # Fill sigma with clipping\n        sub[:, 283] = np.clip(sigma_fgs_val, 1e-6, 0.1)        # sigma_1\n        sub[:, 284:] = np.clip(sigma_air_val.reshape(-1, 1), 1e-6, 0.1)  # broadcast to 282 cols\n\n        submission_df = pd.DataFrame(sub, columns=all_cols)\n        # Put planet_id as first column to total 567 cols\n        submission_df.insert(0, 'planet_id', val_labels.index.values)\n\n        # Solution: prefer named wl_1..wl_283; fallback to first 283 cols if names missing\n        if set(mu_cols).issubset(val_labels.columns):\n            solution_df = val_labels.loc[:, mu_cols].copy()\n        else:\n            solution_df = pd.DataFrame(val_labels.iloc[:, :283].to_numpy(), columns=mu_cols, index=val_labels.index)\n        solution_df.insert(0, 'planet_id', val_labels.index.values)\n        # print(submission_df.shape, submission_df.head(1))\n        # Compute competition score\n        return score(\n            solution=solution_df.copy(),\n            submission=submission_df.copy(),\n            row_id_column_name='planet_id',\n            naive_mean=naive_mean,\n            naive_sigma=naive_sigma,\n            fsg_sigma_true=1e-6,\n            airs_sigma_true=1e-5,\n            fgs_weight=fgs_weight\n        )\n    except Exception as e:\n        print(f\"Error calculating competition score: {e}\")\n        return 0.0\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:02.927081Z","iopub.execute_input":"2025-09-24T17:37:02.927455Z","iopub.status.idle":"2025-09-24T17:37:02.938537Z","shell.execute_reply.started":"2025-09-24T17:37:02.927426Z","shell.execute_reply":"2025-09-24T17:37:02.937765Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ...existing code...\ndef calculate_competition_score_from_full(pred_full, sigma_fgs_val, sigma_air_val, val_labels, naive_mean, naive_sigma, use_full_sigma_airs=False):\n    \"\"\"\n    pred_full: (n, 283) where [:,0] is FGS, [:,1:] is AIRS\n    \"\"\"\n    pred_full = np.asarray(pred_full, dtype=float)\n    assert pred_full.ndim == 2 and pred_full.shape[1] == 283, f\"expected (n,283), got {pred_full.shape}\"\n\n    n = pred_full.shape[0]\n    mu_cols = [f'wl_{i}' for i in range(1, 284)]\n    sigma_cols = [f'sigma_{i}' for i in range(1, 284)]\n    all_cols = mu_cols + sigma_cols\n\n    sub = np.zeros((n, 566), dtype=float)\n    sub[:, :283] = np.clip(pred_full, 0.0, None)\n    sub[:, 283] = np.clip(np.asarray(sigma_fgs_val).reshape(-1), 1e-6, 0.1)\n    if use_full_sigma_airs:\n        assert sigma_air_val.ndim == 2 and sigma_air_val.shape == (n, 282), f\"expected (n,282) for full_sigma_airs, got {sigma_air_val.shape}\"\n        sub[:, 284:] = np.clip(sigma_air_val, 1e-6, 0.1)\n    else:\n        sub[:, 284:] = np.clip(np.asarray(sigma_air_val).reshape(-1, 1), 1e-6, 0.1)\n\n    submission_df = pd.DataFrame(sub, columns=all_cols, index=val_labels.index)\n    submission_df.insert(0, 'planet_id', val_labels.index.values)\n\n    if set(mu_cols).issubset(val_labels.columns):\n        solution_df = val_labels.loc[:, mu_cols].copy()\n    else:\n        solution_df = pd.DataFrame(val_labels.iloc[:, :283].to_numpy(), columns=mu_cols, index=val_labels.index)\n    solution_df.insert(0, 'planet_id', val_labels.index.values)\n\n    return score(\n        solution=solution_df.copy(),\n        submission=submission_df.copy(),\n        row_id_column_name='planet_id',\n        naive_mean=naive_mean,\n        naive_sigma=naive_sigma,\n        fsg_sigma_true=1e-6,\n        airs_sigma_true=1e-5,\n        fgs_weight=fgs_weight\n    )\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:02.939369Z","iopub.execute_input":"2025-09-24T17:37:02.939568Z","iopub.status.idle":"2025-09-24T17:37:02.960859Z","shell.execute_reply.started":"2025-09-24T17:37:02.939552Z","shell.execute_reply":"2025-09-24T17:37:02.960183Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nset_seed(42)","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:02.961662Z","iopub.execute_input":"2025-09-24T17:37:02.962439Z","iopub.status.idle":"2025-09-24T17:37:02.984022Z","shell.execute_reply.started":"2025-09-24T17:37:02.962416Z","shell.execute_reply":"2025-09-24T17:37:02.983429Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_labels = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/train.csv\", index_col=\"planet_id\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:02.984819Z","iopub.execute_input":"2025-09-24T17:37:02.985058Z","iopub.status.idle":"2025-09-24T17:37:03.088667Z","shell.execute_reply.started":"2025-09-24T17:37:02.985038Z","shell.execute_reply":"2025-09-24T17:37:03.087833Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"resnet2 = ResNetMLP2(num_blocks=80, dropout_rate=0.3)","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.089445Z","iopub.execute_input":"2025-09-24T17:37:03.08966Z","iopub.status.idle":"2025-09-24T17:37:03.174972Z","shell.execute_reply.started":"2025-09-24T17:37:03.089643Z","shell.execute_reply":"2025-09-24T17:37:03.174187Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Prepare training data\n# Extract AIRS channels (columns 1:283) as targets - these are the mu values for wavelengths\ny_targets = train_labels.iloc[:, :283].values.astype(np.float32)  # Shape: (1100, 283)\ny_tensor = torch.tensor(y_targets, dtype=torch.float32)\n\nprint(f\"Input features shape: {X_tensor.shape}\")\nprint(f\"Target labels shape: {y_tensor.shape}\")\nprint(f\"Features: {features}\")\n\n# Scale targets for training (same scaling as used in inference)\ny_tensor_scaled = y_tensor * 10000\n\n# Create dataset and split\ndataset = TensorDataset(X_tensor, y_tensor_scaled)\n\n\n\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.175713Z","iopub.execute_input":"2025-09-24T17:37:03.175929Z","iopub.status.idle":"2025-09-24T17:37:03.185232Z","shell.execute_reply.started":"2025-09-24T17:37:03.175913Z","shell.execute_reply":"2025-09-24T17:37:03.184473Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_training_history(train_losses, val_losses):\n    # Plot training history\n    plt.figure(figsize=(12, 4))\n\n    plt.subplot(1, 2, 1)\n    plt.plot(train_losses, label='Training Loss')\n    plt.plot(val_losses, label='Validation Loss')\n    plt.title('Model Loss')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.legend()\n    plt.grid(True)\n\n    plt.subplot(1, 2, 2)\n    plt.plot(train_losses, label='Training Loss')\n    plt.plot(val_losses, label='Validation Loss')\n    plt.title('Model Loss (Log Scale)')\n    plt.xlabel('Epoch')\n    plt.ylabel('Loss')\n    plt.yscale('log')\n    plt.legend()\n    plt.grid(True)\n\n    plt.tight_layout()\n    plt.show()\n\n    # Load best model for evaluation\n    # resnet2.load_state_dict(torch.load('best_resnet_model.pth'))\n    # resnet2.eval()\n\n    # print(f\"Model loaded with best validation loss: {best_val_loss:.6f}\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.186264Z","iopub.execute_input":"2025-09-24T17:37:03.186737Z","iopub.status.idle":"2025-09-24T17:37:03.19954Z","shell.execute_reply.started":"2025-09-24T17:37:03.186713Z","shell.execute_reply":"2025-09-24T17:37:03.198822Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\nimport torch\ndef cleanup_memory():\n    \"\"\"Clean up GPU memory between folds\"\"\"\n    gc.collect()\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.200322Z","iopub.execute_input":"2025-09-24T17:37:03.200609Z","iopub.status.idle":"2025-09-24T17:37:03.215375Z","shell.execute_reply.started":"2025-09-24T17:37:03.200587Z","shell.execute_reply":"2025-09-24T17:37:03.214671Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# EMA utility for smoother weights\nclass EMA:\n    def __init__(self, model, decay=0.999):\n        self.decay = decay\n        self.shadow = {}\n        self.backup = {}\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.shadow[name] = param.data.clone()\n\n    @torch.no_grad()\n    def update(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                new_avg = self.decay * self.shadow[name] + (1.0 - self.decay) * param.data\n                self.shadow[name] = new_avg.clone()\n\n    def apply_shadow(self, model):\n        self.backup = {}\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.backup[name] = param.data.clone()\n                param.data = self.shadow[name].clone()\n\n    def restore(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad and name in self.backup:\n                param.data = self.backup[name].clone()\n        self.backup = {}","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.21606Z","iopub.execute_input":"2025-09-24T17:37:03.216323Z","iopub.status.idle":"2025-09-24T17:37:03.240698Z","shell.execute_reply.started":"2025-09-24T17:37:03.216297Z","shell.execute_reply":"2025-09-24T17:37:03.239988Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def weighted_huber(pred, tgt,col_w):\n    l = F.huber_loss(pred, tgt, reduction='none')  # (B, 283)\n    return (l * col_w).mean()\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.241474Z","iopub.execute_input":"2025-09-24T17:37:03.241695Z","iopub.status.idle":"2025-09-24T17:37:03.256322Z","shell.execute_reply.started":"2025-09-24T17:37:03.241676Z","shell.execute_reply":"2025-09-24T17:37:03.255722Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train function","metadata":{}},{"cell_type":"code","source":"def train_model(train_dataset, val_dataset, epochs, num_blocks=80, hidden_dim=128, batch_size=32, \n                dropout_rate=0.3, lr=0.001, patience=20, plot_training=False, \n                save_path='best_resnet_model.pth', val_indices=None, use_competition_scoring=False,\n                use_amp=True, warmup_epochs=3, max_grad_norm=1.0, ema_decay=0.999, pretrained_path=None):\n    \"\"\"\n    Train ResNet with AMP, gradient clipping, LR warmup, and EMA.\n    - If use_competition_scoring=True and val_indices is provided, early stopping tracks competition score (higher better).\n      Otherwise tracks validation loss (lower better).\n    - Best checkpoint saves EMA weights when EMA is enabled.\n    \"\"\"\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"Using device: {device}\")\n\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n    \n    # Model\n    resnet2 = ResNetMLP2(num_blocks=num_blocks, dropout_rate=dropout_rate, \n                         hidden_dim=hidden_dim,use_noise=True, noise_std=0.01)\n    if pretrained_path is not None:\n        resnet2.load_state_dict(torch.load(pretrained_path))\n    resnet2.to(device)\n\n    # Loss and optimizer\n    criterion = weighted_huber\n    col_w = np.ones(283, dtype=np.float32)\n    col_w[0] = fgs_weight  # e.g., 0.4/1.95*282\n    col_w = col_w / col_w.mean()\n    col_w = torch.tensor(col_w, dtype=torch.float32, device=device)\n    optimizer = torch.optim.AdamW(resnet2.parameters(), lr=lr, weight_decay=1e-4)\n\n    # Schedulers: warmup + plateau\n    def lr_lambda(epoch):\n        if warmup_epochs and epoch < warmup_epochs:\n            return float(epoch + 1) / float(max(1, warmup_epochs))\n        return 1.0\n    warmup_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n    plateau_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer, mode='min', factor=0.5, patience=10\n    )\n\n    # AMP scaler\n    use_cuda_amp = use_amp and (device.type == 'cuda')\n    scaler = torch.amp.GradScaler(device.type, enabled=use_cuda_amp)\n\n    # EMA\n    ema = EMA(resnet2, decay=ema_decay) if ema_decay and ema_decay < 1.0 else None\n\n    # Early stopping bookkeeping\n    best_metric = -float('inf') if use_competition_scoring else float('inf')\n    no_improve_count = 0\n    best_state = None\n\n    train_losses, val_losses, competition_scores = [], [], []\n    print(\"Starting training...\")\n\n    for epoch in range(epochs):\n        resnet2.train()\n        epoch_train_loss = 0.0\n\n        for data, target in train_loader:\n            data, target = data.to(device), target.to(device)\n\n            optimizer.zero_grad(set_to_none=True)\n            with torch.amp.autocast(device.type, enabled=use_cuda_amp):\n                output = resnet2(data)\n                loss = criterion(output, target, col_w)\n            scaler.scale(loss).backward()\n\n            if max_grad_norm and max_grad_norm > 0:\n                scaler.unscale_(optimizer)\n                torch.nn.utils.clip_grad_norm_(resnet2.parameters(), max_grad_norm)\n\n            scaler.step(optimizer)\n            scaler.update()\n\n            if ema is not None:\n                ema.update(resnet2)\n\n            epoch_train_loss += loss.item()\n\n        avg_train_loss = epoch_train_loss / max(1, len(train_loader))\n        train_losses.append(avg_train_loss)\n\n        # Validation\n        resnet2.eval()\n        epoch_val_loss = 0.0\n        val_predictions = []\n        if ema is not None:\n            ema.apply_shadow(resnet2)\n        with torch.no_grad():\n            for data, target in val_loader:\n                data, target = data.to(device), target.to(device)\n                with torch.amp.autocast(device.type, enabled=use_cuda_amp):\n                    output = resnet2(data)\n                    vloss = criterion(output, target, col_w)\n                epoch_val_loss += vloss.item()\n                val_predictions.append(output.detach().cpu().numpy())\n        if ema is not None:\n            ema.restore(resnet2)\n        avg_val_loss = epoch_val_loss / max(1, len(val_loader))\n        val_losses.append(avg_val_loss)\n        plateau_scheduler.step(avg_val_loss)\n        warmup_scheduler.step()\n\n        # Optional competition score on EMA weights for stability\n        comp_score = None\n        if use_competition_scoring and val_indices is not None:\n            if ema is not None:\n                ema.apply_shadow(resnet2)\n            try:\n                all_val_predictions = np.concatenate(val_predictions) / 10000.0\n                airs = all_val_predictions[:, 1:]\n                airs_s = savgol_filter(airs, window_length=13, polyorder=2, axis=1)\n                all_val_predictions[:, 1:] = 0.65 * airs + 0.35 * airs_s\n                # predictions_val = predictions[val_indices]\n                sigma_fgs_val = sigma_fgs_vec[val_indices]\n                sigma_air_val = sigma_air_vec[val_indices]\n                train_labels_val = train_labels.iloc[val_indices]\n                naive_mean = np.mean(train_labels.values[:, :283])\n                naive_sigma = np.std(train_labels.values[:, :283])\n                comp_score = calculate_competition_score_from_full(\n                    all_val_predictions,\n                    sigma_fgs_val,\n                    sigma_air_val,\n                    train_labels_val,\n                    naive_mean,\n                    naive_sigma\n                )\n                competition_scores.append(comp_score)\n            finally:\n                if ema is not None:\n                    ema.restore(resnet2)\n\n        # Select metric and check improvement\n        if use_competition_scoring and val_indices is not None:\n            current_metric = comp_score if comp_score is not None else 0.0\n            is_better = current_metric > best_metric\n        else:\n            current_metric = avg_val_loss\n            is_better = current_metric < best_metric\n\n        # Save best (EMA) weights\n        if is_better:\n            best_metric = current_metric\n            no_improve_count = 0\n            if ema is not None:\n                ema.apply_shadow(resnet2)\n                best_state = {k: v.cpu().clone() for k, v in resnet2.state_dict().items()}\n                ema.restore(resnet2)\n            else:\n                best_state = {k: v.cpu().clone() for k, v in resnet2.state_dict().items()}\n            torch.save(best_state, save_path)\n            if use_competition_scoring and val_indices is not None:\n                print(f'Epoch {epoch+1:3d}: lr={optimizer.param_groups[0][\"lr\"]:.2e} | train {avg_train_loss:.6f} | val {avg_val_loss:.6f} | comp {current_metric:.6f} [best]')\n            else:\n                print(f'Epoch {epoch+1:3d}: lr={optimizer.param_groups[0][\"lr\"]:.2e} | train {avg_train_loss:.6f} | val {avg_val_loss:.6f} [best]')\n                \n        else:\n            no_improve_count += 1\n            if use_competition_scoring and val_indices is not None:\n                print(f'Epoch {epoch+1:3d}: lr={optimizer.param_groups[0][\"lr\"]:.2e} | train {avg_train_loss:.6f} | val {avg_val_loss:.6f} | comp {current_metric if comp_score is not None else float(\"nan\"):.6f}')\n            else:\n                if (epoch + 1) % 5 == 0:\n                    print(f'Epoch {epoch+1:3d}: lr={optimizer.param_groups[0][\"lr\"]:.2e} | train {avg_train_loss:.6f} | val {avg_val_loss:.6f}')\n\n        if no_improve_count >= patience:\n            print(f'Early stopping at epoch {epoch+1}')\n            break\n        if (epoch + 1) % 10 == 0:\n            if comp_score is not None:\n                print(f\"Epoch {epoch+1}: Current competition score: {comp_score:.6f})\")\n    print(\"Training completed!\")\n    if use_competition_scoring and val_indices is not None:\n        print(f\"Best competition score: {best_metric:.6f}\")\n    else:\n        print(f\"Best validation loss: {best_metric:.6f}\")\n\n    # Load best weights into model\n    if best_state is not None:\n        resnet2.load_state_dict(best_state)\n\n    # Apply EMA shadow for returned model for smoother inference\n    if ema is not None:\n        ema.apply_shadow(resnet2)\n\n    if plot_training:\n        plot_training_history(train_losses, val_losses)\n\n    return resnet2, best_metric","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.257298Z","iopub.execute_input":"2025-09-24T17:37:03.25754Z","iopub.status.idle":"2025-09-24T17:37:03.279636Z","shell.execute_reply.started":"2025-09-24T17:37:03.257523Z","shell.execute_reply":"2025-09-24T17:37:03.278945Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data for train","metadata":{}},{"cell_type":"code","source":"# do 5 folds spliting\n# import subset\nfrom torch.utils.data import Subset\nfrom sklearn.model_selection import KFold, StratifiedKFold\n# Use a 1D proxy target for stratification (mean spectrum depth)\ntransit_depths = y_targets.mean(axis=1)\n\ndef make_strat_bins(y, n_splits=5, max_bins=12, min_bins=6):\n    y = pd.Series(y)\n    for n_bins in range(max_bins, min_bins - 1, -1):\n        y_binned = pd.qcut(y, q=n_bins, labels=False, duplicates='drop')\n        cnt = y_binned.value_counts().to_dict()\n        if min(cnt.values()) >= n_splits:\n            print(f\"Using {len(cnt)} quantile bins for stratification. Min per bin = {min(cnt.values())}\")\n            return y_binned.to_numpy(dtype=int), cnt\n    print(\"Warning: could not find valid quantile bins; falling back to a single bin.\")\n    return np.zeros(len(y), dtype=int), {0: len(y)}\n\ny_binned, bin_counts = make_strat_bins(transit_depths, n_splits=5, max_bins=12, min_bins=6)\nprint(\"Bin counts:\", dict(bin_counts))\n\nskf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)\nfold_competition_scores = []\nraw_X = X_tensor.numpy()","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.280473Z","iopub.execute_input":"2025-09-24T17:37:03.280784Z","iopub.status.idle":"2025-09-24T17:37:03.309732Z","shell.execute_reply.started":"2025-09-24T17:37:03.280692Z","shell.execute_reply":"2025-09-24T17:37:03.309056Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## params tuning for 5 folds","metadata":{}},{"cell_type":"code","source":"# Install optuna if not available\n# !pip install optuna\n\nimport optuna\nfrom sklearn.model_selection import KFold\nfrom torch.utils.data import Subset\nimport tempfile\nimport os\n\ndef objective(trial):\n    \"\"\"\n    Objective function for Optuna hyperparameter optimization using official competition scoring.\n    \"\"\"\n    # Suggest hyperparameters\n    lr = trial.suggest_float('lr', 1e-5, 1e-2, log=True)\n    batch_size = trial.suggest_categorical('batch_size', [16, 32, 64])\n    num_blocks = trial.suggest_int('num_blocks', 40, 120)\n    dropout_rate = trial.suggest_float('dropout_rate', 0.1, 0.5)\n    hidden_dim = trial.suggest_categorical('hidden_dim', [64, 128, 256])\n    ema_decay = trial.suggest_float('ema_decay', 0.9, 0.9999)\n    \n    print(f\"\\nTrial {trial.number}: lr={lr:.6f}, batch_size={batch_size}, \"\n          f\"num_blocks={num_blocks}, dropout_rate={dropout_rate:.3f}, hidden_dim={hidden_dim}\")\n    \n    # 5-fold cross-validation\n    kf = KFold(n_splits=5, shuffle=True, random_state=42)\n    best_scores = []\n    \n    for fold, (train_idx, val_idx) in enumerate(kf.split(raw_X)):\n        print(f\"  Fold {fold + 1}/5\", end=\" \")\n        cleanup_memory()\n\n        # Per-fold scaler and datasets\n        scaler_fold = StandardScaler().fit(raw_X[train_idx])\n        X_train = torch.tensor(scaler_fold.transform(raw_X[train_idx]), dtype=torch.float32)\n        X_val   = torch.tensor(scaler_fold.transform(raw_X[val_idx]),   dtype=torch.float32)\n        y_train = y_tensor_scaled[train_idx]\n        y_val   = y_tensor_scaled[val_idx]\n\n        train_subset = TensorDataset(X_train, y_train)\n        val_subset   = TensorDataset(X_val,   y_val)\n\n        # Unique temp model path per trial+fold\n        temp_save_path = os.path.join(tempfile.gettempdir(), f'temp_model_t{trial.number}_f{fold}.pth')\n\n        # Train (minimize val loss; skip comp score to save time)\n        _, best_score = train_model(\n            train_subset, \n            val_subset, \n            epochs=20,\n            batch_size=batch_size,\n            lr=lr,\n            num_blocks=num_blocks,\n            dropout_rate=dropout_rate,\n            hidden_dim=hidden_dim,\n            patience=10,\n            save_path=temp_save_path,\n            val_indices=None,                 # <-- skip comp score compute\n            use_competition_scoring=False,    # <-- minimize val loss\n            ema_decay=ema_decay,\n        )\n        best_scores.append(best_score)\n        \n        # with torch.no_grad():\n        #     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n        #     model_fold = model_fold.to(device)\n            \n        #     # Extract validation data only\n        #     X_val = X_tensor[val_idx].to(device)\n        #     outputs_val = model_fold(X_val).cpu().numpy() / 10000  # Scale back predictions\n            \n        #     # Get corresponding validation predictions and sigma values\n        #     predictions_val = predictions[val_idx]  # FGS predictions for validation samples\n        #     sigma_fgs_val = sigma_fgs_vec[val_idx]\n        #     sigma_air_val = sigma_air_vec[val_idx]\n            \n        #     # Get validation labels\n        #     train_labels_val = train_labels.iloc[val_idx]\n            \n        #     # Calculate competition score for validation samples\n        #     competition_score = calculate_competition_score_for_validation(\n        #         outputs_val,           # AIRS predictions for validation\n        #         predictions_val,       # FGS predictions for validation  \n        #         sigma_fgs_val,         # FGS sigma for validation\n        #         sigma_air_val,         # AIRS sigma for validation\n        #         train_labels_val,      # True labels for validation\n        #         naive_mean,\n        #         naive_sigma\n        #     )\n            \n        #     fold_competition_scores.append(competition_score)\n        #     print(f\"  Fold {fold + 1}/5 Competition Score: {competition_score:.6f}\")\n\n    # Return average competition score across folds (higher is better)\n    avg_score = np.mean(best_scores)\n    print(f\"  Average Competition Score: {avg_score:.6f}\")\n    \n    return avg_score","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.310624Z","iopub.execute_input":"2025-09-24T17:37:03.311321Z","iopub.status.idle":"2025-09-24T17:37:03.329294Z","shell.execute_reply.started":"2025-09-24T17:37:03.311296Z","shell.execute_reply":"2025-09-24T17:37:03.328599Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ...existing code...\n# study = optuna.create_study(\n#     direction='minimize',\n#     study_name='resnet_hyperparameter_tuning',\n#     sampler=optuna.samplers.TPESampler(seed=42)\n# )\n\n# n_trials = 30\n# study.optimize(objective, n_trials=n_trials, gc_after_trial=True)\n\n# # ...existing code...\n# print(f\"Best trial: {study.best_trial.number}\")\n# print(f\"Best val loss: {study.best_value:.6f}\")\n# print(\"Best parameters:\")\n# for key, value in study.best_params.items():\n#     print(f\"  {key}: {value}\")\n\n# print(f\"\\nTrial history:\")\n# for trial in study.trials[-5:]:\n#     print(f\"Trial {trial.number}: val_loss={trial.value:.6f}, params={trial.params}\")\n\n# best_params = study.best_params\n# best_ema_decay = best_params.get('ema_decay', 0.999)\n# print(f\"\\nBest parameters stored for final training: {best_params}\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.330539Z","iopub.execute_input":"2025-09-24T17:37:03.330857Z","iopub.status.idle":"2025-09-24T17:37:03.353597Z","shell.execute_reply.started":"2025-09-24T17:37:03.330834Z","shell.execute_reply":"2025-09-24T17:37:03.352988Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# best_params = {\n#     'lr': 0.004669,\n#     'batch_size': 32,\n#     'num_blocks': 72,\n#     'dropout_rate': 0.35,\n#     'hidden_dim': 64\n# }","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.354443Z","iopub.execute_input":"2025-09-24T17:37:03.354729Z","iopub.status.idle":"2025-09-24T17:37:03.378531Z","shell.execute_reply.started":"2025-09-24T17:37:03.354709Z","shell.execute_reply":"2025-09-24T17:37:03.377822Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train final model with best hyperparameters\n# print(\"\\n\" + \"=\"*50)\n# print(\"Training Final Model with Best Hyperparameters\")\n# print(\"=\"*50)\n\n# # Use best parameters from Optuna\n# best_lr = best_params['lr']\n# best_batch_size = best_params['batch_size']\n# best_num_blocks = best_params['num_blocks']\n# best_dropout_rate = best_params['dropout_rate']\n# best_hidden_dim = best_params['hidden_dim']\n# best_ema_decay = best_params.get('ema_decay', 0.999)\n# # save best params\n# with open('best_hyperparameters.txt', 'w') as f:\n#     for key, value in best_params.items():\n#         f.write(f\"{key}: {value}\\n\")\n\n# print(f\"Using best parameters:\")\n# print(f\"  Learning rate: {best_lr:.6f}\")\n# print(f\"  Batch size: {best_batch_size}\")\n# print(f\"  Number of blocks: {best_num_blocks}\")\n# print(f\"  Dropout rate: {best_dropout_rate:.3f}\")\n# print(f\"  Hidden dimension: {best_hidden_dim}\")\n\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.379285Z","iopub.execute_input":"2025-09-24T17:37:03.38023Z","iopub.status.idle":"2025-09-24T17:37:03.397061Z","shell.execute_reply.started":"2025-09-24T17:37:03.380206Z","shell.execute_reply":"2025-09-24T17:37:03.396452Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5 folds","metadata":{}},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\", category=UserWarning, module=\"torch.amp\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.397935Z","iopub.execute_input":"2025-09-24T17:37:03.398231Z","iopub.status.idle":"2025-09-24T17:37:03.417606Z","shell.execute_reply.started":"2025-09-24T17:37:03.39821Z","shell.execute_reply":"2025-09-24T17:37:03.416819Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# for fold, (train_idx, val_idx) in enumerate(skf.split(np.zeros(len(y_binned)), y_binned)):\n#     print(f\"Fold {fold + 1}\")\n#     #only train fold 4\n#     # if (fold != 4): \n#     #     continue\n#     os.makedirs(f'fold_{fold+1}', exist_ok=True)\n    \n#     # Per-fold scaler to avoid leakage\n#     scaler_fold = StandardScaler().fit(raw_X[train_idx])\n#     joblib.dump(scaler_fold, f'fold_{fold+1}/scaler.joblib')\n\n#     X_train = torch.tensor(scaler_fold.transform(raw_X[train_idx]), dtype=torch.float32)\n#     X_val   = torch.tensor(scaler_fold.transform(raw_X[val_idx]),   dtype=torch.float32)\n#     y_train = y_tensor_scaled[train_idx]\n#     y_val   = y_tensor_scaled[val_idx]\n\n#     train_subset = TensorDataset(X_train, y_train)\n#     val_subset   = TensorDataset(X_val,   y_val)\n\n#     #make directory for each fold\n#     os.makedirs(f'fold_{fold+1}', exist_ok=True)\n#     model_fold,fold_score=train_model(train_subset, val_subset, epochs=200, plot_training=True,\n#                            num_blocks=best_num_blocks, lr=best_lr, hidden_dim= best_hidden_dim, dropout_rate=best_dropout_rate,\n#                            batch_size=best_batch_size, patience=30,val_indices=val_idx,\n#                            save_path = f'fold_{fold+1}/best_resnet_model.pth',ema_decay=best_ema_decay)\n    \n#     #calculate gll for each fold\n#     naive_mean = np.mean(train_labels.values[:, :283])\n#     naive_sigma = np.std(train_labels.values[:, :283])\n    \n#     # ema = EMA(model_fold, decay=ema_decay) if ema_decay and ema_decay < 1.0 else None\n#     with torch.no_grad():\n#         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n#         model_fold = model_fold.to(device)\n#         # ema = EMA(model_fold, decay=ema_decay) if ema_decay and ema_decay < 1.0 else None\n#         # if ema is not None:\n#         #     ema.apply_shadow(model_fold)\n\n#         model_fold.eval()\n#         # Extract validation data only\n#         X_val = X_tensor[val_idx].to(device)\n#         outputs_val = model_fold(X_val).cpu().numpy() / 10000.0  # Scale back predictions\n        \n#         # Get corresponding validation predictions and sigma values\n#         predictions_val = predictions[val_idx]  # FGS predictions for validation samples\n#         sigma_fgs_val = sigma_fgs_vec[val_idx]\n#         sigma_air_val = sigma_air_vec[val_idx]\n        \n#         # Get validation labels\n#         train_labels_val = train_labels.iloc[val_idx]\n#         print('after model eval')\n#         # Calculate competition score for validation samples\n#         competition_score = calculate_competition_score_from_full(\n#             outputs_val, sigma_fgs_val, sigma_air_val, train_labels_val,\n#             naive_mean, naive_sigma\n#         )\n        \n#         fold_competition_scores.append(competition_score)\n#         print(f\"  Fold {fold + 1}/5 Competition Score: {competition_score:.6f}\")\n\n\n\n# print(f\"\\n=== Cross-Validation Results ===\")\n# avg_score = np.mean(fold_competition_scores)\n# for fold, fold_score in enumerate(fold_competition_scores):\n#     print(f\"Fold {fold + 1}: Competition Score = {fold_score:.6f}\")\n# print(f\"Average Competition Score: {avg_score:.6f}\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.418357Z","iopub.execute_input":"2025-09-24T17:37:03.418594Z","iopub.status.idle":"2025-09-24T17:37:03.435112Z","shell.execute_reply.started":"2025-09-24T17:37:03.418573Z","shell.execute_reply":"2025-09-24T17:37:03.434426Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # ...existing code...\n# for fold, (train_idx, val_idx) in enumerate(skf.split(np.zeros(len(y_binned)), y_binned)):\n\n#     model_path = f'fold_{fold+1}/best_resnet_model.pth'\n#     model = ResNetMLP2(num_blocks=best_num_blocks, dropout_rate=best_dropout_rate, \n#                          hidden_dim=best_hidden_dim, use_noise=True, noise_std=0.01)\n#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n#     model.load_state_dict(torch.load(model_path, map_location=device))\n#     model.eval()\n    \n#     model.to(device)\n#     scaler_fold = joblib.load(f'fold_{fold+1}/scaler.joblib')\n#     # Evaluate on THIS fold's validation indices\n#     with torch.no_grad():\n#         X_val = X_tensor[val_idx].to(device)\n#         X_val_scaled = torch.tensor(scaler_fold.transform(raw_X[val_idx]), dtype=torch.float32).to(device)\n#         ctx = torch.amp.autocast(device.type, enabled=torch.cuda.is_available())\n#         with ctx:\n#             out = model(X_val_scaled)\n#         # If model returns (mu, sigma), take mu\n#         if isinstance(out, (tuple, list)):\n#             out = out[0]\n#         outputs_val = out.detach().cpu().numpy() / 10000.0\n\n#     # Sanity check to avoid silent shape bugs\n#     assert outputs_val.shape == (len(val_idx), 283), f\"pred shape {outputs_val.shape} != ({len(val_idx)}, 283)\"\n#     airs = outputs_val[:, 1:]\n#     airs_s = savgol_filter(airs, window_length=13, polyorder=2, axis=1)\n#     outputs_val[:, 1:] = 0.65 * airs + 0.35 * airs_s\n#     # Calculate score for this fold\n#     predictions_val = predictions[val_idx]\n#     sigma_fgs_val   = sigma_fgs_vec[val_idx]\n#     sigma_air_val   = sigma_air_vec[val_idx]\n#     train_labels_val = train_labels.iloc[val_idx]\n#     naive_mean = np.mean(train_labels.values[:, :283])\n#     naive_sigma = np.std(train_labels.values[:, :283])\n\n#     competition_score = calculate_competition_score_from_full(\n#         outputs_val, sigma_fgs_val, sigma_air_val, train_labels_val,\n#         naive_mean, naive_sigma\n#     )\n#     print(f\"Fold {fold + 1}: Competition Score = {competition_score:.6f}\")\n# # ...existing code...\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.435986Z","iopub.execute_input":"2025-09-24T17:37:03.436271Z","iopub.status.idle":"2025-09-24T17:37:03.458382Z","shell.execute_reply.started":"2025-09-24T17:37:03.436247Z","shell.execute_reply":"2025-09-24T17:37:03.457658Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Make sigma model","metadata":{}},{"cell_type":"markdown","source":"## load mu models","metadata":{}},{"cell_type":"code","source":"#read hypperparameters from file\n# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n# print(f\"Using device: {device}\")\n# best_params = {}\n# with open(f'{config.PRETRAIN_PATH}/best_hyperparameters.txt', 'r') as f:\n#     for line in f:\n#         key, value = line.strip().split(': ')\n#         try:\n#             best_params[key] = float(value) if '.' in value else int(value)\n#         except ValueError:\n#             best_params[key] = value","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.459248Z","iopub.execute_input":"2025-09-24T17:37:03.459512Z","iopub.status.idle":"2025-09-24T17:37:03.480394Z","shell.execute_reply.started":"2025-09-24T17:37:03.45949Z","shell.execute_reply":"2025-09-24T17:37:03.479708Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# best_params","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.481185Z","iopub.execute_input":"2025-09-24T17:37:03.481431Z","iopub.status.idle":"2025-09-24T17:37:03.496481Z","shell.execute_reply.started":"2025-09-24T17:37:03.481407Z","shell.execute_reply":"2025-09-24T17:37:03.495827Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# from optuna\nbest_lr= 0.009946836274714647\nbatch_size= 32\nbest_num_blocks= 48\nbest_dropout_rate= 0.490927644772368\nbest_hidden_dim= 128\nbest_ema_decay= 0.9379123061634838\n# best_lr = best_params.get('lr', 0.004669)\n# best_batch_size = best_params.get('batch_size', 32)\n# best_num_blocks = best_params.get('num_blocks', 72)\n# best_dropout_rate = best_params.get('dropout_rate', 0.35)\n# best_hidden_dim = best_params.get('hidden_dim', 64)","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.497294Z","iopub.execute_input":"2025-09-24T17:37:03.497622Z","iopub.status.idle":"2025-09-24T17:37:03.512199Z","shell.execute_reply.started":"2025-09-24T17:37:03.497601Z","shell.execute_reply":"2025-09-24T17:37:03.511553Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load 5 folds models\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodels = []\nscalers = []\npretrain = config.PRETRAIN_PATH\nfor fold in range(5):\n    model_fold = ResNetMLP2(num_blocks=best_num_blocks, dropout_rate=best_dropout_rate, hidden_dim=best_hidden_dim)\n    model_fold.load_state_dict(torch.load(f'{pretrain}/fold_{fold+1}/best_resnet_model.pth', map_location=device))\n    model_fold = model_fold.to(device)\n    model_fold.eval()\n    models.append(model_fold)\n    scaler_fold = joblib.load(f'{pretrain}/fold_{fold+1}/scaler.joblib')\n    scalers.append(scaler_fold)\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:03.512809Z","iopub.execute_input":"2025-09-24T17:37:03.512982Z","iopub.status.idle":"2025-09-24T17:37:04.726846Z","shell.execute_reply.started":"2025-09-24T17:37:03.512968Z","shell.execute_reply":"2025-09-24T17:37:04.726071Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ...existing code...\n\nfor fold, (train_idx, val_idx) in enumerate(skf.split(np.zeros(len(y_binned)), y_binned)):\n\n        # model_path = f'fold_{fold+1}/best_resnet_model.pth'\n        model = models[fold]\n        scaler_fold = scalers[fold]\n        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n        # model.load_state_dict(torch.load(model_path, map_location=device))\n        model.eval()\n        \n        model.to(device)\n        # scaler_fold = joblib.load(f'fold_{fold+1}/scaler.joblib')\n        # Evaluate on THIS fold's validation indices\n        with torch.no_grad():\n            X_val = X_tensor[val_idx].to(device)\n            X_val_scaled = torch.tensor(scaler_fold.transform(raw_X[val_idx]), dtype=torch.float32).to(device)\n            ctx = torch.amp.autocast(device.type, enabled=torch.cuda.is_available())\n            with ctx:\n                out = model(X_val_scaled)\n            # If model returns (mu, sigma), take mu\n            if isinstance(out, (tuple, list)):\n                out = out[0]\n            outputs_val = out.detach().cpu().numpy() / 10000.0\n\n        # Sanity check to avoid silent shape bugs\n        assert outputs_val.shape == (len(val_idx), 283), f\"pred shape {outputs_val.shape} != ({len(val_idx)}, 283)\"\n        airs = outputs_val[:, 1:]\n        airs_s = savgol_filter(airs, window_length=13, polyorder=2, axis=1)\n        outputs_val[:, 1:] = 0.65 * airs + 0.35 * airs_s\n        # Calculate score for this fold\n        predictions_val = predictions[val_idx]\n        sigma_fgs_val   = sigma_fgs_vec[val_idx]\n        sigma_air_val   = sigma_air_vec[val_idx]\n        # repeat 282 for full_sigma_airs\n        full_sigma_airs_val = np.repeat(sigma_air_val[:, None], 282, axis=1)\n        train_labels_val = train_labels.iloc[val_idx]\n        naive_mean = np.mean(train_labels.values[:, :283])\n        naive_sigma = np.std(train_labels.values[:, :283])\n        # print(sigma_air_val)\n        competition_score = calculate_competition_score_from_full(\n            outputs_val, sigma_fgs_val, full_sigma_airs_val, train_labels_val,\n            naive_mean, naive_sigma, use_full_sigma_airs=True\n        )\n        print(f\"Fold {fold + 1}: Competition Score = {competition_score:.6f}\")\n# ...existing code...\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:04.727579Z","iopub.execute_input":"2025-09-24T17:37:04.727784Z","iopub.status.idle":"2025-09-24T17:37:05.25307Z","shell.execute_reply.started":"2025-09-24T17:37:04.727768Z","shell.execute_reply":"2025-09-24T17:37:05.252274Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Sigma net","metadata":{}},{"cell_type":"code","source":"# === 3) SigmaNet and differentiable competition score (torch) ===\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import TensorDataset, DataLoader\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import StandardScaler\nimport joblib\n\nclass SigmaNet(nn.Module):\n    def __init__(self, input_dim, hidden=128, blocks=3, drop=0.2):\n        super().__init__()\n        self.inp = nn.Sequential(\n            nn.Linear(input_dim, hidden),\n            nn.BatchNorm1d(hidden),\n            nn.ReLU()\n        )\n        layers = []\n        for _ in range(blocks):\n            layers += [\n                nn.Linear(hidden, hidden),\n                nn.BatchNorm1d(hidden),\n                nn.ReLU(),\n                nn.Dropout(drop)\n            ]\n        self.body = nn.Sequential(*layers)\n        self.out = nn.Linear(hidden, 283)  # [sigma_fgs, sigma_air]\n\n    def forward(self, x):\n        x = self.inp(x)\n        x = self.body(x)\n        s = self.out(x)\n        # Positive and bounded\n        s = F.softplus(s) + 1e-6\n        s = torch.clamp(s, 1e-6, 1e-1)\n        return s[:, :1], s[:, 1:]  # (B,1), (B,282)\n# ...existing code...\nclass SigmaNetMul(nn.Module):\n    \"\"\"\n    Predict multiplicative residual r over baseline sigmas (in features):\n      sigma_pred = clip(baseline * r, 1e-6, 1e-1)\n    We constrain r in [0.5, 1.5] via tanh for stability.\n    \"\"\"\n    def __init__(self, input_dim, hidden=128, blocks=3, drop=0.2):\n        super().__init__()\n        self.inp = nn.Sequential(\n            nn.Linear(input_dim, hidden),\n            nn.BatchNorm1d(hidden),\n            nn.ReLU()\n        )\n        layers = []\n        for _ in range(blocks):\n            layers += [\n                nn.Linear(hidden, hidden),\n                nn.BatchNorm1d(hidden),\n                nn.ReLU(),\n                nn.Dropout(drop)\n            ]\n        self.body = nn.Sequential(*layers)\n        self.out = nn.Linear(hidden, 283)  # residual multipliers per channel\n        # start at baseline: r=1.0 => tanh(0)=0; so zero init\n        nn.init.zeros_(self.out.weight)\n        nn.init.zeros_(self.out.bias)\n\n    def forward(self, x, base_fgs, base_air):\n        h = self.body(self.inp(x))\n        delta = self.out(h)                        # (B,283), unconstrained\n        r = 1.0 + 0.25 * torch.tanh(delta)         # r in [0.5, 1.5]\n        base_full = torch.cat(\n            [base_fgs, base_air.repeat(1, 282)], dim=1\n        )                                          # (B,283)\n        s = (base_full * r).clamp(1e-6, 1e-1)\n        return s[:, :1], s[:, 1:]                 # (B,1), (B,282)\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:05.25398Z","iopub.execute_input":"2025-09-24T17:37:05.254281Z","iopub.status.idle":"2025-09-24T17:37:05.264495Z","shell.execute_reply.started":"2025-09-24T17:37:05.254255Z","shell.execute_reply":"2025-09-24T17:37:05.263647Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sigma_model = SigmaNet(input_dim=X_tensor.shape[1], hidden=128, blocks=3, drop=0.2).to(device)\n#example output\nwith torch.no_grad():\n    example_out = sigma_model(X_tensor[:5].to(device))\n    print(f\"SigmaNet example output shapes: {example_out[0].shape}, {example_out[1].shape}\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:05.265279Z","iopub.execute_input":"2025-09-24T17:37:05.265557Z","iopub.status.idle":"2025-09-24T17:37:05.294461Z","shell.execute_reply.started":"2025-09-24T17:37:05.265533Z","shell.execute_reply":"2025-09-24T17:37:05.2937Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## data for training","metadata":{}},{"cell_type":"code","source":"# ...existing code...\n# === 1) Build OOF mu for train and ensemble stats from your 5 folds ===\nimport numpy as np\nimport torch\nfrom scipy.signal import savgol_filter\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nsigma_fgs_vec = estimate_sigma_fgs(preprocessed_data, config)\nsigma_air_vec = estimate_sigma_air(preprocessed_data, config)\nn = len(raw_X)\noof_mu = np.zeros((n, 283), dtype=np.float32)  # [:,0]=FGS, [:,1:]=AIRS\nseen = np.zeros(n, dtype=bool)\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:05.295207Z","iopub.execute_input":"2025-09-24T17:37:05.29547Z","iopub.status.idle":"2025-09-24T17:37:07.376411Z","shell.execute_reply.started":"2025-09-24T17:37:05.295449Z","shell.execute_reply":"2025-09-24T17:37:07.37565Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_data():\n    \n    # Use the same skf and fold order as your pretrained models/scalers\n    for fold, (train_idx, val_idx) in enumerate(skf.split(np.zeros(len(y_binned)), y_binned)):\n        model = models[fold].to(device).eval()\n        scaler = scalers[fold]\n        X_val_scaled = scaler.transform(raw_X[val_idx])\n        with torch.no_grad(), torch.amp.autocast(device.type, enabled=(device.type=='cuda')):\n            out = model(torch.tensor(X_val_scaled, dtype=torch.float32, device=device)).float().cpu().numpy()\n        mu_val = out / 10000.0\n        # Optional: keep your light smoothing for AIRS (same as in eval)\n        airs = mu_val[:, 1:]\n        airs_s = savgol_filter(airs, window_length=13, polyorder=2, axis=1)\n        mu_val[:, 1:] = 0.65 * airs + 0.35 * airs_s\n        oof_mu[val_idx] = mu_val\n        seen[val_idx] = True\n\n    assert seen.all(), \"OOF mu is missing some rows.\"\n\n    # Also compute per-sample ensemble stats across all 5 models (used both train/test)\n    per_model_preds = []\n    for fold in range(5):\n        model = models[fold].to(device).eval()\n        scaler = scalers[fold]\n        X_all_scaled = scaler.transform(raw_X)\n        with torch.no_grad(), torch.amp.autocast(device.type, enabled=(device.type=='cuda')):\n            out = model(torch.tensor(X_all_scaled, dtype=torch.float32, device=device)).float().cpu().numpy()\n        pmu = out / 10000.0\n        airs = pmu[:, 1:]\n        airs_s = savgol_filter(airs, window_length=13, polyorder=2, axis=1)\n        pmu[:, 1:] = 0.65 * airs + 0.35 * airs_s\n        per_model_preds.append(pmu)\n\n    per_model_preds = np.stack(per_model_preds, axis=0)  # (5, n, 283)\n    ens_mean = per_model_preds.mean(axis=0)             # (n, 283)\n    ens_std  = per_model_preds.std(axis=0)              # (n, 283)\n\n    # === 2) Build sigma-model features for train ===\n    # Base astro features (already created above as input_df and features list)\n    X_base = input_df[features].values.astype(np.float32)\n\n    # Simple mu-derived features (no labels):\n    mu_fgs_oof   = oof_mu[:, 0:1]\n    air_mean_oof = oof_mu[:, 1:].mean(axis=1, keepdims=True)\n    air_std_oof  = oof_mu[:, 1:].std(axis=1, keepdims=True)\n\n    # Ensemble dispersion features (across 5 models):\n    fgs_ens_std      = ens_std[:, 0:1]\n    air_mean_ens     = ens_mean[:, 1:].mean(axis=1, keepdims=True)\n    air_mean_ens_std = per_model_preds[:, :, 1:].mean(axis=2).std(axis=0, keepdims=False).reshape(-1, 1)\n\n    # Estimated sigma as features (not targets):\n    sigma_fgs_est = sigma_fgs_vec.reshape(-1, 1).astype(np.float32)\n    sigma_air_est = sigma_air_vec.reshape(-1, 1).astype(np.float32)\n    # sigma_air_est repeated for each wavelength channel\n    # sigma_air_est = np.repeat(sigma_air_vec[:, None], 282, axis=1).astype(np.float32)\n    print(sigma_air_est.shape)\n    # Final feature matrix for the sigma model\n    X_sigma = np.hstack([\n        X_base,\n        mu_fgs_oof, air_mean_oof, air_std_oof,\n        fgs_ens_std, air_mean_ens, air_mean_ens_std,\n        sigma_fgs_est, sigma_air_est\n    ]).astype(np.float32)\n\n    # Targets for sigma model training: we do NOT have sigma targets.\n    # We need y_true (train_labels) and fixed mu (oof_mu) for the loss.\n    Y_true = train_labels.iloc[:, :283].values.astype(np.float32)  # (n,283)\n    MU_fix = oof_mu.astype(np.float32)                             # (n,283)\n\n    print(\"Sigma features shape:\", X_sigma.shape)\n    print(\"Y_true shape:\", Y_true.shape, \"MU_fix shape:\", MU_fix.shape)\n    return X_sigma, Y_true, MU_fix\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.377266Z","iopub.execute_input":"2025-09-24T17:37:07.377643Z","iopub.status.idle":"2025-09-24T17:37:07.389299Z","shell.execute_reply.started":"2025-09-24T17:37:07.377615Z","shell.execute_reply":"2025-09-24T17:37:07.388719Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"X_sigma, Y_true, MU_fix = make_data()","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.389981Z","iopub.execute_input":"2025-09-24T17:37:07.390414Z","iopub.status.idle":"2025-09-24T17:37:07.794359Z","shell.execute_reply.started":"2025-09-24T17:37:07.39039Z","shell.execute_reply":"2025-09-24T17:37:07.793559Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def torch_competition_score(mu, sigma_fgs, full_sigma_air, y_true, naive_mean, naive_sigma, fgs_weight):\n    \"\"\"\n    mu:        (B, 283)\n    sigma_fgs: (B, 1)\n    full_sigma_air: (B, 282)\n    y_true:    (B, 283)\n    \"\"\"\n    log2pi = np.log(2.0 * np.pi)\n    device = mu.device\n\n    # sigma prediction (FGS 1 col + AIRS 282 cols)\n    sigma_pred = torch.cat([sigma_fgs, full_sigma_air], dim=1)  # (B,283)\n    # print(sigma_pred.shape)\n    sigma_pred = torch.clamp(sigma_pred, 1e-6, 0.1)\n\n    # logpdf(x|mu,sigma) = -0.5*((x-mu)/sigma)^2 - log(sigma) - 0.5*log(2π)\n    z = (y_true - mu) / sigma_pred\n    GLL_pred = -0.5 * z.pow(2) - torch.log(sigma_pred) - 0.5 * log2pi\n\n    # True sigma for normalization term (constants as in the comp)\n    sigma_true = torch.ones_like(y_true, device=device) * 1e-5\n    sigma_true[:, 0] = 1e-6\n    zt = (y_true - y_true) / sigma_true\n    GLL_true = -0.5 * zt.pow(2) - torch.log(sigma_true) - 0.5 * log2pi\n\n    # Naive baseline\n    y_mean = torch.ones_like(y_true, device=device) * naive_mean\n    y_std  = torch.ones_like(y_true, device=device) * naive_sigma\n    zn = (y_true - y_mean) / y_std\n    GLL_mean = -0.5 * zn.pow(2) - torch.log(y_std) - 0.5 * log2pi\n\n    denom = (GLL_true - GLL_mean).clamp_min(1e-12)\n    ind_scores = (GLL_pred - GLL_mean) / denom  # (B,283)\n\n    # wavelength weights\n    w = torch.ones((1, 283), device=device)\n    w[:, 0] = fgs_weight\n    # average over wavelengths then over batch\n    score_w = (ind_scores * w).sum(dim=1) / w.sum(dim=1)  # (B,)\n    return score_w.mean()  # scalar","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.795231Z","iopub.execute_input":"2025-09-24T17:37:07.795546Z","iopub.status.idle":"2025-09-24T17:37:07.802223Z","shell.execute_reply.started":"2025-09-24T17:37:07.795528Z","shell.execute_reply":"2025-09-24T17:37:07.80152Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train sigma","metadata":{}},{"cell_type":"code","source":"# ...existing code...\n# Train/val split for sigma model\nidx_train, idx_val = train_test_split(np.arange(n), test_size=0.2, random_state=42, shuffle=True)\n\n# Scale sigma features (save scaler for test)\n# >>> only scale all columns EXCEPT the last two baseline columns <<<\nIDX_FGS = X_sigma.shape[1] - 2\nIDX_AIR = X_sigma.shape[1] - 1\nif config.PRETRAIN_PATH is not None:\n    sigma_feat_scaler = joblib.load(f'{config.PRETRAIN_PATH}/sigma_features_scaler.joblib')\nelse:\n    sigma_feat_scaler = StandardScaler().fit(X_sigma[idx_train, :IDX_FGS])\njoblib.dump(sigma_feat_scaler, 'sigma_features_scaler.joblib')\n\nX_sigma_tr = np.hstack([\n    sigma_feat_scaler.transform(X_sigma[idx_train, :IDX_FGS]),\n    X_sigma[idx_train, IDX_FGS:]  # keep baseline sigmas unscaled\n]).astype(np.float32)\n\nX_sigma_va = np.hstack([\n    sigma_feat_scaler.transform(X_sigma[idx_val, :IDX_FGS]),\n    X_sigma[idx_val, IDX_FGS:]    # keep baseline sigmas unscaled\n]).astype(np.float32)\n\nds_tr = TensorDataset(\n    torch.tensor(X_sigma_tr, dtype=torch.float32),\n    torch.tensor(MU_fix[idx_train], dtype=torch.float32),\n    torch.tensor(Y_true[idx_train], dtype=torch.float32),\n)\nds_va = TensorDataset(\n    torch.tensor(X_sigma_va, dtype=torch.float32),\n    torch.tensor(MU_fix[idx_val], dtype=torch.float32),\n    torch.tensor(Y_true[idx_val], dtype=torch.float32),\n)\n# Globals for score normalization\nnaive_mean = float(Y_true.mean())\nnaive_sigma = float(Y_true.std())\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.803133Z","iopub.execute_input":"2025-09-24T17:37:07.803863Z","iopub.status.idle":"2025-09-24T17:37:07.83168Z","shell.execute_reply.started":"2025-09-24T17:37:07.803845Z","shell.execute_reply":"2025-09-24T17:37:07.830919Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# === 4) Train sigma model (maximize comp score) and print score per epoch ===\ndef train_sigma(ds_tr, ds_va, epochs=40, batch_size=64, lr=1e-3, hidden=128, blocks=3, drop=0.2, patience=8, ema_decay=None):\n    model = SigmaNetMul(input_dim=X_sigma.shape[1], hidden=hidden, blocks=blocks, drop=drop).to(device)\n    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)\n    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode='max', factor=0.5, patience=3)\n    use_amp = (device.type == 'cuda')\n    scaler = torch.amp.GradScaler(device.type, enabled=use_amp)\n    tr_loader = DataLoader(ds_tr, batch_size=batch_size, shuffle=True, drop_last=False)\n    va_loader = DataLoader(ds_va, batch_size=512, shuffle=False, drop_last=False)\n    if ema_decay is not None:\n        ema = EMA(model, decay=ema_decay)\n    best_score = -1e9\n    best_state = None\n    no_improve = 0\n    IDX_FGS = X_sigma.shape[1] - 2\n    IDX_AIR = X_sigma.shape[1] - 1\n    # print(IDX_FGS, IDX_AIR)\n    lam_reg = 1e-4  # tiny regularization to keep sigma near baseline early on\n\n    for epoch in range(1, epochs+1):\n        model.train()\n        epoch_scores = []\n\n        for xb, mu_b, yt_b in tr_loader:\n            # last two columns are raw baseline sigmas (not scaled)\n            base_fgs = xb[:, IDX_FGS:IDX_FGS+1]\n            base_air = xb[:, IDX_AIR:IDX_AIR+1]\n\n            xb = xb.to(device); mu_b = mu_b.to(device); yt_b = yt_b.to(device)\n            base_fgs = base_fgs.to(device); base_air = base_air.to(device)\n\n            opt.zero_grad(set_to_none=True)\n            with torch.amp.autocast(device.type, enabled=use_amp):\n                s_fgs, s_air = model(xb, base_fgs, base_air)\n                comp = torch_competition_score(mu_b, s_fgs, s_air, yt_b, naive_mean, naive_sigma, fgs_weight)\n\n                # Optional: small penalty to keep σ close to baseline at start\n                base_full = torch.cat([base_fgs, base_air.repeat(1, 282)], dim=1)\n                ratio = torch.clamp((torch.cat([s_fgs, s_air], dim=1) / (base_full + 1e-12)), 1e-6, 1e6)\n                reg = lam_reg * (torch.log(ratio).pow(2).mean())\n\n                loss = -(comp) + reg\n\n            scaler.scale(loss).backward()\n            scaler.step(opt)\n            scaler.update()\n            epoch_scores.append(comp.detach().item())\n\n        # Validation comp score\n        model.eval()\n        with torch.no_grad():\n            if ema_decay is not None:\n                ema = EMA(model, decay=ema_decay)\n                ema.apply_shadow(model)\n            va_scores = []\n            for xb, mu_b, yt_b in va_loader:\n                base_fgs = xb[:, IDX_FGS:IDX_FGS+1]  # (batch, 1)\n                base_air = xb[:, IDX_AIR:IDX_AIR+1]  # (batch, 1)\n                base_fgs = base_fgs.to(device)\n                base_air = base_air.to(device)\n                xb = xb.to(device); mu_b = mu_b.to(device); yt_b = yt_b.to(device)\n                s_fgs, s_air = model(xb, base_fgs, base_air)\n                comp = torch_competition_score(mu_b, s_fgs, s_air, yt_b, naive_mean, naive_sigma, fgs_weight)\n                va_scores.append(comp.item())\n            val_comp = float(np.mean(va_scores))\n            if ema_decay is not None:\n                ema.restore(model)\n        sched.step(val_comp)\n        \n\n        if val_comp > best_score + 1e-6:\n            if ema_decay is not None:\n                ema.apply_shadow(model)\n            best_score = val_comp\n            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}\n            no_improve = 0\n            torch.save(best_state, 'best_sigma_model.pth')\n            if ema_decay is not None:\n                ema.restore(model)\n            print(f\"Epoch {epoch:03d} | train_comp={np.mean(epoch_scores):.6f} | val_comp={val_comp:.6f} | lr={opt.param_groups[0]['lr']:.2e} [best]\")\n            \n        else:\n            no_improve += 1\n            if no_improve >= patience:\n                print(f\"Early stopping at epoch {epoch}\")\n                break\n            print(f\"Epoch {epoch:03d} | train_comp={np.mean(epoch_scores):.6f} | val_comp={val_comp:.6f} | lr={opt.param_groups[0]['lr']:.2e}\")\n    if best_state is not None:\n        model.load_state_dict(best_state)\n    print(f\"Best val competition score (sigma model): {best_score:.6f}\")\n    return model\nbest_sigma_params ={'lr': 0.0036867363843357476,\n                     'batch_size': 64, \n                     'hidden': 256, 'blocks': 11, \n                     'drop': 0.06418692636577976, \n                     'ema_decay': 0.9782537826235022}\nif config.PRETRAIN_PATH is not None:\n    #load sigma model, scaler, and skip training\n    print(\"Loading pretrained sigma model...\")\n    sigma_model = SigmaNetMul(input_dim=X_sigma.shape[1], hidden=256, blocks=11, drop=0.06418692636577976).to(device)\n    sigma_model.load_state_dict(torch.load(f'{config.PRETRAIN_PATH}/best_sigma_model.pth', map_location=device))\n# sigma_model = train_sigma(ds_tr, ds_va, epochs=200, batch_size=64, lr=0.001, hidden=128, blocks=10, drop=0.3, patience=70,ema_decay=0.999)\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.832476Z","iopub.execute_input":"2025-09-24T17:37:07.83303Z","iopub.status.idle":"2025-09-24T17:37:07.886678Z","shell.execute_reply.started":"2025-09-24T17:37:07.833012Z","shell.execute_reply":"2025-09-24T17:37:07.886164Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ...existing code...\n# === 5) Validate sigma model with the official score() on the held-out validation ===\ndef score_from_sigma_model(model, X_sigma_feats, mu_full, labels_df):\n    model.eval()\n    with torch.no_grad():\n        X_ = torch.tensor(X_sigma_feats, dtype=torch.float32, device=device)\n        # last two columns are raw baseline sigmas (unscaled)\n        IDX_FGS = X_sigma_feats.shape[1] - 2\n        IDX_AIR = X_sigma_feats.shape[1] - 1\n        base_fgs = X_[:, IDX_FGS:IDX_FGS+1]\n        base_air = X_[:, IDX_AIR:IDX_AIR+1]\n\n        s_fgs, s_air = model(X_, base_fgs, base_air)\n        s_fgs_np = s_fgs.squeeze(1).cpu().numpy()\n        s_air_np = s_air.cpu().numpy()\n        #print range for sigmas\n        print(f\"Sigma FGS: min={s_fgs_np.min():.6f}, max={s_fgs_np.max():.6f}, mean={s_fgs_np.mean():.6f}\")\n        print(f\"Sigma AIR: min={s_air_np.min():.6f}, max={s_air_np.max():.6f}, mean={s_air_np.mean():.6f}\")\n    val_score = calculate_competition_score_from_full(\n        mu_full, s_fgs_np, s_air_np, labels_df, naive_mean, naive_sigma, use_full_sigma_airs=True\n    )\n    return val_score\n\ntrain_labels_val = train_labels.iloc[idx_val]\nval_official = score_from_sigma_model(sigma_model, X_sigma_va, MU_fix[idx_val], train_labels_val)\nprint(f\"(Official) validation score with sigma model: {val_official:.6f}\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.887369Z","iopub.execute_input":"2025-09-24T17:37:07.887595Z","iopub.status.idle":"2025-09-24T17:37:07.924567Z","shell.execute_reply.started":"2025-09-24T17:37:07.887571Z","shell.execute_reply":"2025-09-24T17:37:07.923991Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#score from with sigma_fgs_vec, sigma_air_vec\ncalculate_competition_score_from_full(\n        MU_fix, sigma_fgs_vec, sigma_air_vec, train_labels, naive_mean, naive_sigma\n)","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:07.925222Z","iopub.execute_input":"2025-09-24T17:37:07.925477Z","iopub.status.idle":"2025-09-24T17:37:08.00074Z","shell.execute_reply.started":"2025-09-24T17:37:07.925451Z","shell.execute_reply":"2025-09-24T17:37:07.999949Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Infer","metadata":{}},{"cell_type":"code","source":"# Switch to test mode for inference\nMODE = \"test\"\n\n# Reload config and signal processor for test mode\nconfig_test = Config()\nconfig_test.DATASET = \"test\"\nsignal_processor_test = SignalProcessor(config_test)\n\n# Load test star info\nStarInfo_test = pd.read_csv(config_test.DATA_PATH + f\"/{config_test.DATASET}_star_info.csv\")\nStarInfo_test[\"planet_id\"] = StarInfo_test[\"planet_id\"].astype(int)\nPlanetIds_test = StarInfo_test[\"planet_id\"].tolist()\nStarInfo_test = StarInfo_test.set_index(\"planet_id\")\n\nprint(f\"Processing {len(PlanetIds_test)} test samples...\")\n\n# Process test data\npreprocessed_data_test = signal_processor_test.process_all_data()\n\n# Predict transit depths for test data\nmodel_test = TransitModel(config_test)\npredictions_test = model_test.predict_all(preprocessed_data_test)\n\nprint(f\"Transit depth predictions completed for test set\")\nprint(f\"Test predictions range: {predictions_test.min():.6f} to {predictions_test.max():.6f}\")\n# Estimate sigma values for test data\nsigma_fgs_vec_test = estimate_sigma_fgs(preprocessed_data_test, config_test)\nsigma_air_vec_test = estimate_sigma_air(preprocessed_data_test, config_test)\n\n# Create test features for ResNet model\ntransit_df_test = pd.DataFrame({\n    \"planet_id\": PlanetIds_test,\n    \"transit_depth\": predictions_test\n})\n\ninput_df_test = pd.merge(transit_df_test, StarInfo_test, on=\"planet_id\", how=\"left\")\ninput_df_test[\"transit_depth\"] *= 10000\n\n# Use same features as training\nfeatures = ['transit_depth','Rs','i','Ms','Ts','Mp','e','P','sma']\nX_tensor_test = torch.tensor(input_df_test[features].values.astype(np.float32), dtype=torch.float32)\n\nprint(f\"Test features shape: {X_tensor_test.shape}\")\nprint(f\"Features: {features}\")\nprint(\"\\nTest feature statistics:\")\nfor i, feature in enumerate(features):\n    print(f\"{feature}: mean={X_tensor_test[:, i].mean():.4f}, std={X_tensor_test[:, i].std():.4f}\")\n# Load the trained ResNet model for inference\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:08.001606Z","iopub.execute_input":"2025-09-24T17:37:08.001925Z","iopub.status.idle":"2025-09-24T17:37:11.929434Z","shell.execute_reply.started":"2025-09-24T17:37:08.001905Z","shell.execute_reply":"2025-09-24T17:37:11.928599Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ...existing code...\n# Use SAME feature order as training\n#make wl columns\nfeatures = ['transit_depth', 'Rs', 'Ms', 'Ts', 'Mp', 'e', 'P', 'sma', 'i']\nX_test_raw = input_df_test[features].values.astype(np.float32)\n\npredictions_list = []\nfor scaler, model in zip(scalers, models):\n    with torch.no_grad():\n        X_scaled = scaler.transform(X_test_raw)\n        preds = model(torch.tensor(X_scaled, dtype=torch.float32).to(device)).cpu().numpy()\n        predictions_list.append(preds)\n\npred_full_test = np.mean(predictions_list, axis=0) / 10000.0  # (n, 283)\n# Smooth/blend AIRS only (cols 1:)\nairs = pred_full_test[:, 1:]\nairs_s = savgol_filter(airs, window_length=13, polyorder=2, axis=1)\npred_full_test[:, 1:] = 0.65 * airs + 0.35 * airs_s\n\nprint(f\"Predictions shape: {pred_full_test.shape}\")\nprint(f\"AIRS range: {pred_full_test[:,1:].min():.6f}..{pred_full_test[:,1:].max():.6f} | FGS range: {pred_full_test[:,0].min():.6f}..{pred_full_test[:,0].max():.6f}\")","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:11.930343Z","iopub.execute_input":"2025-09-24T17:37:11.930614Z","iopub.status.idle":"2025-09-24T17:37:12.044253Z","shell.execute_reply.started":"2025-09-24T17:37:11.930597Z","shell.execute_reply":"2025-09-24T17:37:12.043628Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#make sigma features for test\ndef make_test_sigma_features(X_test_raw,pred_full_test):\n    # Base astro features (already created above as input_df and features list)\n    X_base = input_df_test[features].values.astype(np.float32)\n\n    # Simple mu-derived features (no labels):\n    mu_fgs_test   = pred_full_test[:, 0:1]\n    air_mean_test = pred_full_test[:, 1:].mean(axis=1, keepdims=True)\n    air_std_test  = pred_full_test[:, 1:].std(axis=1, keepdims=True)\n\n    # Ensemble dispersion features (across 5 models):\n    # For test, we don't have per-model preds, so we use zeros\n    fgs_ens_std      = np.zeros((len(X_test_raw), 1), dtype=np.float32)\n    air_mean_ens     = np.zeros((len(X_test_raw), 1), dtype=np.float32)\n    air_mean_ens_std = np.zeros((len(X_test_raw), 1), dtype=np.float32)\n\n    # Estimated sigma as features (not targets):\n    sigma_fgs_est = sigma_fgs_vec_test.reshape(-1, 1).astype(np.float32)\n    sigma_air_est = sigma_air_vec_test.reshape(-1, 1).astype(np.float32)\n\n    # Final feature matrix for the sigma model\n    X_sigma_test = np.hstack([\n        X_base,\n        mu_fgs_test, air_mean_test, air_std_test,\n        fgs_ens_std, air_mean_ens, air_mean_ens_std,\n        sigma_fgs_est, sigma_air_est\n    ]).astype(np.float32)\n\n    print(\"Test Sigma features shape:\", X_sigma_test.shape)\n    return X_sigma_test\nX_sigma_test = make_test_sigma_features(X_test_raw, pred_full_test)\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.044971Z","iopub.execute_input":"2025-09-24T17:37:12.045323Z","iopub.status.idle":"2025-09-24T17:37:12.052931Z","shell.execute_reply.started":"2025-09-24T17:37:12.045303Z","shell.execute_reply":"2025-09-24T17:37:12.052291Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sigma_feat_scaler = joblib.load('sigma_features_scaler.joblib')\nIDX_FGS = X_sigma_test.shape[1] - 2\nIDX_AIR = X_sigma_test.shape[1] - 1\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.053992Z","iopub.execute_input":"2025-09-24T17:37:12.054344Z","iopub.status.idle":"2025-09-24T17:37:12.077887Z","shell.execute_reply.started":"2025-09-24T17:37:12.054319Z","shell.execute_reply":"2025-09-24T17:37:12.077208Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"base_fgs_test = X_sigma_test[:, IDX_FGS:IDX_FGS+1]\nbase_air_test = X_sigma_test[:, IDX_AIR:IDX_AIR+1]\n","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.078976Z","iopub.execute_input":"2025-09-24T17:37:12.079285Z","iopub.status.idle":"2025-09-24T17:37:12.094156Z","shell.execute_reply.started":"2025-09-24T17:37:12.079263Z","shell.execute_reply":"2025-09-24T17:37:12.093623Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"base_fgs_test","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.094916Z","iopub.execute_input":"2025-09-24T17:37:12.095146Z","iopub.status.idle":"2025-09-24T17:37:12.112361Z","shell.execute_reply.started":"2025-09-24T17:37:12.095114Z","shell.execute_reply":"2025-09-24T17:37:12.111602Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Predict sigma","metadata":{}},{"cell_type":"code","source":"#make sigma predictions for test\nX_sigma_test_scaled = np.hstack([\n    sigma_feat_scaler.transform(X_sigma_test[:, :IDX_FGS]),\n    X_sigma_test[:, IDX_FGS:]  # keep baseline sigmas raw\n]).astype(np.float32)\nwith torch.no_grad():\n    X_ = torch.tensor(X_sigma_test_scaled, dtype=torch.float32, device=device)\n    base_fgs = torch.tensor(base_fgs_test, dtype=torch.float32, device=device)\n    base_air = torch.tensor(base_air_test, dtype=torch.float32, device=device)\n    s_fgs_test, s_air_test = sigma_model(X_, base_fgs, base_air)\n    s_fgs_np_test = s_fgs_test.squeeze(1).cpu().numpy()\n    s_air_np_test = s_air_test.cpu().numpy()","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.113183Z","iopub.execute_input":"2025-09-24T17:37:12.113696Z","iopub.status.idle":"2025-09-24T17:37:12.132479Z","shell.execute_reply.started":"2025-09-24T17:37:12.113672Z","shell.execute_reply":"2025-09-24T17:37:12.131914Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"s_fgs_np_test,s_air_np_test","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.133191Z","iopub.execute_input":"2025-09-24T17:37:12.133404Z","iopub.status.idle":"2025-09-24T17:37:12.154877Z","shell.execute_reply.started":"2025-09-24T17:37:12.133389Z","shell.execute_reply":"2025-09-24T17:37:12.154367Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Make final submission","metadata":{}},{"cell_type":"code","source":"submission_generator = SubmissionGenerator(config_test)\nsubmission_test = submission_generator.create_from_full(\n    pred_full_test,\n    sigma_fgs=s_fgs_np_test,\n    sigma_air=s_air_np_test\n)\nprint(\"Submission file 'submission.csv' generated successfully!\")\nprint(f\"Submission shape: {submission_test.shape}\")\nprint(\"\\nSubmission preview:\")\nprint(submission_test.head())\n# ...existing code...","metadata":{"execution":{"iopub.status.busy":"2025-09-24T17:37:12.155494Z","iopub.execute_input":"2025-09-24T17:37:12.155717Z","iopub.status.idle":"2025-09-24T17:37:12.190657Z","shell.execute_reply.started":"2025-09-24T17:37:12.155693Z","shell.execute_reply":"2025-09-24T17:37:12.189914Z"},"trusted":true},"outputs":[],"execution_count":null}]}