Expose multivariate normal method argument in post-estimation tasks (#484)

* Expose `method` argument to MvNormals used in statespace distributions when doing post-estimation tasks

* Use keyword arguments when calling post-estimation functions internally

* Fix typo, more clear argument name

* improve type hint

* Add `method` to `build_statespace_graph` and include some general advice.

* Fix incorrect signature in call to _sample_unconditional

Author: jessegrabowski

pymc_extras/statespace/core/statespace.py

@@ -1,7 +1,7 @@
 import logging
 
 from collections.abc import Callable, Sequence
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import pandas as pd
@@ -822,6 +822,7 @@ def build_statespace_graph(
         mode: str | None = None,
         missing_fill_value: float | None = None,
         cov_jitter: float | None = JITTER_DEFAULT,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         save_kalman_filter_outputs_in_idata: bool = False,
     ) -> None:
         """
@@ -865,6 +866,14 @@ def build_statespace_graph(
 
                 - The Univariate Filter is more robust than other filters, and can tolerate a lower jitter value
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         save_kalman_filter_outputs_in_idata: bool, optional, default=False
             If True, Kalman Filter outputs will be saved in the model as deterministics. Useful for debugging, but
             should not be necessary for the majority of users.
@@ -915,6 +924,7 @@ def build_statespace_graph(
             logp=logp,
             observed=data,
             dims=obs_dims,
+            method=mvn_method,
         )
 
         self._fit_coords = pm_mod.coords.copy()
@@ -1109,6 +1119,7 @@ def _sample_conditional(
         group: str,
         random_seed: RandomState | None = None,
         data: pt.TensorLike | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         **kwargs,
     ):
         """
@@ -1130,6 +1141,14 @@ def _sample_conditional(
             Observed data on which to condition the model. If not provided, the function will use the data that was
             provided when the model was built.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -1181,6 +1200,7 @@ def _sample_conditional(
                     covs=cov,
                     logp=dummy_ll,
                     dims=state_dims,
+                    method=mvn_method,
                 )
 
                 obs_mu = (Z @ mu[..., None]).squeeze(-1)
@@ -1192,6 +1212,7 @@ def _sample_conditional(
                     covs=obs_cov,
                     logp=dummy_ll,
                     dims=obs_dims,
+                    method=mvn_method,
                 )
 
         # TODO: Remove this after pm.Flat initial values are fixed
@@ -1222,6 +1243,7 @@ def _sample_unconditional(
         steps: int | None = None,
         use_data_time_dim: bool = False,
         random_seed: RandomState | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         **kwargs,
     ):
         """
@@ -1251,6 +1273,14 @@ def _sample_unconditional(
         random_seed : int, RandomState or Generator, optional
             Seed for the random number generator.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -1309,6 +1339,7 @@ def _sample_unconditional(
                 steps=steps,
                 dims=dims,
                 mode=self._fit_mode,
+                method=mvn_method,
                 sequence_names=self.kalman_filter.seq_names,
                 k_endog=self.k_endog,
             )
@@ -1331,7 +1362,11 @@ def _sample_unconditional(
         return idata_unconditional.posterior_predictive
 
     def sample_conditional_prior(
-        self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs
+        self,
+        idata: InferenceData,
+        random_seed: RandomState | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
+        **kwargs,
     ) -> InferenceData:
         """
         Sample from the conditional prior; that is, given parameter draws from the prior distribution,
@@ -1347,6 +1382,14 @@ def sample_conditional_prior(
         random_seed : int, RandomState or Generator, optional
             Seed for the random number generator.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -1358,10 +1401,16 @@ def sample_conditional_prior(
              "predicted_prior", and "smoothed_prior".
         """
 
-        return self._sample_conditional(idata, "prior", random_seed, **kwargs)
+        return self._sample_conditional(
+            idata=idata, group="prior", random_seed=random_seed, mvn_method=mvn_method, **kwargs
+        )
 
     def sample_conditional_posterior(
-        self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs
+        self,
+        idata: InferenceData,
+        random_seed: RandomState | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
+        **kwargs,
     ):
         """
         Sample from the conditional posterior; that is, given parameter draws from the posterior distribution,
@@ -1376,6 +1425,14 @@ def sample_conditional_posterior(
         random_seed : int, RandomState or Generator, optional
             Seed for the random number generator.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -1387,14 +1444,17 @@ def sample_conditional_posterior(
              "predicted_posterior", and "smoothed_posterior".
         """
 
-        return self._sample_conditional(idata, "posterior", random_seed, **kwargs)
+        return self._sample_conditional(
+            idata=idata, group="posterior", random_seed=random_seed, mvn_method=mvn_method, **kwargs
+        )
 
     def sample_unconditional_prior(
         self,
         idata: InferenceData,
         steps: int | None = None,
         use_data_time_dim: bool = False,
         random_seed: RandomState | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         **kwargs,
     ) -> InferenceData:
         """
@@ -1423,6 +1483,14 @@ def sample_unconditional_prior(
         random_seed : int, RandomState or Generator, optional
             Seed for the random number generator.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -1439,7 +1507,13 @@ def sample_unconditional_prior(
         """
 
         return self._sample_unconditional(
-            idata, "prior", steps, use_data_time_dim, random_seed, **kwargs
+            idata=idata,
+            group="prior",
+            steps=steps,
+            use_data_time_dim=use_data_time_dim,
+            random_seed=random_seed,
+            mvn_method=mvn_method,
+            **kwargs,
         )
 
     def sample_unconditional_posterior(
@@ -1448,6 +1522,7 @@ def sample_unconditional_posterior(
         steps: int | None = None,
         use_data_time_dim: bool = False,
         random_seed: RandomState | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         **kwargs,
     ) -> InferenceData:
         """
@@ -1477,6 +1552,14 @@ def sample_unconditional_posterior(
         random_seed : int, RandomState or Generator, optional
             Seed for the random number generator.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         Returns
         -------
         InferenceData
@@ -1490,7 +1573,13 @@ def sample_unconditional_posterior(
         """
 
         return self._sample_unconditional(
-            idata, "posterior", steps, use_data_time_dim, random_seed, **kwargs
+            idata=idata,
+            group="posterior",
+            steps=steps,
+            use_data_time_dim=use_data_time_dim,
+            random_seed=random_seed,
+            mvn_method=mvn_method,
+            **kwargs,
         )
 
     def sample_statespace_matrices(
@@ -1933,6 +2022,7 @@ def forecast(
         filter_output="smoothed",
         random_seed: RandomState | None = None,
         verbose: bool = True,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         **kwargs,
     ) -> InferenceData:
         """
@@ -1989,6 +2079,14 @@ def forecast(
         verbose: bool, default=True
             Whether to print diagnostic information about forecasting.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -2098,6 +2196,7 @@ def forecast(
                 sequence_names=self.kalman_filter.seq_names,
                 k_endog=self.k_endog,
                 append_x0=False,
+                method=mvn_method,
             )
 
         forecast_model.rvs_to_initial_values = {
@@ -2126,6 +2225,7 @@ def impulse_response_function(
         shock_trajectory: np.ndarray | None = None,
         orthogonalize_shocks: bool = False,
         random_seed: RandomState | None = None,
+        mvn_method: Literal["cholesky", "eigh", "svd"] = "svd",
         **kwargs,
     ):
         """
@@ -2177,6 +2277,14 @@ def impulse_response_function(
         random_seed : int, RandomState or Generator, optional
             Seed for the random number generator.
 
+        mvn_method: str, default "svd"
+            Method used to invert the covariance matrix when calculating the pdf of a multivariate normal
+            (or when generating samples). One of "cholesky", "eigh", or "svd". "cholesky" is fastest, but least robust
+            to ill-conditioned matrices, while "svd" is slow but extremely robust.
+
+            In general, if your model has measurement error, "cholesky" will be safe to use. Otherwise, "svd" is
+            recommended. "eigh" can also be tried if sampling with "svd" is very slow, but it is not as robust as "svd".
+
         kwargs:
             Additional keyword arguments are passed to pymc.sample_posterior_predictive
 
@@ -2236,7 +2344,7 @@ def impulse_response_function(
                 shock_trajectory = pt.zeros((n_steps, self.k_posdef))
                 if Q is not None:
                     init_shock = pm.MvNormal(
-                        "initial_shock", mu=0, cov=Q, dims=[SHOCK_DIM], method="svd"
+                        "initial_shock", mu=0, cov=Q, dims=[SHOCK_DIM], method=mvn_method
                     )
                 else:
                     init_shock = pm.Deterministic(