From 1ce723db22301d68337c249c6546d39797042032 Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Thu, 22 Aug 2024 21:50:16 +0000 Subject: [PATCH] have an option for truncating negative values --- src/timesfm/timesfm.py | 334 ++++++++++++++++++----------------------- 1 file changed, 147 insertions(+), 187 deletions(-) diff --git a/src/timesfm/timesfm.py b/src/timesfm/timesfm.py index f1f8b79..3e7f6ff 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 collections @@ -60,23 +59,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 @@ -202,9 +194,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,), ), ), ) @@ -222,34 +212,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. @@ -269,8 +263,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, @@ -284,8 +277,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}.") @@ -298,8 +290,7 @@ 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): @@ -335,34 +326,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 @@ -383,24 +376,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) @@ -426,6 +416,7 @@ class TimesFm: window_size: int | None = None, forecast_context_len: int | None = None, return_forecast_on_context: bool = False, + truncate_negative: bool = False, ) -> tuple[JTensor, JTensor]: """Forecasts on a list of time series. @@ -440,6 +431,8 @@ class TimesFm: 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. + truncate_negative: truncate to only non-negative values if all the contexts + have non-negative values. Returns: A tuple for JTensors: @@ -453,8 +446,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] @@ -476,50 +468,48 @@ 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) 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 = 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) @@ -535,7 +525,7 @@ class TimesFm: if window_size is not None: mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...] full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...] - if inp_min >= 0: + if inp_min >= 0 and truncate_negative: mean_outputs = np.maximum(mean_outputs, 0.0) full_outputs = np.maximum(full_outputs, 0.0) return mean_outputs, full_outputs @@ -543,16 +533,13 @@ class TimesFm: 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, + 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, + 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, @@ -594,17 +581,12 @@ class TimesFm: """ # Verify and bookkeep covariates. - if not ( - dynamic_numerical_covariates - or dynamic_categorical_covariates - or static_numerical_covariates - or static_categorical_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." - ) + " 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. @@ -624,20 +606,18 @@ class TimesFm: if dynamic_numerical_covariates: test_lens.append( - len(list(dynamic_numerical_covariates.values())[0][i]) - input_len - ) + 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 - ) + 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}." - ) + f"supports: {test_lens[-1]} vs {self.horizon_len}.") # Prepare the covariates into train and test. train_dynamic_numerical_covariates = collections.defaultdict(list) @@ -660,11 +640,9 @@ class TimesFm: continue for covariate_name, covariate_values in covariates.items(): for input_len, train_len, covariate_value in zip( - input_lens, train_lens, covariate_values - ): + input_lens, train_lens, covariate_values): train_covariates[covariate_name].append( - covariate_value[(input_len - train_len) : input_len] - ) + covariate_value[(input_len - train_len):input_len]) test_covariates[covariate_name].append(covariate_value[input_len:]) # Fit models. @@ -678,15 +656,10 @@ class TimesFm: 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 - ) + (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: @@ -697,8 +670,10 @@ class TimesFm: 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, + 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( @@ -713,12 +688,8 @@ class TimesFm: if normalize_xreg_target_per_input: xregs = _renormalize(xregs, per_instance_stats) outputs = [ - ( - mean_output[ - self._horizon_start : (self._horizon_start + test_len) - ] - + xreg - ) + (mean_output[self._horizon_start:(self._horizon_start + test_len)] + + xreg) for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs) ] @@ -737,8 +708,10 @@ class TimesFm: 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, + 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( @@ -761,12 +734,8 @@ class TimesFm: return_forecast_on_context=True, ) outputs = [ - ( - mean_output[ - self._horizon_start : (self._horizon_start + test_len) - ] - + xreg - ) + (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: @@ -806,14 +775,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.") @@ -840,36 +805,31 @@ 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) if verbose: - print("Finished preprocessing dataframe.") + 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) if verbose: - print("Finished forecasting.") + print("Finished forecasting.") fcst_df = make_future_dataframe( uids=uids, last_times=df_sorted.groupby("unique_id")["ds"].tail(1), 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.")