From 41929ba643f0b2639d50bb4a7088cc3fad5e5f40 Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Tue, 9 Jul 2024 17:10:18 +0000 Subject: [PATCH 1/2] standardizing styles accross all files. --- experiments/baselines/timegpt_pipeline.py | 364 +++++----- .../extended_benchmarks/run_timegpt.py | 73 +- .../extended_benchmarks/run_timesfm.py | 14 +- experiments/extended_benchmarks/utils.py | 60 +- .../long_horizon_benchmarks/run_eval.py | 55 +- src/timesfm/__init__.py | 1 - src/timesfm/data_loader.py | 22 +- src/timesfm/patched_decoder.py | 641 +++++++++--------- src/timesfm/time_features.py | 28 +- src/timesfm/timesfm.py | 241 +++---- 10 files changed, 718 insertions(+), 781 deletions(-) diff --git a/experiments/baselines/timegpt_pipeline.py b/experiments/baselines/timegpt_pipeline.py index c5a38f5..5ea18ad 100644 --- a/experiments/baselines/timegpt_pipeline.py +++ b/experiments/baselines/timegpt_pipeline.py @@ -31,191 +31,191 @@ from nixtla import NixtlaClient def get_seasonality(freq: str) -> int: - return _get_seasonality(freq, seasonalities={"D": 7}) + 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 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 + """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 = [] + # 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 + 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 - ) + # 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")) + # 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 + # 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 + """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 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 + 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 + """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 __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 _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, ) - - 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 + 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( @@ -227,25 +227,27 @@ def run_timegpt( 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 + 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/extended_benchmarks/run_timegpt.py b/experiments/extended_benchmarks/run_timegpt.py index cceab37..455782d 100644 --- a/experiments/extended_benchmarks/run_timegpt.py +++ b/experiments/extended_benchmarks/run_timegpt.py @@ -11,7 +11,6 @@ # 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 @@ -25,7 +24,6 @@ from ..baselines.timegpt_pipeline import run_timegpt from .utils import ExperimentHandler - dataset_names = [ "m1_monthly", "m1_quarterly", @@ -63,46 +61,45 @@ _MODEL_NAME = flags.DEFINE_string( ) _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") + 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() + FLAGS = flags.FLAGS + FLAGS(sys.argv) + main() diff --git a/experiments/extended_benchmarks/run_timesfm.py b/experiments/extended_benchmarks/run_timesfm.py index 23d19d4..f01bf39 100644 --- a/experiments/extended_benchmarks/run_timesfm.py +++ b/experiments/extended_benchmarks/run_timesfm.py @@ -11,7 +11,6 @@ # 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 @@ -26,7 +25,6 @@ import timesfm from .utils import ExperimentHandler - dataset_names = [ "m1_monthly", "m1_quarterly", @@ -74,16 +72,14 @@ 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) @@ -127,9 +123,9 @@ def main(): ) total_time = time.time() - init_time time_df = pd.DataFrame({"time": [total_time], "model": model_name}) - results = exp.evaluate_from_predictions( - models=[model_name], fcsts_df=fcsts_df, times_df=time_df - ) + results = exp.evaluate_from_predictions(models=[model_name], + fcsts_df=fcsts_df, + times_df=time_df) print(results, flush=True) results_list.append(results) results_full = pd.concat(results_list) diff --git a/experiments/extended_benchmarks/utils.py b/experiments/extended_benchmarks/utils.py index de0368b..c9307f4 100644 --- a/experiments/extended_benchmarks/utils.py +++ b/experiments/extended_benchmarks/utils.py @@ -11,7 +11,6 @@ # 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 @@ -46,11 +45,9 @@ 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 @@ -66,10 +63,8 @@ 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) @@ -80,10 +75,8 @@ 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 @@ -122,9 +115,8 @@ 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 @@ -153,9 +145,8 @@ 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 @@ -177,9 +168,8 @@ 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", @@ -215,23 +205,21 @@ 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") @@ -262,9 +250,9 @@ class ExperimentHandler: eval_prob_df["metric"] = "scaled_crps" eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True) eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index() - eval_df = eval_df.melt( - id_vars="metric", value_name="value", var_name="model" - ) + eval_df = eval_df.melt(id_vars="metric", + value_name="value", + var_name="model") times_df.insert(0, "metric", "time") times_df = times_df.rename(columns={"time": "value"}) eval_df = pd.concat([eval_df, times_df]) diff --git a/experiments/long_horizon_benchmarks/run_eval.py b/experiments/long_horizon_benchmarks/run_eval.py index a079496..981f11f 100644 --- a/experiments/long_horizon_benchmarks/run_eval.py +++ b/experiments/long_horizon_benchmarks/run_eval.py @@ -11,7 +11,6 @@ # 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 @@ -28,38 +27,28 @@ 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": { @@ -176,9 +165,8 @@ 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( @@ -213,10 +201,9 @@ def eval(): for batch in tqdm.tqdm(eval_itr): past = batch[0] actuals = batch[3] - forecasts = get_forecasts( - model_path, model, past, int_freq, _PRED_LEN.value - ) - forecasts = forecasts[:, 0 : actuals.shape[1]] + forecasts = get_forecasts(model_path, model, past, int_freq, + _PRED_LEN.value) + forecasts = forecasts[:, 0:actuals.shape[1]] mae_run_losses.append(_mae(forecasts, actuals).sum()) mse_run_losses.append(_mse(forecasts, actuals).sum()) smape_run_losses.append(_smape(forecasts, actuals).sum()) diff --git a/src/timesfm/__init__.py b/src/timesfm/__init__.py index 866e848..e794e45 100644 --- a/src/timesfm/__init__.py +++ b/src/timesfm/__init__.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - """TimesFM init file.""" from .timesfm import TimesFm, freq_map diff --git a/src/timesfm/data_loader.py b/src/timesfm/data_loader.py index eeace7b..b524bd4 100644 --- a/src/timesfm/data_loader.py +++ b/src/timesfm/data_loader.py @@ -11,13 +11,11 @@ # 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 @@ -79,9 +77,8 @@ 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 @@ -94,18 +91,16 @@ 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) @@ -135,7 +130,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() @@ -253,9 +248,8 @@ class TimeSeriesdata(object): gen_fn = self.train_gen else: gen_fn = lambda: self.test_val_gen(mode, shift) - output_types = tuple( - [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2 - ) + output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + + [tf.int32] * 2) dataset = tf.data.Dataset.from_generator(gen_fn, output_types) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) return dataset diff --git a/src/timesfm/patched_decoder.py b/src/timesfm/patched_decoder.py index 5dbdcbc..4e2600a 100644 --- a/src/timesfm/patched_decoder.py +++ b/src/timesfm/patched_decoder.py @@ -11,7 +11,6 @@ # 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. @@ -36,7 +35,6 @@ from praxis.layers import normalizations from praxis.layers import stochastics from praxis.layers import transformers - # PAX shortcuts NestedMap = py_utils.NestedMap JTensor = pytypes.JTensor @@ -44,7 +42,6 @@ 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] @@ -57,36 +54,35 @@ _FREQ = "freq" _OUTPUT_TOKENS = "output_tokens" _STATS = "stats" - # Small numerical value. _TOLERANCE = 1e-7 def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor: - """Shifts rows of seq based on the first 0 in each row of the mask.""" - num = seq.shape[1] + """Shifts rows of seq based on the first 0 in each row of the mask.""" + num = seq.shape[1] - # Find the index of the first 0 in each row of the mask - first_zero_idx = jnp.argmin(mask, axis=1) + # Find the index of the first 0 in each row of the mask + first_zero_idx = jnp.argmin(mask, axis=1) - # Create a range array for indexing - idx_range = jnp.arange(num) + # Create a range array for indexing + idx_range = jnp.arange(num) - def shift_row(carry, x): - seq_row, shift = x - shifted_idx = (idx_range - shift) % num - shifted_row = seq_row[shifted_idx] - return carry, shifted_row + def shift_row(carry, x): + seq_row, shift = x + shifted_idx = (idx_range - shift) % num + shifted_row = seq_row[shifted_idx] + return carry, shifted_row - # Use lax.scan to shift each row of seq based on the corresponding - # first_zero_idx. - _, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx)) + # Use lax.scan to shift each row of seq based on the corresponding + # first_zero_idx. + _, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx)) - return shifted_seq + return shifted_seq class ResidualBlock(base_layer.BaseLayer): - """Simple feedforward block with residual connection. + """Simple feedforward block with residual connection. Attributes: input_dims: input dimension. @@ -99,67 +95,68 @@ class ResidualBlock(base_layer.BaseLayer): act_tpl: config for activation in hidden layer. """ - input_dims: int = 0 - hidden_dims: int = 0 - output_dims: int = 0 - dropout_prob: float = 0.0 - layer_norm: bool = False - dropout_tpl: LayerTpl = template_field(stochastics.Dropout) - ln_tpl: LayerTpl = template_field(normalizations.LayerNorm) - act_tpl: LayerTpl = template_field(activations.Swish) + input_dims: int = 0 + hidden_dims: int = 0 + output_dims: int = 0 + dropout_prob: float = 0.0 + layer_norm: bool = False + dropout_tpl: LayerTpl = template_field(stochastics.Dropout) + ln_tpl: LayerTpl = template_field(normalizations.LayerNorm) + act_tpl: LayerTpl = template_field(activations.Swish) - def setup(self): - lnorm_tpl = self.ln_tpl.clone() - lnorm_tpl.dim = self.output_dims - self.create_child("ln_layer", lnorm_tpl) + def setup(self): + lnorm_tpl = self.ln_tpl.clone() + lnorm_tpl.dim = self.output_dims + self.create_child("ln_layer", lnorm_tpl) - dropout_tpl = self.dropout_tpl.clone() - dropout_tpl.keep_prob = 1.0 - self.dropout_prob - self.create_child("dropout", dropout_tpl) + dropout_tpl = self.dropout_tpl.clone() + dropout_tpl.keep_prob = 1.0 - self.dropout_prob + self.create_child("dropout", dropout_tpl) - self.create_child( - "hidden_layer", - pax_fiddle.Config( - linears.FeedForward, - input_dims=self.input_dims, - output_dims=self.hidden_dims, - activation_tpl=self.act_tpl.clone(), - ), - ) + self.create_child( + "hidden_layer", + pax_fiddle.Config( + linears.FeedForward, + input_dims=self.input_dims, + output_dims=self.hidden_dims, + activation_tpl=self.act_tpl.clone(), + ), + ) - self.create_child( - "output_layer", - pax_fiddle.Config( - linears.FeedForward, - input_dims=self.hidden_dims, - output_dims=self.output_dims, - activation_tpl=pax_fiddle.Config(activations.Identity), - ), - ) + self.create_child( + "output_layer", + pax_fiddle.Config( + linears.FeedForward, + input_dims=self.hidden_dims, + output_dims=self.output_dims, + activation_tpl=pax_fiddle.Config(activations.Identity), + ), + ) - self.create_child( - "residual_layer", - pax_fiddle.Config( - linears.FeedForward, - input_dims=self.input_dims, - output_dims=self.output_dims, - activation_tpl=pax_fiddle.Config(activations.Identity), - ), - ) + self.create_child( + "residual_layer", + pax_fiddle.Config( + linears.FeedForward, + input_dims=self.input_dims, + output_dims=self.output_dims, + activation_tpl=pax_fiddle.Config(activations.Identity), + ), + ) - def __call__(self, inputs: JTensor) -> JTensor: - hidden = self.hidden_layer(inputs) - output = self.output_layer(hidden) - output = self.dropout(output) - residual = self.residual_layer(inputs) - if self.layer_norm: - return self.ln_layer(output + residual) - else: - return output + residual + def __call__(self, inputs: JTensor) -> JTensor: + hidden = self.hidden_layer(inputs) + output = self.output_layer(hidden) + output = self.dropout(output) + residual = self.residual_layer(inputs) + if self.layer_norm: + return self.ln_layer(output + residual) + else: + return output + residual -def _masked_mean_std(inputs: JTensor, padding: JTensor) -> Tuple[JTensor, JTensor]: - """Calculates mean and standard deviation of arr across axis 1. +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. @@ -171,48 +168,48 @@ def _masked_mean_std(inputs: JTensor, padding: JTensor) -> Tuple[JTensor, JTenso 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) + # Selecting the first pad with more than 3 unpadded values. + pad_sum = jnp.sum(1 - padding, axis=2) - def _get_patch_index(arr: JTensor): - indices = jnp.argmax(arr >= 3, axis=1) - row_sum = (arr >= 3).sum(axis=1) - return jnp.where(row_sum == 0, arr.shape[1] - 1, indices) + def _get_patch_index(arr: JTensor): + indices = jnp.argmax(arr >= 3, axis=1) + row_sum = (arr >= 3).sum(axis=1) + return jnp.where(row_sum == 0, arr.shape[1] - 1, indices) - patch_indices = _get_patch_index(pad_sum) - bidxs = jnp.arange(inputs.shape[0]) + patch_indices = _get_patch_index(pad_sum) + bidxs = jnp.arange(inputs.shape[0]) - arr = inputs[bidxs, patch_indices, :] - pad = padding[bidxs, patch_indices, :] + arr = inputs[bidxs, patch_indices, :] + pad = padding[bidxs, patch_indices, :] - # Create a mask where P is 0 - mask = 1 - pad + # Create a mask where P is 0 + mask = 1 - pad - # Calculate the number of valid elements - num_valid_elements = jnp.sum(mask, axis=1) + # Calculate the number of valid elements + num_valid_elements = jnp.sum(mask, axis=1) - num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements) + num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements) - # 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) + # 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) - # Calculate the masked mean and standard deviation - masked_mean = masked_sum / num_valid_elements - masked_var = masked_squared_sum / num_valid_elements - masked_mean**2 - masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var) - masked_std = jnp.sqrt(masked_var) + # Calculate the masked mean and standard deviation + masked_mean = masked_sum / num_valid_elements + masked_var = masked_squared_sum / num_valid_elements - masked_mean**2 + masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var) + masked_std = jnp.sqrt(masked_var) - return masked_mean, masked_std + return masked_mean, masked_std def _create_quantiles() -> list[float]: - """Returns the quantiles for forecasting.""" - return DEFAULT_QUANTILES + """Returns the quantiles for forecasting.""" + return DEFAULT_QUANTILES class PatchedTimeSeriesDecoder(base_layer.BaseLayer): - """Patch decoder layer for time-series foundation model. + """Patch decoder layer for time-series foundation model. Attributes: patch_len: length of input patches. @@ -231,137 +228,137 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): number of output logits. D is model dimension. """ - patch_len: int = 0 - horizon_len: int = 0 - model_dims: int = 0 - hidden_dims: int = 0 - 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 + patch_len: int = 0 + horizon_len: int = 0 + model_dims: int = 0 + hidden_dims: int = 0 + 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) + use_freq: bool = True + + def setup(self) -> None: + """Construct the model.""" + num_outputs = len(self.quantiles) + 1 + + stl = self.stacked_transformer_params_tpl.clone() + stl.model_dims = self.model_dims + stl.hidden_dims = self.hidden_dims + stl.mask_self_attention = True + + self.create_child("stacked_transformer_layer", stl) + + input_resl = self.residual_block_tpl.clone() + ff_in_dims = 2 * self.patch_len + input_resl.input_dims = ff_in_dims + input_resl.hidden_dims = self.hidden_dims + input_resl.output_dims = self.model_dims + self.create_child( + "input_ff_layer", + input_resl, ) - use_freq: bool = True - def setup(self) -> None: - """Construct the model.""" - num_outputs = len(self.quantiles) + 1 + horizon_resl = self.residual_block_tpl.clone() + horizon_resl.input_dims = self.model_dims + horizon_resl.hidden_dims = self.hidden_dims + horizon_resl.output_dims = self.horizon_len * num_outputs + self.create_child( + "horizon_ff_layer", + horizon_resl, + ) - stl = self.stacked_transformer_params_tpl.clone() - stl.model_dims = self.model_dims - stl.hidden_dims = self.hidden_dims - stl.mask_self_attention = True + self.create_child( + "position_emb", + pax_fiddle.Config(layers.PositionalEmbedding, + embedding_dims=self.model_dims), + ) - self.create_child("stacked_transformer_layer", stl) + if self.use_freq: + self.create_child( + "freq_emb", + pax_fiddle.Config( + embedding_softmax.Embedding, + num_classes=3, + input_dims=self.model_dims, + ), + ) - input_resl = self.residual_block_tpl.clone() - ff_in_dims = 2 * self.patch_len - input_resl.input_dims = ff_in_dims - input_resl.hidden_dims = self.hidden_dims - input_resl.output_dims = self.model_dims - self.create_child( - "input_ff_layer", - input_resl, - ) + def transform_decode_state( + 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) - horizon_resl = self.residual_block_tpl.clone() - horizon_resl.input_dims = self.model_dims - horizon_resl.hidden_dims = self.hidden_dims - horizon_resl.output_dims = self.horizon_len * num_outputs - self.create_child( - "horizon_ff_layer", - horizon_resl, - ) + def _forward_transform( + 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) + return outputs, (mu, sigma) - self.create_child( - "position_emb", - pax_fiddle.Config( - layers.PositionalEmbedding, embedding_dims=self.model_dims - ), - ) + 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] - if self.use_freq: - self.create_child( - "freq_emb", - pax_fiddle.Config( - embedding_softmax.Embedding, - num_classes=3, - input_dims=self.model_dims, - ), - ) + def _preprocess_input( + self, + input_ts: JTensor, + input_padding: JTensor, + pos_emb: Optional[JTensor] = None, + ) -> Tuple[JTensor, JTensor, Optional[Tuple[JTensor, JTensor]], JTensor]: + """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) + # B x N x D + patched_inputs = patched_inputs * (1.0 - patched_pads) + concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1) + model_input = self.input_ff_layer(concat_inputs) + # A patch should not be padded even if there is at least one zero. + patched_padding = jnp.min(patched_pads, axis=-1) - def transform_decode_state( - 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) + if pos_emb is None: + position_emb = self.position_emb(seq_length=model_input.shape[1]) + else: + position_emb = pos_emb + if self.do_eval: + if position_emb.shape[0] != model_input.shape[0]: + position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0) + position_emb = _shift_padded_seq(patched_padding, position_emb) + model_input += position_emb - def _forward_transform( - 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) - return outputs, (mu, sigma) + return model_input, patched_padding, stats, patched_inputs - 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] + def _postprocess_output( + self, + model_output: JTensor, + num_outputs: int, + stats: Tuple[JTensor, JTensor], + ) -> JTensor: + """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) + return self._reverse_transform(output_ts, stats) - def _preprocess_input( - self, - input_ts: JTensor, - input_padding: JTensor, - pos_emb: Optional[JTensor] = None, - ) -> Tuple[JTensor, JTensor, Optional[Tuple[JTensor, JTensor]], JTensor]: - """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) - # B x N x D - patched_inputs = patched_inputs * (1.0 - patched_pads) - concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1) - model_input = self.input_ff_layer(concat_inputs) - # A patch should not be padded even if there is at least one zero. - patched_padding = jnp.min(patched_pads, axis=-1) - - if pos_emb is None: - position_emb = self.position_emb(seq_length=model_input.shape[1]) - else: - position_emb = pos_emb - if self.do_eval: - if position_emb.shape[0] != model_input.shape[0]: - position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0) - position_emb = _shift_padded_seq(patched_padding, position_emb) - model_input += position_emb - - return model_input, patched_padding, stats, patched_inputs - - def _postprocess_output( - self, - model_output: JTensor, - num_outputs: int, - stats: Tuple[JTensor, JTensor], - ) -> JTensor: - """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 - ) - return self._reverse_transform(output_ts, stats) - - def __call__(self, inputs: NestedMap) -> NestedMap: - """PatchTST call. + def __call__(self, inputs: NestedMap) -> NestedMap: + """PatchTST call. Args: inputs: A NestedMap containing (1) input_ts: input sequence of shape [B, @@ -374,32 +371,34 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): (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( - input_ts=input_ts, - input_padding=input_padding, - ) - if self.use_freq: - freq = inputs[_FREQ].astype(jnp.int32) - f_emb = self.freq_emb(freq) # B x 1 x D - f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1) - model_input += f_emb - model_output = self.stacked_transformer_layer(model_input, patched_padding) + input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING] + num_outputs = len(self.quantiles) + 1 + model_input, patched_padding, stats, _ = self._preprocess_input( + input_ts=input_ts, + input_padding=input_padding, + ) + if self.use_freq: + freq = inputs[_FREQ].astype(jnp.int32) + f_emb = self.freq_emb(freq) # B x 1 x D + f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1) + model_input += f_emb + 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} - ) + output_ts = self._postprocess_output(model_output, num_outputs, stats) + return NestedMap({ + _OUTPUT_TOKENS: model_output, + _OUTPUT_TS: output_ts, + _STATS: stats + }) - def decode( - self, - inputs: NestedMap, - horizon_len: int, - output_patch_len: Optional[int] = None, - max_len: int = 512, - ) -> tuple[JTensor, JTensor]: - """Auto-regressive decoding without caching. + def decode( + self, + inputs: NestedMap, + horizon_len: int, + output_patch_len: Optional[int] = None, + max_len: int = 512, + ) -> tuple[JTensor, JTensor]: + """Auto-regressive decoding without caching. Args: inputs: input time-series and paddings. Time-series shape B x C, padding @@ -415,83 +414,80 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer): - Full predictions (mean and quantiles) as a tensor with shape B x H x (1 + # quantiles). """ - final_out = inputs[_INPUT_TS] - inp_time_len = final_out.shape[1] - paddings = inputs[_INPUT_PADDING] - if self.use_freq: - freq = inputs[_FREQ].astype(jnp.int32) - else: - freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32) - full_outputs = [] - 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}" - ) - 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]] - input_ts = final_out[:, -max_len:] - input_padding = current_padding[:, -max_len:] - model_input = NestedMap( - input_ts=input_ts, - input_padding=input_padding, - freq=freq, - ) - fprop_outputs = self(model_input)[_OUTPUT_TS] - # (full batch, last patch, output_patch_len, index of mean forecast = 0) - new_ts = fprop_outputs[:, -1, :output_patch_len, 0] - # (full batch, last patch, output_patch_len, all output indices) - full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :]) - final_out = jnp.concatenate([final_out, new_ts], axis=-1) + final_out = inputs[_INPUT_TS] + inp_time_len = final_out.shape[1] + paddings = inputs[_INPUT_PADDING] + if self.use_freq: + freq = inputs[_FREQ].astype(jnp.int32) + else: + freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32) + full_outputs = [] + 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}") + 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]] + input_ts = final_out[:, -max_len:] + input_padding = current_padding[:, -max_len:] + model_input = NestedMap( + input_ts=input_ts, + input_padding=input_padding, + freq=freq, + ) + fprop_outputs = self(model_input)[_OUTPUT_TS] + # (full batch, last patch, output_patch_len, index of mean forecast = 0) + new_ts = fprop_outputs[:, -1, :output_patch_len, 0] + # (full batch, last patch, output_patch_len, all output indices) + full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :]) + 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, :], - ) + return ( + final_out[:, inp_time_len:inp_time_len + horizon_len], + jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :], + ) class PatchedDecoderFinetuneModel(base_model.BaseModel): - """Model class for finetuning patched time-series decoder. + """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 + core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder) + freq: int = 0 - def setup(self) -> None: - self.create_child("core_layer", self.core_layer_tpl) + 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 + 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) + 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. + def _quantile_loss(self, pred: JTensor, actual: JTensor, + quantile: float) -> JTensor: + """Calculates quantile loss. Args: pred: B x T @@ -501,21 +497,20 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel): 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) + 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 + 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/src/timesfm/time_features.py b/src/timesfm/time_features.py index 0bd90a9..0cc0861 100644 --- a/src/timesfm/time_features.py +++ b/src/timesfm/time_features.py @@ -11,7 +11,6 @@ # 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. @@ -36,7 +35,6 @@ 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 @@ -50,9 +48,8 @@ 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 @@ -60,16 +57,19 @@ def _distance_to_holiday(holiday): return _distance_to_day -EasterSunday = Holiday( - "Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)] -) +EasterSunday = Holiday("Easter Sunday", + month=1, + day=1, + offset=[Easter(), Day(0)]) NewYearsDay = Holiday("New Years Day", month=1, day=1) -SuperBowl = Holiday( - "Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1)) -) -MothersDay = Holiday( - "Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2)) -) +SuperBowl = Holiday("Superbowl", + month=2, + day=1, + offset=DateOffset(weekday=SU(1))) +MothersDay = Holiday("Mothers Day", + month=5, + day=1, + offset=DateOffset(weekday=SU(2))) IndependenceDay = Holiday("Independence Day", month=7, day=4) ChristmasEve = Holiday("Christmas", month=12, day=24) ChristmasDay = Holiday("Christmas", month=12, day=25) diff --git a/src/timesfm/timesfm.py b/src/timesfm/timesfm.py index 5dc7611..70c5292 100644 --- a/src/timesfm/timesfm.py +++ b/src/timesfm/timesfm.py @@ -11,7 +11,6 @@ # 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 logging @@ -52,23 +51,16 @@ 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 @@ -179,9 +171,7 @@ 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,), ), ), ) @@ -199,34 +189,38 @@ class TimesFm: def _get_sample_inputs(self): return { - "input_ts": jnp.zeros( - ( - self.per_core_batch_size, - self.context_len + self.output_patch_len, + "input_ts": + 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, + "input_padding": + jnp.zeros( + ( + self.per_core_batch_size, + self.context_len + self.output_patch_len, + ), + dtype=jnp.float32, ), - dtype=jnp.float32, - ), - "freq": jnp.zeros( - ( - self.per_core_batch_size, - 1, + "freq": + jnp.zeros( + ( + self.per_core_batch_size, + 1, + ), + dtype=jnp.int32, ), - 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. @@ -246,8 +240,7 @@ 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, @@ -261,8 +254,7 @@ 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}.") @@ -275,12 +267,12 @@ 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): """Jitting decoding function.""" + # Initialize and jit the decode fn. def _decode(inputs): assert self._model is not None @@ -310,34 +302,36 @@ 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, + "input_ts": + jnp.zeros( + ( + self.num_devices, + self.per_core_batch_size, + self.context_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, + "input_padding": + jnp.zeros( + ( + self.num_devices, + self.per_core_batch_size, + self.context_len + self.horizon_len, + ), + dtype=jnp.float32, ), - dtype=jnp.float32, - ), - "date_features": None, - "freq": jnp.zeros( - (self.num_devices, self.per_core_batch_size, 1), - dtype=jnp.int32, - ), - }) - ) + "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 @@ -358,24 +352,21 @@ 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) @@ -425,8 +416,7 @@ 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] @@ -448,47 +438,45 @@ 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 - ) + 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) @@ -539,14 +527,10 @@ 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.") @@ -571,17 +555,15 @@ 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, @@ -589,16 +571,13 @@ 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.") From 954b7ff4ac031c1ffaf9fed10465d75b9ecafed0 Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Tue, 9 Jul 2024 17:13:51 +0000 Subject: [PATCH 2/2] adding formatting instructions to README --- README.md | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 1f4f18e..b252802 100644 --- a/README.md +++ b/README.md @@ -164,4 +164,19 @@ forecast_df = tfm.forecast_on_df( ## Finetuning -We have provided an example of finetuning the model on a new dataset in `notebooks/finetuning.ipynb`. \ No newline at end of file +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. \ No newline at end of file