{"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"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np  # linear algebra\nimport pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\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\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-19T18:08:32.135611Z","iopub.execute_input":"2024-10-19T18:08:32.136068Z","iopub.status.idle":"2024-10-19T18:08:35.142162Z","shell.execute_reply.started":"2024-10-19T18:08:32.136024Z","shell.execute_reply":"2024-10-19T18:08:35.140478Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport itertools\nfrom concurrent.futures import ThreadPoolExecutor\nfrom scipy.optimize import minimize\nfrom tqdm import tqdm\nfrom astropy.stats import sigma_clip\nfrom pqdm.processes import pqdm\n\nclass Calibrator:\n    cut_inf = 39\n    cut_sup = 321\n    sensor_to_sizes_dict = {\n        \"AIRS-CH0\": [[11250, 32, 356], [1, 32, cut_sup - cut_inf]],\n        \"FGS1\": [[135000, 32, 32], [1, 32, 32]],\n    }\n    sensor_to_linear_corr_dict = {\"AIRS-CH0\": (6, 32, 356), \"FGS1\": (6, 32, 32)}\n\n    def __init__(self, dataset, planet_id, sensor, adc_info):\n        self.dataset = dataset\n        self.planet_id = planet_id\n        self.sensor = sensor\n        self.adc_info = adc_info\n\n    def _apply_linear_corr(self, linear_corr, clean_signal):\n        linear_corr = np.flip(linear_corr, axis=0)\n        for x, y in itertools.product(range(clean_signal.shape[1]), range(clean_signal.shape[2])):\n            poli = np.poly1d(linear_corr[:, x, y])\n            clean_signal[:, x, y] = poli(clean_signal[:, x, y])\n        return clean_signal\n\n    def _clean_dark(self, 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\n    def get_calibrated_signal(self):\n        # Load data\n        signal = pd.read_parquet(f\"/kaggle/input/ariel-data-challenge-2024/{self.dataset}/{self.planet_id}/{self.sensor}_signal.parquet\").to_numpy()\n        dark_frame = pd.read_parquet(f\"/kaggle/input/ariel-data-challenge-2024/{self.dataset}/{self.planet_id}/{self.sensor}_calibration/dark.parquet\", engine=\"pyarrow\").to_numpy()\n        dead_frame = pd.read_parquet(f\"/kaggle/input/ariel-data-challenge-2024/{self.dataset}/{self.planet_id}/{self.sensor}_calibration/dead.parquet\", engine=\"pyarrow\").to_numpy()\n        flat_frame = pd.read_parquet(f\"/kaggle/input/ariel-data-challenge-2024/{self.dataset}/{self.planet_id}/{self.sensor}_calibration/flat.parquet\", engine=\"pyarrow\").to_numpy()\n        linear_corr = pd.read_parquet(f\"/kaggle/input/ariel-data-challenge-2024/{self.dataset}/{self.planet_id}/{self.sensor}_calibration/linear_corr.parquet\").values.astype(np.float64).reshape(self.sensor_to_linear_corr_dict[self.sensor])\n\n        signal = signal.reshape(self.sensor_to_sizes_dict[self.sensor][0])\n        gain = self.adc_info.loc[self.planet_id, f\"{self.sensor}_adc_gain\"]\n        offset = self.adc_info.loc[self.planet_id, f\"{self.sensor}_adc_offset\"]\n        signal = signal / gain + offset\n\n        # Hot pixel removal\n        hot = sigma_clip(dark_frame, sigma=5, maxiters=5).mask\n\n        if self.sensor == \"AIRS-CH0\":\n            signal = signal[:, :, self.cut_inf : self.cut_sup]\n            dt = np.ones(len(signal)) * 0.1\n            dt[1::2] += 4.5\n            linear_corr = linear_corr[:, :, self.cut_inf : self.cut_sup]\n            dark_frame = dark_frame[:, self.cut_inf : self.cut_sup]\n            dead_frame = dead_frame[:, self.cut_inf : self.cut_sup]\n            flat_frame = flat_frame[:, self.cut_inf : self.cut_sup]\n            hot = hot[:, self.cut_inf : self.cut_sup]\n        elif self.sensor == \"FGS1\":\n            dt = np.ones(len(signal)) * 0.1\n            dt[1::2] += 0.1\n\n        signal = signal.clip(0)\n        linear_corr_signal = self._apply_linear_corr(linear_corr, signal)\n        signal = self._clean_dark(linear_corr_signal, dark_frame, dt)\n\n        flat = flat_frame.reshape(self.sensor_to_sizes_dict[self.sensor][1])\n        flat[dead_frame.reshape(self.sensor_to_sizes_dict[self.sensor][1])] = np.nan\n        flat[hot.reshape(self.sensor_to_sizes_dict[self.sensor][1])] = np.nan\n        signal = signal / flat\n        return signal\n\n\nclass Preprocessor:\n    sensor_to_binning = {\"AIRS-CH0\": 30, \"FGS1\": 30 * 12}\n    sensor_to_binned_dict = {\n        \"AIRS-CH0\": [11250 // sensor_to_binning[\"AIRS-CH0\"] // 2, 282],\n        \"FGS1\": [135000 // sensor_to_binning[\"FGS1\"] // 2],\n    }\n\n    def __init__(self, dataset, planet_id, sensor, adc_info):\n        self.dataset = dataset\n        self.planet_id = planet_id\n        self.sensor = sensor\n        self.binning = self.sensor_to_binning[sensor]\n\n    def preprocess_signal(self):\n        signal = Calibrator(self.dataset, self.planet_id, self.sensor, adc_info).get_calibrated_signal()\n\n        if self.sensor == \"AIRS-CH0\":\n            signal = signal[:, 10:22, :]\n        elif self.sensor == \"FGS1\":\n            signal = signal[:, 10:22, 10:22]\n            signal = signal.reshape(signal.shape[0], signal.shape[1] * signal.shape[2])\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((self.sensor_to_binned_dict[self.sensor]))\n        for j in range(cds_signal.shape[0] // self.binning):\n            binned[j] = cds_signal[j * self.binning : j * self.binning + self.binning].mean(axis=0)\n\n        if self.sensor == \"FGS1\":\n            binned = binned.reshape((binned.shape[0], 1))\n\n        return binned\n\n\ndef phase_detector(signal):\n    MIN = np.argmin(signal[30:140]) + 30\n    signal1 = signal[:MIN]\n    signal2 = signal[MIN:]\n\n    first_derivative1 = np.gradient(signal1)\n    first_derivative1 /= first_derivative1.max()\n    first_derivative2 = np.gradient(signal2)\n    first_derivative2 /= first_derivative2.max()\n\n    phase1 = np.argmin(first_derivative1)\n    phase2 = np.argmax(first_derivative2) + MIN\n\n    return phase1, phase2\n\n\ndef predict_spectra(signal):\n    def objective_to_minimize(s):\n        delta = 2\n        x = list(range(signal.shape[0] - delta * 4))\n        y = (\n            signal[:phase1 - delta].tolist()\n            + (signal[phase1 + delta:phase2 - delta] * (1 + s)).tolist()\n            + signal[phase2 + delta:].tolist()\n        )\n\n        # Fit polynomial and return mean error\n        z = np.polyfit(x, y, deg=3)\n        p = np.poly1d(z)\n        q = np.abs(p(x) - y).mean()\n        return q\n        \n\n    signal = signal[:, 1:].mean(axis=1)\n    phase1, phase2 = phase_detector(signal)\n    s = minimize(fun=objective_to_minimize, x0=[0.0001], method=\"Nelder-Mead\").x[0]\n    return s\n\n\n# Main Execution\ndataset = \"test\"\nadc_info = pd.read_csv(f\"/kaggle/input/ariel-data-challenge-2024/{dataset}_adc_info.csv\", index_col=\"planet_id\")\nplanet_ids = adc_info.index\n\n# Use ThreadPoolExecutor for preprocessing\ndef preprocess_for_all_sensors(sensor):\n    args = [dict(dataset=dataset, planet_id=planet_id, sensor=sensor, adc_info=adc_info) for planet_id in planet_ids]\n    with ThreadPoolExecutor() as executor:\n        return list(executor.map(lambda x: Preprocessor(**x).preprocess_signal(), args))\n\npreprocessed_signal_fgs1 = preprocess_for_all_sensors(\"FGS1\")\npreprocessed_signal_airs_ch0 = preprocess_for_all_sensors(\"AIRS-CH0\")\n\n# Combine preprocessed signals\npreprocessed_signal = np.concatenate([np.stack(preprocessed_signal_fgs1), np.stack(preprocessed_signal_airs_ch0)], axis=2)\n\n# Predict spectra for all signals\npredictions_spectra = [predict_spectra(signal) for signal in preprocessed_signal]\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-10-19T18:08:35.145364Z","iopub.execute_input":"2024-10-19T18:08:35.14641Z","iopub.status.idle":"2024-10-19T18:08:49.458814Z","shell.execute_reply.started":"2024-10-19T18:08:35.14635Z","shell.execute_reply":"2024-10-19T18:08:49.457514Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nsample_submission = pd.read_csv(\n    \"/kaggle/input/ariel-data-challenge-2024/sample_submission.csv\",\n    index_col=\"planet_id\",\n)\n\npredictions_spectra = np.array(predictions_spectra)\n\n# Ensure that the total size of predictions_spectra matches what is required for reshaping\nrequired_size = len(predictions_spectra) * 283\nif predictions_spectra.size < required_size:\n    # Repeat elements if needed\n    predictions_spectra = np.repeat(predictions_spectra, 283).reshape(\n        (len(predictions_spectra), 283)\n    )\nelif predictions_spectra.size == required_size:\n    predictions_spectra = predictions_spectra.reshape((len(predictions_spectra), 283))\nelse:\n    # If there are too many elements, truncate the array\n    predictions_spectra = predictions_spectra[:required_size].reshape((len(predictions_spectra), 283))\n\n# Clip the predictions_spectra values to be non-negative\npredictions_spectra = predictions_spectra.clip(0)\n\n# Create sigmas array with the same shape as predictions_spectra\nsigmas = np.ones_like(predictions_spectra) * 0.000145\n\n# Create the final submission DataFrame\nsubmission = pd.DataFrame(\n    np.concatenate([predictions_spectra, sigmas], axis=1),\n    columns=sample_submission.columns,\n)\nsubmission.index = sample_submission.index\n\nsubmission.to_csv(\"submission.csv\")\n\n# Display the submission DataFrame\nsubmission\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-10-19T18:08:49.460348Z","iopub.execute_input":"2024-10-19T18:08:49.460877Z","iopub.status.idle":"2024-10-19T18:08:49.535796Z","shell.execute_reply.started":"2024-10-19T18:08:49.460821Z","shell.execute_reply":"2024-10-19T18:08:49.534203Z"}},"outputs":[],"execution_count":null}]}