pymc-extras 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

pymc_extras/statespace/filters/distributions.py

@@ -197,10 +197,9 @@ class _LinearGaussianStateSpace(Continuous):
         n_seq = len(sequence_names)
 
         def step_fn(*args):
-            seqs, state, non_seqs = args[:n_seq], args[n_seq], args[n_seq + 1 :]
-            non_seqs, rng = non_seqs[:-1], non_seqs[-1]
+            seqs, (rng, state, *non_seqs) = args[:n_seq], args[n_seq:]
 
-            c, d, T, Z, R, H, Q = sort_args(seqs + non_seqs)
+            c, d, T, Z, R, H, Q = sort_args((*seqs, *non_seqs))
             k = T.shape[0]
             a = state[:k]
 
@@ -219,7 +218,7 @@ class _LinearGaussianStateSpace(Continuous):
 
             next_state = pt.concatenate([a_next, y_next], axis=0)
 
-            return next_state, {rng: next_rng}
+            return next_rng, next_state
 
         Z_init = Z_ if Z_ in non_sequences else Z_[0]
         H_init = H_ if H_ in non_sequences else H_[0]
@@ -229,13 +228,14 @@ class _LinearGaussianStateSpace(Continuous):
 
         init_dist_ = pt.concatenate([init_x_, init_y_], axis=0)
 
-        statespace, updates = pytensor.scan(
+        ss_rng, statespace = pytensor.scan(
             step_fn,
-            outputs_info=[init_dist_],
+            outputs_info=[rng, init_dist_],
             sequences=None if len(sequences) == 0 else sequences,
-            non_sequences=[*non_sequences, rng],
+            non_sequences=[*non_sequences],
             n_steps=steps,
             strict=True,
+            return_updates=False,
         )
 
         if append_x0:
@@ -245,7 +245,6 @@ class _LinearGaussianStateSpace(Continuous):
             statespace_ = statespace
             statespace_ = pt.specify_shape(statespace_, (steps, None))
 
-        (ss_rng,) = tuple(updates.values())
         linear_gaussian_ss_op = LinearGaussianStateSpaceRV(
             inputs=[a0_, P0_, c_, d_, T_, Z_, R_, H_, Q_, steps, rng],
             outputs=[ss_rng, statespace_],
@@ -385,10 +384,15 @@ class SequenceMvNormal(Continuous):
 
         def step(mu, cov, rng):
             new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method=method).owner.outputs
-            return mvn, {rng: new_rng}
+            return new_rng, mvn
 
-        mvn_seq, updates = pytensor.scan(
-            step, sequences=[mus_, covs_], non_sequences=[rng], strict=True, n_steps=mus_.shape[0]
+        seq_mvn_rng, mvn_seq = pytensor.scan(
+            step,
+            sequences=[mus_, covs_],
+            outputs_info=[rng, None],
+            strict=True,
+            n_steps=mus_.shape[0],
+            return_updates=False,
         )
         mvn_seq = pt.specify_shape(mvn_seq, mus.type.shape)
 
@@ -396,8 +400,6 @@ class SequenceMvNormal(Continuous):
         if mvn_seq.ndim > 2:
             mvn_seq = pt.moveaxis(mvn_seq, 0, -2)
 
-        (seq_mvn_rng,) = tuple(updates.values())
-
         mvn_seq_op = KalmanFilterRV(
             inputs=[mus_, covs_, logp_, rng], outputs=[seq_mvn_rng, mvn_seq], ndim_supp=2
         )

pymc_extras/statespace/filters/kalman_filter.py

@@ -148,10 +148,9 @@ class BaseFilter(ABC):
         R,
         H,
         Q,
-        return_updates=False,
         missing_fill_value=None,
         cov_jitter=None,
-    ) -> list[TensorVariable] | tuple[list[TensorVariable], dict]:
+    ) -> list[TensorVariable]:
         """
         Construct the computation graph for the Kalman filter. See [1] for details.
 
@@ -200,7 +199,7 @@ class BaseFilter(ABC):
         self.n_endog = Z_shape[-2]
 
         data, a0, P0, *params = self.check_params(data, a0, P0, c, d, T, Z, R, H, Q)
-
+        data = pt.specify_shape(data, (data.type.shape[0], self.n_endog))
         sequences, non_sequences, seq_names, non_seq_names = split_vars_into_seq_and_nonseq(
             params, PARAM_NAMES
         )
@@ -211,20 +210,17 @@ class BaseFilter(ABC):
         if len(sequences) > 0:
             sequences = self.add_check_on_time_varying_shapes(data, sequences)
 
-        results, updates = pytensor.scan(
+        results = pytensor.scan(
             self.kalman_step,
             sequences=[data, *sequences],
             outputs_info=[None, a0, None, None, P0, None, None],
             non_sequences=non_sequences,
             name="forward_kalman_pass",
             strict=False,
+            return_updates=False,
         )
 
-        filter_results = self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
-
-        if return_updates:
-            return filter_results, updates
-        return filter_results
+        return self._postprocess_scan_results(results, a0, P0, n=data.type.shape[0])
 
     def _postprocess_scan_results(self, results, a0, P0, n) -> list[TensorVariable]:
         """
@@ -393,7 +389,7 @@ class BaseFilter(ABC):
         .. [1] Durbin, J., and S. J. Koopman. Time Series Analysis by State Space Methods.
                2nd ed, Oxford University Press, 2012.
         """
-        a_hat = T.dot(a) + c
+        a_hat = T @ a + c
         P_hat = quad_form_sym(T, P) + quad_form_sym(R, Q)
 
         return a_hat, P_hat
@@ -580,16 +576,16 @@ class StandardFilter(BaseFilter):
         .. [1] Durbin, J., and S. J. Koopman. Time Series Analysis by State Space Methods.
                2nd ed, Oxford University Press, 2012.
         """
-        y_hat = d + Z.dot(a)
+        y_hat = d + Z @ a
         v = y - y_hat
 
-        PZT = P.dot(Z.T)
+        PZT = P.dot(Z.mT)
         F = Z.dot(PZT) + stabilize(H, self.cov_jitter)
 
-        K = pt.linalg.solve(F.T, PZT.T, assume_a="pos", check_finite=False).T
+        K = pt.linalg.solve(F.mT, PZT.mT, assume_a="pos", check_finite=False).mT
         I_KZ = pt.eye(self.n_states) - K.dot(Z)
 
-        a_filtered = a + K.dot(v)
+        a_filtered = a + K @ v
         P_filtered = quad_form_sym(I_KZ, P) + quad_form_sym(K, H)
 
         F_inv_v = pt.linalg.solve(F, v, assume_a="pos", check_finite=False)
@@ -630,9 +626,9 @@ class SquareRootFilter(BaseFilter):
         a_hat = T.dot(a) + c
         Q_chol = pt.linalg.cholesky(Q, lower=True)
 
-        M = pt.horizontal_stack(T @ P_chol, R @ Q_chol).T
+        M = pt.horizontal_stack(T @ P_chol, R @ Q_chol).mT
         R_decomp = pt.linalg.qr(M, mode="r")
-        P_chol_hat = R_decomp[: self.n_states, : self.n_states].T
+        P_chol_hat = R_decomp[..., : self.n_states, : self.n_states].mT
 
         return a_hat, P_chol_hat
 
@@ -652,20 +648,22 @@ class SquareRootFilter(BaseFilter):
         y_hat = Z.dot(a) + d
         v = y - y_hat
 
-        H_chol = pytensor.ifelse(pt.all(pt.eq(H, 0.0)), H, pt.linalg.cholesky(H, lower=True))
+        H_chol = pytensor.ifelse(
+            pt.all(pt.eq(H, 0.0)), H, pt.linalg.cholesky(H, lower=True, on_error="nan")
+        )
 
         # The following notation comes from https://ipnpr.jpl.nasa.gov/progress_report/42-233/42-233A.pdf
         # Construct upper-triangular block matrix A = [[chol(H), Z @ L_pred],
         #                                              [0,           L_pred]]
         # The Schur decomposition of this matrix will be B (upper triangular). We are
-        # more insterested in B^T:
+        # more interested in B^T:
         # Structure of B^T = [[chol(F),     0              ],
         #                    [K @ chol(F), chol(P_filtered)]
         zeros = pt.zeros((self.n_states, self.n_endog))
         upper = pt.horizontal_stack(H_chol, Z @ P_chol)
         lower = pt.horizontal_stack(zeros, P_chol)
         A_T = pt.vertical_stack(upper, lower)
-        B = pt.linalg.qr(A_T.T, mode="r").T
+        B = pt.linalg.qr(A_T.mT, mode="r").mT
 
         F_chol = B[: self.n_endog, : self.n_endog]
         K_F_chol = B[self.n_endog :, : self.n_endog]
@@ -677,6 +675,7 @@ class SquareRootFilter(BaseFilter):
             inner_term = solve_triangular(
                 F_chol, solve_triangular(F_chol, v, lower=True), lower=True
             )
+
             loss = (v.T @ inner_term).ravel()
 
             # abs necessary because we're not guaranteed a positive diagonal from the schur decomposition
@@ -693,8 +692,10 @@ class SquareRootFilter(BaseFilter):
             """
             return [a, P_chol, pt.zeros(())]
 
+        degenerate = pt.eq(all_nan_flag, 1.0)
+        F_chol = pytensor.ifelse(degenerate, pt.eye(*F_chol.shape), F_chol)
         [a_filtered, P_chol_filtered, ll] = pytensor.ifelse(
-            pt.eq(all_nan_flag, 1.0),
+            degenerate,
             compute_degenerate(P_chol_filtered, F_chol, K_F_chol, v),
             compute_non_degenerate(P_chol_filtered, F_chol, K_F_chol, v),
         )
@@ -785,11 +786,12 @@ class UnivariateFilter(BaseFilter):
         H_masked = W.dot(H)
         y_masked = pt.set_subtensor(y[nan_mask], 0.0)
 
-        result, updates = pytensor.scan(
+        result = pytensor.scan(
             self._univariate_inner_filter_step,
             sequences=[y_masked, Z_masked, d, pt.diag(H_masked), nan_mask],
             outputs_info=[a, P, None, None, None],
             name="univariate_inner_scan",
+            return_updates=False,
         )
 
         a_filtered, P_filtered, obs_mu, obs_cov, ll_inner = result
@@ -800,7 +802,7 @@ class UnivariateFilter(BaseFilter):
             obs_cov[-1],
         )
 
-        P_filtered = stabilize(0.5 * (P_filtered + P_filtered.T), self.cov_jitter)
+        P_filtered = stabilize(0.5 * (P_filtered + P_filtered.mT), 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())

pymc_extras/statespace/filters/kalman_smoother.py

@@ -1,8 +1,6 @@
 import pytensor
 import pytensor.tensor as pt
 
-from pytensor.tensor.nlinalg import matrix_dot
-
 from pymc_extras.statespace.filters.utilities import (
     quad_form_sym,
     split_vars_into_seq_and_nonseq,
@@ -78,16 +76,16 @@ class KalmanSmoother:
         self.seq_names = seq_names
         self.non_seq_names = non_seq_names
 
-        smoother_result, updates = pytensor.scan(
+        smoothed_states, smoothed_covariances = pytensor.scan(
             self.smoother_step,
             sequences=[filtered_states[:-1], filtered_covariances[:-1], *sequences],
             outputs_info=[a_last, P_last],
             non_sequences=non_sequences,
             go_backwards=True,
             name="kalman_smoother",
+            return_updates=False,
         )
 
-        smoothed_states, smoothed_covariances = smoother_result
         smoothed_states = pt.concatenate(
             [smoothed_states[::-1], pt.expand_dims(a_last, axis=(0,))], axis=0
         )
@@ -105,7 +103,7 @@ class KalmanSmoother:
         a_hat, P_hat = self.predict(a, P, T, R, Q)
 
         # Use pinv, otherwise P_hat is singular when there is missing data
-        smoother_gain = matrix_dot(pt.linalg.pinv(P_hat, hermitian=True), T, P).T
+        smoother_gain = (pt.linalg.pinv(P_hat, hermitian=True) @ T @ P).mT
         a_smooth_next = a + smoother_gain @ (a_smooth - a_hat)
 
         P_smooth_next = P + quad_form_sym(smoother_gain, P_smooth - P_hat)

pymc_extras/statespace/filters/utilities.py

@@ -1,7 +1,5 @@
 import pytensor.tensor as pt
 
-from pytensor.tensor.nlinalg import matrix_dot
-
 from pymc_extras.statespace.utils.constants import JITTER_DEFAULT, NEVER_TIME_VARYING, VECTOR_VALUED
 
 
@@ -48,12 +46,11 @@ def split_vars_into_seq_and_nonseq(params, param_names):
 
 
 def stabilize(cov, jitter=JITTER_DEFAULT):
-    # Ensure diagonal is non-zero
     cov = cov + pt.identity_like(cov) * jitter
 
     return cov
 
 
 def quad_form_sym(A, B):
-    out = matrix_dot(A, B, A.T)
-    return 0.5 * (out + out.T)
+    out = A @ B @ A.mT
+    return 0.5 * (out + out.mT)

pymc_extras/statespace/models/DFM.py

@@ -5,7 +5,7 @@ import pytensor
 import pytensor.tensor as pt
 
 from pymc_extras.statespace.core.statespace import PyMCStateSpace
-from pymc_extras.statespace.models.utilities import make_default_coords
+from pymc_extras.statespace.models.utilities import make_default_coords, validate_names
 from pymc_extras.statespace.utils.constants import (
     ALL_STATE_AUX_DIM,
     ALL_STATE_DIM,
@@ -224,9 +224,7 @@ class BayesianDynamicFactor(PyMCStateSpace):
         self,
         k_factors: int,
         factor_order: int,
-        k_endog: int | None = None,
         endog_names: Sequence[str] | None = None,
-        k_exog: int | None = None,
         exog_names: Sequence[str] | None = None,
         shared_exog_states: bool = False,
         exog_innovations: bool = False,
@@ -249,19 +247,11 @@ class BayesianDynamicFactor(PyMCStateSpace):
             and are modeled as a white noise process, i.e., :math:`f_t = \varepsilon_{f,t}`.
             Therefore, the state vector will include one state per factor and "factor_ar" will not exist.
 
-        k_endog : int, optional
-            Number of observed time series. If not provided, the number of observed series will be inferred from `endog_names`.
-            At least one of `k_endog` or `endog_names` must be provided.
-
         endog_names : list of str, optional
-            Names of the observed time series. If not provided, default names will be generated as `endog_1`, `endog_2`, ..., `endog_k` based on `k_endog`.
-            At least one of `k_endog` or `endog_names` must be provided.
-
-        k_exog : int, optional
-            Number of exogenous variables. If not provided, the model will not have exogenous variables.
+            Names of the observed time series.
 
         exog_names : Sequence[str], optional
-            Names of the exogenous variables. If not provided, but `k_exog` is specified, default names will be generated as `exog_1`, `exog_2`, ..., `exog_k`.
+            Names of the exogenous variables.
 
         shared_exog_states: bool, optional
             Whether exogenous latent states are shared across the observed states. If True, there will be only one set of exogenous latent
@@ -289,13 +279,8 @@ class BayesianDynamicFactor(PyMCStateSpace):
 
         """
 
-        if k_endog is None and endog_names is None:
-            raise ValueError("Either k_endog or endog_names must be provided.")
-        if k_endog is None:
-            k_endog = len(endog_names)
-        if endog_names is None:
-            endog_names = [f"endog_{i}" for i in range(k_endog)]
-
+        validate_names(endog_names, var_name="endog_names", optional=False)
+        k_endog = len(endog_names)
         self.endog_names = endog_names
         self.k_endog = k_endog
         self.k_factors = k_factors
@@ -304,17 +289,17 @@ class BayesianDynamicFactor(PyMCStateSpace):
         self.error_var = error_var
         self.error_cov_type = error_cov_type
 
-        if k_exog is None and exog_names is None:
-            self.k_exog = 0
-        else:
+        if exog_names is not None:
             self.shared_exog_states = shared_exog_states
             self.exog_innovations = exog_innovations
-            if k_exog is None:
-                k_exog = len(exog_names) if exog_names is not None else 0
-            elif exog_names is None:
-                exog_names = [f"exog_{i}" for i in range(k_exog)] if k_exog > 0 else None
+            validate_names(
+                exog_names, var_name="exog_names", optional=True
+            )  # Not sure if this adds anything
+            k_exog = len(exog_names)
             self.k_exog = k_exog
             self.exog_names = exog_names
+        else:
+            self.k_exog = 0
 
         self.k_exog_states = self.k_exog * self.k_endog if not shared_exog_states else self.k_exog
         self.exog_flag = self.k_exog > 0