freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #91 from google-research/test_650786820

Browse Source 
In context XReg.
This commit is contained in:
Rajat Sen
2024-07-09 16:56:08 -07:00
committed by GitHub
parent 28dcfa669b 179d20df4a
commit ccd67b9221
17 changed files with 1241 additions and 334 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2 -2
View File
@@ -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 <filename>` on all affected files.
Please run `yapf --in-place --recursive <filename>` on all affected files.
environment.yml
+1
View File
@@ -15,3 +15,4 @@ dependencies:
- paxml
- jax[cuda12]==0.4.26
- einshape
- scikit-learn
environment_cpu.yml
+1
View File
@@ -15,3 +15,4 @@ dependencies:
- paxml
- jax[cpu]==0.4.26
- einshape
- scikit-learn
experiments/baselines/timegpt_pipeline.py
+35 -29
View File
@@ -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
experiments/environment.yml
+1
View File
@@ -25,3 +25,4 @@ dependencies:
- python-dotenv
- nixtla>=0.5.1
- rich
- scikit-learn
experiments/environment_cpu.yml
+1
View File
@@ -25,3 +25,4 @@ dependencies:
- python-dotenv
- nixtla>=0.5.1
- rich
- scikit-learn
experiments/extended_benchmarks/run_timegpt.py
+7 -4
View File
@@ -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)
experiments/extended_benchmarks/run_timesfm.py
+9 -5
View File
@@ -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)
experiments/extended_benchmarks/utils.py
+36 -24
View File
@@ -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])
experiments/long_horizon_benchmarks/run_eval.py
+35 -22
View File
@@ -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())
pyproject.toml
+2 -1
View File
@@ -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"},
src/timesfm/__init__.py
+1
View File
@@ -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
src/timesfm/data_loader.py
+14 -8
View File
@@ -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
src/timesfm/patched_decoder.py
+154 -115
View File
@@ -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)
src/timesfm/time_features.py
+14 -14
View File
@@ -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)
src/timesfm/timesfm.py
+396 -110
View File
@@ -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.")
src/timesfm/xreg_lib.py
+532
View File
@@ -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