{"metadata":{"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"}],"dockerImageVersionId":30775,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3 (ipykernel)","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.12"},"papermill":{"default_parameters":{},"duration":1914.629019,"end_time":"2024-09-16T00:50:15.655732","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-09-16T00:18:21.026713","version":"2.5.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# NeurIPS - Ariel Data Challenge 2024\n\n## Kaggle competitions\n\nCompleted by Rostislav Epifanov for the Master's Research Seminar at Novosibirsk State University","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":0.007855,"end_time":"2024-09-16T00:18:24.782854","exception":false,"start_time":"2024-09-16T00:18:24.774999","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import pickle\nimport numpy as np\nimport scipy.stats\nimport pandas as pd\nimport polars as pl\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\n\nfrom sklearn.model_selection import KFold, cross_val_predict\nfrom sklearn.metrics import r2_score, mean_squared_error","metadata":{"_kg_hide-input":true,"papermill":{"duration":3.333514,"end_time":"2024-09-16T00:18:28.123927","exception":false,"start_time":"2024-09-16T00:18:24.790413","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data","metadata":{"papermill":{"duration":0.006563,"end_time":"2024-09-16T00:18:28.137578","exception":false,"start_time":"2024-09-16T00:18:28.131015","status":"completed"},"tags":[]}},{"cell_type":"code","source":"from pathlib import Path\n\ngpath = Path('/kaggle/input/ariel-data-challenge-2024/')\n\nwavelengths = pd.read_csv(gpath / 'wavelengths.csv')\naxis_info = pd.read_parquet(gpath / 'axis_info.parquet')\n\ntrain_labels = pd.read_csv(gpath / 'train_labels.csv', index_col='planet_id')\ntrain_adc_info = pd.read_csv(gpath / 'train_adc_info.csv', index_col='planet_id')","metadata":{"papermill":{"duration":0.332549,"end_time":"2024-09-16T00:18:28.476896","exception":false,"start_time":"2024-09-16T00:18:28.144347","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def f_read_and_preprocess(dataset, adc_info, planet_ids):\n    f_raw_train = np.full((len(planet_ids), 67500), np.nan, dtype=np.float32)\n\n    for idx, planet_idx in tqdm(enumerate(planet_ids), total=len(planet_ids)):\n        signal = pl.read_parquet(gpath / f'{dataset}/{planet_idx}/FGS1_signal.parquet')\n        signal = signal.to_numpy().reshape(135000, 32, 32)\n\n        signal = signal.astype(np.float32)\n        mean_signal = signal.mean((1, 2))\n\n        net_signal = mean_signal[1::2] - mean_signal[0::2]\n        f_raw_train[idx] = net_signal\n\n    return f_raw_train\n\nf_raw_train = f_read_and_preprocess('train', train_adc_info, train_labels.index)\n\nwith open('f_raw_train.pickle', 'wb') as f:\n    pickle.dump(f_raw_train, f)","metadata":{"papermill":{"duration":0.018905,"end_time":"2024-09-16T00:18:28.502823","exception":false,"start_time":"2024-09-16T00:18:28.483918","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def a_read_and_preprocess(dataset, adc_info, planet_ids):\n    a_raw_train = np.zeros((len(planet_ids), 5625), dtype=np.float32)\n\n    for idx, planet_idx in tqdm(enumerate(planet_ids), total=len(planet_ids)):\n        signal = pl.read_parquet(gpath / f'{dataset}/{planet_idx}/AIRS-CH0_signal.parquet')\n        signal = signal.to_numpy().reshape(11250, 32, 356)\n\n        signal = signal.astype(np.float32)[:, :, 39:321]\n        mean_signal = signal.mean((1, 2))\n        net_signal = mean_signal[1::2] - mean_signal[0::2]\n        a_raw_train[idx] = net_signal\n\n    return a_raw_train\n\na_raw_train = a_read_and_preprocess('train', train_adc_info, train_labels.index)\n\nwith open('a_raw_train.pickle', 'wb') as f:\n    pickle.dump(a_raw_train, f)","metadata":{"papermill":{"duration":1184.072902,"end_time":"2024-09-16T00:50:07.033765","exception":false,"start_time":"2024-09-16T00:30:22.960863","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Feature engineering","metadata":{"papermill":{"duration":0.131487,"end_time":"2024-09-16T00:50:07.298045","exception":false,"start_time":"2024-09-16T00:50:07.166558","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def feature_engineering(f_raw, a_raw):\n    obscured = f_raw[:, 23500:44000].mean(axis=1)\n    unobscured = (f_raw[:, :20500].mean(axis=1) + f_raw[:, 47000:].mean(axis=1)) / 2\n    f_relative_reduction = 1 - obscured / unobscured\n\n    obscured = a_raw[:, 1958:3666].mean(axis=1)\n    unobscured = (a_raw[:, :1708].mean(axis=1) + a_raw[:, 3916:].mean(axis=1)) / 2\n    a_relative_reduction = 1 - obscured / unobscured\n\n    return f_relative_reduction[:, None], a_relative_reduction[:, None]\n\ntrain_data_fgs1, train_data_airs_ch0 = feature_engineering(f_raw_train, a_raw_train)","metadata":{"papermill":{"duration":0.168278,"end_time":"2024-09-16T00:50:07.876531","exception":false,"start_time":"2024-09-16T00:50:07.708253","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Modelling","metadata":{"papermill":{"duration":0.133879,"end_time":"2024-09-16T00:50:08.1416","exception":false,"start_time":"2024-09-16T00:50:08.007721","status":"completed"},"tags":[]}},{"cell_type":"code","source":"class ParticipantVisibleError(Exception):\n    pass\n\ndef competition_score(\n        solution: pd.DataFrame,\n        submission: pd.DataFrame,\n        naive_mean: float,\n        naive_sigma: float,\n        sigma_true: float,\n        row_id_column_name='planet_id',\n    ) -> float:\n\n    del solution[row_id_column_name]\n    del submission[row_id_column_name]\n\n    if submission.min().min() < 0:\n        raise ParticipantVisibleError('Negative values in the submission')\n    for col in submission.columns:\n        if not pd.api.types.is_numeric_dtype(submission[col]):\n            raise ParticipantVisibleError(f'Submission column {col} must be a number')\n\n    n_wavelengths = len(solution.columns)\n    if len(submission.columns) != n_wavelengths*2:\n        raise ParticipantVisibleError('Wrong number of columns in the submission')\n\n    y_pred = submission.iloc[:, :n_wavelengths].values\n    # Set a non-zero minimum sigma pred to prevent division by zero errors.\n    sigma_pred = np.clip(submission.iloc[:, n_wavelengths:].values, a_min=10**-15, a_max=None)\n    y_true = solution.values\n\n    GLL_pred = np.sum(scipy.stats.norm.logpdf(y_true, loc=y_pred, scale=sigma_pred))\n    GLL_true = np.sum(scipy.stats.norm.logpdf(y_true, loc=y_true, scale=sigma_true * np.ones_like(y_true)))\n    GLL_mean = np.sum(scipy.stats.norm.logpdf(y_true, loc=naive_mean * np.ones_like(y_true), scale=naive_sigma * np.ones_like(y_true)))\n\n    submit_score = (GLL_pred - GLL_mean)/(GLL_true - GLL_mean)\n    return float(np.clip(submit_score, 0.0, 1.0))\n\ndef postprocessing(pred_array, index, sigma_pred):\n    return pd.concat([pd.DataFrame(pred_array.clip(0, None), index=index, columns=wavelengths.columns),\n                      pd.DataFrame(sigma_pred, index=index, columns=[f\"sigma_{i}\" for i in range(1, 284)])],\n                     axis=1)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch, torch.nn as nn\n\nclass EncoderFgs1(nn.Module):\n    def __init__(self, in_fgs1_channels=1, out_channels=283):\n        super().__init__()\n\n        self.in_fgs1_channels = in_fgs1_channels\n        self.out_channels = out_channels\n\n        self.encoder = nn.Identity()\n\n    def forward(self, x):\n        x = self.encoder(x)\n\n        return x\n\nclass EncoderAirsCh0(nn.Module):\n    def __init__(self, in_airs_ch0_channels=1, out_channels=283):\n        super().__init__()\n\n        self.in_airs_ch0_channels = in_airs_ch0_channels\n        self.out_channels = out_channels\n\n        self.encoder = nn.Identity()\n\n    def forward(self, x):\n        x = self.encoder(x)\n\n        return x\n\nclass Model(nn.Module):\n    def __init__(self, in_fgs1_channels=1, in_airs_ch0_channels=1, out_channels=283):\n        super().__init__()\n\n        self.in_fgs1_channels = in_fgs1_channels\n        self.in_airs_ch0_channels = in_airs_ch0_channels\n        self.out_channels = out_channels\n\n        self.encoder_fgs1 = EncoderFgs1(in_fgs1_channels)\n        self.encoder_airs_ch0 = EncoderAirsCh0(in_airs_ch0_channels)\n\n        self.bone = nn.Sequential(\n            nn.Linear(in_fgs1_channels+in_airs_ch0_channels, 512, bias=False),\n            nn.GELU(),\n            nn.Linear(512, 512, bias=False),\n            nn.GELU(),\n            nn.Linear(512, 512, bias=False),\n            nn.GELU(),\n            nn.Linear(512, out_channels, bias=False),\n        )\n\n    def forward(self, fx, ax):\n        fx = self.encoder_fgs1(fx)\n        ax = self.encoder_airs_ch0(ax)\n\n        x = torch.concat([fx, ax], dim=-1)\n        x = self.bone(x)\n\n        return x\n\nclass Estimator(object):\n    def __init__(self, in_fgs1_channels, in_airs_ch0_channels, out_channels, seed, device):\n        self.model = Model(\n            in_fgs1_channels=in_fgs1_channels,\n            in_airs_ch0_channels=in_airs_ch0_channels,\n            out_channels=out_channels\n        ).to(device)\n\n        self.seed = seed\n        self._init_seed()\n\n        self.device = device\n\n    def _init_seed(self):\n        torch.backends.cudnn.benchmark = False\n        torch.backends.cudnn.determenistic = True\n        torch.backends.cudnn.enabled = False\n    \n        torch.manual_seed(self.seed)\n        torch.cuda.manual_seed(self.seed)\n        torch.cuda.manual_seed_all(self.seed)\n    \n    def fit(self, X_fgs1, X_airs_ch0, y):\n        self.model.train()\n\n        X_fgs1 = torch.tensor(X_fgs1, dtype=torch.float32, device=self.device)\n        X_airs_ch0 = torch.tensor(X_airs_ch0, dtype=torch.float32, device=self.device)\n        y = torch.tensor(y.values, dtype=torch.float32, device=self.device)\n\n        opt = torch.optim.AdamW(self.model.parameters(), lr=1e-3)\n\n        for _ in range(1000):\n            p = self.model(X_fgs1, X_airs_ch0)\n    \n            loss = (y - p)**2\n            loss = loss.mean()\n            loss.backward()\n            \n            opt.step()\n            opt.zero_grad(set_to_none=True)\n\n    def predict(self, X_fgs1, X_airs_ch0):\n        self.model.eval()\n\n        X_fgs1 = torch.tensor(X_fgs1, dtype=torch.float32, device=self.device)\n        X_airs_ch0 = torch.tensor(X_airs_ch0, dtype=torch.float32, device=self.device)\n\n        with torch.no_grad():\n            return self.model(X_fgs1, X_airs_ch0).cpu()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SEED = 1996\n\nfolds = KFold(n_splits=2)\nfolds = folds.split(train_labels, train_labels)\n\noof_pred = np.zeros_like(train_labels)\n\nfor fold in folds:\n    train_ids, test_ids = fold\n\n    estimator = Estimator(\n        in_fgs1_channels=1,\n        in_airs_ch0_channels=1,\n        out_channels=283,\n        seed=SEED,\n        device=torch.device('cuda'),\n    )\n\n    X_train_fgs1_fold = train_data_fgs1[train_ids]\n    X_train_airs_ch0_fold = train_data_airs_ch0[train_ids]\n    y_train_fold = train_labels.iloc[train_ids]\n\n    X_test_fgs1_fold = train_data_fgs1[test_ids]\n    X_test_airs_ch0_fold = train_data_airs_ch0[test_ids]\n    y_test_fold = train_labels.iloc[test_ids]\n\n    estimator.fit(X_train_fgs1_fold, X_train_airs_ch0_fold, y_train_fold)\n    p_test_fold = estimator.predict(X_test_fgs1_fold, X_test_airs_ch0_fold)\n\n    oof_pred[test_ids] = p_test_fold\n\nprint(f\"R2 score: {r2_score(train_labels, oof_pred):.3f}\")\nsigma_pred = mean_squared_error(train_labels, oof_pred, squared=False)\nprint(f\"Root mean squared error: {sigma_pred:.6f}\")\n\noof_df = postprocessing(oof_pred, train_adc_info.index, sigma_pred)\n\ngll_score = competition_score(train_labels.copy().reset_index(),\n                              oof_df.copy().reset_index(),\n                              naive_mean=train_labels.values.mean(),\n                              naive_sigma=train_labels.values.std(),\n                              sigma_true=0.000003)\n\nprint(f\"Estimated competition score: {gll_score:.3f}\")","metadata":{"papermill":{"duration":1.0781,"end_time":"2024-09-16T00:50:09.350051","exception":false,"start_time":"2024-09-16T00:50:08.271951","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"estimator.fit(train_data_fgs1, train_data_airs_ch0, train_labels)\n\nwith open('model.pickle', 'wb') as f:\n    pickle.dump(estimator, f)\n\nwith open('sigma_pred.pickle', 'wb') as f:\n    pickle.dump(sigma_pred, f)","metadata":{"papermill":{"duration":0.394735,"end_time":"2024-09-16T00:50:10.785637","exception":false,"start_time":"2024-09-16T00:50:10.390902","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","metadata":{"papermill":{"duration":0.131579,"end_time":"2024-09-16T00:50:11.047888","exception":false,"start_time":"2024-09-16T00:50:10.916309","status":"completed"},"tags":[]}},{"cell_type":"code","source":"test_adc_info = pd.read_csv(gpath / 'test_adc_info.csv', index_col='planet_id')\nsample_submission = pd.read_csv(gpath / 'sample_submission.csv', index_col='planet_id')\n\nf_raw_test = f_read_and_preprocess('test', test_adc_info, sample_submission.index)\na_raw_test = a_read_and_preprocess('test', test_adc_info, sample_submission.index)\n\ntest_data_fgs1, test_data_airs_ch0 = feature_engineering(f_raw_test, a_raw_test)\n\nwith open('model.pickle', 'rb') as f:\n    model = pickle.load(f)\n\nwith open('sigma_pred.pickle', 'rb') as f:\n    sigma_pred = pickle.load(f)\n\ntest_pred = model.predict(test_data_fgs1, test_data_airs_ch0)\n\nsub_df = postprocessing(test_pred, sample_submission.index, sigma_pred)\ndisplay(sub_df)\n\nsub_df.to_csv('submission.csv')","metadata":{"papermill":{"duration":3.132362,"end_time":"2024-09-16T00:50:14.313882","exception":false,"start_time":"2024-09-16T00:50:11.18152","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}