Add innovation hidden state, data is observed state

Author: jessegrabowski

pymc_experimental/statespace/models/ETS.py

@@ -272,7 +272,7 @@ def param_info(self) -> dict[str, dict[str, Any]]:
 
     @property
     def state_names(self):
-        states = ["data", "level"]
+        states = ["innovation", "level"]
         if self.trend:
             states += ["trend"]
         if self.seasonal:
@@ -282,7 +282,7 @@ def state_names(self):
 
     @property
     def observed_states(self):
-        return [self.state_names[0]]
+        return ["data"]
 
     @property
     def shock_names(self):
@@ -326,20 +326,15 @@ def make_symbolic_graph(self) -> None:
         self.ssm["initial_state", :] = x0
         self.ssm["initial_state_cov"] = P0
 
-        # Only the first state is ever observed
-        Z = np.zeros((self.k_endog, self.k_states))
-        Z[0, 0] = 1
-        self.ssm["design"] = Z
-
         # The shape of R can be pre-allocated, then filled with the required parameters
         R = pt.zeros((self.k_states, self.k_posdef))
         R = pt.set_subtensor(R[0, :], 1.0)  # We will always have y_t = ... + e_t
 
         alpha = self.make_and_register_variable("alpha", shape=(), dtype=floatX)
         R = pt.set_subtensor(R[1, 0], alpha)  # and l_t = ... + alpha * e_t
 
-        # Data and level component always exists, the base case is y_t = l_{t-1} and l_t = l_{t-1}
-        T_base = pt.as_tensor_variable(np.array([[0.0, 1.0], [0.0, 1.0]]))
+        # Shock and level component always exists, the base case is e_t = e_t and l_t = l_{t-1}
+        T_base = pt.as_tensor_variable(np.array([[0.0, 0.0], [0.0, 1.0]]))
 
         if self.trend:
             beta = self.make_and_register_variable("beta", shape=(), dtype=floatX)
@@ -349,7 +344,7 @@ def make_symbolic_graph(self) -> None:
             # y_t = l_{t-1} + b_{t-1}
             # l_t = l_{t-1} + b_{t-1}
             # b_t = b_{t-1}
-            T_base = pt.as_tensor_variable(([0.0, 1.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]))
+            T_base = pt.as_tensor_variable(([0.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]))
 
         if self.damped_trend:
             phi = self.make_and_register_variable("phi", shape=(), dtype=floatX)
@@ -358,7 +353,7 @@ def make_symbolic_graph(self) -> None:
             # y_t = l_{t-1} + phi * b_{t-1}
             # l_t = l_{t-1} + phi * b_{t-1}
             # b_t = phi * b_{t-1}
-            T_base = pt.set_subtensor(T_base[:, 2], phi)
+            T_base = pt.set_subtensor(T_base[1:, 2], phi)
 
         T_components = [T_base]
 
@@ -375,10 +370,17 @@ def make_symbolic_graph(self) -> None:
         self.ssm["selection"] = R
 
         T = pt.linalg.block_diag(*T_components)
-        if self.seasonal:
-            T = pt.set_subtensor(T[0, 2 + int(self.trend) + 1], 1.0)
         self.ssm["transition"] = pt.specify_shape(T, (self.k_states, self.k_states))
 
+        Z = np.zeros((self.k_endog, self.k_states))
+        Z[0, 0] = 1.0  # innovation
+        Z[0, 1] = 1.0  # level
+        if self.trend:
+            Z[0, 2] = 1.0
+        if self.seasonal:
+            Z[0, 2 + int(self.trend)] = 1.0
+        self.ssm["design"] = Z
+
         # Set up the state covariance matrix
         state_cov_idx = ("state_cov",) + np.diag_indices(self.k_posdef)
         state_cov = self.make_and_register_variable(

tests/statespace/test_ETS.py

@@ -153,38 +153,40 @@ def test_statespace_matrices(order: tuple[str, str, str], expected_params: list[
     assert_allclose(H, np.eye(1) * test_values["sigma_obs"] ** 2)
     assert_allclose(Q, np.eye(1) * test_values["sigma_state"] ** 2)
 
-    Z_val = np.zeros((1, expected_states))
-    Z_val[0, 0] = 1.0
-    assert_allclose(Z, Z_val)
-
     R_val = np.zeros((expected_states, 1))
     R_val[0] = 1.0
     R_val[1] = test_values["alpha"]
 
+    Z_val = np.zeros((1, expected_states))
+    Z_val[0, 0] = 1.0
+    Z_val[0, 1] = 1.0
+
     if order[1] == "N":
-        T_val = np.array([[0.0, 1.0], [0.0, 1.0]])
+        T_val = np.array([[0.0, 0.0], [0.0, 1.0]])
     else:
         R_val[2] = test_values["beta"]
-        T_val = np.array([[0.0, 1.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]])
+        T_val = np.array([[0.0, 0.0, 0.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]])
+        Z_val[0, 2] = 1.0
 
     if order[1] == "Ad":
-        T_val[:, -1] *= test_values["phi"]
+        T_val[1:, -1] *= test_values["phi"]
 
     if order[2] == "A":
         R_val[3] = test_values["gamma"]
         S = np.eye(seasonal_periods, k=-1)
         S[0, :] = -1
+        Z_val[0, 2 + int(order[1] != "N")] = 1.0
     else:
         S = np.eye(0)
+
     T_val = linalg.block_diag(T_val, S)
-    if order[2] != "N":
-        T_val[0, 3 + int(order[1] != "N")] = 1.0
 
     assert_allclose(T, T_val)
     assert_allclose(R, R_val)
+    assert_allclose(Z, Z_val)
 
 
-def test_simulate_model():
+def test_deterministic_simulation_matches_statsmodels():
     mod = BayesianETS(order=("A", "Ad", "A"), seasonal_periods=4, measurement_error=False)
 
     rng = np.random.default_rng()