{"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":"gpu","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"sourceType":"competition"},{"sourceId":13133519,"sourceType":"datasetVersion","datasetId":7873586},{"sourceId":13171574,"sourceType":"datasetVersion","datasetId":8302646}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%load_ext autoreload\n%autoreload 2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:32:10.024557Z","iopub.execute_input":"2025-09-25T16:32:10.025133Z","iopub.status.idle":"2025-09-25T16:32:10.054117Z","shell.execute_reply.started":"2025-09-25T16:32:10.025108Z","shell.execute_reply":"2025-09-25T16:32:10.053539Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index --find-links=/kaggle/input/ariel25-batman-minuit/packages  batman-package iminuit pqdm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:20:24.619973Z","iopub.execute_input":"2025-09-25T16:20:24.620303Z","iopub.status.idle":"2025-09-25T16:20:29.09627Z","shell.execute_reply.started":"2025-09-25T16:20:24.620261Z","shell.execute_reply":"2025-09-25T16:20:29.09529Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nos.environ['DATASET']='train'\nos.environ['RPATH']= '/kaggle/input/ariel-data-challenge-2025/'\n!cp   /kaggle/input/ariel25-source-and-models/*  /kaggle/working\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:51:13.406905Z","iopub.execute_input":"2025-09-25T16:51:13.407697Z","iopub.status.idle":"2025-09-25T16:51:13.680164Z","shell.execute_reply.started":"2025-09-25T16:51:13.40767Z","shell.execute_reply":"2025-09-25T16:51:13.679379Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport os\nimport glob\nimport numpy as np\nfrom sklearn.decomposition import PCA\nimport matplotlib.pyplot as plt\n\npath = os.environ['RPATH']\ndataset = os.environ['DATASET']\n\ntrain_labels= pd.read_csv(f'{path}train.csv')\n\ndtest = glob.glob(f'{path}{dataset}/*/AIRS-CH0_signal_*')\ndt = [d.split('/') for d in dtest]\nnpa = path.count('/')\ndt = [(d[npa+1],d[npa+2][16]) for d in dt]\ndf_t = pd.DataFrame(dt)\ndf_t.columns = ['planet_id','rep']\ndf_t['planet_id'] = df_t.planet_id.astype('int64')\n\n   \nif dataset == 'train':\n    df_labels = pd.merge(df_t,train_labels,how = 'left')\n    y_true=df_labels.iloc[:,2:].values\n\ndf_star_info = pd.read_csv(f'{path}{dataset}_star_info.csv')\npid = df_star_info.planet_id.astype(int).values\ndf_star_info.drop('planet_id', axis = 1)\ndf_star_info['planet_id'] = pid\ndf_star = pd.merge(df_t,df_star_info)\n\n\n\n# Generate PCA templates from true spectra   \n\nncomp = 7\ntl = train_labels.iloc[:,1:]\n\ntls = (tl.T/tl.mean(axis=1)).T-1\npca = PCA(n_components=ncomp,random_state = 42)\npca.fit(tls)\n\nnp.save('components',pca.components_)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:20:29.439638Z","iopub.execute_input":"2025-09-25T16:20:29.439943Z","iopub.status.idle":"2025-09-25T16:20:35.315769Z","shell.execute_reply.started":"2025-09-25T16:20:29.439906Z","shell.execute_reply":"2025-09-25T16:20:35.314979Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Introduction","metadata":{"execution":{"iopub.status.busy":"2025-09-25T08:52:06.962638Z","iopub.execute_input":"2025-09-25T08:52:06.963441Z","iopub.status.idle":"2025-09-25T08:52:06.96728Z","shell.execute_reply.started":"2025-09-25T08:52:06.963402Z","shell.execute_reply":"2025-09-25T08:52:06.966566Z"}}},{"cell_type":"markdown","source":"This solution presents a comprehensive, three-stage pipeline designed to extract high-fidelity transmission spectra from the FGS1 and AIR-CH0 Ariel instrument data. The approach combines physics-based modeling with a machine learning-inspired post-processing framework to tackle the key challenges of instrumental systematics and uncertainty estimation.\n\n1. Stage 1: GPU-Accelerated Preprocessing. Raw sensor data is converted into clean, calibrated light curves using a CuPy-based pipeline. This stage handles standard instrumental corrections, including non-linearity, dark current, flat-fielding, and jitter, all performed efficiently on the GPU.\n\n2. Stage 2: Hierarchical Transit Fitting. I use the **batman** transit modeling library and the **iminuit** optimizer to perform a multi-step fit. A robust global model is first fit to the combined FGS and AIRS light curves to constrain key orbital parameters. This is followed by a 2D drift correction and a wavelength-by-wavelength fit to extract the initial transmission spectrum.\n\n3. Stage 3: Cross-Validated Post-Processing Ensemble. The initial spectra are refined using an ensemble of models trained with Grouped K-Fold Cross-Validation. This final stage blends a PCA-regularized signal with a smoothed signal and employs a parameterized model to predict the final uncertainties (sigmas), optimizing all parameters directly against the competition score.\n\n\nThis multi-stage design ensures that physical constraints are respected in the initial fit, while the final model has the flexibility to learn and correct for residual systematic errors and produce well-calibrated uncertainties.\n","metadata":{}},{"cell_type":"markdown","source":"# Stage 1: Preprocessing Raw Data","metadata":{}},{"cell_type":"markdown","source":"\nThe first step in the pipeline is to transform the raw sensor readouts into scientifically useful light curves. This entire process is executed on the GPU using CuPy for maximum efficiency.\n\nKey Preprocessing Steps:\n\n• Initial Calibrations: The pipeline begins by applying standard instrumental corrections. This includes an ADC correction for gain and offset, a 5th-degree polynomial correction for detector non-linearity (apply_linear_corr_fast), and subtraction of a scaled master dark frame (clean_dark).\n\n• Flat-Fielding & Masking: A master flat frame is applied to correct for pixel-to-pixel sensitivity variations. Pixels identified as \"hot\" (via sigma-clipping the dark frame) or \"dead\" are masked.\n\n• Jitter Correction: For the FGS1 sensor, we apply a center-of-mass  regression to correct for flux variations caused by image motion on the detector. The flux is de-correlated from the normalized X and Y  positions of the stellar image. While a PCA-based jitter correction method was developed for the AIRS sensor, it was found to be less effective and is disabled in the final pipeline.\n\n• Signal Extraction & Cleaning:\n\n   \n* For FGS1, the final flux is extracted using simple aperture photometry. PSF photometry gives lower SNR signals. probably due to the jitter.\n* For AIRS, the signal is extracted from the central detector region, and a background signal, calculated from the top and bottom edges of the detector, is subtracted. This is not optimal because the spectroscopy signal extends i a difractive way to the whole sensor, making dificult to  identify the aditive frequency  dependent background and  making necessary a global  scale factor and an spectrum slope correction in the post processing stage\n* A spike-cleaning algorithm  is applied to the final time series to remove cosmic ray hits by identifying and replacing outliers in the frame-to-frame difference.\n\n    \n• Binning: To improve the signal-to-noise ratio, the final time series for both instruments are binned by averaging consecutive frames.","metadata":{}},{"cell_type":"code","source":"from preprocess import reduce\n\ni0 = 499\n\nbinning = 1\nfile = df_t.iloc[i0]\nplanet_id = file[0]\nfgs0,_ = reduce(file,dataset,'FGS1',binning,dejitter = False)\nfgs1,_ = reduce(file,dataset,'FGS1',binning,dejitter = True)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:20:35.316939Z","iopub.execute_input":"2025-09-25T16:20:35.317157Z","iopub.status.idle":"2025-09-25T16:21:03.899679Z","shell.execute_reply.started":"2025-09-25T16:20:35.317138Z","shell.execute_reply":"2025-09-25T16:21:03.899022Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"FGS1 dejitter clearly improve the Signa to Noise Ratio (SNR) in the measured flux. In the AIR-CH0, a similar jitter appear with different intensities and sign for any of the 32 horizontal channels, but they cancel out if we integrate the image in a simerit interva fro ror r to row 32-r ","metadata":{}},{"cell_type":"code","source":"fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))\n\nax1.plot(fgs0,alpha = .7)\nax1.plot(fgs1,alpha = .7)\nax1.set_title('Signal Flux over time')\nax1.set_xlabel('Flux')\nax1.set_ylabel('Time id')\nax1.legend()\n\nax2.hist(fgs0,bins = 100,alpha =.7);\nax2.hist(fgs1,bins = 100,alpha =.7);\nax2.set_title('Signal Flux histogram')\nax2.set_xlabel('Flux')\nax2.set_ylabel('Count')\nax2.legend()\n\n\nplt.tight_layout()\n\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:21:03.900394Z","iopub.execute_input":"2025-09-25T16:21:03.900615Z","iopub.status.idle":"2025-09-25T16:21:04.691365Z","shell.execute_reply.started":"2025-09-25T16:21:03.900595Z","shell.execute_reply":"2025-09-25T16:21:04.690528Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"binning = 12\nfgs,_ = reduce(file,dataset,'FGS1',12*binning)\nair,_  = reduce(file,dataset,'AIRS-CH0',binning)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:21:04.692137Z","iopub.execute_input":"2025-09-25T16:21:04.692482Z","iopub.status.idle":"2025-09-25T16:21:12.51199Z","shell.execute_reply.started":"2025-09-25T16:21:04.692464Z","shell.execute_reply":"2025-09-25T16:21:12.511205Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Stage 2: Physics-Based Spectral Fitting\n\nWith clean light curves for both instruments, we extract the initial transmission spectrum. This process is parallelized using pqdm to efficiently process all observations.\n\nHierarchical Fitting Strategy:\n\n1. **Global Combined Fit**: We first fit the FGS light curve and the wavelength-averaged AIRS light curve **simultaneously** using fit_combined_curves. This crucial step provides robust constraints on the shared physical parameters of the system (e.g., orbital period per, inclination inc, transit time t0), which are initialized from the provided star_info.csv file. The transit is modeled using batman, asuming a **quadratic limb darkening model**, using 2 parameters, and instrumental trends are modeled with a polynomial baseline in time.\n   \n2. **Refined Fit & Drift Correction**: The results from the first fit are used to initialize a second, refined fit on normalized data. Subsequently, a 2D instrumental drift model is fit to the out-of-transit portion of the AIRS data cube. This model (Drift class) consists of two 4th-order polynomials, one for the spatial and one for the temporal axis, and it corrects for slow-varying systematic patterns across the detector. The data cube is then divided by this drift model.\n\n\n3. **Wavelength-by-Wavelength Fit**: The final step is to measure the transit depth in each individual wavelength channel of the drift-corrected AIRS data. The fit_w function iterates through the channels, fitting the data within a small moving window (deltaw=2) to boost signal-to-noise. In this fit, most orbital parameters are fixed to the values from the global fit, and only the transit depth (dipa) and baseline normalization (A0a) are allowed to vary.\nThis process results in the \"raw\" spectrum  and its associated uncertainties , which serve as the primary inputs for our final modeling stage.","metadata":{}},{"cell_type":"code","source":"from fit2 import get_flux_error,fit_combined_curves\n\n# 1. Initial fit on the mean light curves\nconstants = df_star.iloc[i0]\nflux_f, err_f, d0g, snrf = get_flux_error(fgs[:, 0])\nflux_a, err_a, d1g, snra = get_flux_error(air.mean(axis=1))\n\n\nminuit_result, chi2 = fit_combined_curves(constants, flux_f, err_f/np.sqrt(2), flux_a, err_a/np.sqrt(2), d0g, d1g, 1)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"minuit_result","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from fit2 import get_spectrum \n\n#full fit\n\nallf,alle,n_out,fvalw,vals_,fval,snrf,snra  = get_spectrum(fgs,air,constants,deltaw = 1,fitw =True,dedrift = True,flag_plot = True,ret_all=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:21:12.513721Z","iopub.execute_input":"2025-09-25T16:21:12.514239Z","iopub.status.idle":"2025-09-25T16:21:19.017179Z","shell.execute_reply.started":"2025-09-25T16:21:12.514205Z","shell.execute_reply":"2025-09-25T16:21:19.016481Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"plt.plot(allf)\nplt.plot(y_true[i0])","metadata":{"execution":{"iopub.status.busy":"2025-09-25T16:21:19.018064Z","iopub.execute_input":"2025-09-25T16:21:19.018404Z","iopub.status.idle":"2025-09-25T16:21:19.18178Z","shell.execute_reply.started":"2025-09-25T16:21:19.018382Z","shell.execute_reply":"2025-09-25T16:21:19.181047Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Preprocess and fit the full train set","metadata":{}},{"cell_type":"code","source":"!python process_all.py","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Stage 3: Post-Processing and Uncertainty Modeling\n\nThe final stage of our solution is a post-processing model that refines the raw spectra and, critically, learns a robust model for the final uncertainties. This is where the bulk of the \"machine learning\" occurs.\n\nModel Architecture (build_preds): The model generates a final prediction by blending two different representations of the raw spectrum (preds1):\n\n**PCA-Regularized Prediction (preds0)**: The raw spectrum undergoes a slope correction and is then projected onto a pre-defined PCA basis (components.npy). This de-noises the spectrum and imposes a strong regularization prior, capturing the most common modes of variation.\n**Smoothed Prediction (preds2)**: The raw AIRS spectrum is smoothed using a Savitzky-Golay filter to reduce high-frequency noise while preserving broader spectral features.\n\nThe final mean prediction is a weighted average of these two models: preds = app0 + (1-ap)preds2, where ap is a learned parameter.\n\n**Sigma Model**: Predicting accurate uncertainties is key to maximizing the competition score. I developed a complex, empirically-derived formula to model the final sigma for each data point. This formula combines multiple sources of uncertainty:\n\n \n* a1* err^e1 : The propagated error from the initial fit, with a learned exponent.\n* a01/p1^e2  : A term related to the signal magnitude.\n* b0 * sigma0: A term proportional to the raw signal's volatility.\n* ae * np.abs(p0-preds2): A term that increases uncertainty where the PCA and smoothed models disagree, capturing model \nuncertainty or unknow spectra.\n\n#Training and Ensembling:\n\n  **Objective**: The free parameters of the model (ap, e1, e2, a1, a0, b0, ae, etc.) are optimized using iminuit to directly maximize the competition score.\n\n  **Cross-Validation**: To build a robust model that generalizes well, we employ a 10-fold Grouped K-Fold Cross-Validation strategy, using planet_id to group the data. This ensures that all observations of a single planet are kept within the same fold, preventing data leakage.\n\n  **Calibration**: Within each fold, after the primary parameters are optimized, we calculate a wavelength-dependent calibration factor (alpha) that scales the predicted sigmas to best match the variance observed in the training data. This factor is then applied to the validation set predictions.\n  \n   **Ensembling**: The final submission is generated by averaging the calibrated predictions and sigmas from the 10 models trained during the cross-validation process. This ensembling technique reduces variance and improves the final score.\n\nThe entire training process, including the parallel optimization of each fold, is managed by the CV_model function.","metadata":{}},{"cell_type":"code","source":"%%time\nfrom model import CV_model\n\nsigma, preds, models = CV_model(y_true,df_t,beta=.71,n_splits=4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-25T16:51:33.554206Z","iopub.execute_input":"2025-09-25T16:51:33.554527Z","iopub.status.idle":"2025-09-25T17:01:24.316914Z","shell.execute_reply.started":"2025-09-25T16:51:33.554501Z","shell.execute_reply":"2025-09-25T17:01:24.316099Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}