{"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"},{"sourceId":12857711,"sourceType":"datasetVersion","datasetId":8132602},{"sourceId":12857722,"sourceType":"datasetVersion","datasetId":8132608}],"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 itertools\nfrom tqdm import tqdm\nimport multiprocessing as mp\nfrom numpy.polynomial import Polynomial\nfrom astropy.stats import sigma_clip\nimport numba\nfrom sklearn.preprocessing import StandardScaler\nfrom scipy.ndimage import gaussian_filter1d\nfrom scipy.signal import savgol_filter\nfrom sklearn.linear_model import LinearRegression\nimport matplotlib.pyplot as plt\n\nlinear = LinearRegression()\n\nROOT = \"/kaggle/input/ariel-data-challenge-2025\"\nVERSION = \"v3\"\n\nBINNING = 15\n\nsensor_sizes_dict = {\n    \"AIRS-CH0\": [[11250, 32, 356], [32, 356]],\n    \"FGS1\": [[135000, 32, 32], [32, 32]],\n}  # input, mask\n\n# 16 center pixels, rest is just noise\ncl = 8\ncr = 24","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:06.050135Z","iopub.execute_input":"2025-09-01T21:42:06.050352Z","iopub.status.idle":"2025-09-01T21:42:08.955056Z","shell.execute_reply.started":"2025-09-01T21:42:06.050326Z","shell.execute_reply":"2025-09-01T21:42:08.954496Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"planet_id = 1042982756\nsensor = \"FGS1\"\n\nMODE = \"train\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:08.9564Z","iopub.execute_input":"2025-09-01T21:42:08.956727Z","iopub.status.idle":"2025-09-01T21:42:08.960308Z","shell.execute_reply.started":"2025-09-01T21:42:08.956709Z","shell.execute_reply":"2025-09-01T21:42:08.959591Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:08.96108Z","iopub.execute_input":"2025-09-01T21:42:08.961388Z","iopub.status.idle":"2025-09-01T21:42:08.978588Z","shell.execute_reply.started":"2025-09-01T21:42:08.961361Z","shell.execute_reply":"2025-09-01T21:42:08.977988Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"get_gain_offset()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:08.979199Z","iopub.execute_input":"2025-09-01T21:42:08.979413Z","iopub.status.idle":"2025-09-01T21:42:08.992638Z","shell.execute_reply.started":"2025-09-01T21:42:08.979397Z","shell.execute_reply":"2025-09-01T21:42:08.991891Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nget_gain_offset()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:08.993331Z","iopub.execute_input":"2025-09-01T21:42:08.993544Z","iopub.status.idle":"2025-09-01T21:42:14.149114Z","shell.execute_reply.started":"2025-09-01T21:42:08.993528Z","shell.execute_reply":"2025-09-01T21:42:14.148328Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def read_data(planet_id, sensor, mode=\"train\"):\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    # read_frame = pd.read_parquet(\n    #     f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/read.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    # read_frame = read_frame.values.reshape(sensor_sizes_dict[sensor][1])\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        # read_frame,\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:14.150011Z","iopub.execute_input":"2025-09-01T21:42:14.150283Z","iopub.status.idle":"2025-09-01T21:42:14.157459Z","shell.execute_reply.started":"2025-09-01T21:42:14.150256Z","shell.execute_reply":"2025-09-01T21:42:14.156902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nread_data(planet_id, sensor, mode=MODE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:14.159432Z","iopub.execute_input":"2025-09-01T21:42:14.159614Z","iopub.status.idle":"2025-09-01T21:42:21.28093Z","shell.execute_reply.started":"2025-09-01T21:42:14.159599Z","shell.execute_reply":"2025-09-01T21:42:21.28001Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"AIRS data is larger and takes longer.\n\nNote: it will already take about 1100 * (1.46s + 0.47s) = ~35 minutes to just load the data into RAM for further processing! (single core estimate here and below)","metadata":{}},{"cell_type":"code","source":"def 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    return signal / gain + offset","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:21.281793Z","iopub.execute_input":"2025-09-01T21:42:21.282419Z","iopub.status.idle":"2025-09-01T21:42:21.286874Z","shell.execute_reply.started":"2025-09-01T21:42:21.282388Z","shell.execute_reply":"2025-09-01T21:42:21.286141Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal, dark_frame, dead_frame, linear_corr, flat_frame = read_data(\n    planet_id, sensor, mode=MODE\n)\ngain, offset = get_gain_offset()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:21.287597Z","iopub.execute_input":"2025-09-01T21:42:21.287863Z","iopub.status.idle":"2025-09-01T21:42:21.951224Z","shell.execute_reply.started":"2025-09-01T21:42:21.287841Z","shell.execute_reply":"2025-09-01T21:42:21.950572Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:21.951931Z","iopub.execute_input":"2025-09-01T21:42:21.952204Z","iopub.status.idle":"2025-09-01T21:42:22.165389Z","shell.execute_reply.started":"2025-09-01T21:42:21.952183Z","shell.execute_reply":"2025-09-01T21:42:22.164807Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nADC_convert(signal, gain, offset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:22.166093Z","iopub.execute_input":"2025-09-01T21:42:22.166305Z","iopub.status.idle":"2025-09-01T21:42:23.720152Z","shell.execute_reply.started":"2025-09-01T21:42:22.166285Z","shell.execute_reply":"2025-09-01T21:42:23.719465Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = ADC_convert(signal, gain, offset)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:23.72091Z","iopub.execute_input":"2025-09-01T21:42:23.721186Z","iopub.status.idle":"2025-09-01T21:42:23.919697Z","shell.execute_reply.started":"2025-09-01T21:42:23.721162Z","shell.execute_reply":"2025-09-01T21:42:23.918894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:23.92049Z","iopub.execute_input":"2025-09-01T21:42:23.920771Z","iopub.status.idle":"2025-09-01T21:42:24.129996Z","shell.execute_reply.started":"2025-09-01T21:42:23.920744Z","shell.execute_reply":"2025-09-01T21:42:24.129381Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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    # Set values to np.nan where dead or hot pixels are found\n    signal[dead] = np.nan\n    signal[hot] = np.nan\n    return signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:24.130729Z","iopub.execute_input":"2025-09-01T21:42:24.130994Z","iopub.status.idle":"2025-09-01T21:42:24.135465Z","shell.execute_reply.started":"2025-09-01T21:42:24.130976Z","shell.execute_reply":"2025-09-01T21:42:24.134872Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nmask_hot_dead(signal, dead_frame, dark_frame)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:24.136179Z","iopub.execute_input":"2025-09-01T21:42:24.136442Z","iopub.status.idle":"2025-09-01T21:42:30.552279Z","shell.execute_reply.started":"2025-09-01T21:42:24.136418Z","shell.execute_reply":"2025-09-01T21:42:30.551597Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = mask_hot_dead(signal, dead_frame, dark_frame)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:30.553051Z","iopub.execute_input":"2025-09-01T21:42:30.553355Z","iopub.status.idle":"2025-09-01T21:42:30.636965Z","shell.execute_reply.started":"2025-09-01T21:42:30.553335Z","shell.execute_reply":"2025-09-01T21:42:30.636071Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:30.637864Z","iopub.execute_input":"2025-09-01T21:42:30.638137Z","iopub.status.idle":"2025-09-01T21:42:30.84979Z","shell.execute_reply.started":"2025-09-01T21:42:30.638112Z","shell.execute_reply":"2025-09-01T21:42:30.849186Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:30.850466Z","iopub.execute_input":"2025-09-01T21:42:30.850708Z","iopub.status.idle":"2025-09-01T21:42:30.85524Z","shell.execute_reply.started":"2025-09-01T21:42:30.85069Z","shell.execute_reply":"2025-09-01T21:42:30.85449Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cupy as cp\n\ndef apply_linear_corr_gpu(linear_corr, clean_signal):\n    \"\"\"\n    Step 3: linearity Correction on GPU\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    \"\"\"\n    # Convert the input arrays to CuPy arrays\n    linear_corr_gpu = cp.asarray(linear_corr)\n    clean_signal_gpu = cp.asarray(clean_signal)\n\n    corrected_signal_gpu = (\n        (\n            (\n                (linear_corr_gpu[5] * clean_signal_gpu + linear_corr_gpu[4])\n                * clean_signal_gpu\n                + linear_corr_gpu[3]\n            )\n            * clean_signal_gpu\n            + linear_corr_gpu[2]\n        )\n        * clean_signal_gpu\n        + linear_corr_gpu[1]\n    ) * clean_signal_gpu + linear_corr_gpu[0]\n\n    # Convert the result back to a NumPy array (if needed)\n    corrected_signal = cp.asnumpy(corrected_signal_gpu)\n\n    return corrected_signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:30.856053Z","iopub.execute_input":"2025-09-01T21:42:30.856265Z","iopub.status.idle":"2025-09-01T21:42:33.106666Z","shell.execute_reply.started":"2025-09-01T21:42:30.856239Z","shell.execute_reply":"2025-09-01T21:42:33.10605Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.nanmax(signal), np.nanmin(signal), np.nanmean(signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:33.107456Z","iopub.execute_input":"2025-09-01T21:42:33.107847Z","iopub.status.idle":"2025-09-01T21:42:33.737275Z","shell.execute_reply.started":"2025-09-01T21:42:33.107827Z","shell.execute_reply":"2025-09-01T21:42:33.736665Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\napply_linear_corr(linear_corr, signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:33.737989Z","iopub.execute_input":"2025-09-01T21:42:33.738323Z","iopub.status.idle":"2025-09-01T21:42:42.448254Z","shell.execute_reply.started":"2025-09-01T21:42:33.738302Z","shell.execute_reply":"2025-09-01T21:42:42.447434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"assert np.allclose(apply_linear_corr(linear_corr, signal), apply_linear_corr_gpu(linear_corr, signal), equal_nan=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:42.451948Z","iopub.execute_input":"2025-09-01T21:42:42.452174Z","iopub.status.idle":"2025-09-01T21:42:48.280556Z","shell.execute_reply.started":"2025-09-01T21:42:42.452158Z","shell.execute_reply":"2025-09-01T21:42:48.27983Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = apply_linear_corr(linear_corr, signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:48.281277Z","iopub.execute_input":"2025-09-01T21:42:48.28148Z","iopub.status.idle":"2025-09-01T21:42:49.383297Z","shell.execute_reply.started":"2025-09-01T21:42:48.281464Z","shell.execute_reply":"2025-09-01T21:42:49.382568Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:49.384148Z","iopub.execute_input":"2025-09-01T21:42:49.384397Z","iopub.status.idle":"2025-09-01T21:42:49.594835Z","shell.execute_reply.started":"2025-09-01T21:42:49.384371Z","shell.execute_reply":"2025-09-01T21:42:49.594192Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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 = np.tile(dark, (signal.shape[0], 1, 1))\n    signal -= dark * dt[:, np.newaxis, np.newaxis]\n    return signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:49.595513Z","iopub.execute_input":"2025-09-01T21:42:49.595745Z","iopub.status.idle":"2025-09-01T21:42:49.60026Z","shell.execute_reply.started":"2025-09-01T21:42:49.595727Z","shell.execute_reply":"2025-09-01T21:42:49.599581Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"axis_info = pd.read_parquet(ROOT + \"/axis_info.parquet\")\ndt_airs = axis_info[\"AIRS-CH0-integration_time\"].dropna().values\n\ndt_fgs = np.ones(len(signal)) * 0.1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:49.600945Z","iopub.execute_input":"2025-09-01T21:42:49.601589Z","iopub.status.idle":"2025-09-01T21:42:49.644187Z","shell.execute_reply.started":"2025-09-01T21:42:49.601562Z","shell.execute_reply":"2025-09-01T21:42:49.643475Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if sensor == \"AIRS-CH0\":\n    dt = dt_airs\nelif sensor == \"FGS1\":\n    dt = dt_fgs\n\ndt[1::2] += 0.1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:49.644934Z","iopub.execute_input":"2025-09-01T21:42:49.645174Z","iopub.status.idle":"2025-09-01T21:42:49.649162Z","shell.execute_reply.started":"2025-09-01T21:42:49.645152Z","shell.execute_reply":"2025-09-01T21:42:49.648384Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nclean_dark(signal, dark_frame, dt)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:49.649954Z","iopub.execute_input":"2025-09-01T21:42:49.650237Z","iopub.status.idle":"2025-09-01T21:42:55.242761Z","shell.execute_reply.started":"2025-09-01T21:42:49.650214Z","shell.execute_reply":"2025-09-01T21:42:55.242015Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = clean_dark(signal, dark_frame, dt)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:55.243532Z","iopub.execute_input":"2025-09-01T21:42:55.243815Z","iopub.status.idle":"2025-09-01T21:42:55.952863Z","shell.execute_reply.started":"2025-09-01T21:42:55.243792Z","shell.execute_reply":"2025-09-01T21:42:55.951951Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:55.953788Z","iopub.execute_input":"2025-09-01T21:42:55.954053Z","iopub.status.idle":"2025-09-01T21:42:56.164451Z","shell.execute_reply.started":"2025-09-01T21:42:55.954025Z","shell.execute_reply":"2025-09-01T21:42:56.163697Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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 np.subtract(signal[1::2, :, :], signal[::2, :, :])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:56.165309Z","iopub.execute_input":"2025-09-01T21:42:56.165855Z","iopub.status.idle":"2025-09-01T21:42:56.169693Z","shell.execute_reply.started":"2025-09-01T21:42:56.165827Z","shell.execute_reply":"2025-09-01T21:42:56.168976Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nget_cds(signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:42:56.170354Z","iopub.execute_input":"2025-09-01T21:42:56.170532Z","iopub.status.idle":"2025-09-01T21:43:04.633374Z","shell.execute_reply.started":"2025-09-01T21:42:56.170516Z","shell.execute_reply":"2025-09-01T21:43:04.632603Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = get_cds(signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:04.634116Z","iopub.execute_input":"2025-09-01T21:43:04.634384Z","iopub.status.idle":"2025-09-01T21:43:04.742989Z","shell.execute_reply.started":"2025-09-01T21:43:04.634365Z","shell.execute_reply":"2025-09-01T21:43:04.742167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:04.743913Z","iopub.execute_input":"2025-09-01T21:43:04.744386Z","iopub.status.idle":"2025-09-01T21:43:04.852397Z","shell.execute_reply.started":"2025-09-01T21:43:04.744353Z","shell.execute_reply":"2025-09-01T21:43:04.851851Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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    assert signal.shape[0] % binning == 0  # Ensure the binning is possible\n\n    # cds_transposed = signal.transpose(0, 2, 1)\n    cds_binned = torch.zeros(\n        (\n            signal.shape[0] // binning,\n            signal.shape[1],\n            signal.shape[2],\n        ),\n        device=\"cuda:0\",\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:04.853163Z","iopub.execute_input":"2025-09-01T21:43:04.853419Z","iopub.status.idle":"2025-09-01T21:43:04.859429Z","shell.execute_reply.started":"2025-09-01T21:43:04.853399Z","shell.execute_reply":"2025-09-01T21:43:04.858929Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:04.860212Z","iopub.execute_input":"2025-09-01T21:43:04.860505Z","iopub.status.idle":"2025-09-01T21:43:09.281764Z","shell.execute_reply.started":"2025-09-01T21:43:04.860489Z","shell.execute_reply":"2025-09-01T21:43:09.281186Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nbin_obs(torch.tensor(signal), binning=BINNING * 12)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:09.282458Z","iopub.execute_input":"2025-09-01T21:43:09.282894Z","iopub.status.idle":"2025-09-01T21:43:10.863284Z","shell.execute_reply.started":"2025-09-01T21:43:09.282875Z","shell.execute_reply":"2025-09-01T21:43:10.86246Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = bin_obs(torch.tensor(signal), binning=BINNING * 12)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:10.8641Z","iopub.execute_input":"2025-09-01T21:43:10.864469Z","iopub.status.idle":"2025-09-01T21:43:11.052146Z","shell.execute_reply.started":"2025-09-01T21:43:10.864448Z","shell.execute_reply":"2025-09-01T21:43:11.051571Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:11.052812Z","iopub.execute_input":"2025-09-01T21:43:11.053063Z","iopub.status.idle":"2025-09-01T21:43:11.057256Z","shell.execute_reply.started":"2025-09-01T21:43:11.053038Z","shell.execute_reply":"2025-09-01T21:43:11.05651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\ncorrect_flat_field(torch.tensor(flat_frame).to(\"cuda:0\"), signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:11.058143Z","iopub.execute_input":"2025-09-01T21:43:11.058403Z","iopub.status.idle":"2025-09-01T21:43:14.381241Z","shell.execute_reply.started":"2025-09-01T21:43:11.05838Z","shell.execute_reply":"2025-09-01T21:43:14.380546Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = correct_flat_field(torch.tensor(flat_frame).to(\"cuda:0\"), signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:14.381947Z","iopub.execute_input":"2025-09-01T21:43:14.382171Z","iopub.status.idle":"2025-09-01T21:43:14.386438Z","shell.execute_reply.started":"2025-09-01T21:43:14.382146Z","shell.execute_reply":"2025-09-01T21:43:14.385717Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.isnan(signal).sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:14.38727Z","iopub.execute_input":"2025-09-01T21:43:14.387493Z","iopub.status.idle":"2025-09-01T21:43:14.493222Z","shell.execute_reply.started":"2025-09-01T21:43:14.387478Z","shell.execute_reply":"2025-09-01T21:43:14.492646Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch.nn.functional as F\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:14.493954Z","iopub.execute_input":"2025-09-01T21:43:14.494212Z","iopub.status.idle":"2025-09-01T21:43:14.500392Z","shell.execute_reply.started":"2025-09-01T21:43:14.494193Z","shell.execute_reply":"2025-09-01T21:43:14.499645Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\nnan_interpolation(signal)\nnan_interpolation(signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:14.500985Z","iopub.execute_input":"2025-09-01T21:43:14.501147Z","iopub.status.idle":"2025-09-01T21:43:14.969865Z","shell.execute_reply.started":"2025-09-01T21:43:14.501134Z","shell.execute_reply":"2025-09-01T21:43:14.969111Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = nan_interpolation(signal)\nsignal = nan_interpolation(signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:14.970551Z","iopub.execute_input":"2025-09-01T21:43:14.970828Z","iopub.status.idle":"2025-09-01T21:43:14.975643Z","shell.execute_reply.started":"2025-09-01T21:43:14.97081Z","shell.execute_reply":"2025-09-01T21:43:14.975098Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%timeit\n\ntorch.nanmean(signal, axis=[1, 2]).cpu().numpy()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:14.976339Z","iopub.execute_input":"2025-09-01T21:43:14.976539Z","iopub.status.idle":"2025-09-01T21:43:21.722476Z","shell.execute_reply.started":"2025-09-01T21:43:14.976515Z","shell.execute_reply":"2025-09-01T21:43:21.721852Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal = torch.nanmean(signal, axis=[1, 2]).cpu().numpy()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:21.723129Z","iopub.execute_input":"2025-09-01T21:43:21.723334Z","iopub.status.idle":"2025-09-01T21:43:21.727511Z","shell.execute_reply.started":"2025-09-01T21:43:21.723317Z","shell.execute_reply":"2025-09-01T21:43:21.726802Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:21.728227Z","iopub.execute_input":"2025-09-01T21:43:21.728441Z","iopub.status.idle":"2025-09-01T21:43:21.745338Z","shell.execute_reply.started":"2025-09-01T21:43:21.728426Z","shell.execute_reply":"2025-09-01T21:43:21.744837Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### With all speedups, we can process all planets including Linear Correction!","metadata":{}},{"cell_type":"code","source":"import os\n\n\nos.environ[\"PREPROCESS_MODE\"] = \"train\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:21.745896Z","iopub.execute_input":"2025-09-01T21:43:21.746064Z","iopub.status.idle":"2025-09-01T21:43:21.758517Z","shell.execute_reply.started":"2025-09-01T21:43:21.74605Z","shell.execute_reply":"2025-09-01T21:43:21.757897Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport itertools\nfrom tqdm import tqdm\nfrom multiprocessing.dummy import Pool as ThreadPool\nfrom astropy.stats import sigma_clip\nimport os\nfrom numpy.polynomial import Polynomial\nimport torch\nimport torch.nn.functional as F\n\n\nROOT = \"/kaggle/input/ariel-data-challenge-2025/\"\nVERSION = \"v2\"\n\nBINNING = 15\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\n# 16 center pixels, rest is just noise\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\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    # read_frame = pd.read_parquet(\n    #     f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/read.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    # read_frame = read_frame.values.reshape(sensor_sizes_dict[sensor][1])\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        # read_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    return signal / gain + offset\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    # Set values to np.nan where dead or hot pixels are found\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 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    assert signal.shape[0] % binning == 0  # Ensure the binning is possible\n\n    # cds_transposed = signal.transpose(0, 2, 1)\n    cds_binned = torch.zeros(\n        (\n            signal.shape[0] // binning,\n            signal.shape[1],\n            signal.shape[2],\n        ),\n        device=\"cuda:0\",\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 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 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\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 [\"FGS1\", \"AIRS-CH0\"]:\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 = signal.clip(0)\n        \n\n        # Step 3: linearity Correction\n        signal = apply_linear_corr(\n            torch.tensor(linear_corr).to(\"cuda:0\"), torch.tensor(signal).to(\"cuda:0\")\n        )\n\n        # Step 4: dark current subtraction\n        if sensor == \"FGS1\":\n            dt = torch.ones(len(signal), device=\"cuda:0\") * 0.1\n        elif sensor == \"AIRS-CH0\":\n            dt = torch.tensor(dt_airs).to(\"cuda:0\")\n\n        dt[1::2] += 0.1\n\n        signal = clean_dark(signal, torch.tensor(dark_frame).to(\"cuda:0\"), dt)\n\n        # Step 5: Get Correlated Double Sampling (CDS)\n        signal = get_cds(signal)\n\n        # Step 5.1: Bin Observations\n        if sensor == \"FGS1\":\n            signal = bin_obs(signal, binning=BINNING * 12)\n        elif sensor == \"AIRS-CH0\":\n            signal = bin_obs(signal, binning=BINNING)\n\n        # Step 6: Flat Field Correction\n        signal = correct_flat_field(torch.tensor(flat_frame).to(\"cuda:0\"), 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            f\"{planet_id}_{sensor}_signal_{VERSION}.npz\",\n            signal.astype(np.float64),\n        )\n\n\nstar_info = pd.read_csv(ROOT + f\"/{MODE}_star_info.csv\", index_col=\"planet_id\")\nplanet_ids = star_info.index.tolist()\nsample_submission = pd.read_csv(ROOT + \"/sample_submission.csv\", index_col=\"planet_id\")\ntest_ids = sample_submission.index.tolist()\n\n    # Use up to 4 threads!\nwith ThreadPool(4) as pool:\n    list(tqdm(pool.imap(process_planet, planet_ids), total=len(planet_ids)))\n\nMODE = \"test\"\nwith ThreadPool(4) as pool:\n    list(tqdm(pool.imap(process_planet, test_ids), total=len(test_ids)))\n\n\nsignal_train = []\nfor i, planet_id in tqdm(enumerate(planet_ids)):\n    f_raw = np.load(f\"{planet_id}_FGS1_signal_{VERSION}.npz.npy\")\n    a_raw = np.load(f\"{planet_id}_AIRS-CH0_signal_{VERSION}.npz.npy\")\n\n    signal = np.concatenate([f_raw[:, None], a_raw[:, ::-1]], axis=1)\n    signal_train.append(signal)\nsignal_train = np.array(signal_train)\n\n\ntest_train = []\nfor i, planet_id in tqdm(enumerate(test_ids)):\n    f_raw = np.load(f\"{planet_id}_FGS1_signal_{VERSION}.npz.npy\")\n    a_raw = np.load(f\"{planet_id}_AIRS-CH0_signal_{VERSION}.npz.npy\")\n\n    signal = np.concatenate([f_raw[:, None], a_raw[:, ::-1]], axis=1)\n    test_train.append(signal)\n\n\n\ntest_train = np.array(test_train)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:43:21.759223Z","iopub.execute_input":"2025-09-01T21:43:21.759388Z","iopub.status.idle":"2025-09-01T21:48:18.562519Z","shell.execute_reply.started":"2025-09-01T21:43:21.759372Z","shell.execute_reply":"2025-09-01T21:48:18.561929Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Transit Model ","metadata":{}},{"cell_type":"code","source":"from scipy.ndimage import uniform_filter1d\nfrom sklearn.decomposition import PCA\nimport pymc as pm\nimport pytensor.tensor as pt\nfrom sklearn.gaussian_process import GaussianProcessRegressor\nfrom sklearn.gaussian_process.kernels import Matern, WhiteKernel, Product, RationalQuadratic, RBF\n\n\ndef detrend(average_wl):\n    t = np.arange(average_wl.shape[0])\n    linear.fit(t[:, None], average_wl)\n    return linear.coef_\n\ndef egress_and_ingress_finder(signal, star_data):\n    feature_engineering = []\n    wl_spectrums = []\n    for idx, i in enumerate(tqdm(signal)):\n        averaged_signal = np.nanmean(i, axis=1)\n        averaged_signal /= np.nanmean(averaged_signal)\n        \n        kernel = Matern(length_scale=0.05, nu=1.5) + WhiteKernel(noise_level=1e-4)        \n        gpr = GaussianProcessRegressor(kernel=kernel, random_state=42)\n        gpr.fit(np.arange(len(averaged_signal))[:, None], averaged_signal)\n        gaussian_smooth = gpr.predict(np.arange(len(averaged_signal))[:, None])\n        \n        wl_average = uniform_filter1d(gaussian_smooth, size=20)\n        smoothed = savgol_filter(wl_average, 30, 2)\n        first_deriv = np.gradient(smoothed)\n        second_deriv = savgol_filter(np.gradient(first_deriv), 30, 2)\n\n        smoothed_deriv = savgol_filter(first_deriv, 30, 2)\n        ingress_d = np.argmin(smoothed_deriv)\n        egress_d = np.argmax(smoothed_deriv)\n        rough_middle = int((ingress_d + egress_d)/2)\n        \n        fall, rise = second_deriv[:rough_middle], second_deriv[rough_middle:]\n        \n        transition_1_start = np.argmin(fall) + 1\n        transition_1_end = np.argmax(fall)\n        transition_2_start = np.argmax(rise) + rough_middle\n        transition_2_end = rough_middle + np.argmin(rise) - 1\n                \n        t_left = i[:transition_1_start]\n        t_transit = i[transition_1_end:transition_2_start]\n        t_right = i[transition_2_end:]\n        ingress = i[transition_1_start:transition_1_end]\n        egress = i[transition_2_start:transition_2_end]\n\n        trend1 = float(detrend(smoothed[:transition_1_start]))\n        trend2 = float(detrend(smoothed[transition_1_end:transition_2_start]))\n        trend3 = float(detrend(smoothed[transition_2_end:]))\n\n        wierd_shape = False\n        if not transition_1_start < ingress_d < transition_1_end or ingress.shape[0] == 0:\n            \n            # Recalculate areas\n            ingress_d = np.argmin(smoothed_deriv[:180])\n            egress_d = np.argmax(smoothed_deriv[180:]) + 180\n            \n            transition_1_start = 1\n            transition_1_end = ingress_d + 10\n\n            rough_middle = 180\n            fall, rise = second_deriv[:rough_middle], second_deriv[rough_middle:]\n\n            transition_2_start = np.argmax(rise) + rough_middle\n            transition_2_end = rough_middle + np.argmin(rise) - 1\n            \n            t_left = i[:transition_1_start]\n            t_transit = i[transition_1_end:transition_2_start]\n            t_right = i[transition_2_end:]\n            egress = i[transition_2_start:transition_2_end]\n            \n            if transition_2_start < egress_d < transition_2_end and not egress.shape[0] == 0:\n                # Possibility 1: Ingress is not visible but egress is, in which case, flip the egress and make it this egress\n                ingress = np.flip(egress)\n                trend3 = float(detrend(smoothed[transition_2_end:]))\n                trend1 = -1 * trend3\n                \n            else:\n                # Possibility 2: Ingress and Egress are both not visible, in which case, just use best approximation\n                ingress = i[transition_1_start:transition_1_end]\n                trend1 = float(detrend(smoothed[:transition_1_start]))\n            \n\n            \n            print(\"No ingress detected\")\n            wierd_shape = True\n\n        if not transition_2_start < egress_d < transition_2_end or egress.shape[0] == 0:\n            # Recalculate areas\n\n            transition_2_start = egress_d - 10\n            transition_2_end = i.shape[0] - 1\n                    \n            t_left = i[:transition_1_start]\n            t_transit = i[transition_1_end:transition_2_start]\n            t_right = i[transition_2_end:]\n            ingress = i[transition_1_start:transition_1_end]\n            if not wierd_shape:\n                # Possibility 3: Egress not fully visible but ingress is, in which case, substitute egress for flipped engress\n                egress = np.flip(ingress)\n                trend3 = -1 * trend1\n            else:\n                # Posibility 2 again: Just use best approximation\n                egress = i[transition_2_start:transition_2_end]\n                trend3 = float(detrend(smoothed[transition_2_end:]))\n            print(\"No egress detected\")\n            wierd_shape = True\n\n        if wierd_shape:\n            print(ingress_d, egress_d)\n            plt.clf()\n            plt.plot(averaged_signal, color=\"green\", label=\"Noisy\")\n            plt.plot(gaussian_smooth, color=\"purple\", label=\"Gaussian\")\n            plt.plot(smoothed, color=\"red\", label=\"Savgol\")\n            for idx_ in [transition_1_start, transition_2_end, transition_1_end, transition_2_start]:\n                plt.axvline(x=idx_, color=\"red\")\n\n            plt.axvspan(transition_1_start, transition_1_end, color=\"blue\", alpha=0.2)\n            plt.axvspan(transition_2_start, transition_2_end, color=\"blue\", alpha=0.2)\n            plt.legend()\n            plt.show()\n            plt.plot(second_deriv)\n            plt.show()\n            \n        if not wierd_shape:\n            trend1 = float(detrend(smoothed[:transition_1_start]))\n            trend2 = float(detrend(smoothed[transition_1_end:transition_2_start]))\n            trend3 = float(detrend(smoothed[transition_2_end:]))\n\n        \n        \n        not_transit = np.concatenate((t_left, t_right), axis=0)\n        spectrum = (np.nanmean(not_transit, axis=0) - np.nanmean(t_transit, axis=0))/np.nanmean(not_transit)\n\n        entire_curve_shape_2 = np.polyfit(np.arange(len(smoothed)), smoothed, 2)\n        entire_curve_shape_4 = np.polyfit(np.arange(len(smoothed)), smoothed, 4)\n        entire_curve_shape_6 = np.polyfit(np.arange(len(smoothed)), smoothed, 6)\n\n        avg_transmission_depth = np.nanmean(spectrum)\n        SNR = avg_transmission_depth/np.std(np.nanmean(not_transit, axis=1))\n                \n        features = [\n            ingress_d,\n            egress_d,\n            trend1,\n            trend2,\n            trend3,\n            avg_transmission_depth,\n            SNR,\n            transition_1_start,\n            transition_1_end,\n            transition_2_start,\n            transition_2_end,\n            transition_2_end - transition_2_start,\n            transition_1_end - transition_1_start,\n            np.std(np.mean(t_left, axis=1)),\n            np.std(np.mean(t_left, axis=0)),\n            np.std(np.mean(t_right, axis=1)),\n            np.std(np.mean(t_right, axis=0)),\n            np.std(np.mean(t_transit, axis=1)),\n            np.std(np.mean(t_transit, axis=0)),\n            np.std(np.mean(not_transit, axis=0)),\n            np.std(np.mean(not_transit, axis=1)),\n            transition_2_start - transition_1_end,\n            t_left.shape[0],\n            t_right.shape[0],\n            np.max(np.mean(ingress, axis=0)),\n            np.max(np.mean(egress, axis=0)),\n            np.min(np.mean(ingress, axis=0)),\n            np.min(np.mean(egress, axis=0)),\n            np.max(np.mean(ingress, axis=1)),\n            np.max(np.mean(egress, axis=1)),\n            np.min(np.mean(ingress, axis=1)),\n            np.min(np.mean(egress, axis=1)),\n        ]\n        features.extend(entire_curve_shape_2)\n        features.extend(entire_curve_shape_4)\n        features.extend(entire_curve_shape_6)\n        \n        feature_engineering.append(features)\n        wl_spectrums.append(spectrum)\n    return np.concatenate((wl_spectrums, feature_engineering, star_data), axis=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:51:32.218723Z","iopub.execute_input":"2025-09-01T21:51:32.219019Z","iopub.status.idle":"2025-09-01T21:51:32.239855Z","shell.execute_reply.started":"2025-09-01T21:51:32.219Z","shell.execute_reply":"2025-09-01T21:51:32.238881Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Testing Function","metadata":{}},{"cell_type":"code","source":"from sklearn.metrics import mean_absolute_error\nfrom sklearn.model_selection import KFold\nfrom sklearn.preprocessing import StandardScaler\n\nsc = StandardScaler()\nkf = KFold()\n\ndef display_test_results(train_x, train_y):\n    kernel = Product(Matern(nu=2.5) + RBF(), RationalQuadratic(alpha_bounds=(1e-20, 1e2))) \\\n             + WhiteKernel(noise_level_bounds=(1e-20, 1e2))\n\n    kf = KFold(n_splits=5, shuffle=True, random_state=42)\n\n    if False:\n        errors = []\n        for fold, (train_idx, val_idx) in enumerate(kf.split(train_x)):\n            print(f\"fold {fold}\" )\n            x_train, x_val = train_x[train_idx], train_x[val_idx]\n            y_train, y_val = train_y[train_idx], train_y[val_idx]\n    \n            x_train = sc.fit_transform(x_train)\n            x_val = sc.transform(x_val)\n    \n            gpr = GaussianProcessRegressor(\n                random_state=42,\n                normalize_y=True,\n                kernel=kernel,\n                n_restarts_optimizer=20\n            )\n    \n            gpr.fit(x_train, y_train)\n            pred = gpr.predict(x_val)\n    \n            errors.append(mean_absolute_error(y_val, pred))\n    \n        print(\"CV MAE: \", np.mean(errors))\n\n    # Final fit on all data with scaling\n    train_x = sc.fit_transform(train_x)\n    gpr = GaussianProcessRegressor(\n        random_state=42,\n        normalize_y=True,\n        kernel=kernel,\n        n_restarts_optimizer=20\n    )\n    return gpr.fit(train_x, train_y)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:48:23.593392Z","iopub.execute_input":"2025-09-01T21:48:23.59389Z","iopub.status.idle":"2025-09-01T21:48:23.600558Z","shell.execute_reply.started":"2025-09-01T21:48:23.59387Z","shell.execute_reply":"2025-09-01T21:48:23.59991Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Run Model","metadata":{}},{"cell_type":"code","source":"import warnings\n\nwarnings.filterwarnings(\"ignore\")\n\n\ntest_star_info = pd.read_csv(ROOT + f\"/{MODE}_star_info.csv\", index_col=\"planet_id\")\nsubmission_df = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/sample_submission.csv\", index_col=\"planet_id\")\ntrain_df = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/train.csv\", index_col=\"planet_id\")\nstar_feats = star_info.values\ntest_star_feats = test_star_info.values\n\ntrain_x = egress_and_ingress_finder(signal_train.copy(), star_feats)\ntrain_y = np.array(train_df.values)\n\ngpr = display_test_results(train_x, train_y)\n\ntest_x = egress_and_ingress_finder(test_train, test_star_feats)\ntest_x = sc.transform(test_x)\n\ny_pred, sigma_pred = gpr.predict(test_x, return_std=True)\nsubmission = np.concatenate((y_pred, sigma_pred), axis=1).reshape(submission_df.shape)\nsubmission = np.nan_to_num(submission, 0.0)\nnew_submit_df = pd.DataFrame(submission, columns=submission_df.columns, index=submission_df.index)\nnew_submit_df.to_csv(\"submission.csv\")\nnew_submit_df.head(5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-01T21:51:37.395288Z","iopub.execute_input":"2025-09-01T21:51:37.395886Z","iopub.status.idle":"2025-09-01T21:52:31.313044Z","shell.execute_reply.started":"2025-09-01T21:51:37.395848Z","shell.execute_reply":"2025-09-01T21:52:31.312034Z"}},"outputs":[],"execution_count":null}]}