Fix shape of sigma_irregular in create_structural_model_and_equivalent_statsmodel

Add measurement noise to (un)conditional observed distributions

SARIMAX tests don't depend on parameters being wrapped in `atleast_1d` inside `PyMCStateSpace.insert_random_variables()`

Infer parameter shapes from symbolic placeholders rather than `self.param_info`

Don't wrap all placeholders with `pt.atleast_1d` before applying `clone_replace`

Add properties for exogenous data, insert exogenous data into model separately from parameters

Use absolute path to test data

Refactor `PyMCStateSpace` methods to no longer expect matrices in provided `idata`

Don't add statespace matrices or outputs to PyMC graph by default.

Author: jessegrabowski

pymc_experimental/statespace/core/statespace.py

@@ -290,16 +290,16 @@ def dist(cls, a0, P0, c, d, T, Z, R, H, Q, *, steps=None, **kwargs):
         return latent_states, obs_states
 
 
-class SequenceMvNormalRV(SymbolicRandomVariable):
+class KalmanFilterRV(SymbolicRandomVariable):
     default_output = 1
-    _print_name = ("SequenceMvNormal", "\\operatorname{SequenceMvNormal}")
+    _print_name = ("KalmanFilter", "\\operatorname{KalmanFilter}")
 
     def update(self, node: Node):
         return {node.inputs[-1]: node.outputs[0]}
 
 
 class SequenceMvNormal(Continuous):
-    rv_op = SequenceMvNormalRV
+    rv_op = KalmanFilterRV
 
     def __new__(cls, *args, **kwargs):
         support_shape = get_support_shape(
@@ -336,6 +336,7 @@ def rv_op(cls, mus, covs, logp, steps, support_shape, size=None):
         else:
             batch_size = support_shape
 
+        # mus_, covs_ = mus.type(), covs.type()
         mus_, covs_, support_shape_ = mus.type(), covs.type(), support_shape.type()
         steps_ = steps.type()
         logp_ = logp.type()
@@ -352,15 +353,15 @@ def step(mu, cov, rng):
 
         (seq_mvn_rng,) = tuple(updates.values())
 
-        mvn_seq_op = SequenceMvNormalRV(
+        mvn_seq_op = KalmanFilterRV(
             inputs=[mus_, covs_, logp_, steps_], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2
         )
 
         mvn_seq = mvn_seq_op(mus, covs, logp, steps)
         return mvn_seq
 
 
-@_logprob.register(SequenceMvNormalRV)
+@_logprob.register(KalmanFilterRV)
 def sequence_mvnormal_logp(op, values, mus, covs, logp, steps, rng, **kwargs):
     return check_parameters(
         logp,

pymc_experimental/statespace/filters/distributions.py

@@ -100,6 +100,9 @@ def build_graph(
             [smoothed_covariances[::-1], pt.expand_dims(P_last, axis=(0,))], axis=0
         )
 
+        smoothed_states.name = "smoothed_states"
+        smoothed_covariances.name = "smoothed_covariances"
+
         return smoothed_states, smoothed_covariances
 
     def smoother_step(self, *args):

pymc_experimental/statespace/filters/kalman_smoother.py

@@ -269,11 +269,11 @@ def param_info(self) -> Dict[str, Dict[str, Any]]:
                 "constraints": "Positive Semi-definite",
             },
             "sigma_obs": {
-                "shape": (self.k_endog,),
+                "shape": None if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "Positive",
             },
             "sigma_state": {
-                "shape": (self.k_posdef,),
+                "shape": None if self.k_posdef == 1 else (self.k_posdef,),
                 "constraints": "Positive",
             },
             "ar_params": {
@@ -330,8 +330,9 @@ def param_dims(self):
             "seasonal_ar_params": (SEASONAL_AR_PARAM_DIM,),
             "seasonal_ma_params": (SEASONAL_MA_PARAM_DIM,),
         }
-
-        if not self.measurement_error:
+        if self.k_endog == 1:
+            del coord_map["sigma_state"]
+        if not self.measurement_error or self.k_endog == 1:
             del coord_map["sigma_obs"]
         if self.p == 0:
             del coord_map["ar_params"]
@@ -512,14 +513,14 @@ def make_symbolic_graph(self) -> None:
         # Set up the state covariance matrix
         state_cov_idx = ("state_cov",) + np.diag_indices(self.k_posdef)
         state_cov = self.make_and_register_variable(
-            "sigma_state", shape=(self.k_posdef,), dtype=floatX
+            "sigma_state", shape=() if self.k_posdef == 1 else (self.k_posdef,), dtype=floatX
         )
         self.ssm[state_cov_idx] = state_cov**2
 
         if self.measurement_error:
             obs_cov_idx = ("obs_cov",) + np.diag_indices(self.k_endog)
             obs_cov = self.make_and_register_variable(
-                "sigma_obs", shape=(self.k_endog,), dtype=floatX
+                "sigma_obs", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX
             )
             self.ssm[obs_cov_idx] = obs_cov**2
 

pymc_experimental/statespace/models/SARIMAX.py

@@ -81,15 +81,18 @@ def __init__(
         self,
         ssm: PytensorRepresentation,
         state_names,
+        data_names,
         shock_names,
         param_names,
         exog_names,
         param_dims,
         coords,
         param_info,
+        data_info,
         component_info,
         measurement_error,
         name_to_variable,
+        name_to_data,
         name=None,
         verbose=True,
         filter_type: str = "standard",
@@ -104,6 +107,7 @@ def __init__(
             param_names, param_dims, param_info, k_states
         )
         self._state_names = state_names
+        self._data_names = data_names
         self._shock_names = shock_names
         self._param_names = param_names
         self._param_dims = param_dims
@@ -113,6 +117,7 @@ def __init__(
 
         self._coords = coords
         self._param_info = param_info
+        self._data_info = data_info
         self.measurement_error = measurement_error
 
         super().__init__(
@@ -126,7 +131,10 @@ def __init__(
 
         self.ssm = ssm
         self._component_info = component_info
+
         self._name_to_variable = name_to_variable
+        self._name_to_data = name_to_data
+
         self._exog_names = exog_names
         self._needs_exog_data = len(exog_names) > 0
 
@@ -149,6 +157,10 @@ def _add_inital_state_cov_to_properties(param_names, param_dims, param_info, k_s
     def param_names(self):
         return self._param_names
 
+    @property
+    def data_names(self) -> List[str]:
+        return self._data_names
+
     @property
     def state_names(self):
         return self._state_names
@@ -173,6 +185,10 @@ def coords(self) -> Dict[str, Sequence]:
     def param_info(self) -> Dict[str, Dict[str, Any]]:
         return self._param_info
 
+    @property
+    def data_info(self) -> dict[str, dict[str, Any]]:
+        return self._data_info
+
     def make_symbolic_graph(self) -> None:
         """
         Assign placeholder pytensor variables among statespace matrices in positions where PyMC variables will go.
@@ -338,6 +354,7 @@ def __init__(
         k_states,
         k_posdef,
         state_names=None,
+        data_names=None,
         shock_names=None,
         param_names=None,
         exog_names=None,
@@ -354,14 +371,18 @@ def __init__(
         self.measurement_error = measurement_error
 
         self.state_names = state_names if state_names is not None else []
+        self.data_names = data_names if data_names is not None else []
         self.shock_names = shock_names if shock_names is not None else []
         self.param_names = param_names if param_names is not None else []
         self.exog_names = exog_names if exog_names is not None else []
 
         self.needs_exog_data = len(self.exog_names) > 0
         self.coords = {}
         self.param_dims = {}
+
         self.param_info = {}
+        self.data_info = {}
+
         self.param_counts = {}
 
         if representation is None:
@@ -370,6 +391,7 @@ def __init__(
             self.ssm = representation
 
         self._name_to_variable = {}
+        self._name_to_data = {}
 
         if not component_from_sum:
             self.populate_component_properties()
@@ -429,6 +451,43 @@ def make_and_register_variable(self, name, shape, dtype=floatX) -> Variable:
         self._name_to_variable[name] = placeholder
         return placeholder
 
+    def make_and_register_data(self, name, shape, dtype=floatX) -> Variable:
+        r"""
+        Helper function to create a pytensor symbolic variable and register it in the _name_to_data dictionary
+
+        Parameters
+        ----------
+        name : str
+            The name of the placeholder data. Must be the name of an expected data variable.
+        shape : int or tuple of int
+            Shape of the parameter
+        dtype : str, default pytensor.config.floatX
+            dtype of the parameter
+
+        Notes
+        -----
+        See docstring for make_and_register_variable for more details. This function is similar, but handles data
+        inputs instead of model parameters.
+
+        An error is raised if the provided name has already been registered, or if the name is not present in the
+        ``data_names`` property.
+        """
+        if name not in self.data_names:
+            raise ValueError(
+                f"{name} is not a model parameter. All placeholder variables should correspond to model "
+                f"parameters."
+            )
+
+        if name in self._name_to_data.keys():
+            raise ValueError(
+                f"{name} is already a registered placeholder variable with shape "
+                f"{self._name_to_data[name].type.shape}"
+            )
+
+        placeholder = pt.tensor(name, shape=shape, dtype=dtype)
+        self._name_to_data[name] = placeholder
+        return placeholder
+
     def make_symbolic_graph(self) -> None:
         raise NotImplementedError
 
@@ -481,7 +540,6 @@ def make_slice(name, x, o_x):
         transition.name = T.name
 
         design = pt.concatenate(conform_time_varying_and_time_invariant_matrices(Z, o_Z), axis=-1)
-
         design.name = Z.name
 
         selection = block_diagonal([R, o_R])
@@ -542,14 +600,18 @@ def _make_combined_name(self):
 
     def __add__(self, other):
         state_names = self._combine_property(other, "state_names")
+        data_names = self._combine_property(other, "data_names")
         param_names = self._combine_property(other, "param_names")
         shock_names = self._combine_property(other, "shock_names")
         param_info = self._combine_property(other, "param_info")
+        data_info = self._combine_property(other, "data_info")
         param_dims = self._combine_property(other, "param_dims")
         coords = self._combine_property(other, "coords")
         exog_names = self._combine_property(other, "exog_names")
 
         _name_to_variable = self._combine_property(other, "_name_to_variable")
+        _name_to_data = self._combine_property(other, "_name_to_data")
+
         measurement_error = any([self.measurement_error, other.measurement_error])
 
         k_states, k_posdef, k_endog = self._get_combined_shapes(other)
@@ -567,30 +629,22 @@ def __add__(self, other):
         new_comp._component_info = self._combine_component_info(other)
         new_comp.name = new_comp._make_combined_name()
 
-        property_names = [
-            "state_names",
-            "param_names",
-            "shock_names",
-            "state_dims",
-            "coords",
-            "param_dims",
-            "param_info",
-            "exog_names",
-            "_name_to_variable",
-        ]
-        property_values = [
-            state_names,
-            param_names,
-            shock_names,
-            param_dims,
-            coords,
-            param_dims,
-            param_info,
-            exog_names,
-            _name_to_variable,
+        names_and_props = [
+            ("state_names", state_names),
+            ("data_names", data_names),
+            ("param_names", param_names),
+            ("shock_names", shock_names),
+            ("param_dims", param_dims),
+            ("coords", coords),
+            ("param_dims", param_dims),
+            ("param_info", param_info),
+            ("data_info", data_info),
+            ("exog_names", exog_names),
+            ("_name_to_variable", _name_to_variable),
+            ("_name_to_data", _name_to_data),
         ]
 
-        for prop, value in zip(property_names, property_values):
+        for prop, value in names_and_props:
             setattr(new_comp, prop, value)
 
         return new_comp
@@ -622,15 +676,18 @@ def build(self, name=None, filter_type="standard", verbose=True):
             self.ssm,
             name=name,
             state_names=self.state_names,
+            data_names=self.data_names,
             shock_names=self.shock_names,
             param_names=self.param_names,
             param_dims=self.param_dims,
             coords=self.coords,
             param_info=self.param_info,
+            data_info=self.data_info,
             component_info=self._component_info,
             measurement_error=self.measurement_error,
             exog_names=self.exog_names,
             name_to_variable=self._name_to_variable,
+            name_to_data=self._name_to_data,
             filter_type=filter_type,
             verbose=verbose,
         )
@@ -881,7 +938,8 @@ def populate_component_properties(self):
         }
 
     def make_symbolic_graph(self) -> None:
-        error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=(self.k_endog,))
+        sigma_shape = () if self.k_endog == 1 else (self.k_endog,)
+        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
@@ -1541,7 +1599,7 @@ def _handle_input_data(self, k_exog: int, state_names: Optional[List[str]], name
 
     def make_symbolic_graph(self) -> None:
         betas = self.make_and_register_variable(f"beta_{self.name}", shape=(self.k_states,))
-        regression_data = self.make_and_register_variable(
+        regression_data = self.make_and_register_data(
             f"data_{self.name}", shape=(None, self.k_states)
         )
 
@@ -1560,17 +1618,19 @@ def make_symbolic_graph(self) -> None:
     def populate_component_properties(self) -> None:
         self.shock_names = self.state_names
 
-        self.param_names = [f"beta_{self.name}", f"data_{self.name}"]
+        self.param_names = [f"beta_{self.name}"]
+        self.data_names = [f"data_{self.name}"]
         self.param_dims = {
             f"beta_{self.name}": ("exog_state",),
-            f"data_{self.name}": (TIME_DIM, "exog_state"),
         }
 
         self.param_info = {
             f"beta_{self.name}": {"shape": (1,), "constraints": None, "dims": ("exog_state",)},
+        }
+
+        self.data_info = {
             f"data_{self.name}": {
                 "shape": (None, self.k_states),
-                "constraints": None,
                 "dims": (TIME_DIM, "exog_state"),
             },
         }

pymc_experimental/statespace/models/structural.py

@@ -180,7 +180,7 @@ def pymc_mod(arima_mod):
         ar_params = pm.Normal("ar_params", sigma=0.1, dims=["ar_lag"])
         ma_params = pm.Normal("ma_params", sigma=1, dims=["ma_lag"])
         sigma_state = pm.Exponential("sigma_state", 0.5)
-        arima_mod.build_statespace_graph(data=data)
+        arima_mod.build_statespace_graph(data=data, save_kalman_filter_outputs_in_idata=True)
 
     return pymc_mod
 
@@ -210,7 +210,8 @@ def pymc_mod_interp(arima_mod_interp):
         ma_params = pm.Normal("ma_params", sigma=1, dims=["ma_lag"])
         sigma_state = pm.Exponential("sigma_state", 0.5)
         sigma_obs = pm.Exponential("sigma_obs", 0.1)
-        arima_mod_interp.build_statespace_graph(data=data)
+
+        arima_mod_interp.build_statespace_graph(data=data, save_kalman_filter_outputs_in_idata=True)
 
     return pymc_mod