{"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":"none","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom tqdm import tqdm\nimport joblib\n\nfrom sklearn.linear_model import Ridge\nfrom sklearn.metrics import r2_score, mean_squared_error\nimport itertools\n\nfrom scipy.optimize import minimize\nfrom scipy import optimize\n\nfrom astropy.stats import sigma_clip","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:39:36.648146Z","iopub.execute_input":"2024-11-08T13:39:36.648569Z","iopub.status.idle":"2024-11-08T13:39:37.903013Z","shell.execute_reply.started":"2024-11-08T13:39:36.648525Z","shell.execute_reply":"2024-11-08T13:39:37.901801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DEBUG = False","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:47:56.737365Z","iopub.execute_input":"2024-11-08T13:47:56.738165Z","iopub.status.idle":"2024-11-08T13:47:56.7453Z","shell.execute_reply.started":"2024-11-08T13:47:56.738102Z","shell.execute_reply":"2024-11-08T13:47:56.743982Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"nums_parallel = 4\n\ncut_inf, cut_sup = 39, 321","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:40:02.514437Z","iopub.execute_input":"2024-11-08T13:40:02.515068Z","iopub.status.idle":"2024-11-08T13:40:02.520656Z","shell.execute_reply.started":"2024-11-08T13:40:02.515022Z","shell.execute_reply":"2024-11-08T13:40:02.519476Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DATA_ROOT = '/kaggle/input/ariel-data-challenge-2024/'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:40:37.552286Z","iopub.execute_input":"2024-11-08T13:40:37.553375Z","iopub.status.idle":"2024-11-08T13:40:37.558518Z","shell.execute_reply.started":"2024-11-08T13:40:37.553325Z","shell.execute_reply":"2024-11-08T13:40:37.557092Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataset = 'train'\nadc_info = pd.read_csv(f'{DATA_ROOT}/{dataset}_adc_info.csv',index_col='planet_id')\naxis_info = pd.read_parquet(f'{DATA_ROOT}/axis_info.parquet')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:41:52.032418Z","iopub.execute_input":"2024-11-08T13:41:52.032935Z","iopub.status.idle":"2024-11-08T13:41:52.237662Z","shell.execute_reply.started":"2024-11-08T13:41:52.032892Z","shell.execute_reply":"2024-11-08T13:41:52.236504Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if DEBUG:\n    adc_info = adc_info.head().copy()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:44:19.852018Z","iopub.execute_input":"2024-11-08T13:44:19.852525Z","iopub.status.idle":"2024-11-08T13:44:19.859236Z","shell.execute_reply.started":"2024-11-08T13:44:19.85248Z","shell.execute_reply":"2024-11-08T13:44:19.857619Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"adc_info.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:44:25.287576Z","iopub.execute_input":"2024-11-08T13:44:25.288024Z","iopub.status.idle":"2024-11-08T13:44:25.295927Z","shell.execute_reply.started":"2024-11-08T13:44:25.287981Z","shell.execute_reply":"2024-11-08T13:44:25.294611Z"}},"outputs":[],"execution_count":null},{"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\ndef preproc(dataset, adc_info, sensor, binning = 15):\n    cut_inf, cut_sup = 0, 356\n    sensor_sizes_dict = {\"AIRS-CH0\":[[11250, 32, 356], [1, 32, cut_sup-cut_inf]], \"FGS1\":[[135000, 32, 32], [1, 32, 32]]}\n    binned_dict = {\"AIRS-CH0\":[11250 // binning // 2, cut_sup-cut_inf], \"FGS1\":[135000 // binning // 2]}\n    linear_corr_dict = {\"AIRS-CH0\":(6, 32, 356), \"FGS1\":(6, 32, 32)}\n    planet_ids = adc_info.index\n    \n    feats = []\n    for i, planet_id in tqdm(list(enumerate(planet_ids))):\n        signal = pd.read_parquet(f'{DATA_ROOT}/{dataset}/{planet_id}/{sensor}_signal.parquet').to_numpy()\n        dark_frame = pd.read_parquet(f'{DATA_ROOT}/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/dark.parquet', engine='pyarrow').to_numpy()\n        dead_frame = pd.read_parquet(f'{DATA_ROOT}/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/dead.parquet', engine='pyarrow').to_numpy()\n        flat_frame = pd.read_parquet(f'{DATA_ROOT}/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/flat.parquet', engine='pyarrow').to_numpy()\n        linear_corr = pd.read_parquet(f'{DATA_ROOT}/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/linear_corr.parquet').values.astype(np.float64).reshape(linear_corr_dict[sensor])\n\n        signal = signal.reshape(sensor_sizes_dict[sensor][0]) \n        gain = adc_info[f'{sensor}_adc_gain'].values[i]\n        offset = adc_info[f'{sensor}_adc_offset'].values[i]\n        signal = signal / gain + offset\n        \n        hot = sigma_clip(\n            dark_frame, sigma=5, maxiters=5\n        ).mask\n        \n        if sensor != \"FGS1\":\n            signal = signal[:, :, cut_inf:cut_sup] \n            dt = np.ones(len(signal))*0.1 \n            dt[1::2] += 4.5 #@bilzard idea\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) #@graySnow idea\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        \n        if sensor == \"FGS1\":\n            signal = signal[:,10:22,10:22] # **** updates ****\n            signal = signal.reshape(sensor_sizes_dict[sensor][0][0],144) # # **** updates ****\n\n        if sensor != \"FGS1\":\n            signal = signal[:,10:22,:] # **** updates ****\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":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:45:27.094102Z","iopub.execute_input":"2024-11-08T13:45:27.094574Z","iopub.status.idle":"2024-11-08T13:45:27.122771Z","shell.execute_reply.started":"2024-11-08T13:45:27.094532Z","shell.execute_reply":"2024-11-08T13:45:27.121341Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"f_raw_train = preproc(f'{dataset}', adc_info, \"FGS1\", 30*12)\nf_raw_train.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:45:28.225309Z","iopub.execute_input":"2024-11-08T13:45:28.225751Z","iopub.status.idle":"2024-11-08T13:45:57.605349Z","shell.execute_reply.started":"2024-11-08T13:45:28.22571Z","shell.execute_reply":"2024-11-08T13:45:57.604171Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"a_raw_train = preproc(f'{dataset}', adc_info, \"AIRS-CH0\", 30)\na_raw_train.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:46:25.65944Z","iopub.execute_input":"2024-11-08T13:46:25.659898Z","iopub.status.idle":"2024-11-08T13:46:56.81169Z","shell.execute_reply.started":"2024-11-08T13:46:25.659855Z","shell.execute_reply":"2024-11-08T13:46:56.810503Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.savez('processed_train.npz', f_raw_train=f_raw_train, a_raw_train=a_raw_train)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:47:02.71727Z","iopub.execute_input":"2024-11-08T13:47:02.717749Z","iopub.status.idle":"2024-11-08T13:47:02.730189Z","shell.execute_reply.started":"2024-11-08T13:47:02.717705Z","shell.execute_reply":"2024-11-08T13:47:02.728866Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-08T13:47:05.329045Z","iopub.execute_input":"2024-11-08T13:47:05.329607Z","iopub.status.idle":"2024-11-08T13:47:06.516721Z","shell.execute_reply.started":"2024-11-08T13:47:05.329551Z","shell.execute_reply":"2024-11-08T13:47:06.515058Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}