{"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":"markdown","source":"# Setting","metadata":{}},{"cell_type":"code","source":"!pip install astropy","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:14:59.243262Z","iopub.execute_input":"2024-11-17T06:14:59.243724Z","iopub.status.idle":"2024-11-17T06:15:10.715017Z","shell.execute_reply.started":"2024-11-17T06:14:59.24366Z","shell.execute_reply":"2024-11-17T06:15:10.713624Z"}},"outputs":[{"name":"stdout","text":"Requirement already satisfied: astropy in /opt/conda/lib/python3.10/site-packages (6.1.6)\nRequirement already satisfied: numpy>=1.23 in /opt/conda/lib/python3.10/site-packages (from astropy) (1.26.4)\nRequirement already satisfied: pyerfa>=2.0.1.1 in /opt/conda/lib/python3.10/site-packages (from astropy) (2.0.1.5)\nRequirement already satisfied: astropy-iers-data>=0.2024.10.28.0.34.7 in /opt/conda/lib/python3.10/site-packages (from astropy) (0.2024.11.11.0.32.38)\nRequirement already satisfied: PyYAML>=3.13 in /opt/conda/lib/python3.10/site-packages (from astropy) (6.0.2)\nRequirement already satisfied: packaging>=19.0 in /opt/conda/lib/python3.10/site-packages (from astropy) (21.3)\nRequirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=19.0->astropy) (3.1.2)\n","output_type":"stream"}],"execution_count":28},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport glob\nimport matplotlib.pyplot as plt \nfrom astropy.stats import sigma_clip\nfrom tqdm import tqdm\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.717215Z","iopub.execute_input":"2024-11-17T06:15:10.717616Z","iopub.status.idle":"2024-11-17T06:15:10.725335Z","shell.execute_reply.started":"2024-11-17T06:15:10.717573Z","shell.execute_reply":"2024-11-17T06:15:10.724256Z"}},"outputs":[],"execution_count":29},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"code","source":"path_folder = '/kaggle/input/ariel-data-challenge-2024/'\npath_out = '/kaggle/tmp/data_light/raw/'\nouput_dir = '/kaggle/tmp/data_light_raw/'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.727212Z","iopub.execute_input":"2024-11-17T06:15:10.727547Z","iopub.status.idle":"2024-11-17T06:15:10.737764Z","shell.execute_reply.started":"2024-11-17T06:15:10.727511Z","shell.execute_reply":"2024-11-17T06:15:10.736558Z"}},"outputs":[],"execution_count":30},{"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.')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.740368Z","iopub.execute_input":"2024-11-17T06:15:10.740766Z","iopub.status.idle":"2024-11-17T06:15:10.748937Z","shell.execute_reply.started":"2024-11-17T06:15:10.74072Z","shell.execute_reply":"2024-11-17T06:15:10.747754Z"}},"outputs":[{"name":"stdout","text":"Directory /kaggle/tmp/data_light/raw/ already exists.\n","output_type":"stream"}],"execution_count":31},{"cell_type":"markdown","source":"## From here:","metadata":{}},{"cell_type":"code","source":"CHUNKS_SIZE = 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.75018Z","iopub.execute_input":"2024-11-17T06:15:10.750503Z","iopub.status.idle":"2024-11-17T06:15:10.759443Z","shell.execute_reply.started":"2024-11-17T06:15:10.750465Z","shell.execute_reply":"2024-11-17T06:15:10.758397Z"}},"outputs":[],"execution_count":32},{"cell_type":"markdown","source":"# Functions","metadata":{}},{"cell_type":"code","source":"def ADC_convert(signal, gain, offset):\n    signal = signal.astype(np.float64)\n    signal /= gain\n    signal += offset\n    return signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.761225Z","iopub.execute_input":"2024-11-17T06:15:10.7616Z","iopub.status.idle":"2024-11-17T06:15:10.770391Z","shell.execute_reply.started":"2024-11-17T06:15:10.761555Z","shell.execute_reply":"2024-11-17T06:15:10.769288Z"}},"outputs":[],"execution_count":33},{"cell_type":"code","source":"def 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    signal = np.ma.masked_where(dead, signal)\n    signal = np.ma.masked_where(hot, signal)\n    return signal","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.772102Z","iopub.execute_input":"2024-11-17T06:15:10.772506Z","iopub.status.idle":"2024-11-17T06:15:10.782913Z","shell.execute_reply.started":"2024-11-17T06:15:10.772464Z","shell.execute_reply":"2024-11-17T06:15:10.781942Z"}},"outputs":[],"execution_count":34},{"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.784321Z","iopub.execute_input":"2024-11-17T06:15:10.784867Z","iopub.status.idle":"2024-11-17T06:15:10.795779Z","shell.execute_reply.started":"2024-11-17T06:15:10.784814Z","shell.execute_reply":"2024-11-17T06:15:10.794228Z"}},"outputs":[],"execution_count":35},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.797485Z","iopub.execute_input":"2024-11-17T06:15:10.798025Z","iopub.status.idle":"2024-11-17T06:15:10.809087Z","shell.execute_reply.started":"2024-11-17T06:15:10.797973Z","shell.execute_reply":"2024-11-17T06:15:10.807642Z"}},"outputs":[],"execution_count":36},{"cell_type":"code","source":"def get_cds(signal):\n    cds = signal[:,1::2,:,:] - signal[:,::2,:,:]\n    return cds","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.814282Z","iopub.execute_input":"2024-11-17T06:15:10.814662Z","iopub.status.idle":"2024-11-17T06:15:10.82039Z","shell.execute_reply.started":"2024-11-17T06:15:10.814623Z","shell.execute_reply":"2024-11-17T06:15:10.81949Z"}},"outputs":[],"execution_count":37},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.821816Z","iopub.execute_input":"2024-11-17T06:15:10.822865Z","iopub.status.idle":"2024-11-17T06:15:10.832254Z","shell.execute_reply.started":"2024-11-17T06:15:10.822823Z","shell.execute_reply":"2024-11-17T06:15:10.83082Z"}},"outputs":[],"execution_count":38},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.834458Z","iopub.execute_input":"2024-11-17T06:15:10.834941Z","iopub.status.idle":"2024-11-17T06:15:10.84616Z","shell.execute_reply.started":"2024-11-17T06:15:10.8349Z","shell.execute_reply":"2024-11-17T06:15:10.844793Z"}},"outputs":[],"execution_count":39},{"cell_type":"code","source":"## we will start by getting the index of the training data:\ndef get_index(files, CHUNKS_SIZE):\n    index = []\n    for file in files :\n        file_name = file.split('/')[-1]\n        if file_name.split('_')[0] == 'AIRS-CH0' and file_name.split('_')[1] == 'signal.parquet':\n            file_index = os.path.basename(os.path.dirname(file))\n            index.append(int(file_index))\n    index = np.array(index)\n    index = np.sort(index) \n    # credit to DennisSakva\n    index=np.array_split(index, len(index)//CHUNKS_SIZE)\n    \n    return index","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.847808Z","iopub.execute_input":"2024-11-17T06:15:10.84895Z","iopub.status.idle":"2024-11-17T06:15:10.858492Z","shell.execute_reply.started":"2024-11-17T06:15:10.848894Z","shell.execute_reply":"2024-11-17T06:15:10.857286Z"}},"outputs":[],"execution_count":40},{"cell_type":"markdown","source":"# Load Data","metadata":{}},{"cell_type":"code","source":"files = glob.glob(os.path.join(path_folder + 'train/', '*/*'))\n\nindex = get_index(files[:22],CHUNKS_SIZE)  ## 48 is hardcoded here but please feel free to remove it if you want to do it for the entire dataset\n\ntrain_adc_info = pd.read_csv(os.path.join(path_folder, 'train_adc_info.csv'))\ntrain_adc_info = train_adc_info.set_index('planet_id')\naxis_info = pd.read_parquet(os.path.join(path_folder,'axis_info.parquet'))\nDO_MASK = True\nDO_THE_NL_CORR = False\nDO_DARK = True\nDO_FLAT = True\nTIME_BINNING = True\n\ncut_inf, cut_sup = 39, 321\nl = cut_sup - cut_inf\n\nfor n, index_chunk in enumerate(tqdm(index)):\n    AIRS_CH0_clean = np.zeros((CHUNKS_SIZE, 11250, 32, l))\n    FGS1_clean = np.zeros((CHUNKS_SIZE, 135000, 32, 32))\n    \n    for i in range (CHUNKS_SIZE) : \n        df = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/AIRS-CH0_signal.parquet'))\n        signal = df.values.astype(np.float64).reshape((df.shape[0], 32, 356))\n        gain = train_adc_info['AIRS-CH0_adc_gain'].loc[index_chunk[i]]\n        offset = train_adc_info['AIRS-CH0_adc_offset'].loc[index_chunk[i]]\n        signal = ADC_convert(signal, gain, offset)\n        dt_airs = axis_info['AIRS-CH0-integration_time'].dropna().values\n        dt_airs[1::2] += 0.1\n        chopped_signal = signal[:, :, cut_inf:cut_sup]\n        del signal, df\n        \n        # CLEANING THE DATA: AIRS\n        flat = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/AIRS-CH0_calibration/flat.parquet')).values.astype(np.float64).reshape((32, 356))[:, cut_inf:cut_sup]\n        dark = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/AIRS-CH0_calibration/dark.parquet')).values.astype(np.float64).reshape((32, 356))[:, cut_inf:cut_sup]\n        dead_airs = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/AIRS-CH0_calibration/dead.parquet')).values.astype(np.float64).reshape((32, 356))[:, cut_inf:cut_sup]\n        linear_corr = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/AIRS-CH0_calibration/linear_corr.parquet')).values.astype(np.float64).reshape((6, 32, 356))[:, :, cut_inf:cut_sup]\n        \n        if DO_MASK:\n            chopped_signal = mask_hot_dead(chopped_signal, dead_airs, dark)\n            AIRS_CH0_clean[i] = chopped_signal\n        else:\n            AIRS_CH0_clean[i] = chopped_signal\n            \n        if DO_THE_NL_CORR: \n            linear_corr_signal = apply_linear_corr(linear_corr,AIRS_CH0_clean[i])\n            AIRS_CH0_clean[i,:, :, :] = linear_corr_signal\n        del linear_corr\n        \n        if DO_DARK: \n            cleaned_signal = clean_dark(AIRS_CH0_clean[i], dead_airs, dark, dt_airs)\n            AIRS_CH0_clean[i] = cleaned_signal\n        else: \n            pass\n        del dark\n        \n        df = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/FGS1_signal.parquet'))\n        fgs_signal = df.values.astype(np.float64).reshape((df.shape[0], 32, 32))\n        \n        FGS1_gain = train_adc_info['FGS1_adc_gain'].loc[index_chunk[i]]\n        FGS1_offset = train_adc_info['FGS1_adc_offset'].loc[index_chunk[i]]\n        \n        fgs_signal = ADC_convert(fgs_signal, FGS1_gain, FGS1_offset)\n        dt_fgs1 = np.ones(len(fgs_signal))*0.1\n        dt_fgs1[1::2] += 0.1\n        chopped_FGS1 = fgs_signal\n        del fgs_signal, df\n        \n        # CLEANING THE DATA: FGS1\n        flat = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/FGS1_calibration/flat.parquet')).values.astype(np.float64).reshape((32, 32))\n        dark = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/FGS1_calibration/dark.parquet')).values.astype(np.float64).reshape((32, 32))\n        dead_fgs1 = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/FGS1_calibration/dead.parquet')).values.astype(np.float64).reshape((32, 32))\n        linear_corr = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/FGS1_calibration/linear_corr.parquet')).values.astype(np.float64).reshape((6, 32, 32))\n        \n        if DO_MASK:\n            chopped_FGS1 = mask_hot_dead(chopped_FGS1, dead_fgs1, dark)\n            FGS1_clean[i] = chopped_FGS1\n        else:\n            FGS1_clean[i] = chopped_FGS1\n\n        if DO_THE_NL_CORR: \n            linear_corr_signal = apply_linear_corr(linear_corr,FGS1_clean[i])\n            FGS1_clean[i,:, :, :] = linear_corr_signal\n        del linear_corr\n        \n        if DO_DARK: \n            cleaned_signal = clean_dark(FGS1_clean[i], dead_fgs1, dark,dt_fgs1)\n            FGS1_clean[i] = cleaned_signal\n        else: \n            pass\n        del dark\n        \n    # SAVE DATA AND FREE SPACE\n    AIRS_cds = get_cds(AIRS_CH0_clean)\n    FGS1_cds = get_cds(FGS1_clean)\n    \n    del AIRS_CH0_clean, FGS1_clean\n    \n    ## (Optional) Time Binning to reduce space\n    if TIME_BINNING:\n        AIRS_cds_binned = bin_obs(AIRS_cds,binning=30)\n        FGS1_cds_binned = bin_obs(FGS1_cds,binning=30*12)\n    else:\n        AIRS_cds = AIRS_cds.transpose(0,1,3,2) ## this is important to make it consistent for flat fielding, but you can always change it\n        AIRS_cds_binned = AIRS_cds\n        FGS1_cds = FGS1_cds.transpose(0,1,3,2)\n        FGS1_cds_binned = FGS1_cds\n    \n    del AIRS_cds, FGS1_cds\n    \n    for i in range (CHUNKS_SIZE):\n        flat_airs = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/AIRS-CH0_calibration/flat.parquet')).values.astype(np.float64).reshape((32, 356))[:, cut_inf:cut_sup]\n        flat_fgs = pd.read_parquet(os.path.join(path_folder,f'train/{index_chunk[i]}/FGS1_calibration/flat.parquet')).values.astype(np.float64).reshape((32, 32))\n        if DO_FLAT:\n            corrected_AIRS_cds_binned = correct_flat_field(flat_airs,dead_airs, AIRS_cds_binned[i])\n            AIRS_cds_binned[i] = corrected_AIRS_cds_binned\n            corrected_FGS1_cds_binned = correct_flat_field(flat_fgs,dead_fgs1, FGS1_cds_binned[i])\n            FGS1_cds_binned[i] = corrected_FGS1_cds_binned\n        else:\n            pass\n\n    ## save data\n    np.save(os.path.join(path_out, 'AIRS_clean_train_{}.npy'.format(n)), AIRS_cds_binned)\n    np.save(os.path.join(path_out, 'FGS1_train_{}.npy'.format(n)), FGS1_cds_binned)\n    del AIRS_cds_binned\n    del FGS1_cds_binned","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:15:10.860425Z","iopub.execute_input":"2024-11-17T06:15:10.861165Z","iopub.status.idle":"2024-11-17T06:17:05.274822Z","shell.execute_reply.started":"2024-11-17T06:15:10.861112Z","shell.execute_reply":"2024-11-17T06:17:05.273607Z"}},"outputs":[{"name":"stderr","text":"100%|██████████| 6/6 [01:48<00:00, 18.12s/it]\n","output_type":"stream"}],"execution_count":41},{"cell_type":"code","source":"def load_data (file, chunk_size, nb_files) : \n    data0 = np.load(file + '_0.npy')\n    data_all = np.zeros((nb_files*chunk_size, data0.shape[1], data0.shape[2], data0.shape[3]))\n    data_all[:chunk_size] = data0\n    for i in range (1, nb_files) : \n        data_all[i*chunk_size:(i+1)*chunk_size] = np.load(file + '_{}.npy'.format(i))\n    return data_all \n\ndata_train = load_data(path_out + 'AIRS_clean_train', CHUNKS_SIZE, len(index)) \ndata_train_FGS = load_data(path_out + 'FGS1_train', CHUNKS_SIZE, len(index))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:17:05.276373Z","iopub.execute_input":"2024-11-17T06:17:05.276781Z","iopub.status.idle":"2024-11-17T06:17:05.386955Z","shell.execute_reply.started":"2024-11-17T06:17:05.276742Z","shell.execute_reply":"2024-11-17T06:17:05.385983Z"}},"outputs":[],"execution_count":42},{"cell_type":"code","source":"np.save('/kaggle/working/' + 'data_train.npy', data_train)\nnp.save('/kaggle/working/' + 'data_train_FGS.npy', data_train_FGS)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:17:05.388202Z","iopub.execute_input":"2024-11-17T06:17:05.388548Z","iopub.status.idle":"2024-11-17T06:17:05.458921Z","shell.execute_reply.started":"2024-11-17T06:17:05.388512Z","shell.execute_reply":"2024-11-17T06:17:05.457333Z"}},"outputs":[],"execution_count":43},{"cell_type":"code","source":"print('Shape of the training datasset: \\t')\nprint('\\n For AIRS-CH0:', data_train.shape)\nprint('\\n For FGS1:', data_train_FGS.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-17T06:17:05.460458Z","iopub.execute_input":"2024-11-17T06:17:05.460986Z","iopub.status.idle":"2024-11-17T06:17:05.466899Z","shell.execute_reply.started":"2024-11-17T06:17:05.46093Z","shell.execute_reply":"2024-11-17T06:17:05.465812Z"}},"outputs":[{"name":"stdout","text":"Shape of the training datasset: \t\n\n For AIRS-CH0: (6, 187, 282, 32)\n\n For FGS1: (6, 187, 32, 32)\n","output_type":"stream"}],"execution_count":44}]}