diff --git a/README.md b/README.md index b252802..3d065eb 100644 --- a/README.md +++ b/README.md @@ -164,7 +164,7 @@ forecast_df = tfm.forecast_on_df( ## Finetuning -We have provided an example of finetuning the model on a new dataset in `notebooks/finetuning.ipynb`. +We have provided an example of finetuning the model on a new dataset in `notebooks/finetuning.ipynb`. ## Contribution Style guide @@ -179,4 +179,4 @@ spaces_before_comment = 2 ``` -Please run `yapf --in-place --recursive ` on all affected files. \ No newline at end of file +Please run `yapf --in-place --recursive ` on all affected files. diff --git a/environment.yml b/environment.yml index e3b0f1e..7a1d6fe 100644 --- a/environment.yml +++ b/environment.yml @@ -15,3 +15,4 @@ dependencies: - paxml - jax[cuda12]==0.4.26 - einshape + - scikit-learn diff --git a/environment_cpu.yml b/environment_cpu.yml index 772c4b7..d808176 100644 --- a/environment_cpu.yml +++ b/environment_cpu.yml @@ -15,3 +15,4 @@ dependencies: - paxml - jax[cpu]==0.4.26 - einshape + - scikit-learn diff --git a/experiments/baselines/timegpt_pipeline.py b/experiments/baselines/timegpt_pipeline.py index 5ea18ad..8a2bfbd 100644 --- a/experiments/baselines/timegpt_pipeline.py +++ b/experiments/baselines/timegpt_pipeline.py @@ -12,10 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional, Tuple import os -import pandas as pd +from time import time +from typing import List, Optional, Tuple +from dotenv import load_dotenv from gluonts.time_feature.seasonality import get_seasonality as _get_seasonality +from nixtla import NixtlaClient +import pandas as pd from tqdm import tqdm from utilsforecast.processing import ( backtest_splits, @@ -25,17 +28,15 @@ from utilsforecast.processing import ( take_rows, vertical_concat, ) -from time import time -from dotenv import load_dotenv -from nixtla import NixtlaClient def get_seasonality(freq: str) -> int: return _get_seasonality(freq, seasonalities={"D": 7}) -def maybe_convert_col_to_datetime(df: pd.DataFrame, - col_name: str) -> pd.DataFrame: +def maybe_convert_col_to_datetime( + df: pd.DataFrame, col_name: str +) -> pd.DataFrame: if not pd.api.types.is_datetime64_any_dtype(df[col_name]): df = df.copy() df[col_name] = pd.to_datetime(df[col_name]) @@ -63,15 +64,14 @@ def zero_pad_time_series(df, freq, min_length=36): end=start_date, periods=min_length - len(subset) + 1, freq=freq, # 'MS' for month start - )[:-1] # Exclude the start_date itself + )[ + :-1 + ] # Exclude the start_date itself # 2c. Create padding data - padding_df = pd.DataFrame({ - "ds": padding_dates, - "unique_id": unique_id, - "y": 0 - } # Zero padding - ) + padding_df = pd.DataFrame( + {"ds": padding_dates, "unique_id": unique_id, "y": 0} # Zero padding + ) # 2d. Combine original and padding data, and append to the list padded_data.append(pd.concat([padding_df, subset]).sort_values("ds")) @@ -83,8 +83,9 @@ def zero_pad_time_series(df, freq, min_length=36): class Forecaster: """Borrowed from - https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. - """ + + https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. + """ def forecast( self, @@ -120,7 +121,8 @@ class Forecaster: for _, (cutoffs, train, valid) in tqdm(enumerate(splits)): if len(valid.columns) > 3: raise NotImplementedError( - "Cross validation with exogenous variables is not yet supported.") + "Cross validation with exogenous variables is not yet supported." + ) y_pred = self.forecast( df=train, h=h, @@ -134,9 +136,10 @@ class Forecaster: ) if result.shape[0] < valid.shape[0]: raise ValueError( - "Cross validation result produced less results than expected. " - "Please verify that the frequency parameter (freq) matches your series' " - "and that there aren't any missing periods.") + "Cross validation result produced less results than expected." + " Please verify that the frequency parameter (freq) matches your" + " series' and that there aren't any missing periods." + ) results.append(result) out = vertical_concat(results) out = drop_index_if_pandas(out) @@ -148,9 +151,10 @@ class Forecaster: class TimeGPT(Forecaster): """Borrowed from - https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. - We modify the class to take care of edge cases. - """ + + https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. + We modify the class to take care of edge cases. + """ def __init__( self, @@ -199,7 +203,7 @@ class TimeGPT(Forecaster): all_unique_ids = df["unique_id"].unique() all_fcst_df = [] for i in range(0, len(all_unique_ids), chunk_size): - chunk_ids = all_unique_ids[i:i + chunk_size] + chunk_ids = all_unique_ids[i : i + chunk_size] chunk_df = df[df["unique_id"].isin(chunk_ids)] fct_chunk_df = client.forecast( df=chunk_df, @@ -236,11 +240,13 @@ def run_timegpt( chunk_size = 5000 else: chunk_size = None - fcsts_df = model.forecast(df=padded_train_df, - h=horizon, - level=level, - freq=freq, - chunk_size=chunk_size) + fcsts_df = model.forecast( + df=padded_train_df, + h=horizon, + level=level, + freq=freq, + chunk_size=chunk_size, + ) total_time = time() - init_time # In case levels are not returned we replace the levels with the mean predictions. # Note that this does not affect the results table as we only compare on point diff --git a/experiments/environment.yml b/experiments/environment.yml index ea7e746..32e83ee 100644 --- a/experiments/environment.yml +++ b/experiments/environment.yml @@ -25,3 +25,4 @@ dependencies: - python-dotenv - nixtla>=0.5.1 - rich + - scikit-learn diff --git a/experiments/environment_cpu.yml b/experiments/environment_cpu.yml index 7d5a0f6..3f63013 100644 --- a/experiments/environment_cpu.yml +++ b/experiments/environment_cpu.yml @@ -25,3 +25,4 @@ dependencies: - python-dotenv - nixtla>=0.5.1 - rich + - scikit-learn diff --git a/experiments/extended_benchmarks/run_timegpt.py b/experiments/extended_benchmarks/run_timegpt.py index 455782d..38d964d 100644 --- a/experiments/extended_benchmarks/run_timegpt.py +++ b/experiments/extended_benchmarks/run_timegpt.py @@ -11,6 +11,7 @@ # 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. + """Evaluation script for timegpt.""" import os @@ -20,10 +21,11 @@ import time from absl import flags import numpy as np import pandas as pd -from ..baselines.timegpt_pipeline import run_timegpt +from ..baselines.timegpt_pipeline import run_timegpt from .utils import ExperimentHandler + dataset_names = [ "m1_monthly", "m1_quarterly", @@ -61,6 +63,7 @@ _MODEL_NAME = flags.DEFINE_string( ) _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory") + QUANTILES = list(np.arange(1, 10) / 10.0) @@ -87,9 +90,9 @@ def main(): ) time_df = pd.DataFrame({"time": [total_time], "model": model_name}) fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name) - results = exp.evaluate_from_predictions(models=[model_name], - fcsts_df=fcsts_df, - times_df=time_df) + results = exp.evaluate_from_predictions( + models=[model_name], fcsts_df=fcsts_df, times_df=time_df + ) print(results, flush=True) results_list.append(results) results_full = pd.concat(results_list) diff --git a/experiments/extended_benchmarks/run_timesfm.py b/experiments/extended_benchmarks/run_timesfm.py index f01bf39..23d19d4 100644 --- a/experiments/extended_benchmarks/run_timesfm.py +++ b/experiments/extended_benchmarks/run_timesfm.py @@ -11,6 +11,7 @@ # 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. + """Evaluation script for timesfm.""" import os @@ -25,6 +26,7 @@ import timesfm from .utils import ExperimentHandler + dataset_names = [ "m1_monthly", "m1_quarterly", @@ -72,14 +74,16 @@ context_dict = { "m4_yearly": 64, } -_MODEL_PATH = flags.DEFINE_string("model_path", "/home/timesfm_q10_20240501", - "Path to model") +_MODEL_PATH = flags.DEFINE_string( + "model_path", "/home/timesfm_q10_20240501", "Path to model" +) _BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size") _HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon") _BACKEND = flags.DEFINE_string("backend", "gpu", "Backend") _NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs") _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory") + QUANTILES = list(np.arange(1, 10) / 10.0) @@ -123,9 +127,9 @@ def main(): ) total_time = time.time() - init_time time_df = pd.DataFrame({"time": [total_time], "model": model_name}) - results = exp.evaluate_from_predictions(models=[model_name], - fcsts_df=fcsts_df, - times_df=time_df) + results = exp.evaluate_from_predictions( + models=[model_name], fcsts_df=fcsts_df, times_df=time_df + ) print(results, flush=True) results_list.append(results) results_full = pd.concat(results_list) diff --git a/experiments/extended_benchmarks/utils.py b/experiments/extended_benchmarks/utils.py index c9307f4..de0368b 100644 --- a/experiments/extended_benchmarks/utils.py +++ b/experiments/extended_benchmarks/utils.py @@ -11,6 +11,7 @@ # 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. + """Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py.""" from functools import partial @@ -45,9 +46,11 @@ def quantile_loss( target_col: str = "y", ) -> pd.DataFrame: delta_y = df[models].sub(df[target_col], axis=0) - res = (np.maximum(q * delta_y, - (q - 1) * delta_y).groupby(df[id_col], - observed=True).mean()) + res = ( + np.maximum(q * delta_y, (q - 1) * delta_y) + .groupby(df[id_col], observed=True) + .mean() + ) res.index.name = id_col res = res.reset_index() return res @@ -63,8 +66,10 @@ class ExperimentHandler: models_dir: str = "./models", ): if dataset not in gluonts_datasets: - raise Exception(f"dataset {dataset} not found in gluonts " - f"available datasets: {', '.join(gluonts_datasets)}") + raise Exception( + f"dataset {dataset} not found in gluonts " + f"available datasets: {', '.join(gluonts_datasets)}" + ) self.dataset = dataset self.quantiles = quantiles self.level = self._transform_quantiles_to_levels(quantiles) @@ -75,8 +80,10 @@ class ExperimentHandler: gluonts_dataset = get_dataset(self.dataset) self.horizon = gluonts_dataset.metadata.prediction_length if self.horizon is None: - raise Exception(f"horizon not found for dataset {self.dataset} " - "experiment cannot be run") + raise Exception( + f"horizon not found for dataset {self.dataset} " + "experiment cannot be run" + ) self.freq = gluonts_dataset.metadata.freq # get_seasonality() returns 1 for freq='D', override this to 7. This significantly improves the accuracy of # statistical models on datasets like m5/nn5_daily. The models like AutoARIMA/AutoETS can still set @@ -115,8 +122,9 @@ class ExperimentHandler: @staticmethod def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]: - level = [int(100 - 200 * q) for q in quantiles if q < 0.5 - ] # in this case mean=mediain + level = [ + int(100 - 200 * q) for q in quantiles if q < 0.5 + ] # in this case mean=mediain level = sorted(list(set(level))) return level @@ -145,8 +153,9 @@ class ExperimentHandler: last_n: int | None = None, ) -> pd.DataFrame: with multiprocessing.Pool(os.cpu_count()) as pool: # Create a process pool - results = pool.map(parallel_transform, zip(gluonts_dataset, - repeat(last_n))) + results = pool.map( + parallel_transform, zip(gluonts_dataset, repeat(last_n)) + ) df = pd.concat(results) df = df.reset_index(drop=True) return df @@ -168,8 +177,9 @@ class ExperimentHandler: def save_dataframe(self, df: pd.DataFrame, file_name: str): df.to_csv(f"{self.results_dir}/{file_name}", index=False) - def save_results(self, fcst_df: pd.DataFrame, total_time: float, - model_name: str): + def save_results( + self, fcst_df: pd.DataFrame, total_time: float, model_name: str + ): self.save_dataframe( fcst_df, f"{model_name}-{self.dataset}-fcst.csv", @@ -205,21 +215,23 @@ class ExperimentHandler: times_df = [] for model in models: fcst_method_df = pd.read_csv( - f"{self.results_dir}/{model}-{self.dataset}-fcst.csv").set_index( - ["unique_id", "ds"]) + f"{self.results_dir}/{model}-{self.dataset}-fcst.csv" + ).set_index(["unique_id", "ds"]) fcsts_df.append(fcst_method_df) time_method_df = pd.read_csv( - f"{self.results_dir}/{model}-{self.dataset}-time.csv") + f"{self.results_dir}/{model}-{self.dataset}-time.csv" + ) times_df.append(time_method_df) fcsts_df = pd.concat(fcsts_df, axis=1).reset_index() fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"]) times_df = pd.concat(times_df) - return self.evaluate_from_predictions(models=models, - fcsts_df=fcsts_df, - times_df=times_df) + return self.evaluate_from_predictions( + models=models, fcsts_df=fcsts_df, times_df=times_df + ) - def evaluate_from_predictions(self, models: List[str], fcsts_df: pd.DataFrame, - times_df: pd.DataFrame) -> pd.DataFrame: + def evaluate_from_predictions( + self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame + ) -> pd.DataFrame: test_df = self.test_df train_df = self.train_df test_df = test_df.merge(fcsts_df, how="left") @@ -250,9 +262,9 @@ class ExperimentHandler: eval_prob_df["metric"] = "scaled_crps" eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True) eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index() - eval_df = eval_df.melt(id_vars="metric", - value_name="value", - var_name="model") + eval_df = eval_df.melt( + id_vars="metric", value_name="value", var_name="model" + ) times_df.insert(0, "metric", "time") times_df = times_df.rename(columns={"time": "value"}) eval_df = pd.concat([eval_df, times_df]) diff --git a/experiments/long_horizon_benchmarks/run_eval.py b/experiments/long_horizon_benchmarks/run_eval.py index 981f11f..12a4846 100644 --- a/experiments/long_horizon_benchmarks/run_eval.py +++ b/experiments/long_horizon_benchmarks/run_eval.py @@ -11,6 +11,7 @@ # 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. + """Eval pipeline.""" import json @@ -23,32 +24,42 @@ import numpy as np import pandas as pd from paxml import checkpoints import timesfm +from timesfm import data_loader import torch import tqdm -from timesfm import data_loader + FLAGS = flags.FLAGS -_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, - "Batch size for the randomly sampled batch") +_BATCH_SIZE = flags.DEFINE_integer( + "batch_size", 64, "Batch size for the randomly sampled batch" +) _DATASET = flags.DEFINE_string("dataset", "etth1", "The name of the dataset.") -_MODEL_PATH = flags.DEFINE_string("model_path", "./timesfm_q10_20240501", - "The name of the dataset.") -_DATETIME_COL = flags.DEFINE_string("datetime_col", "date", - "Column having datetime.") -_NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None, - "Column having numerical features.") -_CAT_COV_COLS = flags.DEFINE_list("cat_cov_cols", None, - "Column having categorical features.") +_MODEL_PATH = flags.DEFINE_string( + "model_path", "./timesfm_q10_20240501", "The name of the dataset." +) +_DATETIME_COL = flags.DEFINE_string( + "datetime_col", "date", "Column having datetime." +) +_NUM_COV_COLS = flags.DEFINE_list( + "num_cov_cols", None, "Column having numerical features." +) +_CAT_COV_COLS = flags.DEFINE_list( + "cat_cov_cols", None, "Column having categorical features." +) _TS_COLS = flags.DEFINE_list("ts_cols", None, "Columns of time-series features") -_NORMALIZE = flags.DEFINE_bool("normalize", True, - "normalize data for eval or not") -_CONTEXT_LEN = flags.DEFINE_integer("context_len", 512, - "Length of the context window") +_NORMALIZE = flags.DEFINE_bool( + "normalize", True, "normalize data for eval or not" +) +_CONTEXT_LEN = flags.DEFINE_integer( + "context_len", 512, "Length of the context window" +) _PRED_LEN = flags.DEFINE_integer("pred_len", 96, "prediction length.") _BACKEND = flags.DEFINE_string("backend", "gpu", "backend to use") -_RESULTS_DIR = flags.DEFINE_string("results_dir", "./results/long_horizon", - "results directory") +_RESULTS_DIR = flags.DEFINE_string( + "results_dir", "./results/long_horizon", "results directory" +) + DATA_DICT = { "ettm2": { @@ -165,8 +176,9 @@ def eval(): holiday=False, permute=False, ) - eval_itr = dtl.tf_dataset(mode="test", - shift=_PRED_LEN.value).as_numpy_iterator() + eval_itr = dtl.tf_dataset( + mode="test", shift=_PRED_LEN.value + ).as_numpy_iterator() model_path = _MODEL_PATH.value if model_path.startswith("amazon"): model = chronos.ChronosPipeline.from_pretrained( @@ -201,9 +213,10 @@ def eval(): for batch in tqdm.tqdm(eval_itr): past = batch[0] actuals = batch[3] - forecasts = get_forecasts(model_path, model, past, int_freq, - _PRED_LEN.value) - forecasts = forecasts[:, 0:actuals.shape[1]] + forecasts = get_forecasts( + model_path, model, past, int_freq, _PRED_LEN.value + ) + forecasts = forecasts[:, 0 : actuals.shape[1]] mae_run_losses.append(_mae(forecasts, actuals).sum()) mse_run_losses.append(_mse(forecasts, actuals).sum()) smape_run_losses.append(_smape(forecasts, actuals).sum()) diff --git a/pyproject.toml b/pyproject.toml index 2a34f2e..4aa1d4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ [project] name = "timesfm" description = "Open weights time-series foundation model from Google Research." -version = "0.0.1" +version = "1.0.0" dependencies = [ "einshape>=1.0.0", "paxml>=1.4.0", @@ -11,6 +11,7 @@ dependencies = [ "jax>=0.4.26", "numpy>=1.26.4", "pandas>=2.1.4", + "sklearn>=1.15.1", ] authors = [ {name = "Rajat Sen", email = "senrajat@google.com"}, diff --git a/src/timesfm/__init__.py b/src/timesfm/__init__.py index e794e45..866e848 100644 --- a/src/timesfm/__init__.py +++ b/src/timesfm/__init__.py @@ -11,6 +11,7 @@ # 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. + """TimesFM init file.""" from .timesfm import TimesFm, freq_map diff --git a/src/timesfm/data_loader.py b/src/timesfm/data_loader.py index b524bd4..eeace7b 100644 --- a/src/timesfm/data_loader.py +++ b/src/timesfm/data_loader.py @@ -11,11 +11,13 @@ # 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. + """TF dataloaders for general timeseries datasets. The expected input format is csv file with a datetime index. """ + from absl import logging import numpy as np import pandas as pd @@ -77,8 +79,9 @@ class TimeSeriesdata(object): self.data_df['ccol'] = np.zeros(self.data_df.shape[0]) cat_cov_cols = ['ccol'] self.data_df.fillna(0, inplace=True) - self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]), - inplace=True) + self.data_df.set_index( + pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True + ) self.num_cov_cols = num_cov_cols self.cat_cov_cols = cat_cov_cols self.ts_cols = ts_cols @@ -91,16 +94,18 @@ class TimeSeriesdata(object): data_df_idx[-1] + pd.Timedelta(1, freq=freq), periods=pred_len + 1, freq=freq, - )) + ) + ) self.time_df = time_features.TimeCovariates( - date_index, holiday=holiday).get_covariates() + date_index, holiday=holiday + ).get_covariates() self.hist_len = hist_len self.pred_len = pred_len self.batch_size = batch_size self.freq = freq self.normalize = normalize self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose() - self.data_mat = self.data_mat[:, 0:self.test_range[1]] + self.data_mat = self.data_mat[:, 0 : self.test_range[1]] self.time_mat = self.time_df.to_numpy().transpose() self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose() self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols) @@ -130,7 +135,7 @@ class TimeSeriesdata(object): def _normalize_data(self): self.scaler = StandardScaler() - train_mat = self.data_mat[:, self.train_range[0]:self.train_range[1]] + train_mat = self.data_mat[:, self.train_range[0] : self.train_range[1]] self.scaler = self.scaler.fit(train_mat.transpose()) self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose() @@ -248,8 +253,9 @@ class TimeSeriesdata(object): gen_fn = self.train_gen else: gen_fn = lambda: self.test_val_gen(mode, shift) - output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + - [tf.int32] * 2) + output_types = tuple( + [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2 + ) dataset = tf.data.Dataset.from_generator(gen_fn, output_types) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) return dataset diff --git a/src/timesfm/patched_decoder.py b/src/timesfm/patched_decoder.py index 4e2600a..c7c3758 100644 --- a/src/timesfm/patched_decoder.py +++ b/src/timesfm/patched_decoder.py @@ -11,6 +11,7 @@ # 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. + """Pax ML model for patched time-series decoder. The file implements Residual MLPs, Patched Decoder layers and PAX ML models. @@ -35,6 +36,7 @@ from praxis.layers import normalizations from praxis.layers import stochastics from praxis.layers import transformers + # PAX shortcuts NestedMap = py_utils.NestedMap JTensor = pytypes.JTensor @@ -42,6 +44,7 @@ JTensor = pytypes.JTensor LayerTpl = pax_fiddle.Config[base_layer.BaseLayer] template_field = base_layer.template_field + PAD_VAL = 1123581321.0 DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] @@ -54,6 +57,7 @@ _FREQ = "freq" _OUTPUT_TOKENS = "output_tokens" _STATS = "stats" + # Small numerical value. _TOLERANCE = 1e-7 @@ -84,16 +88,16 @@ def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor: class ResidualBlock(base_layer.BaseLayer): """Simple feedforward block with residual connection. - Attributes: - input_dims: input dimension. - hidden_dims: hidden dimension. - output_dims: output dimension. - dropout_prob: dropout probability. - layer_norm: whether to use layer norm or not. - dropout_tpl: config for dropout. - ln_tpl: config for layer norm. - act_tpl: config for activation in hidden layer. - """ + Attributes: + input_dims: input dimension. + hidden_dims: hidden dimension. + output_dims: output dimension. + dropout_prob: dropout probability. + layer_norm: whether to use layer norm or not. + dropout_tpl: config for dropout. + ln_tpl: config for layer norm. + act_tpl: config for activation in hidden layer. + """ input_dims: int = 0 hidden_dims: int = 0 @@ -154,20 +158,21 @@ class ResidualBlock(base_layer.BaseLayer): return output + residual -def _masked_mean_std(inputs: JTensor, - padding: JTensor) -> Tuple[JTensor, JTensor]: +def _masked_mean_std( + inputs: JTensor, padding: JTensor +) -> Tuple[JTensor, JTensor]: """Calculates mean and standard deviation of arr across axis 1. - It should exclude values where pad is 1. + It should exclude values where pad is 1. - Args: - inputs: A JAX array of shape [b, n, p]. - padding: A JAX array of shape [b, n, p] with values 0 or 1. + Args: + inputs: A JAX array of shape [b, n, p]. + padding: A JAX array of shape [b, n, p] with values 0 or 1. - Returns: - A tuple containing the mean and standard deviation of arr. We return the - statistics of the first patch with more than three non-padded values. - """ + Returns: + A tuple containing the mean and standard deviation of arr. We return the + statistics of the first patch with more than three non-padded values. + """ # Selecting the first pad with more than 3 unpadded values. pad_sum = jnp.sum(1 - padding, axis=2) @@ -192,7 +197,7 @@ def _masked_mean_std(inputs: JTensor, # Calculate the masked sum and squared sum of M masked_sum = jnp.sum(arr * mask, axis=1) - masked_squared_sum = jnp.sum((arr * mask)**2, axis=1) + masked_squared_sum = jnp.sum((arr * mask) ** 2, axis=1) # Calculate the masked mean and standard deviation masked_mean = masked_sum / num_valid_elements @@ -211,22 +216,22 @@ def _create_quantiles() -> list[float]: class PatchedTimeSeriesDecoder(base_layer.BaseLayer): """Patch decoder layer for time-series foundation model. - Attributes: - patch_len: length of input patches. - horizon_len: length of output patches. Referred to as `output_patch_len` - during inference. - model_dims: model dimension of stacked transformer layer. - hidden_dims: hidden dimensions in fully connected layers. - quantiles: list of quantiles for non prob model. - residual_block_tpl: config for residual block. - stacked_transformer_params_tpl: config for stacked transformer. - use_freq: whether to use frequency encoding. + Attributes: + patch_len: length of input patches. + horizon_len: length of output patches. Referred to as `output_patch_len` + during inference. + model_dims: model dimension of stacked transformer layer. + hidden_dims: hidden dimensions in fully connected layers. + quantiles: list of quantiles for non prob model. + residual_block_tpl: config for residual block. + stacked_transformer_params_tpl: config for stacked transformer. + use_freq: whether to use frequency encoding. - In all of what followed, except specified otherwise, B is batch size, T is - sequence length of time-series. N is the number of input patches that can be - obtained from T. P is the input patch length and H is the horizon length. Q is - number of output logits. D is model dimension. - """ + In all of what followed, except specified otherwise, B is batch size, T is + sequence length of time-series. N is the number of input patches that can be + obtained from T. P is the input patch length and H is the horizon length. Q is + number of output logits. D is model dimension. + """ patch_len: int = 0 horizon_len: int = 0 @@ -235,7 +240,8 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles) residual_block_tpl: LayerTpl = template_field(ResidualBlock) stacked_transformer_params_tpl: LayerTpl = template_field( - transformers.StackedTransformer) + transformers.StackedTransformer + ) use_freq: bool = True def setup(self) -> None: @@ -270,8 +276,9 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): self.create_child( "position_emb", - pax_fiddle.Config(layers.PositionalEmbedding, - embedding_dims=self.model_dims), + pax_fiddle.Config( + layers.PositionalEmbedding, embedding_dims=self.model_dims + ), ) if self.use_freq: @@ -285,24 +292,27 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): ) def transform_decode_state( - self, transform_fn: base_layer.DecodeStateTransformFn) -> None: + self, transform_fn: base_layer.DecodeStateTransformFn + ) -> None: """Transforms all decode state variables based on transform_fn.""" self.stacked_transformer_layer.transform_decode_state(transform_fn) def _forward_transform( - self, inputs: JTensor, - patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]: + self, inputs: JTensor, patched_pads: JTensor + ) -> Tuple[JTensor, Tuple[JTensor, JTensor]]: """Input is of shape [B, N, P].""" mu, sigma = _masked_mean_std(inputs, patched_pads) sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma) # Normalize each patch. outputs = (inputs - mu[:, None, None]) / sigma[:, None, None] outputs = jnp.where( - jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs) + jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs + ) return outputs, (mu, sigma) - def _reverse_transform(self, outputs: JTensor, - stats: Tuple[JTensor, JTensor]) -> JTensor: + def _reverse_transform( + self, outputs: JTensor, stats: Tuple[JTensor, JTensor] + ) -> JTensor: """Output is of shape [B, N, P, Q].""" mu, sigma = stats return outputs * sigma[:, None, None, None] + mu[:, None, None, None] @@ -316,13 +326,19 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): """Preprocess input for stacked transformer.""" # Reshape into patches. patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len) - input_padding = jnp.where( - jnp.abs(input_ts - PAD_VAL) < _TOLERANCE, 1, input_padding) - patched_pads = es.jax_einshape("b(np)->bnp", - input_padding, - p=self.patch_len) - patched_inputs, stats = self._forward_transform(patched_inputs, - patched_pads) + patched_pads = es.jax_einshape( + "b(np)->bnp", input_padding, p=self.patch_len + ) + patched_inputs = jnp.where( + jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs + ) + patched_pads = jnp.where( + jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads + ) + patched_inputs, stats = self._forward_transform( + patched_inputs, patched_pads + ) + # B x N x D patched_inputs = patched_inputs * (1.0 - patched_pads) concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1) @@ -351,26 +367,25 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): """Postprocess output of stacked transformer.""" # B x N x (H.Q) output_ts = self.horizon_ff_layer(model_output) - output_ts = es.jax_einshape("bn(hq)->bnhq", - output_ts, - q=num_outputs, - h=self.horizon_len) + output_ts = es.jax_einshape( + "bn(hq)->bnhq", output_ts, q=num_outputs, h=self.horizon_len + ) return self._reverse_transform(output_ts, stats) def __call__(self, inputs: NestedMap) -> NestedMap: """PatchTST call. - Args: - inputs: A NestedMap containing (1) input_ts: input sequence of shape [B, - T] where T must be multiple of patch_length; (2) input_padding: that - contains padding map. + Args: + inputs: A NestedMap containing (1) input_ts: input sequence of shape [B, + T] where T must be multiple of patch_length; (2) input_padding: that + contains padding map. - Returns: - A nested map with two keys: - (1) 'output_tokens' of shape [B, N, D]. - (2) 'output_ts' of shape [B, N, H, Q] - (3) 'stats' a Tuple of statistics for renormalization. - """ + Returns: + A nested map with two keys: + (1) 'output_tokens' of shape [B, N, D]. + (2) 'output_ts' of shape [B, N, H, Q] + (3) 'stats' a Tuple of statistics for renormalization. + """ input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING] num_outputs = len(self.quantiles) + 1 model_input, patched_padding, stats, _ = self._preprocess_input( @@ -385,11 +400,9 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): model_output = self.stacked_transformer_layer(model_input, patched_padding) output_ts = self._postprocess_output(model_output, num_outputs, stats) - return NestedMap({ - _OUTPUT_TOKENS: model_output, - _OUTPUT_TS: output_ts, - _STATS: stats - }) + return NestedMap( + {_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats} + ) def decode( self, @@ -397,25 +410,30 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): horizon_len: int, output_patch_len: Optional[int] = None, max_len: int = 512, + return_forecast_on_context: bool = False, ) -> tuple[JTensor, JTensor]: """Auto-regressive decoding without caching. - Args: - inputs: input time-series and paddings. Time-series shape B x C, padding - shape shape B x (C + H) where H is the prediction length. - horizon_len: prediction length. - output_patch_len: output length to be fetched from one step of - auto-regressive decoding. - max_len: maximum training context length. + Args: + inputs: input time-series and paddings. Time-series shape B x C, padding + shape shape B x (C + H) where H is the prediction length. + horizon_len: prediction length. + output_patch_len: output length to be fetched from one step of + auto-regressive decoding. + max_len: maximum training context length. + return_forecast_on_context: whether to return the model forecast on the + context except the first input patch. - Returns: - Tuple of two forecasting results: - - Point (mean) output predictions as a tensor with shape B x H. - - Full predictions (mean and quantiles) as a tensor with shape - B x H x (1 + # quantiles). - """ + Returns: + Tuple of two forecasting results: + - Point (mean) output predictions as a tensor with shape B x H'. + - Full predictions (mean and quantiles) as a tensor with shape + B x H' x (1 + # quantiles). + In particular, if return_forecast_on_context is True, H' is H plus + the forecastable context length, i.e. context_len - (first) patch_len. + """ final_out = inputs[_INPUT_TS] - inp_time_len = final_out.shape[1] + context_len = final_out.shape[1] paddings = inputs[_INPUT_PADDING] if self.use_freq: freq = inputs[_FREQ].astype(jnp.int32) @@ -425,13 +443,15 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): if paddings.shape[1] != final_out.shape[1] + horizon_len: raise ValueError( "Length of paddings must match length of input + horizon_len:" - f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}") + f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}" + ) if output_patch_len is None: output_patch_len = self.horizon_len - num_decode_patches = (horizon_len + output_patch_len - - 1) // output_patch_len - for _ in range(num_decode_patches): - current_padding = paddings[:, 0:final_out.shape[1]] + num_decode_patches = ( + horizon_len + output_patch_len - 1 + ) // output_patch_len + for step_index in range(num_decode_patches): + current_padding = paddings[:, 0 : final_out.shape[1]] input_ts = final_out[:, -max_len:] input_padding = current_padding[:, -max_len:] model_input = NestedMap( @@ -440,25 +460,40 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): freq=freq, ) fprop_outputs = self(model_input)[_OUTPUT_TS] + if return_forecast_on_context and step_index == 0: + # For the first decodings step, collect the model forecast on the + # context except the unavailable first input batch forecast. + new_full_ts = fprop_outputs[:, :-1, : self.patch_len, :] + new_full_ts = es.jax_einshape("bnph->b(np)h", new_full_ts) + + full_outputs.append(new_full_ts) + # (full batch, last patch, output_patch_len, index of mean forecast = 0) new_ts = fprop_outputs[:, -1, :output_patch_len, 0] + new_full_ts = fprop_outputs[:, -1, :output_patch_len, :] # (full batch, last patch, output_patch_len, all output indices) - full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :]) + full_outputs.append(new_full_ts) final_out = jnp.concatenate([final_out, new_ts], axis=-1) - return ( - final_out[:, inp_time_len:inp_time_len + horizon_len], - jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :], - ) + if return_forecast_on_context: + # `full_outputs` indexing starts at after the first input patch. + full_outputs = jnp.concatenate(full_outputs, axis=1)[ + :, : (context_len - self.patch_len + horizon_len), : + ] + else: + # `full_outputs` indexing starts at the forecast horizon. + full_outputs = jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :] + + return (full_outputs[:, :, 0], full_outputs) class PatchedDecoderFinetuneModel(base_model.BaseModel): """Model class for finetuning patched time-series decoder. - Attributes: - core_layer_tpl: config for core layer. - freq: freq to finetune on. - """ + Attributes: + core_layer_tpl: config for core layer. + freq: freq to finetune on. + """ core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder) freq: int = 0 @@ -471,12 +506,14 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel): input_padding = jnp.zeros_like(input_ts) context_len = input_ts.shape[1] input_patch_len = self.core_layer_tpl.patch_len - context_pad = ((context_len + input_patch_len - 1) // - input_patch_len) * input_patch_len - context_len + context_pad = ( + (context_len + input_patch_len - 1) // input_patch_len + ) * input_patch_len - context_len input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)]) - input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)], - constant_values=1) + input_padding = jnp.pad( + input_padding, [(0, 0), (context_pad, 0)], constant_values=1 + ) freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq new_input_batch = NestedMap( input_ts=input_ts, @@ -485,28 +522,30 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel): ) return self.core_layer(new_input_batch) - def _quantile_loss(self, pred: JTensor, actual: JTensor, - quantile: float) -> JTensor: + def _quantile_loss( + self, pred: JTensor, actual: JTensor, quantile: float + ) -> JTensor: """Calculates quantile loss. - Args: - pred: B x T - actual: B x T - quantile: quantile at which loss is computed. + Args: + pred: B x T + actual: B x T + quantile: quantile at which loss is computed. - Returns: - per coordinate loss. - """ + Returns: + per coordinate loss. + """ dev = actual - pred loss_first = dev * quantile loss_second = -dev * (1.0 - quantile) return 2 * jnp.where(loss_first >= 0, loss_first, loss_second) - def compute_loss(self, prediction_output: NestedMap, - input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]: + def compute_loss( + self, prediction_output: NestedMap, input_batch: NestedMap + ) -> Tuple[NestedMap, NestedMap]: output_ts = prediction_output[_OUTPUT_TS] actual_ts = input_batch[_TARGET_FUTURE] - pred_ts = output_ts[:, -1, 0:actual_ts.shape[1], :] + pred_ts = output_ts[:, -1, 0 : actual_ts.shape[1], :] loss = jnp.square(pred_ts[:, :, 0] - actual_ts) for i, quantile in enumerate(self.core_layer.quantiles): loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile) diff --git a/src/timesfm/time_features.py b/src/timesfm/time_features.py index 0cc0861..0bd90a9 100644 --- a/src/timesfm/time_features.py +++ b/src/timesfm/time_features.py @@ -11,6 +11,7 @@ # 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. + """Directory to extract time covariates. Extract time covariates from datetime. @@ -35,6 +36,7 @@ from pandas.tseries.offsets import Easter from sklearn.preprocessing import StandardScaler from tqdm import tqdm + # This is 183 to cover half a year (in both directions), also for leap years # + 17 as Eastern can be between March, 22 - April, 25 MAX_WINDOW = 183 + 17 @@ -48,8 +50,9 @@ def _distance_to_holiday(holiday): index - pd.Timedelta(days=MAX_WINDOW), index + pd.Timedelta(days=MAX_WINDOW), ) - assert (len(holiday_date) != 0 # pylint: disable=g-explicit-length-test - ), f"No closest holiday for the date index {index} found." + assert ( + len(holiday_date) != 0 # pylint: disable=g-explicit-length-test + ), f"No closest holiday for the date index {index} found." # It sometimes returns two dates if it is exactly half a year after the # holiday. In this case, the smaller distance (182 days) is returned. return (index - holiday_date[0]).days @@ -57,19 +60,16 @@ def _distance_to_holiday(holiday): return _distance_to_day -EasterSunday = Holiday("Easter Sunday", - month=1, - day=1, - offset=[Easter(), Day(0)]) +EasterSunday = Holiday( + "Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)] +) NewYearsDay = Holiday("New Years Day", month=1, day=1) -SuperBowl = Holiday("Superbowl", - month=2, - day=1, - offset=DateOffset(weekday=SU(1))) -MothersDay = Holiday("Mothers Day", - month=5, - day=1, - offset=DateOffset(weekday=SU(2))) +SuperBowl = Holiday( + "Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1)) +) +MothersDay = Holiday( + "Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2)) +) IndependenceDay = Holiday("Independence Day", month=7, day=4) ChristmasEve = Holiday("Christmas", month=12, day=24) ChristmasDay = Holiday("Christmas", month=12, day=25) diff --git a/src/timesfm/timesfm.py b/src/timesfm/timesfm.py index 70c5292..ec63a20 100644 --- a/src/timesfm/timesfm.py +++ b/src/timesfm/timesfm.py @@ -11,8 +11,10 @@ # 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. + """TimesFM forecast API for inference.""" +import collections import logging import multiprocessing from os import path @@ -20,11 +22,11 @@ import time from typing import Any, Literal, Optional, Sequence import einshape as es +from huggingface_hub import snapshot_download import jax import jax.numpy as jnp import numpy as np import pandas as pd -from huggingface_hub import snapshot_download from paxml import checkpoints from paxml import tasks_lib from praxis import base_hyperparams @@ -34,12 +36,19 @@ from praxis import py_utils from praxis import pytypes from praxis.layers import normalizations from praxis.layers import transformers -from . import patched_decoder + from utilsforecast.processing import make_future_dataframe +from . import patched_decoder +from . import xreg_lib + instantiate = base_hyperparams.instantiate NestedMap = py_utils.NestedMap JTensor = pytypes.JTensor +Category = xreg_lib.Category +XRegMode = xreg_lib.XRegMode + +_TOL = 1e-6 def process_group(key, group, value_name, forecast_context_len): @@ -51,16 +60,23 @@ def moving_average(arr, window_size): """Calculates the moving average using NumPy's convolution function.""" # Pad with zeros to handle initial window positions arr_padded = np.pad(arr, (window_size - 1, 0), "constant") - smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / - window_size) + smoothed_arr = ( + np.convolve(arr_padded, np.ones(window_size), "valid") / window_size + ) return [smoothed_arr, arr - smoothed_arr] def freq_map(freq: str): """Returns the frequency map for the given frequency string.""" freq = str.upper(freq) - if (freq.endswith("H") or freq.endswith("T") or freq.endswith("MIN") or - freq.endswith("D") or freq.endswith("B") or freq.endswith("U")): + if ( + freq.endswith("H") + or freq.endswith("T") + or freq.endswith("MIN") + or freq.endswith("D") + or freq.endswith("B") + or freq.endswith("U") + ): return 0 elif freq.endswith(("W", "M", "MS")): return 1 @@ -70,6 +86,20 @@ def freq_map(freq: str): raise ValueError(f"Invalid frequency: {freq}") +# 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 TimesFm: """TimesFM forecast API for inference. @@ -149,6 +179,7 @@ class TimesFm: self.horizon_len = horizon_len self.input_patch_len = input_patch_len self.output_patch_len = output_patch_len + self._horizon_start = self.context_len - self.input_patch_len self.mesh_shape = [1, self.num_devices, 1] self.mesh_name = ["replica", "data", "mdl"] @@ -171,7 +202,9 @@ class TimesFm: num_layers=num_layers, transformer_layer_params_tpl=pax_fiddle.Config( transformers.Transformer, - ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,), + ln_tpl=pax_fiddle.Config( + normalizations.RmsNorm, + ), ), ), ) @@ -189,38 +222,34 @@ class TimesFm: def _get_sample_inputs(self): return { - "input_ts": - jnp.zeros( - ( - self.per_core_batch_size, - self.context_len + self.output_patch_len, - ), - dtype=jnp.float32, + "input_ts": jnp.zeros( + ( + self.per_core_batch_size, + self.context_len + self.output_patch_len, ), - "input_padding": - jnp.zeros( - ( - self.per_core_batch_size, - self.context_len + self.output_patch_len, - ), - dtype=jnp.float32, + dtype=jnp.float32, + ), + "input_padding": jnp.zeros( + ( + self.per_core_batch_size, + self.context_len + self.output_patch_len, ), - "freq": - jnp.zeros( - ( - self.per_core_batch_size, - 1, - ), - dtype=jnp.int32, + dtype=jnp.float32, + ), + "freq": jnp.zeros( + ( + self.per_core_batch_size, + 1, ), + dtype=jnp.int32, + ), } def load_from_checkpoint( self, checkpoint_path: Optional[str] = None, repo_id: str = "google/timesfm-1.0-200m", - checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType. - FLAX, + checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX, step: int | None = None, ) -> None: """Loads a checkpoint and compiles the decoder. @@ -240,7 +269,8 @@ class TimesFm: start_time = time.time() self._model = instantiate(self.model_p) var_weight_hparams = self._model.abstract_init_with_metadata( - self._get_sample_inputs(), do_eval=True) + self._get_sample_inputs(), do_eval=True + ) train_state_partition_specs = tasks_lib.create_state_partition_specs( var_weight_hparams, mesh_shape=self.mesh_shape, @@ -254,7 +284,8 @@ class TimesFm: learners=None, ) self._logging( - f"Constructed model weights in {time.time() - start_time:.2f} seconds.") + f"Constructed model weights in {time.time() - start_time:.2f} seconds." + ) # Load the model weights. self._logging(f"Restoring checkpoint from {checkpoint_path}.") @@ -267,7 +298,8 @@ class TimesFm: step=step, ) self._logging( - f"Restored checkpoint in {time.time() - start_time:.2f} seconds.") + f"Restored checkpoint in {time.time() - start_time:.2f} seconds." + ) self.jit_decode() def jit_decode(self): @@ -283,6 +315,7 @@ class TimesFm: horizon_len=self.horizon_len, output_patch_len=self.output_patch_len, max_len=self.context_len, + return_forecast_on_context=True, rngs={ base_layer.PARAMS: self._key1, base_layer.RANDOM: self._key2, @@ -302,36 +335,34 @@ class TimesFm: with base_layer.JaxContext.new_context(hparams=self._eval_context): _ = self._pmapped_decode( NestedMap({ - "input_ts": - jnp.zeros( - ( - self.num_devices, - self.per_core_batch_size, - self.context_len, - ), - dtype=jnp.float32, + "input_ts": jnp.zeros( + ( + self.num_devices, + self.per_core_batch_size, + self.context_len, ), - "input_padding": - jnp.zeros( - ( - self.num_devices, - self.per_core_batch_size, - self.context_len + self.horizon_len, - ), - dtype=jnp.float32, + dtype=jnp.float32, + ), + "input_padding": jnp.zeros( + ( + self.num_devices, + self.per_core_batch_size, + self.context_len + self.horizon_len, ), - "date_features": - None, - "freq": - jnp.zeros( - (self.num_devices, self.per_core_batch_size, 1), - dtype=jnp.int32, - ), - })) + dtype=jnp.float32, + ), + "date_features": None, + "freq": jnp.zeros( + (self.num_devices, self.per_core_batch_size, 1), + dtype=jnp.int32, + ), + }) + ) self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.") - def _preprocess(self, inputs: Sequence[np.array], - freq: Sequence[int]) -> tuple[np.array, np.array, int]: + def _preprocess( + self, inputs: Sequence[np.array], freq: Sequence[int] + ) -> tuple[np.array, np.array, int]: """Formats and pads raw inputs to feed into the model. This function both pads each time series to match the context length, and @@ -352,21 +383,24 @@ class TimesFm: input_ts, input_padding, inp_freq = [], [], [] - pmap_pad = ((len(inputs) - 1) // self.global_batch_size + - 1) * self.global_batch_size - len(inputs) + pmap_pad = ( + (len(inputs) - 1) // self.global_batch_size + 1 + ) * self.global_batch_size - len(inputs) for i, ts in enumerate(inputs): input_len = ts.shape[0] padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float) if input_len < self.context_len: num_front_pad = self.context_len - input_len - ts = np.concatenate([np.zeros(shape=(num_front_pad,), dtype=float), ts], - axis=0) + ts = np.concatenate( + [np.zeros(shape=(num_front_pad,), dtype=float), ts], axis=0 + ) padding = np.concatenate( - [np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0) + [np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0 + ) elif input_len > self.context_len: - ts = ts[-self.context_len:] - padding = padding[-(self.context_len + self.horizon_len):] + ts = ts[-self.context_len :] + padding = padding[-(self.context_len + self.horizon_len) :] input_ts.append(ts) input_padding.append(padding) @@ -391,6 +425,7 @@ class TimesFm: freq: Sequence[int] | None = None, window_size: int | None = None, forecast_context_len: int | None = None, + return_forecast_on_context: bool = False, ) -> tuple[JTensor, JTensor]: """Forecasts on a list of time series. @@ -403,6 +438,8 @@ class TimesFm: window_size: window size of trend + residual decomposition. If None then we do not do decomposition. forecast_context_len: optional max context length. + return_forecast_on_context: True to return the forecast on the context + when available, i.e. after the first input patch. Returns: A tuple for JTensors: @@ -416,7 +453,8 @@ class TimesFm: if not self._train_state or not self._model: raise ValueError( "Checkpoint not loaded. Call `load_from_checkpoint` before" - " `forecast`.") + " `forecast`." + ) if forecast_context_len is None: forecast_context_len = self.context_len inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs] @@ -438,45 +476,50 @@ class TimesFm: full_outputs = [] assert input_ts.shape[0] % self.global_batch_size == 0 for i in range(input_ts.shape[0] // self.global_batch_size): - input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) * - self.global_batch_size]) + input_ts_in = jnp.array( + input_ts[ + i * self.global_batch_size : (i + 1) * self.global_batch_size + ] + ) input_padding_in = jnp.array( - input_padding[i * self.global_batch_size:(i + 1) * - self.global_batch_size],) + input_padding[ + i * self.global_batch_size : (i + 1) * self.global_batch_size + ], + ) inp_freq_in = jnp.array( - inp_freq[i * self.global_batch_size:(i + 1) * - self.global_batch_size, :], + inp_freq[ + i * self.global_batch_size : (i + 1) * self.global_batch_size, : + ], dtype=jnp.int32, ) pmapped_inputs = NestedMap({ - "input_ts": - es.jax_einshape( - "(db)...->db...", - input_ts_in, - d=self.num_devices, - ), - "input_padding": - es.jax_einshape( - "(db)...->db...", - input_padding_in, - d=self.num_devices, - ), - "date_features": - None, - "freq": - es.jax_einshape( - "(db)...->db...", - inp_freq_in, - d=self.num_devices, - ), + "input_ts": es.jax_einshape( + "(db)...->db...", + input_ts_in, + d=self.num_devices, + ), + "input_padding": es.jax_einshape( + "(db)...->db...", + input_padding_in, + d=self.num_devices, + ), + "date_features": None, + "freq": es.jax_einshape( + "(db)...->db...", + inp_freq_in, + d=self.num_devices, + ), }) mean_output, full_output = self._pmapped_decode(pmapped_inputs) - mean_output = es.jax_einshape("db...->(db)...", - mean_output, - d=self.num_devices) - full_output = es.jax_einshape("db...->(db)...", - full_output, - d=self.num_devices) + if not return_forecast_on_context: + mean_output = mean_output[:, :, self._horizon_start :, ...] + full_output = full_output[:, :, self._horizon_start :, ...] + mean_output = es.jax_einshape( + "db...->(db)...", mean_output, d=self.num_devices + ) + full_output = es.jax_einshape( + "db...->(db)...", full_output, d=self.num_devices + ) mean_output = np.array(mean_output) full_output = np.array(full_output) mean_outputs.append(mean_output) @@ -497,6 +540,240 @@ class TimesFm: full_outputs = np.maximum(full_outputs, 0.0) return mean_outputs, full_outputs + 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, + freq: Sequence[int] | None = None, + window_size: int | None = None, + forecast_context_len: int | 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. + freq: frequency of each context time series. 0 for high frequency + (default), 1 for medium, and 2 for low. Notice this is different from + the `freq` required by `forecast_on_df`. + window_size: window size of trend + residual decomposition. If None then + we do not do decomposition. + forecast_context_len: optional max context length. + 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. + """ + + # 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.input_patch_len)) + 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.horizon_len) + + if test_lens[-1] > self.horizon_len: + raise ValueError( + "Forecast requested longer horizon than the model definition " + f"supports: {test_lens[-1]} vs {self.horizon_len}." + ) + + # 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. + mean_outputs, _ = self.forecast( + inputs, + freq, + window_size, + forecast_context_len, + return_forecast_on_context=True, + ) + targets = [ + ( + np.array(input_ts)[-train_len:] + - mean_output[ + (self._horizon_start - train_len) : self._horizon_start + ] + ) + for input_ts, mean_output, train_len in zip( + inputs, mean_outputs, train_lens + ) + ] + per_instance_stats = None + if normalize_xreg_target_per_input: + targets, per_instance_stats = _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 = _renormalize(xregs, per_instance_stats) + outputs = [ + ( + mean_output[ + self._horizon_start : (self._horizon_start + test_len) + ] + + xreg + ) + for mean_output, test_len, xreg in zip(mean_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 = _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, + ) + mean_outputs, _ = self.forecast( + [ + target - xreg_on_context + for target, xreg_on_context in zip(targets, xregs_on_context) + ], + freq, + window_size, + forecast_context_len, + return_forecast_on_context=True, + ) + outputs = [ + ( + mean_output[ + self._horizon_start : (self._horizon_start + test_len) + ] + + xreg + ) + for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs) + ] + if normalize_xreg_target_per_input: + outputs = _renormalize(outputs, per_instance_stats) + + return outputs, xregs + def forecast_on_df( self, inputs: pd.DataFrame, @@ -527,10 +804,14 @@ class TimesFm: Returns: Future forecasts dataframe. """ - if not ("unique_id" in inputs.columns and "ds" in inputs.columns and - value_name in inputs.columns): + if not ( + "unique_id" in inputs.columns + and "ds" in inputs.columns + and value_name in inputs.columns + ): raise ValueError( - f"DataFrame must have unique_id, ds and {value_name} columns.") + f"DataFrame must have unique_id, ds and {value_name} columns." + ) if not forecast_context_len: forecast_context_len = self.context_len logging.info("Preprocessing dataframe.") @@ -555,15 +836,17 @@ class TimesFm: with multiprocessing.Pool(processes=num_jobs) as pool: results = pool.starmap( process_group, - [(key, group, value_name, forecast_context_len) - for key, group in df_sorted.groupby("unique_id")], + [ + (key, group, value_name, forecast_context_len) + for key, group in df_sorted.groupby("unique_id") + ], ) new_inputs, uids = zip(*results) print("Finished preprocessing dataframe.") freq_inps = [freq_map(freq)] * len(new_inputs) - _, full_forecast = self.forecast(new_inputs, - freq=freq_inps, - window_size=window_size) + _, full_forecast = self.forecast( + new_inputs, freq=freq_inps, window_size=window_size + ) print("Finished forecasting.") fcst_df = make_future_dataframe( uids=uids, @@ -571,13 +854,16 @@ class TimesFm: h=self.horizon_len, freq=freq, ) - fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1) + fcst_df[model_name] = full_forecast[:, 0 : self.horizon_len, 0].reshape( + -1, 1 + ) if self._model.quantiles is not None: for i, q in enumerate(self._model.quantiles): q_col = f"{model_name}-q-{q}" - fcst_df[q_col] = full_forecast[:, 0:self.horizon_len, - 1 + i].reshape(-1, 1) + fcst_df[q_col] = full_forecast[:, 0 : self.horizon_len, 1 + i].reshape( + -1, 1 + ) if q == 0.5: fcst_df[model_name] = fcst_df[q_col] logging.info("Finished creating output dataframe.") diff --git a/src/timesfm/xreg_lib.py b/src/timesfm/xreg_lib.py new file mode 100644 index 0000000..c083bcd --- /dev/null +++ b/src/timesfm/xreg_lib.py @@ -0,0 +1,532 @@ +# Copyright 2024 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 + +import jax +import jax.numpy as jnp +import numpy as np +from sklearn import preprocessing + +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}") + + +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=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