docstrings

pymc_experimental/statespace/core/representation.py

@@ -41,6 +41,53 @@ def _preprocess_data(data: Union[DataFrame, np.ndarray], expected_dims=3):
 
 
 class PytensorRepresentation:
+    r"""
+    Core class to hold all objects required by linear gaussian statespace models
+
+    Shamelessly copied from the Statsmodels.api implementationfound here:
+
+    https://github.com/statsmodels/statsmodels/blob/main/statsmodels/tsa/statespace/representation.py
+
+    Parameters
+    ----------
+    k_endog: int
+        Number of observed states (called "endogeous states" in statsmodels)
+    k_states: int
+        Number of hidden states
+    k_posdef: int
+        Number of states that have exogenous shocks; also the rank of the selection matrix R.
+    design: ArrayLike, optional
+        Design matrix, denoted 'Z' in [1].
+    obs_intercept: ArrayLike, optional
+        Constant vector in the observation equation, denoted 'd' in [1]. Currently
+        not used.
+    obs_cov: ArrayLike, optional
+        Covariance matrix for multivariate-normal errors in the observation equation. Denoted 'H' in
+        [1].
+    transition: ArrayLike, optional
+        Transition equation that updates the hidden state between time-steps. Denoted 'T' in [1].
+    state_intercept: ArrayLike, optional
+        Constant vector for the observation equation, denoted 'c' in [1]. Currently not used.
+    selection: ArrayLike, optional
+        Selection matrix that matches shocks to hidden states, denoted 'R' in [1]. This is the identity
+        matrix when k_posdef = k_states.
+    state_cov: ArrayLike, optional
+        Covariance matrix for state equations, denoted 'Q' in [1]. Null matrix when there is no observation
+        noise.
+    initial_state: ArrayLike, optional
+        Experimental setting to allow for Bayesian estimation of the initial state, denoted `alpha_0` in [1]. Default
+        It should potentially be removed in favor of the closed-form diffuse initialization.
+    initial_state_cov: ArrayLike, optional
+        Experimental setting to allow for Bayesian estimation of the initial state, denoted `P_0` in [1]. Default
+        It should potentially be removed in favor of the closed-form diffuse initialization.
+
+    References
+    ----------
+    .. [1] Durbin, James, and Siem Jan Koopman. 2012.
+        Time Series Analysis by State Space Methods: Second Edition.
+        Oxford University Press.
+    """
+
     def __init__(
         self,
         k_endog: int,
@@ -56,51 +103,6 @@ def __init__(
         initial_state=None,
         initial_state_cov=None,
     ) -> None:
-        """
-        A representation of a State Space model, in Pytensor. Shamelessly copied from the Statsmodels.api implementation
-        found here:
-
-        https://github.com/statsmodels/statsmodels/blob/main/statsmodels/tsa/statespace/representation.py
-
-        Parameters
-        ----------
-        k_endog: int
-            Number of observed states (called "endogeous states" in statsmodels)
-        k_states: int
-            Number of hidden states
-        k_posdef: int
-            Number of states that have exogenous shocks; also the rank of the selection matrix R.
-        design: ArrayLike, optional
-            Design matrix, denoted 'Z' in [1].
-        obs_intercept: ArrayLike, optional
-            Constant vector in the observation equation, denoted 'd' in [1]. Currently
-            not used.
-        obs_cov: ArrayLike, optional
-            Covariance matrix for multivariate-normal errors in the observation equation. Denoted 'H' in
-            [1].
-        transition: ArrayLike, optional
-            Transition equation that updates the hidden state between time-steps. Denoted 'T' in [1].
-        state_intercept: ArrayLike, optional
-            Constant vector for the observation equation, denoted 'c' in [1]. Currently not used.
-        selection: ArrayLike, optional
-            Selection matrix that matches shocks to hidden states, denoted 'R' in [1]. This is the identity
-            matrix when k_posdef = k_states.
-        state_cov: ArrayLike, optional
-            Covariance matrix for state equations, denoted 'Q' in [1]. Null matrix when there is no observation
-            noise.
-        initial_state: ArrayLike, optional
-            Experimental setting to allow for Bayesian estimation of the initial state, denoted `alpha_0` in [1]. Default
-            It should potentially be removed in favor of the closed-form diffuse initialization.
-        initial_state_cov: ArrayLike, optional
-            Experimental setting to allow for Bayesian estimation of the initial state, denoted `P_0` in [1]. Default
-            It should potentially be removed in favor of the closed-form diffuse initialization.
-
-        References
-        ----------
-        .. [1] Durbin, James, and Siem Jan Koopman. 2012.
-            Time Series Analysis by State Space Methods: Second Edition.
-            Oxford University Press.
-        """
         self.k_states = k_states
         self.k_endog = k_endog
         self.k_posdef = k_posdef if k_posdef is not None else k_states

pymc_experimental/statespace/core/statespace.py

@@ -67,6 +67,28 @@ def _verify_group(group):
 
 
 class PyMCStateSpace:
+    r"""
+    Base class for Linear Gaussian Statespace models in PyMC.
+
+    Parameters
+    ----------
+    k_endog : int
+        The number of endogenous variables (observed time series).
+
+    k_states : int
+        The number of state variables.
+
+    k_posdef : int
+        The number of shocks in the model
+
+    filter_type : str, optional
+        The type of Kalman filter to use. Valid options are "standard", "univariate", "single", "cholesky", and
+        "steady_state". For more information, see the docs for each filter. Default is "standard".
+
+    verbose : bool, optional
+        If True, displays information about the initialized model. Defaults to True.
+    """
+
     def __init__(
         self,
         k_endog: int,
@@ -75,29 +97,6 @@ def __init__(
         filter_type: str = "standard",
         verbose: bool = True,
     ):
-
-        """
-        Base class for Linear Gaussian Statespace models in PyMC.
-
-        Parameters
-        ----------
-        k_endog : int
-            The number of endogenous variables (observed time series).
-
-        k_states : int
-            The number of state variables.
-
-        k_posdef : int
-            The number of shocks in the model
-
-        filter_type : str, optional
-            The type of Kalman filter to use. Valid options are "standard", "univariate", "single", "cholesky", and
-            "steady_state". For more information, see the docs for each filter. Default is "standard".
-
-        verbose : bool, optional
-            If True, displays information about the initialized model. Defaults to True.
-        """
-
         self._fit_mode = None
         self._fit_coords = None
         self._prior_mod = pm.Model()

pymc_experimental/statespace/filters/distributions.py

@@ -135,6 +135,11 @@ def step_fn(*args):
 
 
 class LinearGaussianStateSpace:
+    """
+    Linear Gaussian Statespace distribution
+
+    """
+
     def __new__(cls, name, a0, P0, c, d, T, Z, R, H, Q, *, steps=None, mode=None, **kwargs):
         dims = kwargs.pop("dims", None)
         latent_dims = None

pymc_experimental/statespace/filters/kalman_filter.py

@@ -31,7 +31,7 @@
 class BaseFilter(ABC):
     def __init__(self, mode=None):
         """
-        Kalman Filter Abstract Class.
+        Kalman Filter.
 
         Parameters
         ----------

pymc_experimental/statespace/filters/kalman_smoother.py

@@ -11,6 +11,11 @@
 
 
 class KalmanSmoother:
+    """
+    Kalman Smoother
+
+    """
+
     def __init__(self, mode: Optional[str] = None):
         self.mode = mode
         self.seq_names = []

pymc_experimental/statespace/models/VARMAX.py

@@ -22,6 +22,10 @@
 
 
 class BayesianVARMAX(PyMCStateSpace):
+    r"""
+    Vector AutoRegressive Moving Average with eXogenous Regressors
+    """
+
     def __init__(
         self,
         order: Tuple[int, int],

pymc_experimental/statespace/models/local_level.py

@@ -20,6 +20,11 @@
 
 
 class BayesianLocalLevel(PyMCStateSpace):
+    r"""
+    Local Level Model
+
+    """
+
     def __init__(self, verbose=False, filter_type: str = "standard"):
         k_states = 2
         k_posdef = 2