{"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":"none","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport itertools\nimport os\nimport glob \nfrom astropy.stats import sigma_clip\nfrom concurrent.futures import ProcessPoolExecutor, as_completed\n\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nfrom matplotlib.animation import FuncAnimation\nfrom IPython.display import HTML\nfrom scipy.signal import medfilt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T02:57:06.287668Z","iopub.execute_input":"2025-09-13T02:57:06.287973Z","iopub.status.idle":"2025-09-13T02:57:08.814104Z","shell.execute_reply.started":"2025-09-13T02:57:06.287932Z","shell.execute_reply":"2025-09-13T02:57:08.813146Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"path_folder = '/kaggle/input/ariel-data-challenge-2025/' # path to the folder containing the data\npath_out = '/kaggle/tmp/data_light_raw/' # path to the folder to store the light data\nout_dir = 'data/' # path for the output directory","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T02:59:51.996189Z","iopub.execute_input":"2025-09-13T02:59:51.996523Z","iopub.status.idle":"2025-09-13T02:59:52.002389Z","shell.execute_reply.started":"2025-09-13T02:59:51.996498Z","shell.execute_reply":"2025-09-13T02:59:52.001044Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Useful functions\ndef ADC_convert(signal, gain=0.4369, offset=-1000):\n    \"\"\"The Analog-to-Digital Conversion (adc) is performed by the detector to convert\n    the pixel voltage into an integer number. Since we are using the same conversion number \n    this year, we have simply hard-coded it inside. \"\"\"\n    signal /= gain\n    signal += offset\n    return signal\n\n# Mask hot and dead pixels\ndef mask_hot_dead(signal, dead, dark):\n    hot = sigma_clip(\n        dark, sigma=5, maxiters=5\n    ).mask\n    hot = np.tile(hot, (signal.shape[0], 1, 1))\n    dead = np.tile(dead, (signal.shape[0], 1, 1))\n\n    # Combine masks\n    combined_mask = np.logical_or(dead, hot)\n\n    signal = np.ma.masked_where(combined_mask, signal).astype(signal.dtype, copy=False)\n    return signal\n\ndef apply_linear_corr(linear_corr, signal):\n    \"\"\"\n    Vectorized Horner evaluation.\n    linear_corr: (K, Y, X) coefficients c0..c_{K-1} (ascending)\n    clean_signal: (T, Y, X)\n    Returns: (T, Y, X)\n    \"\"\"\n    # xp = type(signal)  # works if you have xp in outer scope; else pass xp in\n\n    # Optional: switch to float32 for speed (if accuracy allows)\n    # signal = signal.astype(xp.float32, copy=False)\n    # linear_corr = linear_corr.astype(xp.float32, copy=False)\n\n    # Start from highest coefficient and fold down: v = c_{K-1}; v = v*s + c_{K-2}; ...\n    c = linear_corr  # ascending\n    s = signal\n\n    # Broadcast c_{K-1} to (T,Y,X), then in-place Horner\n    out = np.broadcast_to(c[-1][None, ...], s.shape).copy()\n    for k in range(c.shape[0] - 2, -1, -1):\n        out *= s           # v *= s\n        out += c[k]        # v += c_k\n    return out\n\n# Dark current subtraction function\ndef clean_dark(signal, dead, dark, dt):\n\n    dark = np.ma.masked_where(dead, dark)\n    dark = np.tile(dark, (signal.shape[0], 1, 1))\n\n    signal -= dark* dt[:, np.newaxis, np.newaxis]\n    return signal\n\n# Co-related Double Sampling (CDS) function\ndef get_cds(signal):\n    cds = signal[1::2,:,:] - signal[::2,:,:]\n    return cds\n\n# Flat field correction function\ndef correct_flat_field(flat, dead, signal):\n    # flat = flat.transpose(1, 0)\n    # dead = dead.transpose(1, 0)\n    flat = np.ma.masked_where(dead, flat)\n    flat = np.tile(flat, (signal.shape[0], 1, 1))\n    signal = signal / flat\n    return signal\n\n# (Optional) Time Binning function (May need some modifications)\ndef bin_obs(cds_signal, binning):\n    cds_transposed = cds_signal.transpose(0,1,3,2)\n    cds_binned = np.zeros((cds_transposed.shape[0], cds_transposed.shape[1]//binning, cds_transposed.shape[2], cds_transposed.shape[3]))\n    for i in range(cds_transposed.shape[1]//binning):\n        cds_binned[:,i,:,:] = np.sum(cds_transposed[:,i*binning:(i+1)*binning,:,:], axis=1)\n    return cds_binned\n\n# Cosmic ray removal function\n# Credit: lordpatil\ndef mad_clip_data(data, window_size=51, sigma=3):\n    \"\"\"Simple sigma clipping function.\"\"\"\n    local_median = medfilt(data, kernel_size=window_size)\n    residual = data - local_median\n    mad = np.median(np.abs(residual))\n    robust_std = mad * 1.4826\n    outliers = np.abs(residual) > (sigma * robust_std)\n\n    masked_data = np.ma.array(data, mask=outliers)\n    return masked_data, outliers","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T02:59:54.228265Z","iopub.execute_input":"2025-09-13T02:59:54.228991Z","iopub.status.idle":"2025-09-13T02:59:54.241568Z","shell.execute_reply.started":"2025-09-13T02:59:54.228962Z","shell.execute_reply":"2025-09-13T02:59:54.240439Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Code to process the signal from both detectors\ndef process_signal(path_folder, index, instrument, obs_count=0, x_slice=slice(None), y_slice=slice(None),\n                   do_mask=True, do_nl_corr=True, do_dark=True, do_flat=True, binning_factor=1):\n    \"\"\"Process the signal for AIRS-CH0 or FGS1 with the given parameters.\n    \n    Args:\n        path_folder (str): Base path to dataset folder.\n        index (str): Planet ID (e.g., planet_id).\n        instrument (str): 'AIRS-CH0' or 'FGS1'.\n        obs_count (int): Observation count (default: 0).\n        x_slice (slice): Slice for x-axis (default: instrument specific).\n        y_slice (slice): Slice for y-axis (default: instrument specific).\n        do_mask (bool): Apply hot/dead pixel masking.\n        do_nl_corr (bool): Apply non-linearity correction.\n        do_dark (bool): Apply dark subtraction.\n        do_flat (bool): Apply flat-field correction.\n        time_binning (bool): Apply time binning (default: False).\n    \n    Returns:\n        times (np.ndarray): Time array for the observation.\n        signal (np.ndarray): Processed 3D signal (n_times//2 x spatial x dispersion).\n    \"\"\"\n    if instrument == 'AIRS-CH0':\n        signal_file = f'{instrument}_signal_{obs_count}.parquet'\n        calib_folder = f'{instrument}_calibration_{obs_count}'\n        calib_shape = (32, 356)\n        calib_linear_shape = (6, 32, 356)\n        if x_slice == slice(None):\n            x_slice = slice(39, 321)   # wavelength slice for AIRS-CH0\n        if y_slice == slice(None):\n            y_slice = slice(10, 22)  # spatial slice for AIRS-CH0\n        axis_key = 'AIRS-CH0-axis0-h'\n        dt_key = 'AIRS-CH0-integration_time'\n    elif instrument == 'FGS1':\n        signal_file = f'{instrument}_signal_{obs_count}.parquet'\n        calib_folder = f'{instrument}_calibration_{obs_count}'\n        calib_shape = (32, 32)\n        calib_linear_shape = (6, 32, 32)\n        if x_slice == slice(None):\n            x_slice = slice(8, 24)\n        if y_slice == slice(None):\n            y_slice = slice(8, 24)\n        axis_key = 'FGS1-axis0-h'\n        # dt_key = 'FGS1-integration_time'\n    else:\n        raise ValueError(\"Instrument must be 'AIRS-CH0' or 'FGS1'\")\n\n    # Load signal data first\n    df = pd.read_parquet(os.path.join(path_folder, f'train/{index}/{signal_file}'))\n    reshape_dims = (df.shape[0], calib_shape[0], calib_shape[1])\n\n    # Load calibrations\n    flat = pd.read_parquet(os.path.join(path_folder, f'train/{index}/{calib_folder}/flat.parquet')).values.astype(np.float32).reshape(calib_shape)[y_slice, x_slice]\n    dark = pd.read_parquet(os.path.join(path_folder, f'train/{index}/{calib_folder}/dark.parquet')).values.astype(np.float32).reshape(calib_shape)[y_slice, x_slice]\n    dead = pd.read_parquet(os.path.join(path_folder, f'train/{index}/{calib_folder}/dead.parquet')).values.astype(np.float32).reshape(calib_shape)[y_slice, x_slice]\n    linear_corr = pd.read_parquet(os.path.join(path_folder, f'train/{index}/{calib_folder}/linear_corr.parquet')).values.astype(np.float32).reshape(calib_linear_shape)[:, y_slice, x_slice]\n    # Have not used the read data yet, probably used for handling uncertainty\n    read = pd.read_parquet(os.path.join(path_folder, f'train/{index}/{calib_folder}/read.parquet')).values.astype(np.float32).reshape(calib_shape)[y_slice, x_slice]\n    axis_info = pd.read_parquet(os.path.join(path_folder, 'axis_info.parquet'))\n    if instrument == 'AIRS-CH0':\n        dt = axis_info[dt_key].dropna().values\n    else:\n        dt = np.ones(df.shape[0]) * 0.1\n    dt[1::2] += 0.1\n\n    # Reshape and process signal\n    signal = df.values.astype(np.float32).reshape(reshape_dims)[:, y_slice, x_slice]\n    signal = ADC_convert(signal)\n\n    if do_mask:\n        signal = mask_hot_dead(signal, dead, dark)\n    \n    if do_nl_corr:\n        signal = apply_linear_corr(linear_corr, signal)\n\n    if do_dark:\n        signal = clean_dark(signal, dead, dark, dt)\n    \n    # Compute time array\n    times = axis_info[axis_key].dropna().values\n    if len(times) % 2 != 0:   # Ensure even number of time points\n        times = times[:-1]\n    times = 0.5 * (times[1::2] + times[0::2])\n\n    # Correlated Double Sampling (CDS)\n    signal = get_cds(signal)\n\n    if binning_factor > 1:\n        signal = signal[:(signal.shape[0] // binning_factor) * binning_factor, :, :]\n        signal = signal.reshape(signal.shape[0] // binning_factor, binning_factor, signal.shape[1], signal.shape[2]).mean(axis=1)\n        times = times[:(len(times) // binning_factor) * binning_factor]\n        times = times.reshape(-1, binning_factor).mean(axis=1)\n\n    # Flat field correction was applied after time binning in the original code\n    if do_flat:\n        signal = correct_flat_field(flat, dead, signal).data\n\n    # Returning read data for uncertainty estimation\n    return times, signal, read","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T02:59:58.017769Z","iopub.execute_input":"2025-09-13T02:59:58.018231Z","iopub.status.idle":"2025-09-13T02:59:58.034645Z","shell.execute_reply.started":"2025-09-13T02:59:58.018202Z","shell.execute_reply":"2025-09-13T02:59:58.033194Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_index_obsCount(files):\n    index_obsCount = []\n    obs_count = 0\n    for file in files:\n        index = file.split('/')[-1]\n        airs_files = glob.glob(os.path.join(file, 'AIRS-CH0_signal_*.parquet'))\n        obs_count = len(airs_files)\n        index_obsCount.append((int(index), obs_count))\n\n    return index_obsCount\n\n# Get list of planet IDs with their observation counts\nfiles = glob.glob(os.path.join(path_folder + 'train/', '*'))\n\n# planet IDs with their corresponding observation counts\nindices_w_obsCount = get_index_obsCount(files)\nindices_w_obsCount.sort(key=lambda x: x[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:00:22.726639Z","iopub.execute_input":"2025-09-13T03:00:22.727028Z","iopub.status.idle":"2025-09-13T03:00:27.54332Z","shell.execute_reply.started":"2025-09-13T03:00:22.727002Z","shell.execute_reply":"2025-09-13T03:00:27.542357Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not os.path.exists(path_out):\n    os.makedirs(path_out)\n    print(f\"Directory {path_out} created.\")\nelse:\n    print(f\"Directory {path_out} already exists.\")\n\nif not os.path.exists(out_dir):\n    os.makedirs(out_dir)\n    print(f\"Directory {out_dir} created.\")\nelse:\n    print(f\"Directory {out_dir} already exists.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:00:34.692897Z","iopub.execute_input":"2025-09-13T03:00:34.693234Z","iopub.status.idle":"2025-09-13T03:00:34.701668Z","shell.execute_reply.started":"2025-09-13T03:00:34.693209Z","shell.execute_reply":"2025-09-13T03:00:34.700633Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_star_info = pd.read_csv(os.path.join(path_folder, 'train_star_info.csv'))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:00:40.517179Z","iopub.execute_input":"2025-09-13T03:00:40.517477Z","iopub.status.idle":"2025-09-13T03:00:40.544188Z","shell.execute_reply.started":"2025-09-13T03:00:40.517455Z","shell.execute_reply":"2025-09-13T03:00:40.54334Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create a new DataFrame to store the expanded metadata (repeating rows based on observation counts)\ndf_new = pd.DataFrame({'planet_id': pd.Series(dtype=\"int\"), \n                       'Rs': pd.Series(dtype=\"float32\"), \n                       'Ms': pd.Series(dtype=\"float32\"), \n                       'Ts': pd.Series(dtype=\"float32\"),\n                       'Mp': pd.Series(dtype=\"float32\"),\n                       'e': pd.Series(dtype=\"float32\"),\n                       'P': pd.Series(dtype=\"float32\"),\n                       'sma': pd.Series(dtype=\"float32\"),\n                       'i': pd.Series(dtype=\"float32\")})\n\nfor idx, obs_count in tqdm(indices_w_obsCount):\n    for obs in range(obs_count):\n        df_new.loc[len(df_new)] = train_star_info[train_star_info['planet_id'] == idx][['planet_id', 'Rs', 'Ms', 'Ts', 'Mp', 'e', 'P', 'sma', 'i']].values[0]\n\ndf_new['planet_id'] = df_new['planet_id'].astype(int)\n\n# Save the new DataFrame to a CSV file\ndf_new.to_csv(os.path.join(out_dir,\"star_planet_metadata.csv\"), index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:00:41.928695Z","iopub.execute_input":"2025-09-13T03:00:41.929176Z","iopub.status.idle":"2025-09-13T03:00:43.623288Z","shell.execute_reply.started":"2025-09-13T03:00:41.92912Z","shell.execute_reply":"2025-09-13T03:00:43.622Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# WRAPPING THE WORK IN A FUNCTION\ndef process_obs(path_folder, index, obs, overall_binning, path_out=None, count=None):\n    # AIRS\n    times_airs, airs_cds, airs_read = process_signal(\n        path_folder, index, 'AIRS-CH0', obs_count=obs,\n        do_mask=True, do_nl_corr=True, do_dark=True,\n        do_flat=True, binning_factor=overall_binning\n    )\n\n    # FGS\n    times_fgs, fgs_cds, fgs_read = process_signal(\n        path_folder, index, 'FGS1', obs_count=obs,\n        do_mask=True, do_nl_corr=True, do_dark=True,\n        do_flat=True, binning_factor=12 * overall_binning\n    )\n\n    # Optional saving\n    if path_out is not None and count is not None:\n        np.save(os.path.join(path_out, f'AIRS_train_{count}.npy'), airs_cds)\n        np.save(os.path.join(path_out, f'FGS1_train_{count}.npy'), fgs_cds)\n\n    # Clean memory\n    del airs_cds, fgs_cds  \n\n    return (index, obs, times_airs, times_fgs)  # return what you need","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:17:31.940021Z","iopub.execute_input":"2025-09-13T03:17:31.940421Z","iopub.status.idle":"2025-09-13T03:17:31.950337Z","shell.execute_reply.started":"2025-09-13T03:17:31.94039Z","shell.execute_reply":"2025-09-13T03:17:31.949181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# BUILDING THE TASK LIST\n\n# Counter for number of processed observations\ncount = 0\n# A binning factor of 12 is used to match the time resolution of FGS1 with AIRS-CH0\n# Additional overall binning is applied to both detectors to reduce noise\noverall_binning = 10\n\ntasks = []\nfor index, obs_count in indices_w_obsCount:\n    for obs in range(obs_count):\n        tasks.append((path_folder, index, obs, overall_binning, path_out, count))\n        count += 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:17:33.692391Z","iopub.execute_input":"2025-09-13T03:17:33.692705Z","iopub.status.idle":"2025-09-13T03:17:33.699392Z","shell.execute_reply.started":"2025-09-13T03:17:33.692684Z","shell.execute_reply":"2025-09-13T03:17:33.698161Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"n_workers = 4 # Number of workers\n\nresults = []\nwith ProcessPoolExecutor(max_workers=n_workers) as executor:\n    futures = [executor.submit(process_obs, *task) for task in tasks]\n    for fut in tqdm(as_completed(futures), total=len(futures)):\n        results.append(fut.result())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:17:34.067852Z","iopub.execute_input":"2025-09-13T03:17:34.068233Z","iopub.status.idle":"2025-09-13T03:19:21.826999Z","shell.execute_reply.started":"2025-09-13T03:17:34.068205Z","shell.execute_reply":"2025-09-13T03:19:21.825221Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Concatenate all data into a single .npy file\ndef load_data(path, detector):\n    files = glob.glob(os.path.join(path, f'{detector}_train_*.npy'))\n    data_tmp = np.load(files[0])\n    data_all = np.zeros((len(files), *data_tmp.shape))\n    for i in range(len(files)):\n        data_all[i] = np.load(os.path.join(path, f'{detector}_train_{i}.npy'))\n    return data_all","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:22:01.093738Z","iopub.execute_input":"2025-09-13T03:22:01.09422Z","iopub.status.idle":"2025-09-13T03:22:01.099928Z","shell.execute_reply.started":"2025-09-13T03:22:01.094194Z","shell.execute_reply":"2025-09-13T03:22:01.098941Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls /kaggle/tmp/data_light_raw/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:22:01.366319Z","iopub.execute_input":"2025-09-13T03:22:01.367433Z","iopub.status.idle":"2025-09-13T03:22:01.495635Z","shell.execute_reply.started":"2025-09-13T03:22:01.367383Z","shell.execute_reply":"2025-09-13T03:22:01.4944Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Combined ARIS and FGS data\ndata_train_airs = load_data(path_out, 'AIRS')\nnp.save(os.path.join(out_dir, 'data_train_airs.npy'), data_train_airs)\ndel data_train_airs  # Free up memory\ndata_train_fgs = load_data(path_out, 'FGS1')\nnp.save(os.path.join(out_dir, 'data_train_fgs.npy'), data_train_fgs)\ndel data_train_fgs  # Free up memory","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:22:03.534389Z","iopub.execute_input":"2025-09-13T03:22:03.535371Z","iopub.status.idle":"2025-09-13T03:22:04.241966Z","shell.execute_reply.started":"2025-09-13T03:22:03.535332Z","shell.execute_reply":"2025-09-13T03:22:04.240412Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"index, obs = indices_w_obsCount[0]\nobs -= 1\n\ntimes_airs, airs_cds, airs_read = process_signal(\n        path_folder, index, 'AIRS-CH0', obs_count=obs,\n        do_mask=True, do_nl_corr=True, do_dark=True,\n        do_flat=True, binning_factor=overall_binning\n    )\n\ntimes_fgs, fgs_cds, fgs_read = process_signal(\n        path_folder, index, 'FGS1', obs_count=obs,\n        do_mask=True, do_nl_corr=True, do_dark=True,\n        do_flat=True, binning_factor=12 * overall_binning\n    )\n\nnp.save(os.path.join(out_dir, 'times_airs.npy'), times_airs)\nnp.save(os.path.join(out_dir, 'times_fgs.npy'), times_fgs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-13T03:22:17.788454Z","iopub.execute_input":"2025-09-13T03:22:17.788765Z","iopub.status.idle":"2025-09-13T03:22:23.717981Z","shell.execute_reply.started":"2025-09-13T03:22:17.788744Z","shell.execute_reply":"2025-09-13T03:22:23.716995Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}