{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"},{"sourceId":9798794,"sourceType":"datasetVersion","datasetId":6005206},{"sourceId":9799118,"sourceType":"datasetVersion","datasetId":6005399}],"dockerImageVersionId":30786,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Kaggle  \n# He HAO, Instituto Superior Técnico (https://github.com/HAOHE123)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Ariel Data Challenge 2024: Spectra Prediction\nIn this notebook, we aim to process satellite sensor data to predict planetary spectra as part of the Ariel Data Challenge 2024. We will perform data preprocessing, calibration, and optimization to generate accurate predictions of the planetary spectra.","metadata":{}},{"cell_type":"markdown","source":"## Introduction\nThe goal of this notebook is to process raw sensor data from the Ariel satellite's observations of exoplanets and predict their spectra. The process involves:\n\n- Calibrating the raw signals from different sensors.\n- Detecting the transit phases where the planet passes in front of its star.\n- Optimizing a scaling factor to adjust the signal during the transit phase.\n- Generating predictions for the spectra with associated uncertainties.## Libraries and Setup","metadata":{}},{"cell_type":"markdown","source":"## Libraries and Setup","metadata":{}},{"cell_type":"code","source":"import sys\nprint(sys.version)","metadata":{"execution":{"iopub.status.busy":"2024-11-04T02:18:19.371794Z","iopub.execute_input":"2024-11-04T02:18:19.372372Z","iopub.status.idle":"2024-11-04T02:18:19.426185Z","shell.execute_reply.started":"2024-11-04T02:18:19.372322Z","shell.execute_reply":"2024-11-04T02:18:19.424421Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Install missing packages\n!pip install  /kaggle/input/d/hehao111/nips-lib/pyerfa-2.0.1.4-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n!pip install  /kaggle/input/d/hehao111/nips-lib/astropy_iers_data-0.2024.11.4.0.33.34-py3-none-any.whl\n!pip install  /kaggle/input/d/hehao111/nips-lib/astropy-6.1.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\nfrom astropy.stats import sigma_clip","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Import required libraries\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport itertools\nfrom tqdm import tqdm\nfrom scipy.optimize import minimize\nfrom functools import partial","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Loading\n\nWe will load the test data and auxiliary information needed for preprocessing and calibration.","metadata":{}},{"cell_type":"code","source":"# Read the test_adc_info.csv and axis_info.parquet files\ntest_adc_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/test_adc_info.csv',\n                            index_col='planet_id')\naxis_info = pd.read_parquet('/kaggle/input/ariel-data-challenge-2024/axis_info.parquet')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Preprocessing","metadata":{}},{"cell_type":"code","source":"def apply_linear_corr(linear_corr, clean_signal):\n    linear_corr = np.flip(linear_corr, axis=0)\n    for x, y in itertools.product(\n                range(clean_signal.shape[1]), range(clean_signal.shape[2])\n            ):\n        poli = np.poly1d(linear_corr[:, x, y])\n        clean_signal[:, x, y] = poli(clean_signal[:, x, y])\n    return clean_signal\n\ndef clean_dark(signal, dark, dt):\n    dark = np.tile(dark, (signal.shape[0], 1, 1))\n    signal -= dark * dt[:, np.newaxis, np.newaxis]\n    return signal\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **Preprocessing Function**\n","metadata":{}},{"cell_type":"code","source":"\ndef preproc(dataset, adc_info, sensor, binning=15):\n    cut_inf, cut_sup = 39, 321\n    sensor_sizes_dict = {\n        \"AIRS-CH0\": [[11250, 32, 356], [1, 32, cut_sup - cut_inf]],\n        \"FGS1\": [[135000, 32, 32], [1, 32, 32]]\n    }\n    binned_dict = {\n        \"AIRS-CH0\": [11250 // binning // 2, 282],\n        \"FGS1\": [135000 // binning // 2]\n    }\n    linear_corr_dict = {\n        \"AIRS-CH0\": (6, 32, 356),\n        \"FGS1\": (6, 32, 32)\n    }\n    planet_ids = adc_info.index\n\n    feats = []\n    for i, planet_id in tqdm(list(enumerate(planet_ids))):\n        # Load signal and calibration data\n        signal = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_signal.parquet').to_numpy()\n        dark_frame = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_calibration/dark.parquet').to_numpy()\n        dead_frame = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_calibration/dead.parquet').to_numpy()\n        flat_frame = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_calibration/flat.parquet').to_numpy()\n        linear_corr = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_calibration/linear_corr.parquet').values.astype(np.float64).reshape(linear_corr_dict[sensor])\n\n        # Reshape and calibrate the signal\n        signal = signal.reshape(sensor_sizes_dict[sensor][0]) \n        gain = adc_info.loc[planet_id, f'{sensor}_adc_gain']\n        offset = adc_info.loc[planet_id, f'{sensor}_adc_offset']\n        signal = signal / gain + offset\n\n        hot = sigma_clip(dark_frame, sigma=5, maxiters=5).mask\n\n        if sensor != \"FGS1\":\n            signal = signal[:, :, cut_inf:cut_sup]  # 11250 x 32 x 282\n            dt = np.ones(len(signal)) * 0.1\n            dt[1::2] += 4.5  # Adjust integration time\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 = np.ones(len(signal)) * 0.1\n            dt[1::2] += 0.1\n\n        signal = signal.clip(0)\n        linear_corr_signal = apply_linear_corr(linear_corr, signal)\n        signal = clean_dark(linear_corr_signal, dark_frame, dt)\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 == \"FGS1\":\n            signal = signal.reshape(\n                (sensor_sizes_dict[sensor][0][0],\n                 sensor_sizes_dict[sensor][0][1] * sensor_sizes_dict[sensor][0][2])\n            )\n\n        mean_signal = np.nanmean(signal, axis=1)\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.append(binned)\n\n    return np.stack(feats)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **Preprocessing the Data**","metadata":{}},{"cell_type":"code","source":"pre_train = np.concatenate([\n    preproc('test', test_adc_info, \"FGS1\", 30 * 12),\n    preproc('test', test_adc_info, \"AIRS-CH0\", 30)\n], axis=2)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Signal Calibration and Optimization\n\nIn this section, we calibrate the signals by detecting the transit phases and optimizing a scaling factor `s` that minimizes the mean absolute error between the polynomial fit and the adjusted signal.","metadata":{}},{"cell_type":"markdown","source":"### **Phase Detection Function**","metadata":{}},{"cell_type":"code","source":"def phase_detector(signal):\n    phase1, phase2 = None, None\n    best_drop = 0\n    for i in range(25, 75):\n        t1 = signal[i:i + 10].max() - signal[i:i + 10].min()\n        if t1 > best_drop:\n            phase1 = i + 12  # Adjusted for center of the window\n            best_drop = t1\n\n    best_drop = 0\n    for i in range(100, 125):\n        t1 = signal[i:i + 10].max() - signal[i:i + 10].min()\n        if t1 > best_drop:\n            phase2 = i - 2  # Adjusted for center of the window\n            best_drop = t1\n\n    return phase1, phase2","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### **Optimization Functions**","metadata":{}},{"cell_type":"code","source":"def try_s(signal, p1, p2, deg, s):\n    out = list(range(p1 - 30)) + list(range(p2 + 30, signal.shape[0]))\n    x, y = out, signal[out].tolist()\n    x = x + list(range(p1, p2))\n    y = y + (signal[p1:p2] * (1 + s[0])).tolist()\n    z = np.polyfit(x, y, deg)\n    p = np.poly1d(z)\n    q = np.abs(p(x) - y).mean()\n\n    if s[0] < 1e-4:\n        return q + 1e3\n\n    return q\n\ndef calibrate_train(signal):\n    p1, p2 = phase_detector(signal)\n\n    best_deg, best_score = 1, 1e12\n    for deg in range(1, 6):\n        f = partial(try_s, signal, p1, p2, deg)\n        r = minimize(f, [0.001], method='Nelder-Mead')\n        s = r.x[0]\n\n        out = list(range(p1 - 30)) + list(range(p2 + 30, signal.shape[0]))\n        x, y = out, signal[out].tolist()\n        x = x + list(range(p1, p2))\n        y = y + (signal[p1:p2] * (1 + s)).tolist()\n\n        z = np.polyfit(x, y, deg)\n        p = np.poly1d(z)\n        q = np.abs(p(x) - y).mean()\n\n        if q < best_score:\n            best_score = q\n            best_deg = deg\n\n    z = np.polyfit(x, y, best_deg)\n    p = np.poly1d(z)\n\n    return s, p(np.arange(signal.shape[0])), p1, p2","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prediction Generation\nWe apply the calibration function to each planet's signal to obtain the scaling factors used for prediction.","metadata":{}},{"cell_type":"code","source":"# Calibrate the signals and collect scaling factors\ntrain = pre_train.copy()\nall_s = []\nfor i in range(len(test_adc_info)):\n    signal = train[i, :, 1:].mean(axis=1)\n    s, p, p1, p2 = calibrate_train(signal)\n    all_s.append(s)\n\n# Generate predictions\ntrain_s = np.repeat(np.array(all_s), 283).reshape((len(all_s), 283))\ntrain_sigma = np.ones_like(train_s) * 0.00016","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Visualize the calibration for the first planet\nn = 0\nsignal = pre_train[n, :, 1:].mean(axis=1)\ns, p, p1, p2 = calibrate_train(signal)\nplt.figure(figsize=(10, 6))\nplt.scatter(range(len(signal)), signal, label='Original Data', alpha=0.5)\nplt.plot(range(len(p)), p, color='red', label='Fitted Polynomial')\nplt.axvspan(p1, p2, color='yellow', alpha=0.3, label='Transit Phase')\nplt.xlabel('Time Index')\nplt.ylabel('Signal')\nplt.title('Calibration Fit for Planet Index 0')\nplt.legend()\nplt.show()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submission ","metadata":{}},{"cell_type":"code","source":"# Load the sample submission file to use as a template\nsample_submission = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/sample_submission.csv')\n\n# Ensure predictions are non-negative (assumes 'train_s' contains the prediction values)\npredictions = train_s.clip(0)\n\n# Assume 'train_sigma' contains the uncertainties for each prediction\nuncertainties = train_sigma\n\n# Interleave predictions and uncertainties to match the expected submission format\n# Concatenate along axis=1 (column-wise) and use column names from the sample submission starting from the second column\ninterleaved_data = np.column_stack((predictions, uncertainties))\nsubmission = pd.DataFrame(interleaved_data, columns=sample_submission.columns[1:])\n\n# Ensure the index matches the test information index to correspond with submission requirements\nsubmission.index = test_adc_info.index\n\n# Save the DataFrame to a CSV file for submission\nsubmission.to_csv('submission.csv', index=False)  # Including index=False to avoid writing row indices\n\n# Display the first few rows to verify correctness\nprint(submission.head())","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Conclusion\n\nIn this notebook, we processed raw sensor data, performed calibration and optimization, and generated predictions for the planetary spectra. This approach aims to accurately model the underlying trends in the signal and provide reliable predictions for the Ariel Data Challenge 2024.","metadata":{}}]}