{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":84493,"databundleVersionId":11037875,"sourceType":"competition"}],"dockerImageVersionId":30886,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\n\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\nimport keras\nimport numpy as np\nimport pandas as pd\nimport polars as pl\n\nfrom typing import Dict, List, Tuple\n\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\n\nfrom keras import layers\nfrom keras import Model\nfrom keras import ops\nfrom keras import regularizers\n\nfrom sklearn.metrics import r2_score\n\nfrom pathlib import Path\nimport gc\n\nDATA_DIR = Path(\"/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:05.122994Z","iopub.execute_input":"2025-02-23T21:46:05.123286Z","iopub.status.idle":"2025-02-23T21:46:08.438943Z","shell.execute_reply.started":"2025-02-23T21:46:05.123263Z","shell.execute_reply":"2025-02-23T21:46:08.438227Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DataProcessor:\n    ALL_FEATURES = [f\"feature_{i:02}\" for i in range(79)]\n    FEATURES = [x for x in ALL_FEATURES if x not in\n     [\"feature_09\", \"feature_10\", \"feature_11\"]]\n    # selected features from\n    # https://github.com/evgeniavolkova/kagglejanestreet/blob/master/janestreet/data_processor.py\n    SELECT_FEATURES = [\n        'feature_06',\n        'feature_04',\n        'feature_07',\n        'feature_36',\n        'feature_60',\n        'feature_45',\n        'feature_56',\n        'feature_05',\n        'feature_51',\n        'feature_19',\n        'feature_66',\n        'feature_59',\n        'feature_54',\n        'feature_70',\n        'feature_71',\n        'feature_72',\n    ]\n    CAT_FEATURES = [\"feature_09\", \"feature_10\", \"feature_11\"]\n    RESPONDERS = [f\"responder_{i}\" for i in range(9)]\n    FEATURE_09_DICT = {\n        2: 0,\n        4: 1,\n        9: 2,\n        11: 3,\n        12: 4,\n        14: 5,\n        15: 6,\n        25: 7,\n        26: 8,\n        30: 9,\n        34: 10,\n        42: 11,\n        44: 12,\n        46: 13,\n        49: 14,\n        50: 15,\n        57: 16,\n        64: 17,\n        68: 18,\n        70: 19,\n        81: 20,\n        82: 21\n        }\n    FEATURE_10_DICT = {1: 0, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 7: 6, 10: 7, 12: 8}\n    FEATURE_11_DICT = {\n        9: 0,\n        11: 1,\n        13: 2,\n        16: 3,\n        24: 4,\n        25: 5,\n        34: 6,\n        40: 7,\n        48: 8,\n        50: 9,\n        59: 10,\n        62: 11,\n        63: 12,\n        66: 13,\n        76: 14,\n        150: 15,\n        158: 16,\n        159: 17,\n        171: 18,\n        195: 19,\n        214: 20,\n        230: 21,\n        261: 22,\n        297: 23,\n        336: 24,\n        376: 25,\n        388: 26,\n        410: 27,\n        522: 28,\n        534: 29,\n        539: 30\n    }\n\n    def __init__(self, df: pl.LazyFrame):\n        self.df = df\n\n    def reduce_memory_usage(self) -> list[pl.Expr]:\n        expressions = [\n            pl.col(pl.Float64).cast(pl.Float32),\n            pl.col(\"date_id\").cast(pl.Int16),\n            pl.col(\"time_id\").cast(pl.Int16),\n            pl.col(\"symbol_id\").cast(pl.Int8),\n        ]\n        return expressions\n\n    def map_category(self) -> list[pl.Expr]:\n        expressions = [\n            pl.col('symbol_id').replace({i: i for i in range(39)}).fill_null(99),\n            pl.col('feature_09').replace(DataProcessor.FEATURE_09_DICT).fill_null(99),\n            pl.col('feature_10').replace(DataProcessor.FEATURE_10_DICT).fill_null(99),\n            pl.col('feature_11').replace(DataProcessor.FEATURE_11_DICT).fill_null(99),\n        ]\n        return expressions\n\n    def get_lag_responders(self) -> list[pl.Expr]:\n        cols = [f\"responder_{i}\" for i in range(9)]\n        expressions = [\n            pl.col(col)\n            .shift(i)\n            .fill_null(0)\n            .over('symbol_id', 'time_id')\n            .alias(f\"{col}_lag_{i}\")\n            for col in cols for i in [1]  # lags\n        ]\n        return expressions\n\n    def get_temporal_features(self) -> list[pl.Expr]:\n        expressions = [\n            (pl.col('date_id') % 170).alias('day'),\n            (pl.col('date_id') * 2 * np.pi / 170).sin()\\\n            .cast(pl.Float32).alias('date_sin'),\n            (pl.col('date_id') * 2 * np.pi / 170).cos()\\\n            .cast(pl.Float32).alias('date_cos'),\n        ]\n        return expressions\n\n    def get_tag_average(self) -> list[pl.Expr]:\n        expressions = [\n            pl.mean_horizontal(['feature_67', 'feature_68', 'feature_69'])\\\n            .alias('tag_11_mean'),\n        ]\n        return expressions\n\n    def get_select_lag_features(self) -> list[pl.Expr]:\n        group = ['symbol_id', 'date_id']\n        expressions = [\n            pl.col('feature_07').shift(2).fill_null(0)\\\n            .over(group).alias('feature_07_lag'),\n            pl.col('feature_06').shift(1).fill_null(0)\\\n            .over(group).alias('feature_06_lag'),\n            pl.col('feature_60').shift(2).fill_null(0)\\\n            .over(group).alias('feature_60_lag'),\n        ]\n        return expressions\n\n    def get_lag_responder_stats_per_day(self) -> list[pl.Expr]:\n        group = [\"date_id\", \"symbol_id\"]\n        cols = ['responder_6_lag_1']\n        expressions = []\n        for col in cols:\n            exprs = [\n                pl.col(col).max().over(group).alias(f\"{col}_max\"),\n                # pl.col(col).min().over(group).alias(f\"{col}_min\"),\n                # pl.col(col).mean().over(group).alias(f\"{col}_mean\"),\n                # pl.col(col).std().over(group).alias(f\"{col}_std\")\n            ]\n            expressions.extend(exprs)\n        return expressions\n\n    def get_stats_per_date_time(self) -> list[pl.Expr]:\n        group = [\"date_id\", \"time_id\"]\n        expressions = []\n        for col in DataProcessor.SELECT_FEATURES:\n            exprs =[\n                pl.col(col).mean().over(group).alias(f\"{col}_mean\"),\n                pl.col(col).std().over(group).alias(f\"{col}_std\"),\n                # pl.col(col).skew().over(group).alias(f\"{col}_skew\"),\n                # pl.col(col).kurtosis().over(group).alias(f\"{col}_kurtosis\")\n            ]\n            expressions.extend(exprs)\n        return expressions\n\n    def get_diff_means(self) -> list[pl.Expr]:\n        expressions = [\n            (pl.col(col) - pl.col(f\"{col}_mean\")).alias(f\"{col}_diff_mean\")\n            for col in DataProcessor.SELECT_FEATURES\n        ]\n        return expressions\n\n    def min_max_scaler(self, df: pl.DataFrame, columns: List) -> pl.DataFrame:\n        for col in columns:\n            col_min = df.select(col).min()\n            col_max = df.select(col).max()\n            df = df.with_columns(\n                ((pl.col(col) - col_min) / ((col_max - col_min)+1e-10))\n                .alias(col)\n                )\n        return df\n\n    def generate_features(self) -> pl.DataFrame:\n        exprs_1 = [\n            self.reduce_memory_usage(),\n            self.get_temporal_features(),\n            self.get_tag_average(),\n            self.get_select_lag_features(),\n            self.get_stats_per_date_time(),\n            # self.get_lag_responders(),\n        ]\n        expressions_1 = [e for sublist in exprs_1 for e in sublist]  # Flatten the list\n        df = self.df.with_columns(expressions_1)\n\n        exprs_2 = [\n            self.map_category(),\n            # self.get_lag_responder_stats_per_day(),\n            [pl.col('time_id') / 968, pl.col('day') / 170]\n        ]\n        expressions_2 = [e for sublist in exprs_2 for e in sublist]\n        df = df.with_columns(expressions_2)\n        df = df.with_columns(self.map_category())\n\n        ign_cols = [x for x in DataProcessor.RESPONDERS if x != 'responder_6']\\\n                  + ['partition_id'] \\\n                  # + [f\"responder_{i}_lag_1\" for i in range(9)] \\\n                  # + [f\"{col}_mean\" for col in DataProcessor.SELECT_FEATURES] \\\n                  # + [f\"{col}_rmean\" for col in DataProcessor.SELECT_FEATURES] \\\n        df = df.drop(ign_cols).fill_null(0)\n\n        return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:08.440071Z","iopub.execute_input":"2025-02-23T21:46:08.440572Z","iopub.status.idle":"2025-02-23T21:46:08.619132Z","shell.execute_reply.started":"2025-02-23T21:46:08.440538Z","shell.execute_reply":"2025-02-23T21:46:08.618192Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# pl.scan_parquet(DATA_DIR).filter(pl.col('partition_id')>4).select('date_id').first().collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:08.620567Z","iopub.execute_input":"2025-02-23T21:46:08.620864Z","iopub.status.idle":"2025-02-23T21:46:08.645391Z","shell.execute_reply.started":"2025-02-23T21:46:08.620842Z","shell.execute_reply":"2025-02-23T21:46:08.644766Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_initial_lags(df: pl.LazyFrame, date_id: int) -> pl.DataFrame:\n    df_lags = (\n        df.filter(pl.col('date_id')==date_id)\n        .select(['time_id', 'symbol_id'] + [f\"responder_{i}\" for i in range(9)])\n        .collect()\n        )\n    df_lags = df_lags.rename(dict(zip([f\"responder_{i}\" for i in range(9)],\n                            [f\"responder_{i}_lag_1\" for i in range(9)])))\n    return df_lags\n\n\ndf = pl.scan_parquet(DATA_DIR).filter(pl.col('date_id')>=700)\ndf_train = df.filter(pl.col('date_id') < 1572, pl.col('date_id')!=700)\ndf_valid = df.filter(pl.col('date_id') > 1576)\n\ntrain_lags = get_initial_lags(df, 700)\nvalid_lags = get_initial_lags(df, 1576)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:08.646525Z","iopub.execute_input":"2025-02-23T21:46:08.646769Z","iopub.status.idle":"2025-02-23T21:46:08.810920Z","shell.execute_reply.started":"2025-02-23T21:46:08.646749Z","shell.execute_reply":"2025-02-23T21:46:08.810269Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_embeddings(df: pl.LazyFrame, col: str, embd_dim: int):\n    df_col = df.filter(pl.col('time_id')==0).select(col).collect()\n    num_unique = df_col.n_unique()\n    model = keras.Sequential() # keras embedding model\n    model.add(keras.layers.Embedding(num_unique, embd_dim))\n    model.compile('rmsprop', 'mse')\n    embeddings = model.predict(df_col).reshape(-1, embd_dim)\n    embeddings = pl.DataFrame(\n        np.repeat(embeddings, 968, axis=0),\n        schema=[f\"{col}_embd_{i}\" for i in range(embd_dim)]\n        )\n    return embeddings\n\ndf_symbol = get_embeddings(df, 'symbol_id', 20)\ndf_feat_09 = get_embeddings(df, 'feature_09', 11)\ndf_feat_10 = get_embeddings(df, 'feature_10', 9)\ndf_feat_11 = get_embeddings(df, 'feature_11', 15)\n\n#  # one-hot encoding\n# df_symbol = df.select('symbol_id').collect().to_dummies()\n# df_feat_09 = df.select('feature_09').collect().to_dummies()\n# df_feat_10 = df.select('feature_10').collect().to_dummies()\n# df_feat_11 = df.select('feature_11').collect().to_dummies()\n\ndf_embeddings = pl.concat([df_symbol, df_feat_09, df_feat_10, df_feat_11],\n                          how='horizontal')\ndf_embeddings.columns = [f\"embd_{i}\" for i in range(df_embeddings.shape[1])]\ndf_embeddings = pl.concat(\n    [df.select('date_id').collect(), df_embeddings], how='horizontal'\n    )\n\ndel df_symbol, df_feat_09, df_feat_10, df_feat_11\ngc.collect()\n\nprint(df_embeddings.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:08.811759Z","iopub.execute_input":"2025-02-23T21:46:08.811967Z","iopub.status.idle":"2025-02-23T21:46:26.303874Z","shell.execute_reply.started":"2025-02-23T21:46:08.811950Z","shell.execute_reply":"2025-02-23T21:46:26.303027Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DataGenerator(tf.keras.utils.Sequence):\n\n    def __init__(\n        self, df: pl.LazyFrame, lags: pl.DataFrame, embeddings: pl.DataFrame,\n        **kwargs\n        ):\n        super().__init__(**kwargs)\n        self.df = df\n        self.lags = lags\n        self.embeddings = embeddings\n        self.date_ids = df.select('date_id').unique().sort('date_id').collect()\n        self.targets = ['responder_6']\n\n    def __len__(self):\n        return len(self.date_ids) - 1\n\n    def __getitem__(self, idx):\n        RESPONDERS = [f\"responder_{i}\" for i in range(9)]\n        LAG_RESPONDERS = [f\"responder_{i}_lag_1\" for i in range(9)]\n        EMBEDDING_COLS = self.embeddings.columns\n        TEMPORAL_COLS = ['time_id', 'day', 'date_sin', 'date_cos']\n        IGNORED_COLS = ['date_id', 'symbol_id', 'feature_09', 'feature_10',\n                       'feature_11', 'partition_id']\n\n        date_id = self.date_ids[idx]\n        df_cur = self.df.filter(pl.col('date_id') == date_id).collect()\n        df_embeddings_cur = self.embeddings.filter(pl.col('date_id') == date_id)\n\n        if idx > 0:\n            date_id_prev = self.date_ids[idx - 1]\n            df_prev_responders = (\n                self.df.filter(pl.col('date_id') == date_id_prev)\n                .select(['time_id', 'symbol_id'] + RESPONDERS)\n                .collect()\n                )\n            df_prev_responders = df_prev_responders.rename(\n                dict(zip(RESPONDERS, LAG_RESPONDERS)))\n        else:\n            df_prev_responders = self.lags\n\n        df_cur = df_cur.join(\n            df_prev_responders, on=['symbol_id', 'time_id'], how='left'\n            )\n        df_cur = pl.concat([df_cur, df_embeddings_cur.drop('date_id')], how='horizontal')\n\n        processor = DataProcessor(df_cur)\n        df_cur = processor.generate_features()\n\n        df_cur = df_cur.fill_null(0)\n        df_cur = df_cur.fill_nan(0)\n\n        trend_data = df_cur.select(pl.all().exclude(IGNORED_COLS + RESPONDERS))\n        scaled_cols = [col for col in trend_data.columns if col not in\n                       EMBEDDING_COLS + TEMPORAL_COLS]\n        trend_data = processor.min_max_scaler(trend_data, scaled_cols)\n        num_feats = trend_data.shape[1]\n        trend_data = trend_data.to_numpy().reshape(-1, 968, num_feats)\n        target_data = df_cur.select(self.targets).to_numpy().reshape(-1, 968, 1)\n\n        return trend_data, target_data\n\n    @property\n    def num_batches(self):\n        return len(self)\n\ntrain_generator = DataGenerator(df_train, train_lags, df_embeddings)\nvalid_generator = DataGenerator(df_valid, valid_lags, df_embeddings)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:26.304979Z","iopub.execute_input":"2025-02-23T21:46:26.305294Z","iopub.status.idle":"2025-02-23T21:46:26.683107Z","shell.execute_reply.started":"2025-02-23T21:46:26.305260Z","shell.execute_reply":"2025-02-23T21:46:26.682432Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for x, y in train_generator:\n    print(x.shape, y.shape)\n    break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:26.683988Z","iopub.execute_input":"2025-02-23T21:46:26.684209Z","iopub.status.idle":"2025-02-23T21:46:27.121276Z","shell.execute_reply.started":"2025-02-23T21:46:26.684185Z","shell.execute_reply":"2025-02-23T21:46:27.120289Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"keras.config.set_dtype_policy(\"mixed_float16\")\n\ndef build_model(input_sequence_length=968):\n    inputs = layers.Input(shape=(input_sequence_length, 181))\n    print(\"Expected Model Input Shape:\", inputs.shape)\n\n    x = layers.GRU(128, return_sequences=True)(inputs)\n    # x = layers.Dropout(0.1)(x)\n    # x = layers.GRU(128, return_sequences=True)(x)\n    # x = layers.Dropout(0.1)(x)\n    x = layers.GRU(128)(x)\n    x = layers.Dropout(0.1)(x)\n    outputs = layers.Dense(input_sequence_length, activation='linear')(x)\n\n    model = Model(inputs, outputs)\n    # optimizer = keras.optimizers.Adam(learning_rate=1e-3, clipnorm=1.0)\n\n    model.compile(\n        optimizer='adam',\n        loss='mse',\n        metrics=['mae']\n    )\n    return model\n\nmodel = build_model()\nmodel.summary()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:27.123457Z","iopub.execute_input":"2025-02-23T21:46:27.123732Z","iopub.status.idle":"2025-02-23T21:46:28.569117Z","shell.execute_reply.started":"2025-02-23T21:46:27.123709Z","shell.execute_reply":"2025-02-23T21:46:28.568402Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\ncallback = keras.callbacks.EarlyStopping(\n        monitor='val_loss', patience=2, restore_best_weights=True\n        )\n\nhistory = model.fit(\n        train_generator,\n        epochs=5,\n        validation_data=(valid_generator),\n        callbacks=[callback],\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T21:46:28.570011Z","iopub.execute_input":"2025-02-23T21:46:28.570350Z","iopub.status.idle":"2025-02-23T22:38:04.454858Z","shell.execute_reply.started":"2025-02-23T21:46:28.570314Z","shell.execute_reply":"2025-02-23T22:38:04.454091Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"preds = model.predict(valid_generator).reshape(-1)\nprint(preds.min())\nprint(preds.max())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T22:38:04.455726Z","iopub.execute_input":"2025-02-23T22:38:04.456014Z","iopub.status.idle":"2025-02-23T22:39:12.838597Z","shell.execute_reply.started":"2025-02-23T22:38:04.455991Z","shell.execute_reply":"2025-02-23T22:39:12.837775Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def weighted_r2(y_true, y_pred, sample_weight):\n    r2 = (1 - np.average((y_true - y_pred)**2, weights=sample_weight) /\n          (np.average(y_true**2, weights=sample_weight) + 1e-38)\n    )\n    return r2\n\n\ntargets = []\nfor i in range(len(valid_generator)):\n    targets.append(valid_generator[i][1].reshape(-1))\ntargets = np.concatenate(targets)\n\nweights = []\nfor i in range(len(valid_generator)):\n    weights.append(valid_generator[i][0].reshape(-1, 181)[:, 1])\nweights = np.concatenate(weights)\n\nweighted_r2(targets, preds, weights)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-23T22:39:12.839385Z","iopub.execute_input":"2025-02-23T22:39:12.839669Z","iopub.status.idle":"2025-02-23T22:41:09.005617Z","shell.execute_reply.started":"2025-02-23T22:39:12.839647Z","shell.execute_reply":"2025-02-23T22:41:09.004825Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}