{"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":30804,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Background\n\n[David Dirring](https://www.kaggle.com/romandovega) created a great heatmap visualization in Tableau and shared it in this post: [Why Does Responder 6 Reset at Time_ID 369?](https://www.kaggle.com/competitions/jane-street-real-time-market-data-forecasting/discussion/549568). I wanted to re-create this in Python, sharing the code below.\n\n## Imports","metadata":{}},{"cell_type":"code","source":"# Data processing imports\nimport numpy as np\nimport pandas as pd\nimport os\nimport polars as pl\nfrom pathlib import Path\n\n# Visualization imports\nimport matplotlib.pyplot as plt\nimport seaborn as sns","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:47:39.240481Z","iopub.execute_input":"2024-12-03T06:47:39.240938Z","iopub.status.idle":"2024-12-03T06:47:39.246729Z","shell.execute_reply.started":"2024-12-03T06:47:39.240894Z","shell.execute_reply":"2024-12-03T06:47:39.245473Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Heatmap Function","metadata":{}},{"cell_type":"code","source":"def plot_heatmap(parquet_files=None, symbol_id=33, column='responder_6', palette='RdYlBu'):\n    \"\"\"\n    Create a heatmap for a column's values filtered by symbol_id across a list of parquet files.\n\n    Parameters:\n        parquet_files (list): List of Parquet file paths to process.\n        symbol_id (int): The symbol ID to filter on (defaults to 33).\n        column (str): The column to plot (defaults to 'responder_6').\n        palette (str): The color palette to use for the heatmap (defaults to 'RdYlBu').\n    \"\"\"\n    if parquet_files is None:\n        DATA_DIR = Path('/kaggle/input/jane-street-real-time-market-data-forecasting')\n        N_PARTITION = 10\n        parquet_files = [DATA_DIR / f'train.parquet/partition_id={i}/part-0.parquet' for i in range(N_PARTITION)]\n\n    heatmap_stats = None\n\n    for parquet_file in parquet_files:\n        df = pl.scan_parquet(parquet_file)\n        df = df.filter(pl.col('symbol_id') == symbol_id)\n\n        stats = df.select(['time_id', 'date_id', column]).collect()\n        stats = stats.with_columns(pl.col(column).cast(pl.Float64))\n\n        if heatmap_stats is None:\n            heatmap_stats = stats\n        else:\n            heatmap_stats = heatmap_stats.vstack(stats)\n\n    heatmap_data = heatmap_stats.to_pandas().set_index(['time_id', 'date_id'])[column].unstack()\n\n    plt.figure(figsize=(35, 30))\n    sns.heatmap(heatmap_data, cmap='RdYlBu', annot=False, fmt='.2f', cbar=False)\n    plt.title(f'Heatmap of {column} for Symbol {symbol_id} by Date and Time')\n    plt.xlabel('Date ID')\n    plt.ylabel('Time ID')\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:21:00.267813Z","iopub.execute_input":"2024-12-03T06:21:00.268239Z","iopub.status.idle":"2024-12-03T06:21:00.277648Z","shell.execute_reply.started":"2024-12-03T06:21:00.268207Z","shell.execute_reply":"2024-12-03T06:21:00.276411Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Example Plots","metadata":{}},{"cell_type":"code","source":"# Default is all parquet files, symbol 33, responder_6\nplot_heatmap()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:10:44.856371Z","iopub.execute_input":"2024-12-03T06:10:44.856769Z","iopub.status.idle":"2024-12-03T06:10:53.821668Z","shell.execute_reply.started":"2024-12-03T06:10:44.856735Z","shell.execute_reply":"2024-12-03T06:10:53.820247Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Plot a different symbol\nplot_heatmap(symbol_id=38)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:12:59.419244Z","iopub.execute_input":"2024-12-03T06:12:59.419707Z","iopub.status.idle":"2024-12-03T06:13:06.414830Z","shell.execute_reply.started":"2024-12-03T06:12:59.419667Z","shell.execute_reply":"2024-12-03T06:13:06.412947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Plot a different feature/responder\nplot_heatmap(symbol_id=38, column='feature_01')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:13:49.493712Z","iopub.execute_input":"2024-12-03T06:13:49.494141Z","iopub.status.idle":"2024-12-03T06:13:57.103508Z","shell.execute_reply.started":"2024-12-03T06:13:49.494106Z","shell.execute_reply":"2024-12-03T06:13:57.101735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Plot another symbol and feature/responder\nplot_heatmap(symbol_id=1, column='feature_51')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:22:02.339986Z","iopub.execute_input":"2024-12-03T06:22:02.340382Z","iopub.status.idle":"2024-12-03T06:22:10.788691Z","shell.execute_reply.started":"2024-12-03T06:22:02.340349Z","shell.execute_reply":"2024-12-03T06:22:10.786838Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Another interesting feature\nplot_heatmap(symbol_id=38, column='feature_02')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:48:00.041000Z","iopub.execute_input":"2024-12-03T06:48:00.041392Z","iopub.status.idle":"2024-12-03T06:48:06.940285Z","shell.execute_reply.started":"2024-12-03T06:48:00.041359Z","shell.execute_reply":"2024-12-03T06:48:06.938497Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Another interesting feature\nplot_heatmap(symbol_id=38, column='feature_04')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-03T06:49:03.585556Z","iopub.execute_input":"2024-12-03T06:49:03.585958Z","iopub.status.idle":"2024-12-03T06:49:11.125324Z","shell.execute_reply.started":"2024-12-03T06:49:03.585925Z","shell.execute_reply":"2024-12-03T06:49:11.123444Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}