{"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":30761,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport itertools\nfrom astropy.stats import sigma_clip\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nfrom scipy.optimize import minimize\nfrom functools import partial\n\n\n# 1. Define a Class for Data Handling\n\nclass DataProcessor:\n    linear_corr_dict = {\n        \"AIRS-CH0\": (6, 32, 356),\n        \"FGS1\": (6, 32, 32)\n    }\n\n    def __init__(self, adc_info_path, axis_info_path):\n        self.adc_info = pd.read_csv(adc_info_path, index_col='planet_id')\n        self.axis_info = pd.read_parquet(axis_info_path)\n        \n    def load_data(self, dataset, planet_id, sensor):\n        base_path = f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_calibration/'\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(base_path + 'dark.parquet', engine='pyarrow').to_numpy()\n        dead_frame = pd.read_parquet(base_path + 'dead.parquet', engine='pyarrow').to_numpy()\n        flat_frame = pd.read_parquet(base_path + 'flat.parquet', engine='pyarrow').to_numpy()\n        linear_corr = pd.read_parquet(base_path + 'linear_corr.parquet').values.astype(np.float64)\n\n        # Print dimensions to debug\n        print(f\"Linear Corr shape after loading: {linear_corr.shape}\")\n\n        if sensor in self.linear_corr_dict:\n            expected_shape = self.linear_corr_dict[sensor]\n            if linear_corr.shape != expected_shape:\n                print(f\"Reshaping linear_corr from {linear_corr.shape} to {expected_shape}\")\n                linear_corr = linear_corr.reshape(expected_shape)\n\n        return signal, dark_frame, dead_frame, flat_frame, linear_corr\n\n    def preprocess_signal(self, signal, dark_frame, dead_frame, flat_frame, linear_corr, sensor, binning):\n        sensor_sizes = {\"AIRS-CH0\": [[11250, 32, 356], [1, 32, 282]], \"FGS1\": [[135000, 32, 32], [1, 32, 32]]}\n        binned_dict = {\"AIRS-CH0\": [11250 // binning // 2, 282], \"FGS1\": [135000 // binning // 2]}\n\n        signal = signal.reshape(sensor_sizes[sensor][0])\n        gain = self.adc_info[f'{sensor}_adc_gain'].values[0]\n        offset = self.adc_info[f'{sensor}_adc_offset'].values[0]\n        signal = signal / gain + offset\n        \n        if sensor != \"FGS1\":\n            cut_inf, cut_sup = 39, 321\n            signal = signal[:, :, cut_inf:cut_sup]\n            dt = np.ones(len(signal)) * 0.1\n            dt[1::2] += 4.5\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        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 = self.apply_linear_corr(linear_corr, signal)\n        signal = self.clean_dark(linear_corr_signal, dark_frame, dt)\n        \n        flat = flat_frame.reshape(sensor_sizes[sensor][1])\n        flat[dead_frame.reshape(sensor_sizes[sensor][1])] = np.nan\n        flat[sigma_clip(dark_frame, sigma=5, maxiters=5).mask.reshape(sensor_sizes[sensor][1])] = np.nan\n        signal = signal / flat\n        \n        if sensor == \"FGS1\":\n            signal = signal.reshape((sensor_sizes[sensor][0][0], sensor_sizes[sensor][0][1] * sensor_sizes[sensor][0][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((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        return binned\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 preprocess_dataset(self, dataset, sensor, binning=15):\n        planet_ids = self.adc_info.index\n        feats = []\n        for i, planet_id in tqdm(list(enumerate(planet_ids))):\n            signal, dark_frame, dead_frame, flat_frame, linear_corr = self.load_data(dataset, planet_id, sensor)\n            binned = self.preprocess_signal(signal, dark_frame, dead_frame, flat_frame, linear_corr, sensor, binning)\n            feats.append(binned)\n        return np.stack(feats)\n\n\n\n    \n# 2. Define a Class for Calibration\nclass SignalCalibrator:\n    @staticmethod\n    def phase_detector(signal):\n        phase1, phase2 = None, None\n        best_drop = 0\n        for i in range(50 // 2, 150 // 2):\n            t1 = signal[i:i + 20 // 2].max() - signal[i:i + 20 // 2].min()\n            if t1 > best_drop:\n                phase1 = i + (20 + 5) // 2\n                best_drop = t1\n\n        best_drop = 0\n        for i in range(200 // 2, 250 // 2):\n            t1 = signal[i:i + 20 // 2].max() - signal[i:i + 20 // 2].min()\n            if t1 > best_drop:\n                phase2 = i - 5 // 2\n                best_drop = t1\n\n        return phase1, phase2\n\n    @staticmethod\n    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        return q + 1e3 if s < 1e-4 else q\n\n    @staticmethod\n    def calibrate_signal(signal):\n        p1, p2 = SignalCalibrator.phase_detector(signal)\n        best_deg, best_score = 1, 1e12\n        for deg in range(1, 6):\n            f = partial(SignalCalibrator.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, x, y, p(x)\n\n    @staticmethod\n    def calibrate_train(signal):\n        p1, p2 = SignalCalibrator.phase_detector(signal)\n        best_deg, best_score = 1, 1e12\n        for deg in range(1, 6):\n            f = partial(SignalCalibrator.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\n\n# 3. Putting Everything Together\ndef main():\n    # Initialize DataProcessor and SignalCalibrator\n    data_processor = DataProcessor(\n        adc_info_path='/kaggle/input/ariel-data-challenge-2024/test_adc_info.csv',\n        axis_info_path='/kaggle/input/ariel-data-challenge-2024/axis_info.parquet'\n    )\n\n    # Preprocess data\n    pre_train = np.concatenate([\n        data_processor.preprocess_dataset('test', \"FGS1\", 30*12),\n        data_processor.preprocess_dataset('test', \"AIRS-CH0\", 30)\n    ], axis=2)\n\n    # Calibrate signals\n    all_s = []\n    for i in range(len(data_processor.adc_info)):\n        signal = pre_train[i, :, 1:].mean(axis=1)\n        s, p, p1, p2 = SignalCalibrator.calibrate_train(pre_train[i, :, 1:].mean(axis=1))\n        all_s.append(s)\n\n    train_s = np.repeat(np.array(all_s), 283).reshape((len(all_s), 283))\n    train_sigma = np.ones_like(train_s) * 0.00016\n\n    # Plot an example\n    n = 0\n    s, x, y, y_new = SignalCalibrator.calibrate_signal(pre_train[n, :, 1:].mean(axis=1))\n    plt.scatter(x, y)\n    plt.scatter(x, y_new)\n    plt.show()\n\n    # Prepare submission\n    ss = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/sample_submission.csv')\n    preds = train_s.clip(0)\n    submission = pd.DataFrame(np.concatenate([preds, train_sigma], axis=1), columns=ss.columns[1:])\n    submission.index = data_processor.adc_info.index\n    submission.to_csv('submission.csv')\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-09-03T07:52:32.361318Z","iopub.execute_input":"2024-09-03T07:52:32.361846Z","iopub.status.idle":"2024-09-03T07:52:43.986891Z","shell.execute_reply.started":"2024-09-03T07:52:32.361797Z","shell.execute_reply":"2024-09-03T07:52:43.985211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}