diff --git a/README.md b/README.md index 72d0de8..e59ad5d 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ timesfm-1.0-200m is the first open model checkpoint: ## Benchmarks -Please refer to our result tables on the [extended benchmarks](./experiments/extended_benchmarks/tfm_results.png) and the [long horizon benchmarks](./experiments/long_horizon_benchmarks/tfm_long_horizon.png). +Please refer to our result tables on the [extended benchmarks](https://github.com/google-research/timesfm/tree/master/experiments/extended_benchmarks) and the [long horizon benchmarks](https://github.com/google-research/timesfm/tree/master/experiments/long_horizon_benchmarks). Please look into the README files in the respective benchmark directories within `experiments/` for instructions for running TimesFM on the respective benchmarks. @@ -51,6 +51,8 @@ This will install the environment in the local .venv folder (depends on the conf ### Conda / GPU installation +We recommend at least 16GB RAM to load TimesFM dependencies. + For calling TimesFM, We have two environment files. Inside `timesfm`, for GPU installation (assuming CUDA 12 has been setup), you can create a conda environment `tfm_env` from the base folder through: @@ -184,4 +186,24 @@ forecast_df = tfm.forecast_on_df( freq="M", # monthly value_name="y", num_jobs=-1, -)``` +) +``` + +## Finetuning + +We have provided an example of finetuning the model on a new dataset in `notebooks/finetuning.ipynb`. + +## Contribution Style guide + +If you would like to submit a PR please make sure that you use our formatting style. We use [yapf](https://github.com/google/yapf) for formatting with the following options, + +``` +[style] +based_on_style = google +# Add your custom style rules here +indent_width = 2 +spaces_before_comment = 2 + +``` + +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/__init__.py b/experiments/baselines/__init__.py new file mode 100644 index 0000000..2d9f7b3 --- /dev/null +++ b/experiments/baselines/__init__.py @@ -0,0 +1,13 @@ +# 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. \ No newline at end of file diff --git a/experiments/baselines/timegpt_pipeline.py b/experiments/baselines/timegpt_pipeline.py new file mode 100644 index 0000000..8a2bfbd --- /dev/null +++ b/experiments/baselines/timegpt_pipeline.py @@ -0,0 +1,259 @@ +# 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. + +import os +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, + drop_index_if_pandas, + join, + maybe_compute_sort_indices, + take_rows, + vertical_concat, +) + + +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: + if not pd.api.types.is_datetime64_any_dtype(df[col_name]): + df = df.copy() + df[col_name] = pd.to_datetime(df[col_name]) + return df + + +def zero_pad_time_series(df, freq, min_length=36): + """If time_series length is less than min_length, front pad it with zeros.""" + # 1. Calculate required padding for each unique_id + value_counts = df["unique_id"].value_counts() + to_pad = value_counts[value_counts < min_length].index + + # 2. Create a new DataFrame to hold padded data + padded_data = [] + + for unique_id in to_pad: + # 2a. Filter data for the specific unique_id + subset = df[df["unique_id"] == unique_id] + if len(subset) > min_length: + padded_data.append(subset) + else: + # 2b. Determine earliest date and calculate padding dates + start_date = subset["ds"].min() + padding_dates = pd.date_range( + end=start_date, + periods=min_length - len(subset) + 1, + freq=freq, # 'MS' for month start + )[ + :-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 + ) + + # 2d. Combine original and padding data, and append to the list + padded_data.append(pd.concat([padding_df, subset]).sort_values("ds")) + + # 3. Combine all padded data and original data (unchanged) + result_df = pd.concat(padded_data + [df[~df["unique_id"].isin(to_pad)]]) + return result_df + + +class Forecaster: + """Borrowed from + + https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models. + """ + + def forecast( + self, + df: pd.DataFrame, + h: int, + freq: str, + ) -> pd.DataFrame: + raise NotImplementedError + + def cross_validation( + self, + df: pd.DataFrame, + h: int, + freq: str, + n_windows: int = 1, + step_size: int | None = None, + ) -> pd.DataFrame: + df = maybe_convert_col_to_datetime(df, "ds") + # mlforecast cv code + results = [] + sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds") + if sort_idxs is not None: + df = take_rows(df, sort_idxs) + splits = backtest_splits( + df, + n_windows=n_windows, + h=h, + id_col="unique_id", + time_col="ds", + freq=pd.tseries.frequencies.to_offset(freq), + step_size=h if step_size is None else step_size, + ) + for _, (cutoffs, train, valid) in tqdm(enumerate(splits)): + if len(valid.columns) > 3: + raise NotImplementedError( + "Cross validation with exogenous variables is not yet supported." + ) + y_pred = self.forecast( + df=train, + h=h, + freq=freq, + ) + y_pred = join(y_pred, cutoffs, on="unique_id", how="left") + result = join( + valid[["unique_id", "ds", "y"]], + y_pred, + on=["unique_id", "ds"], + ) + 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." + ) + results.append(result) + out = vertical_concat(results) + out = drop_index_if_pandas(out) + first_out_cols = ["unique_id", "ds", "cutoff", "y"] + remaining_cols = [c for c in out.columns if c not in first_out_cols] + fcst_cv_df = out[first_out_cols + remaining_cols] + return fcst_cv_df + + +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. + """ + + def __init__( + self, + api_key: str | None = None, + base_url: Optional[str] = None, + max_retries: int = 1, + model: str = "timegpt-1", + alias: str = "TimeGPT", + ): + self.api_key = api_key + self.base_url = base_url + self.max_retries = max_retries + self.model = model + self.alias = alias + + def _get_client(self) -> NixtlaClient: + if self.api_key is None: + api_key = os.environ["NIXTLA_API_KEY"] + else: + api_key = self.api_key + return NixtlaClient( + api_key=api_key, + base_url=self.base_url, + max_retries=self.max_retries, + ) + + def forecast( + self, + df: pd.DataFrame, + h: int, + freq: str, + level: List = [90.0], + chunk_size: Optional[int] = None, + ) -> pd.DataFrame: + client = self._get_client() + fcst_df = None + if chunk_size is None: + fcst_df = client.forecast( + df=df, + h=h, + freq=freq, + level=level, + model=self.model, + ) + else: + 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_df = df[df["unique_id"].isin(chunk_ids)] + fct_chunk_df = client.forecast( + df=chunk_df, + h=h, + freq=freq, + level=level, + ) + all_fcst_df.append(fct_chunk_df) + fcst_df = pd.concat(all_fcst_df) + fcst_df["ds"] = pd.to_datetime(fcst_df["ds"]) + replace_dict = {} + for col in fcst_df.columns: + if col.startswith("TimeGPT"): + replace_dict[col] = col.replace("TimeGPT", self.alias) + fcst_df = fcst_df.rename(columns=replace_dict) + return fcst_df + + +def run_timegpt( + train_df: pd.DataFrame, + horizon: int, + freq: str, + seasonality: int, + level: List[int], + dataset: str, + model: str = "timegpt-1", +) -> Tuple[pd.DataFrame, float, str]: + os.environ["NIXTLA_ID_AS_COL"] = "true" + model = TimeGPT(model="timegpt-1", alias=model) + padded_train_df = zero_pad_time_series(train_df, freq) + init_time = time() + # For these datasets the API fails if we do not chunk. + if dataset in ["m5", "m4_quarterly"]: + 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, + ) + 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 + # forecastign metrics. + for lvl in level: + if f"{model.alias}-lo-{lvl}" not in fcsts_df.columns: + fcsts_df[f"{model.alias}-lo-{lvl}"] = fcsts_df[model.alias] + if f"{model.alias}-hi-{lvl}" not in fcsts_df.columns: + fcsts_df[f"{model.alias}-hi-{lvl}"] = fcsts_df[model.alias] + return fcsts_df, total_time, model.alias diff --git a/experiments/environment.yml b/experiments/environment.yml index 25b4b7b..32e83ee 100644 --- a/experiments/environment.yml +++ b/experiments/environment.yml @@ -22,3 +22,7 @@ dependencies: - paxml - jax[cuda12]==0.4.26 - einshape + - python-dotenv + - nixtla>=0.5.1 + - rich + - scikit-learn diff --git a/experiments/environment_cpu.yml b/experiments/environment_cpu.yml index 2301b8e..3f63013 100644 --- a/experiments/environment_cpu.yml +++ b/experiments/environment_cpu.yml @@ -22,3 +22,7 @@ dependencies: - paxml - jax[cpu]==0.4.26 - einshape + - python-dotenv + - nixtla>=0.5.1 + - rich + - scikit-learn diff --git a/experiments/extended_benchmarks/README.md b/experiments/extended_benchmarks/README.md index f37f22c..1491364 100644 --- a/experiments/extended_benchmarks/README.md +++ b/experiments/extended_benchmarks/README.md @@ -2,7 +2,6 @@ The benchmark setting has been borrowed from Nixtla's original [benchmarking](https://github.com/AzulGarza/nixtla/tree/main/experiments/amazon-chronos) of time-series foundation models against a strong statistical ensemble. Later more datasets were added by the Chronos team in this [pull request](https://github.com/shchur/nixtla/tree/chronos-full-eval/experiments/amazon-chronos). We compare on all the datasets in this extended benchmarks. -All experiments were performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). ## Running TimesFM on the benchmark @@ -19,7 +18,11 @@ Note: In the current version of TimesFM we focus on point forecasts and therefor ## Benchmark Results -![Benchmark Results Table](./tfm_results.png) +![Benchmark Results Table](./tfm_extended_new.png) + +__Update:__ We have added TimeGPT-1 to the benchmark results. We had to remove the Dominick dataset as we were not able to run TimeGPT-1 on this benchmark. Note that the previous results including Dominick remain available at `./tfm_results.png`. In order to reproduce the results for TimeGPT-1, please run `run_timegpt.py`. + +_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT. We can see that TimesFM performs the best in terms of both mase and smape. More importantly it is much faster than the other methods, in particular it is more than 600x faster than StatisticalEnsemble and 80x faster than Chronos (Large). diff --git a/experiments/extended_benchmarks/run_timegpt.py b/experiments/extended_benchmarks/run_timegpt.py new file mode 100644 index 0000000..38d964d --- /dev/null +++ b/experiments/extended_benchmarks/run_timegpt.py @@ -0,0 +1,108 @@ +# 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. + +"""Evaluation script for timegpt.""" + +import os +import sys +import time + +from absl import flags +import numpy as np +import pandas as pd + +from ..baselines.timegpt_pipeline import run_timegpt +from .utils import ExperimentHandler + + +dataset_names = [ + "m1_monthly", + "m1_quarterly", + "m1_yearly", + "m3_monthly", + "m3_other", + "m3_quarterly", + "m3_yearly", + "m4_quarterly", + "m4_yearly", + "tourism_monthly", + "tourism_quarterly", + "tourism_yearly", + "nn5_daily_without_missing", + "m5", + "nn5_weekly", + "traffic", + "weather", + "australian_electricity_demand", + "car_parts_without_missing", + "cif_2016", + "covid_deaths", + "ercot", + "ett_small_15min", + "ett_small_1h", + "exchange_rate", + "fred_md", + "hospital", +] + +_MODEL_NAME = flags.DEFINE_string( + "model_name", + "timegpt-1-long-horizon", + "Path to model, can also be set to timegpt-1", +) +_SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory") + + +QUANTILES = list(np.arange(1, 10) / 10.0) + + +def main(): + results_list = [] + run_id = np.random.randint(100000) + model_name = _MODEL_NAME.value + for dataset in dataset_names: + print(f"Evaluating model {model_name} on dataset {dataset}", flush=True) + exp = ExperimentHandler(dataset, quantiles=QUANTILES) + train_df = exp.train_df + horizon = exp.horizon + seasonality = exp.seasonality + freq = exp.freq + level = exp.level + fcsts_df, total_time, model_name = run_timegpt( + train_df=train_df, + horizon=exp.horizon, + model=model_name, + seasonality=seasonality, + freq=freq, + dataset=dataset, + level=level, + ) + 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 + ) + print(results, flush=True) + results_list.append(results) + results_full = pd.concat(results_list) + save_path = os.path.join(_SAVE_DIR.value, str(run_id)) + print(f"Saving results to {save_path}", flush=True) + os.makedirs(save_path, exist_ok=True) + results_full.to_csv(f"{save_path}/results.csv") + + +if __name__ == "__main__": + FLAGS = flags.FLAGS + FLAGS(sys.argv) + main() diff --git a/experiments/extended_benchmarks/run_timesfm.py b/experiments/extended_benchmarks/run_timesfm.py index bbf1ec9..23d19d4 100644 --- a/experiments/extended_benchmarks/run_timesfm.py +++ b/experiments/extended_benchmarks/run_timesfm.py @@ -45,7 +45,6 @@ dataset_names = [ "nn5_weekly", "traffic", "weather", - "dominick", "australian_electricity_demand", "car_parts_without_missing", "cif_2016", diff --git a/experiments/extended_benchmarks/tfm_extended_new.png b/experiments/extended_benchmarks/tfm_extended_new.png new file mode 100644 index 0000000..ee1b1a1 Binary files /dev/null and b/experiments/extended_benchmarks/tfm_extended_new.png differ diff --git a/experiments/long_horizon_benchmarks/run_eval.py b/experiments/long_horizon_benchmarks/run_eval.py index accfbd8..12a4846 100644 --- a/experiments/long_horizon_benchmarks/run_eval.py +++ b/experiments/long_horizon_benchmarks/run_eval.py @@ -24,9 +24,9 @@ import numpy as np import pandas as pd from paxml import checkpoints import timesfm +from timesfm import data_loader import torch import tqdm -from . import data_loader FLAGS = flags.FLAGS diff --git a/notebooks/finetuning.ipynb b/notebooks/finetuning.ipynb new file mode 100644 index 0000000..31da8f3 --- /dev/null +++ b/notebooks/finetuning.ipynb @@ -0,0 +1,612 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Importing relevant packages for finetuning" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'\n", + "os.environ['JAX_PMAP_USE_TENSORSTORE'] = 'false'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import timesfm\n", + "import gc\n", + "import numpy as np\n", + "import pandas as pd\n", + "from timesfm import patched_decoder\n", + "from timesfm import data_loader" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tqdm import tqdm\n", + "import dataclasses\n", + "import IPython\n", + "import IPython.display\n", + "import matplotlib as mpl\n", + "import matplotlib.pyplot as plt\n", + "mpl.rcParams['figure.figsize'] = (8, 6)\n", + "mpl.rcParams['axes.grid'] = False" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Loading TimesFM pretrained checkpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tfm = timesfm.TimesFm(\n", + " context_len=512,\n", + " horizon_len=128,\n", + " input_patch_len=32,\n", + " output_patch_len=128,\n", + " num_layers=20,\n", + " model_dims=1280,\n", + " backend=\"gpu\",\n", + ")\n", + "tfm.load_from_checkpoint(repo_id=\"google/timesfm-1.0-200m\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluating pretrained checkpoint on ETT datasets" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "DATA_DICT = {\n", + " \"ettm2\": {\n", + " \"boundaries\": [34560, 46080, 57600],\n", + " \"data_path\": \"../datasets/ETT-small/ETTm2.csv\",\n", + " \"freq\": \"15min\",\n", + " },\n", + " \"ettm1\": {\n", + " \"boundaries\": [34560, 46080, 57600],\n", + " \"data_path\": \"../datasets/ETT-small/ETTm1.csv\",\n", + " \"freq\": \"15min\",\n", + " },\n", + " \"etth2\": {\n", + " \"boundaries\": [8640, 11520, 14400],\n", + " \"data_path\": \"../datasets/ETT-small/ETTh2.csv\",\n", + " \"freq\": \"H\",\n", + " },\n", + " \"etth1\": {\n", + " \"boundaries\": [8640, 11520, 14400],\n", + " \"data_path\": \"../datasets/ETT-small/ETTh1.csv\",\n", + " \"freq\": \"H\",\n", + " },\n", + " \"elec\": {\n", + " \"boundaries\": [18413, 21044, 26304],\n", + " \"data_path\": \"../datasets/electricity/electricity.csv\",\n", + " \"freq\": \"H\",\n", + " },\n", + " \"traffic\": {\n", + " \"boundaries\": [12280, 14036, 17544],\n", + " \"data_path\": \"../datasets/traffic/traffic.csv\",\n", + " \"freq\": \"H\",\n", + " },\n", + " \"weather\": {\n", + " \"boundaries\": [36887, 42157, 52696],\n", + " \"data_path\": \"../datasets/weather/weather.csv\",\n", + " \"freq\": \"10min\",\n", + " },\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dataset = \"ettm1\"\n", + "data_path = DATA_DICT[dataset][\"data_path\"]\n", + "freq = DATA_DICT[dataset][\"freq\"]\n", + "int_freq = timesfm.freq_map(freq)\n", + "boundaries = DATA_DICT[dataset][\"boundaries\"]\n", + "\n", + "data_df = pd.read_csv(open(data_path, \"r\"))\n", + "\n", + "\n", + "ts_cols = [col for col in data_df.columns if col != \"date\"]\n", + "num_cov_cols = None\n", + "cat_cov_cols = None\n", + "\n", + "context_len = 512\n", + "pred_len = 96\n", + "\n", + "num_ts = len(ts_cols)\n", + "batch_size = 16\n", + "\n", + "dtl = data_loader.TimeSeriesdata(\n", + " data_path=data_path,\n", + " datetime_col=\"date\",\n", + " num_cov_cols=num_cov_cols,\n", + " cat_cov_cols=cat_cov_cols,\n", + " ts_cols=np.array(ts_cols),\n", + " train_range=[0, boundaries[0]],\n", + " val_range=[boundaries[0], boundaries[1]],\n", + " test_range=[boundaries[1], boundaries[2]],\n", + " hist_len=context_len,\n", + " pred_len=pred_len,\n", + " batch_size=num_ts,\n", + " freq=freq,\n", + " normalize=True,\n", + " epoch_len=None,\n", + " holiday=False,\n", + " permute=True,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "train_batches = dtl.tf_dataset(mode=\"train\", shift=1).batch(batch_size)\n", + "val_batches = dtl.tf_dataset(mode=\"val\", shift=pred_len)\n", + "test_batches = dtl.tf_dataset(mode=\"test\", shift=pred_len)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for tbatch in tqdm(train_batches.as_numpy_iterator()):\n", + " pass\n", + "print(tbatch[0].shape)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### MAE on the test split for the pretrained TimesFM model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mae_losses = []\n", + "for batch in tqdm(test_batches.as_numpy_iterator()):\n", + " past = batch[0]\n", + " actuals = batch[3]\n", + " _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n", + " forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n", + " mae_losses.append(np.abs(forecasts - actuals).mean())\n", + "\n", + "print(f\"MAE: {np.mean(mae_losses)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Finetuning the model on the ETT dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "from jax import numpy as jnp\n", + "from praxis import pax_fiddle\n", + "from praxis import py_utils\n", + "from praxis import pytypes\n", + "from praxis import base_model\n", + "from praxis import optimizers\n", + "from praxis import schedules\n", + "from praxis import base_hyperparams\n", + "from praxis import base_layer\n", + "from paxml import tasks_lib\n", + "from paxml import trainer_lib\n", + "from paxml import checkpoints\n", + "from paxml import learners\n", + "from paxml import partitioning\n", + "from paxml import checkpoint_types" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# PAX shortcuts\n", + "NestedMap = py_utils.NestedMap\n", + "WeightInit = base_layer.WeightInit\n", + "WeightHParams = base_layer.WeightHParams\n", + "InstantiableParams = py_utils.InstantiableParams\n", + "JTensor = pytypes.JTensor\n", + "NpTensor = pytypes.NpTensor\n", + "WeightedScalars = pytypes.WeightedScalars\n", + "instantiate = base_hyperparams.instantiate\n", + "LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]\n", + "AuxLossStruct = base_layer.AuxLossStruct\n", + "\n", + "AUX_LOSS = base_layer.AUX_LOSS\n", + "template_field = base_layer.template_field\n", + "\n", + "# Standard prng key names\n", + "PARAMS = base_layer.PARAMS\n", + "RANDOM = base_layer.RANDOM\n", + "\n", + "key = jax.random.PRNGKey(seed=1234)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model = pax_fiddle.Config(\n", + " patched_decoder.PatchedDecoderFinetuneModel,\n", + " name='patched_decoder_finetune',\n", + " core_layer_tpl=tfm.model_p,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### We will hold the transformer layers fixed while finetuning, while training all other components." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@pax_fiddle.auto_config\n", + "def build_learner() -> learners.Learner:\n", + " return pax_fiddle.Config(\n", + " learners.Learner,\n", + " name='learner',\n", + " loss_name='avg_qloss',\n", + " optimizer=optimizers.Adam(\n", + " epsilon=1e-7,\n", + " clip_threshold=1e2,\n", + " learning_rate=1e-2,\n", + " lr_schedule=pax_fiddle.Config(\n", + " schedules.Cosine,\n", + " initial_value=1e-3,\n", + " final_value=1e-4,\n", + " total_steps=40000,\n", + " ),\n", + " ema_decay=0.9999,\n", + " ),\n", + " # Linear probing i.e we hold the transformer layers fixed.\n", + " bprop_variable_exclusion=['.*/stacked_transformer_layer/.*'],\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "task_p = tasks_lib.SingleTask(\n", + " name='ts-learn',\n", + " model=model,\n", + " train=tasks_lib.SingleTask.Train(\n", + " learner=build_learner(),\n", + " ),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "task_p.model.ici_mesh_shape = [1, 1, 1]\n", + "task_p.model.mesh_axis_names = ['replica', 'data', 'mdl']\n", + "\n", + "DEVICES = np.array(jax.devices()).reshape([1, 1, 1])\n", + "MESH = jax.sharding.Mesh(DEVICES, ['replica', 'data', 'mdl'])\n", + "\n", + "num_devices = jax.local_device_count()\n", + "print(f'num_devices: {num_devices}')\n", + "print(f'device kind: {jax.local_devices()[0].device_kind}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "jax_task = task_p\n", + "key, init_key = jax.random.split(key)\n", + "\n", + "# To correctly prepare a batch of data for model initialization (now that shape\n", + "# inference is merged), we take one devices*batch_size tensor tuple of data,\n", + "# slice out just one batch, then run the prepare_input_batch function over it.\n", + "\n", + "\n", + "def process_train_batch(batch):\n", + " past_ts = batch[0].reshape(batch_size * num_ts, -1)\n", + " actual_ts = batch[3].reshape(batch_size * num_ts, -1)\n", + " return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n", + "\n", + "\n", + "def process_eval_batch(batch):\n", + " past_ts = batch[0]\n", + " actual_ts = batch[3]\n", + " return NestedMap(input_ts=past_ts, actual_ts=actual_ts)\n", + "\n", + "\n", + "jax_model_states, _ = trainer_lib.initialize_model_state(\n", + " jax_task,\n", + " init_key,\n", + " process_train_batch(tbatch),\n", + " checkpoint_type=checkpoint_types.CheckpointType.GDA,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Setting the initial model weights to the pretrained TimesFM parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "jax_model_states.mdl_vars['params']['core_layer'] = tfm._train_state.mdl_vars['params']\n", + "jax_vars = jax_model_states.mdl_vars\n", + "gc.collect()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Training loop" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "jax_task = task_p\n", + "\n", + "\n", + "def train_step(states, prng_key, inputs):\n", + " return trainer_lib.train_step_single_learner(\n", + " jax_task, states, prng_key, inputs\n", + " )\n", + "\n", + "\n", + "def eval_step(states, prng_key, inputs):\n", + " states = states.to_eval_state()\n", + " return trainer_lib.eval_step_single_learner(\n", + " jax_task, states, prng_key, inputs\n", + " )\n", + "\n", + "key, train_key, eval_key = jax.random.split(key, 3)\n", + "train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())\n", + "eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())\n", + "\n", + "p_train_step = jax.pmap(train_step, axis_name='batch')\n", + "p_eval_step = jax.pmap(eval_step, axis_name='batch')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)\n", + "replicated_jax_vars = replicated_jax_states.mdl_vars" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "best_eval_loss = 1e7\n", + "step_count = 0\n", + "patience = 0\n", + "NUM_EPOCHS = 100\n", + "PATIENCE = 5\n", + "TRAIN_STEPS_PER_EVAL = 1000\n", + "CHECKPOINT_DIR='/home/senrajat_google_com/ettm1_finetune'" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def reshape_batch_for_pmap(batch, num_devices):\n", + " def _reshape(input_tensor):\n", + " bsize = input_tensor.shape[0]\n", + " residual_shape = list(input_tensor.shape[1:])\n", + " nbsize = bsize // num_devices\n", + " return jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)\n", + "\n", + " return jax.tree.map(_reshape, batch)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for epoch in range(NUM_EPOCHS):\n", + " print(f\"__________________Epoch: {epoch}__________________\", flush=True)\n", + " train_its = train_batches.as_numpy_iterator()\n", + " if patience >= PATIENCE:\n", + " print(\"Early stopping.\", flush=True)\n", + " break\n", + " for batch in tqdm(train_its):\n", + " train_losses = []\n", + " if patience >= PATIENCE:\n", + " print(\"Early stopping.\", flush=True)\n", + " break\n", + " tbatch = process_train_batch(batch)\n", + " tbatch = reshape_batch_for_pmap(tbatch, num_devices)\n", + " replicated_jax_states, step_fun_out = p_train_step(\n", + " replicated_jax_states, train_prng_seed, tbatch\n", + " )\n", + " train_losses.append(step_fun_out.loss[0])\n", + " if step_count % TRAIN_STEPS_PER_EVAL == 0:\n", + " print(\n", + " f\"Train loss at step {step_count}: {np.mean(train_losses)}\",\n", + " flush=True,\n", + " )\n", + " train_losses = []\n", + " print(\"Starting eval.\", flush=True)\n", + " val_its = val_batches.as_numpy_iterator()\n", + " eval_losses = []\n", + " for ev_batch in tqdm(val_its):\n", + " ebatch = process_eval_batch(ev_batch)\n", + " ebatch = reshape_batch_for_pmap(ebatch, num_devices)\n", + " _, step_fun_out = p_eval_step(\n", + " replicated_jax_states, eval_prng_seed, ebatch\n", + " )\n", + " eval_losses.append(step_fun_out.loss[0])\n", + " mean_loss = np.mean(eval_losses)\n", + " print(f\"Eval loss at step {step_count}: {mean_loss}\", flush=True)\n", + " if mean_loss < best_eval_loss or np.isnan(mean_loss):\n", + " best_eval_loss = mean_loss\n", + " print(\"Saving checkpoint.\")\n", + " jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(\n", + " replicated_jax_states\n", + " )\n", + " checkpoints.save_checkpoint(\n", + " jax_state_for_saving, CHECKPOINT_DIR, overwrite=True\n", + " )\n", + " patience = 0\n", + " del jax_state_for_saving\n", + " gc.collect()\n", + " else:\n", + " patience += 1\n", + " print(f\"patience: {patience}\")\n", + " step_count += 1" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Loading and evaluating the best (according to validation loss) finetuned checkpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "train_state = checkpoints.restore_checkpoint(jax_model_states, CHECKPOINT_DIR)\n", + "print(train_state.step)\n", + "tfm._train_state.mdl_vars['params'] = train_state.mdl_vars['params']['core_layer']\n", + "tfm.jit_decode()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mae_losses = []\n", + "for batch in tqdm(test_batches.as_numpy_iterator()):\n", + " past = batch[0]\n", + " actuals = batch[3]\n", + " _, forecasts = tfm.forecast(list(past), [0] * past.shape[0])\n", + " forecasts = forecasts[:, 0 : actuals.shape[1], 5]\n", + " mae_losses.append(np.abs(forecasts - actuals).mean())\n", + "\n", + "print(f\"MAE: {np.mean(mae_losses)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## There is around a __9%__ reduction in MAE from finetuning." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "tfm_env_v3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/pyproject.toml b/pyproject.toml index 168ddcd..0bac628 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,25 @@ -[tool.poetry] -name = "timesfm" +# This project can be installed with `python3 -m pip install -e .` from the main directory. + +[project] +name = "timesfm-jax" packages = [ { include = "*", from = "src" }, ] version = "0.0.1" description = "Open weights time-series foundation model from Google Research." +version = "1.0.1" +dependencies = [ + "jax==0.4.26", + "paxml==1.4.0", + "praxis==1.4.0", + "jaxlib==0.4.26", + "numpy==1.26.4", + "pandas==2.1.4", + "einshape==1.0.0", + "utilsforecast==0.1.10", + "huggingface_hub[cli]==0.23.0", + "scikit-learn==1.5.1", +] authors = [ "Rajat Sen ", "Yichen Zhou ", diff --git a/src/__init__.py b/src/timesfm/__init__.py similarity index 87% rename from src/__init__.py rename to src/timesfm/__init__.py index 646426d..866e848 100644 --- a/src/__init__.py +++ b/src/timesfm/__init__.py @@ -14,7 +14,4 @@ """TimesFM init file.""" -from __future__ import absolute_import - -from .patched_decoder import PatchedTimeSeriesDecoder from .timesfm import TimesFm, freq_map diff --git a/experiments/long_horizon_benchmarks/data_loader.py b/src/timesfm/data_loader.py similarity index 100% rename from experiments/long_horizon_benchmarks/data_loader.py rename to src/timesfm/data_loader.py diff --git a/src/patched_decoder.py b/src/timesfm/patched_decoder.py similarity index 79% rename from src/patched_decoder.py rename to src/timesfm/patched_decoder.py index b0decf8..c7c3758 100644 --- a/src/patched_decoder.py +++ b/src/timesfm/patched_decoder.py @@ -24,6 +24,7 @@ import einshape as es from jax import lax import jax.numpy as jnp from praxis import base_layer +from praxis import base_model from praxis import layers from praxis import pax_fiddle from praxis import py_utils @@ -49,6 +50,7 @@ DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] # NestedMap keys _INPUT_TS = "input_ts" +_TARGET_FUTURE = "actual_ts" _INPUT_PADDING = "input_padding" _OUTPUT_TS = "output_ts" _FREQ = "freq" @@ -324,15 +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 = 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) @@ -404,6 +410,7 @@ 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. @@ -414,15 +421,19 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): 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. + - 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). + 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) @@ -439,7 +450,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): num_decode_patches = ( horizon_len + output_patch_len - 1 ) // output_patch_len - for _ in range(num_decode_patches): + 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:] @@ -449,13 +460,96 @@ 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. + """ + + core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder) + freq: int = 0 + + def setup(self) -> None: + self.create_child("core_layer", self.core_layer_tpl) + + def compute_predictions(self, input_batch: NestedMap) -> NestedMap: + input_ts = input_batch[_INPUT_TS] + 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 + + 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 ) + freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq + new_input_batch = NestedMap( + input_ts=input_ts, + input_padding=input_padding, + freq=freq, + ) + return self.core_layer(new_input_batch) + + 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. + + 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]: + output_ts = prediction_output[_OUTPUT_TS] + actual_ts = input_batch[_TARGET_FUTURE] + 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) + loss = loss.mean() + loss_weight = jnp.array(1.0, dtype=jnp.float32) + per_example_out = NestedMap() + return {"avg_qloss": (loss, loss_weight)}, per_example_out diff --git a/experiments/long_horizon_benchmarks/time_features.py b/src/timesfm/time_features.py similarity index 100% rename from experiments/long_horizon_benchmarks/time_features.py rename to src/timesfm/time_features.py diff --git a/src/timesfm.py b/src/timesfm/timesfm.py similarity index 66% rename from src/timesfm.py rename to src/timesfm/timesfm.py index 7dc9d0a..ec63a20 100644 --- a/src/timesfm.py +++ b/src/timesfm/timesfm.py @@ -14,6 +14,7 @@ """TimesFM forecast API for inference.""" +import collections import logging import multiprocessing from os import path @@ -21,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 @@ -35,12 +36,19 @@ from praxis import py_utils from praxis import pytypes from praxis.layers import normalizations from praxis.layers import transformers -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): @@ -78,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. @@ -157,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"] @@ -277,6 +300,10 @@ class TimesFm: self._logging( f"Restored checkpoint in {time.time() - start_time:.2f} seconds." ) + self.jit_decode() + + def jit_decode(self): + """Jitting decoding function.""" # Initialize and jit the decode fn. def _decode(inputs): @@ -288,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, @@ -397,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. @@ -409,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: @@ -480,6 +511,9 @@ class TimesFm: ), }) mean_output, full_output = self._pmapped_decode(pmapped_inputs) + 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 ) @@ -506,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, 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