diff --git a/pymc_extras/statespace/core/representation.py b/pymc_extras/statespace/core/representation.py index 450108d38..f563ca861 100644 --- a/pymc_extras/statespace/core/representation.py +++ b/pymc_extras/statespace/core/representation.py @@ -298,7 +298,7 @@ def _validate_matrix_shape(self, name: str, X: np.ndarray | pt.TensorVariable) - ) # Time varying vector case, check only the static shapes - if X.ndim == 2 and X.shape[1:] != expected_shape: + if X.ndim == 2 and shape[1:] != expected_shape: raise ValueError( f"Last dimension of array provided for {name} has shape {X.shape[1]}, " f"expected {expected_shape}" @@ -334,7 +334,7 @@ def _check_provided_tensor(self, name: str, X: pt.TensorVariable) -> pt.TensorVa X = pt.expand_dims(X, (0,)) X = pt.specify_shape(X, self.shapes[name]) - elif X.ndim == 2: + elif X.ndim == 2 and name not in VECTOR_VALUED: X = pt.expand_dims(X, (0,)) X = pt.specify_shape(X, self.shapes[name]) diff --git a/pymc_extras/statespace/models/structural.py b/pymc_extras/statespace/models/structural.py index a982366c3..c7b52d893 100644 --- a/pymc_extras/statespace/models/structural.py +++ b/pymc_extras/statespace/models/structural.py @@ -25,6 +25,8 @@ ALL_STATE_DIM, AR_PARAM_DIM, LONG_MATRIX_NAMES, + OBS_STATE_AUX_DIM, + OBS_STATE_DIM, POSITION_DERIVATIVE_NAMES, TIME_DIM, ) @@ -78,6 +80,7 @@ def __init__( component_info: dict[str, dict[str, Any]], measurement_error: bool, name_to_variable: dict[str, Variable], + endog_names: list[str] | None = None, name_to_data: dict[str, Variable] | None = None, name: str | None = None, verbose: bool = True, @@ -88,6 +91,7 @@ def __init__( if name is None: name = "data" self._name = name + self._endog_names = endog_names if endog_names is not None else [name] k_states, k_posdef, k_endog = ssm.k_states, ssm.k_posdef, ssm.k_endog param_names, param_dims, param_info = self._add_inital_state_cov_to_properties( @@ -165,7 +169,7 @@ def state_names(self): @property def observed_states(self): - return [self._name] + return self._endog_names @property def shock_names(self): @@ -351,6 +355,7 @@ def __init__( k_endog, k_states, k_posdef, + endog_names=None, state_names=None, data_names=None, shock_names=None, @@ -367,6 +372,7 @@ def __init__( self.k_states = k_states self.k_posdef = k_posdef self.measurement_error = measurement_error + self.endog_names = [self.name] if endog_names is None else endog_names self.state_names = state_names if state_names is not None else [] self.data_names = data_names if data_names is not None else [] @@ -597,6 +603,9 @@ def _make_combined_name(self): return name def __add__(self, other): + if self.endog_names != other.endog_names: + raise ValueError("The endog names must be the same.") + state_names = self._combine_property(other, "state_names") data_names = self._combine_property(other, "data_names") param_names = self._combine_property(other, "param_names") @@ -617,7 +626,8 @@ def __add__(self, other): new_comp = Component( name="", - k_endog=1, + endog_names=self.endog_names, + k_endog=k_endog, k_states=k_states, k_posdef=k_posdef, measurement_error=measurement_error, @@ -682,6 +692,7 @@ def build( return StructuralTimeSeries( self.ssm, name=name, + endog_names=self.endog_names, state_names=self.state_names, data_names=self.data_names, shock_names=self.shock_names, @@ -806,9 +817,11 @@ class LevelTrendComponent(Component): def __init__( self, + endog_names, order: int | list[int] = 2, innovations_order: int | list[int] | None = None, name: str = "LevelTrend", + **kwargs, ): if innovations_order is None: innovations_order = order @@ -840,14 +853,24 @@ def __init__( super().__init__( name, - k_endog=1, + endog_names=endog_names, + k_endog=len(endog_names), 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)), + **kwargs, ) + @property + def observed_states(self): + return self.endog_names + + @property + def k_endog(self): + return len(self.endog_names) + def populate_component_properties(self): name_slice = POSITION_DERIVATIVE_NAMES[: self.k_states] self.param_names = ["initial_trend"] @@ -878,7 +901,9 @@ def make_symbolic_graph(self) -> None: R = R[:, self.innovations_order] self.ssm["selection", :, :] = R - self.ssm["design", 0, :] = np.array([1.0] + [0.0] * (self.k_states - 1)) + self.ssm["design", :, :] = np.tile( + np.array([1.0] + [0.0] * (self.k_states - 1)), (self.k_endog, 1) + ) if self.k_posdef > 0: sigma_trend = self.make_and_register_variable("sigma_trend", shape=(self.k_posdef,)) @@ -925,32 +950,43 @@ class MeasurementError(Component): idata = pm.sample(nuts_sampler='numpyro') """ - def __init__(self, name: str = "MeasurementError"): - k_endog = 1 + def __init__(self, endog_names, time_dim, name: str = "MeasurementError", **kwargs): k_states = 0 k_posdef = 0 + self._time_dim = time_dim super().__init__( - name, k_endog, k_states, k_posdef, measurement_error=True, combine_hidden_states=False + name, + k_states, + k_posdef, + endog_names=endog_names, + k_endog=len(endog_names), + measurement_error=True, + combine_hidden_states=False, + **kwargs, ) def populate_component_properties(self): - self.param_names = [f"sigma_{self.name}"] - self.param_dims = {} self.param_info = { - f"sigma_{self.name}": { - "shape": (), - "constraints": "Positive", - "dims": None, + f"{self.name}_covariance": { + "shape": (self.k_endog, self.k_endog), + "constraints": "Positive semi-definite", + "dims": ( + OBS_STATE_DIM, + OBS_STATE_AUX_DIM, + ), } } + self.param_names = list(self.param_info.keys()) + self.param_dims = {name: val["dims"] for name, val in self.param_info.items()} def make_symbolic_graph(self) -> None: - sigma_shape = () - error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=sigma_shape) - diag_idx = np.diag_indices(self.k_endog) - idx = np.s_["obs_cov", diag_idx[0], diag_idx[1]] - self.ssm[idx] = error_sigma**2 + sigma_shape = (self.k_endog, self.k_endog) + obs_covariance = self.make_and_register_variable( + f"{self.name}_covariance", shape=sigma_shape + ) + + self.ssm["obs_cov"] = obs_covariance class AutoregressiveComponent(Component): @@ -1009,7 +1045,7 @@ class AutoregressiveComponent(Component): """ - def __init__(self, order: int = 1, name: str = "AutoRegressive"): + def __init__(self, endog_names, order: int = 1, name: str = "AutoRegressive"): order = order_to_mask(order) ar_lags = np.flatnonzero(order).ravel().astype(int) + 1 k_states = len(order) @@ -1019,7 +1055,8 @@ def __init__(self, order: int = 1, name: str = "AutoRegressive"): super().__init__( name=name, - k_endog=1, + endog_names=endog_names, + k_endog=len(endog_names), k_states=k_states, k_posdef=1, measurement_error=True, @@ -1051,7 +1088,8 @@ def make_symbolic_graph(self) -> None: T = np.eye(self.k_states, k=-1) self.ssm["transition", :, :] = T self.ssm["selection", 0, 0] = 1 - self.ssm["design", 0, 0] = 1 + # We want all the endogenous variables to load on this + self.ssm["design", :, 0] = 1 ar_idx = ("transition", np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0]) self.ssm[ar_idx] = ar_params @@ -1315,7 +1353,7 @@ class FrequencySeasonality(Component): to isolate and identify a "Monday" effect, for instance. """ - def __init__(self, season_length, n=None, name=None, innovations=True): + def __init__(self, endog_names, season_length, n=None, name=None, innovations=True, **kwargs): if n is None: n = int(season_length // 2) if name is None: @@ -1337,16 +1375,18 @@ def __init__(self, season_length, n=None, name=None, innovations=True): super().__init__( name=name, - k_endog=1, + endog_names=endog_names, + k_endog=len(endog_names), k_states=k_states, k_posdef=k_states * int(self.innovations), measurement_error=False, combine_hidden_states=True, obs_state_idxs=obs_state_idx, + **kwargs, ) def make_symbolic_graph(self) -> None: - self.ssm["design", 0, slice(0, self.k_states, 2)] = 1 + self.ssm["design", :, slice(0, self.k_states, 2)] = 1 init_state = self.make_and_register_variable(f"{self.name}", shape=(self.n_coefs,)) @@ -1677,3 +1717,46 @@ def populate_component_properties(self) -> None: "constraints": "Positive", "dims": ("exog_state",), } + + +class ObsIntercept(Component): + def __init__(self, endog_names, time_dim, name: str = "y"): + self._time_dim = time_dim + super().__init__( + name=name, + endog_names=endog_names, + k_endog=len(endog_names), + k_states=0, + k_posdef=0, + measurement_error=False, + combine_hidden_states=False, + ) + + def populate_component_properties(self): + self.param_info = { + "mu": { + "shape": ( + # TODO This should not be necessary. + self._time_dim, + self.k_endog, + ), + "constraints": None, + "dims": ( + TIME_DIM, + OBS_STATE_DIM, + ), + }, + } + self.param_dims = {name: val["dims"] for name, val in self.param_info.items()} + self.param_names = list(self.param_info.keys()) + + def make_symbolic_graph(self) -> None: + mu = self.make_and_register_variable( + "mu", + shape=( + self._time_dim, + self.k_endog, + ), + ) + + self.ssm["obs_intercept"] = mu