{"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":"gpu","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"}],"dockerImageVersionId":30762,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import itertools\nimport os\nimport random\nfrom functools import partial\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport scipy.stats\nfrom scipy.optimize import minimize\nfrom sklearn.linear_model import Ridge\nfrom sklearn.metrics import mean_squared_error, r2_score\nfrom sklearn.model_selection import cross_val_predict\nfrom tqdm.auto import tqdm","metadata":{"execution":{"iopub.status.busy":"2024-09-03T05:28:40.787964Z","iopub.execute_input":"2024-09-03T05:28:40.788258Z","iopub.status.idle":"2024-09-03T05:28:42.52095Z","shell.execute_reply.started":"2024-09-03T05:28:40.788225Z","shell.execute_reply":"2024-09-03T05:28:42.519978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cupy as cp\nimport polars as pl\nfrom astropy.stats import sigma_clip\n\n\ndef ADC_convert(signal, gain, offset):\n    signal /= gain\n    signal += offset\n    return signal.clip(0)\n\n\ndef mask_hot_dead(signal, dead, dark):\n    hot = sigma_clip(dark.get(), sigma=5, maxiters=5).mask\n    hot_mask = cp.asarray(hot.reshape((-1,)))\n    dead_mask = cp.asarray((dead == 1.0).reshape((-1,)))\n\n    signal[:, hot_mask] = cp.nan\n    signal[:, dead_mask] = cp.nan\n    return signal\n\n\ndef clean_dark(signal, dark, dt):\n    dark_current = dt[:, cp.newaxis] * dark\n    signal -= dark_current\n    return signal\n\n\ndef clean_flat(signal, flat):\n    signal = signal / flat\n    return signal\n\n\ndef apply_linear_corr(linear_corr, signal):\n    result = linear_corr[5, :]\n    for i in range(4, -1, -1):\n        result = result * signal + linear_corr[i, :]\n    return result\n\n\ndef bin_obs(signal, binning):\n    signal_binned = cp.zeros((signal.shape[0] // binning, signal.shape[1]))\n    for i in range(signal.shape[0] // binning):\n        signal_binned[i, :] = cp.mean(signal[i * binning : (i + 1) * binning, :], axis=0)\n    return signal_binned\n\n\ndef load_signal_data(planet_id, dataset, instrument, adc_info, step_size=30):\n    binning = step_size if instrument == \"AIRS-CH0\" else step_size * 12  # 4642\n    img_size = (32, 356) if instrument == \"AIRS-CH0\" else (32, 32)\n    gain, offset = (\n        adc_info.loc[planet_id][f\"{instrument}_adc_gain\"],\n        adc_info.loc[planet_id][f\"{instrument}_adc_offset\"],\n    )\n    signal_file = f\"{PATH}/{dataset}/{planet_id}/{instrument}_signal.parquet\"\n    calibration_file = f\"{PATH}/{dataset}/{planet_id}/{instrument}_calibration\"\n\n    signal = cp.array(pl.read_parquet(signal_file).cast(pl.Float32).to_numpy())\n    flat = cp.array(pl.read_parquet(f\"{calibration_file}/flat.parquet\").cast(pl.Float32).to_numpy().reshape((1, -1)))\n    dark = cp.array(pl.read_parquet(f\"{calibration_file}/dark.parquet\").cast(pl.Float32).to_numpy().reshape((1, -1)))\n    dead = cp.array(pl.read_parquet(f\"{calibration_file}/dead.parquet\").cast(pl.Float32).to_numpy().reshape((1, -1)))\n    linear_corr = cp.array(\n        pl.read_parquet(f\"{calibration_file}/linear_corr.parquet\").cast(pl.Float32).to_numpy().reshape(6, -1)\n    )\n    if instrument == \"AIRS-CH0\":\n        dt = cp.array(airs_dt_info)\n        dt[1::2] += 0.1\n    else:\n        dt = cp.ones(len(signal)) * 0.1\n        dt[1::2] += 4.5\n\n    signal = ADC_convert(signal, gain, offset)\n    signal = clean_dark(signal, dark, dt)\n    signal = apply_linear_corr(linear_corr, signal)\n    signal = mask_hot_dead(signal, dead, dark)\n    signal = clean_flat(signal, flat)\n    \n    signal = signal.reshape((signal.shape[0], img_size[0], img_size[1]))\n    signal = cp.nanmean(signal, axis=1)\n    signal = signal[1::2] - signal[0::2]\n    signal = bin_obs(signal, binning)\n\n    return signal\n\n\nPATH = \"/kaggle/input/ariel-data-challenge-2024/\"\ntrain_adc_info = pd.read_csv(f\"{PATH}/train_adc_info.csv\", index_col=\"planet_id\")\ntest_adc_info = pd.read_csv(f\"{PATH}/test_adc_info.csv\", index_col=\"planet_id\")\naxis_info = pl.read_parquet(f\"{PATH}/axis_info.parquet\").to_pandas()\nairs_dt_info = axis_info[\"AIRS-CH0-integration_time\"].dropna().to_numpy()","metadata":{"execution":{"iopub.status.busy":"2024-09-03T05:30:09.047252Z","iopub.execute_input":"2024-09-03T05:30:09.04822Z","iopub.status.idle":"2024-09-03T05:30:09.091164Z","shell.execute_reply.started":"2024-09-03T05:30:09.04818Z","shell.execute_reply":"2024-09-03T05:30:09.090158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def generate_signal(mode, adc_info):\n    signal = []\n    for planet_id in tqdm(adc_info.index):\n        signal.append(\n            np.concatenate(\n                (\n                    load_signal_data(planet_id, mode, \"FGS1\", adc_info),\n                    load_signal_data(planet_id, mode, \"AIRS-CH0\", adc_info)\n                ), axis=-1,\n            )\n         )\n    return np.stack(signal)","metadata":{"execution":{"iopub.status.busy":"2024-09-03T05:31:27.444778Z","iopub.execute_input":"2024-09-03T05:31:27.445163Z","iopub.status.idle":"2024-09-03T05:31:27.451097Z","shell.execute_reply.started":"2024-09-03T05:31:27.445125Z","shell.execute_reply":"2024-09-03T05:31:27.449937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_signal = generate_signal(\"train\", train_adc_info)\ntest_signal = generate_signal(\"test\", test_adc_info)","metadata":{"execution":{"iopub.status.busy":"2024-09-03T05:31:27.626196Z","iopub.execute_input":"2024-09-03T05:31:27.626603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"np.save(\"train_signal.npy\", train_signal)\nnp.save(\"test_signal.npy\", test_signal)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]}]}