From 92cb49cab6143e04d3f3179de7f8451152289ae2 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 19:47:31 +0100 Subject: [PATCH 1/9] Test structural models against statsmodels --- .../statespace/core/statespace.py | 96 +++++- .../statespace/filters/distributions.py | 39 ++- .../statespace/models/structural.py | 53 +-- .../statespace/test_data/airpassangers.csv | 145 ++++++++ .../tests/statespace/test_structural.py | 314 +++++++++++++++--- .../statespace/utilities/test_helpers.py | 7 +- 6 files changed, 572 insertions(+), 82 deletions(-) create mode 100644 pymc_experimental/tests/statespace/test_data/airpassangers.csv diff --git a/pymc_experimental/statespace/core/statespace.py b/pymc_experimental/statespace/core/statespace.py index ba40c715..a68a858f 100644 --- a/pymc_experimental/statespace/core/statespace.py +++ b/pymc_experimental/statespace/core/statespace.py @@ -23,6 +23,7 @@ ) from pymc_experimental.statespace.filters.distributions import ( LinearGaussianStateSpace, + LinearGaussianStateSpaceRV, SequenceMvNormal, ) from pymc_experimental.statespace.filters.utilities import stabilize @@ -636,6 +637,37 @@ def _insert_random_variables(self) -> List[Variable]: } self.subbed_ssm = graph_replace(matrices, replace=replacement_dict, strict=True) + def _register_matrices_with_pymc_model(self) -> List[pt.TensorVariable]: + """ + Add all statespace matrices to the PyMC model currently on the context stack as pm.Deterministic nodes, and + adds named dimensions if they are found. + + Returns + ------- + registered_matrices: list of pt.TensorVariable + List of statespace matrices, wrapped in pm.Deterministic + """ + + pm_mod = modelcontext(None) + matrices = self.unpack_statespace() + + registered_matrices = [] + for i, (matrix, name) in enumerate(zip(matrices, MATRIX_NAMES)): + time_varying_ndim = 2 if name in VECTOR_VALUED else 3 + if not getattr(pm_mod, name, None): + shape, dims = self._get_matrix_shape_and_dims(name) + has_dims = dims is not None + + if matrix.ndim == time_varying_ndim and has_dims: + dims = (TIME_DIM,) + dims + + x = pm.Deterministic(name, matrix, dims=dims) + registered_matrices.append(x) + else: + registered_matrices.append(matrices[i]) + + return registered_matrices + def add_exogenous(self, exog: pt.TensorVariable) -> None: """ Add an exogenous process to the statespace model @@ -746,7 +778,6 @@ def build_statespace_graph( pm_mod = modelcontext(None) self._insert_random_variables() - matrices = self.unpack_statespace() obs_coords = pm_mod.coords.get(OBS_STATE_DIM, None) self.data_len = data.shape[0] @@ -758,20 +789,7 @@ def build_statespace_graph( missing_fill_value=missing_fill_value, ) - registered_matrices = [] - for i, (matrix, name) in enumerate(zip(matrices, MATRIX_NAMES)): - time_varying_ndim = 2 if name in VECTOR_VALUED else 3 - if not getattr(pm_mod, name, None): - shape, dims = self._get_matrix_shape_and_dims(name) - has_dims = dims is not None - - if matrix.ndim == time_varying_ndim and has_dims: - dims = (TIME_DIM,) + dims - - x = pm.Deterministic(name, matrix, dims=dims) - registered_matrices.append(x) - else: - registered_matrices.append(matrices[i]) + registered_matrices = self._register_matrices_with_pymc_model() filter_outputs = self.kalman_filter.build_graph( pt.as_tensor_variable(data), @@ -833,6 +851,54 @@ def build_statespace_graph( if return_updates: return updates + def build_as_prior( + self, + name: str, + n_steps: int, + register_statespace_matrices: Optional[bool] = False, + k_endog: Optional[int] = None, + mode: Optional[str] = None, + **kwargs, + ) -> LinearGaussianStateSpaceRV: + """ + Construct a random variable over + + Parameters + ---------- + name: str + Name of the random variable + n_steps: int + Length of the prior sequence, in time steps + register_statespace_matrices: bool + If True, statespace matrices used to construct the prior sequence will be wrapped in pm.Deterministic and + saved to the active model. + k_endog: int, optional + The number of observed states in the statespace. For debugging purposes only, should be automatically + inferred in most cases. + mode: str, optional + Pytensor compile mode. + kwargs: + PyMC distribution arguments, passed to LinearGaussianStateSpace + + Returns + ------- + obs_states: LinearGaussianStateSpaceRV + A random variable representing realizations of the observation equations of the statespace. + """ + + pm_mod = modelcontext(None) + self._insert_random_variables() + + if register_statespace_matrices: + matrices = self._register_matrices_with_pymc_model() + else: + matrices = self.unpack_statespace() + + latent_states, obs_states = LinearGaussianStateSpace( + name, *matrices, steps=n_steps, k_endog=k_endog, mode=mode, **kwargs + ) + return obs_states + def _build_smoother_graph( self, filtered_states: pt.TensorVariable, diff --git a/pymc_experimental/statespace/filters/distributions.py b/pymc_experimental/statespace/filters/distributions.py index ecc497be..45eedd90 100644 --- a/pymc_experimental/statespace/filters/distributions.py +++ b/pymc_experimental/statespace/filters/distributions.py @@ -10,6 +10,7 @@ from pytensor.graph.basic import Node floatX = pytensor.config.floatX +COV_ZERO_TOL = 0 lgss_shape_message = ( "The LinearGaussianStateSpace distribution needs shape information to be constructed. " @@ -157,8 +158,11 @@ def step_fn(*args): middle_rng, a_innovation = pm.MvNormal.dist(mu=0, cov=Q, rng=rng).owner.outputs next_rng, y_innovation = pm.MvNormal.dist(mu=0, cov=H, rng=middle_rng).owner.outputs - a_next = c + T @ a + R @ a_innovation - y_next = d + Z @ a_next + y_innovation + a_mu = c + T @ a + a_next = pt.switch(pt.all(pt.le(Q, COV_ZERO_TOL)), a_mu, a_mu + R @ a_innovation) + + y_mu = d + Z @ a_next + y_next = pt.switch(pt.all(pt.le(H, COV_ZERO_TOL)), y_mu, y_mu + y_innovation) next_state = pt.concatenate([a_next, y_next], axis=0) @@ -168,7 +172,11 @@ def step_fn(*args): Z_init = Z_ if Z_ in non_sequences else Z_[0] H_init = H_ if H_ in non_sequences else H_[0] - init_y_ = pm.MvNormal.dist(Z_init @ init_x_, H_init, rng=rng) + init_y_ = pt.switch( + pt.all(pt.le(H_init, COV_ZERO_TOL)), + Z_init @ init_x_, + pm.MvNormal.dist(Z_init @ init_x_, H_init, rng=rng), + ) init_dist_ = pt.concatenate([init_x_, init_y_], axis=0) statespace, updates = pytensor.scan( @@ -216,6 +224,7 @@ def __new__( steps=None, mode=None, sequence_names=None, + k_endog=None, **kwargs, ): dims = kwargs.pop("dims", None) @@ -239,11 +248,29 @@ def __new__( sequence_names=sequence_names, **kwargs, ) - k_states = T.type.shape[0] - latent_states = latent_obs_combined[..., :k_states] - obs_states = latent_obs_combined[..., k_states:] + if k_endog is None and k_states is None: + raise ValueError("Could not infer number of observed states, explicitly pass k_endog.") + if k_endog is not None and k_states is not None: + total_shape = latent_obs_combined.type.shape[-1] + inferred_endog = total_shape - k_states + if inferred_endog != k_endog: + raise ValueError( + f"Inferred k_endog does not agree with provided value ({inferred_endog} != {k_endog}). " + f"It is not necessary to provide k_endog when the value can be inferred." + ) + latent_slice = slice(None, -k_endog) + obs_slice = slice(-k_endog, None) + elif k_endog is None: + latent_slice = slice(None, k_states) + obs_slice = slice(k_states, None) + else: + latent_slice = slice(None, -k_endog) + obs_slice = slice(-k_endog, None) + + latent_states = latent_obs_combined[..., latent_slice] + obs_states = latent_obs_combined[..., obs_slice] latent_states = pm.Deterministic(f"{name}_latent", latent_states, dims=latent_dims) obs_states = pm.Deterministic(f"{name}_observed", obs_states, dims=obs_dims) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index 3ee926c6..72bc56f6 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -1,7 +1,7 @@ import functools as ft import logging from abc import ABC -from typing import Any, Dict, List, Optional, Sequence +from typing import Any, Dict, List, Optional, Sequence, Union import numpy as np import pytensor @@ -737,13 +737,25 @@ class LevelTrendComponent(Component): """ def __init__( - self, order: int = 2, innovations_order: Optional[int] = None, name: str = "LevelTrend" + self, + order: Union[int | list[int]] = 2, + innovations_order: Optional[Union[int, list[int]]] = None, + name: str = "LevelTrend", ): if innovations_order is None: innovations_order = order - order = order_to_mask(order) - k_states = int(sum(order)) + self._order_mask = order_to_mask(order) + max_state = np.flatnonzero(self._order_mask)[-1].item() + 1 + + # If the user passes excess zeros, raise an error. The alternative is to prune them, but this would cause + # the shape of the state to be different to what the user expects. + if len(self._order_mask) > max_state: + raise ValueError( + f"order={order} is invalid. The highest derivative should not be set to zero. If you want a " + f"lower order model, explicitly omit the zeros." + ) + k_states = max_state if isinstance(innovations_order, int): n = innovations_order @@ -755,29 +767,32 @@ def __init__( else: innovations_order = order_to_mask(innovations_order) - self.innovations_order = innovations_order + self.innovations_order = innovations_order[:max_state] k_posdef = int(sum(innovations_order)) super().__init__( name, - 1, - k_states, - k_posdef, + k_endog=1, + k_states=k_states, + k_posdef=k_posdef, measurement_error=False, combine_hidden_states=False, obs_state_idxs=np.array([1.0] + [0.0] * (k_states - 1)), ) def populate_component_properties(self): + name_slice = POSITION_DERIVATIVE_NAMES[: self.k_states] self.param_names = ["initial_trend"] - self.state_names = POSITION_DERIVATIVE_NAMES[: self.k_states] + self.state_names = [name for name, mask in zip(name_slice, self._order_mask) if mask] self.param_dims = {"initial_trend": ("trend_state",)} self.coords = {"trend_state": self.state_names} self.param_info = {"initial_trend": {"shape": (self.k_states,), "constraints": "None"}} if self.k_posdef > 0: self.param_names += ["sigma_trend"] - self.shock_names = list(np.array(self.state_names)[self.innovations_order]) + self.shock_names = [ + name for name, mask in zip(name_slice, self.innovations_order) if mask + ] self.param_dims["sigma_trend"] = ("trend_shock",) self.coords["trend_shock"] = self.shock_names self.param_info["sigma_trend"] = {"shape": (self.k_posdef,), "constraints": "Positive"} @@ -1106,7 +1121,7 @@ def __init__( name=name, k_endog=1, k_states=k_states, - k_posdef=1, # TODO: Why not int(self.innovation)? + k_posdef=int(self.innovations), measurement_error=False, combine_hidden_states=True, obs_state_idxs=np.r_[[1.0], np.zeros(k_states - 1)], @@ -1227,7 +1242,7 @@ def __init__(self, season_length, n=None, name=None, innovations=True): name=name, k_endog=1, k_states=k_states, - k_posdef=k_states, + k_posdef=k_states * int(self.innovations), measurement_error=False, combine_hidden_states=True, obs_state_idxs=obs_state_idx, @@ -1235,7 +1250,6 @@ def __init__(self, season_length, n=None, name=None, innovations=True): def make_symbolic_graph(self) -> None: self.ssm["design", 0, slice(0, self.k_states, 2)] = 1 - self.ssm["selection", :, :] = np.eye(self.k_states) init_state = self.make_and_register_variable(f"{self.name}", shape=(self.n_coefs,)) @@ -1249,11 +1263,11 @@ def make_symbolic_graph(self) -> None: if self.innovations: sigma_season = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,)) self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season + self.ssm["selection", :, :] = np.eye(self.k_states) def populate_component_properties(self): self.state_names = [f"{self.name}_{f}_{i}" for i in range(self.n) for f in ["Cos", "Sin"]] self.param_names = [f"{self.name}"] - self.shock_names = self.state_names.copy() self.param_dims = {self.name: (f"{self.name}_initial_state",)} self.param_info = { @@ -1270,6 +1284,7 @@ def populate_component_properties(self): self.coords = {f"{self.name}_initial_state": [self.state_names[i] for i in init_state_idx]} if self.innovations: + self.shock_names = self.state_names.copy() self.param_names += [f"sigma_{self.name}"] self.param_info[f"sigma_{self.name}"] = { "shape": (1,), @@ -1296,7 +1311,7 @@ def __init__( if cycle_length is not None and estimate_cycle_length: raise ValueError("Cannot specify cycle_length if estimate_cycle_length is True") if name is None: - cycle = cycle_length if cycle_length is not None else "Estimate" + cycle = int(cycle_length) if cycle_length is not None else "Estimate" name = f"Cycle[s={cycle}, dampen={dampen}, innovations={innovations}]" self.estimate_cycle_length = estimate_cycle_length @@ -1326,7 +1341,7 @@ def make_symbolic_graph(self) -> None: self.ssm["design", 0, slice(0, self.k_states, 2)] = 1 self.ssm["selection", :, :] = np.eye(self.k_states) - init_state = self.make_and_register_variable(f"{self.name}", shape=(1,)) + init_state = self.make_and_register_variable(f"initial_{self.name}", shape=(1,)) self.ssm["initial_state", 0] = init_state @@ -1340,7 +1355,7 @@ def make_symbolic_graph(self) -> None: else: rho = 1 - T = rho * _frequency_transition_block(s=lamb, j=1) + T = rho * _frequency_transition_block(lamb, j=1) self.ssm["transition", :, :] = T if self.innovations: @@ -1348,8 +1363,8 @@ def make_symbolic_graph(self) -> None: self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season def populate_component_properties(self): - self.state_names = [f"{self.name}_{f}" for f in ["Cos", "Sin"]] - self.param_names = [f"{self.name}"] + self.state_names = [f"{self.name}_{f}" for f in ["Sin", "Cos"]] + self.param_names = [f"initial_{self.name}"] self.param_dims = {self.name: (f"{self.name}_initial_state",)} self.param_info = { diff --git a/pymc_experimental/tests/statespace/test_data/airpassangers.csv b/pymc_experimental/tests/statespace/test_data/airpassangers.csv new file mode 100644 index 00000000..7014d867 --- /dev/null +++ b/pymc_experimental/tests/statespace/test_data/airpassangers.csv @@ -0,0 +1,145 @@ +Month,#Passengers +1949-01,112 +1949-02,118 +1949-03,132 +1949-04,129 +1949-05,121 +1949-06,135 +1949-07,148 +1949-08,148 +1949-09,136 +1949-10,119 +1949-11,104 +1949-12,118 +1950-01,115 +1950-02,126 +1950-03,141 +1950-04,135 +1950-05,125 +1950-06,149 +1950-07,170 +1950-08,170 +1950-09,158 +1950-10,133 +1950-11,114 +1950-12,140 +1951-01,145 +1951-02,150 +1951-03,178 +1951-04,163 +1951-05,172 +1951-06,178 +1951-07,199 +1951-08,199 +1951-09,184 +1951-10,162 +1951-11,146 +1951-12,166 +1952-01,171 +1952-02,180 +1952-03,193 +1952-04,181 +1952-05,183 +1952-06,218 +1952-07,230 +1952-08,242 +1952-09,209 +1952-10,191 +1952-11,172 +1952-12,194 +1953-01,196 +1953-02,196 +1953-03,236 +1953-04,235 +1953-05,229 +1953-06,243 +1953-07,264 +1953-08,272 +1953-09,237 +1953-10,211 +1953-11,180 +1953-12,201 +1954-01,204 +1954-02,188 +1954-03,235 +1954-04,227 +1954-05,234 +1954-06,264 +1954-07,302 +1954-08,293 +1954-09,259 +1954-10,229 +1954-11,203 +1954-12,229 +1955-01,242 +1955-02,233 +1955-03,267 +1955-04,269 +1955-05,270 +1955-06,315 +1955-07,364 +1955-08,347 +1955-09,312 +1955-10,274 +1955-11,237 +1955-12,278 +1956-01,284 +1956-02,277 +1956-03,317 +1956-04,313 +1956-05,318 +1956-06,374 +1956-07,413 +1956-08,405 +1956-09,355 +1956-10,306 +1956-11,271 +1956-12,306 +1957-01,315 +1957-02,301 +1957-03,356 +1957-04,348 +1957-05,355 +1957-06,422 +1957-07,465 +1957-08,467 +1957-09,404 +1957-10,347 +1957-11,305 +1957-12,336 +1958-01,340 +1958-02,318 +1958-03,362 +1958-04,348 +1958-05,363 +1958-06,435 +1958-07,491 +1958-08,505 +1958-09,404 +1958-10,359 +1958-11,310 +1958-12,337 +1959-01,360 +1959-02,342 +1959-03,406 +1959-04,396 +1959-05,420 +1959-06,472 +1959-07,548 +1959-08,559 +1959-09,463 +1959-10,407 +1959-11,362 +1959-12,405 +1960-01,417 +1960-02,391 +1960-03,419 +1960-04,461 +1960-05,472 +1960-06,535 +1960-07,622 +1960-08,606 +1960-09,508 +1960-10,461 +1960-11,390 +1960-12,432 diff --git a/pymc_experimental/tests/statespace/test_structural.py b/pymc_experimental/tests/statespace/test_structural.py index 510a35bd..0d4dc00d 100644 --- a/pymc_experimental/tests/statespace/test_structural.py +++ b/pymc_experimental/tests/statespace/test_structural.py @@ -1,3 +1,7 @@ +import functools as ft +import warnings +from typing import Optional + import numpy as np import pandas as pd import pymc as pm @@ -24,6 +28,275 @@ RTOL = 0 if floatX.endswith("64") else 1e-6 +def _assert_all_statespace_matrices_match(mod, params, sm_mod): + x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params) + + for name, matrix in zip(["T", "R", "Z", "Q"], [T, R, Z, Q]): + long_name = SHORT_NAME_TO_LONG[name] + if np.any([x == 0 for x in matrix.shape]): + continue + assert_allclose( + sm_mod.ssm[long_name], + matrix, + err_msg=f"matrix {name} does not match statsmodels", + atol=ATOL, + rtol=RTOL, + ) + + +def create_structural_model_and_equivalent_statsmodel( + rng, + level: Optional[bool] = False, + trend: Optional[bool] = False, + seasonal: Optional[int] = None, + freq_seasonal: Optional[list[dict]] = None, + cycle: bool = False, + autoregressive: Optional[int] = None, + exog: Optional[np.ndarray] = None, + irregular: Optional[bool] = False, + stochastic_level: Optional[bool] = True, + stochastic_trend: Optional[bool] = False, + stochastic_seasonal: Optional[bool] = True, + stochastic_freq_seasonal: Optional[list[bool]] = None, + stochastic_cycle: Optional[bool] = False, + damped_cycle: Optional[bool] = False, +): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + mod = ft.partial( + sm.tsa.UnobservedComponents, + level=level, + trend=trend, + seasonal=seasonal, + freq_seasonal=freq_seasonal, + cycle=cycle, + autoregressive=autoregressive, + exog=exog, + irregular=irregular, + stochastic_level=stochastic_level, + stochastic_trend=stochastic_trend, + stochastic_seasonal=stochastic_seasonal, + stochastic_freq_seasonal=stochastic_freq_seasonal, + stochastic_cycle=stochastic_cycle, + damped_cycle=damped_cycle, + mle_regression=False, + ) + + params = {} + sm_params = {} + components = [] + + if irregular: + sigma = np.abs(rng.normal(size=(1,))) + params["sigma_irregular"] = sigma + sm_params["sigma2.irregular"] = sigma.item() + comp = st.MeasurementError("irregular") + components.append(comp) + + level_trend_order = [0, 0] + level_trend_innov_order = [0, 0] + + if level: + level_trend_order[0] = 1 + if stochastic_level: + level_trend_innov_order[0] = 1 + + if trend: + level_trend_order[1] = 1 + if stochastic_trend: + level_trend_innov_order[1] = 1 + + if level or trend: + level_value = np.where( + level_trend_order, + rng.normal( + size=2, + ), + np.zeros( + 2, + ), + ) + sigma_level_value = np.abs(rng.normal(size=(2,)))[ + np.array(level_trend_innov_order, dtype="bool") + ] + max_order = np.flatnonzero(level_value)[-1].item() + 1 + level_trend_order = level_trend_order[:max_order] + + params["initial_trend"] = level_value[:max_order] + if sum(level_trend_innov_order) > 0: + params["sigma_trend"] = sigma_level_value + + sigma_level_value = sigma_level_value.tolist() + if stochastic_level: + sigma = sigma_level_value.pop(0) + sm_params["sigma2.level"] = sigma + if stochastic_trend: + sigma = sigma_level_value.pop(0) + sm_params[f"sigma2.trend"] = sigma + + comp = st.LevelTrendComponent( + name="level", order=level_trend_order, innovations_order=level_trend_innov_order + ) + components.append(comp) + + if seasonal is not None: + params["seasonal_coefs"] = rng.normal(size=(seasonal - 1,)) + + if stochastic_seasonal: + sigma = np.abs(rng.normal(size=(1,))) + params["sigma_seasonal"] = sigma + sm_params["sigma2.seasonal"] = sigma + + comp = st.TimeSeasonality( + name="seasonal", season_length=seasonal, innovations=stochastic_seasonal + ) + components.append(comp) + + if freq_seasonal is not None: + for d, has_innov in zip(freq_seasonal, stochastic_freq_seasonal): + n = d["harmonics"] + s = d["period"] + + last_state_not_identified = s / n == 2.0 + + params[f"seasonal_{s}"] = rng.normal(size=(2 * n - int(last_state_not_identified))) + if has_innov: + sigma = np.abs(rng.normal(size=(1,))) + params[f"sigma_seasonal_{s}"] = sigma + sm_params[f"sigma2.freq_seasonal_{s}({n})"] = sigma + + comp = st.FrequencySeasonality( + name=f"seasonal_{s}", season_length=s, n=n, innovations=has_innov + ) + components.append(comp) + + if cycle: + cycle_length = pm.math.floatX(np.random.choice(np.arange(2, 12))) + + # Statsmodels takes the frequency not the cycle length, so convert it. + sm_params["frequency.cycle"] = 2.0 * np.pi / cycle_length + params["cycle_cycle_length"] = np.atleast_1d(cycle_length) + params["initial_cycle"] = np.ones((1,)) + + if stochastic_cycle: + sigma = np.abs(rng.normal(size=(1,))) + params["sigma_cycle"] = sigma + sm_params["sigma2.cycle"] = sigma + + if damped_cycle: + rho = rng.beta(1, 1, size=(1,)) + params["cycle_dampening_factor"] = rho + sm_params["damping.cycle"] = rho + + comp = st.CycleComponent( + name="cycle", + dampen=damped_cycle, + innovations=stochastic_cycle, + estimate_cycle_length=True, + ) + + components.append(comp) + + if autoregressive is not None: + ar_params = rng.normal(size=(autoregressive,)) + sigma = np.abs(rng.normal(size=(1,))) + params["ar_params"] = ar_params + params["sigma_ar"] = sigma + + sm_params["sigma2.ar"] = sigma + for i, rho in enumerate(ar_params): + sm_params[f"ar.L{i + 1}"] = rho + + comp = st.AutoregressiveComponent(name="ar", order=autoregressive) + components.append(comp) + + if exog is not None: + names = [f"x{i + 1}" for i in range(exog.shape[1])] + betas = rng.normal(size=(exog.shape[1],)) + params["beta_exog"] = betas + params["data_exog"] = exog + + for i, beta in enumerate(betas): + sm_params[f"beta.x{i + 1}"] = beta + comp = st.RegressionComponent(name="exog", state_names=names) + components.append(comp) + + st_mod = components.pop(0) + for comp in components: + st_mod += comp + return mod, st_mod, params, sm_params + + +@pytest.mark.parametrize( + "level, trend, stochastic_level, stochastic_trend, irregular", + [ + (False, False, False, False, True), + (True, True, True, True, True), + (True, True, False, True, False), + ], +) +@pytest.mark.parametrize("autoregressive", [None, 3]) +@pytest.mark.parametrize("seasonal, stochastic_seasonal", [(None, False), (12, False), (12, True)]) +@pytest.mark.parametrize( + "freq_seasonal, stochastic_freq_seasonal", + [ + (None, None), + ([{"period": 12, "harmonics": 2}], [False]), + ([{"period": 12, "harmonics": 6}], [True]), + ], +) +@pytest.mark.parametrize( + "cycle, damped_cycle, stochastic_cycle", + [(False, False, False), (True, False, True), (True, True, True)], +) +@pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.ConvergenceWarning") +@pytest.mark.filterwarnings("ignore::statsmodels.tools.sm_exceptions.SpecificationWarning") +def test_structural_model_against_statsmodels( + level, + trend, + stochastic_level, + stochastic_trend, + irregular, + autoregressive, + seasonal, + stochastic_seasonal, + freq_seasonal, + stochastic_freq_seasonal, + cycle, + damped_cycle, + stochastic_cycle, + rng, +): + f_sm_mod, mod, params, sm_params = create_structural_model_and_equivalent_statsmodel( + rng, + level=level, + trend=trend, + seasonal=seasonal, + freq_seasonal=freq_seasonal, + cycle=cycle, + damped_cycle=damped_cycle, + autoregressive=autoregressive, + irregular=irregular, + stochastic_level=stochastic_level, + stochastic_trend=stochastic_trend, + stochastic_seasonal=stochastic_seasonal, + stochastic_freq_seasonal=stochastic_freq_seasonal, + stochastic_cycle=stochastic_cycle, + ) + + data = rng.normal(size=(100,)) + sm_mod = f_sm_mod(data) + sm_mod.initialize_default() + + if len(sm_params) > 0: + param_array = np.concatenate( + [np.atleast_1d(sm_params[k]).ravel() for k in sm_mod.param_names] + ) + sm_mod.update(param_array, transformed=True) + + _assert_all_statespace_matrices_match(mod, params, sm_mod) + + def test_deterministic_constant_model(rng): mod = st.LevelTrendComponent(order=1, innovations_order=0) params = {"initial_trend": [1.0]} @@ -66,25 +339,6 @@ def test_time_seasonality(s, innovations, rng): if not innovations: assert_pattern_repeats(y, s, atol=ATOL, rtol=RTOL) - mod2 = sm.tsa.UnobservedComponents( - endog=rng.normal(size=100), - seasonal=s, - stochastic_seasonal=innovations, - # Silence a warning about no innovations when innovations = False - irregular=True, - ) - x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params) - - for name, matrix in zip(["T", "R", "Z", "Q"], [T, R, Z, Q]): - long_name = SHORT_NAME_TO_LONG[name] - assert_allclose( - mod2.ssm[long_name], - matrix, - err_msg=f"matrix {name} does not match statsmodels", - atol=ATOL, - rtol=RTOL, - ) - def get_shift_factor(s): s_str = str(s) @@ -106,26 +360,6 @@ def test_frequency_seasonality(n, s, rng): x, y = simulate_from_numpy_model(mod, rng, params, 2 * T) assert_pattern_repeats(y, T, atol=ATOL, rtol=RTOL) - init_dict = {"period": s} - if n is not None: - init_dict["harmonics"] = n - - mod2 = sm.tsa.UnobservedComponents(endog=rng.normal(size=100), freq_seasonal=[init_dict]) - mod2.initialize_default() - x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params) - - for name, matrix in zip(["T", "Z", "R", "Q"], [T, Z, R, Q]): - name = SHORT_NAME_TO_LONG[name] - assert_allclose( - mod2.ssm[name], - matrix, - err_msg=f"matrix {name} does not match statsmodels", - rtol=RTOL, - atol=ATOL, - ) - - assert mod2.initialization.constant.shape == mod.ssm["initial_state"].type.shape - @pytest.mark.parametrize("s,n", [(10, 5), pytest.param(10.2, 5, marks=pytest.mark.xfail)]) def test_state_removed_when_freq_seasonality_is_saturated(s, n): @@ -175,7 +409,7 @@ def test_cycle_component(innovations, dampen, cycle_length, estimate_cycle_lengt innovations=innovations, ) - params = {"cycle": rng.normal(size=(1,)).astype(floatX)} + params = {"initial_cycle": rng.normal(size=(1,)).astype(floatX)} if estimate_cycle_length: params["cycle_cycle_length"] = np.array([7], dtype=floatX) diff --git a/pymc_experimental/tests/statespace/utilities/test_helpers.py b/pymc_experimental/tests/statespace/utilities/test_helpers.py index e8f58d05..e72b9416 100644 --- a/pymc_experimental/tests/statespace/utilities/test_helpers.py +++ b/pymc_experimental/tests/statespace/utilities/test_helpers.py @@ -207,9 +207,12 @@ def unpack_statespace(ssm): return [ssm[SHORT_NAME_TO_LONG[x]] for x in MATRIX_NAMES] -def unpack_symbolic_matrices_with_params(mod, param_dict): +def unpack_symbolic_matrices_with_params(mod, param_dict, mode="FAST_COMPILE"): f_matrices = pytensor.function( - list(mod._name_to_variable.values()), unpack_statespace(mod.ssm), on_unused_input="ignore" + list(mod._name_to_variable.values()), + unpack_statespace(mod.ssm), + on_unused_input="raise", + mode=mode, ) x0, P0, c, d, T, Z, R, H, Q = f_matrices(**param_dict) return x0, P0, c, d, T, Z, R, H, Q From 1751091ac1cebe66325f7ea8cd23b5eacfc9e87b Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 19:58:14 +0100 Subject: [PATCH 2/9] Fix `TimeSeasonality` coords --- pymc_experimental/statespace/models/structural.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index 72bc56f6..bacadbdf 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -1137,7 +1137,7 @@ def populate_component_properties(self): } } self.param_dims = {f"{self.name}_coefs": (f"{self.name}_periods",)} - self.coords = {f"{self.name}_periods": self.state_names} + self.coords = {f"{self.name}_state": self.state_names} if self.innovations: self.param_names += [f"sigma_{self.name}"] From 1d78b62c00edd52c32bde40748e4983538e5c102 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 22:22:21 +0100 Subject: [PATCH 3/9] Add tests of coords in `test_structural.py` --- .../statespace/models/structural.py | 8 +- .../tests/statespace/test_structural.py | 199 ++++++++++-------- .../statespace/utilities/test_helpers.py | 7 +- 3 files changed, 126 insertions(+), 88 deletions(-) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index bacadbdf..cea02b4b 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -600,7 +600,7 @@ def build(self, name=None, filter_type="standard", verbose=True): Parameters ---------- name: str, optional - Name of the exogenous data being modeled. Default is "obs" + Name of the exogenous data being modeled. Default is "data" filter_type : str, optional The type of Kalman filter to use. Valid options are "standard", "univariate", "single", "cholesky", and @@ -1109,19 +1109,19 @@ def __init__( f"state_names must be a list of length season_length, got {len(state_names)}" ) state_names = state_names.copy() - self.state_names = state_names self.innovations = innovations # The first state doesn't get a coefficient, it is defined as -sum(state_coefs) # TODO: Can I stash that information in the model somewhere so users don't have to know that? - state_0 = state_names.pop(-1) + state_0 = state_names.pop(0) k_states = season_length - 1 super().__init__( name=name, k_endog=1, k_states=k_states, - k_posdef=int(self.innovations), + k_posdef=int(innovations), + state_names=state_names, measurement_error=False, combine_hidden_states=True, obs_state_idxs=np.r_[[1.0], np.zeros(k_states - 1)], diff --git a/pymc_experimental/tests/statespace/test_structural.py b/pymc_experimental/tests/statespace/test_structural.py index 0d4dc00d..b5077dba 100644 --- a/pymc_experimental/tests/statespace/test_structural.py +++ b/pymc_experimental/tests/statespace/test_structural.py @@ -44,6 +44,15 @@ def _assert_all_statespace_matrices_match(mod, params, sm_mod): ) +def _assert_basic_coords_correct(mod): + assert mod.coords["state"] == mod.state_names + assert mod.coords["state_aux"] == mod.state_names + assert mod.coords["shock"] == mod.shock_names + assert mod.coords["shock_aux"] == mod.shock_names + assert mod.coords["observed_state"] == ["data"] + assert mod.coords["observed_state_aux"] == ["data"] + + def create_structural_model_and_equivalent_statsmodel( rng, level: Optional[bool] = False, @@ -297,36 +306,55 @@ def test_structural_model_against_statsmodels( _assert_all_statespace_matrices_match(mod, params, sm_mod) -def test_deterministic_constant_model(rng): - mod = st.LevelTrendComponent(order=1, innovations_order=0) - params = {"initial_trend": [1.0]} - x, y = simulate_from_numpy_model(mod, rng, params) - - assert_allclose(y, 1, atol=ATOL, rtol=RTOL) - - -def test_deterministic_slope_model(rng): +def test_level_trend_model(rng): mod = st.LevelTrendComponent(order=2, innovations_order=0) params = {"initial_trend": [0.0, 1.0]} x, y = simulate_from_numpy_model(mod, rng, params) assert_allclose(np.diff(y), 1, atol=ATOL, rtol=RTOL) + # Check coords + mod = mod.build(verbose=False) + _assert_basic_coords_correct(mod) + assert mod.coords["trend_state"] == ["level", "trend"] + + +def test_measurement_error(rng): + mod = st.MeasurementError("obs") + st.LevelTrendComponent(order=2) + mod = mod.build(verbose=False) + + _assert_basic_coords_correct(mod) + assert "sigma_obs" in mod.param_names + @pytest.mark.parametrize("order", [1, 2, [1, 0, 1]], ids=["AR1", "AR2", "AR(1,0,1)"]) def test_autoregressive_model(order, rng): ar = st.AutoregressiveComponent(order=order) params = { - "ar_params": np.full((sum(ar.order),), 0.95, dtype=floatX), - "sigma_ar": np.array([0.1], dtype=floatX), + "ar_params": np.full((sum(ar.order),), 0.5, dtype=floatX), + "sigma_ar": np.array([0.0], dtype=floatX), } - x, y = simulate_from_numpy_model(ar, rng, params) + x, y = simulate_from_numpy_model(ar, rng, params, steps=100) + + # Check coords + ar.build(verbose=False) + _assert_basic_coords_correct(ar) + lags = np.arange(len(order) if isinstance(order, list) else order, dtype="int") + 1 + if isinstance(order, list): + lags = lags[np.flatnonzero(order)] + assert_allclose(ar.coords["ar_lags"], lags) @pytest.mark.parametrize("s", [10, 25, 50]) @pytest.mark.parametrize("innovations", [True, False]) def test_time_seasonality(s, innovations, rng): - mod = st.TimeSeasonality(season_length=s, innovations=innovations, name="season") + def random_word(rng): + return "".join(rng.choice(list("abcdefghijklmnopqrstuvwxyz")) for _ in range(5)) + + state_names = [random_word(rng) for _ in range(s)] + mod = st.TimeSeasonality( + season_length=s, innovations=innovations, name="season", state_names=state_names + ) x0 = np.zeros(mod.k_states, dtype=floatX) x0[0] = 1 @@ -339,6 +367,11 @@ def test_time_seasonality(s, innovations, rng): if not innovations: assert_pattern_repeats(y, s, atol=ATOL, rtol=RTOL) + # Check coords + mod.build(verbose=False) + _assert_basic_coords_correct(mod) + assert mod.coords["season_state"] == state_names[1:] + def get_shift_factor(s): s_str = str(s) @@ -360,75 +393,89 @@ def test_frequency_seasonality(n, s, rng): x, y = simulate_from_numpy_model(mod, rng, params, 2 * T) assert_pattern_repeats(y, T, atol=ATOL, rtol=RTOL) + # Check coords + mod.build(verbose=False) + _assert_basic_coords_correct(mod) + if n is None: + n = int(s // 2) + states = [f"season_{f}_{i}" for i in range(n) for f in ["Cos", "Sin"]] -@pytest.mark.parametrize("s,n", [(10, 5), pytest.param(10.2, 5, marks=pytest.mark.xfail)]) -def test_state_removed_when_freq_seasonality_is_saturated(s, n): - mod = st.FrequencySeasonality(season_length=s, n=n, name="test") - init_params = mod._name_to_variable["test"] - - assert init_params.type.shape[0] == (n * 2 - 1) + # Remove the last state when the model is completely saturated + if s / n == 2.0: + states.pop() + assert mod.coords["season_initial_state"] == states cycle_test_vals = zip([None, None, 3, 5, 10], [False, True, True, False, False]) -@pytest.mark.parametrize("innovations", [True, False], ids=[f"innov={x}" for x in [True, False]]) -@pytest.mark.parametrize("dampen", [True, False], ids=[f"dampen={x}" for x in [True, False]]) -@pytest.mark.parametrize( - "cycle_length, estimate_cycle_length", - list(cycle_test_vals), - ids=[f"cycle_len={a}-est_cycle_len={b}" for a, b in cycle_test_vals], -) -def test_cycle_component(innovations, dampen, cycle_length, estimate_cycle_length, rng): - if estimate_cycle_length and (cycle_length is not None): - with pytest.raises(ValueError, match="Cannot specify cycle_length"): - st.CycleComponent( - name="cycle", - cycle_length=cycle_length, - estimate_cycle_length=estimate_cycle_length, - dampen=dampen, - innovations=innovations, - ) +def test_cycle_component_deterministic(rng): + cycle = st.CycleComponent( + name="cycle", cycle_length=12, estimate_cycle_length=False, innovations=False + ) + params = {"initial_cycle": np.array([1.0], dtype=floatX)} + x, y = simulate_from_numpy_model(cycle, rng, params, steps=12 * 12) - elif not estimate_cycle_length and (cycle_length is None): - with pytest.raises(ValueError, match="Must specify cycle_length"): - st.CycleComponent( - name="cycle", - cycle_length=cycle_length, - estimate_cycle_length=estimate_cycle_length, - dampen=dampen, - innovations=innovations, - ) + assert_pattern_repeats(y, 12, atol=ATOL, rtol=RTOL) - else: - cycle = st.CycleComponent( - name="cycle", - cycle_length=cycle_length, - estimate_cycle_length=estimate_cycle_length, - dampen=dampen, - innovations=innovations, - ) - params = {"initial_cycle": rng.normal(size=(1,)).astype(floatX)} +def test_cycle_component_with_dampening(rng): + cycle = st.CycleComponent( + name="cycle", cycle_length=12, estimate_cycle_length=False, innovations=False, dampen=True + ) + params = { + "initial_cycle": np.array([10.0], dtype=floatX), + "cycle_dampening_factor": np.array([0.75], dtype=floatX), + } + x, y = simulate_from_numpy_model(cycle, rng, params, steps=100) - if estimate_cycle_length: - params["cycle_cycle_length"] = np.array([7], dtype=floatX) - cycle_len = params.get("cycle_cycle_length", cycle_length) - fit_params = {"frequency.cycle": cycle_len, "sigma2.irregular": 0.0} - if dampen: - params["cycle_dampening_factor"] = np.array([0.95], dtype=floatX) - fit_params["damping.cycle"] = 0.95 - if innovations: - params["sigma_cycle"] = np.ones((1,), dtype=floatX) - fit_params["sigma2.cycle"] = 1.0 + # Check that the cycle dampens to zero over time + assert_allclose(y[-1], 0.0, atol=ATOL, rtol=RTOL) - x, y = simulate_from_numpy_model(cycle, rng, params) - if not innovations and not estimate_cycle_length: - if dampen: - # undo the dampening so it's all constant - y = y / 0.95 ** np.arange(y.shape[0]) - assert_pattern_repeats(y, cycle_length, atol=ATOL, rtol=RTOL) +def test_cycle_component_with_innovations_and_cycle_length(rng): + cycle = st.CycleComponent( + name="cycle", estimate_cycle_length=True, innovations=True, dampen=True + ) + params = { + "initial_cycle": np.array([1.0], dtype=floatX), + "cycle_cycle_length": np.array([12], dtype=floatX), + "cycle_dampening_factor": np.array([0.95], dtype=floatX), + "sigma_cycle": np.array([1.0], dtype=floatX), + } + + x, y = simulate_from_numpy_model(cycle, rng, params) + + cycle.build(verbose=False) + _assert_basic_coords_correct(cycle) + assert cycle.coords["cycle_initial_state"] == ["cycle_Sin", "cycle_Cos"] + + +def test_exogenous_component(rng): + data = rng.normal(size=(100, 2)).astype(floatX) + mod = st.RegressionComponent(state_names=["feature_1", "feature_2"], name="exog") + + params = {"beta_exog": np.array([1.0, 2.0], dtype=floatX), "data_exog": data} + x, y = simulate_from_numpy_model(mod, rng, params) + + # Check that the generated data is just a linear regression + assert_allclose(y, data @ params["beta_exog"], atol=ATOL, rtol=RTOL) + + mod.build(verbose=False) + _assert_basic_coords_correct(mod) + assert mod.coords["exog_state"] == ["feature_1", "feature_2"] + + +def test_adding_exogenous_component(rng): + data = rng.normal(size=(100, 2)).astype(floatX) + reg = st.RegressionComponent(state_names=["a", "b"], name="exog") + ll = st.LevelTrendComponent(name="level") + + seasonal = st.FrequencySeasonality(name="annual", season_length=12, n=4) + mod = reg + ll + seasonal + + assert mod.ssm["design"].eval({"data_exog": data}).shape == (100, 1, 2 + 2 + 8) + assert_allclose(mod.ssm["design", 5, 0, :2].eval({"data_exog": data}), data[5]) def test_add_components(): @@ -475,18 +522,6 @@ def test_add_components(): assert_allclose(all_mat, np.concatenate([ll_mat, se_mat], axis=axis), atol=ATOL, rtol=RTOL) -def test_adding_exogenous_component(rng): - data = rng.normal(size=(100, 2)).astype(floatX) - reg = st.RegressionComponent(state_names=["a", "b"], name="exog") - ll = st.LevelTrendComponent(name="level") - - seasonal = st.FrequencySeasonality(name="annual", season_length=12, n=4) - mod = reg + ll + seasonal - - assert mod.ssm["design"].eval({"data_exog": data}).shape == (100, 1, 2 + 2 + 8) - assert_allclose(mod.ssm["design", 5, 0, :2].eval({"data_exog": data}), data[5]) - - def test_filter_scans_time_varying_design_matrix(rng): time_idx = pd.date_range(start="2000-01-01", freq="D", periods=100) data = pd.DataFrame(rng.normal(size=(100, 2)), columns=["a", "b"], index=time_idx) diff --git a/pymc_experimental/tests/statespace/utilities/test_helpers.py b/pymc_experimental/tests/statespace/utilities/test_helpers.py index e72b9416..1ff9fd29 100644 --- a/pymc_experimental/tests/statespace/utilities/test_helpers.py +++ b/pymc_experimental/tests/statespace/utilities/test_helpers.py @@ -230,7 +230,7 @@ def simulate_from_numpy_model(mod, rng, param_dict, steps=100): y = np.zeros(steps) x[0] = x0 - y[0] = (Z @ x0).squeeze() + y[0] = (Z @ x0).squeeze() if Z.ndim == 2 else (Z[0] @ x0).squeeze() if not np.allclose(H, 0): y[0] += rng.multivariate_normal(mean=np.zeros(1), cov=H) @@ -248,7 +248,10 @@ def simulate_from_numpy_model(mod, rng, param_dict, steps=100): error = 0 x[t] = c + T @ x[t - 1] + innov - y[t] = (d + Z @ x[t] + error).squeeze() + if Z.ndim == 2: + y[t] = (d + Z @ x[t] + error).squeeze() + else: + y[t] = (d + Z[t] @ x[t] + error).squeeze() return x, y From 6cb4cdd0da1aebe704a1d07f92ebbb0c613c40fd Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 22:59:40 +0100 Subject: [PATCH 4/9] Fix coords generated for cycle component Verify coord lengths against matrix shapes in `test_structural.py` --- .../statespace/models/structural.py | 10 +++--- .../tests/statespace/test_structural.py | 32 ++++++++++++++++--- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index cea02b4b..3fdc0951 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -1341,7 +1341,7 @@ def make_symbolic_graph(self) -> None: self.ssm["design", 0, slice(0, self.k_states, 2)] = 1 self.ssm["selection", :, :] = np.eye(self.k_states) - init_state = self.make_and_register_variable(f"initial_{self.name}", shape=(1,)) + init_state = self.make_and_register_variable(f"{self.name}", shape=(1,)) self.ssm["initial_state", 0] = init_state @@ -1364,17 +1364,15 @@ def make_symbolic_graph(self) -> None: def populate_component_properties(self): self.state_names = [f"{self.name}_{f}" for f in ["Sin", "Cos"]] - self.param_names = [f"initial_{self.name}"] + self.param_names = [f"{self.name}"] - self.param_dims = {self.name: (f"{self.name}_initial_state",)} self.param_info = { f"{self.name}": { "shape": (1,), "constraints": "None", - "dims": f"({self.name}_initial_state, )", + "dims": None, } } - self.coords = {f"{self.name}_initial_state": self.state_names} if self.estimate_cycle_length: self.param_names += [f"{self.name}_cycle_length"] @@ -1399,7 +1397,7 @@ def populate_component_properties(self): "constraints": "Positive", "dims": "None", } - self.shock_names = [f"{self.name}"] + self.shock_names = self.state_names.copy() class RegressionComponent(Component): diff --git a/pymc_experimental/tests/statespace/test_structural.py b/pymc_experimental/tests/statespace/test_structural.py index b5077dba..60dcba3a 100644 --- a/pymc_experimental/tests/statespace/test_structural.py +++ b/pymc_experimental/tests/statespace/test_structural.py @@ -44,6 +44,26 @@ def _assert_all_statespace_matrices_match(mod, params, sm_mod): ) +def _assert_coord_shapes_match_matrices(mod, params): + if "initial_state_cov" not in params: + params["initial_state_cov"] = np.eye(mod.k_states) + + x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params) + n_states = len(mod.coords["state"]) + n_shocks = len(mod.coords["shock"]) + n_obs = len(mod.coords["observed_state"]) + + assert x0.shape[-1:] == (n_states,) + assert P0.shape[-2:] == (n_states, n_states) + assert c.shape[-1:] == (n_states,) + assert d.shape[-1:] == (n_obs,) + assert T.shape[-2:] == (n_states, n_states) + assert Z.shape[-2:] == (n_obs, n_states) + assert R.shape[-2:] == (n_states, n_shocks) + assert H.shape[-2:] == (n_obs, n_obs) + assert Q.shape[-2:] == (n_shocks, n_shocks) + + def _assert_basic_coords_correct(mod): assert mod.coords["state"] == mod.state_names assert mod.coords["state_aux"] == mod.state_names @@ -185,7 +205,7 @@ def create_structural_model_and_equivalent_statsmodel( # Statsmodels takes the frequency not the cycle length, so convert it. sm_params["frequency.cycle"] = 2.0 * np.pi / cycle_length params["cycle_cycle_length"] = np.atleast_1d(cycle_length) - params["initial_cycle"] = np.ones((1,)) + params["cycle"] = np.ones((1,)) if stochastic_cycle: sigma = np.abs(rng.normal(size=(1,))) @@ -305,6 +325,9 @@ def test_structural_model_against_statsmodels( _assert_all_statespace_matrices_match(mod, params, sm_mod) + mod.build(verbose=False) + _assert_coord_shapes_match_matrices(mod, params) + def test_level_trend_model(rng): mod = st.LevelTrendComponent(order=2, innovations_order=0) @@ -413,7 +436,7 @@ def test_cycle_component_deterministic(rng): cycle = st.CycleComponent( name="cycle", cycle_length=12, estimate_cycle_length=False, innovations=False ) - params = {"initial_cycle": np.array([1.0], dtype=floatX)} + params = {"cycle": np.array([1.0], dtype=floatX)} x, y = simulate_from_numpy_model(cycle, rng, params, steps=12 * 12) assert_pattern_repeats(y, 12, atol=ATOL, rtol=RTOL) @@ -424,7 +447,7 @@ def test_cycle_component_with_dampening(rng): name="cycle", cycle_length=12, estimate_cycle_length=False, innovations=False, dampen=True ) params = { - "initial_cycle": np.array([10.0], dtype=floatX), + "cycle": np.array([10.0], dtype=floatX), "cycle_dampening_factor": np.array([0.75], dtype=floatX), } x, y = simulate_from_numpy_model(cycle, rng, params, steps=100) @@ -438,7 +461,7 @@ def test_cycle_component_with_innovations_and_cycle_length(rng): name="cycle", estimate_cycle_length=True, innovations=True, dampen=True ) params = { - "initial_cycle": np.array([1.0], dtype=floatX), + "cycle": np.array([1.0], dtype=floatX), "cycle_cycle_length": np.array([12], dtype=floatX), "cycle_dampening_factor": np.array([0.95], dtype=floatX), "sigma_cycle": np.array([1.0], dtype=floatX), @@ -448,7 +471,6 @@ def test_cycle_component_with_innovations_and_cycle_length(rng): cycle.build(verbose=False) _assert_basic_coords_correct(cycle) - assert cycle.coords["cycle_initial_state"] == ["cycle_Sin", "cycle_Cos"] def test_exogenous_component(rng): From cb9b700ea874ec407fb79560f780e1c15f4af403 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 23:24:37 +0100 Subject: [PATCH 5/9] Write `CycleComponent` docstring --- .../statespace/models/structural.py | 105 ++++++++++++++++-- .../tests/statespace/test_structural.py | 38 +++---- 2 files changed, 112 insertions(+), 31 deletions(-) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index 3fdc0951..a31b6022 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -1053,9 +1053,10 @@ class TimeSeasonality(Component): And so on. So for interpretation, the ``season_length - 1`` initial states are, when reversed, the coefficients associated with ``state_names[1:]``. - .. warning:: Although the ``season_names`` argument expects a list of length ``season_length``, only - ``season_names[1:]`` will be saved as model dimensions, since the 1st coefficient is not estimated (it is the sum - of the other 11). + .. warning:: + Although the ``state_names`` argument expects a list of length ``season_length``, only ``state_names[1:]`` + will be saved as model dimensions, since the 1st coefficient is not identified (it is defined as + :math:`-\sum_{i=1}^{s} \gamma_{t-i}`). Examples -------- @@ -1295,16 +1296,96 @@ def populate_component_properties(self): class CycleComponent(Component): r""" - # TODO: WRITEME + A component for modeling longer-term cyclical effects + + Parameters + ---------- + name: str + Name of the component. Used in generated coordinates and state names. If None, a descriptive name will be + used. + + cycle_length: int, optional + The length of the cycle, in the calendar units of your data. For example, if your data is monthly, and you + want to model a 12-month cycle, use ``cycle_length=12``. You cannot specify both ``cycle_length`` and + ``estimate_cycle_length``. + + estimate_cycle_length: bool, default False + Whether to estimate the cycle length. If True, an additional parameter, ``cycle_length`` will be added to the + model. You cannot specify both ``cycle_length`` and ``estimate_cycle_length``. + + dampen: bool, default False + Whether to dampen the cycle by multiplying by a dampening factor :math:`\rho` at every timestep. If true, + an additional parameter, ``dampening_factor`` will be added to the model. + + innovations: bool, default True + Whether to include stochastic innovations in the strength of the seasonal effect. If True, an additional + parameter, ``sigma_{name}`` will be added to the model. + + Notes + ----- + The cycle component is very similar in implementation to the frequency domain seasonal component, expect that it + is restricted to n=1. The cycle component can be expressed: + + .. math:: + \begin{align} + \gamma_t &= \rho \gamma_{t-1} \cos \lambda + \rho \gamma_{t-1}^\star \sin \lambda + \omega_{t} \\ + \gamma_{t}^\star &= -\rho \gamma_{t-1} \sin \lambda + \rho \gamma_{t-1}^\star \cos \lambda + \omega_{t}^\star + \lambda &= \frac{2\pi}{s} + \end{align} + + Where :math:`s` is the ``cycle_length``. [1] recommend that this component be used for longer term cyclical + effects, such as business cycles, and that the seasonal component be used for shorter term effects, such as + weekly or monthly seasonality. + + Unlike a FrequencySeasonality component, the length of a CycleComponent can be estimated. + + Examples + -------- + Estimate a business cycle with length between 6 and 12 years: + + .. code:: python + + from pymc_experimental.statespace import structural as st + import pymc as pm + import pytensor.tensor as pt + import pandas as pd + import numpy as np + + data = np.random.normal(size=(100, 1)) + + # Build the structural model + grw = st.LevelTrendComponent(order=1, innovations_order=1) + cycle = st.CycleComponent('business_cycle', estimate_cycle_length=True, dampen=False) + ss_mod = (grw + cycle).build() + + # Estimate with PyMC + with pm.Model(coords=ss_mod.coords) as model: + P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states), dims=ss_mod.param_dims['P0']) + intitial_trend = pm.Normal('initial_trend', dims=ss_mod.param_dims['initial_trend']) + sigma_trend = pm.HalfNormal('sigma_trend', dims=ss_mod.param_dims['sigma_trend']) + + cycle_strength = pm.Normal('business_cycle') + cycle_length = pm.Uniform('business_cycle_length', lower=6, upper=12) + + sigma_cycle = pm.HalfNormal('sigma_business_cycle', sigma=1) + ss_mod.build_statespace_graph(data, mode='JAX') + + idata = pm.sample(nuts_sampler='numpyro') + + References + ---------- + .. [1] Durbin, James, and Siem Jan Koopman. 2012. + Time Series Analysis by State Space Methods: Second Edition. + Oxford University Press. """ def __init__( self, - name=None, - cycle_length=None, - estimate_cycle_length=False, - dampen=False, - innovations=True, + name: str = None, + cycle_length: int = None, + estimate_cycle_length: bool = False, + dampen: bool = False, + innovations: bool = True, ): if cycle_length is None and not estimate_cycle_length: raise ValueError("Must specify cycle_length if estimate_cycle_length is False") @@ -1346,7 +1427,7 @@ def make_symbolic_graph(self) -> None: self.ssm["initial_state", 0] = init_state if self.estimate_cycle_length: - lamb = self.make_and_register_variable(f"{self.name}_cycle_length", shape=(1,)) + lamb = self.make_and_register_variable(f"{self.name}_length", shape=(1,)) else: lamb = self.cycle_length @@ -1375,8 +1456,8 @@ def populate_component_properties(self): } if self.estimate_cycle_length: - self.param_names += [f"{self.name}_cycle_length"] - self.param_info[f"{self.name}_cycle_length"] = { + self.param_names += [f"{self.name}_length"] + self.param_info[f"{self.name}_length"] = { "shape": (1,), "constraints": "Positive, non-zero", "dims": None, diff --git a/pymc_experimental/tests/statespace/test_structural.py b/pymc_experimental/tests/statespace/test_structural.py index 60dcba3a..4002d5aa 100644 --- a/pymc_experimental/tests/statespace/test_structural.py +++ b/pymc_experimental/tests/statespace/test_structural.py @@ -116,7 +116,7 @@ def create_structural_model_and_equivalent_statsmodel( components = [] if irregular: - sigma = np.abs(rng.normal(size=(1,))) + sigma = np.abs(rng.normal(size=(1,))).astype(floatX) params["sigma_irregular"] = sigma sm_params["sigma2.irregular"] = sigma.item() comp = st.MeasurementError("irregular") @@ -140,10 +140,8 @@ def create_structural_model_and_equivalent_statsmodel( level_trend_order, rng.normal( size=2, - ), - np.zeros( - 2, - ), + ).astype(floatX), + np.zeros(2, dtype=floatX), ) sigma_level_value = np.abs(rng.normal(size=(2,)))[ np.array(level_trend_innov_order, dtype="bool") @@ -169,10 +167,10 @@ def create_structural_model_and_equivalent_statsmodel( components.append(comp) if seasonal is not None: - params["seasonal_coefs"] = rng.normal(size=(seasonal - 1,)) + params["seasonal_coefs"] = rng.normal(size=(seasonal - 1,)).astype(floatX) if stochastic_seasonal: - sigma = np.abs(rng.normal(size=(1,))) + sigma = np.abs(rng.normal(size=(1,))).astype(floatX) params["sigma_seasonal"] = sigma sm_params["sigma2.seasonal"] = sigma @@ -188,9 +186,11 @@ def create_structural_model_and_equivalent_statsmodel( last_state_not_identified = s / n == 2.0 - params[f"seasonal_{s}"] = rng.normal(size=(2 * n - int(last_state_not_identified))) + params[f"seasonal_{s}"] = rng.normal( + size=(2 * n - int(last_state_not_identified)) + ).astype(floatX) if has_innov: - sigma = np.abs(rng.normal(size=(1,))) + sigma = np.abs(rng.normal(size=(1,))).astype(floatX) params[f"sigma_seasonal_{s}"] = sigma sm_params[f"sigma2.freq_seasonal_{s}({n})"] = sigma @@ -200,20 +200,20 @@ def create_structural_model_and_equivalent_statsmodel( components.append(comp) if cycle: - cycle_length = pm.math.floatX(np.random.choice(np.arange(2, 12))) + cycle_length = np.random.choice(np.arange(2, 12)).astype(floatX) # Statsmodels takes the frequency not the cycle length, so convert it. sm_params["frequency.cycle"] = 2.0 * np.pi / cycle_length - params["cycle_cycle_length"] = np.atleast_1d(cycle_length) - params["cycle"] = np.ones((1,)) + params["cycle_length"] = np.atleast_1d(cycle_length) + params["cycle"] = np.ones((1,), dtype=floatX) if stochastic_cycle: - sigma = np.abs(rng.normal(size=(1,))) + sigma = np.abs(rng.normal(size=(1,))).astype(floatX) params["sigma_cycle"] = sigma sm_params["sigma2.cycle"] = sigma if damped_cycle: - rho = rng.beta(1, 1, size=(1,)) + rho = rng.beta(1, 1, size=(1,)).astype(floatX) params["cycle_dampening_factor"] = rho sm_params["damping.cycle"] = rho @@ -227,8 +227,8 @@ def create_structural_model_and_equivalent_statsmodel( components.append(comp) if autoregressive is not None: - ar_params = rng.normal(size=(autoregressive,)) - sigma = np.abs(rng.normal(size=(1,))) + ar_params = rng.normal(size=(autoregressive,)).astype(floatX) + sigma = np.abs(rng.normal(size=(1,))).astype(floatX) params["ar_params"] = ar_params params["sigma_ar"] = sigma @@ -241,7 +241,7 @@ def create_structural_model_and_equivalent_statsmodel( if exog is not None: names = [f"x{i + 1}" for i in range(exog.shape[1])] - betas = rng.normal(size=(exog.shape[1],)) + betas = rng.normal(size=(exog.shape[1],)).astype(floatX) params["beta_exog"] = betas params["data_exog"] = exog @@ -313,7 +313,7 @@ def test_structural_model_against_statsmodels( stochastic_cycle=stochastic_cycle, ) - data = rng.normal(size=(100,)) + data = rng.normal(size=(100,)).astype(floatX) sm_mod = f_sm_mod(data) sm_mod.initialize_default() @@ -462,7 +462,7 @@ def test_cycle_component_with_innovations_and_cycle_length(rng): ) params = { "cycle": np.array([1.0], dtype=floatX), - "cycle_cycle_length": np.array([12], dtype=floatX), + "cycle_length": np.array([12], dtype=floatX), "cycle_dampening_factor": np.array([0.95], dtype=floatX), "sigma_cycle": np.array([1.0], dtype=floatX), } From 99df6f6d0134a557b9e54514fd6870d3b0b99de0 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 23:41:32 +0100 Subject: [PATCH 6/9] Remove pipe from typehint --- pymc_experimental/__init__.py | 2 +- pymc_experimental/statespace/models/structural.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc_experimental/__init__.py b/pymc_experimental/__init__.py index e319ef35..7b9cf7bb 100644 --- a/pymc_experimental/__init__.py +++ b/pymc_experimental/__init__.py @@ -23,7 +23,7 @@ handler = logging.StreamHandler() _log.addHandler(handler) -from pymc_experimental import distributions, gp, utils +from pymc_experimental import distributions, gp, statespace, utils from pymc_experimental.inference.fit import fit from pymc_experimental.model.marginal_model import MarginalModel from pymc_experimental.model.model_api import as_model diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index a31b6022..d6486cc5 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -738,7 +738,7 @@ class LevelTrendComponent(Component): def __init__( self, - order: Union[int | list[int]] = 2, + order: Union[int, list[int]] = 2, innovations_order: Optional[Union[int, list[int]]] = None, name: str = "LevelTrend", ): From dc304c7fafcff775ad8e6915eaf3e0e0a1010b72 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Sun, 17 Dec 2023 23:57:52 +0100 Subject: [PATCH 7/9] Add `CycleComponent` and `BayesianSARIMA` to doctree --- docs/statespace/models.rst | 2 +- docs/statespace/models/structural.rst | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/docs/statespace/models.rst b/docs/statespace/models.rst index 59517f3d..6c3d1370 100644 --- a/docs/statespace/models.rst +++ b/docs/statespace/models.rst @@ -6,7 +6,7 @@ Statespace Models .. autosummary:: :toctree: generated - BayesianARIMA + BayesianSARIMA BayesianVARMAX ********************* diff --git a/docs/statespace/models/structural.rst b/docs/statespace/models/structural.rst index d19c0ee3..5582e5a1 100644 --- a/docs/statespace/models/structural.rst +++ b/docs/statespace/models/structural.rst @@ -11,6 +11,4 @@ Structural Components TimeSeasonality FrequencySeasonality MeasurementError - - StructuralTimeSeries - Component + CycleComponent From 73810fd40a3d9657cb0d99a386c77e634adbbe56 Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 18 Dec 2023 06:31:34 +0100 Subject: [PATCH 8/9] Fix typo in CycleComponent docstring --- pymc_experimental/statespace/models/structural.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index d6486cc5..c2942d03 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -1329,7 +1329,7 @@ class CycleComponent(Component): .. math:: \begin{align} \gamma_t &= \rho \gamma_{t-1} \cos \lambda + \rho \gamma_{t-1}^\star \sin \lambda + \omega_{t} \\ - \gamma_{t}^\star &= -\rho \gamma_{t-1} \sin \lambda + \rho \gamma_{t-1}^\star \cos \lambda + \omega_{t}^\star + \gamma_{t}^\star &= -\rho \gamma_{t-1} \sin \lambda + \rho \gamma_{t-1}^\star \cos \lambda + \omega_{t}^\star \\ \lambda &= \frac{2\pi}{s} \end{align} From 330a77d36d4309d435159253c1336c847c617fdb Mon Sep 17 00:00:00 2001 From: jessegrabowski Date: Mon, 18 Dec 2023 17:59:59 +0100 Subject: [PATCH 9/9] Remove `build_from_prior` function from unrelated branch --- .../statespace/core/statespace.py | 49 ------------------- 1 file changed, 49 deletions(-) diff --git a/pymc_experimental/statespace/core/statespace.py b/pymc_experimental/statespace/core/statespace.py index a68a858f..5bcce7a3 100644 --- a/pymc_experimental/statespace/core/statespace.py +++ b/pymc_experimental/statespace/core/statespace.py @@ -23,7 +23,6 @@ ) from pymc_experimental.statespace.filters.distributions import ( LinearGaussianStateSpace, - LinearGaussianStateSpaceRV, SequenceMvNormal, ) from pymc_experimental.statespace.filters.utilities import stabilize @@ -851,54 +850,6 @@ def build_statespace_graph( if return_updates: return updates - def build_as_prior( - self, - name: str, - n_steps: int, - register_statespace_matrices: Optional[bool] = False, - k_endog: Optional[int] = None, - mode: Optional[str] = None, - **kwargs, - ) -> LinearGaussianStateSpaceRV: - """ - Construct a random variable over - - Parameters - ---------- - name: str - Name of the random variable - n_steps: int - Length of the prior sequence, in time steps - register_statespace_matrices: bool - If True, statespace matrices used to construct the prior sequence will be wrapped in pm.Deterministic and - saved to the active model. - k_endog: int, optional - The number of observed states in the statespace. For debugging purposes only, should be automatically - inferred in most cases. - mode: str, optional - Pytensor compile mode. - kwargs: - PyMC distribution arguments, passed to LinearGaussianStateSpace - - Returns - ------- - obs_states: LinearGaussianStateSpaceRV - A random variable representing realizations of the observation equations of the statespace. - """ - - pm_mod = modelcontext(None) - self._insert_random_variables() - - if register_statespace_matrices: - matrices = self._register_matrices_with_pymc_model() - else: - matrices = self.unpack_statespace() - - latent_states, obs_states = LinearGaussianStateSpace( - name, *matrices, steps=n_steps, k_endog=k_endog, mode=mode, **kwargs - ) - return obs_states - def _build_smoother_graph( self, filtered_states: pt.TensorVariable,