{"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"},{"sourceId":8963053,"sourceType":"datasetVersion","datasetId":5394843}],"dockerImageVersionId":30747,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pwd","metadata":{"execution":{"iopub.status.busy":"2024-07-16T04:22:43.432203Z","iopub.execute_input":"2024-07-16T04:22:43.432514Z","iopub.status.idle":"2024-07-16T04:22:44.416412Z","shell.execute_reply.started":"2024-07-16T04:22:43.432487Z","shell.execute_reply":"2024-07-16T04:22:44.415402Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pandas as pd\nimport os\nimport sys\n\nx_cols = [f'state_t_{i}' for i in range(0, 60)] + \\\n         [f'state_q0001_{i}' for i in range(0, 60)] + \\\n         [f'state_q0002_{i}' for i in range(0, 60)] + \\\n         [f'state_q0003_{i}' for i in range(0, 60)] + \\\n         [f'state_u_{i}' for i in range(0, 60)] + \\\n         [f'state_v_{i}' for i in range(0, 60)] + \\\n         [f'pbuf_ozone_{i}' for i in range(0, 60)] + \\\n         [f'pbuf_CH4_{i}' for i in range(0, 27)] + \\\n         [f'pbuf_N2O_{i}' for i in range(0, 27)] + \\\n         ['state_ps', 'pbuf_SOLIN', 'pbuf_LHFLX', 'pbuf_SHFLX', 'pbuf_TAUX', 'pbuf_TAUY', 'pbuf_COSZRS', 'cam_in_ALDIF',\n          'cam_in_ALDIR', 'cam_in_ASDIF', 'cam_in_ASDIR', 'cam_in_LWUP', 'cam_in_ICEFRAC', 'cam_in_LANDFRAC',\n          'cam_in_OCNFRAC', 'cam_in_SNOWHLAND']\n\ny_cols = [f'ptend_t_{i}' for i in range(0, 60)] + \\\n         [f'ptend_q0001_{i}' for i in range(12, 60)] + \\\n         [f'ptend_q0002_{i}' for i in range(27, 60)] + \\\n         [f'ptend_q0003_{i}' for i in range(12, 60)] + \\\n         [f'ptend_u_{i}' for i in range(12, 60)] + \\\n         [f'ptend_v_{i}' for i in range(12, 60)] + \\\n         ['cam_out_NETSW', 'cam_out_FLWDS', 'cam_out_PRECSC', 'cam_out_PRECC', 'cam_out_SOLS', 'cam_out_SOLL',\n          'cam_out_SOLSD', 'cam_out_SOLLD']\n\nvalid_y_cols = [f'ptend_t_{i}' for i in range(0, 60)] + \\\n               [f'ptend_q0001_{i}' for i in range(0, 4)] + \\\n               [f'ptend_q0001_{i}' for i in range(12, 60)] + \\\n               [f'ptend_q0002_{i}' for i in range(12, 60)] + \\\n               [f'ptend_q0003_{i}' for i in range(12, 60)] + \\\n               [f'ptend_u_{i}' for i in range(12, 60)] + \\\n               [f'ptend_v_{i}' for i in range(12, 60)] + \\\n               ['cam_out_NETSW', 'cam_out_FLWDS', 'cam_out_PRECSC', 'cam_out_PRECC', 'cam_out_SOLS', 'cam_out_SOLL',\n                'cam_out_SOLSD', 'cam_out_SOLLD']\n\ny_cols_index = [valid_y_cols.index(v) for v in y_cols]\n\n\nmean = pd.read_parquet('/kaggle/input/leap-final/mean.parquet').iloc[0]\nstd = pd.read_parquet('/kaggle/input/leap-final/std_v2.parquet').iloc[0]\nlog_mean = pd.read_parquet('/kaggle/input/leap-final/log_mean.parquet').iloc[0]\nlog_std = pd.read_parquet('/kaggle/input/leap-final/log_std.parquet').iloc[0]","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-07-16T01:57:25.494668Z","iopub.execute_input":"2024-07-16T01:57:25.494931Z","iopub.status.idle":"2024-07-16T01:57:26.962901Z","shell.execute_reply.started":"2024-07-16T01:57:25.494906Z","shell.execute_reply":"2024-07-16T01:57:26.961951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import re\nimport os\nimport sys\n\n\nDEBUG = True if sys.gettrace() else False\nKAGGLE = bool([k for k in os.environ.keys() if 'KAGGLE' in k])\nTRAINING = False\nSERVER = False\n\n\nclass CFG_base:\n    fc_dropout = 0.2\n    batch_size = 1024\n    valid_batch_size = batch_size\n    num_workers = 2\n\n    freeze_n = 0\n    gradient_accumulation_steps = 1\n    gradient_checkpointing = False\n    mixed_precision = 'no'  # 控制训练时的精度设置, 推理时不要设置\n\n    custom_tokens = [\"\\n\", \" \" * 2]\n    fp16_infer = False\n\nclass CFG4_1_n17_d16_clip_v2(CFG_base):\n    folds = [0, ]\n\n    base_dir = '/kaggle/input/leap-final/server/v4_1_n17_d16_clip_v2'\n    clip_xy = False\n    model_name = 'v4_1'\n    layer_n = 17\n    emb_dim = 16\n    use_lognorm = True\n\n\nclass CFG8_1_n17_d8_v3(CFG_base):\n    folds = [0, ]\n\n    base_dir = '/kaggle/input/leap-final/server/v8_1_n17_d8_v3'\n    clip_xy = False\n    model_name = 'v8_1'\n    layer_n = 17\n    emb_dim = 8\n    use_lognorm = True\n\n\nclass CFG5_1_n9_d8_clip_v3(CFG_base):\n    folds = [0, ]\n\n    base_dir = '/kaggle/input/leap-final/server/v5_1_n9_d8_clip_v3'\n    clip_xy = False\n    model_name = 'v5_1'\n    layer_n = 9\n    emb_dim = 8\n    use_lognorm = True\n\n\nclass CFG4_8_n7_d8_clip_v3(CFG_base):\n    folds = [0, ]\n\n    base_dir = '/kaggle/input/leap-final/server/v4.8_n7_d8_clip_v3'\n    clip_xy = False\n    model_name = 'v4_8'\n    layer_n = 7\n    emb_dim = 8\n    use_lognorm = True\n\n\nclass CFG4_1_n7_d8_10M_100M(CFG_base):\n    folds = [0, ]\n\n    base_dir = '/kaggle/input/leap-final/server2/v4.1_n7_d8_10M_100M'\n    clip_xy = False\n    model_name = 'v4_1'\n    layer_n = 7\n    emb_dim = 8\n    use_lognorm = True\n\n\nclass CFG6_n11_d8_100M(CFG_base):\n    folds = [0, ]\n\n    base_dir = '/kaggle/input/leap-final/server2/v6_n11_d8_100M'\n    clip_xy = False\n    model_name = 'v6'\n    layer_n = 11\n    emb_dim = 8\n    use_lognorm = True\n\n\n# 后处理: 自动生成相关值\ndef post_process_CFGs(CFGs):\n    for CFG in CFGs:\n        CFG.DEBUG = DEBUG\n        CFG.KAGGLE = KAGGLE\n        CFG.TRAINING = TRAINING\n        CFG.SERVER = SERVER\n\n        # save as main.py\n        CFG.input_size = len(x_cols)\n        CFG.target_size = len(y_cols)\n\n        CFG.x_cols = x_cols\n        CFG.y_cols = y_cols\n        CFG.x_mean = mean[x_cols].values\n        CFG.x_log_mean = log_mean[x_cols].values\n        CFG.y_mean = mean[y_cols].values\n        CFG.x_std = std[x_cols].values\n        CFG.x_log_std = log_std[x_cols].values\n        CFG.y_std = std[y_cols].values\n        CFG.y_cols_index = y_cols_index\n\n        CFG.model_path = os.path.join(CFG.base_dir, re.sub(\n            r'fold\\d+', 'fold{fold}',\n            [v for v in os.listdir(CFG.base_dir) if v.startswith('model') and v.endswith('pth')][0],\n        ))","metadata":{"execution":{"iopub.status.busy":"2024-07-16T01:57:26.965154Z","iopub.execute_input":"2024-07-16T01:57:26.965546Z","iopub.status.idle":"2024-07-16T01:57:26.984804Z","shell.execute_reply.started":"2024-07-16T01:57:26.96552Z","shell.execute_reply":"2024-07-16T01:57:26.983629Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nimport numpy as np\nfrom sklearn.model_selection import KFold, StratifiedKFold\nimport torch\nimport os\nfrom torch.utils.data.distributed import DistributedSampler\n\n\nclass Collate:\n    def __init__(self, CFG):\n        self.cfg = CFG\n\n    def __call__(self, batch):\n        # UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor.\n        X = torch.tensor(np.array([sample[\"X\"] for sample in batch]), dtype=torch.float32)\n        if self.cfg.TRAINING:\n            y = torch.tensor(np.array([sample[\"y\"] for sample in batch]), dtype=torch.float32)\n            return {'X': X}, y\n        else:\n            return {'X': X}\n\n\nclass CustomDataset(Dataset):\n    def __init__(self, CFG, df, training):\n        self.cfg = CFG\n        self.df = df\n        self.training = training\n        self.index_end = []\n        for i, arr in enumerate(self.df):\n            if i == 0:\n                self.index_end.append(arr.shape[0])\n            else:\n                self.index_end.append(arr.shape[0] + self.index_end[-1])\n        # [100, 200, 300...]\n        self.index_end = tuple(self.index_end)\n\n    def __len__(self):\n        return self.index_end[-1]\n\n    def __getitem__(self, idx):\n        df_i = 0\n        while idx >= self.index_end[df_i]:\n            df_i += 1\n        if df_i != 0:\n            idx -= self.index_end[df_i - 1]\n\n        row = self.df[df_i][idx]\n        X = row[:490]\n\n        log_X = np.log10(np.abs(X) + np.finfo(np.float64).tiny)\n        log_X = (log_X - self.cfg.x_log_mean) / self.cfg.x_log_std\n        X = (X - self.cfg.x_mean) / self.cfg.x_std\n\n        X = X.astype(np.float32)\n        log_X = log_X.astype(np.float32)\n\n        if self.cfg.use_lognorm:\n            X = np.concatenate([X, log_X], axis=0)\n\n        if self.cfg.clip_xy:\n            X = np.clip(X, -10, 10)\n\n        if self.training:\n            valid_y = row[490:]\n            y = valid_y[self.cfg.y_cols_index]\n            y = (y - self.cfg.y_mean) / self.cfg.y_std\n            y = y.astype(np.float32)\n            if self.cfg.clip_xy:\n                y = np.clip(y, -10, 10)\n            return {'X': X, 'y': y}\n        else:\n            return {'X': X}\n\n\ndef get_test_dataloaders(CFG):\n    test_df = [pd.read_csv(r\"/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv\")[x_cols].values, ]\n\n    test_dataset = CustomDataset(CFG, test_df, training=False)\n    collate_fn = Collate(CFG)\n\n    test_loader = DataLoader(\n        test_dataset,\n        batch_size=CFG.batch_size,\n        shuffle=False,\n        collate_fn=collate_fn,\n        num_workers=CFG.num_workers, pin_memory=True, drop_last=False,\n    )\n    return test_loader\n","metadata":{"execution":{"iopub.status.busy":"2024-07-16T01:57:26.986246Z","iopub.execute_input":"2024-07-16T01:57:26.986635Z","iopub.status.idle":"2024-07-16T01:57:31.762196Z","shell.execute_reply.started":"2024-07-16T01:57:26.986604Z","shell.execute_reply":"2024-07-16T01:57:31.761314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\nos.chdir('..')\nKAGGLE = bool([k for k in os.environ.keys() if 'KAGGLE' in k])\n\nif not KAGGLE:\n    os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0'\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport torch\nfrom torch import nn\nimport importlib\nimport polars as pl\n\n\ndef load_custom_model(model_name):\n    try:\n        module = importlib.import_module(f'input.leap-final.model.{model_name}')\n        CustomModel = getattr(module, 'CustomModel')\n        return CustomModel\n    except ModuleNotFoundError:\n        print(f\"Module model.{model_name} not found\")\n    except AttributeError:\n        print(f\"'CustomModel' not found in model.{model_name}\")\n\n\ndef test(model, test_loader, device, CFG):\n    all_y_pred = []\n\n    model.eval()\n    for step, X in enumerate(tqdm(test_loader, ncols=70)):\n        for k, v in X.items():\n            X[k] = v.to(device)\n\n        with torch.no_grad():\n            y_pred = model(X).detach()\n            all_y_pred.append(y_pred.cpu())\n\n    y_pred = torch.cat(all_y_pred, dim=0)\n\n    return y_pred\n\n\ndef CFG2pred(CFG, CustomModel):\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n    y_preds = []\n    test_loader = get_test_dataloaders(CFG)\n    for fold in CFG.folds:\n        model = CustomModel(CFG)\n        state_dict = torch.load(CFG.model_path.format(fold=fold))\n        weights_dict = {}\n        for k, v in state_dict.items():\n            k = k.replace('module.', '') if 'module.' in k else k\n            k = k.replace('_orig_mod.', '') if '_orig_mod.' in k else k\n            weights_dict[k] = v\n        state_dict = weights_dict\n        model.load_state_dict(state_dict)\n        if CFG.fp16_infer:\n            model.half()\n        model.to(device)\n        # model = torch.jit.script(model)\n        if torch.cuda.device_count() > 1:\n            model = nn.DataParallel(model, device_ids=[0, 1])\n\n        y_pred = test(model, test_loader, device, CFG)\n        y_preds.append(y_pred)\n\n    return torch.mean(torch.stack(y_preds, dim=0), dim=0).numpy()\n\n\ndef test_main():\n    CFGs = [\n        CFG8_1_n17_d8_v3,\n        CFG4_1_n7_d8_10M_100M,\n        CFG4_1_n17_d16_clip_v2,\n        CFG4_8_n7_d8_clip_v3,\n        CFG5_1_n9_d8_clip_v3,\n        CFG6_n11_d8_100M,\n    ]\n    for CFG in CFGs:\n        CFG.fp16_infer = False\n        CFG.batch_size = 2048\n    post_process_CFGs(CFGs)\n\n    weights = np.load('/kaggle/input/leap-final/hill_climbing_by_col_clip_10091520.npy')\n\n    col_min_values = np.load(r\"/kaggle/input/leap-final/col_min_values.npy\")\n    col_max_values = np.load(r\"/kaggle/input/leap-final/col_max_values.npy\")\n\n    y_preds = []\n    for i, CFG in enumerate(CFGs):\n        y_preds.append(\n            np.clip(\n                (CFG2pred(CFG, load_custom_model(CFG.model_name)) * CFGs[0].y_std) + CFGs[0].y_mean,\n                a_min=col_min_values, a_max=col_max_values,\n            )\n        )\n    final_y_pred = []\n    for j in range(293):\n        temp = []\n        for i, y_pred in enumerate(y_preds):\n            temp.append(y_pred[:, j] * weights[j, i])\n        final_y_pred.append(np.sum(temp, axis=0))\n    final_y_pred = np.stack(final_y_pred, axis=1)\n    print(final_y_pred.shape)\n    y_pred = final_y_pred\n\n    del y_preds\n\n    y_pred = pd.DataFrame(y_pred, columns=y_cols)\n\n    # convert to final submission\n    train1 = pd.read_csv('/kaggle/input/leap-atmospheric-physics-ai-climsim/train.csv', nrows=1)\n    df_subm = pd.read_csv(r\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\", nrows=1)\n\n    sample_id = pd.read_csv(\n        r\"/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv\",\n        usecols=['sample_id', ],\n    )\n    y_pred['sample_id'] = sample_id.iloc[:, 0]\n\n    for col in df_subm.columns:\n        if col not in y_pred.columns:\n            y_pred[col] = train1[col].iloc[0] * df_subm[col].iloc[0]\n\n    y_pred = y_pred[df_subm.columns]  # 交换列顺序\n\n    # weight is 0\n    for col in [f'ptend_q0001_{i}' for i in range(0, 12)] + \\\n               [f'ptend_q0002_{i}' for i in range(0, 15)] + \\\n               [f'ptend_q0003_{i}' for i in range(0, 12)] + \\\n               [f'ptend_u_{i}' for i in range(0, 12)] + \\\n               [f'ptend_v_{i}' for i in range(0, 12)]:\n        y_pred.loc[:, col] = 0\n\n    # target that can infer from input\n    REPLACE_TO = [f'state_q0002_{i}' for i in range(27)]  # x\n    REPLACE_FROM = [f'ptend_q0002_{i}' for i in range(27)]  # y\n    df_test = pl.read_csv('/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv')\n    df_test = df_test.to_pandas()\n    df_test = df_test.set_index(\"sample_id\")\n    y_pred[REPLACE_FROM] = -df_test[REPLACE_TO].to_numpy() * df_subm[REPLACE_FROM].to_numpy() / 1200\n\n    y_pred.to_csv(r'/kaggle/working/submission.csv', index=False)\n\n\nif __name__ == '__main__':\n    test_main()","metadata":{"execution":{"iopub.status.busy":"2024-07-16T01:57:31.764405Z","iopub.execute_input":"2024-07-16T01:57:31.764918Z"},"trusted":true},"execution_count":null,"outputs":[]}]}