{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":101849,"databundleVersionId":13093295}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install pqdm astropy","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-03-25T06:18:16.398788Z","iopub.execute_input":"2026-03-25T06:18:16.399457Z","iopub.status.idle":"2026-03-25T06:18:20.693691Z","shell.execute_reply.started":"2026-03-25T06:18:16.399429Z","shell.execute_reply":"2026-03-25T06:18:20.692785Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport itertools\nfrom pathlib import Path\nfrom typing import Literal\n\nimport numpy as np\nimport pandas as pd\nfrom scipy.optimize import minimize\nfrom scipy.signal import savgol_filter\nfrom astropy.stats import sigma_clip\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.model_selection import KFold\nfrom pqdm.processes import pqdm\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T06:18:20.695304Z","iopub.execute_input":"2026-03-25T06:18:20.695568Z","iopub.status.idle":"2026-03-25T06:18:26.156364Z","shell.execute_reply.started":"2026-03-25T06:18:20.695538Z","shell.execute_reply":"2026-03-25T06:18:26.155773Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!kaggle kernels output ibrahimhabibeg/calibration-minimal-binning-and-no-channel-cut -p .","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T06:18:26.157309Z","iopub.execute_input":"2026-03-25T06:18:26.157779Z","iopub.status.idle":"2026-03-25T06:19:34.998369Z","shell.execute_reply.started":"2026-03-25T06:18:26.157752Z","shell.execute_reply":"2026-03-25T06:19:34.99737Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"INPUT_DATA_FOLDER = Path(\"/kaggle/input/competitions/ariel-data-challenge-2025\")\nCALIBRATED_DATA_FOLDER = Path(\"/kaggle/working/calibrated\")\n# if you're generating calibration on the fly:\n# CALIBRATED_DATA_FOLDER = Path(\"/kaggle/working/calibrated\")\nOUTPUT_FOLDER = Path(\"/kaggle/working\")\nMODEL_FOLDER = Path(\"/kaggle/working/models\")\nSUBMISSION_FILE = Path(\"/kaggle/working/submission.csv\")\n\nBINNING = 4\nMEAN_FGS_SIGMA = 0.00085\nMEAN_AIRS_SIGMA = 0.00065\nEPOCHS = 300\n\nos.makedirs(MODEL_FOLDER, exist_ok=True)\nos.makedirs(OUTPUT_FOLDER, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T06:19:35.000992Z","iopub.execute_input":"2026-03-25T06:19:35.001671Z","iopub.status.idle":"2026-03-25T06:19:35.009595Z","shell.execute_reply.started":"2026-03-25T06:19:35.00162Z","shell.execute_reply":"2026-03-25T06:19:35.008886Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Pre-Processing","metadata":{}},{"cell_type":"code","source":"class SGSmoothing:\n    def __init__(self, window_size: int = 5, poly_order: int = 2):\n        self.window_size = window_size\n        self.poly_order = poly_order\n\n    def smooth(self, signal: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Smooths the input signal using Savitzky-Golay filter.\n\n        Args:\n            signal (np.ndarray): 1D array with shape (num_time_steps)\n\n        Returns:\n            np.ndarray: Smoothed signal with the shape (num_time_steps)\n        \"\"\"\n        assert len(signal.shape) == 1, \"Expecting White curve (1D array)\"\n        return savgol_filter(signal, self.window_size, self.poly_order)\n\n\nclass FunctionFittingBasedPhaseDetector:\n    def __init__(self, window_size: int = 100, width: int = 300):\n        self.window_size = window_size\n        self.width = width\n\n    def _get_middle_of_phases_estimate(self, data: np.ndarray) -> tuple[int, int]:\n        kernel = np.ones(self.window_size) / self.window_size\n        moving_average = np.convolve(data, kernel, mode=\"valid\")\n        diff = np.diff(moving_average)\n        moving_average_diff = np.convolve(diff, kernel, mode=\"valid\")\n        min_idx = np.argmin(moving_average_diff)\n        max_idx = np.argmax(moving_average_diff)\n        min_idx_on_data = min_idx + self.window_size\n        max_idx_on_data = max_idx + self.window_size\n        return int(min_idx_on_data), int(max_idx_on_data)\n\n    def _cost_function(\n        self, params: tuple[float, float, float, float], data: np.ndarray, is_drop=True\n    ) -> float:\n        t1, t2, a, b = params\n        t1 = int(t1)\n        t2 = int(t2)\n        # Constrain Violation\n        if t1 > t2:\n            return (t1 - t2) * 1e9\n        if is_drop and a < b:\n            return (b - a) * 1e9\n        if not is_drop and a > b:\n            return (a - b) * 1e9\n        y = np.full((data.shape[0],), a)\n        y[t1:t2] = np.linspace(a, b, t2 - t1)\n        y[t2:] = b\n        cost = np.sum((data - y) ** 2)\n        return cost\n\n    def _get_phase_boundaries(\n        self, data: np.ndarray, region_estimate: int, is_drop=True\n    ) -> tuple[int, int]:\n        data = data[\n            max(0, region_estimate - self.width) : min(len(data), region_estimate + self.width)\n        ]\n        initial_params = [\n            max(0, self.width // 2),\n            min(self.width * 3 // 2, len(data) - 1),\n            np.mean(data[: max(1, self.width // 2)]),\n            np.mean(data[min(len(data) - 1, self.width * 3 // 2) :]),\n        ]\n        bounds = [\n            (1, len(data) - 2),\n            (2, len(data) - 1),\n            (min(data), max(data)),\n            (min(data), max(data)),\n        ]\n        result = minimize(\n            self._cost_function,\n            initial_params,\n            args=(data, is_drop),\n            bounds=bounds,\n            method=\"Nelder-Mead\",\n        )\n        phase_begin, phase_end, _, _ = result.x\n        phase_begin = int(phase_begin) + max(0, region_estimate - self.width)\n        phase_end = int(phase_end) + max(0, region_estimate - self.width)\n        return phase_begin, phase_end\n\n    def phase_detect(self, data: np.ndarray) -> tuple[int, int, int, int]:\n        assert len(data.shape) == 1, \"Expecting White Curve. Average over wavelengths first.\"\n        min_idx, max_idx = self._get_middle_of_phases_estimate(data)\n        drop_begin, drop_end = self._get_phase_boundaries(data.copy(), min_idx, is_drop=True)\n        rise_begin, rise_end = self._get_phase_boundaries(data.copy(), max_idx, is_drop=False)\n        return drop_begin, drop_end, rise_begin, rise_end\n\n    def phase_detect_multiple_planets(self, data: np.ndarray):\n        assert data.ndim == 2, \"Expecting 2D array with shape (num_planets, num_time_steps)\"\n        result = np.zeros((data.shape[0], 4), dtype=int)\n        for i in tqdm(range(data.shape[0])):\n            result[i] = self.phase_detect(data[i])\n        return result\n\nclass TransitMultiplicationFactorFinder:\n    def __init__(self, poly_degree: int = 3, error_degree: int = 1):\n        self.poly_degree = poly_degree\n        self.error_degree = error_degree\n\n    def _cost_function(\n        self, params: tuple[float], signal: np.ndarray, t1: int, t2: int, t3: int, t4: int\n    ) -> float:\n        s = params[0]\n        y = np.concatenate([signal[:t1], signal[t2:t3] * (s + 1.0), signal[t4:]])\n        x = np.arange(len(y))\n        coeffs = np.polyfit(x, y, deg=self.poly_degree)\n        poly = np.poly1d(coeffs)\n        fitted = poly(x)\n        cost = np.mean(np.abs(y - fitted) ** self.error_degree)\n        return float(cost)\n\n    def predict(self, signal: np.ndarray, t1: int, t2: int, t3: int, t4: int) -> float:\n        assert len(signal.shape) == 1, (\n            \"Signal must be a 1D array. Average across wavelengths before passing.\"\n        )\n        assert 0 <= t1 < t2 < t3 < t4 < len(signal), (\n            \"t1, t2, t3, t4 must satisfy 0 <= t1 < t2 < t3 < t4 < signal length.\"\n        )\n        initial_s = (\n            np.mean(np.concatenate([signal[:t1], signal[t4:]])) / np.mean(signal[t2:t3])\n        ) - 1.0\n        result = minimize(\n            self._cost_function,\n            x0=[initial_s],\n            args=(signal, t1, t2, t3, t4),\n            method=\"Nelder-Mead\",\n        )\n        return result.x[0]\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DataLoaderAndCalibrator:\n    def __init__(\n        self,\n        data_path: Path,\n        output_path: Path,\n        force_recalibration: bool = False,\n        data_cutoff: int | None = None,\n        binning: int = 4,\n        cut_airs_channels: bool = True,\n        n_jobs: int = 4,\n    ):\n        assert data_path.exists(), f\"Data path {data_path} does not exist.\"\n        self.data_path = data_path\n        self.adc_info = pd.read_csv(data_path / \"adc_info.csv\")\n        self.data_cutoff = data_cutoff\n        os.makedirs(output_path, exist_ok=True)\n        self.train_data_file = output_path / \"calibrated_train_data.npy\"\n        self.test_data_file = output_path / \"calibrated_test_data.npy\"\n        self.force_recalibration = force_recalibration\n        self.binning = binning\n        self.cut_airs_channels = cut_airs_channels\n        self.n_jobs = n_jobs\n        self.airs_first_channel_index = 39\n        self.airs_last_channel_index = 321\n        self.sensor_config = {\n            \"AIRS-CH0\": {\n                \"raw_shape\": [11250, 32, 356],\n                \"calibrated_shape\": [1, 32, 282 if cut_airs_channels else 356],\n                \"linear_corr_shape\": (6, 32, 356),\n                \"dt_pattern\": (0.1, 4.5),\n                \"binning\": binning,\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\": binning * 12,\n            },\n        }\n\n    def _apply_linear_corr(self, linear_corr, signal):\n        linear_corr_flipped = np.flip(linear_corr, axis=0)\n        corrected_signal = signal.copy()\n\n        for x, y in itertools.product(range(signal.shape[1]), range(signal.shape[2])):\n            poly = np.poly1d(linear_corr_flipped[:, x, y])\n            corrected_signal[:, x, y] = poly(corrected_signal[:, x, y])\n\n        return corrected_signal\n\n    def _calibrate_single_signal(self, planet_id, dataset, sensor):\n        sensor_cfg = self.sensor_config[sensor]\n\n        signal = pd.read_parquet(\n            f\"{self.data_path}/{dataset}/{planet_id}/{sensor}_signal_0.parquet\"\n        ).to_numpy()\n        dark = pd.read_parquet(\n            f\"{self.data_path}/{dataset}/{planet_id}/{sensor}_calibration_0/dark.parquet\"\n        ).to_numpy()\n        dead = pd.read_parquet(\n            f\"{self.data_path}/{dataset}/{planet_id}/{sensor}_calibration_0/dead.parquet\"\n        ).to_numpy()\n        flat = pd.read_parquet(\n            f\"{self.data_path}/{dataset}/{planet_id}/{sensor}_calibration_0/flat.parquet\"\n        ).to_numpy()\n        linear_corr = (\n            pd.read_parquet(\n                f\"{self.data_path}/{dataset}/{planet_id}/{sensor}_calibration_0/linear_corr.parquet\"\n            )\n            .values.astype(np.float64)\n            .reshape(sensor_cfg[\"linear_corr_shape\"])\n        )\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  # type: ignore\n\n        if sensor == \"AIRS-CH0\" and self.cut_airs_channels:\n            signal = signal[:, :, self.airs_first_channel_index : self.airs_last_channel_index]\n            linear_corr = linear_corr[\n                :, :, self.airs_first_channel_index : self.airs_last_channel_index\n            ]\n            dark = dark[:, self.airs_first_channel_index : self.airs_last_channel_index]\n            dead = dead[:, self.airs_first_channel_index : self.airs_last_channel_index]\n            flat = flat[:, self.airs_first_channel_index : self.airs_last_channel_index]\n            hot = hot[:, self.airs_first_channel_index : self.airs_last_channel_index]\n\n        base_dt, increment = sensor_cfg[\"dt_pattern\"]\n        dt = np.ones(len(signal)) * base_dt\n        dt[1::2] += increment\n\n        signal = signal.clip(0)\n        signal = self._apply_linear_corr(linear_corr, signal)\n        signal -= dark * dt[:, np.newaxis, np.newaxis]\n\n        flat = flat.reshape(sensor_cfg[\"calibrated_shape\"])\n        flat[dead.reshape(sensor_cfg[\"calibrated_shape\"])] = np.nan\n        flat[hot.reshape(sensor_cfg[\"calibrated_shape\"])] = np.nan\n\n        signal = signal / flat\n\n        return signal\n\n    def _preprocess_calibrated_signal(self, calibrated_signal, sensor):\n        sensor_cfg = self.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)  # type: ignore\n\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) for j in range(n_bins)]\n        )\n\n        if sensor == \"FGS1\":\n            binned = binned.reshape((binned.shape[0], 1))\n\n        return binned\n\n    def _process_planet_sensor(self, args):\n        planet_id, sensor, dataset = args[\"planet_id\"], args[\"sensor\"], args[\"dataset\"]\n        calibrated = self._calibrate_single_signal(planet_id, dataset, sensor)\n        preprocessed = self._preprocess_calibrated_signal(calibrated, sensor)\n        return preprocessed\n\n    def _process_all_data(self, dataset: Literal[\"train\", \"test\"]) -> np.ndarray:\n        \"\"\"\n        Process and preprocess signals for all planets and sensors.\n\n        Returns data of shape (num_planets, num_time_bins, total_channels)\n\n        Currently, num_time_bins is 187 (notice the binning) and total channels is 283 with the first being FGS1 and the rest being AIRS-CH0.\n\n        Returns:\n            np.ndarray: Preprocessed signals with shape (num_planets, num_time_bins, total_channels)\n        \"\"\"\n\n        # NOTE: The order of sensors is important here for the shape of the output array.\n        # NOTE: Currently, process only the first observation (index 0) of each planet.\n        planet_ids = pd.read_csv(f\"{self.data_path}/{dataset}_star_info.csv\")[\n            \"planet_id\"\n        ].values.astype(int)\n        if self.data_cutoff is not None:\n            planet_ids = planet_ids[: self.data_cutoff]\n\n        args_fgs1 = [\n            dict(planet_id=planet_id, dataset=dataset, sensor=\"FGS1\") for planet_id in planet_ids\n        ]\n\n        preprocessed_fgs1 = pqdm(args_fgs1, self._process_planet_sensor, n_jobs=self.n_jobs)\n\n        args_airs_ch0 = [\n            dict(planet_id=planet_id, dataset=dataset, sensor=\"AIRS-CH0\")\n            for planet_id in planet_ids\n        ]\n        preprocessed_airs_ch0 = pqdm(\n            args_airs_ch0, self._process_planet_sensor, n_jobs=self.n_jobs\n        )\n\n        preprocessed_signal = np.concatenate(\n            [np.stack(preprocessed_fgs1), np.stack(preprocessed_airs_ch0)], axis=2\n        )\n        return preprocessed_signal\n\n    def _load_train_labels(self) -> np.ndarray:\n        labels_df = pd.read_csv(self.data_path / \"train.csv\")\n        labels = labels_df.drop(columns=[\"planet_id\"]).to_numpy()\n        if self.data_cutoff is not None:\n            labels = labels[: self.data_cutoff]\n        return labels\n\n    def load_all_train_data(self) -> tuple[np.ndarray, np.ndarray]:\n        if not self.train_data_file.exists() or self.force_recalibration:\n            print(\"Calibrating and saving train data...\")\n            train_data = self._process_all_data(\"train\")\n            np.save(self.train_data_file, train_data)\n        else:\n            print(\"Loading calibrated train data...\")\n            train_data = np.load(self.train_data_file, allow_pickle=True)\n            if self.data_cutoff is not None:\n                train_data = train_data[: self.data_cutoff]\n        labels = self._load_train_labels()\n        assert len(train_data) == len(labels), (\n            \"Mismatch between data and labels lengths. Please recheck.\"\n        )\n        assert train_data.shape[2] == 283 if self.cut_airs_channels else 357, (\n            \"Unexpected number of channels in train data. This is probably due to previously saved data \"\n            \"with different channel cutting settings. Please delete the existing calibrated data file or set \"\n            \"force_recalibration to True.\"\n        )\n        return train_data, labels\n\n    def load_all_test_data(self) -> np.ndarray:\n        test_data = self._process_all_data(\"test\")\n        # np.save(self.test_data_file, test_data)\n        assert test_data.shape[2] == 283 if self.cut_airs_channels else 357, (\n            \"Unexpected number of channels in test data. Unknown reason. Please recheck.\"\n        )\n        return test_data\n\n\n\n\nclass WavelengthsGroupsMultiplierFinder:\n    def __init__(\n        self,\n        transit_finder: FunctionFittingBasedPhaseDetector = FunctionFittingBasedPhaseDetector(),\n        multiplier: TransitMultiplicationFactorFinder = TransitMultiplicationFactorFinder(),\n        smoother: SGSmoothing | None = SGSmoothing(window_size=150, poly_order=2),\n    ):\n        self.transit_finder = transit_finder\n        self.multiplier = multiplier\n        self.smoother = smoother\n\n    def extract_features(\n        self,\n        all_data: np.ndarray,\n        wavelengths_groups: list[int] = [1, 2, 4, 8, 16, 32, 64],\n        weights: list[float] = [1, 1, 1, 1, 1, 1, 1],\n        average_cross_groups: bool = True,\n        return_transit_locations: bool = False,\n    ) -> np.ndarray:\n        assert len(all_data.shape) == 3, (\n            \"Expecting 3D array: (num_planets, num_time_steps, num_wavelengths)\"\n        )\n        assert all_data.shape[2] == 283, \"Expecting the AIRS channels to be cut\"\n        assert len(wavelengths_groups) == len(weights), (\n            \"wavelengths_groups and weights must have the same length\"\n        )\n\n        num_planets = all_data.shape[0]\n        num_wavelengths = all_data.shape[2]\n\n        num_groups = len(wavelengths_groups)\n        feats = np.zeros((num_planets, num_wavelengths, num_groups))\n        ranges_list = [np.array_split(np.arange(num_wavelengths), cb) for cb in wavelengths_groups]\n\n        transit_location = (\n            np.zeros((num_planets, 4), dtype=int) if return_transit_locations else None\n        )\n\n        for i in tqdm(range(num_planets)):\n            data = all_data[i]\n            t1, t2, t3, t4 = self.transit_finder.phase_detect(\n                self.smoother.smooth(data.mean(axis=1)) if self.smoother else data.mean(axis=1)\n            )\n\n            if return_transit_locations:\n                transit_location[i] = (t1, t2, t3, t4)  # type: ignore\n\n            for j, ranges in enumerate(ranges_list):\n                for r in ranges:\n                    signal = data[:, r].mean(axis=1)\n                    signal = self.smoother.smooth(signal) if self.smoother else signal\n                    s = self.multiplier.predict(signal, t1, t2, t3, t4)\n                    feats[i, r, j] = s\n\n        if average_cross_groups:\n            feats = np.average(feats, axis=2, weights=weights)\n\n        if return_transit_locations:\n            return feats, transit_location  # type: ignore\n\n        return feats\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model Classes","metadata":{}},{"cell_type":"code","source":"\nclass ResNetBlock(nn.Module):\n    def __init__(self, in_features, hidden_features=None):\n        super().__init__()\n        hidden_features = hidden_features or in_features\n        self.fc1 = nn.Linear(in_features, hidden_features)\n        self.fc2 = nn.Linear(hidden_features, in_features)\n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        identity = x\n        out = self.relu(self.fc1(x))\n        out = self.fc2(out)\n        out += identity\n        out = self.relu(out)\n        return out\n\n\nclass SValuesCNNWithStarInfoTorch(nn.Module):\n    def __init__(\n        self,\n        num_channels=6,\n        wavelengths=283,\n        num_star_features=3,\n        num_resnet_blocks=3,\n        resnet_hidden=512,\n    ):\n        super().__init__()\n        self.num_channels = num_channels\n        self.wavelengths = wavelengths\n        self.num_star_features = num_star_features\n\n        # CNN layers for s-values\n        self.cnn1 = nn.Conv1d(num_channels, 32, kernel_size=5, padding=\"same\")\n        self.cnn2 = nn.Conv1d(32, 16, kernel_size=7, padding=\"same\")\n        self.cnn3 = nn.Conv1d(16, 1, kernel_size=9, padding=\"same\")\n\n        # ResNet blocks for fully connected layers\n        self.input_fc = nn.Linear(wavelengths + num_star_features, resnet_hidden)\n        self.resnet_blocks = nn.Sequential(\n            *[ResNetBlock(resnet_hidden) for _ in range(num_resnet_blocks)]\n        )\n        self.output_fc = nn.Linear(resnet_hidden, wavelengths)\n\n    def forward(self, s_values, star_info):\n        # s_values: (batch_size, num_channels, wavelengths)\n        x = F.relu(self.cnn1(s_values))\n        x = F.relu(self.cnn2(x))\n        x = self.cnn3(x)  # (batch_size, 1, wavelengths)\n        x = x.squeeze(1)  # (batch_size, wavelengths)\n\n        # star_info: (batch_size, num_star_features)\n        # Concatenate along feature axis\n        x = torch.cat([x, star_info], dim=1)  # (batch_size, wavelengths + num_star_features)\n\n        x = F.relu(self.input_fc(x))\n        x = self.resnet_blocks(x)\n        x = self.output_fc(x)  # (batch_size, wavelengths)\n        return x\n\n\nclass _SValuesDataset(torch.utils.data.Dataset):\n    def __init__(self, s_values, star_info, targets=None):\n        self.s_values = torch.tensor(s_values, dtype=torch.float32)\n        self.star_info = torch.tensor(star_info, dtype=torch.float32)\n        self.targets = torch.tensor(targets, dtype=torch.float32) if targets is not None else None\n\n    def __len__(self):\n        return len(self.s_values)\n\n    def __getitem__(self, idx):\n        if self.targets is not None:\n            return self.s_values[idx], self.star_info[idx], self.targets[idx]\n        else:\n            return self.s_values[idx], self.star_info[idx]\n\n\nclass SValuesCNNWithStarInfoModel:\n    def __init__(\n        self,\n        weights_path: Path,\n        num_channels: int = 6,\n        wavelengths: int = 283,\n        num_star_features: int = 3,\n        num_resnet_blocks: int = 3,\n        resnet_hidden: int = 512,\n        learning_rate: float = 1e-3,\n        batch_size: int = 32,\n        device: torch.device = torch.device(\n            \"cuda\"\n            if torch.cuda.is_available()\n            else \"mps\"\n            if torch.backends.mps.is_available()\n            else \"cpu\"\n        ),\n    ):\n        self.weights_path = weights_path\n        self.device = device\n        self.model = SValuesCNNWithStarInfoTorch(\n            num_channels, wavelengths, num_star_features, num_resnet_blocks, resnet_hidden\n        ).to(self.device)\n        self.learning_rate = learning_rate\n        self.batch_size = batch_size\n        self.wavelengths = wavelengths\n\n    def _create_dataloader(\n        self,\n        s_values: np.ndarray,\n        star_info: np.ndarray,\n        targets: np.ndarray = None,\n        shuffle: bool = True,\n    ):\n        dataset = _SValuesDataset(s_values, star_info, targets)\n        return torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=shuffle)\n\n    def _rmse(self, preds: torch.Tensor, targets: torch.Tensor) -> float:\n        return torch.sqrt(torch.mean((preds - targets) ** 2)).item()\n\n    def train(\n        self,\n        train_s_values: np.ndarray,\n        train_star_info: np.ndarray,\n        train_targets: np.ndarray,\n        val_s_values: np.ndarray,\n        val_star_info: np.ndarray,\n        val_targets: np.ndarray,\n        epochs: int = 200,\n        force_training: bool = False,\n    ) -> tuple[float, float, np.ndarray]:\n        if os.path.exists(self.weights_path) and not force_training:\n            print(f\"Weights found at {self.weights_path}. Skipping training.\")\n            self.model.load_state_dict(torch.load(self.weights_path, map_location=self.device))\n            val_preds = self.predict(val_s_values, val_star_info)\n            val_loss = self._rmse(torch.tensor(val_preds), torch.tensor(val_targets))\n            train_loss = None\n            return val_loss, train_loss, val_preds\n\n        train_loader = self._create_dataloader(train_s_values, train_star_info, train_targets)\n        val_loader = self._create_dataloader(\n            val_s_values, val_star_info, val_targets, shuffle=False\n        )\n\n        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, \"min\", patience=10, factor=0.5\n        )\n        criterion = nn.MSELoss()\n\n        best_val_loss = float(\"inf\")\n        train_losses, val_losses = [], []\n\n\n        progress_bar = tqdm(range(epochs), desc=\"Training\", leave=True)\n        for epoch in progress_bar:\n            self.model.train()\n            running_train_loss = 0.0\n            for s_vals, star_inf, targets in train_loader:\n                s_vals = s_vals.to(self.device)\n                star_inf = star_inf.to(self.device)\n                targets = targets.to(self.device)\n                optimizer.zero_grad()\n                outputs = self.model(s_vals, star_inf)\n                loss = criterion(outputs, targets)\n                loss.backward()\n                optimizer.step()\n                running_train_loss += loss.item() * s_vals.size(0)\n            train_loss = running_train_loss / len(train_loader.dataset)\n            train_losses.append(train_loss)\n\n            self.model.eval()\n            running_val_loss = 0.0\n            val_preds = []\n            for s_vals, star_inf, targets in val_loader:\n                s_vals = s_vals.to(self.device)\n                star_inf = star_inf.to(self.device)\n                targets = targets.to(self.device)\n                with torch.no_grad():\n                    outputs = self.model(s_vals, star_inf)\n                val_preds.append(outputs.cpu())\n                loss = criterion(outputs, targets)\n                running_val_loss += loss.item() * s_vals.size(0)\n            val_loss = running_val_loss / len(val_loader.dataset)\n            val_losses.append(val_loss)\n            scheduler.step(val_loss)\n            val_preds = torch.cat(val_preds, dim=0).numpy()\n            rmse_val = self._rmse(torch.tensor(val_preds), torch.tensor(val_targets))\n            progress_bar.set_description(f\"Epoch {epoch+1}/{epochs}\")\n            progress_bar.set_postfix({\n                \"train_loss\": f\"{train_loss:.4f}\",\n                \"val_loss\": f\"{val_loss:.4f}\",\n                \"val_rmse\": f\"{rmse_val:.4f}\"\n            })\n            if val_loss < best_val_loss:\n                best_val_loss = val_loss\n                torch.save(self.model.state_dict(), self.weights_path)\n\n        # Load best weights\n        self.model.load_state_dict(torch.load(self.weights_path, map_location=self.device))\n        val_preds = self.predict(val_s_values, val_star_info)\n        final_val_loss = self._rmse(torch.tensor(val_preds), torch.tensor(val_targets))\n        final_train_loss = train_losses[-1] if train_losses else None\n        return final_val_loss, final_train_loss, val_preds\n\n    def predict(self, s_values: np.ndarray, star_info: np.ndarray) -> np.ndarray:\n        self.model.eval()\n        if os.path.exists(self.weights_path):\n            print(f\"Loading weights from {self.weights_path}\")\n            self.model.load_state_dict(torch.load(self.weights_path, map_location=self.device))\n        else:\n            raise FileNotFoundError(\n                f\"Weights file not found at {self.weights_path}. Please train the model first.\"\n            )\n        loader = self._create_dataloader(s_values, star_info, targets=None, shuffle=False)\n        preds = []\n        with torch.no_grad():\n            for s_vals, star_inf in loader:\n                s_vals = s_vals.to(self.device)\n                star_inf = star_inf.to(self.device)\n                outputs = self.model(s_vals, star_inf)\n                preds.append(outputs.cpu())\n        preds = torch.cat(preds, dim=0).numpy()\n        return preds  # (num_samples, wavelengths)\n\n\nclass SValuesCNNWithStarInfoTrainer:\n    def __init__(\n        self,\n        weights_dir: Path,\n        n_splits: int = 5,\n        num_channels: int = 6,\n        wavelengths: int = 283,\n        num_star_features: int = 3,\n        num_resnet_blocks: int = 3,\n        resnet_hidden: int = 512,\n        learning_rate: float = 1e-3,\n        batch_size: int = 32,\n        device: torch.device = torch.device(\n            \"cuda\"\n            if torch.cuda.is_available()\n            else \"mps\"\n            if torch.backends.mps.is_available()\n            else \"cpu\"\n        ),\n    ):\n        self.n_splits = n_splits\n        self.weights_dir = Path(weights_dir)\n        self.weights_dir.mkdir(parents=True, exist_ok=True)\n        self.model_params = dict(\n            num_channels=num_channels,\n            wavelengths=wavelengths,\n            num_star_features=num_star_features,\n            num_resnet_blocks=num_resnet_blocks,\n            resnet_hidden=resnet_hidden,\n            learning_rate=learning_rate,\n            batch_size=batch_size,\n            device=device,\n        )\n        self.models = []\n\n    def train(\n        self,\n        s_values: np.ndarray,\n        star_info: np.ndarray,\n        targets: np.ndarray,\n        epochs: int = 200,\n        force_training: bool = False,\n    ) -> tuple[list[float], list[float], np.ndarray]:\n        kf = KFold(n_splits=self.n_splits, shuffle=True, random_state=42)\n        val_losses, train_losses = [], []\n        num_planets = s_values.shape[0]\n        wavelengths = self.model_params[\"wavelengths\"]\n        val_preds_full = np.zeros((num_planets, wavelengths), dtype=np.float32)\n        self.models = []\n        for fold, (train_idx, val_idx) in enumerate(kf.split(s_values)):\n            print(f\"Fold {fold + 1}/{self.n_splits}\")\n            train_s, val_s = s_values[train_idx], s_values[val_idx]\n            train_star, val_star = star_info[train_idx], star_info[val_idx]\n            train_t, val_t = targets[train_idx], targets[val_idx]\n            weights_path = self.weights_dir / f\"model_fold_{fold + 1}.pt\"\n            model = SValuesCNNWithStarInfoModel(weights_path, **self.model_params)\n            val_loss, train_loss, val_preds = model.train(\n                train_s,\n                train_star,\n                train_t,\n                val_s,\n                val_star,\n                val_t,\n                epochs=epochs,\n                force_training=force_training,\n            )\n            val_losses.append(val_loss)\n            train_losses.append(train_loss)\n            val_preds_full[val_idx] = val_preds\n            self.models.append(model)\n        return val_losses, train_losses, val_preds_full\n\n    def predict(self, s_values: np.ndarray, star_info: np.ndarray) -> np.ndarray:\n        if not self.models:\n            # Load models from weights_dir\n            self.models = []\n            for fold in range(1, self.n_splits + 1):\n                weights_path = self.weights_dir / f\"model_fold_{fold}.pt\"\n                model = SValuesCNNWithStarInfoModel(weights_path, **self.model_params)\n                self.models.append(model)\n        preds = [model.predict(s_values, star_info) for model in self.models]\n        preds = np.stack(preds, axis=0)  # (n_splits, num_samples, wavelengths)\n        return np.mean(preds, axis=0)  # (num_samples, wavelengths)\n\nSValuesCNNWithStarInfoModel(weights_path, **model_params)\n\n\n\nclass LROnVariousFeaturesSigmaCalculator:\n    def __init__(\n        self,\n        mean_fgs_sigma: float = 9.0e-4,\n        mean_airs_sigma: float = 5.5e-4,\n        fgs_min_sigma: float = 1e-6,\n        airs_min_sigma: float = 1e-6,\n        eps: float = 1e-12,\n        smoother: SGSmoothing = SGSmoothing(window_size=150, poly_order=2),\n        transit_finder: FunctionFittingBasedPhaseDetector = FunctionFittingBasedPhaseDetector(),\n    ):\n        self.mean_fgs_sigma = mean_fgs_sigma\n        self.mean_airs_sigma = mean_airs_sigma\n        self.fgs_min_sigma = fgs_min_sigma\n        self.airs_min_sigma = airs_min_sigma\n        self.eps = eps\n        self.smoother = smoother\n        self.transit_finder = transit_finder\n        self.models = []\n\n    def _create_model(self):\n        model = LinearRegression()\n        return model\n\n    def _get_white_curve_var(\n        self,\n        data: np.ndarray,\n        transit_locations: np.ndarray,\n        mean_sigma: float = 5.5e-4,\n        min_sigma: float = 1e-6,\n        eps=1e-12,\n    ) -> np.ndarray:\n        sigma = np.zeros(data.shape[0])\n\n        for i in tqdm(range(data.shape[0]), desc=\"Calculating AIRS sigma\"):\n            signal = data[i]\n            t1, t2, t3, t4 = transit_locations[i]\n\n            airs = signal[:, 1:].mean(axis=1)\n\n            oot = np.concatenate([airs[:t1], airs[t4:]], axis=0)\n            inn = airs[t2:t3]\n\n            var_oot = np.var(oot)\n            var_inn = np.var(inn)\n            oot_mean = np.mean(oot)\n\n            sigma[i] = np.sqrt((var_inn / len(inn) + var_oot / len(oot)).mean()) / max(\n                oot_mean, eps\n            )\n\n        sigma_scaled = sigma / np.mean(sigma) * mean_sigma\n        sigma_scaled = sigma_scaled.clip(min=min_sigma)\n        return sigma_scaled\n\n    def _extract_features(self, data: np.ndarray, spectrum: np.ndarray):\n        spectrum_mean = spectrum.mean(axis=1)\n        spectrum_std = spectrum.std(axis=1)\n        # spectrum_var = spectrum.var(axis=1)\n\n        smoothed_white_curves = np.zeros((data.shape[0], data.shape[1]))\n        for i in range(data.shape[0]):\n            smoothed_white_curves[i] = self.smoother.smooth(data[i, :, 1:].mean(axis=1))\n        transit_locations = self.transit_finder.phase_detect_multiple_planets(\n            smoothed_white_curves\n        )\n\n        white_curve_var = self._get_white_curve_var(\n            data=data.copy(),\n            transit_locations=transit_locations,\n            mean_sigma=self.mean_airs_sigma,\n            eps=self.eps,\n        )\n\n        features = np.concatenate(\n            [\n                spectrum_mean[:, np.newaxis],\n                spectrum_std[:, np.newaxis],\n                # spectrum_var[:, np.newaxis],\n                white_curve_var[:, np.newaxis],\n            ],\n            axis=1,\n        )\n        return features\n\n    def _get_optimal_sigma(self, spectrum: np.ndarray, labels: np.ndarray):\n        def cost_function(params, x, mu):\n            sigma = params[0]\n            return (\n                0.5 * (np.log(2 * np.pi) + np.log(sigma**2) + ((x - mu) ** 2) / (sigma**2))\n            ).mean()\n\n        optimal_sigma = np.zeros(spectrum.shape[0])\n\n        for i in tqdm(range(spectrum.shape[0])):\n            res = minimize(\n                cost_function, [0.00060], args=(labels[i], spectrum[i]), bounds=[(1e-6, 0.01)]\n            )\n            optimal_sigma[i] = res.x[0]\n        \n        return optimal_sigma\n\n    def train(self, data: np.ndarray, spectrum: np.ndarray, labels: np.ndarray):\n        X = self._extract_features(data, spectrum)\n        y = self._get_optimal_sigma(spectrum, labels)\n\n        kf = KFold(n_splits=5, shuffle=True)\n        cv_preds = np.zeros_like(y)\n        self.models = []\n        for i, (train_index, val_index) in enumerate(kf.split(X)):\n            X_train, X_val = X[train_index], X[val_index]\n            y_train, y_val = y[train_index], y[val_index]\n\n            model = self._create_model()\n            model.fit(X_train, y_train)\n            self.models.append(model)\n            cv_preds[val_index] = model.predict(X_val)\n\n            print(\n                f\"Fold {i + 1}, Validation RMSE: {np.sqrt(np.mean((y_val - cv_preds[val_index]) ** 2))}\"\n            )\n\n        sigma = np.repeat(cv_preds[:, np.newaxis], 283, axis=1)\n        sigma[:, 0] = sigma[:, 0] / max(np.mean(sigma[:, 0]), self.eps) * self.mean_fgs_sigma\n        sigma[:, 0] = sigma[:, 0].clip(min=self.fgs_min_sigma)\n        sigma[:, 1:] = sigma[:, 1:] / max(np.mean(sigma[:, 1:]), self.eps) * self.mean_airs_sigma\n        sigma[:, 1:] = sigma[:, 1:].clip(min=self.airs_min_sigma)\n        return sigma\n\n    def get_sigma(self, data: np.ndarray, spectrum: np.ndarray):\n        assert data.ndim == 3 and data.shape[2] == 283, (\n            \"Data should be of shape (num_planets, num_time_steps, 283)\"\n        )\n        assert len(self.models) > 0, \"Models are not trained. Please call the train method first.\"\n\n        X = self._extract_features(data, spectrum)\n        preds = np.zeros((X.shape[0], len(self.models)))\n        for i, model in enumerate(self.models):\n            preds[:, i] = model.predict(X)\n        sigma_pred = preds.mean(axis=1)\n\n        sigma = np.repeat(sigma_pred[:, np.newaxis], 283, axis=1)\n        sigma[:, 0] = sigma[:, 0] / max(np.mean(sigma[:, 0]), self.eps) * self.mean_fgs_sigma\n        sigma[:, 0] = sigma[:, 0].clip(min=self.fgs_min_sigma)\n        sigma[:, 1:] = sigma[:, 1:] / max(np.mean(sigma[:, 1:]), self.eps) * self.mean_airs_sigma\n        sigma[:, 1:] = sigma[:, 1:].clip(min=self.airs_min_sigma)\n        return sigma\n\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T05:29:26.581461Z","iopub.execute_input":"2026-03-25T05:29:26.582133Z","iopub.status.idle":"2026-03-25T05:29:26.632664Z","shell.execute_reply.started":"2026-03-25T05:29:26.582101Z","shell.execute_reply":"2026-03-25T05:29:26.631608Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Evaluation Metric","metadata":{}},{"cell_type":"code","source":"import scipy.stats\n\ndef compute_score(\n    train_labels: np.ndarray,      # (num_planets, 283) — ground truth\n    predicted_means: np.ndarray,   # (num_planets, 283) — your μ\n    predicted_sigmas: np.ndarray,  # (num_planets, 283) — your σ\n    fgs_weight: float = 57.846,\n    fsg_sigma_true: float = 1e-6,\n    airs_sigma_true: float = 1e-5,\n) -> float:\n\n    n_wavelengths = train_labels.shape[1]  # 283\n\n    # Naive baseline: predict training mean and std for everything\n    naive_mean  = train_labels.mean()\n    naive_sigma = train_labels.std()\n\n    y_true    = train_labels\n    y_pred    = predicted_means\n    sigma_pred = np.clip(predicted_sigmas, a_min=1e-15, a_max=None)\n\n    # Ideal sigma: 1e-6 for FGS, 1e-5 for AIRS\n    sigma_true = np.concatenate([\n        np.array([fsg_sigma_true]),\n        np.ones(n_wavelengths - 1) * airs_sigma_true\n    ])  # shape (283,)\n\n    # GLL for your submission\n    GLL_pred = scipy.stats.norm.logpdf(y_true, loc=y_pred,  scale=sigma_pred)\n\n    print(f\"Raw GLL:    {GLL_pred.mean():.4f}\")\n\n\n    # GLL for ideal submission (perfect μ, ideal σ)\n    GLL_true = scipy.stats.norm.logpdf(\n        y_true, loc=y_true,\n        scale=sigma_true * np.ones_like(y_true)\n    )\n\n    # GLL for naive baseline (training mean/std for everything)\n    GLL_mean = scipy.stats.norm.logpdf(\n        y_true,\n        loc=naive_mean  * np.ones_like(y_true),\n        scale=naive_sigma * np.ones_like(y_true)\n    )\n\n    # Per-sample, per-wavelength normalized score\n    ind_scores = (GLL_pred - GLL_mean) / (GLL_true - GLL_mean)\n\n    # Weights: FGS gets 57.846, all AIRS get 1.0\n    weights = np.ones(n_wavelengths)\n    weights[0] = fgs_weight\n    weights = weights * np.ones_like(ind_scores)  # broadcast to (planets, 283)\n\n    final_score = float(np.clip(np.average(ind_scores, weights=weights), 0.0, 1.0))\n    return final_score","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T05:29:26.633304Z","iopub.status.idle":"2026-03-25T05:29:26.633539Z","shell.execute_reply.started":"2026-03-25T05:29:26.633427Z","shell.execute_reply":"2026-03-25T05:29:26.633441Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_loader = DataLoaderAndCalibrator(\n    data_path=INPUT_DATA_FOLDER,\n    output_path=CALIBRATED_DATA_FOLDER,\n    force_recalibration=False,   # set True if you want to regenerate\n    cut_airs_channels=True,\n    binning=BINNING,\n    n_jobs=2,                    # Kaggle usually has 2 CPU cores available\n)\n\ntrain_data, train_labels = data_loader.load_all_train_data()\nprint(f\"Train data shape:   {train_data.shape}\")    # (num_planets, ~187, 283)\nprint(f\"Train labels shape: {train_labels.shape}\")  # (num_planets, 283)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T05:29:26.634898Z","iopub.status.idle":"2026-03-25T05:29:26.635218Z","shell.execute_reply.started":"2026-03-25T05:29:26.635042Z","shell.execute_reply":"2026-03-25T05:29:26.635064Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"features_extractor = WavelengthsGroupsMultiplierFinder()\n\ntrain_features = features_extractor.extract_features(\n    train_data,\n    average_cross_groups=False,\n    wavelengths_groups=[1, 2, 4, 8, 16, 32, 64],\n    weights=[1, 1, 1, 1, 1, 1, 1],\n)\n# shape: (num_planets, 283, 7)\n\nmodel_data = train_features.transpose(0, 2, 1)\n# shape: (num_planets, 7, 283) — ready for CNN\n\ntrain_star_info = pd.read_csv(INPUT_DATA_FOLDER / \"train_star_info.csv\")[[\"Rs\", \"i\"]].values\nprint(f\"Model data shape:     {model_data.shape}\")\nprint(f\"Star info shape:      {train_star_info.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T05:29:26.636317Z","iopub.status.idle":"2026-03-25T05:29:26.636728Z","shell.execute_reply.started":"2026-03-25T05:29:26.636465Z","shell.execute_reply":"2026-03-25T05:29:26.63649Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trainer = SValuesCNNWithStarInfoTrainer(\n    weights_dir=MODEL_FOLDER,\n    n_splits=5,\n    num_channels=model_data.shape[1],    # 7\n    wavelengths=model_data.shape[2],     # 283\n    num_star_features=train_star_info.shape[1],  # 2\n)\n\n_, _, val_preds = trainer.train(\n    model_data * 1e3,\n    train_star_info,\n    train_labels * 1e3,\n    epochs=EPOCHS,\n    force_training=False,   # set True to retrain from scratch\n)\n\ntrain_spectrum = val_preds / 1e3\nprint(f\"Train spectrum shape: {train_spectrum.shape}\")  # (num_planets, 283)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T05:29:26.637603Z","iopub.status.idle":"2026-03-25T05:29:26.637848Z","shell.execute_reply.started":"2026-03-25T05:29:26.637736Z","shell.execute_reply":"2026-03-25T05:29:26.637749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sigma_calculator = LROnVariousFeaturesSigmaCalculator(\n    mean_fgs_sigma=MEAN_FGS_SIGMA,\n    mean_airs_sigma=MEAN_AIRS_SIGMA,\n    fgs_min_sigma=1e-6,\n    airs_min_sigma=1e-6,\n)\n\nval_sigma = sigma_calculator.train(\n    data=train_data,\n    spectrum=train_spectrum,\n    labels=train_labels,\n)\nprint(f\"Val sigma shape: {val_sigma.shape}\")  # (num_planets, 283)\n\n# After sigma training\nval_submission = np.concatenate([train_spectrum, val_sigma], axis=1)\n\n\n# Real competition score [0, 1]\nreal_score = compute_score(\n    train_labels=train_labels,\n    predicted_means=train_spectrum,\n    predicted_sigmas=val_sigma,\n)\nprint(f\"Real Score: {real_score:.4f}\")   # this is what the leaderboard shows","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_data = data_loader.load_all_test_data()\nprint(f\"Test data shape: {test_data.shape}\")\n\ntest_features = features_extractor.extract_features(\n    test_data,\n    average_cross_groups=False,\n    wavelengths_groups=[1, 2, 4, 8, 16, 32, 64],\n    weights=[1, 1, 1, 1, 1, 1, 1],\n)\n\ntest_model_data = test_features.transpose(0, 2, 1)\ntest_star_info = pd.read_csv(INPUT_DATA_FOLDER / \"test_star_info.csv\")[[\"Rs\", \"i\"]].values\n\ntest_spectrum = trainer.predict(test_model_data * 1e3, test_star_info) / 1e3\npredicted_sigma = sigma_calculator.get_sigma(data=test_data, spectrum=test_spectrum)\n\nprint(f\"Test spectrum shape: {test_spectrum.shape}\")\nprint(f\"Test sigma shape:    {predicted_sigma.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"result = np.concatenate([test_spectrum, predicted_sigma], axis=1)\nassert result.shape == (test_data.shape[0], 283 * 2), \\\n    f\"Unexpected shape: {result.shape}\"\n\nsubmission_df = pd.read_csv(INPUT_DATA_FOLDER / \"sample_submission.csv\")\nsubmission_df.iloc[:, 1:] = result\nsubmission_df.to_csv(SUBMISSION_FILE, index=False)\nprint(f\"Submission saved → {SUBMISSION_FILE}\")\nprint(submission_df.head())","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}