Different stability strategy for UnivariateFilter

pymc-devs · Aug 5, 2023 · 3a143a3 · 3a143a3
1 parent b0cf3de
commit 3a143a3
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/pymc_experimental/statespace/filters/kalman_filter.py b/pymc_experimental/statespace/filters/kalman_filter.py
@@ -862,7 +862,6 @@ def _univariate_inner_filter_step(self, y, Z_row, d_row, sigma_H, nan_flag, a, P
 
         a_filtered = a + K * v
         P_filtered = P - pt.outer(K, K) * F
-        P_filtered = stabilize(0.5 * (P_filtered + P_filtered.T))
 
         ll_inner = pt.switch(F_zero_flag, 0.0, pt.log(F) + v**2 / F)
 
@@ -892,7 +891,7 @@ def kalman_step(self, y, a, P, c, d, T, Z, R, H, Q):
             obs_cov[-1],
         )
 
-        P_filtered = stabilize(P_filtered, self.cov_jitter)
+        P_filtered = stabilize(0.5 * (P_filtered + P_filtered.T), self.cov_jitter)
         a_hat, P_hat = self.predict(a=a_filtered, P=P_filtered, c=c, T=T, R=R, Q=Q)
 
         ll = -0.5 * ((pt.neq(ll_inner, 0).sum()) * MVN_CONST + ll_inner.sum())