{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":56537,"databundleVersionId":8877088,"sourceType":"competition"}],"dockerImageVersionId":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-06-24T02:40:50.311858Z","iopub.execute_input":"2024-06-24T02:40:50.312674Z","iopub.status.idle":"2024-06-24T02:40:50.69043Z","shell.execute_reply.started":"2024-06-24T02:40:50.312629Z","shell.execute_reply":"2024-06-24T02:40:50.689472Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\nimport os\nimport random\nimport time\nimport torch\nimport datetime\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset , DataLoader\nfrom torchmetrics.regression import R2Score","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:40:50.692384Z","iopub.execute_input":"2024-06-24T02:40:50.693164Z","iopub.status.idle":"2024-06-24T02:41:00.987057Z","shell.execute_reply.started":"2024-06-24T02:40:50.693129Z","shell.execute_reply":"2024-06-24T02:41:00.986256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_PATH = \"/kaggle/input/\"\nBATCH_SIZE = 12288\nMIN_STD = 1e-6\nSCHEDULER_PATIENCE = 3\nSCHEDULER_FACTOR = 10**(-0.5)\nEPOCHS = 50\nPATIENCE = 6\nPRINT_FREQ = 50","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:41:00.988282Z","iopub.execute_input":"2024-06-24T02:41:00.988797Z","iopub.status.idle":"2024-06-24T02:41:00.994154Z","shell.execute_reply.started":"2024-06-24T02:41:00.988765Z","shell.execute_reply":"2024-06-24T02:41:00.993096Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def format_time(elapsed):\n    elapsed_rounded = int(round((elapsed)))\n    return str(datetime.timedelta(seconds=elapsed_rounded))","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:41:00.996748Z","iopub.execute_input":"2024-06-24T02:41:00.997107Z","iopub.status.idle":"2024-06-24T02:41:01.008826Z","shell.execute_reply.started":"2024-06-24T02:41:00.997077Z","shell.execute_reply":"2024-06-24T02:41:01.008125Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def seed_everything(seed_val = 1331):\n    random.seed(seed_val)\n    np.random.seed(seed_val)\n    torch.manual_seed(seed_val)\n    torch.cuda.manual_seed_all(seed_val)","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:41:01.009831Z","iopub.execute_input":"2024-06-24T02:41:01.010167Z","iopub.status.idle":"2024-06-24T02:41:01.021359Z","shell.execute_reply.started":"2024-06-24T02:41:01.010138Z","shell.execute_reply":"2024-06-24T02:41:01.020512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ts = time.time()\n\nweights = pd.read_csv(DATA_PATH + \"leap-atmospheric-physics-ai-climsim/sample_submission.csv\", nrows=1)\ndel weights['sample_id']\nweights = weights.T\nweights = weights.to_dict()[0]\ndf_train = pl.read_csv(DATA_PATH + \"leap-atmospheric-physics-ai-climsim/train.csv\", n_rows = 2_500_000)\n\n\nfor target in weights:\n    df_train = df_train.with_columns(pl.col(target).mul(weights[target]))\n\nprint(\"Time to read dataset:\",format_time(time.time()-ts),flush=True)\n\nFEAT_COLS = df_train.columns[1:557]\nTARGET_COLS = df_train.columns[557:]\n\nfor col in FEAT_COLS:\n    df_train = df_train.with_columns(pl.col(col).cast(pl.Float32))\nfor col in TARGET_COLS:\n    df_train = df_train.with_columns(pl.col(col).cast(pl.Float32))\n\n    \nx_train = df_train.select(FEAT_COLS).to_numpy()\ny_train = df_train.select(TARGET_COLS).to_numpy()\n\ndel df_train\ngc.collect()\n\nmeanx = x_train.mean(axis=0)\nstdx = np.maximum(x_train.std(axis=0),MIN_STD)\nx_train = (x_train - meanx.reshape(1,-1)) / stdx.reshape(1,-1)\n\nmeany = y_train.mean(axis=0)\nstdy = np.maximum(np.sqrt((y_train*y_train).mean(axis=0)),MIN_STD)\ny_train = (y_train - meany.reshape(1,-1)) / stdy.reshape(1,-1)\n\nprint(\"Time after processing data:\", format_time(time.time()-ts),flush = True)","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:41:01.022364Z","iopub.execute_input":"2024-06-24T02:41:01.023461Z","iopub.status.idle":"2024-06-24T02:44:56.190916Z","shell.execute_reply.started":"2024-06-24T02:41:01.023438Z","shell.execute_reply":"2024-06-24T02:44:56.189872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seed_everything()\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:44:56.192506Z","iopub.execute_input":"2024-06-24T02:44:56.193232Z","iopub.status.idle":"2024-06-24T02:44:56.410686Z","shell.execute_reply.started":"2024-06-24T02:44:56.193198Z","shell.execute_reply":"2024-06-24T02:44:56.409655Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class NumpyDataset(Dataset):\n    def __init__(self, x, y):\n        assert x.shape[0] == y.shape[0], \"Features and labels must have the same number of samples\"\n        self.x = x\n        self.y = y\n\n    def __len__(self):\n        return self.x.shape[0]\n\n    def __getitem__(self, index):\n        return torch.from_numpy(self.x[index]).float().to(device), torch.from_numpy(self.y[index]).float().to(device)","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:44:56.412046Z","iopub.execute_input":"2024-06-24T02:44:56.412394Z","iopub.status.idle":"2024-06-24T02:44:56.419183Z","shell.execute_reply.started":"2024-06-24T02:44:56.412364Z","shell.execute_reply":"2024-06-24T02:44:56.418232Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset = NumpyDataset(x_train, y_train)\n\ntrain_size = int(0.9 * len(dataset))\nval_size = len(dataset) - train_size\ntrain_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:44:56.420355Z","iopub.execute_input":"2024-06-24T02:44:56.420627Z","iopub.status.idle":"2024-06-24T02:44:56.732558Z","shell.execute_reply.started":"2024-06-24T02:44:56.420605Z","shell.execute_reply":"2024-06-24T02:44:56.731505Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class FFNN(nn.Module):\n    def __init__(self, input_size, hidden_sizes, output_size):\n        super(FFNN, self).__init__()\n        \n        layers = []\n        previous_size = input_size\n        for hidden_size in hidden_sizes:\n            layers.append(nn.Linear(previous_size, hidden_size))\n            layers.append(nn.LayerNorm(hidden_size))\n            layers.append(nn.PReLU())\n            layers.append(nn.Dropout(p=0.1))\n            previous_size = hidden_size\n        \n        layers.append(nn.Linear(previous_size, output_size))\n        \n        self.layers = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.layers(x)","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:44:56.735161Z","iopub.execute_input":"2024-06-24T02:44:56.735453Z","iopub.status.idle":"2024-06-24T02:44:56.742565Z","shell.execute_reply.started":"2024-06-24T02:44:56.735429Z","shell.execute_reply":"2024-06-24T02:44:56.741605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"input_size = x_train.shape[1]\noutput_size = y_train.shape[1]\nhidden_size = input_size + output_size\nmodel = FFNN(input_size, [3*hidden_size, 2*hidden_size, 2*hidden_size, 2*hidden_size, 3*hidden_size], output_size).to(device)\ncriterion = nn.MSELoss()\noptimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=SCHEDULER_FACTOR, patience=SCHEDULER_PATIENCE)\n\nprint(\"Time after all preparations:\", format_time(time.time()-ts), flush=True)","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:44:56.743661Z","iopub.execute_input":"2024-06-24T02:44:56.743908Z","iopub.status.idle":"2024-06-24T02:44:58.361446Z","shell.execute_reply.started":"2024-06-24T02:44:56.743887Z","shell.execute_reply":"2024-06-24T02:44:58.360467Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"best_val_loss = float('inf')\nbest_model_state = None\npatience_count = 0\nr2score = R2Score(num_outputs=len(TARGET_COLS)).to(device)\nfor epoch in range(EPOCHS):\n    print(\"\")\n    model.train()\n    total_loss = 0\n    steps = 0\n    for batch_idx, (inputs, labels) in enumerate(train_loader):\n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n        steps += 1\n\n        if (batch_idx + 1) % PRINT_FREQ == 0:\n            current_lr = optimizer.param_groups[0][\"lr\"]\n            elapsed_time = format_time(time.time() - ts)\n            print(f'  Epoch: {epoch+1}',\\\n                  f'  Batch: {batch_idx + 1}/{len(train_loader)}',\\\n                  f'  Train Loss: {total_loss / steps:.4f}',\\\n                  f'  LR: {current_lr:.1e}',\\\n                  f'  Time: {elapsed_time}', flush=True)\n            total_loss = 0\n            steps = 0\n    \n\n    model.eval()\n    val_loss = 0\n    y_true = torch.tensor([], device=device)\n    all_outputs = torch.tensor([], device=device)\n    with torch.no_grad():\n        for inputs, labels in val_loader:\n            outputs = model(inputs)\n            val_loss += criterion(outputs, labels).item()\n            y_true = torch.cat((y_true, labels), 0)\n            all_outputs = torch.cat((all_outputs, outputs), 0)\n    r2=0\n    r2_broken = []\n    r2_broken_names = []\n    for i in range(368):\n        r2_i = r2score(all_outputs[:, i], y_true[:, i])\n        if r2_i > 1e-6:\n            r2 += r2_i\n        else:\n            r2_broken.append(i)\n            r2_broken_names.append(FEAT_COLS[i])\n    r2 /= 368\n\n    avg_val_loss = val_loss / len(val_loader)\n    print(f'\\nEpoch: {epoch+1}  Val Loss: {avg_val_loss:.4f}  R2 score: {r2:.4f}')\n    print(f'{len(r2_broken)} targets were excluded during evaluation of R2 score.')\n    # print(r2_broken)\n    # print(r2_broken_names, flush=True)\n   \n    scheduler.step(avg_val_loss)\n\n    if avg_val_loss < best_val_loss:\n        best_val_loss = avg_val_loss\n        best_model_state = model.state_dict()\n        patience_count = 0\n        print(\"Validation loss decreased, saving new best model and resetting patience counter.\")\n    else:\n        patience_count += 1\n        print(f\"No improvement in validation loss for {patience_count} epochs.\")\n        \n    if patience_count >= PATIENCE:\n        print(\"Stopping early due to no improvement in validation loss.\")\n        break\n\ndel x_train, y_train\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-06-24T02:44:58.362618Z","iopub.execute_input":"2024-06-24T02:44:58.363032Z","iopub.status.idle":"2024-06-24T06:58:49.950922Z","shell.execute_reply.started":"2024-06-24T02:44:58.363007Z","shell.execute_reply":"2024-06-24T06:58:49.949977Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.load_state_dict(best_model_state)\nmodel.eval()\n\ndf_test = pl.read_csv(DATA_PATH + \"leap-atmospheric-physics-ai-climsim/test.csv\")\n\nfor col in FEAT_COLS:\n    df_test = df_test.with_columns(pl.col(col).cast(pl.Float32))\n\nx_test = df_test.select(FEAT_COLS).to_numpy()\n\nx_test = (x_test - meanx.reshape(1,-1)) / stdx.reshape(1,-1)\n\npredt = np.zeros([x_test.shape[0], output_size], dtype=np.float32)\n\ni1 = 0\nfor i in range(10000):\n    i2 = np.minimum(i1 + BATCH_SIZE, x_test.shape[0])\n    if i1 == i2:  # Break the loop if range does not change\n        break\n\n    # Convert the current slice of xt to a PyTorch tensor\n    inputs = torch.from_numpy(x_test[i1:i2, :]).float().to(device)\n\n    # No need to track gradients for inference\n    with torch.no_grad():\n        outputs = model(inputs)  # Get model predictions\n        predt[i1:i2, :] = outputs.cpu().numpy()  # Store predictions in predt\n\n    i1 = i2  # Update i1 to the end of the current batch\n\n    if i2 >= x_test.shape[0]:\n        break\n\nfor i in range(stdy.shape[0]):\n    if stdy[i] < MIN_STD * 1.1:\n        predt[:,i] = 0\n\npredt = predt * stdy.reshape(1,-1) + meany.reshape(1,-1)\n\nss = pd.read_csv(DATA_PATH + \"leap-atmospheric-physics-ai-climsim/sample_submission.csv\")\nss.iloc[:,1:] = predt\n\ndel predt\ngc.collect()\n\nuse_cols = []\nfor i in range(27):\n    use_cols.append(f\"ptend_q0002_{i}\")\n\nss2 = pd.read_csv(DATA_PATH + \"leap-atmospheric-physics-ai-climsim/sample_submission.csv\")\ndf_test = df_test.to_pandas()\nfor col in use_cols:\n    ss[col] = -df_test[col.replace(\"ptend\", \"state\")]*ss2[col]/1200.\n\ntest_polars = pl.from_pandas(ss[[\"sample_id\"]+TARGET_COLS])\ntest_polars.write_csv(\"submission.csv\")\n\nprint(\"Total time:\", format_time(time.time()-ts))\n","metadata":{"execution":{"iopub.status.busy":"2024-06-24T06:58:49.952189Z","iopub.execute_input":"2024-06-24T06:58:49.952476Z","iopub.status.idle":"2024-06-24T07:00:39.855474Z","shell.execute_reply.started":"2024-06-24T06:58:49.952452Z","shell.execute_reply":"2024-06-24T07:00:39.854245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### ","metadata":{}}]}