{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"isSourceIdPinned":false,"sourceType":"competition"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport seaborn as sns\nimport scipy.stats\nimport polars as pl\nfrom tqdm import tqdm\nimport pickle\nimport gc\n\n\nimport itertools\nimport torch\nfrom time import time\nfrom scipy.optimize import minimize\nfrom functools import partial\nimport random, os\nfrom astropy.stats import sigma_clip\nfrom scipy.signal import savgol_filter\n\nfrom catboost import CatBoostRegressor","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/test_star_info.csv')\ntrain_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train_star_info.csv')\nadc_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/adc_info.csv')\naxis_info = pd.read_parquet('/kaggle/input/ariel-data-challenge-2025/axis_info.parquet')\nDEVICE = \"cuda\"\nbinn = 10\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def apply_linear_corr(linear_corr, clean_signal):\n    linear_corr = torch.flip(linear_corr, dims=[0])\n\n    x_range = torch.arange(clean_signal.shape[1], device=clean_signal.device)\n    y_range = torch.arange(clean_signal.shape[2], device=clean_signal.device)\n    xx, yy = torch.meshgrid(x_range, y_range, indexing=\"ij\")\n\n    clean_signal = clean_signal.cuda()\n    result_signal = torch.zeros_like(clean_signal)\n\n    for i in range(linear_corr.shape[0]):\n        result_signal += linear_corr[i, xx, yy] * clean_signal ** (linear_corr.shape[0] - 1 - i)\n\n    return result_signal\n\ndef clean_dark(signal, dark, dt):\n    dark_expanded = dark.unsqueeze(0)\n    dt_expanded = dt.unsqueeze(1).unsqueeze(2)\n\n    signal = signal - dark_expanded * dt_expanded\n\n    return signal\n\n\ndef preproc(dataset, star_info, adc_info, sensor, binning = 15):\n    cut_inf, cut_sup = 39, 321\n    sensor_sizes_dict = {\"AIRS-CH0\":[[11250, 32, 356], [1, 32, cut_sup-cut_inf]], \"FGS1\":[[135000, 32, 32], [1, 32, 32]]}\n    binned_dict = {\"AIRS-CH0\":[11250 // binning // 2, 282], \"FGS1\":[135000 // binning // 2]}\n    linear_corr_dict = {\"AIRS-CH0\":(6, 32, 356), \"FGS1\":(6, 32, 32)}\n    planet_ids = star_info['planet_id'].astype(int).tolist()\n    DEVICE = \"cuda:0\"\n\n    feats = []\n    feats_center = []\n    for i, planet_id in tqdm(list(enumerate(planet_ids)), desc=f\"Processing {sensor}\"):\n        signal = torch.Tensor(pl.read_parquet(f'/kaggle/input/ariel-data-challenge-2025/{dataset}/{planet_id}/{sensor}_signal_0.parquet').to_numpy().astype(np.float32)).to(DEVICE)\n        dark_frame = pl.read_parquet(f'/kaggle/input/ariel-data-challenge-2025/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration_0/dark.parquet').to_numpy()\n        dead_frame = pl.read_parquet(f'/kaggle/input/ariel-data-challenge-2025/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration_0/dead.parquet').to_numpy()\n        flat_frame = pl.read_parquet(f'/kaggle/input/ariel-data-challenge-2025/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration_0/flat.parquet').to_numpy()\n        linear_corr = torch.Tensor(pl.read_parquet(f'/kaggle/input/ariel-data-challenge-2025/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration_0/linear_corr.parquet').to_numpy().astype(np.float32).reshape(linear_corr_dict[sensor])).to(DEVICE)\n\n        signal = signal.reshape(sensor_sizes_dict[sensor][0]) \n        gain = adc_info[f'{sensor}_adc_gain'].values[0]\n        offset = adc_info[f'{sensor}_adc_offset'].values[0]\n        signal = signal / gain + offset\n\n        hot = sigma_clip(\n            dark_frame, sigma=5, maxiters=5\n        ).mask\n        dark_frame = torch.Tensor(dark_frame).to(DEVICE)\n\n        if sensor == \"AIRS-CH0\":\n            signal = signal[:, :, cut_inf:cut_sup] #11250 * 32 * 282\n            #dt = axis_info['AIRS-CH0-integration_time'].dropna().values\n            dt = torch.ones(len(signal)).to(DEVICE)*0.1 \n            dt[1::2] += 4.5 #@bilzard idea\n            linear_corr = linear_corr[:, :, cut_inf:cut_sup]\n            dark_frame = dark_frame[:, cut_inf:cut_sup]\n            dead_frame = dead_frame[:, cut_inf:cut_sup]\n            flat_frame = flat_frame[:, cut_inf:cut_sup]\n            hot = hot[:, cut_inf:cut_sup]\n        else:\n            dt = torch.ones(len(signal)).to(DEVICE)*0.1\n            dt[1::2] += 0.1\n\n        signal = signal.clip(0) #@graySnow idea\n        linear_corr_signal = apply_linear_corr(linear_corr, signal)\n        signal = clean_dark(linear_corr_signal, dark_frame, dt).cpu().numpy()\n\n        flat = flat_frame.reshape(sensor_sizes_dict[sensor][1])\n        flat[dead_frame.reshape(sensor_sizes_dict[sensor][1])] = np.nan\n        flat[hot.reshape(sensor_sizes_dict[sensor][1])] = np.nan\n        signal = signal / flat\n        \n        if sensor == \"AIRS-CH0\":\n            signal_center = signal[:, 10:22, :]\n        elif sensor == \"FGS1\":\n            signal_center = signal[:, 10:22, 10:22]\n            signal_center = signal_center.reshape(\n                signal_center.shape[0], signal_center.shape[1] * signal_center.shape[2]\n            )\n            \n        if sensor == \"FGS1\":\n            signal = signal.reshape((sensor_sizes_dict[sensor][0][0], sensor_sizes_dict[sensor][0][1]*sensor_sizes_dict[sensor][0][2]))\n\n        mean_signal = np.nanmean(signal, axis=1) # mean over the 32*32(FGS1) or 32(CH0) pixels\n        cds_signal = (mean_signal[1::2] - mean_signal[0::2])\n        binned = np.zeros((binned_dict[sensor]))\n        for j in range(cds_signal.shape[0] // binning):\n            binned[j] = cds_signal[j*binning:j*binning+binning].mean(axis=0)\n        if sensor == \"FGS1\":\n            binned = binned.reshape((binned.shape[0],1))\n        feats.append(binned)\n        \n        \n        mean_signal = np.nanmean(signal_center, axis=1) # mean over the 32*32(FGS1) or 32(CH0) pixels\n        cds_signal = (mean_signal[1::2] - mean_signal[0::2])\n\n        binned = np.zeros((binned_dict[sensor]))\n        for j in range(cds_signal.shape[0] // binning):\n            binned[j] = cds_signal[j*binning:j*binning+binning].mean(axis=0)\n       \n        if sensor == \"FGS1\":\n            binned = binned.reshape((binned.shape[0],1))\n\n        feats_center.append(binned)\n        \n        \n\n    return np.stack(feats), np.stack(feats_center)\n\nfgs1_feats, fgs1_feats_center = preproc('train', train_star_info, adc_info, \"FGS1\", binn*12)\nairsch0_feats, airsch0_feats_center = preproc('train', train_star_info, adc_info, \"AIRS-CH0\", binn)\n\n\ntrain_binn = np.concatenate([fgs1_feats_center, airsch0_feats_center], axis=2)\ntrain_binn_not_centered = np.concatenate([fgs1_feats, airsch0_feats], axis=2)\n\npickle.dump(train_binn_not_centered, open('train_binn_not_centered.pkl', 'wb'))\n\ndel fgs1_feats\ndel fgs1_feats_center\ndel airsch0_feats\ndel airsch0_feats_center\ndel train_binn_not_centered\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.save('train_binn.npy', train_binn)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_feats, fgs1_feats_center = preproc('test', test_star_info, adc_info, \"FGS1\", binn*12)\nairsch0_feats, airsch0_feats_center = preproc('test', test_star_info, adc_info, \"AIRS-CH0\", binn)\n\n\ntest_binn = np.concatenate([fgs1_feats_center, airsch0_feats_center], axis=2)\ntest_binn_not_centered = np.concatenate([fgs1_feats, airsch0_feats], axis=2)\n\npickle.dump(test_binn_not_centered, open('test_binn_not_centered.pkl', 'wb'))\n\ndel fgs1_feats\ndel fgs1_feats_center\ndel airsch0_feats\ndel airsch0_feats_center\ndel test_binn_not_centered\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.save('test_binn.npy', test_binn)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}