From b6b207a4edc5412b4a5cfc7e74573eeac3a8cef1 Mon Sep 17 00:00:00 2001 From: siriuz42 Date: Wed, 8 Oct 2025 22:59:07 +0000 Subject: [PATCH 1/4] bug fix --- src/timesfm/timesfm_2p5/timesfm_2p5_flax.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py index 51d6a3c..8da19b8 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py @@ -523,7 +523,7 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5): f"Continuous quantile head is not supported for horizons > {self.model.os}." ) - self.forecast_config = forecast_config + self.forecast_config = fc self.model.compile( context=self.forecast_config.max_context, horizon=self.forecast_config.max_horizon, From c24db2afd248a4f63167eb0f572f5c03b09e4bc9 Mon Sep 17 00:00:00 2001 From: siriuz42 Date: Tue, 28 Oct 2025 22:29:40 +0000 Subject: [PATCH 2/4] covariate support --- pyproject.toml | 4 + src/timesfm/timesfm_2p5/timesfm_2p5_base.py | 231 +++++++++++++++++++- 2 files changed, 234 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index facfbdb..0b7f1dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,10 @@ flax = [ "jaxtyping", "jax[cuda]" ] +xreg = [ + "jax[cuda]", + "scikit-learn", +] [tool.ruff] line-length = 88 diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py index f438e40..5b79012 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py @@ -15,8 +15,9 @@ """TimesFM 2p5 base implementation.""" import dataclasses -from typing import Any, Callable +from typing import Any, Callable, Sequence +import collections import numpy as np from .. import configs @@ -25,6 +26,8 @@ ResidualBlockConfig = configs.ResidualBlockConfig StackedTransformersConfig = configs.StackedTransformersConfig TransformerConfig = configs.TransformerConfig ForecastConfig = configs.ForecastConfig +Category = int | str +XRegMode = str def strip_leading_nans(arr): @@ -191,3 +194,229 @@ class TimesFM_2p5: output_points = np.concatenate(output_points, axis=0) output_quantiles = np.concatenate(output_quantiles, axis=0) return output_points[:num_inputs], output_quantiles[:num_inputs] + + def forecast_with_covariates( + self, + inputs: list[Sequence[float]], + dynamic_numerical_covariates: dict[str, Sequence[Sequence[float]]] | None = None, + dynamic_categorical_covariates: ( + dict[str, Sequence[Sequence[Category]]] | None + ) = None, + static_numerical_covariates: dict[str, Sequence[float]] | None = None, + static_categorical_covariates: dict[str, Sequence[Category]] | None = None, + xreg_mode: XRegMode = "xreg + timesfm", + normalize_xreg_target_per_input: bool = True, + ridge: float = 0.0, + max_rows_per_col: int = 0, + force_on_cpu: bool = False, + ): + """Forecasts on a list of time series with covariates. + + To optimize inference speed, avoid string valued categorical covariates. + + Args: + inputs: A list of time series forecast contexts. Each context time series + should be in a format convertible to JTensor by `jnp.array`. + dynamic_numerical_covariates: A dict of dynamic numerical covariates. + dynamic_categorical_covariates: A dict of dynamic categorical covariates. + static_numerical_covariates: A dict of static numerical covariates. + static_categorical_covariates: A dict of static categorical covariates. + xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "xreg + timesfm" + fits a model on the residuals of the TimesFM forecast. "timesfm + xreg" + fits a model on the targets then forecasts on the residuals via TimesFM. + normalize_xreg_target_per_input: whether to normalize the xreg target per + input in the given batch. + ridge: ridge penalty for the linear model. + max_rows_per_col: max number of rows per column for the linear model. + force_on_cpu: whether to force running on cpu for the linear model. + + Returns: + A tuple of two lists. The first is the outputs of the model. The second is + the outputs of the xreg. + """ + if self.forecast_config is None: + raise ValueError("Model is not compiled. Please call compile() first.") + elif not self.forecast_config.return_backcast: + raise ValueError( + "For XReg, `return_backcast` must be set to True in the forecast config. Please recompile the model." + ) + + from ..utils import xreg_lib + + # Verify and bookkeep covariates. + if not ( + dynamic_numerical_covariates + or dynamic_categorical_covariates + or static_numerical_covariates + or static_categorical_covariates + ): + raise ValueError( + "At least one of dynamic_numerical_covariates," + " dynamic_categorical_covariates, static_numerical_covariates," + " static_categorical_covariates must be set." + ) + + # Track the lengths of (1) each input, (2) the part that can be used in the + # linear model, and (3) the horizon. + input_lens, train_lens, test_lens = [], [], [] + + for i, input_ts in enumerate(inputs): + input_len = len(input_ts) + input_lens.append(input_len) + + if xreg_mode == "timesfm + xreg": + # For fitting residuals, no TimesFM forecast on the first patch. + train_lens.append(max(0, input_len - self.model.p)) + elif xreg_mode == "xreg + timesfm": + train_lens.append(input_len) + else: + raise ValueError(f"Unsupported mode: {xreg_mode}") + + if dynamic_numerical_covariates: + test_lens.append( + len(list(dynamic_numerical_covariates.values())[0][i]) - input_len + ) + elif dynamic_categorical_covariates: + test_lens.append( + len(list(dynamic_categorical_covariates.values())[0][i]) - input_len + ) + else: + test_lens.append(self.forecast_config.max_horizon) + + if test_lens[-1] > self.forecast_config.max_horizon: + raise ValueError( + "Forecast horizon length inferred from the dynamic covaraites is longer than the" + f"max_horizon defined in the forecast config: {test_lens[-1]} > {self.forecast_config.max_horizon=}." + ) + + # Prepare the covariates into train and test. + train_dynamic_numerical_covariates = collections.defaultdict(list) + test_dynamic_numerical_covariates = collections.defaultdict(list) + train_dynamic_categorical_covariates = collections.defaultdict(list) + test_dynamic_categorical_covariates = collections.defaultdict(list) + for covariates, train_covariates, test_covariates in ( + ( + dynamic_numerical_covariates, + train_dynamic_numerical_covariates, + test_dynamic_numerical_covariates, + ), + ( + dynamic_categorical_covariates, + train_dynamic_categorical_covariates, + test_dynamic_categorical_covariates, + ), + ): + if not covariates: + continue + for covariate_name, covariate_values in covariates.items(): + for input_len, train_len, covariate_value in zip( + input_lens, train_lens, covariate_values + ): + train_covariates[covariate_name].append( + covariate_value[(input_len - train_len) : input_len] + ) + test_covariates[covariate_name].append(covariate_value[input_len:]) + + # Fit models. + if xreg_mode == "timesfm + xreg": + # Forecast via TimesFM then fit a model on the residuals. + point_outputs, quantile_outputs = self.forecast( + horizon=self.forecast_config.max_horizon, inputs=inputs + ) + targets = [ + ( + np.array(input_ts)[-train_len:] + - point_output[: -self.forecast_config.max_horizon][-train_len:] + ) + for input_ts, point_output, train_len in zip(inputs, point_outputs, train_lens) + ] + per_instance_stats = None + if normalize_xreg_target_per_input: + targets, per_instance_stats = xreg_lib.normalize(targets) + xregs = xreg_lib.BatchedInContextXRegLinear( + targets=targets, + train_lens=train_lens, + test_lens=test_lens, + train_dynamic_numerical_covariates=train_dynamic_numerical_covariates, + test_dynamic_numerical_covariates=test_dynamic_numerical_covariates, + train_dynamic_categorical_covariates=train_dynamic_categorical_covariates, + test_dynamic_categorical_covariates=test_dynamic_categorical_covariates, + static_numerical_covariates=static_numerical_covariates, + static_categorical_covariates=static_categorical_covariates, + ).fit( + ridge=ridge, + one_hot_encoder_drop=None if ridge > 0 else "first", + max_rows_per_col=max_rows_per_col, + force_on_cpu=force_on_cpu, + debug_info=False, + assert_covariates=True, + assert_covariate_shapes=True, + ) + if normalize_xreg_target_per_input: + xregs = xreg_lib.renormalize(xregs, per_instance_stats) + xregs = np.array(xregs) + new_point_outputs = [ + (point_output[-self.forecast_config.max_horizon :][:test_len] + xreg) + for point_output, test_len, xreg in zip(point_outputs, test_lens, xregs) + ] + new_quantile_outputs = [ + ( + quantile_output[-self.forecast_config.max_horizon :][:test_len] + + xreg[..., None] + ) + for quantile_output, test_len, xreg in zip(quantile_outputs, test_lens, xregs) + ] + + else: + # Fit a model on the targets then forecast on the residuals via TimesFM. + targets = [ + np.array(input_ts)[-train_len:] + for input_ts, train_len in zip(inputs, train_lens) + ] + per_instance_stats = None + if normalize_xreg_target_per_input: + targets, per_instance_stats = xreg_lib.normalize(targets) + xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear( + targets=targets, + train_lens=train_lens, + test_lens=test_lens, + train_dynamic_numerical_covariates=train_dynamic_numerical_covariates, + test_dynamic_numerical_covariates=test_dynamic_numerical_covariates, + train_dynamic_categorical_covariates=train_dynamic_categorical_covariates, + test_dynamic_categorical_covariates=test_dynamic_categorical_covariates, + static_numerical_covariates=static_numerical_covariates, + static_categorical_covariates=static_categorical_covariates, + ).fit( + ridge=ridge, + one_hot_encoder_drop=None if ridge > 0 else "first", + max_rows_per_col=max_rows_per_col, + force_on_cpu=force_on_cpu, + debug_info=True, + assert_covariates=True, + assert_covariate_shapes=True, + ) + point_outputs, quantile_outputs = self.forecast( + horizon=self.forecast_config.max_horizon, + inputs=[ + target - xreg_on_context + for target, xreg_on_context in zip(targets, xregs_on_context) + ], + ) + new_point_outputs = [ + (point_output[-self.forecast_config.max_horizon :][:test_len] + xreg) + for point_output, test_len, xreg in zip(point_outputs, test_lens, xregs) + ] + new_quantile_outputs = [ + ( + quantile_output[-self.forecast_config.max_horizon :][:test_len] + + xreg[..., None] + ) + for quantile_output, test_len, xreg in zip(quantile_outputs, test_lens, xregs) + ] + if normalize_xreg_target_per_input: + new_point_outputs = xreg_lib.renormalize(new_point_outputs, per_instance_stats) + new_quantile_outputs = xreg_lib.renormalize( + new_quantile_outputs, per_instance_stats + ) + + return new_point_outputs, new_quantile_outputs From 0d4d5141d795f473e435e77b8eaab521a0b89d28 Mon Sep 17 00:00:00 2001 From: siriuz42 Date: Tue, 28 Oct 2025 22:30:24 +0000 Subject: [PATCH 3/4] covariate support --- src/timesfm/utils/xreg_lib.py | 520 ++++++++++++++++++++++++++++++++++ 1 file changed, 520 insertions(+) create mode 100644 src/timesfm/utils/xreg_lib.py diff --git a/src/timesfm/utils/xreg_lib.py b/src/timesfm/utils/xreg_lib.py new file mode 100644 index 0000000..7a1b19b --- /dev/null +++ b/src/timesfm/utils/xreg_lib.py @@ -0,0 +1,520 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Helper functions for in-context covariates and regression.""" + +import itertools +import math +from typing import Any, Iterable, Literal, Mapping, Sequence + +try: + import jax + import jax.numpy as jnp + import numpy as np + from sklearn import preprocessing +except ImportError: + raise ImportError( + "Failed to load the XReg module. Did you forget to install `timesfm[xreg]`?" + ) + +Category = int | str + +_TOL = 1e-6 +XRegMode = Literal["timesfm + xreg", "xreg + timesfm"] + + +def _unnest(nested: Sequence[Sequence[Any]]) -> np.ndarray: + return np.array(list(itertools.chain.from_iterable(nested))) + + +def _repeat(elements: Iterable[Any], counts: Iterable[int]) -> np.ndarray: + return np.array( + list(itertools.chain.from_iterable(map(itertools.repeat, elements, counts))) + ) + + +def _to_padded_jax_array(x: np.ndarray) -> jax.Array: + if x.ndim == 1: + (i,) = x.shape + di = 2 ** math.ceil(math.log2(i)) - i + return jnp.pad(x, ((0, di),), mode="constant", constant_values=0.0) + elif x.ndim == 2: + i, j = x.shape + di = 2 ** math.ceil(math.log2(i)) - i + dj = 2 ** math.ceil(math.log2(j)) - j + return jnp.pad(x, ((0, di), (0, dj)), mode="constant", constant_values=0.0) + else: + raise ValueError(f"Unsupported array shape: {x.shape}") + + +# Per time series normalization: forward. +def normalize(batch): + stats = [(np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch] + new_batch = [(x - stat[0]) / stat[1] for x, stat in zip(batch, stats)] + return new_batch, stats + + +# Per time series normalization: inverse. +def renormalize(batch, stats): + return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)] + + +class BatchedInContextXRegBase: + """Helper class for in-context regression covariate formatting. + + Attributes: + targets: List of targets (responses) of the in-context regression. + train_lens: List of lengths of each target vector from the context. + test_lens: List of lengths of each forecast horizon. + train_dynamic_numerical_covariates: Dict of covariate names mapping to the + dynamic numerical covariates of each forecast task on the context. Their + lengths should match the corresponding lengths in `train_lens`. + train_dynamic_categorical_covariates: Dict of covariate names mapping to the + dynamic categorical covariates of each forecast task on the context. Their + lengths should match the corresponding lengths in `train_lens`. + test_dynamic_numerical_covariates: Dict of covariate names mapping to the + dynamic numerical covariates of each forecast task on the horizon. Their + lengths should match the corresponding lengths in `test_lens`. + test_dynamic_categorical_covariates: Dict of covariate names mapping to the + dynamic categorical covariates of each forecast task on the horizon. Their + lengths should match the corresponding lengths in `test_lens`. + static_numerical_covariates: Dict of covariate names mapping to the static + numerical covariates of each forecast task. + static_categorical_covariates: Dict of covariate names mapping to the static + categorical covariates of each forecast task. + """ + + def __init__( + self, + targets: Sequence[Sequence[float]], + train_lens: Sequence[int], + test_lens: Sequence[int], + train_dynamic_numerical_covariates: ( + Mapping[str, Sequence[Sequence[float]]] | None + ) = None, + train_dynamic_categorical_covariates: ( + Mapping[str, Sequence[Sequence[Category]]] | None + ) = None, + test_dynamic_numerical_covariates: ( + Mapping[str, Sequence[Sequence[float]]] | None + ) = None, + test_dynamic_categorical_covariates: ( + Mapping[str, Sequence[Sequence[Category]]] | None + ) = None, + static_numerical_covariates: Mapping[str, Sequence[float]] | None = None, + static_categorical_covariates: (Mapping[str, Sequence[Category]] | None) = None, + ) -> None: + """Initializes with the exogenous covariate inputs. + + Here we use model fitting language to refer to the context as 'train' and + the horizon as 'test'. We assume batched inputs. To properly format the + request: + + - `train_lens` represents the contexts in the batch. Targets and all train + dynamic covariates should have the same lengths as the corresponding + elements + in `train_lens`. Notice each `train_len` can be different from the exact + length of the corresponding context depending on how much of the context is + used for fitting the in-context model. + - `test_lens` represents the horizon lengths in the batch. All tesdt + dynamic + covariates should have the same lengths as the corresponding elements in + `test_lens`. + - Static covariates should be one for each input. + - For train and test dynamic covariates, they should have the same + covariate + names. + + Pass an empty dict {} for a covariate type if it is not present. + + Example: + Here is a set of valid inputs whose schema can be used for reference. + ``` + targets = [ + [0.0, 0.1, 0.2], + [0.0, 0.1, 0.2, 0.3], + ] # Two inputs in this batch. + train_lens = [3, 4] + test_lens = [2, 5] # Forecast horizons 2 and 5 respectively. + train_dynamic_numerical_covariates = { + "cov_1_dn": [[0.0, 0.5, 1.0], [0.0, 0.5, 1.0, 1.5]], + "cov_2_dn": [[0.0, 1.5, 1.0], [0.0, 1.5, 1.0, 2.5]], + } # Each train dynamic covariate has 3 and 4 elements respectively. + test_dynamic_numerical_covariates = { + "cov_1_dn": [[0.1, 0.6], [0.1, 0.6, 1.1, 1.6, 2.4]], + "cov_2_dn": [[0.1, 1.1], [0.1, 1.6, 1.1, 2.6, 10.0]], + } # Each test dynamic covariate has 2 and 5 elements respectively. + train_dynamic_categorical_covariates = { + "cov_1_dc": [[0, 1, 0], [0, 1, 2, 3]], + "cov_2_dc": [["good", "bad", "good"], ["good", "good", "bad", + "bad"]], + } + test_dynamic_categorical_covariates = { + "cov_1_dc": [[1, 0], [1, 0, 2, 3, 1]], + "cov_2_dc": [["bad", "good"], ["bad", "bad", "bad", "bad", "bad"]], + } + static_numerical_covariates = { + "cov_1_sn": [0.0, 3.0], + "cov_2_sn": [2.0, 1.0], + "cov_3_sn": [1.0, 2.0], + } # Each static covariate has 1 element for each input. + static_categorical_covariates = { + "cov_1_sc": ["apple", "orange"], + "cov_2_sc": [2, 3], + } + ``` + + Args: + targets: List of targets (responses) of the in-context regression. + train_lens: List of lengths of each target vector from the context. + test_lens: List of lengths of each forecast horizon. + train_dynamic_numerical_covariates: Dict of covariate names mapping to the + dynamic numerical covariates of each forecast task on the context. Their + lengths should match the corresponding lengths in `train_lens`. + train_dynamic_categorical_covariates: Dict of covariate names mapping to + the dynamic categorical covariates of each forecast task on the context. + Their lengths should match the corresponding lengths in `train_lens`. + test_dynamic_numerical_covariates: Dict of covariate names mapping to the + dynamic numerical covariates of each forecast task on the horizon. Their + lengths should match the corresponding lengths in `test_lens`. + test_dynamic_categorical_covariates: Dict of covariate names mapping to + the dynamic categorical covariates of each forecast task on the horizon. + Their lengths should match the corresponding lengths in `test_lens`. + static_numerical_covariates: Dict of covariate names mapping to the static + numerical covariates of each forecast task. + static_categorical_covariates: Dict of covariate names mapping to the + static categorical covariates of each forecast task. + """ + self.targets = targets + self.train_lens = train_lens + self.test_lens = test_lens + self.train_dynamic_numerical_covariates = train_dynamic_numerical_covariates or {} + self.train_dynamic_categorical_covariates = ( + train_dynamic_categorical_covariates or {} + ) + self.test_dynamic_numerical_covariates = test_dynamic_numerical_covariates or {} + self.test_dynamic_categorical_covariates = test_dynamic_categorical_covariates or {} + self.static_numerical_covariates = static_numerical_covariates or {} + self.static_categorical_covariates = static_categorical_covariates or {} + + def _assert_covariates(self, assert_covariate_shapes: bool = False) -> None: + """Verifies the validity of the covariate inputs.""" + + # Check presence. + if ( + self.train_dynamic_numerical_covariates + and not self.test_dynamic_numerical_covariates + ) or ( + not self.train_dynamic_numerical_covariates + and self.test_dynamic_numerical_covariates + ): + raise ValueError( + "train_dynamic_numerical_covariates and" + " test_dynamic_numerical_covariates must be both present or both" + " absent." + ) + + if ( + self.train_dynamic_categorical_covariates + and not self.test_dynamic_categorical_covariates + ) or ( + not self.train_dynamic_categorical_covariates + and self.test_dynamic_categorical_covariates + ): + raise ValueError( + "train_dynamic_categorical_covariates and" + " test_dynamic_categorical_covariates must be both present or both" + " absent." + ) + + # Check keys. + for dict_a, dict_b, dict_a_name, dict_b_name in ( + ( + self.train_dynamic_numerical_covariates, + self.test_dynamic_numerical_covariates, + "train_dynamic_numerical_covariates", + "test_dynamic_numerical_covariates", + ), + ( + self.train_dynamic_categorical_covariates, + self.test_dynamic_categorical_covariates, + "train_dynamic_categorical_covariates", + "test_dynamic_categorical_covariates", + ), + ): + if w := set(dict_a.keys()) - set(dict_b.keys()): + raise ValueError(f"{dict_a_name} has keys not present in {dict_b_name}: {w}") + if w := set(dict_b.keys()) - set(dict_a.keys()): + raise ValueError(f"{dict_b_name} has keys not present in {dict_a_name}: {w}") + + # Check shapes. + if assert_covariate_shapes: + if len(self.targets) != len(self.train_lens): + raise ValueError( + "targets and train_lens must have the same number of elements." + ) + + if len(self.train_lens) != len(self.test_lens): + raise ValueError( + "train_lens and test_lens must have the same number of elements." + ) + + for i, (target, train_len) in enumerate(zip(self.targets, self.train_lens)): + if len(target) != train_len: + raise ValueError( + f"targets[{i}] has length {len(target)} != expected {train_len}." + ) + + for key, values in self.static_numerical_covariates.items(): + if len(values) != len(self.train_lens): + raise ValueError( + f"static_numerical_covariates has key {key} with number of" + f" examples {len(values)} != expected {len(self.train_lens)}." + ) + + for key, values in self.static_categorical_covariates.items(): + if len(values) != len(self.train_lens): + raise ValueError( + f"static_categorical_covariates has key {key} with number of" + f" examples {len(values)} != expected {len(self.train_lens)}." + ) + + for lens, dict_cov, dict_cov_name in ( + ( + self.train_lens, + self.train_dynamic_numerical_covariates, + "train_dynamic_numerical_covariates", + ), + ( + self.train_lens, + self.train_dynamic_categorical_covariates, + "train_dynamic_categorical_covariates", + ), + ( + self.test_lens, + self.test_dynamic_numerical_covariates, + "test_dynamic_numerical_covariates", + ), + ( + self.test_lens, + self.test_dynamic_categorical_covariates, + "test_dynamic_categorical_covariates", + ), + ): + for key, cov_values in dict_cov.items(): + if len(cov_values) != len(lens): + raise ValueError( + f"{dict_cov_name} has key {key} with number of examples" + f" {len(cov_values)} != expected {len(lens)}." + ) + for i, cov_value in enumerate(cov_values): + if len(cov_value) != lens[i]: + raise ValueError( + f"{dict_cov_name} has key {key} with its {i}-th example" + f" length {len(cov_value)} != expected {lens[i]}." + ) + + def create_covariate_matrix( + self, + one_hot_encoder_drop: str | None = "first", + use_intercept: bool = True, + assert_covariates: bool = False, + assert_covariate_shapes: bool = False, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Creates target vector and covariate matrices for in context regression. + + Here we use model fitting language to refer to the context as 'train' and + the horizon as 'test'. + + Args: + one_hot_encoder_drop: Which drop strategy to use for the one hot encoder. + use_intercept: Whether to prepare an intercept (all 1) column in the + matrices. + assert_covariates: Whether to assert the validity of the covariate inputs. + assert_covariate_shapes: Whether to assert the shapes of the covariate + inputs when `assert_covariates` is True. + + Returns: + A tuple of the target vector, the covariate matrix for the context, and + the covariate matrix for the horizon. + """ + if assert_covariates: + self._assert_covariates(assert_covariate_shapes) + + x_train, x_test = [], [] + + # Numerical features. + for name in sorted(self.train_dynamic_numerical_covariates): + x_train.append( + _unnest(self.train_dynamic_numerical_covariates[name])[:, np.newaxis] + ) + x_test.append( + _unnest(self.test_dynamic_numerical_covariates[name])[:, np.newaxis] + ) + + for covs in self.static_numerical_covariates.values(): + x_train.append(_repeat(covs, self.train_lens)[:, np.newaxis]) + x_test.append(_repeat(covs, self.test_lens)[:, np.newaxis]) + + if x_train: + x_train = np.concatenate(x_train, axis=1) + x_test = np.concatenate(x_test, axis=1) + + # Normalize for robustness. + x_mean = np.mean(x_train, axis=0, keepdims=True) + x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0) + x_train = [(x_train - x_mean) / x_std] + x_test = [(x_test - x_mean) / x_std] + + # Categorical features. Encode one by one. + one_hot_encoder = preprocessing.OneHotEncoder( + drop=one_hot_encoder_drop, + sparse_output=False, + handle_unknown="ignore", + ) + for name in sorted(self.train_dynamic_categorical_covariates.keys()): + ohe_train = _unnest(self.train_dynamic_categorical_covariates[name])[ + :, np.newaxis + ] + ohe_test = _unnest(self.test_dynamic_categorical_covariates[name])[:, np.newaxis] + x_train.append(np.array(one_hot_encoder.fit_transform(ohe_train))) + x_test.append(np.array(one_hot_encoder.transform(ohe_test))) + + for covs in self.static_categorical_covariates.values(): + ohe = one_hot_encoder.fit_transform(np.array(covs)[:, np.newaxis]) + x_train.append(_repeat(ohe, self.train_lens)) + x_test.append(_repeat(ohe, self.test_lens)) + + x_train = np.concatenate(x_train, axis=1) + x_test = np.concatenate(x_test, axis=1) + + if use_intercept: + x_train = np.pad(x_train, ((0, 0), (1, 0)), constant_values=1.0) + x_test = np.pad(x_test, ((0, 0), (1, 0)), constant_values=1.0) + + return _unnest(self.targets), x_train, x_test + + def fit(self) -> Any: + raise NotImplementedError("Fit is not implemented.") + + +class BatchedInContextXRegLinear(BatchedInContextXRegBase): + """Linear in-context regression model.""" + + def fit( + self, + ridge: float = 0.0, + one_hot_encoder_drop: str | None = "first", + use_intercept: bool = True, + force_on_cpu: bool = False, + max_rows_per_col: int = 0, + max_rows_per_col_sample_seed: int = 42, + debug_info: bool = False, + assert_covariates: bool = False, + assert_covariate_shapes: bool = False, + ) -> ( + list[np.ndarray] + | tuple[list[np.ndarray], list[np.ndarray], jax.Array, jax.Array, jax.Array] + ): + """Fits a linear model for in-context regression. + + Args: + ridge: A non-negative value for specifying the ridge regression penalty. + If 0 is provided, fallback to ordinary least squares. Note this penalty + is added to the normalized covariate matrix. + one_hot_encoder_drop: Which drop strategy to use for the one hot encoder. + use_intercept: Whether to prepare an intercept (all 1) column in the + matrices. + force_on_cpu: Whether to force execution on cpu for accelerator machines. + max_rows_per_col: How many rows to subsample per column. 0 for no + subsampling. This is for speeding up model fitting. + max_rows_per_col_sample_seed: The seed for the subsampling if needed by + `max_rows_per_col`. + debug_info: Whether to return debug info. + assert_covariates: Whether to assert the validity of the covariate inputs. + assert_covariate_shapes: Whether to assert the shapes of the covariate + inputs when `assert_covariates` is True. + + Returns: + If `debug_info` is False: + The linear fits on the horizon. + If `debug_info` is True: + A tuple of: + - the linear fits on the horizon, + - the linear fits on the context, + - the flattened target vector, + - the covariate matrix for the context, and + - the covariate matrix for the horizon. + """ + flat_targets, x_train_raw, x_test = self.create_covariate_matrix( + one_hot_encoder_drop=one_hot_encoder_drop, + use_intercept=use_intercept, + assert_covariates=assert_covariates, + assert_covariate_shapes=assert_covariate_shapes, + ) + + x_train = x_train_raw.copy() + if max_rows_per_col: + nrows, ncols = x_train.shape + if nrows > (w := ncols * max_rows_per_col): + subsample = jax.random.choice( + jax.random.PRNGKey(max_rows_per_col_sample_seed), + nrows, + (w,), + replace=False, + ) + x_train = x_train[subsample] + flat_targets = flat_targets[subsample] + + device = jax.devices("cpu")[0] if force_on_cpu else None + # Runs jitted version of the solvers which are quicker at the cost of + # running jitting during the first time calling. Re-jitting happens whenever + # new (padded) shapes are encountered. + # Ocassionally it helps with the speed and the accuracy if we force single + # thread execution on cpu for accelerator machines: + # 1. Avoid moving data to accelarator memory. + # 2. Avoid precision loss if any. + with jax.default_device(device): + x_train_raw = _to_padded_jax_array(x_train_raw) + x_train = _to_padded_jax_array(x_train) + flat_targets = _to_padded_jax_array(flat_targets) + x_test = _to_padded_jax_array(x_test) + beta_hat = ( + jnp.linalg.pinv( + x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]), + hermitian=True, + ) + @ x_train.T + @ flat_targets + ) + y_hat = x_test @ beta_hat + y_hat_context = x_train_raw @ beta_hat if debug_info else None + + outputs = [] + outputs_context = [] + + # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits. + train_index, test_index = 0, 0 + for train_index_delta, test_index_delta in zip(self.train_lens, self.test_lens): + outputs.append(np.array(y_hat[test_index : (test_index + test_index_delta)])) + if debug_info: + outputs_context.append( + np.array(y_hat_context[train_index : (train_index + train_index_delta)]) + ) + train_index += train_index_delta + test_index += test_index_delta + + if debug_info: + return outputs, outputs_context, flat_targets, x_train, x_test + else: + return outputs From ba2417e11acb754bc2b1ebc630e387eb6b34e7f8 Mon Sep 17 00:00:00 2001 From: siriuz42 Date: Tue, 28 Oct 2025 22:40:21 +0000 Subject: [PATCH 4/4] update requirement --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index e9e72b1..5c61e68 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,7 +16,7 @@ idna==3.10 # via requests inquirerpy==0.3.4 # via huggingface-hub -numpy==2.3.3 +numpy==2.2.6 # via timesfm (pyproject.toml) packaging==25.0 # via huggingface-hub