{"metadata":{"kernelspec":{"display_name":"jsenv","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.16"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"},{"sourceId":203900450,"sourceType":"kernelVersion"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false},"papermill":{"default_parameters":{},"duration":7.594014,"end_time":"2024-10-10T11:58:36.355301","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-10-10T11:58:28.761287","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pickle\nimport polars as pl\nimport numpy as np\nimport pandas as pd\nimport joblib\nfrom tqdm import tqdm\n\n","metadata":{"execution":{"iopub.execute_input":"2024-11-05T04:58:17.953615Z","iopub.status.busy":"2024-11-05T04:58:17.952915Z","iopub.status.idle":"2024-11-05T04:58:17.962711Z","shell.execute_reply":"2024-11-05T04:58:17.961348Z","shell.execute_reply.started":"2024-11-05T04:58:17.953574Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load Data","metadata":{}},{"cell_type":"code","source":"category_mappings = {'feature_09': {2: 0, 4: 1, 9: 2, 11: 3, 12: 4, 14: 5, 15: 6, 25: 7, 26: 8, 30: 9, 34: 10, 42: 11, 44: 12, 46: 13, 49: 14, 50: 15, 57: 16, 64: 17, 68: 18, 70: 19, 81: 20, 82: 21},\n 'feature_10': {1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 10: 7, 12: 8},\n 'feature_11': {9: 0, 11: 1, 13: 2, 16: 3, 24: 4, 25: 5, 34: 6, 40: 7, 48: 8, 50: 9, 59: 10, 62: 11, 63: 12, 66: 13,\n  76: 14, 150: 15, 158: 16, 159: 17, 171: 18, 195: 19, 214: 20, 230: 21, 261: 22, 297: 23, 336: 24, 376: 25, 388: 26, 410: 27, 522: 28, 534: 29, 539: 30},\n 'symbol_id': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 11, 12: 12, 13: 13, 14: 14, 15: 15, 16: 16, 17: 17, 18: 18, 19: 19,\n  20: 20, 21: 21, 22: 22, 23: 23, 24: 24, 25: 25, 26: 26, 27: 27, 28: 28, 29: 29, 30: 30, 31: 31, 32: 32, 33: 33, 34: 34, 35: 35, 36: 36, 37: 37, 38: 38},\n  'time_id' : {i : i for i in range(968)}}\n\n\ndef encode_column(df, column, mapping):\n    max_value = max(mapping.values())  \n\n    def encode_category(category):\n        return mapping.get(category, max_value + 1)  \n    \n    return df.with_columns(\n        pl.col(column).map_elements(encode_category, return_dtype=pl.Int16).alias(column)\n    )\n\n\nfeature_names = ['time_id', 'symbol_id', 'responder_6', 'responder_3', 'weight'] + [f\"feature_{i:02d}\" for i in range(79) if i != 61] + [f\"responder_{idx}_lag_1\" for idx in range(9)]\\\n     + ['sin_time_id', 'cos_time_id', 'sin_time_id_halfday', 'cos_time_id_halfday', 'sin_feature_61', 'cos_feature_61']\n\n\nlabel_name = 'responder_6'\nlabel2_name = 'responder_3'\nweight_name = 'weight'\nfeature_cat = ['feature_09', 'feature_10', 'feature_11', 'symbol_id', 'time_id']\nfeature_cont = [col for col in feature_names if col not in  feature_cat + [label_name, label2_name, weight_name]]\n\n\nfeature_cont_idx =  [feature_names.index(col) for col in feature_cont]\nfeature_cat_idx = [feature_names.index(col) for col in feature_cat]\nlabel_idx = feature_names.index(label_name)\nlabel2_idx = feature_names.index(label2_name)\nweight_idx = feature_names.index(weight_name)\n\n\n\ninput_path = 'root/data_pre'\ntrain_original = pl.scan_parquet(f\"{input_path}/training.parquet\").sort(['date_id', 'time_id', 'symbol_id'])\nvalid_original = pl.scan_parquet(f\"{input_path}/validation.parquet\").sort(['date_id', 'time_id', 'symbol_id'])\n\n# I use the 1577th day to split training and validation, and data also include the feature : [f\"responder_{idx}_lag_1\" for idx in range(9)]\n# create lag feature just like : https://www.kaggle.com/code/motono0223/js24-preprocessing-create-lags\n# data has also been mean-std normalized\n\nfor col in feature_cat:\n    train_original = encode_column(train_original, col, category_mappings[col])\n    valid_original = encode_column(valid_original, col, category_mappings[col])\n\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"means = joblib.load('root/data_stats.pkl')['mean']\nstds = joblib.load('root/data_stats.pkl')['std']\n\n\ntrain_original = train_original.with_columns([\n    (2 * np.pi * pl.col('time_id') / 967).sin().alias('sin_time_id'),\n    (2 * np.pi * pl.col('time_id') / 967).cos().alias('cos_time_id'),\n    (2 * np.pi * pl.col('time_id') / 483).sin().alias('sin_time_id_halfday'),\n    (2 * np.pi * pl.col('time_id') / 483).cos().alias('cos_time_id_halfday')\n])\n\nvalid_original = valid_original.with_columns([\n    (2 * np.pi * pl.col('time_id') / 967).sin().alias('sin_time_id'),\n    (2 * np.pi * pl.col('time_id') / 967).cos().alias('cos_time_id'),\n    (2 * np.pi * pl.col('time_id') / 483).sin().alias('sin_time_id_halfday'),\n    (2 * np.pi * pl.col('time_id') / 483).cos().alias('cos_time_id_halfday')\n])\n\n# denormalize feature_61 first\ntrain_original = train_original.with_columns([\n    pl.col('feature_61') * stds['feature_61'] + means['feature_61']\n])\n\nvalid_original = valid_original.with_columns([\n    pl.col('feature_61') * stds['feature_61'] + means['feature_61']\n])\n\ntrain_original = train_original.with_columns([\n    (2 * np.pi * pl.col('feature_61') / 20).sin().alias('sin_feature_61'),\n    (2 * np.pi * pl.col('feature_61') / 20).cos().alias('cos_feature_61')\n])\n\nvalid_original = valid_original.with_columns([\n    (2 * np.pi * pl.col('feature_61') / 20).sin().alias('sin_feature_61'),\n    (2 * np.pi * pl.col('feature_61') / 20).cos().alias('cos_feature_61')\n])\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\n# I use the day after 756 to debug the model\n# but finally I will use all the data after 252 day to train the model again\nall_data = False\nif all_data:\n    train_original = pl.concat([train_original, valid_original])\n    df = train_original\\\n        .filter(pl.col('date_id') >= 252)\\\n        .select(feature_names)\\\n        .collect().to_numpy()\n    valid = None\nelse:\n    df = train_original\\\n        .filter(pl.col('date_id') >= 756)\\\n        .select(feature_names)\\\n        .collect().to_numpy()\n    \n    valid = valid_original\\\n        .select(feature_names)\\\n        .collect().to_numpy()\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Configurations","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append('root/')\n\nimport os\nimport warnings\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\nfrom pytorch_lightning import (LightningDataModule, LightningModule, Trainer)\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, Timer\n\nimport pandas as pd\nimport numpy as np\nfrom sklearn.metrics import r2_score\nfrom sklearn.model_selection import train_test_split\nfrom torch.utils.data import Dataset, DataLoader\nfrom tabm_reference import Model, make_parameter_groups\n\n\nclass custom_args():\n    def __init__(self):\n        self.usegpu = True\n        self.gpuid = 0\n        self.seed = 42\n        self.loader_workers = 12   \n        self.bs = 8192\n        self.lr = 1e-3\n        self.weight_decay = 8e-4\n        self.n_cont_features = len(feature_cont)\n        self.n_cat_features = 5\n        self.n_classes = 1 # only use responder_6 or use both responder_6 and responder_3\n        self.cat_cardinalities = None if feature_cat is None else  [23, 10, 32, 40, 969]\n        self.max_epochs = 7\n\n\nmy_args = custom_args()","metadata":{"execution":{"iopub.execute_input":"2024-11-05T04:58:17.965474Z","iopub.status.busy":"2024-11-05T04:58:17.964754Z","iopub.status.idle":"2024-11-05T04:58:17.993759Z","shell.execute_reply":"2024-11-05T04:58:17.992851Z","shell.execute_reply.started":"2024-11-05T04:58:17.965427Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# PyTorch CustomDataset Definition","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n    def __init__(self, array):\n        self.features_cont = torch.FloatTensor(array[:, feature_cont_idx])\n        self.features_cat = torch.LongTensor(array[:, feature_cat_idx])\n        self.labels = torch.FloatTensor(array[:, label_idx] )\n        self.labels2 = torch.FloatTensor(array[:, label2_idx])\n        self.weights = torch.FloatTensor(array[:, weight_idx])\n    \n    def __len__(self):\n        return len(self.labels)\n    \n    def __getitem__(self, idx):\n        x_cont = self.features_cont[idx]\n        x_cat = self.features_cat[idx]\n        y = self.labels[idx]\n        y2 = self.labels2[idx]\n        w = self.weights[idx]\n        return x_cont,x_cat,y, y2,  w","metadata":{"execution":{"iopub.execute_input":"2024-11-05T04:58:17.995267Z","iopub.status.busy":"2024-11-05T04:58:17.994912Z","iopub.status.idle":"2024-11-05T04:58:18.776454Z","shell.execute_reply":"2024-11-05T04:58:18.77564Z","shell.execute_reply.started":"2024-11-05T04:58:17.995223Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model Definition","metadata":{}},{"cell_type":"code","source":"# Custom R2 metric for validation\ndef r2_val(y_true, y_pred, sample_weight):\n    residuals = sample_weight * (y_true - y_pred) ** 2\n    weighted_residual_sum = np.sum(residuals)\n\n    # Calculate weighted sum of squared true values (denominator)\n    weighted_true_sum = np.sum(sample_weight * (y_true) ** 2)\n\n    # Calculate weighted R2\n    r2 = 1 - weighted_residual_sum / weighted_true_sum\n    return r2\n\n\nclass R2Loss(nn.Module):\n    def __init__(self):\n        super(R2Loss, self).__init__()\n\n    def forward(self, y_pred, y_true):\n        mse_loss = torch.sum((y_pred - y_true) ** 2)\n        var_y = torch.sum(y_true ** 2)\n        loss = mse_loss / (var_y + 1e-38)\n        return loss\n\n\nclass NN(LightningModule):\n    def __init__(self, n_cont_features, cat_cardinalities, n_classes, lr, weight_decay):\n        super().__init__()\n        self.save_hyperparameters()\n        self.k = 16\n        self.model = Model(\n                n_num_features=n_cont_features,\n                cat_cardinalities=cat_cardinalities,\n                n_classes=n_classes,\n                backbone={\n                    'type': 'MLP',\n                    'n_blocks': 3 ,\n                    'd_block': [512, 512, 512],\n                    'dropout': 0.25 ,\n                },\n                bins=None,\n                num_embeddings= None,\n                # cat_embeddings={\n                #     None if cat_cardinalities is None else\n                #         'type': 'TrainablePositionEncoding',\n                #         'd_embedding' : [32, 32, 32, 32, 64],\n                #         'cardinality' : cat_cardinalities,\n                # },\n                # cat_dmodel = [32, 32, 32, 32, 64],\n                arch_type='tabm',\n                k=self.k,\n            )\n\n        self.lr = lr\n        self.weight_decay = weight_decay\n        self.training_step_outputs = []\n        self.validation_step_outputs = []\n        self.n_classes = n_classes\n        self.loss_fn = R2Loss()\n        # self.loss_fn = nn.MSELoss()\n\n\n\n    def forward(self, x_cont, x_cat):\n        return self.model(x_cont, x_cat).squeeze(-1)\n\n    def training_step(self, batch):\n        x_cont,x_cat, y, y2, w= batch\n        x_cont = x_cont + torch.randn_like(x_cont) * 0.02\n        y_hat = self(x_cont, x_cat)\n\n\n        if self.n_classes == 1:\n            loss = self.loss_fn(y_hat.flatten(0, 1), y.repeat_interleave(self.k))\n            self.training_step_outputs.append((y_hat.mean(1), y, w))\n        else:\n            loss1 = self.loss_fn(y_hat[:, :, 0].flatten(0, 1), y.repeat_interleave(self.k))\n            loss2 = self.loss_fn(y_hat[:, :, 1].flatten(0, 1), y2.repeat_interleave(self.k))\n            loss = 0.85 * loss1 + 0.15 * loss2\n            self.training_step_outputs.append((y_hat[:, :, 0].mean(1), y, w))\n\n        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=x_cont.size(0))\n        return loss\n\n    def validation_step(self, batch):\n        x_cont,x_cat, y, y2, w = batch\n        if len(x_cat.size()) == 1:\n            x_cat = None\n        # x_cont = x_cont + torch.randn_like(x_cont) * 0.02\n        y_hat = self(x_cont, x_cat)\n\n        if self.n_classes == 1:\n            loss = self.loss_fn(y_hat.flatten(0, 1), y.repeat_interleave(self.k))\n            self.validation_step_outputs.append((y_hat.mean(1), y, w))\n        else:\n            loss1 = self.loss_fn(y_hat[:, :, 0].flatten(0, 1), y.repeat_interleave(self.k))\n            loss2 = self.loss_fn(y_hat[:, :, 1].flatten(0, 1), y2.repeat_interleave(self.k))\n            loss = 0.85 * loss1 + 0.15 * loss2\n            self.validation_step_outputs.append((y_hat[:, :, 0].mean(1), y, w))\n\n        self.log('val_loss', loss, on_step=False, on_epoch=True, prog_bar=False, logger=True, batch_size=x_cont.size(0))\n        return loss\n\n    def on_validation_epoch_end(self):\n        y = torch.cat([x[1] for x in self.validation_step_outputs]).cpu().numpy()\n        if self.trainer.sanity_checking:\n            prob = torch.cat([x[0] for x in self.validation_step_outputs]).cpu().numpy()\n        else:\n            prob = torch.cat([x[0] for x in self.validation_step_outputs]).cpu().numpy()\n            weights = torch.cat([x[2] for x in self.validation_step_outputs]).cpu().numpy()\n            # r2_val\n            val_r_square = r2_val(y, prob, weights)\n\n            val_r_square_adj = 1 - (1 - r2_score(y, prob)) * (len(y) - 1) / (len(y) - 1 - 16)\n            self.log(\"val_r_square\", val_r_square, prog_bar=True, on_step=False, on_epoch=True)\n            self.log(\"val_r_square_adj\", val_r_square_adj, prog_bar=True, on_step=False, on_epoch=True)\n        self.validation_step_outputs.clear()\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(make_parameter_groups(self.model), lr=self.lr, weight_decay=self.weight_decay)\n\n        return {\n            'optimizer': optimizer,\n        }\n\n    def on_train_epoch_end(self):\n        if self.trainer.sanity_checking:\n            return\n\n        y = torch.cat([x[1] for x in self.training_step_outputs]).cpu().numpy()\n        prob = torch.cat([x[0] for x in self.training_step_outputs]).detach().cpu().numpy()\n        weights = torch.cat([x[2] for x in self.training_step_outputs]).cpu().numpy()\n        # r2_training\n        train_r_square = r2_val(y, prob, weights)\n        train_r_square_adj = 1 - (1 - r2_score(y, prob)) * (len(y) - 1) / (len(y) - 1 - 16)\n        # self.log(\"train_r_square\", train_r_square, prog_bar=True, on_step=False, on_epoch=True)\n        self.log(\"train_r_square_adj\", train_r_square_adj, prog_bar=True, on_step=False, on_epoch=True)\n        self.training_step_outputs.clear()\n\n        epoch = self.trainer.current_epoch\n        metrics = {k: v.item() if isinstance(v, torch.Tensor) else v for k, v in self.trainer.logged_metrics.items()}\n        formatted_metrics = {k: f\"{v:.7f}\" for k, v in metrics.items()}\n\n        formatted_metrics.pop('train_loss_step', None)\n        print(f\"Epoch {epoch}: {formatted_metrics}\")","metadata":{"execution":{"iopub.execute_input":"2024-11-05T04:58:18.778953Z","iopub.status.busy":"2024-11-05T04:58:18.778669Z","iopub.status.idle":"2024-11-05T04:58:18.796672Z","shell.execute_reply":"2024-11-05T04:58:18.795818Z","shell.execute_reply.started":"2024-11-05T04:58:18.778923Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Create PyTorch Data Module","metadata":{}},{"cell_type":"code","source":"%%time\n\npl.seed_everything(my_args.seed)\n# checking device\ndevice = torch.device(f'cuda:{my_args.gpuid}' if torch.cuda.is_available() and my_args.usegpu else 'cpu')\naccelerator = 'gpu' if torch.cuda.is_available() and my_args.usegpu else 'cpu'\nloader_device = 'cpu'\n\ntrain_ds = CustomDataset(df)\ntrain_dl = DataLoader(train_ds, batch_size=my_args.bs, shuffle=True, num_workers=my_args.loader_workers)\n\nif valid is not None:\n    valid_ds = CustomDataset(valid)\n    valid_dl = DataLoader(valid_ds, batch_size=my_args.bs, shuffle=False, num_workers=my_args.loader_workers)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Create Model and Training","metadata":{}},{"cell_type":"code","source":"import gc\ndel df\ngc.collect()\n\nmodel = NN(\n    n_cont_features=my_args.n_cont_features,\n    cat_cardinalities=my_args.cat_cardinalities,\n    n_classes=my_args.n_classes,\n    lr=my_args.lr,\n    weight_decay=my_args.weight_decay\n)\n\n# Initialize Callbacks\n\ncheckpoint_callback = ModelCheckpoint(monitor='val_r_square', mode='max', save_top_k=1, verbose=False, filename=f\"./models/tabm.model\") \ncevery_heckpoint_callback = ModelCheckpoint(every_n_epochs=1, save_top_k=-1, verbose=False, filename=\"./models/tabm_{epoch:02d}\") \ntimer = Timer()\n\n\n# Initialize Trainer\n\nprint(\"Training Epoch is \", my_args.max_epochs)\ntrainer = Trainer(\n    default_root_dir = 'root/',\n    max_epochs=my_args.max_epochs,\n    accelerator=accelerator,\n    devices=[my_args.gpuid] if my_args.usegpu else None,\n    callbacks=[ checkpoint_callback, every_heckpoint_callback, timer],\n    enable_progress_bar=True,\n    val_check_interval=0.5,\n)\n# Start Training\ntrainer.fit(model, train_dl, valid_dl)\n\nprint(f'\\Training completed in {timer.time_elapsed(\"train\"):.2f}s')\n\n","metadata":{"execution":{"iopub.execute_input":"2024-11-05T04:58:18.798262Z","iopub.status.busy":"2024-11-05T04:58:18.797992Z","iopub.status.idle":"2024-11-05T04:58:18.818565Z","shell.execute_reply":"2024-11-05T04:58:18.817654Z","shell.execute_reply.started":"2024-11-05T04:58:18.798232Z"},"trusted":true},"outputs":[],"execution_count":null}]}