{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"sourceType":"competition"}],"dockerImageVersionId":31041,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport pickle\nimport os\n\nfrom tqdm import tqdm\nfrom sklearn.model_selection import cross_val_predict\nfrom sklearn.linear_model import Ridge, Lasso\nfrom sklearn.metrics import r2_score, mean_squared_error","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T18:50:30.449062Z","iopub.execute_input":"2025-10-01T18:50:30.449412Z","iopub.status.idle":"2025-10-01T18:50:34.758628Z","shell.execute_reply.started":"2025-10-01T18:50:30.449374Z","shell.execute_reply":"2025-10-01T18:50:34.757568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"VERSION = \"v2\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T18:51:08.481968Z","iopub.execute_input":"2025-10-01T18:51:08.482373Z","iopub.status.idle":"2025-10-01T18:51:08.487082Z","shell.execute_reply.started":"2025-10-01T18:51:08.482344Z","shell.execute_reply":"2025-10-01T18:51:08.486171Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile preprocess.py\n\nimport pandas as pd\nimport numpy as np\nimport multiprocessing as mp\nimport itertools\nimport os\nimport subprocess\nfrom astropy.stats import sigma_clip\nfrom numpy.polynomial import Polynomial\nfrom tqdm import tqdm\nimport torch\nimport torch.nn.functional as F\n\n\nROOT = \"/kaggle/input/ariel-data-challenge-2025/\"\nVERSION = \"v2\"\nA_BINNING = 15\nF_BINNING = 12*15\n\n\ndevice = (\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\nMODE = os.getenv('PREPROCESS_MODE')\n\n\nsensor_sizes_dict = {\n    \"AIRS-CH0\": [[11250, 32, 356], [32, 356]],\n    \"FGS1\": [[135000, 32, 32], [32, 32]],\n}  # input, mask\n\ncl = 8\ncr = 24\n\n\ndef get_gain_offset():\n    \"\"\"\n    Get the gain and offset for a given planet and sensor\n\n    Unlike last year's challenge, all planets use the same adc_info.\n    We can just hard code it.\n    \"\"\"\n    gain = 0.4369\n    offset = -1000.0\n    return gain, offset\n\n\n\ndef read_data(planet_id, sensor, mode):\n    \"\"\"\n    Read the data for a given planet and sensor\n    \"\"\"\n    # get all noise correction frames and signal\n    signal = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_signal_0.parquet\",\n        engine=\"pyarrow\",\n    )\n    dark_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration_0/dark.parquet\",\n        engine=\"pyarrow\",\n    )\n    dead_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration_0/dead.parquet\",\n        engine=\"pyarrow\",\n    )\n    linear_corr_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration_0/linear_corr.parquet\",\n        engine=\"pyarrow\",\n    )\n    flat_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration_0/flat.parquet\",\n        engine=\"pyarrow\",\n    )\n\n    # reshape to sensor shape and cast to float64\n    signal = signal.values.astype(np.float64).reshape(sensor_sizes_dict[sensor][0])[\n        :, cl:cr, :\n    ]\n    dark_frame = dark_frame.values.astype(np.float64).reshape(\n        sensor_sizes_dict[sensor][1]\n    )[cl:cr, :]\n    dead_frame = dead_frame.values.reshape(sensor_sizes_dict[sensor][1])[cl:cr, :]\n    flat_frame = flat_frame.values.astype(np.float64).reshape(\n        sensor_sizes_dict[sensor][1]\n    )[cl:cr, :]\n\n    linear_corr = linear_corr_frame.values.astype(np.float64).reshape(\n        [6] + sensor_sizes_dict[sensor][1]\n    )[:, cl:cr, :]\n\n    return (\n        signal,\n        dark_frame,\n        dead_frame,\n        linear_corr,\n        flat_frame,\n    )\n\n\ndef ADC_convert(signal, gain, offset):\n    \"\"\"\n    Step 1: Analog-to-Digital Conversion (ADC) correction\n\n    The Analog-to-Digital Conversion (adc) is performed by the detector to convert the\n    pixel voltage into an integer number. We revert this operation by using the gain\n    and offset for the calibration files 'train_adc_info.csv'.\n    \"\"\"\n\n    signal /= gain\n    signal += offset\n    return signal\n\n\ndef mask_hot_dead(signal, dead, dark):\n    \"\"\"\n    Step 2: Mask hot/dead pixel\n\n    The dead pixels map is a map of the pixels that do not respond to light and, thus,\n    can't be accounted for any calculation. In all these frames the dead pixels are\n    masked using python masked arrays. The bad pixels are thus masked but left\n    uncorrected. Some methods can be used to correct bad-pixels but this task,\n    if needed, is left to the participants.\n    \"\"\"\n\n    hot = sigma_clip(dark, sigma=5, maxiters=5).mask\n    hot = np.tile(hot, (signal.shape[0], 1, 1))\n    dead = np.tile(dead, (signal.shape[0], 1, 1))\n\n    signal[dead] = np.nan\n    signal[hot] = np.nan\n    return signal\n\n\ndef apply_linear_corr(c, signal):\n    \"\"\"\n    Step 3: linearity Correction\n\n    The non-linearity of the pixels' response can be explained as capacitive leakage\n    on the readout electronics of each pixel during the integration time. The number\n    of electrons in the well is proportional to the number of photons that hit the\n    pixel, with a quantum efficiency coefficient. However, the response of the pixel\n    is not linear with the number of electrons in the well. This effect can be\n    described by a polynomial function of the number of electrons actually in the well.\n    The data is provided with calibration files linear_corr.parquet that are the\n    coefficients of the inverse polynomial function and can be used to correct this\n    non-linearity effect.\n    Using horner's method to evaluate the polynomial\n    \"\"\"\n    assert c.shape[0] == 6  # Ensure the polynomial is of degree 5\n\n    return (\n        (((c[5] * signal + c[4]) * signal + c[3]) * signal + c[2]) * signal + c[1]\n    ) * signal + c[0]\n\n\ndef clean_dark(signal, dark, dt):\n    \"\"\"\n    Step 4: dark current subtraction\n\n    The data provided include calibration for dark current estimation, which can be\n    used to pre-process the observations. Dark current represents a constant signal\n    that accumulates in each pixel during the integration time, independent of the\n    incoming light. To obtain the corrected image, the following conventional approach\n    is applied: The data provided include calibration files such as dark frames or\n    dead pixels' maps. They can be used to pre-process the observations. The dark frame\n    is a map of the detector response to a very short exposure time, to correct for the\n    dark current of the detector.\n\n    image - (dark * dt)\n\n    The corrected image is conventionally obtained via the following: where the dark\n    current map is first corrected for the dead pixel.\n    \"\"\"\n\n    dark = torch.tile(dark, (signal.shape[0], 1, 1))\n    signal -= dark * dt[:, None, None]\n    return signal\n\n\ndef get_cds(signal):\n    \"\"\"\n    Step 5: Get Correlated Double Sampling (CDS)\n\n    The science frames are alternating between the start of the exposure and the end of\n    the exposure. The lecture scheme is a ramp with a double sampling, called\n    Correlated Double Sampling (CDS), the detector is read twice, once at the start\n    of the exposure and once at the end of the exposure. The final CDS is the\n    difference (End of exposure) - (Start of exposure).\n    \"\"\"\n\n    return torch.subtract(signal[1::2, :, :], signal[::2, :, :])\n\n\ndef correct_flat_field(flat, signal):\n    \"\"\"\n    Step 6: Flat Field Correction\n\n    The flat field is a map of the detector response to uniform illumination, to\n    correct for the pixel-to-pixel variations of the detector, for example the\n    different quantum efficiencies of each pixel.\n    \"\"\"\n\n    return signal / flat\n\n\ndef bin_obs(signal, binning):\n    \"\"\"\n    Step 5.1: Bin Observations\n\n    The data provided are binned in the time dimension. The binning is performed by\n    summing the signal over the time dimension.\n    \"\"\"\n\n    cds_binned = torch.zeros(\n        (\n            signal.shape[0] // binning,\n            signal.shape[1],\n            signal.shape[2],\n        ),\n        device=device,\n    )\n    for i in range(signal.shape[0] // binning):\n        cds_binned[i, :, :] = torch.sum(\n            signal[i * binning : (i + 1) * binning, :, :], axis=0\n        )\n    return cds_binned\n\n\ndef nan_interpolation(tensor):\n    # Assume tensor is of shape (batch, height, width)\n    nan_mask = torch.isnan(tensor)\n\n    # Replace NaNs with zero temporarily\n    tensor_filled = torch.where(\n        nan_mask, torch.tensor(0.0, device=tensor.device), tensor\n    )\n\n    # Create a binary mask (0 where NaNs were and 1 elsewhere)\n    ones = torch.ones_like(tensor, device=tensor.device)\n    weight = torch.where(nan_mask, torch.tensor(0.0, device=tensor.device), ones)\n\n    # Perform interpolation by convolving with a kernel\n    # using bilinear interpolation\n    kernel = torch.ones(1, 1, 1, 3, device=tensor.device, dtype=tensor.dtype)\n\n    # Apply padding to the tensor and weight to prevent boundary issues\n    tensor_padded = F.pad(\n        tensor_filled.unsqueeze(1), (1, 1, 0, 0), mode=\"replicate\"\n    ).squeeze(1)\n    weight_padded = F.pad(weight.unsqueeze(1), (1, 1, 0, 0), mode=\"replicate\").squeeze(\n        1\n    )\n\n    # Convolve the filled tensor and the weight mask\n    tensor_conv = F.conv2d(tensor_padded.unsqueeze(1), kernel, stride=1)\n    weight_conv = F.conv2d(weight_padded.unsqueeze(1), kernel, stride=1)\n\n    # Compute interpolated values (normalized by weights)\n    interpolated_tensor = tensor_conv / weight_conv\n\n    # Apply the interpolated values only to the positions of NaNs\n    result = torch.where(nan_mask, interpolated_tensor.squeeze(1), tensor)\n\n    return result\n\n\n\ndef process_planet(planet_id):\n    \"\"\"\n    Process a single planet's data\n    \"\"\"\n    axis_info = pd.read_parquet(ROOT + \"/axis_info.parquet\")\n    dt_airs = axis_info[\"AIRS-CH0-integration_time\"].dropna().values\n\n    for sensor in [\"AIRS-CH0\", \"FGS1\"]:\n        # load all data for this planet and sensor\n        signal, dark_frame, dead_frame, linear_corr, flat_frame = read_data(\n            planet_id, sensor, mode=MODE\n        )\n        gain, offset = get_gain_offset()\n\n        # Step 1: ADC correction\n        signal = ADC_convert(signal, gain, offset)\n\n        # Step 2: Mask hot/dead pixel\n        signal = mask_hot_dead(signal, dead_frame, dark_frame)\n\n        # clip at 0\n        signal = torch.tensor(signal.clip(0)).to(device)    \n        signal = apply_linear_corr(\n            torch.tensor(linear_corr).to(device), signal.clone().detach()\n        )\n\n        if sensor == \"FGS1\":\n            dt = torch.ones(len(signal), device=device) * 0.1\n        elif sensor == \"AIRS-CH0\":\n            dt = torch.tensor(dt_airs).to(device)\n\n        dt[1::2] += 0.1\n\n        signal = clean_dark(signal, torch.tensor(dark_frame).to(device), dt)\n\n        # Step 5: Get Correlated Double Sampling (CDS)\n        signal = get_cds(signal)\n\n        # Step 6: Flat Field Correction\n\n        if sensor == \"AIRS-CH0\":\n            signal = bin_obs(signal, binning=A_BINNING)\n        else:\n            signal = bin_obs(signal, binning=F_BINNING)\n\n        signal = correct_flat_field(torch.tensor(flat_frame).to(device), signal)\n\n        # Step 7: Interpolate NaNs (twice!)\n        signal = nan_interpolation(signal)\n        signal = nan_interpolation(signal)\n\n        if sensor == \"FGS1\":\n            signal = torch.nanmean(signal, axis=[1, 2]).cpu().numpy()\n        elif sensor == \"AIRS-CH0\":\n            signal = torch.nanmean(signal, axis=1).cpu().numpy()\n            \n        # save the processed signal\n        np.save(\n            str(planet_id) + \"_\" + sensor + f\"_signal_{VERSION}.npy\",\n            signal.astype(np.float64),\n        )\n\n\nif __name__ == \"__main__\":\n    star_info = pd.read_csv(ROOT + f\"/{MODE}_star_info.csv\", index_col=\"planet_id\")\n    planet_ids = [int(x) for x in star_info.index.tolist()]\n\n    # Use up to 4 threads!\n    mp.set_start_method('spawn')\n    with mp.Pool(processes=4) as pool:\n        list(tqdm(pool.imap(process_planet, planet_ids), total=len(planet_ids)))\n\n    \n    signal_train = []\n\n    for planet_id in planet_ids:\n        f_raw = np.load(f\"{planet_id}_FGS1_signal_{VERSION}.npy\")\n        a_raw = np.load(f\"{planet_id}_AIRS-CH0_signal_{VERSION}.npy\")\n\n        # flip a_raw\n        signal = np.concatenate([f_raw[:, None], a_raw[:, ::-1]], axis=1)\n        signal_train.append(signal)\n\n        os.remove(\"/kaggle/working/\" + str(planet_id) + f\"_AIRS-CH0_signal_{VERSION}.npy\")\n        os.remove(\"/kaggle/working/\" + str(planet_id) + f\"_FGS1_signal_{VERSION}.npy\")\n\n    signal_train = np.array(signal_train)\n    np.save(f\"signal_{VERSION}.npy\", signal_train, allow_pickle=False)\n\n    print(\"Processing complete!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T18:51:16.934353Z","iopub.execute_input":"2025-10-01T18:51:16.934733Z","iopub.status.idle":"2025-10-01T18:51:16.947508Z","shell.execute_reply.started":"2025-10-01T18:51:16.934704Z","shell.execute_reply":"2025-10-01T18:51:16.946339Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.environ[\"PREPROCESS_MODE\"] = \"train\"\n\n!python preprocess.py\n\ndata_train = np.load(f\"signal_{VERSION}.npy\")\ndata_train.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T18:51:20.362311Z","iopub.execute_input":"2025-10-01T18:51:20.362671Z","iopub.status.idle":"2025-10-01T18:51:22.243773Z","shell.execute_reply.started":"2025-10-01T18:51:20.362645Z","shell.execute_reply":"2025-10-01T18:51:22.242927Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_adc_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/adc_info.csv')\ntrain_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train_star_info.csv')\n\ntrain_labels = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train.csv',\n                           index_col='planet_id')\nwavelengths = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/wavelengths.csv')\naxis_info = pd.read_parquet('/kaggle/input/ariel-data-challenge-2025/axis_info.parquet')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T18:51:25.653001Z","iopub.execute_input":"2025-10-01T18:51:25.653507Z","iopub.status.idle":"2025-10-01T18:51:26.354632Z","shell.execute_reply.started":"2025-10-01T18:51:25.653471Z","shell.execute_reply":"2025-10-01T18:51:26.353563Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from scipy.signal import savgol_filter\n\n\nMODEL_VERSION = \"v1\"\nPRE_BINNED_TIME = 15\n\nROOT = \"/kaggle/input/ariel-data-challenge-2025/\"\n\n\n# find transit zones\ndef phase_detector(signal_orig, smooth_window=11):\n    signal = signal_orig.reshape(-1, 1).mean(-1)\n    signal = savgol_filter(signal, smooth_window, 2)  # smooth\n    first_derivative = np.gradient(signal)\n    second_derivative = savgol_filter(np.gradient(savgol_filter(first_derivative, 41, 2)), 41, 2)\n\n    local_min = (np.diff(np.sign(np.diff(second_derivative))) > 0).nonzero()[0] + 1\n    local_max = (np.diff(np.sign(np.diff(second_derivative))) < 0).nonzero()[0] + 1\n    \n    if len(local_min) >= 2:\n        top2_min_indices = local_min[np.argsort(second_derivative[local_min])[:2]]\n    else:\n        top2_min_indices = local_min\n    \n    if len(local_max) >= 2:\n        top2_max_indices = local_max[np.argsort(second_derivative[local_max])[-2:]]\n    else:\n        top2_max_indices = local_max\n    top2_min_indices.sort()\n    top2_max_indices.sort()\n\n    # 4 extrema of the 2nd derivative and 2 of the 1st\n    phase1 = top2_min_indices[0]\n    phase2 = top2_max_indices[0]\n    phase3 = top2_max_indices[1]\n    phase4 = top2_min_indices[1]\n    phase5 = np.argmin(first_derivative)\n    phase6 = np.argmax(first_derivative)\n\n    return phase1, phase2, phase3, phase4, phase5, phase6\n\n\ndef get_breakpoints(x, smooth=19):\n    bp1 = np.zeros(x.shape[0], dtype=np.int32)\n    bp2 = np.zeros(x.shape[0], dtype=np.int32)\n    bp3 = np.zeros(x.shape[0], dtype=np.int32)\n    bp4 = np.zeros(x.shape[0], dtype=np.int32)\n    bp5 = np.zeros(x.shape[0], dtype=np.int32)\n    bp6 = np.zeros(x.shape[0], dtype=np.int32)\n    for i in range(x.shape[0]):\n        signal = x[i]\n        p1, p2, p3, p4, p5, p6 = phase_detector(\n            signal, smooth_window=smooth\n        )\n        bp1[i] = p1\n        bp2[i] = p2\n        bp3[i] = p3\n        bp4[i] = p4\n        bp5[i] = p5\n        bp6[i] = p6\n        \n    return [bp1, bp2, bp3, bp4, bp5, bp6]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T18:51:29.57994Z","iopub.execute_input":"2025-10-01T18:51:29.580508Z","iopub.status.idle":"2025-10-01T18:51:29.831142Z","shell.execute_reply.started":"2025-10-01T18:51:29.580469Z","shell.execute_reply":"2025-10-01T18:51:29.829914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile feature_engineering.py\nfrom scipy.signal import savgol_filter\nimport warnings\nfrom scipy.stats import kurtosis, skew\nfrom scipy.signal import medfilt\nfrom scipy.optimize import curve_fit\nfrom scipy.interpolate import interp1d\nfrom numpy.polynomial import Polynomial\nfrom scipy.optimize import least_squares, minimize\nfrom sklearn.metrics import mean_squared_error\nfrom scipy.ndimage import median_filter\n\n\nwarnings.simplefilter(\"ignore\")\n\nA_BINNING = 15\n\n# threshold values for identifying outliers\nbad_low = 20\nbad_up = 354\n\nbuf = 15  # for lower1, upper1\nbuf1 = 10  # for lower, mid1, mid2, upper\n\n\ndef calculate_weights(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    \"\"\"\n    Calculating weights for averaging based on SNR\n    \"\"\"\n    max_len = a.shape[0] - 1\n    if is_outlier:\n        return np.ones_like(a[0]) / a.shape[-1]\n    else:\n        y_combined = np.concatenate([a[:max(lower - buf1, 1), :], a[min(upper + buf1, max_len):, :]], axis=0)\n        ratio = y_combined.mean(0) / y_combined.std(0)\n        return ratio / ratio.sum()\n\n\ndef calc_for_outliers(a, lower, upper):\n    \"\"\"\n    Estimating transit depth for outlier cases\n    \"\"\"\n    max_len = len(a) - 1\n            \n    if lower + buf < upper - buf:\n        obs = a[lower + buf : upper - buf].mean()\n    else:\n        obs = a[lower : upper].mean()\n\n    if lower - buf >= 10 and max_len - upper - buf >= 10:\n        unobs = (np.median(a[:max(lower - buf, 1)]) + np.median(a[min(upper + buf, max_len):])) / 2\n    elif lower >= max_len - upper:\n        unobs = np.median(a[:max(lower - buf, 1)])\n    else:\n        unobs = np.median(a[min(upper + buf, max_len):])\n\n    arr1 = 1 - (obs / unobs)\n    arr2 = 1 - a[(lower + upper) // 2] / unobs\n\n    return arr1, arr2\n\n\ndef calc_err(x_combined, y_combined, degree):\n    \"\"\"\n    Calculating the error to obtain the optimal polynomial degree\n    \"\"\"\n    max_len = 374 # hardcoded for BINNING=15\n    \n    poly_guess = np.polyfit(x_combined, y_combined, degree)\n    inter = np.polyval(poly_guess, np.arange(max_len + 1))\n    err = mean_squared_error(y_combined, inter[x_combined], squared=False)\n\n    # penalizing rmse, high polynomial degree and small number of points in curve fitting.\n    return err * degree**(1 - len(x_combined) / max_len)\n\n    \ndef calc_depth_and_detrend(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False, unstable=False, fixed_degree=None):\n    \"\"\"\n    Main function for transit depth estimation and detrending\n    \n    Parameters:\n        - a: 1d numpy array of observation points\n        - lower, mid1, mid2, upper, lower1, upper1: transit boundary points\n        - is_outlier: boolean flag indicating if the data point is an outlier\n        - unstable: if True, don't use curve fitting\n        - fixed_degree: if not None, use provided degree; otherwise, find optimal degree\n    \n    Returns:\n        - (arr1, arr2): tuple containing the averaged transit depth and the transit depth at mid-transit\n    \"\"\"\n    max_len = len(a) - 1\n    degree = 3\n    \n    if is_outlier:\n        a /= a.mean()\n        arr1, arr2 = calc_for_outliers(a, lower1, upper1)\n    else:\n        a /= a.mean()\n\n        # region outside the transit\n        x_combined = np.concatenate([np.arange(max(lower - buf1, 1)), np.arange(min(upper + buf1, max_len), max_len + 1)], axis=-1)\n        y_combined = a[x_combined]\n            \n        if fixed_degree is None: # find the optimal degre\n            best_val = 10**100\n            for j in [1, 2, 3, 4, 5]:\n                if calc_err(x_combined, y_combined, j) < best_val:\n                    best_val = calc_err(x_combined, y_combined, j)\n                    degree = j\n        else:\n            degree = fixed_degree\n\n        obs = a[mid1 : mid2]\n        \n        if unstable: # no curve fitting\n            unobs = y_combined.mean()\n        else:\n            poly_guess = np.polyfit(x_combined, y_combined, degree)         \n            inter = np.polyval(poly_guess, np.arange(max_len + 1))\n                \n            a /= inter\n            inter /= inter\n            unobs = inter[mid1 : mid2]\n\n        arr1 = 1 - np.mean(obs / unobs)\n        arr2 = 1 - a[(lower1 + upper1) // 2] / unobs.mean()\n                \n\n    if np.isnan(arr1):\n        arr1 = 0\n    if np.isnan(arr2):\n        arr2 = 0\n        \n    return arr1, arr2\n\n\ndef calc_slope(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    \"\"\"\n    Calculate transit wall steepness as slope between contact points\n    \"\"\"\n    max_len = len(a) - 1\n    \n    if not (lower < mid1 < mid2 < upper) or is_outlier:\n        return 0\n    else:\n        return ((a[mid1] - a[lower]) / (mid1 - lower) - (a[upper] - a[mid2]) / (upper - mid2)) / 2\n\n\ndef calc_slope_2(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    \"\"\"\n    Calculate transit bottom curvature as slope between mid-transit \n    and contact point\n    \"\"\"\n    max_len = len(a) - 1\n    \n    if not (lower < mid1 < mid2 < upper) or is_outlier:\n        return 0\n    else:\n        mid_ind = (mid1 + mid2) // 2\n        if not (mid1 < mid_ind < mid2):\n            return 0\n        return ((a[mid_ind] - a[mid1]) / (mid_ind - mid1) - (a[mid2] - a[mid_ind]) / (mid2 - mid_ind)) / 2\n\n\n# calcluating slopes for unsimmetric cases\ndef calc_slope_2_left(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    max_len = len(a) - 1\n    \n    if not (lower < mid1 < mid2 < upper):\n        return 0\n    else:\n        mid_ind = (mid1 + mid2) // 2\n        if not (mid1 < mid_ind < mid2):\n            return 0\n        return (a[mid_ind] - a[mid1]) / (mid_ind - mid1)\n\n\ndef calc_slope_2_right(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    max_len = len(a) - 1\n    \n    if not (lower < mid1 < mid2 < upper):\n        return 0\n    else:\n        mid_ind = (mid1 + mid2) // 2\n        if not (mid1 < mid_ind < mid2):\n            return 0\n        return (a[mid2] - a[mid_ind]) / (mid2 - mid_ind)\n\n\n# gradient slopes\ndef calc_curv_left(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    if not (lower < mid1 < mid2 < upper) or is_outlier:\n        return 0\n    else:\n        mid_ind = (mid1 + mid2) // 2\n        a = savgol_filter(np.gradient(a), 41, 1)\n        return (a[mid_ind] - a[mid1]) / (mid_ind - mid1)\n\n\ndef calc_curv_right(a, lower, mid1, mid2, upper, lower1, upper1, is_outlier=False):\n    if not (lower < mid1 < mid2 < upper) or is_outlier:\n        return 0\n    else:\n        mid_ind = (mid1 + mid2) // 2\n        a = savgol_filter(np.gradient(a), 41, 1)\n        return (a[mid2] - a[mid_ind]) / (mid2 - mid_ind)\n\n\ndef calc_perc(a, lower, upper, q, is_outlier=False):\n    \"\"\"\n    Calculate the percentile of the transit depth, assumes the input signal is already detrended\n    \"\"\"\n    if is_outlier:\n        return 0\n    return np.quantile(1 - a[lower : upper], q)\n    \n\ndef feature_engineering(star_info, data):\n    \"\"\"\n    Prepares features for training or inference\n    \n    Parameters:\n        - star_info: star metadata DataFrame\n        - data: 3d data array (samples, time, frequencies)\n    \n    Returns:\n        tuple of DataFrame and outliers mask\n    \"\"\"\n    df = pd.DataFrame()\n\n    cut_inf, cut_sup = 36, 318\n    \n    signal = np.concatenate(\n        [data[:, :, 0][:, :, None], data[:, :, cut_inf:cut_sup]], axis=2\n    )\n    max_len = signal.shape[1] - 1\n        \n    lower, mid1, mid2, upper, lower1, upper1 = get_breakpoints(signal[:, :, 1:].mean(-1))\n    boundaries = (lower, mid1, mid2, upper, lower1, upper1) \n\n    # identifying outliers\n    outliers = (np.array(lower1) < bad_low) | (np.array(upper1) > bad_up) | (np.array(lower) < bad_low) | (np.array(upper) > bad_up)\n    for i in range(signal.shape[0]):\n        if not (lower[i] < mid1[i] < mid2[i] < upper[i]):\n            outliers[i] = 1\n    \n    signal_mean_raw = np.zeros(signal.shape[:2])\n    for i in tqdm(range(signal_mean_raw.shape[0])): # weighted averaging along frequency dimension\n        weights = calculate_weights(signal[i, :, 1:], *(b[i] for b in boundaries), outliers[i])\n        signal_mean_raw[i, :] = (signal[i, :, 1:] @ weights.T).T\n\n    signal_mean = savgol_filter(signal_mean_raw, 11, 1)\n\n    # frequency set for precise depth estimation via curve fitting (less robust)\n    good_waves = [1, 6, 11, 16, 21, 26, 31, 36, 41, 51, 61, 71, 76, 81, 86, 91, 96, 101, 106, 111, 121, 131, 141, 151, 161, 171, 196, 201, 206]\n    for i in tqdm(range(len(signal_mean))):\n\n        # filter very bad cases :) (exclude from training, use larger sigma for prediction)\n        df.loc[i, 'very_bad'] = (lower1[i] < bad_low // 2 or upper1[i] > max_len - (max_len - bad_up) // 2)\n        if os.environ[\"PREPROCESS_MODE\"] == 'train' and star_info.loc[i, 'planet_id'] in [2486733311, 2554492145]:\n            df.loc[i, 'very_bad'] = True\n\n\n        # averaged and mid-transit depth estimation\n        fake_avg, fake_mid = calc_depth_and_detrend(signal_mean[i].copy(), *(b[i] for b in boundaries), outliers[i], unstable=True)\n        fake_avg_2, fake_mid_2 = calc_depth_and_detrend(signal_mean[i].copy(), *(b[i] for b in boundaries), outliers[i], fixed_degree=3)\n        df.loc[i, 'average_depth'], df.loc[i, 'mid_depth'] = calc_depth_and_detrend(signal_mean[i], *(b[i] for b in boundaries), outliers[i])\n\n        norm_coef = (1 - df.loc[i, 'average_depth']) / (1 - fake_avg)\n        norm_coef_mid = (1 - df.loc[i, 'mid_depth']) / (1 - fake_mid)\n        norm_coef_2 = (1 - df.loc[i, 'average_depth']) / (1 - fake_avg_2)\n                        \n        fg1_signal = signal[i, :, 0].copy()\n        fg1_signal = savgol_filter(fg1_signal, 11, 1)\n        fg1_slope = savgol_filter(signal[i, :, 0].copy(), 41, 2)\n        _, _ = calc_depth_and_detrend(fg1_slope, *(b[i] for b in boundaries), outliers[i], fixed_degree=3) # for detrending\n        \n        df.loc[i, 'fg1_average_depth'], df.loc[i, 'fg1_mid_depth'] = calc_depth_and_detrend(fg1_signal, *(b[i] for b in boundaries), \n                                                                                            outliers[i], \n                                                                                            fixed_degree=3)\n\n        \n        # transit depth percentiles\n        for q in [0.01, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]:\n            df.loc[i, f'q_1_{q}'] = calc_perc(signal_mean[i], mid1[i], mid2[i], q, outliers[i])\n            df.loc[i, f'q_2_{q}'] = calc_perc(signal_mean[i], lower[i] - buf1, upper[i] + buf1, q, outliers[i])\n            df.loc[i, f'q_3_{q}'] = calc_perc(signal_mean[i], lower[i], upper[i], q, outliers[i]) \n        for q in [0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6, 0.7, 0.8, 0.9]:\n            df.loc[i, f'fg1_q_1_{q}'] = calc_perc(fg1_signal, mid1[i], mid2[i], q, outliers[i])\n            df.loc[i, f'fg1_q_2_{q}'] = calc_perc(fg1_signal, lower[i], upper[i], q, outliers[i])\n            df.loc[i, f'fg1_q_3_{q}'] = calc_perc(fg1_signal, lower[i] - buf1, upper[i] + buf1, q, outliers[i])\n\n        \n        # slope features\n        df.loc[i, 'slope'] = calc_slope(signal_mean[i], *(b[i] for b in boundaries), outliers[i])\n        df.loc[i, 'slope_2'] = calc_slope_2(signal_mean[i], *(b[i] for b in boundaries), outliers[i])\n        df.loc[i, 'slope_2_left'] = calc_slope_2_left(signal_mean[i], *(b[i] for b in boundaries), outliers[i])\n        df.loc[i, 'slope_2_right'] = calc_slope_2_right(signal_mean[i], *(b[i] for b in boundaries), outliers[i])\n        df.loc[i, 'slope_g'] = max(0, -df.loc[i, 'slope_2'])**0.5\n                          \n        df.loc[i, 'fg1_slope'] = calc_slope(fg1_slope, *(b[i] for b in boundaries), outliers[i])     \n        df.loc[i, 'fg1_slope_2'] = calc_slope_2(fg1_slope, *(b[i] for b in boundaries), outliers[i])      \n        df.loc[i, 'fg1_slope_g'] = max(0, -df.loc[i, 'fg1_slope_2'])**0.5\n        df.loc[i, 'fg1_curv_left'] = calc_curv_left(fg1_slope, *(b[i] for b in boundaries), outliers[i])\n        df.loc[i, 'fg1_curv_right'] = calc_curv_right(fg1_slope, *(b[i] for b in boundaries), outliers[i])\n\n        \n        # combinations with slopes\n        df.loc[i, 'slope_rel'] = df.loc[i, 'slope_2'] * df.loc[i, 'average_depth']\n        df.loc[i, 'fg1_slope_T'] = df.loc[i, 'fg1_slope_2'] * star_info.loc[i, 'Ts'] \n        df.loc[i, 'fg1_slope_rel'] = df.loc[i, 'fg1_slope_2'] * df.loc[i, 'fg1_average_depth']\n        df.loc[i, 'fg1_slope_g_rel'] = df.loc[i, 'fg1_slope_g'] * df.loc[i, 'fg1_average_depth']\n        \n        for q in [0.01, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]:\n            df.loc[i, f'slope_q_{q}'] = df.loc[i, 'slope_2'] * df.loc[i, f'q_1_{q}']\n            df.loc[i, f'slope_q_{q}_2'] = df.loc[i, 'slope_2'] * df.loc[i, f'q_2_{q}']\n\n        \n        # other features\n        df.loc[i, 't14'] = upper[i] - lower[i]\n        df.loc[i, 't23'] = mid2[i] - mid1[i]\n        df.loc[i, 'time'] = (mid1[i] - lower[i]) / (upper[i] - lower[i])        \n        df.loc[i, 'P_mul_Rs'] = star_info.loc[i, 'P'] * star_info.loc[i, 'Rs']\n        df.loc[i, 'P_div_Rs'] = star_info.loc[i, 'P'] / star_info.loc[i, 'Rs']\n        \n        step = 5\n        max_rel = 0\n        min_rel = 1\n        meaning = 60 # window size for frequency averagning\n        \n        for j in range(1, signal.shape[-1] - meaning + 1, step):\n            if j <= 80:\n                meaning = 20\n            elif j <= 180:\n                meaning = 30\n            else:\n                meaning = 60\n\n            cur_mean = signal[i, :, j : min(j + meaning, signal.shape[-1])].mean(-1)\n\n            # median filter\n            if not outliers[i]:\n                if j >= 180:\n                    med_kernel = 31\n                else:\n                    med_kernel = 21\n                cur_mean = median_filter(cur_mean, size=med_kernel, mode=\"constant\")\n                           \n            cur_mean = savgol_filter(cur_mean, 11, 1)\n            \n            df.loc[i, f'averaged_{j}_unstable'], df.loc[i, f'mid_{j}_unstable'] = calc_depth_and_detrend(cur_mean.copy(), *(b[i] for b in boundaries), \n                                                                                                            outliers[i],\n                                                                                                            unstable=True)  \n            if not outliers[i]:\n                df.loc[i, f'averaged_{j}_unstable'] = 1 - (1 - df.loc[i, f'averaged_{j}_unstable']) * norm_coef\n  \n            if j in good_waves:\n                df.loc[i, f'averaged_{j}'], df.loc[i, f'mid_{j}'] = calc_depth_and_detrend(cur_mean, *(b[i] for b in boundaries),\n                                                                                              outliers[i],\n                                                                                              fixed_degree=3)\n                if not outliers[i]:\n                    df.loc[i, f'averaged_{j}'] = 1 - (1 - df.loc[i, f'averaged_{j}']) * norm_coef_2\n\n            \n            # percentiles\n            for q in [0.1, 0.15, 0.2]:\n                if j in good_waves and not outliers[i]:\n                    df.loc[i, f'q_w_{j}_{q}'] = calc_perc(cur_mean, mid1[i], mid2[i], q, outliers[i])\n                elif outliers[i] and mid1[i] < mid2[i]:\n                    x_combined = np.concatenate([np.arange(max(lower[i] - buf1, 1)), np.arange(min(upper[i] + buf1, max_len), max_len + 1)], axis=-1)\n                    mid_q = np.quantile(cur_mean[mid1[i] : mid2[i]], q)\n                    df.loc[i, f'q_w_{j}_{q}'] = 1 - mid_q / cur_mean[x_combined].mean()\n                else:\n                    df.loc[i, f'q_w_{j}_{q}'] = 0\n           \n            \n            max_rel = max(max_rel, df.loc[i, f'averaged_{j}_unstable'])\n            min_rel = min(min_rel, df.loc[i, f'averaged_{j}_unstable'])\n\n  \n            # slope combinations\n            df.loc[i, f'averaged_slope_{j}'] = df.loc[i, f'averaged_{j}_unstable'] * df.loc[i, 'slope_2']\n            df.loc[i, f'averaged_slope_g_{j}'] = df.loc[i, f'averaged_{j}_unstable'] * df.loc[i, 'slope_g']\n\n        \n        # large amplitude     \n        if max_rel - min_rel >= 0.005:\n            df.loc[i, 'very_bad'] = True\n\n    \n    df['Rs'] = star_info['Rs']\n    df['Ms'] = star_info['Ms']\n    df['Ts'] = star_info['Ts']\n    df['sma'] = star_info['sma']  \n    df['g'] = np.log10(star_info['Ms'] / (star_info['Rs']**2))\n    df['g_T'] = df['g'] * star_info['Ts']\n    df['big_rs'] = (star_info['Rs'] > np.quantile(star_info['Rs'].values, 0.97))\n    \n    df['outliers'] = outliers\n\n    df = df.fillna(0)\n    \n    return outliers, df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T21:42:46.974156Z","iopub.execute_input":"2025-10-01T21:42:46.974563Z","iopub.status.idle":"2025-10-01T21:42:46.989879Z","shell.execute_reply.started":"2025-10-01T21:42:46.974537Z","shell.execute_reply":"2025-10-01T21:42:46.988709Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"exec(open('feature_engineering.py', 'r').read())\noutliers, train = feature_engineering(train_star_info, data_train)\noutliers = np.arange(train.shape[0])[outliers]\nlen(outliers)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T21:42:49.24665Z","iopub.execute_input":"2025-10-01T21:42:49.247031Z","iopub.status.idle":"2025-10-01T21:43:26.497924Z","shell.execute_reply.started":"2025-10-01T21:42:49.247007Z","shell.execute_reply":"2025-10-01T21:43:26.496988Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.pipeline import make_pipeline\nfrom sklearn.preprocessing import RobustScaler\nfrom sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin\nfrom sklearn.utils.validation import check_is_fitted\n\n\nclass CustomRidge(BaseEstimator):\n    \"\"\"\n    Provides two models: main model for normal cases and outlier-specific model\n    \"\"\"\n    def __init__(self):\n        self.main = Ridge(alpha=3e-2) \n        self.outliers = Ridge(alpha=3e-1)\n\n        self.main_scaler = RobustScaler()\n        self.outliers_scaler = RobustScaler()\n\n\n    def fit(self, X, y):\n        groups = X['outliers']\n        X = X.drop(columns=['outliers'])\n\n        main_mask = (groups == 0).values\n        main_mask[X.loc[:, 'very_bad'] == True] = 0\n        \n        X_main = self.main_scaler.fit_transform(X[main_mask])\n\n        self.main.fit(self.main_scaler.transform(X[main_mask]), y[main_mask])\n        self.outliers.fit(self.outliers_scaler.fit_transform(X), y)     \n        self.pred_shape = y.shape[-1]\n\n        return self\n\n    def predict(self, X):\n        groups = X['outliers']\n        X = X.drop(columns=['outliers'])\n        \n        predictions = np.zeros((X.shape[0], self.pred_shape))\n\n        main_mask = (groups == 0).values\n        if main_mask.sum():\n            predictions[main_mask] = self.main.predict(self.main_scaler.transform(X[main_mask]))                                                 \n        if (~main_mask).sum():\n            predictions[~main_mask] = self.outliers.predict(self.outliers_scaler.transform(X[~main_mask]))\n        \n        return predictions\n\n\nmodel = CustomRidge()\n\noof_pred = cross_val_predict(model, train, train_labels.values, cv=100)\n\nprint(f\"# R2 score: {r2_score(train_labels, oof_pred):.3f}\")\n\nsigma_pred = mean_squared_error(train_labels, oof_pred, squared=False)\n  \nprint(f\"# Root mean squared error: {sigma_pred:.6f}\")\n\ncol = 1\nplt.scatter(oof_pred[:,col], train_labels.iloc[:,col], s=15, c='lightgreen')\nplt.gca().set_aspect('equal')\nplt.xlabel('y_pred')\nplt.ylabel('y_true')\nplt.title('Comparing y_true and y_pred')\nplt.show()\n\n# clipping\noof_pred = np.maximum(oof_pred, 0.003)\noof_pred = np.minimum(oof_pred, 0.1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T21:43:31.038515Z","iopub.execute_input":"2025-10-01T21:43:31.038945Z","iopub.status.idle":"2025-10-01T21:43:37.130836Z","shell.execute_reply.started":"2025-10-01T21:43:31.038915Z","shell.execute_reply":"2025-10-01T21:43:37.12967Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SigmaPredictor:\n    \"\"\"\n    Class for sigma predicting\n    \"\"\"\n    def __init__(self):\n        self.sigmas = {}\n        \n    def fit(self, y_pred, y_true, outliers, very_bad):        \n        outliers = [i for i in outliers if i not in very_bad]\n\n        self.sigmas['outliers'] = self._calc(y_pred[outliers], y_true[outliers]) * 5\n\n        main = self._del_outliers(np.ones(len(y_pred), dtype=bool), outliers + list(very_bad))\n        self.sigmas['main'] = self._calc(y_pred[main], y_true[main])\n\n        print({ k: v.mean() for k, v in self.sigmas.items() })\n\n    def predict(self, sigma_pred, y_pred, outliers, very_bad, bootstrap_preds=None):\n        if len(outliers) > 0:\n            sigma_pred[outliers] = self.sigmas['outliers']\n\n        main = self._del_outliers(np.ones(len(y_pred), dtype=bool), outliers)\n        if main.sum() > 0:\n            sigma_pred[main] = self.sigmas['main']\n\n        W1 = 0.75\n        W2 = 1.0 - W1\n        \n        sigma_pred[main, :] = bootstrap_preds[main, :] * W1 + sigma_pred[main, :] * W2\n        sigma_pred[outliers] = bootstrap_preds[outliers] * 1.5\n\n        sigma_pred[very_bad] = 0.003 \n        for i in very_bad:\n            if i in outliers:\n                continue\n            sigma_pred[i, :] = 0.5 * bootstrap_preds[i, :] + 0.5 * sigma_pred[i, :]\n\n        return sigma_pred\n        \n\n    def _calc(self, y_pred, y_true): # calculate rmse for each frequency\n        sigmas = []\n        for i in range(y_pred.shape[1]):\n            sigmas.append(mean_squared_error(y_pred[:, i], y_true[:, i], squared=False))\n        return np.array(sigmas)\n\n    def _del_outliers(self, mask, outliers):\n        for i in range(len(mask)):\n            if i in outliers:\n                mask[i] = False\n        return mask                         \n\n\ndef postprocessing(pred_array, index, sigma_pred, sigma_predictor, outliers, very_bad, bootstrap_preds=None, column_names=None):\n    \"\"\"\n    Creates a submission DataFrame with mean predictions and uncertainties.\n\n    Parameters:\n    - pred_array: ndarray of shape (n_samples, 283)\n    - index: pandas.Index of length n_samples\n    - sigma_pred: float or ndarray of shape (n_samples, 283)\n    - column_names: list of wavelength column names (optional)\n\n    Returns:\n    - df: DataFrame of shape (n_samples, 566)\n    \"\"\"\n    n_samples, n_waves = pred_array.shape\n\n    if column_names is None:\n        column_names = [f\"wl_{i+1}\" for i in range(n_waves)]\n    \n    sigma_pred = sigma_predictor.predict(np.zeros_like(pred_array), pred_array, outliers, very_bad, bootstrap_preds=bootstrap_preds)\n\n    # Safety check\n    assert sigma_pred.shape == pred_array.shape, \"sigma_pred must match shape of pred_array\"\n    assert len(index) == n_samples, \"Index length must match number of rows\"\n\n    df_mean = pd.DataFrame(pred_array.clip(0, None), index=index, columns=column_names)\n    df_sigma = pd.DataFrame(sigma_pred, index=index, columns=[f\"sigma_{i+1}\" for i in range(n_waves)])\n\n    return pd.concat([df_mean, df_sigma], axis=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-01T21:31:20.797225Z","iopub.execute_input":"2025-10-01T21:31:20.797661Z","iopub.status.idle":"2025-10-01T21:31:20.821289Z","shell.execute_reply.started":"2025-10-01T21:31:20.797635Z","shell.execute_reply":"2025-10-01T21:31:20.820272Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.fit(train, train_labels)\n\nsigma_predictor = SigmaPredictor()\nvery_bad = np.arange(train_labels.shape[0])[train['very_bad'].values == True]\nsigma_predictor.fit(oof_pred, train_labels.values, outliers, very_bad)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T17:04:43.578506Z","iopub.execute_input":"2025-09-30T17:04:43.578828Z","iopub.status.idle":"2025-09-30T17:04:44.278693Z","shell.execute_reply.started":"2025-09-30T17:04:43.578806Z","shell.execute_reply":"2025-09-30T17:04:44.277596Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def bootstrap_uncertainty_inference( \n    X_train,\n    y_train,\n    X_test,\n    n_bootstraps = 100, \n    random_state = 42,\n):\n    \"\"\"\n    Sigma estimation via bootstrapping\n    \"\"\"\n    if random_state is not None:\n        np.random.seed(random_state)\n    \n    y_train_values = y_train\n        \n    n_test_samples = X_test.shape[0]\n    n_targets = y_train_values.shape[1]\n    \n    predictions = np.full((n_test_samples, n_targets, n_bootstraps), np.nan)\n    \n    bootstrap_iter = range(n_bootstraps)\n    bootstrap_iter = tqdm(bootstrap_iter, desc=\"bootstrap interations\")\n    \n    for b in bootstrap_iter:\n        bootstrap_indices = np.random.choice(\n            len(X_train), size=len(X_train), replace=True\n        )\n        \n        X_bootstrap = X_train.iloc[bootstrap_indices].reset_index(drop=True)\n        y_bootstrap = y_train_values.iloc[bootstrap_indices].reset_index(drop=True)\n        \n        model_bootstrap = CustomRidge()\n        model_bootstrap.fit(X_bootstrap, y_bootstrap)\n        \n        y_pred = model_bootstrap.predict(X_test)   \n        predictions[:, :, b] = y_pred\n     \n    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-30T17:04:54.45816Z","iopub.execute_input":"2025-09-30T17:04:54.458504Z","iopub.status.idle":"2025-09-30T17:04:54.469582Z","shell.execute_reply.started":"2025-09-30T17:04:54.458479Z","shell.execute_reply":"2025-09-30T17:04:54.468309Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pickle\nimport gc\n\n\ntest_adc_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/test_star_info.csv', index_col='planet_id')\nsample_submission = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/sample_submission.csv', index_col='planet_id')\nwavelengths = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/wavelengths.csv')\ntest_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/test_star_info.csv')\n\ndel data_train\n\ngc.collect()\nos.environ[\"PREPROCESS_MODE\"] = \"test\"\n\n!python preprocess.py\n!rm -rf *AIRS-CH0_signal*\n\ndata_test = np.load(f\"signal_{VERSION}.npy\")\n\noutliers, test_features = feature_engineering(test_star_info, data_test)\noutliers = np.arange(test_features.shape[0])[outliers]\nvery_bad = np.arange(test_features.shape[0])[test_features['very_bad'].values == True]\n\ntest_pred = model.predict(test_features)\n\nboot_pred = bootstrap_uncertainty_inference(train, train_labels, test_features, n_bootstraps=1000)\ntest_bootstrap_preds = boot_pred.std(-1) * 2.75\n\ntest_pred = np.maximum(test_pred, 0.003)\ntest_pred = np.minimum(test_pred, 0.1)\n\nprint('sigma', sigma_pred)\n\n\ndef postprocessing(pred_array, index, sigma_pred, sigma_predictor, outliers, very_bad, bootstrap_preds=None, column_names=None):\n    \"\"\"\n    Convert predictions and uncertainty into final submission DataFrame\n    \"\"\"\n\n    sigma_array = sigma_predictor.predict(np.zeros_like(pred_array), pred_array, outliers, very_bad, bootstrap_preds=bootstrap_preds)\n\n    df_pred = pd.DataFrame(pred_array.clip(0, None), index=index, columns=column_names)\n    df_sigma = pd.DataFrame(sigma_array, index=index, columns=[f\"sigma_{i}\" for i in range(1, len(column_names)+1)])\n    return pd.concat([df_pred, df_sigma], axis=1)\n\n\nsubmission_df = postprocessing(\n    pred_array=test_pred,\n    index=sample_submission.index,\n    sigma_pred=sigma_pred,\n    sigma_predictor=sigma_predictor,\n    column_names=wavelengths.columns,\n    bootstrap_preds=test_bootstrap_preds,\n    outliers=outliers,\n    very_bad=very_bad\n)\n\nsubmission_df.to_csv('submission.csv')\n!head submission.csv","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}