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
 
         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