{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"sourceType":"competition"},{"sourceId":12580690,"sourceType":"datasetVersion","datasetId":7925491},{"sourceId":486108,"sourceType":"modelInstanceVersion","modelInstanceId":387768,"modelId":406804}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Necessary imports","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nfrom warnings import filterwarnings\n\nfilterwarnings('ignore')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T11:03:43.796295Z","iopub.execute_input":"2025-07-26T11:03:43.796643Z","iopub.status.idle":"2025-07-26T11:03:46.337981Z","shell.execute_reply.started":"2025-07-26T11:03:43.796612Z","shell.execute_reply":"2025-07-26T11:03:46.337154Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Reading files","metadata":{}},{"cell_type":"markdown","source":"## Getting files' content","metadata":{}},{"cell_type":"code","source":"adc_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/adc_info.csv')\nadc_info","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T05:08:06.408709Z","iopub.execute_input":"2025-07-26T05:08:06.409563Z","iopub.status.idle":"2025-07-26T05:08:06.448734Z","shell.execute_reply.started":"2025-07-26T05:08:06.409522Z","shell.execute_reply":"2025-07-26T05:08:06.447879Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BASE_DIR = '/kaggle/input/ariel-data-challenge-2025/'\n\ntrain_csv = pd.read_csv(f'{BASE_DIR}/train.csv')\ntrain_csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T11:03:51.699004Z","iopub.execute_input":"2025-07-26T11:03:51.699312Z","iopub.status.idle":"2025-07-26T11:03:51.918763Z","shell.execute_reply.started":"2025-07-26T11:03:51.699287Z","shell.execute_reply":"2025-07-26T11:03:51.917845Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_star_info = pd.read_csv(f'{BASE_DIR}/train_star_info.csv')\ntrain_star_info","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:22:27.807903Z","iopub.execute_input":"2025-07-22T07:22:27.808234Z","iopub.status.idle":"2025-07-22T07:22:27.835462Z","shell.execute_reply.started":"2025-07-22T07:22:27.808208Z","shell.execute_reply":"2025-07-22T07:22:27.834598Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_star_info.columns","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:22:31.632917Z","iopub.execute_input":"2025-07-22T07:22:31.633241Z","iopub.status.idle":"2025-07-22T07:22:31.641037Z","shell.execute_reply.started":"2025-07-22T07:22:31.633217Z","shell.execute_reply":"2025-07-22T07:22:31.640113Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_star_info = pd.read_csv(f'{BASE_DIR}/test_star_info.csv')\ntest_star_info","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T05:08:38.062798Z","iopub.execute_input":"2025-07-26T05:08:38.063464Z","iopub.status.idle":"2025-07-26T05:08:38.079149Z","shell.execute_reply.started":"2025-07-26T05:08:38.063431Z","shell.execute_reply":"2025-07-26T05:08:38.078397Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv['planet_id'] == train_star_info['planet_id']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:34:16.16986Z","iopub.execute_input":"2025-07-22T07:34:16.170601Z","iopub.status.idle":"2025-07-22T07:34:16.180451Z","shell.execute_reply.started":"2025-07-22T07:34:16.170574Z","shell.execute_reply":"2025-07-22T07:34:16.179461Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"wl_csv = pd.read_csv(f'{BASE_DIR}/wavelengths.csv')\nwl_csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T08:03:52.844122Z","iopub.execute_input":"2025-07-24T08:03:52.844422Z","iopub.status.idle":"2025-07-24T08:03:52.879236Z","shell.execute_reply.started":"2025-07-24T08:03:52.844399Z","shell.execute_reply":"2025-07-24T08:03:52.878274Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"wl_csv.columns","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T08:04:06.233628Z","iopub.execute_input":"2025-07-24T08:04:06.233986Z","iopub.status.idle":"2025-07-24T08:04:06.246246Z","shell.execute_reply.started":"2025-07-24T08:04:06.233962Z","shell.execute_reply":"2025-07-24T08:04:06.245478Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"axis_info = pd.read_parquet(f'{BASE_DIR}/axis_info.parquet')\naxis_info","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:37:05.220559Z","iopub.execute_input":"2025-07-22T07:37:05.220896Z","iopub.status.idle":"2025-07-22T07:37:05.255119Z","shell.execute_reply.started":"2025-07-22T07:37:05.220875Z","shell.execute_reply":"2025-07-22T07:37:05.2542Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sample_submission = pd.read_csv(f'{BASE_DIR}/sample_submission.csv')\nsample_submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T11:04:03.188197Z","iopub.execute_input":"2025-07-26T11:04:03.188517Z","iopub.status.idle":"2025-07-26T11:04:03.223722Z","shell.execute_reply.started":"2025-07-26T11:04:03.188493Z","shell.execute_reply":"2025-07-26T11:04:03.222384Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sample_submission.columns","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T08:04:44.575314Z","iopub.execute_input":"2025-07-24T08:04:44.576007Z","iopub.status.idle":"2025-07-24T08:04:44.581109Z","shell.execute_reply.started":"2025-07-24T08:04:44.575981Z","shell.execute_reply":"2025-07-24T08:04:44.580435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_DIR = BASE_DIR + \"/train\"\n\nairs_0_signal = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_signal_0.parquet')\nairs_0_signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:22:50.402109Z","iopub.execute_input":"2025-07-22T07:22:50.402494Z","iopub.status.idle":"2025-07-22T07:22:53.380357Z","shell.execute_reply.started":"2025-07-22T07:22:50.402463Z","shell.execute_reply":"2025-07-22T07:22:53.379526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_0 = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_signal_0.parquet')\nfgs1_0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:22:57.74152Z","iopub.execute_input":"2025-07-22T07:22:57.742519Z","iopub.status.idle":"2025-07-22T07:22:59.420776Z","shell.execute_reply.started":"2025-07-22T07:22:57.742485Z","shell.execute_reply":"2025-07-22T07:22:59.419934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"airs_ch0_calibration_dark = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_calibration_0/dark.parquet')\nairs_ch0_calibration_dark","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:03.871445Z","iopub.execute_input":"2025-07-22T07:23:03.87175Z","iopub.status.idle":"2025-07-22T07:23:03.931882Z","shell.execute_reply.started":"2025-07-22T07:23:03.871729Z","shell.execute_reply":"2025-07-22T07:23:03.930892Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"airs_ch0_calib_dead = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_calibration_0/dead.parquet')\nairs_ch0_calib_dead","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:09.298832Z","iopub.execute_input":"2025-07-22T07:23:09.299142Z","iopub.status.idle":"2025-07-22T07:23:09.358261Z","shell.execute_reply.started":"2025-07-22T07:23:09.299118Z","shell.execute_reply":"2025-07-22T07:23:09.357437Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"airs_ch0_flat = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_calibration_0/flat.parquet')\nairs_ch0_flat","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:13.380911Z","iopub.execute_input":"2025-07-22T07:23:13.381763Z","iopub.status.idle":"2025-07-22T07:23:13.439191Z","shell.execute_reply.started":"2025-07-22T07:23:13.381733Z","shell.execute_reply":"2025-07-22T07:23:13.438245Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"airs_ch0_linear_corr = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_calibration_0/linear_corr.parquet')\nairs_ch0_linear_corr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:17.481059Z","iopub.execute_input":"2025-07-22T07:23:17.481369Z","iopub.status.idle":"2025-07-22T07:23:17.539226Z","shell.execute_reply.started":"2025-07-22T07:23:17.481346Z","shell.execute_reply":"2025-07-22T07:23:17.538309Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"airs_ch0_calib_read = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_calibration_0/read.parquet')\nairs_ch0_calib_read","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:22.281041Z","iopub.execute_input":"2025-07-22T07:23:22.281306Z","iopub.status.idle":"2025-07-22T07:23:22.337885Z","shell.execute_reply.started":"2025-07-22T07:23:22.281287Z","shell.execute_reply":"2025-07-22T07:23:22.336952Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_calibration_dark = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_calibration_0/dark.parquet')\nfgs1_calibration_dark","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:27.122779Z","iopub.execute_input":"2025-07-22T07:23:27.123106Z","iopub.status.idle":"2025-07-22T07:23:27.162804Z","shell.execute_reply.started":"2025-07-22T07:23:27.123083Z","shell.execute_reply":"2025-07-22T07:23:27.161942Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_calibration_dead = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_calibration_0/dead.parquet')\nfgs1_calibration_dead","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:38.38843Z","iopub.execute_input":"2025-07-22T07:23:38.388695Z","iopub.status.idle":"2025-07-22T07:23:38.423818Z","shell.execute_reply.started":"2025-07-22T07:23:38.388676Z","shell.execute_reply":"2025-07-22T07:23:38.423029Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_calibration_flat = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_calibration_0/flat.parquet')\nfgs1_calibration_flat","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:43.31592Z","iopub.execute_input":"2025-07-22T07:23:43.316191Z","iopub.status.idle":"2025-07-22T07:23:43.358074Z","shell.execute_reply.started":"2025-07-22T07:23:43.316171Z","shell.execute_reply":"2025-07-22T07:23:43.357209Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_calibration_0_lin_corr = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_calibration_0/linear_corr.parquet')\nfgs1_calibration_0_lin_corr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:47.5438Z","iopub.execute_input":"2025-07-22T07:23:47.544473Z","iopub.status.idle":"2025-07-22T07:23:47.58306Z","shell.execute_reply.started":"2025-07-22T07:23:47.54444Z","shell.execute_reply":"2025-07-22T07:23:47.582139Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_calib0_read = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_calibration_0/read.parquet')\nfgs1_calib0_read","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:51.167912Z","iopub.execute_input":"2025-07-22T07:23:51.168624Z","iopub.status.idle":"2025-07-22T07:23:51.209264Z","shell.execute_reply.started":"2025-07-22T07:23:51.168598Z","shell.execute_reply":"2025-07-22T07:23:51.208216Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"airs_ch0_signal_0 = pd.read_parquet(f'{TRAIN_DIR}/1010375142/AIRS-CH0_signal_0.parquet')\nairs_ch0_signal_0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:23:59.101272Z","iopub.execute_input":"2025-07-22T07:23:59.101587Z","iopub.status.idle":"2025-07-22T07:24:00.174625Z","shell.execute_reply.started":"2025-07-22T07:23:59.101564Z","shell.execute_reply":"2025-07-22T07:24:00.173786Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fgs1_signal_0 = pd.read_parquet(f'{TRAIN_DIR}/1010375142/FGS1_signal_0.parquet')\nfgs1_signal_0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-22T07:24:03.819647Z","iopub.execute_input":"2025-07-22T07:24:03.819965Z","iopub.status.idle":"2025-07-22T07:24:04.311835Z","shell.execute_reply.started":"2025-07-22T07:24:03.819934Z","shell.execute_reply":"2025-07-22T07:24:04.310952Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Creating a regression based Neural network","metadata":{}},{"cell_type":"markdown","source":"## Creating NN class","metadata":{}},{"cell_type":"code","source":"import torch\nfrom torch import nn\nfrom torch.optim import Adam\n\nNUM_WAVELENGTH_OUTPUT_CHANNELS = 283 \n\n# class SpectrumRegressor(torch.nn.Module):\n#     def __init__(self, input_dim):\n#         super().__init__()\n#         # The network should map the 3 input features to a large output vector\n#         # that is 2 * 283 = 566 elements long.\n#         self.net = torch.nn.Sequential(\n#             torch.nn.Linear(input_dim, 256),\n#             torch.nn.ReLU(),\n#             torch.nn.Linear(256, 512),\n#             torch.nn.ReLU(),\n#             torch.nn.Linear(512, NUM_WAVELENGTH_OUTPUT_CHANNELS * 2) # <-- This is the critical line\n#         )\n\n#     def forward(self, x):\n#         out = self.net(x)\n#         # out has shape (batch_size, 566)\n        \n#         # Split the output into two equal chunks for mu and sigma\n#         mu = out[..., :NUM_WAVELENGTH_OUTPUT_CHANNELS]\n#         sigma = torch.nn.functional.softplus(out[..., NUM_WAVELENGTH_OUTPUT_CHANNELS:])\n#         return mu, sigma\n\nclass SpectrumRegressor(torch.nn.Module):\n    def __init__(self, input_dim, output_dim): # Add output_dim as a parameter\n        super().__init__()\n        \n        # The network should map input_dim features to output_dim * 2 elements (mu and sigma)\n        self.net = torch.nn.Sequential(\n            torch.nn.Linear(input_dim, 256), # First layer takes input_dim (e.g., 364)\n            torch.nn.ReLU(),\n            torch.nn.Linear(256, 512),\n            torch.nn.ReLU(),\n            torch.nn.Linear(512, output_dim * 2) # Last layer produces output_dim * 2 (e.g., 283 * 2 = 566)\n        )\n        self.output_dim = output_dim # Store it for the forward pass\n\n    def forward(self, x):\n        out = self.net(x)\n        \n        # Split the output into two equal chunks for mu and sigma\n        mu = out[..., :self.output_dim]\n        # Ensure sigma has a minimum value to prevent numerical instability in loss\n        sigma_raw = out[..., self.output_dim:]\n        sigma = torch.nn.functional.softplus(sigma_raw).clamp(min=1e-6) \n        \n        return mu, sigma","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:36:33.654214Z","iopub.execute_input":"2025-07-26T12:36:33.654468Z","iopub.status.idle":"2025-07-26T12:36:37.488151Z","shell.execute_reply.started":"2025-07-26T12:36:33.654445Z","shell.execute_reply":"2025-07-26T12:36:37.487355Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Required preprocessing of data","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom torch import tensor\nimport torch\nimport os\nfrom warnings import filterwarnings\n\nfilterwarnings('ignore')\n\n# --- 0. Configuration Constants ---\n# These constants define the dimensions based on data formats\nAIRS_CH0_SPECTRAL_DIM = 356 \nAIRS_CH0_SPATIAL_DIM = 32\n\n# Number of features from train_star_info.csv\nNUM_STELLAR_FEATURES = 8 # ('Rs', 'Ms', 'Ts', 'Mp', 'e', 'P', 'sma', 'i')\n\n# Number of features derived from axis_info.parquet\n# We'll use global means of: AIRS-CH0-axis0-h, AIRS-CH0-integration_time, FGS1-axis0-h\nNUM_AXIS_INFO_FEATURES = 3\n\n# Total input features for the model = (features from AIRS-CH0 signal) + (stellar features) + (axis info features)\n# Here, AIRS-CH0 signal features are `AIRS_CH0_SPECTRAL_DIM` (356)\n# TOTAL_MODEL_INPUT_DIM = AIRS_CH0_SPECTRAL_DIM + NUM_STELLAR_FEATURES + NUM_AXIS_INFO_FEATURES # 356 + 3 + 3 = 362\nTOTAL_MODEL_INPUT_DIM = NUM_STELLAR_FEATURES  # The number of features we are using\n\nOUTPUT_DIR = '/kaggle/working/'\nDATA_ROOT = \"/kaggle/input/ariel-data-challenge-2025\"\nTRAIN_DIR = f\"{DATA_ROOT}/train\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:37:42.218596Z","iopub.execute_input":"2025-07-26T12:37:42.219141Z","iopub.status.idle":"2025-07-26T12:37:42.477292Z","shell.execute_reply.started":"2025-07-26T12:37:42.219117Z","shell.execute_reply":"2025-07-26T12:37:42.476755Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_calibration(calib_dir, gain_val, offset_val):\n    calib_data = {}\n    \n    # Iterate through files in the calibration directory\n    for fname in os.listdir(calib_dir):\n        file_path = os.path.join(calib_dir, fname)\n        data = pd.read_parquet(file_path).values\n\n        # Determining the key for the dictionary entry (e.g.,\n        # 'dark', 'flat', 'dead', 'linear_corr'). Being a bit\n        # more robust here, taking the part before the first '.'\n        # and then splitting by '_' just in case (e.g.,\n        # 'linear_corr.parquet' -> 'linear_corr')\n        key = fname.split('.')[0] \n        # Specific reshaping logic for common calibration files like dark,\n        #  flat, dead. These should typically be (spatial_dim, spectral_dim)\n        if data.ndim == 1 and data.size == AIRS_CH0_SPATIAL_DIM * AIRS_CH0_SPECTRAL_DIM:\n            data = data.reshape(AIRS_CH0_SPATIAL_DIM, AIRS_CH0_SPECTRAL_DIM)\n        \n        # For 'dead' mask, ensure it's boolean\n        if key == 'dead':\n            calib_data[key] = data.astype(bool)\n        else:\n            calib_data[key] = data\n\n    # Explicitly add the gain and offset values provided as arguments\n    calib_data['gain'] = gain_val\n    calib_data['offset'] = offset_val\n    \n    return calib_data","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:37:47.739769Z","iopub.execute_input":"2025-07-26T12:37:47.740538Z","iopub.status.idle":"2025-07-26T12:37:47.74557Z","shell.execute_reply.started":"2025-07-26T12:37:47.740514Z","shell.execute_reply":"2025-07-26T12:37:47.744869Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def correct_signal(raw, calib):\n    '''\n    Takes raw 3D signal and 2D calibration data, corrects it, and averages it.\n    Input `raw` shape: (num_timesteps, 32, 356)\n    Output shape: (356,)\n    '''\n    # Applying a linear scaling to the raw sensor data using gain and offset.\n    # This works because NumPy correctly broadcasts the scalar values.\n    signal = raw * calib['gain'] + calib['offset']\n\n    # Subtracting the dark frame.\n    # NumPy correctly broadcasts the 2D calib['dark'] array to subtract\n    # it from each time-step slice of the 3D 'signal' array.\n    signal -= calib['dark']\n\n    # Dividing by the flat field.\n    # Again, NumPy correctly broadcasts the 2D calib['flat'] array.\n    signal /= calib['flat']\n\n    # Creating a mask based on dead pixels.\n    # NumPy correctly broadcasts the 2D calib['dead'] mask.\n    mask = calib['dead']\n    signal = np.where(mask, np.nan, signal)\n\n    # Now, after all the corrections have been applied,\n    # perform the averaging over the time and spatial axes.\n    corrected_and_averaged_signal = np.mean(signal, axis=(0, 1))\n\n    # Replace any nan or inf values in the final, averaged result.\n    return np.nan_to_num(corrected_and_averaged_signal)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:37:49.07399Z","iopub.execute_input":"2025-07-26T12:37:49.074222Z","iopub.status.idle":"2025-07-26T12:37:49.078994Z","shell.execute_reply.started":"2025-07-26T12:37:49.074204Z","shell.execute_reply":"2025-07-26T12:37:49.078237Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\ndef load_planet_instrument_data(planet_root_dir, inst, adc_info_df):\n    \"\"\"\n    Loads and corrects signal data for a single instrument of a given planet.\n    \"\"\"\n    signal_file_path = os.path.join(planet_root_dir, f\"{inst}_signal_0.parquet\")\n    raw_signal_flat = pd.read_parquet(signal_file_path).values\n\n    calib_dir_path = os.path.join(planet_root_dir, f\"{inst}_calibration_0\")\n\n    current_inst_gain = adc_info_df[f\"{inst}_adc_gain\"].iloc[0]\n    current_inst_offset = adc_info_df[f\"{inst}_adc_offset\"].iloc[0]\n\n    calib = load_calibration(calib_dir_path, current_inst_gain, current_inst_offset)\n\n    if inst == 'AIRS-CH0':\n        num_timesteps = raw_signal_flat.shape[0]\n        AIRS_CH0_SPATIAL_DIM = 32\n        AIRS_CH0_SPECTRAL_DIM = 356\n        raw_signal_3d = raw_signal_flat.reshape(\n            (num_timesteps, AIRS_CH0_SPATIAL_DIM, AIRS_CH0_SPECTRAL_DIM)\n        )\n        \n        # Now, call correct_signal with the full 3D array.\n        # This function will now return the final, averaged (356,) vector.\n        return correct_signal(raw_signal_3d, calib)\n        \n    elif inst == 'FGS1':\n        # ... FGS1 logic ...\n        # Assuming FGS1 calibration is simpler, this part might need its own\n        # averaging within correct_signal.\n        return correct_signal(raw_signal_flat, calib)\n\n    else:\n        raise ValueError(f\"Unknown instrument '{inst}'. Cannot determine expected shape.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:37:50.392505Z","iopub.execute_input":"2025-07-26T12:37:50.392767Z","iopub.status.idle":"2025-07-26T12:37:50.398903Z","shell.execute_reply.started":"2025-07-26T12:37:50.392748Z","shell.execute_reply":"2025-07-26T12:37:50.398222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_ground_truth_spectrum(pid, train_labels_df):\n    \"\"\"\n    Loads the true exoplanet transmission spectrum (ground truth, y) for a given planet ID.\n    \"\"\"\n    if pid not in train_labels_df.index:\n        raise ValueError(f\"Ground truth spectrum for planet {pid} not found in train.csv!\")\n    \n    # Select all columns that represent the spectrum (e.g., 'wl_1' to 'wl_283')\n    # and convert to a NumPy array of float32.\n    spectrum_values = train_labels_df.loc[pid].filter(like='wl_').values\n    return spectrum_values.astype(np.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:37:54.329344Z","iopub.execute_input":"2025-07-26T12:37:54.329605Z","iopub.status.idle":"2025-07-26T12:37:54.333998Z","shell.execute_reply.started":"2025-07-26T12:37:54.329584Z","shell.execute_reply":"2025-07-26T12:37:54.333319Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_all_global_data(data_root):\n    \"\"\"\n    Loads all global CSV/Parquet files (adc_info, train_star_info, wavelengths, train.csv, axis_info.parquet) once.\n    \"\"\"\n    print(\"Loading global data files from:\", data_root)\n    \n    adc_info_df = pd.read_csv(os.path.join(data_root, 'adc_info.csv'))\n    train_star_info_df = pd.read_csv(os.path.join(data_root, 'train_star_info.csv'))\n    wavelengths = pd.read_csv(os.path.join(data_root, 'wavelengths.csv')).iloc[0].values\n    train_labels_df = pd.read_csv(os.path.join(data_root, 'train.csv'))\n\n    train_star_info_df.set_index(['planet_id'], inplace=True)\n    train_labels_df.set_index(['planet_id'], inplace=True)\n    \n    # --- NEW: Load axis_info.parquet ---\n    axis_info_df = pd.read_parquet(os.path.join(data_root, 'axis_info.parquet'))\n    # Fill NaNs as you indicated (using column means)\n    axis_info_df = axis_info_df.fillna(axis_info_df.mean(numeric_only=True))\n\n    # Derive global summary features from axis_info.parquet\n    # We'll use the mean of a few relevant columns as features for every planet.\n    # Note: If axis_info.parquet has specific rows per planet, you'd need\n    # a more sophisticated mapping here. This approach uses global averages.\n    global_axis_features = np.array([\n        axis_info_df['AIRS-CH0-axis0-h'].mean(),\n        axis_info_df['AIRS-CH0-integration_time'].mean(),\n        axis_info_df['FGS1-axis0-h'].mean()\n    ]).astype(np.float32)\n\n    print(f\" - ADC Info Shape: {adc_info_df.shape}\")\n    print(f\" - Star Info Shape: {train_star_info_df.shape}\")\n    print(f\" - Wavelengths Count: {len(wavelengths)}\")\n    print(f\" - Train Labels (Ground Truth) Shape: {train_labels_df.shape}\")\n    print(f\" - Axis Info Shape: {axis_info_df.shape}\")\n    print(f\" - Derived Global Axis Features: {global_axis_features.shape} (e.g., {global_axis_features})\")\n\n    # Return all loaded dataframes and the derived global axis features\n    return adc_info_df, train_star_info_df, wavelengths, train_labels_df, global_axis_features","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:37:56.652003Z","iopub.execute_input":"2025-07-26T12:37:56.652655Z","iopub.status.idle":"2025-07-26T12:37:56.659087Z","shell.execute_reply.started":"2025-07-26T12:37:56.65263Z","shell.execute_reply":"2025-07-26T12:37:56.65819Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# def processing_single_planet(pid, train_dir, global_data, is_train=True):\n#     # This is a dummy function. Replace with your actual processing logic.\n#     # It takes the planet ID as a string and the global data tuple.\n#     adc_info_df, train_star_info_df, wavelengths, train_labels_df, global_axis_features = global_data\n#     planet_root_dir = f\"{train_dir}/{pid}\"\n    \n#     # Check for ground truth and star info before attempting to process\n#     if pid not in train_labels_df.index:\n#         raise ValueError(f\"Ground truth spectrum for planet {pid_int} not found in train_labels.csv!\")\n#     if pid not in train_star_info_df.index:\n#         raise ValueError(f\"Star info for planet {pid} not found in train_star_info.csv!\")\n\n#     # FGS1 is not loaded because it does only provide the stability\n#     # AIRS-CH0 is loaded because it provides the spectrographical\n#     # signals\n#     airs_ch0_data = load_planet_instrument_data(\n#         planet_root_dir, 'AIRS-CH0', adc_info_df\n#     )\n#     # Extract Features from AIRS-CH0 signal\n#     signal_features = airs_ch0_data\n    \n#     star_row = train_star_info_df.loc[pid]\n#     stellar_features = np.array([\n#         star_row['Rs'],\n#         star_row['Ms'],\n#         star_row['Ts'],\n#     ], dtype=np.float32)\n\n#     # Combine all input features\n#     combined_input_features = np.concatenate([signal_features, stellar_features, global_axis_features])\n\n#     ground_truth_spectrum = None\n#     if is_train:\n#         # load_ground_truth_spectrum will raise ValueError if not found\n#         ground_truth_spectrum = load_ground_truth_spectrum(pid, train_labels_df)\n\n#     return combined_input_features, ground_truth_spectrum","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-25T10:50:30.786089Z","iopub.execute_input":"2025-07-25T10:50:30.786395Z","iopub.status.idle":"2025-07-25T10:50:30.792168Z","shell.execute_reply.started":"2025-07-25T10:50:30.786376Z","shell.execute_reply":"2025-07-25T10:50:30.791074Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nfrom pathlib import Path\n# from multiprocessing import Pool, cpu_count\n# import functools\n\n# --- Define a helper function to process a single planet's data ---\n# This function does the work for one planet and can be parallelized.\ndef process_single_planet(pid, TRAIN_DIR, train_star_info_df, adc_info_df):\n    \"\"\"Processes all data for a single planet ID.\"\"\"\n    print(f\"Processing planet: {pid}\")\n    planet_root_dir = f\"{TRAIN_DIR}/{pid}\"\n    \n    # Load and extract signal features\n    airs_ch0_data = load_planet_instrument_data(\n        planet_root_dir, 'AIRS-CH0', adc_info_df\n    )\n    signal_features = airs_ch0_data\n    \n    # Get stellar features\n    star_row = train_star_info_df.loc[pid]\n    stellar_features = np.array([\n        star_row['Rs'], star_row['Ms'], star_row['Ts'], star_row['Mp'],\n        star_row['e'], star_row['P'], star_row['sma'], star_row['i'],\n    ], dtype=np.float32)\n    \n    combined_features = np.concatenate([stellar_features, signal_features])\n\n    return combined_features","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:08.172726Z","iopub.execute_input":"2025-07-26T12:38:08.173216Z","iopub.status.idle":"2025-07-26T12:38:08.178315Z","shell.execute_reply.started":"2025-07-26T12:38:08.173193Z","shell.execute_reply":"2025-07-26T12:38:08.177645Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nfrom pathlib import Path\nimport time\n# import functools\n# from tqdm.auto import tqdm\n\n# train_star_info dataframe's features\n# Index(\n#  ['planet_id', 'Rs', 'Ms', 'Ts', 'Mp', 'e', 'P', 'sma', 'i'],\n#  dtype='object'\n#)\n\ndef consolidate_data():\n    print(\"------------- Starting path allocation -------------\")\n\n    start_path_alloc = time.perf_counter()\n    \n    feature_file = f'{OUTPUT_DIR}combined_input_features.npy'\n    ground_truth_file = f'{OUTPUT_DIR}/ground_truth_spectrum.npy'\n    global_axis_file = f\"{OUTPUT_DIR}/global_axis_features.npy\"\n    # numpy_files_dir = os.path.join(OUTPUT_DIR, \"consolidated_data\")\n    \n    # os.makedirs(numpy_files_dir, exist_ok=True)\n    \n    # feature_dir = os.path.join(numpy_files_dir, feature_file)\n    # ground_truth_dir = os.path.join(numpy_files_dir, ground_truth_file)\n    # global_axis_features_dir = os.path.join(numpy_files_dir, \"global_axis_features.npy\")\n    \n    feature_path = Path(feature_file)\n    ground_truth_path = Path(ground_truth_file)\n    global_axis_features_path = Path(global_axis_file)\n\n    if feature_path.is_file() and ground_truth_path.is_file() and global_axis_features_path.is_file():\n        end_path_alloc = time.perf_counter()\n        print(f\"Time consumed for path allocation: {end_path_alloc - start_path_alloc:.2f}\")\n        \n        print(\"Files of features and labels arrays already exist.\")\n        return feature_file, ground_truth_file, global_axis_features_file\n    else:\n        end_path_alloc = time.perf_counter()\n        print(f\"Time consumed for path allocation: {end_path_alloc - start_path_alloc:.2f}\")\n        \n        print(\"------------- Starting arrays creation -------------\")\n        start_arr_creation = time.perf_counter()\n        \n        print(\"------------- Loading global data -------------\")\n        start_load_global = time.perf_counter()\n        try:\n            adc_info_df, train_star_info_df, _, train_labels_df, global_axis_features = load_all_global_data(DATA_ROOT)\n        except FileNotFoundError as e:\n            print(f\"\\nERROR: Required data file not found: {e}\")\n            return\n        end_load_global = time.perf_counter()\n        print(f\"Global data loaded in {end_load_global - start_load_global:.2f} seconds.\")\n        \n        dir_ids_set = set(\n            [int(d) for d in os.listdir(TRAIN_DIR) if d.isdigit()]\n        )\n        labels_ids_set = set(train_labels_df.index.to_list())\n        star_info_ids_set = set(train_star_info_df.index.to_list())\n        training_pids_set = dir_ids_set \\\n                            .intersection(labels_ids_set) \\\n                            .intersection(star_info_ids_set)\n        training_planet_ids = sorted(list(training_pids_set))\n\n        if not training_planet_ids:\n            print(\"\\nERROR: No valid data for training.\")\n            return\n\n        print(f\"------------- Found {len(training_planet_ids)} valid planets. Consolidating data -------------\")\n\n        features = []\n        labels = []\n\n        i = 0\n        start_training_time = time.perf_counter()\n        for pid in training_planet_ids:\n            print(f\"Getting planet id at {i} index\")\n            combined_features = process_single_planet(\n                pid,\n                TRAIN_DIR,\n                train_star_info_df,\n                adc_info_df\n            )\n            features.append(combined_features)\n            print(f\"Getting the truth data of {pid}\")\n            true_label = train_labels_df.loc[pid].values.astype(np.float32)\n            labels.append(true_label)\n            i += 1\n        end_training_time = time.perf_counter()\n        print(f\"Time taken for training {len(training_planet_ids)} planets: {end_training_time - start_training_time:.2f}\")\n\n        # --- Making arrays ---\n        start_arr_time  = time.perf_counter()\n        features_arr = np.array(features)\n        ground_truth_arr = np.array(labels)\n        end_arr_time = time.perf_counter()\n        print(f\"Making arrays took: {end_arr_time - start_arr_time:.2f}\")\n        \n        # --- Save Files ---\n        start_saving_time = time.perf_counter()\n        np.save(feature_file, features_arr)\n        np.save(ground_truth_file, ground_truth_arr)\n        np.save(global_axis_file, global_axis_features)\n        end_saving_time = time.perf_counter()\n        print(f\"Saving files took {end_saving_time - start_saving_time:.2f} seconds.\")\n\n        end_arr_creation = time.perf_counter()\n        print(f\"Total function runtime: {end_arr_creation - start_arr_creation:.2f} seconds.\")\n        \n        print(\"Features and ground truth labels saved successfully!!\")\n        return feature_file, ground_truth_file, global_axis_file","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:12.772354Z","iopub.execute_input":"2025-07-26T12:38:12.772619Z","iopub.status.idle":"2025-07-26T12:38:12.783247Z","shell.execute_reply.started":"2025-07-26T12:38:12.772597Z","shell.execute_reply":"2025-07-26T12:38:12.782511Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Uncomment and run this code if there're no .npy files in the datasets\n# features_dir, ground_truth_dir, global_axis_file = consolidate_data()\n# features_dir, ground_truth_dir, global_axis_file","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T07:36:44.14027Z","iopub.execute_input":"2025-07-24T07:36:44.140534Z","iopub.status.idle":"2025-07-24T07:36:44.144063Z","shell.execute_reply.started":"2025-07-24T07:36:44.140513Z","shell.execute_reply":"2025-07-24T07:36:44.143355Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"features_path = '/kaggle/input/neurlips-2025-exoplanet-arrays/combined_input_features.npy'\nlabels_path = '/kaggle/input/neurlips-2025-exoplanet-arrays/ground_truth_spectrum.npy'\nglobal_axis_path = '//kaggle/input/neurlips-2025-exoplanet-arrays/ground_truth_spectrum.npy'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:18.276767Z","iopub.execute_input":"2025-07-26T12:38:18.277337Z","iopub.status.idle":"2025-07-26T12:38:18.280703Z","shell.execute_reply.started":"2025-07-26T12:38:18.277313Z","shell.execute_reply":"2025-07-26T12:38:18.279884Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\n\nclass PlanetDataset(Dataset):\n    def __init__(self, features_path, labels_path):\n        # Load the entire consolidated dataset into memory once\n        self.features = np.load(features_path)\n        self.labels = np.load(labels_path)\n        \n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, idx):\n        # Simply retrieve the data from the pre-loaded arrays\n        features_tensor = torch.tensor(self.features[idx], dtype=torch.float32)\n        labels_tensor = torch.tensor(self.labels[idx], dtype=torch.float32)\n        \n        return features_tensor, labels_tensor","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:21.40878Z","iopub.execute_input":"2025-07-26T12:38:21.409073Z","iopub.status.idle":"2025-07-26T12:38:21.414247Z","shell.execute_reply.started":"2025-07-26T12:38:21.409055Z","shell.execute_reply":"2025-07-26T12:38:21.413517Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def gll_loss(mu, sigma, y):\n    \"\"\"\n    Calculates the negative log-likelihood of observing y given predicted mu and sigma.\n    \"\"\"\n    # Use torch.mean to get the average loss over the batch\n    return torch.mean(0.5 * torch.log(2 * np.pi * sigma**2) + (y - mu)**2 / (2 * sigma**2))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:24.736026Z","iopub.execute_input":"2025-07-26T12:38:24.736261Z","iopub.status.idle":"2025-07-26T12:38:24.740229Z","shell.execute_reply.started":"2025-07-26T12:38:24.736244Z","shell.execute_reply":"2025-07-26T12:38:24.739503Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"features_arr = np.load(features_path)\nX_shape = features_arr.shape\nX_shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:26.141313Z","iopub.execute_input":"2025-07-26T12:38:26.141837Z","iopub.status.idle":"2025-07-26T12:38:26.194739Z","shell.execute_reply.started":"2025-07-26T12:38:26.141804Z","shell.execute_reply":"2025-07-26T12:38:26.194173Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_shape = np.load(labels_path).shape\ny_shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:27.926102Z","iopub.execute_input":"2025-07-26T12:38:27.926617Z","iopub.status.idle":"2025-07-26T12:38:27.950172Z","shell.execute_reply.started":"2025-07-26T12:38:27.926597Z","shell.execute_reply":"2025-07-26T12:38:27.94963Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training the neural network","metadata":{}},{"cell_type":"code","source":"def run_training_pipeline_with_gpu():\n    print(\"--- Starting Exoplanet Spectrum Training Pipeline (GPU-Optimized) ---\")\n\n    # Checking consolidated data \n    \n    # Check for GPU availability\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {device}\")\n\n    # --- Create the Dataset and DataLoader ---\n    train_dataset = PlanetDataset(\n        features_path,\n        labels_path\n    )\n    \n    # DataLoader handles batching and parallel data loading\n    # Batch size: how many samples to process at once on the GPU\n    # num_workers: how many CPU cores to use for data loading (set to 0 for a simple test)\n    # pin_memory: faster data transfer to the GPU\n    BATCH_SIZE = 32\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=BATCH_SIZE,\n        shuffle=True, # Shuffle the data for training\n        num_workers=4, # Use multiple CPU cores to load data\n        pin_memory=True\n    )\n    print(f\"\\nCreated DataLoader with {len(train_loader)} batches of size {BATCH_SIZE}.\")\n\n    # --- Step 4: Move Model to GPU ---\n    model = SpectrumRegressor(X_shape[1], y_shape[1]).to(device)\n    opt = Adam(model.parameters(), 1e-4)\n\n    # --- Step 5: Training Loop using DataLoader ---\n    print(\"\\n--- Starting Model Training ---\")\n    num_epochs = 1000\n    for epoch in range(num_epochs):\n        print(f\"Going in {epoch}\")\n        model.train()\n        total_loss = 0\n        \n        # Iterate over the DataLoader, which yields batches of data\n        for X_batch, y_batch in train_loader:\n            # Move batch data to the GPU\n            X_batch = X_batch.to(device)\n            y_batch = y_batch.to(device)\n\n            opt.zero_grad()\n            mu, sigma = model(X_batch)\n            loss = gll_loss(mu, sigma, y_batch)\n            loss.backward()\n            opt.step()\n            total_loss += loss.item()\n        \n        avg_loss = total_loss / len(train_loader)\n        if (epoch + 1) % 10 == 0 or epoch == num_epochs:\n            print(f\"Epoch {epoch + 1:02d}/{num_epochs} | Average Loss: {avg_loss:.6f}\")\n\n    print(\"\\n--- Model Training Complete! ---\")\n    print(f\"Final Average Loss: {avg_loss:.6f}\")\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:30.871372Z","iopub.execute_input":"2025-07-26T12:38:30.871961Z","iopub.status.idle":"2025-07-26T12:38:30.878455Z","shell.execute_reply.started":"2025-07-26T12:38:30.871939Z","shell.execute_reply":"2025-07-26T12:38:30.877756Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# No need to run\n# model = run_training_pipeline_with_gpu()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T11:06:08.85547Z","iopub.execute_input":"2025-07-26T11:06:08.856184Z","iopub.status.idle":"2025-07-26T11:06:08.859895Z","shell.execute_reply.started":"2025-07-26T11:06:08.856157Z","shell.execute_reply":"2025-07-26T11:06:08.858952Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Final steps","metadata":{}},{"cell_type":"markdown","source":"## Saving the model with weights","metadata":{}},{"cell_type":"code","source":"'''\nI wrote this code to saved the model as a .pth file with its\nentire learned weights. No need to run this code.\n'''\n\n# import torch\n\n# torch.save(model.state_dict(), \"/kaggle/working/spectrum_regressor_weights.pth\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-23T14:58:35.865898Z","iopub.execute_input":"2025-07-23T14:58:35.866135Z","iopub.status.idle":"2025-07-23T14:58:35.870874Z","shell.execute_reply.started":"2025-07-23T14:58:35.86612Z","shell.execute_reply":"2025-07-23T14:58:35.870313Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Creating submission.csv","metadata":{}},{"cell_type":"code","source":"# import numpy as np\n# import torch\n# from torch.utils.data import Dataset, DataLoader\n\n# class PlanetTestDataset(Dataset):\n#     def __init__(self, test_x, test_y):\n#         self.test_x = test_x\n#         self.test_y = test_y\n\n#     def __len__(self):\n#         return len(self.test_x)\n\n#     def __getitem__(self, idx):\n#         x = torch.tensor(self.test_x[idx], dtype=torch.float32)\n#         y = torch.tensor(self.test_y[idx], dtype=torch.float32)\n\n#         return x, y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T06:26:25.653607Z","iopub.execute_input":"2025-07-26T06:26:25.654018Z","iopub.status.idle":"2025-07-26T06:26:25.659777Z","shell.execute_reply.started":"2025-07-26T06:26:25.65394Z","shell.execute_reply":"2025-07-26T06:26:25.658945Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''\nI wrote this code to create the planets_id.npy.\nYou don't need to run this!!\n'''\n# import numpy as np\n\n# try:\n#     _, train_star_info_df, _, train_labels_df, _ = load_all_global_data(DATA_ROOT)\n# except FileNotFoundError as e:\n#     print(f\"\\nERROR: Required data file not found: {e}\")\n\n# dir_ids_set = set(\n#     [int(d) for d in os.listdir(TRAIN_DIR) if d.isdigit()]\n# )\n# labels_ids_set = set(train_labels_df.index.to_list())\n# star_info_ids_set = set(train_star_info_df.index.to_list())\n# training_pids_set = dir_ids_set \\\n#                     .intersection(labels_ids_set) \\\n#                     .intersection(star_info_ids_set)\n# training_planet_ids = sorted(list(training_pids_set))\n\n# planet_ids_arr = np.array(training_planet_ids)\n\n# np.save('/kaggle/working/planet_ids.npy', planet_ids_arr)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T07:38:00.964498Z","iopub.execute_input":"2025-07-24T07:38:00.965018Z","iopub.status.idle":"2025-07-24T07:38:00.969669Z","shell.execute_reply.started":"2025-07-24T07:38:00.964991Z","shell.execute_reply":"2025-07-24T07:38:00.968964Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sample_submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T06:26:51.598261Z","iopub.execute_input":"2025-07-26T06:26:51.598573Z","iopub.status.idle":"2025-07-26T06:26:51.618003Z","shell.execute_reply.started":"2025-07-26T06:26:51.598548Z","shell.execute_reply":"2025-07-26T06:26:51.617181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv.set_index('planet_id', inplace = True)\ntrain_csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-25T11:45:36.297506Z","iopub.execute_input":"2025-07-25T11:45:36.297835Z","iopub.status.idle":"2025-07-25T11:45:36.326244Z","shell.execute_reply.started":"2025-07-25T11:45:36.297813Z","shell.execute_reply":"2025-07-25T11:45:36.325256Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_csv.loc[4294092928]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-25T11:50:38.940326Z","iopub.execute_input":"2025-07-25T11:50:38.945928Z","iopub.status.idle":"2025-07-25T11:50:38.975238Z","shell.execute_reply.started":"2025-07-25T11:50:38.945807Z","shell.execute_reply":"2025-07-25T11:50:38.973696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\n\ndef process_test_planet(pid, test_dir, train_star_info_df, adc_info_df):\n    test_star_info_dir = '/kaggle/input/ariel-data-challenge-2025/test_star_info.csv'\n    planet_dir = f\"{test_dir}/{pid}\"\n    test_star_info_df = pd.read_csv(test_star_info_dir)\n    test_star_info_df.set_index('planet_id', inplace=True)\n    \n    airs_ch0_data = load_planet_instrument_data(\n        planet_dir, 'AIRS-CH0', adc_info_df\n    )\n    signal_features = airs_ch0_data\n    \n    test_planet_star_info = test_star_info_df.loc[pid]\n    stellar_features = np.array([\n        test_planet_star_info['Rs'], test_planet_star_info['Ms'], test_planet_star_info['Ts'], test_planet_star_info['Mp'],\n        test_planet_star_info['e'], test_planet_star_info['P'], test_planet_star_info['sma'], test_planet_star_info['i']\n    ], dtype=np.float32)\n    combined_features = np.concatenate((signal_features, stellar_features))\n    \n    return combined_features","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:38:55.459572Z","iopub.execute_input":"2025-07-26T12:38:55.45986Z","iopub.status.idle":"2025-07-26T12:38:55.465147Z","shell.execute_reply.started":"2025-07-26T12:38:55.459806Z","shell.execute_reply":"2025-07-26T12:38:55.464408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''No need to run this!'''\n# import os\n# import pandas as pd\n# import numpy as np\n\n# test_feature_file = f'{OUTPUT_DIR}test_feature_file.npy'\n\n# try:\n#     adc_info_df, train_star_info_df, _, _, _ = load_all_global_data(DATA_ROOT)\n# except FileNotFoundError as e:\n#     print(f\"\\nERROR: Required data file not found: {e}\")\n\n# TEST_DIR = '/kaggle/input/ariel-data-challenge-2025/test'\n# test_planets_ids = [int(pid) for pid in os.listdir(TEST_DIR)]\n\n# if not test_planets_ids:\n#     print(\"\\nERROR: No valid data for testing.\")\n\n# print(f\"------------- Found {len(test_planets_ids)} testing planets. Consolidating data -------------\")\n\n# features = []\n# i = 0\n\n# for pid in test_planets_ids:\n#     print(f\"Getting planet id at {i} index\")\n#     combined_features = process_test_planet(\n#         pid,\n#         TEST_DIR,\n#         train_star_info_df,\n#         adc_info_df\n#     )\n#     features.append(combined_features)\n#     i += 1\n\n# features_arr = np.array(features)\n# np.save(test_feature_file, features_arr)\n\n# print(f\"------------- Saved the testing array as {test_feature_file} -------------\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T06:45:28.802625Z","iopub.execute_input":"2025-07-26T06:45:28.802967Z","iopub.status.idle":"2025-07-26T06:45:31.408335Z","shell.execute_reply.started":"2025-07-26T06:45:28.802938Z","shell.execute_reply":"2025-07-26T06:45:31.407304Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\ntest_feature_dir = '/kaggle/input/neurlips-2025-exoplanet-arrays/test_feature_file.npy'\ntest_arr = np.load(test_feature_dir)\ntest_arr.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:39:02.802227Z","iopub.execute_input":"2025-07-26T12:39:02.802561Z","iopub.status.idle":"2025-07-26T12:39:02.811422Z","shell.execute_reply.started":"2025-07-26T12:39:02.802538Z","shell.execute_reply":"2025-07-26T12:39:02.810621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nimport torch\n\nclass PlanetRegressorTestSet(Dataset):\n    def __init__(self):\n        self.test_arr = test_arr\n\n    def __len__(self):\n        return len(self.test_arr)\n\n    def __getitem__(self, idx):\n        return torch.tensor(self.test_arr[idx], dtype=torch.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:39:06.74226Z","iopub.execute_input":"2025-07-26T12:39:06.742808Z","iopub.status.idle":"2025-07-26T12:39:06.747103Z","shell.execute_reply.started":"2025-07-26T12:39:06.742784Z","shell.execute_reply":"2025-07-26T12:39:06.746319Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_planets = np.array([1103775])\ntest_planets.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:39:09.485895Z","iopub.execute_input":"2025-07-26T12:39:09.486118Z","iopub.status.idle":"2025-07-26T12:39:09.49099Z","shell.execute_reply.started":"2025-07-26T12:39:09.486101Z","shell.execute_reply":"2025-07-26T12:39:09.490376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\noutput_dim = 283\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(\"Device at hand: \", device)\nmodel_path = '/kaggle/input/spectrum-regressor/pytorch/default/1/spectrum_regressor_weights.pth'\nloaded_dict = torch.load(model_path, map_location = device)\nmodel = SpectrumRegressor(test_arr.shape[1], output_dim).to(device)\n\nBATCH_SIZE = 32\ntest_dataset = PlanetRegressorTestSet()\ntest_loader = DataLoader(\n    test_dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=False,\n    num_workers=4,\n    pin_memory=True\n)\n\nmodel.load_state_dict(loaded_dict)\n\nall_mu_preds = []\nall_sigma_preds = []\n\nprint(\"\\n--- Generating Predictions for Submission ---\")\ni = 0\nmodel.eval()\n\nwith torch.no_grad(): # Disable gradient calculations\n    for X_batch in test_loader:\n        X_batch = X_batch.to(device)\n        mu, log_sigma = model(X_batch) # Get mu and log_sigma\n        sigma = torch.exp(log_sigma) # Convert log_sigma to sigma\n\n        all_mu_preds.append(mu.cpu().numpy())\n        all_sigma_preds.append(sigma.cpu().numpy())\n        if (i + 1) % 10 == 0:\n            print(f\"Processed {i + 1} batches...\")\n        i += 1\n\nmu_preds_array = np.vstack(all_mu_preds)\nsigma_preds_array = np.vstack(all_sigma_preds)\n\nprint(f\"Shape of aggregated mu predictions: {mu_preds_array.shape}\")\nprint(f\"Shape of aggregated sigma predictions: {sigma_preds_array.shape}\")\n\n# Create column names\noutput_dim = model.output_dim\nwl_cols = [f\"wl_{i}\" for i in range(1, output_dim + 1)] \nsigma_cols = [f\"sigma_{i}\" for i in range(1, output_dim + 1)]\n    \ncombined_preds_array = np.hstack((mu_preds_array, sigma_preds_array))\n\nsubmission_df = pd.DataFrame(combined_preds_array, columns=wl_cols + sigma_cols)\nsubmission_df.insert(0, 'planet_id', test_planets)\n\noutput_filename = f'{OUTPUT_DIR}/submission.csv'\nsubmission_df.to_csv(output_filename, index=False)\n    \nprint(f\"\\nSubmission file '{output_filename}' created successfully!\")\nprint(f\"Final submission file shape: {submission_df.shape}\")\n\nsubmission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:39:36.535988Z","iopub.execute_input":"2025-07-26T12:39:36.536496Z","iopub.status.idle":"2025-07-26T12:39:37.21005Z","shell.execute_reply.started":"2025-07-26T12:39:36.536472Z","shell.execute_reply":"2025-07-26T12:39:37.20934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sample_submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-26T12:33:21.330027Z","iopub.execute_input":"2025-07-26T12:33:21.330374Z","iopub.status.idle":"2025-07-26T12:33:21.351544Z","shell.execute_reply.started":"2025-07-26T12:33:21.330313Z","shell.execute_reply":"2025-07-26T12:33:21.350404Z"}},"outputs":[],"execution_count":null}]}