{"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":"markdown","source":"This notebook implements the context tree weighting (CTW) algorithm and applies it to the training data from the Jane Street Real-Time Market Data Forecasting Kaggle competition. \n\n\nCOMPETITION DESCRIPTION\n\nThe data for the competition was a time series with 79 variables called \"features\" and 9 variables called \"responders.\" There is also a \"weight\" variable that determines how much a prediction for a given row is weighted in the scoring function (see below), and a symbol_id variable that is an integer that indicates what financial instrument the data in each row pertains to. Jane Street didn't provide any additional information about the variables, so it's not clear what exactly the feature and responder variables represent, and it's not clear what \"real world\" financial instrument each symbol_id value maps to. \n\nThe goal of the competition was to forecast the value of one of the responder variables (called \"responder_6\") at each time step given the current and past values of all the feature variables, and the past values of the responder variables. Because the data set contains data for multiple different financial instruments, there is typically more than one row associated with each time step, with each row corresponding to a different symbol_id value. \n\nThe scoring function used to evaluate a model's predictions is the sample weighted zero mean R-squared score for the responder_6 variable. The formula for this scoring function is given by: \n\n$$\nR^2 = 1 - \\frac{\\sum w_i (y_i - \\hat{y}_i)^2}{\\sum w_i y_i^2},\n$$\n\nwhere $y$ and $\\hat{y}$ are the ground-truth and predicted value vectors of responder_6 respectively, and $w$ is the sample weight vector. \n\n\nCTW DESCRIPTION AND WHY I THOUGHT IT WAS A GOOD FIT FOR THIS PROBLEM\n\nThe CTW algorithm is an algorithm that efficiently computes the joint probability of a sequence of symbols according to a probability distribution that is a mixture of all variable order Markov models with depth less than or equal to D, where D is a parameter set by the user. The CTW algorithm is a rare instance of a machine learning algorithm that is both practical/computationally-efficient (at least for some applications) and comes with strong theoretical guarantees. The paper where CTW was originally proposed is available [here](https://www.cs.cmu.edu/~aarti/Class/10704_Fall16/CTW.pdf). For a more accessible introduction to CTW, see Chapter 4 of [this](http://www.hutter1.net/publ/uaibook2.pdf) book. \n\nI thought it would be interesting to try to apply this algorithm to the data from the Jane Street Real-Time Market Data Forecasting Kaggle competition for a couple reasons. First, the variables in the training data are anonymized, so it's not clear what they represent (at least to me). It therefore seemed like it would be difficult to incorporate much prior information into the training/modeling process. CTW mixes over a very large and flexible hypothesis space and comes with strong theoretical guarantees (assuming there is a model in that large hypothesis space that is \"close\" to the true model), so it seemed like it might be a good fit for this problem given that I had very little background knowledge I could use to narrow the hypothesis space. \n\nSecond, the rules of this competition required being able to do model updates and make predictions within relatively tight time constraints. The CTW algorithm is computationally efficient and allows you to perform Bayesian updating and make predictions relatively quickly (in time that is linear in the depth parameter D), so that also made it seem like it might be a good fit for this problem. \n\nThird, the CTW algorithm would enable me to efficiently compute a conditional distribution over the possible values of the responder variable you're supposed to predict in this competition, rather than just computing point estimates. This is potentially useful because the \"best\" prediction to make is arguably the one that minimizes the expected difference between the true value of the responder and the predicted value. Identifying which prediction minimizes the expected value of this difference arguably requires using a conditional distribution over the value of the responder variable, rather than just a point estimate.   \n\nADAPTING THE CTW ALGORITHM FOR THE JANE STREET COMPETITION\n\nIn order to apply the CTW algorithm to the competition dataset in a coherent and computationaly-tractable way, I had to make a couple adjustments/adaptations to both the algorithm and the data. First, the CTW algorithm was developed to model sequences of symbols drawn from a finite alphabet. In contrast, most of the variables in the competition dataset were floating point (i.e. \"continuous\"). Technically, a time series dataset containing floating point numbers (and indeed any dataset that a computer can work with) can be thought of as a sequence of symbols drawn from a binary alphabet, so in theory the CTW algorithm could be applied to this sort of dataset without requiring any modifications. There are a couple reasons why this might not be ideal, however. First, the dataset has about 90 different fp32 variables, so if you used all of these for prediction this would result in a huge number of different possible contexts the CTW algorithm might need to account for (about 2^3000), and there might be relatively little data in the dataset for any given context. This might result in slow learning. Second, the size of the context tree needed to accomodate a context that contains all these fp32 variables would be huge. Although you can initialize nodes in the context tree in an \"online\" fashion so that at any given time the context tree only contains nodes that match contexts you've actually observed in the data, the context tree can still become intractably large when you train on a dataset as large as the one for this competition. \n\nGiven the above considerations, I thought it made sense to try to create reduced precision versions of the variables in the dataset. I did this by computing quantile values for each variable, assigning binary codes to each quantile bin for each variable, and then mapping specific variable values to these binary codes depending on the bin the value falls into. So far I have only done this using 10 bins/quantile ranges for each variable, but this is something that might make sense to experiment more with in the future. \n\nIn addition to creating reduced precision versions of the variables, I also had to make some adjustments to the CTW algorithm to apply it to the competition dataset. As mentioned above, the CTW algorithm was originally developed to model sequences of symbols drawn from a finite alphabet. The standard way to train/update when using the CTW algorithm is to slide the \"context window\" over one symbol at a time. In other words, viewing your data as a sequence of symbols, you update on one symbol using as context the previous D symbols, then slide the context window over by one symbol and update on the next symbol, etc. CTW was originally developed for use in lossless compression, and in this context it makes sense to use this \"sliding window\" approach since for compression you usually want to compress each symbol (or block of symbols) in a dataset in sequence. For this competition, however, we only care about predicting responder_6 (or more precisely the bits that encode the value/range for responder_6). We therefore do not want to waste model capacity trying to learn how to predict any of the other variables. As a result, the code below only performs updates on the bits associated with responder_6, not any of the other variables. In other words, rather than using a sliding window that updates on each bit in the dataset, the window \"jumps\" over all the bits associated with other variables and only updates on the bits associated with responder_6. \n\nAnother issue with the standard CTW algorithm is that the depth D is fixed and doesn't change depending on the bit you are predicting/updating on. This is an issue for this application because I'm trying to use CTW to learn how to predict responder_6, whose value/quantile bin will be encoded using multiple bits. We want to use the same external context (i.e. context from bits associated with variables other than responder_6) to predict each of these bits, but if the context window depth is fixed we will lose one bit of external context every time we slide the window over to predict the next bit of responder_6. Additionally, when predicting the nth bit of responder_6, it probably makes sense to use the preceding n-1 bits of responder_6 as context in addition to the external context. To address both of these issues, I used a separate context tree for each bit of the responder_6 encoding. The context tree associated with the first bit of responder_6 had depth D, the context tree associated with the second bit of responder_6 had depth D+1, and the context tree associated with the kth bit of responder_6 had depth (D + k - 1). This approach is sometimes called \"factored\" CTW (see [this](https://arxiv.org/pdf/0909.0801) paper for more details).    \n\nLast, the ordering of variables in the context window can impact the speed with which the CTW algorithm learns a good model. Roughly speaking, you want variables that are more likely to be useful for predicting the variable you care about to appear \"earlier\" in the context window. As a result, in another notebook I estimated the mutual information between responder_6 and each of the feature variables and ordered them so that the variables with the highest mutual information occur earlier in the context window. Technically, it would have been better to estimate which variable has the highest mutual information with responder_6, then estimate which variable has the highest mutual information with responder_6 when conditioned on the first variable selected, etc. to account for redundancy in the information each variable provides about responder_6. However, it seemed like this would be very computationally expensive and time consuming, so as an approximation I just estimated the unconditional mutual information between each variable and responder_6 and ordered them according to these estimates.\n\n\nRESULTS\n\nUnfortunately, I haven't yet figured out how to use the CTW algorithm to achieve predictive accuracy on par with the best results from the competition. The main issue I've encountered is that the memory footprint of the model quickly becomes too large for the Kaggle environment to handle as additional variables are added to the context window. As a result, I could only include at most 3 or 4 of the feature variables in the context window. This meant I couldn't use the information from the vast majority of the variables to predict responder_6.  \n\nGiven this issue, an obvious question to ask is: is this issue fundamental to the CTW algorithm (i.e. will it crop up in most applications of the algorithm), or is there something about this particular application that makes this problem more severe? \n\nIf you were to initialize all nodes in the context tree up front, the size would be O(2^D), where D is the maximum context window depth of all the models in the mixture CTW is computing. In other words, the size of the context tree is exponential in D, so the size of the context tree can become intractably large for even modest-sized values of D. However, rather than initializing all of the nodes in the context tree up front, you can instead initialize context tree nodes in an \"online\" fashion, meaning that you only initialize a node when you encounter the context that node corresponds to in the data. This is the approach used in the code below. Using this approach, the maximum size of the context tree at any given point would be O(N*D), where N is the number of symbols observed (i.e. the number of symbols you have updated/trained the model on). This will typically be much less than 2^D for even modestly-sized values of D. Moreover, if the data has a lot of structure in it so that many of the contexts occur pretty frequently, the actual size of the online context tree may be much less than this upper bound. \n\nHowever, if you have a lot of data and the data doesn't have a lot of structure in it (i.e. there aren't a lot of repeated contexts), the size of the context tree can still become quite large. This seems to be the case with this dataset. First, the training set has a fairly large number of rows (well over 10 million), and including any significant number of feature variables in the context window would require using a fairly large value of D. Therefore, the O(N*D) upper bound can be fairly large, especially for a Kaggle environment where you don't have access to large amounts of memory. Second, it seems like when you make the context window large enough to include even a modest number of variables, you encounter a lot of unique contexts in the data (i.e. there aren't a lot of contexts that repeat frequently). As a result, the amount of memory used by the context tree ends up being fairly close to the upper bound. Again, this quickly becomes intractable in a Kaggle environment. \n\nSo my sense is that while the size of the context tree is an issue you need to be concerned about whenever using the CTW algorithm, the structure of the training data and the memory limits in the training environment made this problem particularly salient for this application. \n\n\nOTHER THINGS TO POTENTIALLY TRY \n\nIt is possible the CTW algorithm is just not well-suited to modeling the data from this competition and it's better to use other approaches. However, there are several things I can think of to try that might improve the predictive accuracy of the CTW algorithm when applied to the competition data. \n\nFirst, it may be interesting to try different levels of quantization/discretization of the variables in the dataset. In the code below, I compute deciles for each variable and then encode each value of each variable based on the decile it falls into. However, it's possible that using fewer quantiles (e.g. four or five) for each variable would allow you to retain most of the useful information in the variable value, while reducing the number of bits required encode the bin each variable value falls into. This would potentially allow you to include more variables in the context window, which might help improve the accuracy of the model. Alternatively, it might make sense to use more quantiles for variables whose estimated mutual information with responder_6 is high, and then use fewer quantiles (i.e. use a more coarse-grained quantization) for other variables for which the estimated mutual information with responder_6 is lower. \n\nSecond, it might be interesting to investigate whether it's possible to alleviate the memory constraints that I faced when training using the code below. There are a couple approaches that one could take to doing this. First, you could try to make the training code more memory efficient by reducing the memory footprint of the context tree and/or other data structures used during training. For example, the memory footprint of a Python object typically includes a lot of overhead beyond just the raw values the object stores. It may be possible to reduce the size of this overhead for the context tree model itself by storing the numbers associated with all the context tree nodes in NumPy arrays, so that you only incur the overhead cost once (or a small number of times) for each context tree, rather than for each node in the context tree. However, given that we want to use an \"online\" context tree, this would need to be done in a way that enables: (1) the size of the NumPy arrays to grow dynamically over time as new contexts (corresponding to new context tree nodes) are encountered, while, (2) keeping track of which indices in an array are associated with which context tree nodes. I suspect there might be a way to do this, but I haven't spent time trying to figure it out yet. \n\nIn addition to reducing the memory footprint of the context tree and other objects used during training, it may also be possible to rewrite the code to reduce the number of objects that need to be kept in memory at any given time. I noticed that when the training code was running, there were sometimes very large (but temporary) spikes in memory usage of 10 GB or more. I think I would typically run into memory limits during these spikes in memory usage. It might be interesting to investigate what is causing these spikes and if the code can be rewritten in order to avoid them or reduce their magnitude (e.g. by forcing more regular garbage collection).\n\nAside from trying to make the code more memory efficient, it might be interesting to try to train a CTW model on the competition dataset in a computational environment with more memory resources. For example, Kaggle notebooks are limited to 30 GB of system memory, but cutting edge language models typically need 100s of GB of both system and GPU memory just for inference (I'm not sure how much memory is required during training, but I'm guessing it's at least this much). I suspect that scaling the amount of memory the CTW algorithm has access to while training on the competition data would substantially improve its predictive accuracy given CTW's ability to model a wide range of distributions. \n\nLast, there are several variants of the CTW algorithm that might be interesting to try in lieu of the \"vanilla\" CTW algorithm used in the code below. For example, [this](https://papers.nips.cc/paper_files/paper/2017/file/c366c2c97d47b02b24c3ecade4c40a01-Paper.pdf) paper proposes a CTW variant called \"dynamic depth CTW\" where: (1) you can place a constraint on the amount of memory the context tree can use, and (2) rather than setting a fixed maximum depth, the algorithm learns which contexts require more depth to model effectively and which contexts require less depth, and then decides when and where to increase the depth of the tree given this analysis and the user-specified memory constraint. \n\nAnother variant of CTW that might be interesting to try is so-called \"adaptive CTW,\" which is described in [this](https://arxiv.org/pdf/1201.2056) paper. This variant of the CTW algorithm weights more recent data more heavily than older data during updating/training, which allows it to adapt to distributional shifts. In contrast, the standard CTW algorithm assumes that the data generating distribution is stationary (i.e. that the conditional distribution for the next symbol depends only at most the previous D symbols and not on the current time). From what I understand, financial data is often non-stationary. This is sometimes dealt with by applying various transformations to raw financial data to extract features whose distribution is stationary, but it's not clear to me whether or not this was done to the data used for this competition. As a result, it may be interesting to try applying adaptive CTW to this dataset. It is possible that the performance of adaptive CTW would be better than vanilla CTW even if only a few variables can be included in the context window due to memory constraints, though I suspect it still might not be great without alleviating the memory issues I faced with vanilla CTW.  \n\nThere are several other variants of CTW that might be interesting to try applying to the competition data, including [partition tree weighting](https://arxiv.org/pdf/1211.0587), [context tree switching](https://arxiv.org/pdf/1111.3182), and [skip context tree switching](https://proceedings.mlr.press/v32/bellemare14.pdf). However, I suspect most or all of these alternatives wouldn't perform well without figuring out ways to alleviate the memory constraints I faced with the code below. \n\n","metadata":{}},{"cell_type":"code","source":"\"\"\"This cell imports various libraries the cells below use\"\"\"\n\nimport numpy as np \nimport pandas as pd \nimport math\nimport gc \nfrom typing import Optional, List, Dict, Tuple, Union, Sequence\nfrom dataclasses import dataclass\nimport os\nimport time\nimport pickle\nfrom memory_profiler import profile\nimport sys\nimport objsize\nfrom collections import defaultdict\n\n\n# Set pandas to display all columns and rows\npd.set_option('display.max_columns', None)\npd.set_option('display.max_rows', None)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-01-12T03:13:10.784118Z","iopub.execute_input":"2025-01-12T03:13:10.784548Z","iopub.status.idle":"2025-01-12T03:13:12.040652Z","shell.execute_reply.started":"2025-01-12T03:13:10.784471Z","shell.execute_reply":"2025-01-12T03:13:12.039595Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"This cell computes equal-sized quantiles for all variables for the whole \ntraining set. The number of quantiles is specified by the variable n. Note \nthat the training data is partitioned into 10 different Parquet files.\"\"\"\n\n# Define columns of interest\nweight_cols = ['weight']\nfeature_cols = [f'feature_{i:02d}' for i in range(79)]  # feature_00 ... feature_78\nresponder_cols = [f'responder_{i}' for i in range(9)]   # responder_0 ... responder_8\ncols_of_interest = weight_cols + feature_cols + responder_cols\n\n# Number of quantile divisions\nn = 10\nquantile_points = np.linspace(1/n, (n-1)/n, n-1)\n\nquantile_results = {}\npartitions = range(10)  # partition_id=0 through 9\n\nfor col in cols_of_interest:\n    # Load data for this column from each partition\n    data_list = []\n    for i in partitions:\n        file_path = f'/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet/partition_id={i}/part-0.parquet'\n        # Load just the one column\n        chunk = pd.read_parquet(file_path, columns=[col])\n        data_list.append(chunk[col].values)\n    \n    # Concatenate all data for this column\n    full_data = np.concatenate(data_list)\n    \n    # Compute quantiles using nanquantile\n    quantile_values = np.nanquantile(full_data, quantile_points)\n    \n    # Store the results\n    quantile_results[col] = quantile_values\n\n# Print the results\nfor col, vals in quantile_results.items():\n    print(f\"{col}: {vals}\")\n\n\n# Save quantiles for later use\nnp.save(f'/kaggle/working/{n}_quantile_results.npy', quantile_results)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-12T03:13:27.000052Z","iopub.execute_input":"2025-01-12T03:13:27.000665Z","iopub.status.idle":"2025-01-12T03:17:19.922987Z","shell.execute_reply.started":"2025-01-12T03:13:27.000625Z","shell.execute_reply":"2025-01-12T03:17:19.921220Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"This cell defines various classes and methods to implement\nthe factored CTW algorithm\"\"\"\n\nNumber = Union[int, float]\n\n@dataclass\nclass TimePoint:\n    \"\"\"Represents a single row of data.\"\"\" \n    \n    \"\"\"NOTE TO SELF: It would probably be good to rename this class something \n    like 'Row' or 'DataRow' since instances of this class do not actually contain \n    all data from a particular time, but rather data for a particular time and\n    symbol_id (i.e. a row in the competition dataset) . There can be data for \n    multiple symbol_ids associated with a given time (i.e. there can be multiple \n    rows associated with a given time).\"\"\"\n    \n    date_id: int  # Integer representing the date\n    time_id: int  # Index representing time within the date\n    symbol_id: int  # Symbol identifier\n    weight: float  # Weight for this timepoint\n    features: Dict[str, Optional[Number]]  # Features can be int, float, or None (NaN)\n    responders: Optional[Dict[str, float]] = None  # Responders are always float when present\n\nclass Discretizer:\n    \"\"\"Handles discretization of continuous values into binary strings\"\"\"\n    def __init__(self, boundaries: Dict[str, List[Number]]):\n        \"\"\"\n        Initialize discretizer with pre-computed boundaries and mappings.\n       \n        The \"boundaries\" argument should be a dictionary mapping variable \n        names to lists of quantile boundary values.\n        \"\"\"\n        self.boundaries = boundaries\n        self.n_bits = {}\n        self.value_to_encoding = {}  # Maps variable -> {bin_range -> encoding}\n        self.encoding_to_range = {}  # Maps variable -> {encoding tuple -> bin_range}\n        \n        for var in boundaries:\n            # Calculate number of bits needed\n            n_bins = len(boundaries[var]) + 1\n            n_bits = math.ceil(math.log2(n_bins + 1))\n            self.n_bits[var] = n_bits\n            \n            # Initialize mappings for this variable\n            self.value_to_encoding[var] = {}\n            self.encoding_to_range[var] = {}\n            \n            # Create the missing value encoding (all ones)\n            missing_encoding = tuple([1] * n_bits)\n            self.value_to_encoding[var]['missing'] = missing_encoding\n            self.encoding_to_range[var][missing_encoding] = 'missing'\n            \n            # Create encodings for each bin\n            var_boundaries = [-float('inf')] + boundaries[var] + [float('inf')]\n            for bin_idx in range(len(var_boundaries) - 1):\n                # Create binary encoding for this bin\n                bin_encoding = tuple(int(b) for b in format(bin_idx, f'0{n_bits}b'))\n                bin_range = (var_boundaries[bin_idx], var_boundaries[bin_idx + 1])\n                \n                self.value_to_encoding[var][bin_range] = bin_encoding\n                self.encoding_to_range[var][bin_encoding] = bin_range\n        \n    def encode(self, variable: str, value: Optional[Number]) -> List[int]:\n        \"\"\"\n        Convert a value to binary quantile bin encoding using precomputed mappings\n        \n        Args:\n            variable: Name of the variable\n            value: Value to encode, can be None or NaN\n            \n        Returns:\n            Binary encoding using minimal number of bits\n        \"\"\"\n        if variable not in self.value_to_encoding:\n            raise ValueError(f\"Variable {variable} not found\")\n            \n        # Handle both None and NaN as missing values\n        if value is None or (isinstance(value, float) and math.isnan(value)):\n            encoding = self.value_to_encoding[variable]['missing']\n            return list(encoding)\n        \n        # Convert int to float for comparison with boundaries if needed\n        value = float(value)\n            \n        # Find the bin range that contains this value\n        for bin_range, encoding in self.value_to_encoding[variable].items():\n            if bin_range != 'missing':\n                lower, upper = bin_range\n                if lower < value <= upper:\n                    return list(encoding)\n        \n        raise ValueError(f\"Value {value} not found in any bin range\")\n    \n    def decode(self, variable: str, encoding: List[int]) -> Union[Tuple[Number, Number], str]:\n        \"\"\"\n        Convert binary encoding back to value range using precomputed mappings\n        \n        Args:\n            variable: Name of the variable\n            encoding: Binary encoding\n            \n        Returns:\n            Tuple of (lower, upper) bounds or 'missing' if missing value encoding\n        \"\"\"\n        if variable not in self.encoding_to_range:\n            raise ValueError(f\"Variable {variable} not found\")\n            \n        expected_length = self.n_bits[variable]\n        if len(encoding) != expected_length:\n            raise ValueError(f\"Expected encoding length {expected_length}, got {len(encoding)}\")\n        \n        # Convert encoding list to tuple for dictionary lookup\n        encoding_tuple = tuple(encoding)\n        \n        if encoding_tuple not in self.encoding_to_range[variable]:\n            raise ValueError(f\"Invalid encoding {encoding}\")\n            \n        return self.encoding_to_range[variable][encoding_tuple]\n\n\n\nclass Node:\n    \"\"\"Represents a node in a context tree\"\"\"\n    def __init__(self):\n        self.a = 0  # count of zeros\n        self.b = 0  # count of ones\n        self.log_kt = 0  # log of KT estimator\n        self.log_w = 0  # log of weighted probability\n        self.left: Optional[Node] = None  # left child (context ends with 1)\n        self.right: Optional[Node] = None  # right child (context ends with 0)\n        self.parent: Optional[Node] = None\n\n\n\ndef log_sum(a: float, b: float) -> float:\n    \"\"\"\n    Implements the log-sum operator to avoid numerical underflow:\n    a ⊕ b := max(a,b) + log2(1 + 2^(min(a,b) - max(a,b)))\n    \"\"\"\n    if a == float('-inf'):\n        return b\n    if b == float('-inf'):\n        return a\n    \n    max_val = max(a, b)\n    min_val = min(a, b)\n    return max_val + math.log2(1 + pow(2, min_val - max_val))\n\n\n    \nclass CTW:\n    \"\"\"Instances of this class represent context trees. This class\n    also includes various methods for updating context trees and \n    making predictions using context trees.\"\"\"\n    \n    def __init__(self, depth: int):\n        self.depth = depth\n        self.root = Node()\n    \n    def _get_path(self, context: List[int]) -> List[Node]:\n        \"\"\"Helper function to get path through tree following context.\n        Traverses context from most recent bit to least recent.\n        Requires len(context) >= self.depth\"\"\"\n        if len(context) < self.depth:\n            raise ValueError(f\"Context must contain at least {self.depth} bits\")\n            \n        current = self.root\n        nodes = [current]\n        \n        # Get exactly depth bits from end of context and traverse from most recent\n        for ctx_bit in reversed(context[-self.depth:]):\n            if ctx_bit == 1:\n                if current.left is None:\n                    current.left = Node()\n                    current.left.parent = current\n                current = current.left\n            else:\n                if current.right is None:\n                    current.right = Node()\n                    current.right.parent = current\n                current = current.right\n            nodes.append(current)\n        \n        # Verify we processed exactly depth bits\n        assert len(nodes) == self.depth + 1  # +1 for root node\n        return nodes \n\n    def update_log_kt(self, node: Node, bit: int) -> float:\n        \"\"\"\n        Update the log of the KT estimate for a node given a new bit.\n        \n        Args:\n            node: The node to update\n            bit: The new bit (0 or 1)\n            \n        Returns:\n            The new log KT probability\n        \"\"\"\n        if bit == 1:\n            numerator = node.b + 0.5\n        else:\n            numerator = node.a + 0.5\n        \n        denominator = node.a + node.b + 1\n        new_log_kt = math.log2(numerator) - math.log2(denominator) + node.log_kt\n        \n        return new_log_kt\n\n    def update(self, context: List[int], bit: int) -> None:\n        \"\"\"Update context tree with new bit and context.\"\"\"\n        nodes = self._get_path(context)\n        \n        # First update log KT estimates using current counts\n        for node in nodes:\n            node.log_kt = self.update_log_kt(node, bit)\n            if bit == 1:\n                node.b += 1\n            else:\n                node.a += 1\n        \n        # Update weighted probabilities bottom-up\n        for node in reversed(nodes):\n            if node.left is None and node.right is None:\n                node.log_w = node.log_kt\n            else:\n                left_prob = node.left.log_w if node.left else 0\n                right_prob = node.right.log_w if node.right else 0\n                children_prob = left_prob + right_prob\n                node.log_w = log_sum(node.log_kt, children_prob) - 1\n\n    def revert(self, context: List[int], bit: int) -> None:\n        \"\"\"Revert the context tree to state before last update.\"\"\"\n        nodes = self._get_path(context)\n        \n        # First decrement counts and revert KT estimates\n        for node in nodes:\n            if bit == 1:\n                node.b -= 1\n                # Revert log KT estimate (equivalent to multiplying by (a+b+1)/(b+1/2))\n                node.log_kt = (math.log2(node.a + node.b + 1) - \n                             math.log2(node.b + 0.5) + \n                             node.log_kt)\n            else:\n                node.a -= 1\n                # Revert log KT estimate (equivalent to multiplying by (a+b+1)/(a+1/2))\n                node.log_kt = (math.log2(node.a + node.b + 1) - \n                             math.log2(node.a + 0.5) + \n                             node.log_kt)\n        \n        # Revert weighted probabilities bottom-up\n        for node in reversed(nodes):\n            if node.left is None and node.right is None:\n                node.log_w = node.log_kt\n            else:\n                left_prob = node.left.log_w if node.left else 0\n                right_prob = node.right.log_w if node.right else 0\n                children_prob = left_prob + right_prob\n                node.log_w = log_sum(node.log_kt, children_prob) - 1\n\n    def predict(self, context: List[int], bit: int) -> float:\n        \"\"\"\n        Predict probability of next bit being the specified value given context.\n        Returns P(bit|context) = P(context+bit)/P(context)\n        \"\"\"\n        # Store old root probability (log P(context))\n        old_log_prob = self.root.log_w\n        \n        # Update tree with the specified bit\n        self.update(context, bit)\n        \n        # Get new root probability (log P(context+bit))\n        new_log_prob = self.root.log_w\n        \n        # Revert the tree to its previous state\n        self.revert(context, bit)\n        \n        # Convert difference of logs to probability\n        return pow(2, new_log_prob - old_log_prob)\n\n\n\nclass FactoredCTW:\n    \"\"\"Manages multiple CTW trees for predicting responder_6\"\"\"\n    def __init__(self, base_depth: int, n_bits_responder: int):\n        \"\"\"\n        Initialize Factored CTW predictor\n        \n        Args:\n            base_depth: Number of context bits to use from other variables\n            n_bits_responder: Number of bits used to encode responder_6\n        \"\"\"\n        self.base_depth = base_depth\n        self.n_bits_responder = n_bits_responder\n        \n        # Create a CTW tree for each bit position\n        self.trees = []\n        for i in range(n_bits_responder):\n            # Depth for bit i is base_depth + i (preceding bits)\n            depth = base_depth + i\n            self.trees.append(CTW(depth))\n    \n    def _get_context_for_bit(self, \n                            external_context: List[int], \n                            responder_bits: List[int], \n                            bit_position: int) -> List[int]:\n        \"\"\"Get context for a specific bit position\"\"\"\n        # Context is external context + preceding responder bits\n        return external_context + responder_bits[:bit_position]\n\n    def compute_sequence_probability(self, external_context: List[int], bits: List[int]) -> float:\n        \"\"\"Compute probability of a specific binary sequence given context\"\"\"\n        log_prob = 0\n        current_bits = []\n        \n        for i, bit in enumerate(bits):\n            context = self._get_context_for_bit(external_context, current_bits, i)\n            prob = self.trees[i].predict(context, bit)\n            log_prob += math.log2(prob)\n            current_bits.append(bit)\n        \n        return pow(2, log_prob)\n\n    def predict(self, external_context: List[int]) -> Tuple[List[Tuple[List[int], float]], List[int]]:\n        \"\"\"\n        Predict probabilities for all possible binary sequences and find most likely sequence\n        \n        Args:\n            external_context: Context bits from other variables\n            \n        Returns:\n            Tuple of (list of (sequence, probability) tuples, most likely sequence)\n        \"\"\"\n        all_sequences = []\n        max_prob = -float('inf')\n        most_likely_sequence = None\n        \n        valid_encodings = [\n            list(encoding) \n            for encoding, value_range in self.discretizer.encoding_to_range[self.target_variable].items()\n            if value_range != 'missing'  # Exclude the missing value encoding\n        ]\n\n        for bits in valid_encodings:\n            prob = self.compute_sequence_probability(external_context, bits)\n            all_sequences.append((bits, prob))\n            \n            if prob > max_prob:\n                max_prob = prob\n                most_likely_sequence = bits\n        \n        return all_sequences, most_likely_sequence\n    \n    def update(self, external_context: List[int], responder_bits: List[int]):\n        \"\"\"Update all trees with a new observation\"\"\"\n        if len(responder_bits) != self.n_bits_responder:\n            raise ValueError(\"Wrong number of responder bits\")\n            \n        # Update each tree with its bit\n        for i in range(self.n_bits_responder):\n            context = self._get_context_for_bit(external_context, responder_bits, i)\n            self.trees[i].update(context, responder_bits[i])\n\n\n\nclass TimeSeriesPredictor:\n    \"\"\"Main class for time series prediction using Factored CTW\"\"\"\n    def __init__(self, \n                 boundaries: Dict[str, List[Number]],\n                 feature_ordering: List[str],\n                 target_variable: str):\n        \"\"\"\n        Initialize the predictor\n        \n        Args:\n            boundaries: Pre-computed boundaries for each variable\n            feature_ordering: List of feature names in desired order for context\n            target_variable: Name of variable to predict\n        \"\"\"\n        self.discretizer = Discretizer(boundaries)\n        self.feature_ordering = feature_ordering\n        self.target_variable = target_variable\n        \n        # Calculate number of bits needed for each feature (accounting for NaN)\n        def bits_needed(var: str) -> int:\n            n_bins = len(boundaries[var]) + 1  # number of regular bins\n            return math.ceil(math.log2(n_bins + 1))  # +1 for NaN encoding\n        \n        # Calculate total context size from features\n        total_context_bits = sum(bits_needed(f) for f in feature_ordering)\n        \n        # Calculate bits needed for target variable\n        n_bits_target = bits_needed(target_variable)\n        \n        # Initialize factored CTW\n        self.factored_ctw = FactoredCTW(total_context_bits, n_bits_target)\n            \n    def predict(self, timepoint: TimePoint) -> Optional[Number]:\n        \"\"\"\n        Predict value for target variable\n        \n        Args:\n            timepoint: Current timepoint with features\n            \n        Returns:\n            Prediction or None if prediction not possible\n        \"\"\"\n        # Get binary context from features in specified order\n        external_context = []\n        for feature in self.feature_ordering:\n            value = timepoint.features.get(feature)  \n            bits = self.discretizer.encode(feature, value)\n            external_context.extend(bits)\n            \n        # Get all valid encodings for the target variable (excluding 'missing')\n        valid_encodings = [\n            list(encoding) \n            for encoding, value_range in self.discretizer.encoding_to_range[self.target_variable].items()\n            if value_range != 'missing'  # Exclude the missing value encoding\n        ]\n        \n        # Get predictions only for valid encodings\n        all_sequences = []\n        for bits in valid_encodings:\n            prob = self.factored_ctw.compute_sequence_probability(external_context, bits)\n            all_sequences.append((bits, prob))\n        if not all_sequences:\n            return None\n\n        # Decode all of the bit strings into ranges, clip ranges so they\n        # are all between -5 and 5\n        all_ranges = []\n        for bits, probability in all_sequences:\n            (lower,upper) = self.discretizer.decode(self.target_variable, bits)\n            # Clip extremes\n            lower = max(-5, lower)\n            upper = min(5, upper) \n            all_ranges.append(((lower,upper), probability))\n\n        # Find the prediction that minimizes expected error\n        min_exp_error = float('inf')\n        best_prediction = None\n        for (lower,upper), probability in all_ranges:\n            prediction = 0.5 * (lower + upper)\n            expected_error = 0\n            for (lower,upper), probability in all_ranges:\n                error = (0.5 * (lower + upper) - prediction) ** 2\n                expected_error += error * probability\n            if expected_error < min_exp_error:\n                min_exp_error = expected_error\n                best_prediction = prediction\n        \n        # Return prediction that minimizes expected error\n        return best_prediction\n    \n    def update(self, timepoint: TimePoint):\n        \"\"\"\n        Update model with new observation\n        \n        Args:\n            timepoint: Timepoint containing features and target variable value\n        \"\"\"\n        if timepoint.responders is None or self.target_variable not in timepoint.responders:\n            return  # Skip if target variable value not available\n            \n        # Get binary context from features\n        external_context = []\n        for feature in self.feature_ordering:\n            value = timepoint.features.get(feature)\n            bits = self.discretizer.encode(feature, value)\n            external_context.extend(bits)\n            \n        # Get binary encoding of target variable value\n        target_value = timepoint.responders[self.target_variable]\n        target_bits = self.discretizer.encode(self.target_variable, target_value)\n            \n        # Update the model with this observation\n        self.factored_ctw.update(external_context, target_bits)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-12T03:17:46.002264Z","iopub.execute_input":"2025-01-12T03:17:46.002746Z","iopub.status.idle":"2025-01-12T03:17:46.059936Z","shell.execute_reply.started":"2025-01-12T03:17:46.002696Z","shell.execute_reply":"2025-01-12T03:17:46.058346Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"This cell applies the factored CTW algorithm to the Jane Street competition\ntraining data. In particular, it: (1) trains a model on the data, (2) for a subset \nof the data, makes predictions about the value of responder_6 for the date_id that is\none beyond the last date_id the model was trained on, (3) calculates\nthe weighted zero mean r-squared value for the subset of data for which it makes \npredictions, and (4) computes how much time updates and predictions take to see if it \nseems like they will fit within the time limits of the competition.\n\nNote: The code below only makes predictions for a subset of the data because making\npredictions for the whole dataset would increase the execution time of the code beyond \nKaggle's 12 hour session limit.\"\"\"\n\n\n###############################################################################\n# 2) LOAD AND CLEAN QUANTILE BOUNDARIES\n###############################################################################\n\"\"\"For some of the variables, some of the quantile values are the same for\nadjacent quantiles. In other words, for some variables there are enough instances of\na value y that both the k and k+1 quantiles have value y. The code below removes \nthese duplicate quantile boundaries to avoid creating unnecessary/duplicative \nbins, codes, and contexts.\"\"\"\n\nquantile_path = \"/kaggle/working/10_quantile_results.npy\"\nquantile_dict = np.load(quantile_path, allow_pickle=True).item()\n\n# Remove consecutive duplicates to avoid zero-width bins\nfor var_name, q_list in quantile_dict.items():\n    \n    q_list = q_list.tolist()\n    if len(q_list) == 0:\n        continue\n    \n    cleaned = [q_list[0]]\n    for val in q_list[1:]:\n        if val != cleaned[-1]:\n            cleaned.append(val)\n    quantile_dict[var_name] = cleaned\n\n###############################################################################\n# 3) SETUP FEATURE ORDERING & INITIALIZE MODEL\n###############################################################################\nfeature_ordering = [\n    \"feature_14\", \"feature_49\", \"feature_58\", \"feature_60\", \n    \"feature_47\"\n]\n\ntarget_variable = \"responder_6\"\n\npredictor = TimeSeriesPredictor(\n    boundaries=quantile_dict,\n    feature_ordering=feature_ordering,\n    target_variable=target_variable\n)\n\nprint(f\"Initial size of predictor object (GiB): {(objsize.get_deep_size(predictor)/(1024**3))}\")\n\n###############################################################################\n# 4) WEIGHTED ZERO-MEAN R2 (GLOBAL + PER-PARTITION)\n###############################################################################\n# Global accumulators\nsum_w_y_sqr_global = 0.0\nsum_w_res_sqr_global = 0.0\n\n# Per-partition accumulators (partition_stats[partition_id] = {\"sum_w_y_sqr\": ..., \"sum_w_res_sqr\": ...})\npartition_stats = {}\n\n\ndef init_partition_stats(part_id):\n    if part_id not in partition_stats:\n        partition_stats[part_id] = {\"sum_w_y_sqr\": 0.0, \"sum_w_res_sqr\": 0.0}\n\ndef update_r2_stats(part_id, y_true, y_pred, w):\n    global sum_w_y_sqr_global, sum_w_res_sqr_global\n    # Clipped y_true between -5,5 but we trust the data set is already clipped\n    \n    # Update global accumulators\n    sum_w_y_sqr_global += w * (y_true ** 2)\n    diff = (y_true - y_pred)\n    sum_w_res_sqr_global += w * (diff ** 2)\n    \n    # Update partition accumulators\n    init_partition_stats(part_id)\n    partition_stats[part_id][\"sum_w_y_sqr\"] += w * (y_true ** 2)\n    partition_stats[part_id][\"sum_w_res_sqr\"] += w * (diff ** 2)\n\ndef compute_r2(sum_w_y_sqr, sum_w_res_sqr):\n    if sum_w_y_sqr == 0.0:\n        return None\n    return 1.0 - (sum_w_res_sqr / sum_w_y_sqr)\n\ndef compute_global_r2():\n    return compute_r2(sum_w_y_sqr_global, sum_w_res_sqr_global)\n\ndef compute_partition_r2(part_id):\n    stats = partition_stats.get(part_id)\n    if not stats:\n        return None\n    return compute_r2(stats[\"sum_w_y_sqr\"], stats[\"sum_w_res_sqr\"])\n\n###############################################################################\n# 5) HELPER: CONVERT A ROW TO TimePoint\n###############################################################################\ndef row_to_timepoint(row) -> TimePoint:\n    feat_dict = {}\n    resp_dict = {}\n    \n    for col in row._fields:\n        if col.startswith(\"feature_\"):\n            feat_dict[col] = getattr(row, col)  \n        elif col.startswith(\"responder_\"):\n            resp_dict[col] = getattr(row, col)\n    \n    return TimePoint(\n        date_id=row.date_id,\n        time_id=row.time_id,\n        symbol_id=row.symbol_id,\n        weight=row.weight,\n        features=feat_dict,\n        responders=resp_dict\n    )\n\n\n###############################################################################\n# 6) MAIN LOOP\n###############################################################################\ntimings = []  # will store (partition_id, date_id, time_id, update_time_s, predict_time_s)\nprevious_date_id = None\nprevious_date_data = []\n\nfor partition_id in range(10):\n    path = f\"/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet/partition_id={partition_id}/part-0.parquet\"\n    if not os.path.exists(path):\n        print(f\"[WARNING] partition file not found: {path}\")\n        continue\n\n    # Determine which partitions to use to both train and make predictions\n    # and which partitions to use just for training\n    if partition_id == 9: # Only make predictions for last partition \n        predict_resp_6 = True\n    else: \n        predict_resp_6 = False\n    \n    # Read parquet & sort\n    df = pd.read_parquet(path)\n    df = df.sort_values([\"date_id\", \"time_id\", \"symbol_id\"]).reset_index(drop=True)\n    \n    # Process each date in ascending order\n    for date_id in df[\"date_id\"].unique():\n        #     For the *new* date, we first update model with previous_date_data if any\n        #     but the timing of that update will only matter for the *first* time_id we predict.\n        date_df = df.loc[df[\"date_id\"] == date_id]\n        unique_times = date_df[\"time_id\"].unique()\n        \n        # We do a bulk update once per date, but only \"time\" it for the first time_id's deadline\n        update_start = time.time()\n        if previous_date_data:\n            for old_tp in previous_date_data:\n                predictor.update(old_tp)\n        update_end = time.time()\n        update_duration = update_end - update_start\n\n        if date_id % 30 == 0: \n            print(\"Updated on data from 30 more dates!\")\n            print(f\"Current size of predictor object (GiB): {(objsize.get_deep_size(predictor)/(1024**3))}\")\n        \n        if predict_resp_6 == True: # Only make predictions for partition_id values specified above\n        \n            first_time_id = True\n            \n            for t_id in unique_times:\n                # Create sub-dataframe outside timing\n                sub_df = date_df.loc[date_df[\"time_id\"] == t_id]\n                timepoints = [row_to_timepoint(r) for r in sub_df.itertuples()]\n                \n                # Start timing predictions\n                predict_start = time.time()\n                # If this is the first time_id of the date, we've *already* done the update above\n                # so total_time = update_duration + prediction_time\n                i = 0\n                for tp in timepoints:\n                    y_pred = predictor.predict(tp)\n                    y_true = tp.responders.get(\"responder_6\", 0.0)\n                    w = tp.weight\n                    update_r2_stats(partition_id, y_true, y_pred, w)\n                predict_end = time.time()\n                predict_duration = predict_end - predict_start\n                \n                if first_time_id:\n                    # For the first time_id, the total time to respond includes update + predict\n                    total_time = update_duration + predict_duration\n                    # Record in our timings\n                    timings.append((partition_id, date_id, t_id, update_duration, predict_duration))\n                    first_time_id = False\n                    \n                    if total_time > 60.0:\n                        print(f\"[WARNING] For date={date_id}, time={t_id}, update+predict took {total_time:.2f}s > 60s\")\n                else:\n                    # For subsequent time_ids, we only track the prediction time\n                    timings.append((partition_id, date_id, t_id, 0.0, predict_duration))\n                    if predict_duration > 60.0:\n                        print(f\"[WARNING] For date={date_id}, time={t_id}, predict alone took {predict_duration:.2f}s > 60s\")\n\n            if date_id % 30 == 0:\n                print(compute_global_r2())\n                \n        # 7C) Prepare data for the next iteration (the next date)\n        current_date_rows = [row_to_timepoint(r) for r in date_df.itertuples()]\n        previous_date_data = current_date_rows\n        previous_date_id = date_id\n    \n    print(f\"Finished processing partition {partition_id}!\")\n\n\n###############################################################################\n# 8) FINAL METRICS + TIMINGS\n###############################################################################\n# Global R2\nglobal_r2 = compute_global_r2()\nprint(\"====================================================\")\nif global_r2 is not None:\n    print(f\"Global Weighted Zero-Mean R2 on responder_6: {global_r2:.6f}\")\nelse:\n    print(\"No data encountered; global R2 is undefined.\")\n\n# Per-partition R2\nfor p_id in sorted(partition_stats.keys()):\n    part_r2 = compute_partition_r2(p_id)\n    if part_r2 is not None:\n        print(f\"Partition {p_id} R2: {part_r2:.6f}\")\n    else:\n        print(f\"Partition {p_id} R2: No data\")\n\n# Timings summary\ntimings_df = pd.DataFrame(timings, columns=[\"partition_id\", \"date_id\", \"time_id\", \"update_sec\", \"predict_sec\"])\ntimings_df[\"total_sec\"] = timings_df[\"update_sec\"] + timings_df[\"predict_sec\"]\nnum_over_60_update_plus_predict = (timings_df.query(\"update_sec + predict_sec > 60\").shape[0])\nnum_over_60_predict_only = (timings_df.query(\"predict_sec > 60\").shape[0])\n\nprint(\"\\nTIMING SUMMARY\")\nprint(\"--------------\")\nprint(f\"Total # of (date_id, time_id) batches: {len(timings_df)}\")\nprint(f\"Batches with (update + predict) > 60s: {num_over_60_update_plus_predict}\")\nprint(f\"Batches with predict only > 60s:       {num_over_60_predict_only}\")\nprint(timings_df[[\"update_sec\",\"predict_sec\",\"total_sec\"]].describe())\n\n###############################################################################\n# 9) SAVE MODEL\n###############################################################################\nsave_path = \"/kaggle/working/final_ctw_model.pkl\"\nwith open(save_path, \"wb\") as f:\n    pickle.dump(predictor, f)\nprint(f\"\\n[INFO] Model saved to: {save_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-12T03:32:42.915412Z","iopub.execute_input":"2025-01-12T03:32:42.916338Z","execution_failed":"2025-01-12T15:12:50.835Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"This cell computes the number of unique contexts encountered in a dataset \nfor a given discretizer and set of features. This information helps give a sense\nof how large a context tree trained on the dataset using a given discretizer and \nset of features would be.\n\nNOTE: A decent amount of the code below is copied from the cell immediately above. \nThe reason for this is that I didn't want to have to run the cell above in order to \nrun this cell, since it typically takes the cell above several hours to finish \nexecuting. It might be cleaner to put the code from the cell above that I want to\nuse in this cell in a separate cell on its own that can run quickly rather than \nduplicating it here, but duplicating it here made the code easier to follow (for \nme at least) and doesn't seem like it will create too much confusion for others.\"\"\"\n\n\nfeature_ordering = [\n    \"feature_14\", \"feature_49\", \"feature_58\", \"feature_60\", \n    \"feature_47\"\n]\n\n\ndef row_to_timepoint(row) -> TimePoint:\n    \"\"\"Helper function to convert a row from the training dataset to a TimePoint object\"\"\"\n    \n    feat_dict = {}\n    resp_dict = {}\n    \n    for col in row._fields:\n        if col.startswith(\"feature_\"):\n            feat_dict[col] = getattr(row, col)  \n        elif col.startswith(\"responder_\"):\n            resp_dict[col] = getattr(row, col)\n    \n    return TimePoint(\n        date_id=row.date_id,\n        time_id=row.time_id,\n        symbol_id=row.symbol_id,\n        weight=row.weight,\n        features=feat_dict,\n        responders=resp_dict\n    )\n\n###############################################################################\n# 1) LOAD AND CLEAN QUANTILE BOUNDARIES\n###############################################################################\n\"\"\"For some of the variables, some of the quantile values are the same for\nadjacent quantiles. In other words, for some variables there are enough instances of\na value y that both the k and k+1 quantiles have value y. The code below removes \nthese duplicate quantile boundaries to avoid creating unnecessary/duplicative \nbins, codes, and contexts.\"\"\"\n\nquantile_path = \"/kaggle/working/10_quantile_results.npy\"\nquantile_dict = np.load(quantile_path, allow_pickle=True).item()\n\n# Remove consecutive duplicates to avoid zero-width bins\nfor var_name, q_list in quantile_dict.items():\n    \n    q_list = q_list.tolist()\n    if len(q_list) == 0:\n        continue\n    \n    cleaned = [q_list[0]]\n    for val in q_list[1:]:\n        if val != cleaned[-1]:\n            cleaned.append(val)\n    quantile_dict[var_name] = cleaned\n\ndiscretizer = Discretizer(quantile_dict)\n\n\n###############################################################################\n# 2) Compute number of different contexts\n###############################################################################\n\ncontexts_encountered = set()\nfor partition_id in range(10):\n    path = f\"/kaggle/input/jane-street-real-time-market-data-forecasting/train.parquet/partition_id={partition_id}/part-0.parquet\"\n    \n    # Read parquet\n    df = pd.read_parquet(path)\n    \n    for row in df.itertuples():\n        timepoint = row_to_timepoint(row)\n        external_context = []\n        for feature in feature_ordering:\n            value = timepoint.features.get(feature)  # Handle missing features\n            bits = discretizer.encode(feature, value)\n            external_context.extend(bits)\n        contexts_encountered.add(tuple(external_context))\n        if row.Index % 100000 == 0:\n            print(len(contexts_encountered))\n\nprint(f'total contexts encountered: {len(contexts_encountered)}')\n\n    \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-11T16:17:40.554384Z","iopub.execute_input":"2025-01-11T16:17:40.554781Z","iopub.status.idle":"2025-01-11T16:46:23.560095Z","shell.execute_reply.started":"2025-01-11T16:17:40.554748Z","shell.execute_reply":"2025-01-11T16:46:23.558999Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null}]}