{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":84493,"databundleVersionId":9871156,"sourceType":"competition"}],"dockerImageVersionId":30786,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import polars as pl\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom matplotlib import cm\nimport seaborn as sns\nfrom pathlib import Path\nimport gc\nfrom tqdm import tqdm\nimport math\nfrom lightgbm import LGBMRegressor, plot_importance\nfrom sklearn.metrics import r2_score","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:36:52.675169Z","iopub.execute_input":"2024-11-27T23:36:52.675814Z","iopub.status.idle":"2024-11-27T23:36:55.864929Z","shell.execute_reply.started":"2024-11-27T23:36:52.675641Z","shell.execute_reply":"2024-11-27T23:36:55.863352Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def evaluate_model(model, X_train, y_train, X_val, y_val, weights, feature_names):   \n    model.fit(X_train, y_train, feature_name=feature_names)\n    pred_val = model.predict(X_val).clip(-5, 5)\n    return r2_score(y_val, pred_val, sample_weight=weights)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:36:55.866833Z","iopub.execute_input":"2024-11-27T23:36:55.867384Z","iopub.status.idle":"2024-11-27T23:36:55.873210Z","shell.execute_reply.started":"2024-11-27T23:36:55.867349Z","shell.execute_reply":"2024-11-27T23:36:55.872014Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BASE_PATH = Path('/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet')\nPERIOD = 7\nBASE_FEATURES = [f'feature_{i:02d}' for i in range(79)]\nTOTAL_COLUMNS = ['date_id', 'time_id', 'symbol_id', 'weight', 'responder_6'] + BASE_FEATURES\n\ntrain_ds = pl.concat([\n    pl.scan_parquet(BASE_PATH / f'partition_id={i}' / 'part-0.parquet')\n    for i in range(8, 9)\n])\nval_ds = pl.scan_parquet(BASE_PATH / 'partition_id=9' / 'part-0.parquet')\n\nlast_train_date = train_ds.select(pl.col('date_id').max()).collect().item()\nfirst_valid_val_date = last_train_date + PERIOD + 1\n\ncomplete_ds = pl.concat([train_ds, val_ds]).select(TOTAL_COLUMNS)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:36:55.874496Z","iopub.execute_input":"2024-11-27T23:36:55.874839Z","iopub.status.idle":"2024-11-27T23:36:56.121231Z","shell.execute_reply.started":"2024-11-27T23:36:55.874806Z","shell.execute_reply":"2024-11-27T23:36:56.119969Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"params = {'n_estimators': 200, 'max_depth': 3, 'num_leaves': 8, 'learning_rate': 5e-2}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:36:58.603731Z","iopub.execute_input":"2024-11-27T23:36:58.604151Z","iopub.status.idle":"2024-11-27T23:36:58.609988Z","shell.execute_reply.started":"2024-11-27T23:36:58.604113Z","shell.execute_reply":"2024-11-27T23:36:58.608732Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# using partition 7+8 for train yields 0.0061\n# with only partition 8 we have 0.0049\ntrain_data = complete_ds.filter(pl.col('date_id') <= last_train_date) \\\n    .select(BASE_FEATURES + ['responder_6']).collect().to_numpy()\nval_data = complete_ds.filter(pl.col('date_id') >= first_valid_val_date) \\\n    .select(BASE_FEATURES + ['responder_6', 'weight']).collect().to_numpy()\n\nmodel = LGBMRegressor(**params)\n\nevaluate_model(model, train_data[:, :-1], train_data[:, -1], val_data[:, :-2], val_data[:, -2], val_data[:, -1], BASE_FEATURES)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T22:32:41.949350Z","iopub.execute_input":"2024-11-27T22:32:41.949713Z","iopub.status.idle":"2024-11-27T22:36:06.076612Z","shell.execute_reply.started":"2024-11-27T22:32:41.949672Z","shell.execute_reply":"2024-11-27T22:36:06.075369Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(10, 10))\nplot_importance(model, importance_type='gain', title='Feature importance (Gain)', ax=ax)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T22:36:06.078266Z","iopub.execute_input":"2024-11-27T22:36:06.078727Z","iopub.status.idle":"2024-11-27T22:36:07.329708Z","shell.execute_reply.started":"2024-11-27T22:36:06.078677Z","shell.execute_reply":"2024-11-27T22:36:07.328494Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del train_data, val_data, model\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T22:36:07.333432Z","iopub.execute_input":"2024-11-27T22:36:07.333838Z","iopub.status.idle":"2024-11-27T22:36:07.729479Z","shell.execute_reply.started":"2024-11-27T22:36:07.333801Z","shell.execute_reply":"2024-11-27T22:36:07.728338Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Intra-Stock Aggregations and Normalizations","metadata":{}},{"cell_type":"code","source":"MEAN_FEATURES = [0, 2, 3, 5, 6, 7, 18, 19, 34, 35, 36, 37, 38, 41, 43, 44, 48, 53, 55, 59, 62, 65, 68, 73, 74, 75, 76, 77, 78]\nSTD_FEATURES = [39, 42, 46, 53, 57, 66]\nSKEW_FEATURES = [5, 40, 41, 42, 43, 44]\nZSCORE_FEATURES = [1, 36, 40, 45, 48, 49, 51, 52, 53, 54, 55, 59, 60]\n\nlen(MEAN_FEATURES + STD_FEATURES + SKEW_FEATURES + ZSCORE_FEATURES)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:37:06.413469Z","iopub.execute_input":"2024-11-27T23:37:06.413906Z","iopub.status.idle":"2024-11-27T23:37:06.423846Z","shell.execute_reply.started":"2024-11-27T23:37:06.413865Z","shell.execute_reply":"2024-11-27T23:37:06.422586Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"complete_ds = complete_ds.with_columns(\n    pl.col([f'feature_{j:02d}' for j in set(MEAN_FEATURES + ZSCORE_FEATURES)]).mean().over(['date_id', 'time_id']).name.suffix('_mean'),\n    pl.col([f'feature_{j:02d}' for j in set(STD_FEATURES + ZSCORE_FEATURES)]).std().over(['date_id', 'time_id']).name.suffix('_std'),\n    pl.col([f'feature_{j:02d}' for j in SKEW_FEATURES]).skew().over(['date_id', 'time_id']).name.suffix('_skew'),\n).with_columns(\n    pl.col(f'feature_{j:02d}').sub(f'feature_{j:02d}_mean').truediv(f'feature_{j:02d}_std').name.suffix('_zscore') for j in ZSCORE_FEATURES\n).drop([f'feature_{j:02d}_std' for j in ZSCORE_FEATURES if j not in STD_FEATURES] + \\\n       [f'feature_{j:02d}_mean' for j in ZSCORE_FEATURES if j not in MEAN_FEATURES])\n\nfeature_names = BASE_FEATURES + [f'feature_{j:02d}_mean' for j in MEAN_FEATURES] + \\\n    [f'feature_{j:02d}_std' for j in STD_FEATURES] + [f'feature_{j:02d}_skew' for j in SKEW_FEATURES] + \\\n    [f'feature_{j:02d}_zscore' for j in ZSCORE_FEATURES]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:37:08.192737Z","iopub.execute_input":"2024-11-27T23:37:08.193742Z","iopub.status.idle":"2024-11-27T23:37:08.206852Z","shell.execute_reply.started":"2024-11-27T23:37:08.193686Z","shell.execute_reply":"2024-11-27T23:37:08.205318Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_data = complete_ds.filter(pl.col('date_id') <= last_train_date) \\\n    .drop(['date_id', 'time_id', 'symbol_id', 'weight']).collect().to_numpy().astype(np.float32)\nval_data = complete_ds.filter(pl.col('date_id') >= first_valid_val_date) \\\n    .drop(['date_id', 'time_id', 'symbol_id']).collect().to_numpy().astype(np.float32)\n\nmodel = LGBMRegressor(**params)\n\nevaluate_model(model, train_data[:, 1:], train_data[:, 0], val_data[:, 2:], val_data[:, 1], val_data[:, 0], feature_names)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:19:54.174721Z","iopub.execute_input":"2024-11-27T23:19:54.175286Z","iopub.status.idle":"2024-11-27T23:29:40.230522Z","shell.execute_reply.started":"2024-11-27T23:19:54.175244Z","shell.execute_reply":"2024-11-27T23:29:40.228897Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(10, 15))\nplot_importance(model, importance_type='gain', title='Feature importance (Gain)', ax=ax)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:29:40.232621Z","iopub.execute_input":"2024-11-27T23:29:40.233031Z","iopub.status.idle":"2024-11-27T23:29:41.944403Z","shell.execute_reply.started":"2024-11-27T23:29:40.232993Z","shell.execute_reply":"2024-11-27T23:29:41.943236Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del train_data, val_data, model\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:30:14.764825Z","iopub.execute_input":"2024-11-27T23:30:14.765272Z","iopub.status.idle":"2024-11-27T23:30:15.009946Z","shell.execute_reply.started":"2024-11-27T23:30:14.765233Z","shell.execute_reply":"2024-11-27T23:30:15.008850Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Rolling Aggregations and Normalizations","metadata":{}},{"cell_type":"code","source":"# STD_FEATURES = [0, 1, 2, 3, 5, 6, 7, 15, 17, 18, 20, 22, 23, 28, 29, 32, 45, 46, 56, 61, 65]\n# MEDIAN_DEV_FEATURES = [4, 5, 6, 7, 33, 34, 35, 36, 37, 47, 49, 51, 52, 54, 57, 58, 59, 60, 68]\n# ZSCORE_FEATURES = [4, 5, 6, 7, 8, 13, 14, 33, 52, 68]\n\n# len(STD_FEATURES + MEDIAN_DEV_FEATURES + ZSCORE_FEATURES)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:37:12.806450Z","iopub.execute_input":"2024-11-27T23:37:12.806947Z","iopub.status.idle":"2024-11-27T23:37:12.815905Z","shell.execute_reply.started":"2024-11-27T23:37:12.806908Z","shell.execute_reply":"2024-11-27T23:37:12.814525Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# complete_ds = complete_ds.sort(['date_id', 'time_id']).with_columns(\n#         pl.col([f'feature_{j:02d}' for j in set(STD_FEATURES + ZSCORE_FEATURES)]).rolling_std(50, min_periods=1).over('symbol_id').cast(pl.Float32).name.suffix('_rolling_std'), \n#         pl.col([f'feature_{j:02d}' for j in ZSCORE_FEATURES]).rolling_mean(50, min_periods=1).over('symbol_id').cast(pl.Float32).name.suffix('_rolling_mean'), \n#         pl.col([f'feature_{j:02d}' for j in MEDIAN_DEV_FEATURES]).rolling_median(50, min_periods=1).over('symbol_id').cast(pl.Float32).name.suffix('_rolling_median'),\n#     ).with_columns(\n#         [pl.col(f'feature_{j:02d}').sub(f'feature_{j:02d}_rolling_mean') \\\n#             .truediv(f'feature_{j:02d}_rolling_std').cast(pl.Float32) \\\n#             .fill_nan(0.0).fill_null(0.0).alias(f'feature_{j:02d}_rolling_zscore') \n#          for j in ZSCORE_FEATURES] + \\\n#         [pl.col(f'feature_{j:02d}').sub(f'feature_{j:02d}_rolling_median') \\\n#             .cast(pl.Float32).fill_nan(0.0).fill_null(0.0) \\\n#             .alias(f'feature_{j:02d}_rolling_median_deviation') \n#          for j in MEDIAN_DEV_FEATURES]\n#     ).drop([f'feature_{j:02d}_rolling_mean' for j in ZSCORE_FEATURES] + \\\n#            [f'feature_{j:02d}_rolling_median' for j in MEDIAN_DEV_FEATURES] + \\\n#            [f'feature_{j:02d}_rolling_std' for j in ZSCORE_FEATURES if j not in STD_FEATURES])\n\n# feature_names += [f'feature_{j:02d}_rolling_std' for j in STD_FEATURES] + \\\n#     [f'feature_{j:02d}_rolling_zscore' for j in ZSCORE_FEATURES] + \\\n#     [f'feature_{j:02d}_rolling_median_deviation' for j in MEDIAN_DEV_FEATURES]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:37:20.066679Z","iopub.execute_input":"2024-11-27T23:37:20.067064Z","iopub.status.idle":"2024-11-27T23:37:20.082440Z","shell.execute_reply.started":"2024-11-27T23:37:20.067027Z","shell.execute_reply":"2024-11-27T23:37:20.081134Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# train_data = complete_ds.filter(pl.col('date_id') <= last_train_date) \\\n#     .drop(['date_id', 'time_id', 'symbol_id', 'weight']).collect().to_numpy().astype(np.float32)\n# val_data = complete_ds.filter(pl.col('date_id') >= first_valid_val_date) \\\n#     .drop(['date_id', 'time_id', 'symbol_id']).collect().to_numpy().astype(np.float32)\n\n# model = LGBMRegressor(**params)\n\n# evaluate_model(model, train_data[:, 1:], train_data[:, 0], val_data[:, 2:], val_data[:, 1], val_data[:, 0], feature_names)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:37:23.525473Z","iopub.execute_input":"2024-11-27T23:37:23.525875Z","iopub.status.idle":"2024-11-27T23:49:22.566980Z","shell.execute_reply.started":"2024-11-27T23:37:23.525841Z","shell.execute_reply":"2024-11-27T23:49:22.564870Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# fig, ax = plt.subplots(1, 1, figsize=(10, 20))\n# plot_importance(model, importance_type='gain', title='Feature importance (Gain)', ax=ax)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:49:22.570170Z","iopub.execute_input":"2024-11-27T23:49:22.570628Z","iopub.status.idle":"2024-11-27T23:49:24.975470Z","shell.execute_reply.started":"2024-11-27T23:49:22.570570Z","shell.execute_reply":"2024-11-27T23:49:24.974262Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# del train_data, val_data, model\n# gc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:49:24.977441Z","iopub.execute_input":"2024-11-27T23:49:24.977887Z","iopub.status.idle":"2024-11-27T23:49:25.241036Z","shell.execute_reply.started":"2024-11-27T23:49:24.977841Z","shell.execute_reply":"2024-11-27T23:49:25.239835Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Nearest Neighbors","metadata":{}},{"cell_type":"code","source":"schema_dict = {'date_id': pl.Int16, 'time_id': pl.Int16, 'symbol_id': pl.Int8, 'responder_6': pl.Float32}\nschema_list = [('date_id', pl.Int16), ('time_id', pl.Int16), \n               ('symbol_id', pl.Int8), ('responder_6', pl.Float32)]\nfor n in [3, 4, 5, 7, 10, 15, 25, 50]:\n    schema_dict[f'responder_6_knn_{n}'] = pl.Float32\n    schema_list.append((f'responder_6_knn_{n}', pl.Float32))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:49:25.243272Z","iopub.execute_input":"2024-11-27T23:49:25.243626Z","iopub.status.idle":"2024-11-27T23:49:25.260493Z","shell.execute_reply.started":"2024-11-27T23:49:25.243579Z","shell.execute_reply":"2024-11-27T23:49:25.259086Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import pairwise_distances\n\ndef compute_nearest_neighbor_features_batched(group_df, first_valid_date_id):\n    last_date_id = group_df.select(pl.col('date_id').max()).item()\n    if last_date_id < first_valid_date_id:\n        # for interleaving samples between batches where this function is called but is not useful\n        data = {'date_id': [-1], 'time_id': [-1], 'symbol_id': [None], 'responder_6': [None]}\n        for n in [3, 4, 5, 7, 10, 15, 25, 50]:\n            data[f'responder_6_knn_{n}'] = [None]\n        return pl.DataFrame(data, schema=schema_dict)\n    \n    cols = ['date_id', 'time_id', 'symbol_id', 'responder_6']\n    other_data = group_df.filter(pl.col('date_id') == last_date_id).select(cols)    \n    samples = group_df.filter(pl.col('date_id') != last_date_id).select([f'feature_{i:02d}' for i in range(79)] + ['responder_6']).fill_null(0).to_numpy()\n    if samples.shape[0] == 0:\n        # just for the first date\n        return other_data.with_columns(pl.lit(0.0).cast(pl.Float32).alias(f'responder_6_knn_{n}') \n                                       for n in [3, 4, 5, 7, 10, 15, 25, 50])\n        \n    target_samples = group_df.filter(pl.col('date_id') == last_date_id).select([f'feature_{i:02d}' for i in range(79)]).fill_null(0).to_numpy()\n            \n    # remove responder values that are in the last column\n    distances = pairwise_distances(target_samples, samples[:, :-1])\n    sorted_indices = np.argsort(distances, axis=-1)\n    for n in [3, 4, 5, 7, 10, 15, 25, 50]:\n        knn_responders = samples[sorted_indices[:, :n].flatten(), -1].reshape(target_samples.shape[0], n)\n        agg_responders = np.mean(knn_responders, axis=-1).flatten()\n        other_data = other_data.with_columns(pl.Series(agg_responders).cast(pl.Float32).alias(f'responder_6_knn_{n}'))\n    return other_data","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:49:25.262276Z","iopub.execute_input":"2024-11-27T23:49:25.262670Z","iopub.status.idle":"2024-11-27T23:49:25.278776Z","shell.execute_reply.started":"2024-11-27T23:49:25.262595Z","shell.execute_reply":"2024-11-27T23:49:25.277482Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"batch_size_days = 25\nperiod = 7\nstart_date_id = train_ds.select(pl.col('date_id').min()).collect().item()\nnum_date_id = train_ds.select(pl.col('date_id').max() - pl.col('date_id').min()).collect().item()\nmax_date_id = start_date_id + num_date_id\nbatches = None\n\nwith tqdm(total=math.ceil(num_date_id / (batch_size_days - period))) as pbar:\n    while start_date_id <= max_date_id:\n        first_valid_date_id = 0 if start_date_id == 0 else start_date_id + period\n        \n        batch_ds = complete_ds.filter(pl.col('date_id').is_between(start_date_id, start_date_id+25)) \\\n            .with_columns(pl.concat_str(['date_id', 'symbol_id']).alias('day_stock')).collect()\n        batch_ds = batch_ds.group_by('day_stock') \\\n            .map_groups(lambda group_df: compute_nearest_neighbor_features_batched(\n                batch_ds.filter(\n                    pl.col('symbol_id').eq(group_df.select(pl.col('symbol_id').first()).item()) &\n                    pl.col('date_id').ge(group_df.select(pl.col('date_id').first()).item() - period) &\n                    pl.col('date_id').lt(group_df.select(pl.col('date_id').first()).item())\n                ).vstack(group_df),\n                first_valid_date_id=first_valid_date_id\n            )).filter(pl.col('date_id') > -1).drop('responder_6')\n    \n        batches = batch_ds if batches is None else batches.vstack(batch_ds)\n            \n        start_date_id = min(max_date_id + 1, start_date_id + batch_size_days - period)\n        pbar.update(1)\n\nbatches = batches.sort(['date_id', 'time_id', 'symbol_id'])\nbatches.tail()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-27T23:49:25.280448Z","iopub.execute_input":"2024-11-27T23:49:25.280861Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"complete_ds = complete_ds.join(batches.lazy(), how='left', on=['date_id', 'time_id', 'symbol_id'])\n\nfeature_names += [f'responder_6_knn_{n}' for n in [3, 4, 5, 7, 10, 15, 25, 50]]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_data = complete_ds.filter(pl.col('date_id') <= last_train_date) \\\n    .drop(['date_id', 'time_id', 'symbol_id', 'weight']).collect().to_numpy().astype(np.float32)\nval_data = complete_ds.filter(pl.col('date_id') >= first_valid_val_date) \\\n    .drop(['date_id', 'time_id', 'symbol_id']).collect().to_numpy().astype(np.float32)\n\nmodel = LGBMRegressor(**params)\n\nevaluate_model(model, train_data[:, 1:], train_data[:, 0], val_data[:, 2:], val_data[:, 1], val_data[:, 0], feature_names)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, ax = plt.subplots(1, 1, figsize=(10, 22))\nplot_importance(model, importance_type='gain', title='Feature importance (Gain)', ax=ax)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}