pymc-extras 0.4.1__py3-none-any.whl → 0.6.0__py3-none-any.whl

pymc_extras/model_builder.py

@@ -334,7 +334,9 @@ class ModelBuilder:
         >>> model = MyModel(ModelBuilder)
         >>> idata = az.InferenceData(your_dataset)
         >>> model.set_idata_attrs(idata=idata)
-        >>> assert "id" in idata.attrs #this and the following lines are part of doctest, not user manual
+        >>> assert (
+        ...     "id" in idata.attrs
+        ... )  # this and the following lines are part of doctest, not user manual
         >>> assert "model_type" in idata.attrs
         >>> assert "version" in idata.attrs
         >>> assert "sampler_config" in idata.attrs
@@ -382,7 +384,7 @@ class ModelBuilder:
         >>>         super().__init__()
         >>> model = MyModel()
         >>> model.fit(data)
-        >>> model.save('model_results.nc')  # This will call the overridden method in MyModel
+        >>> model.save("model_results.nc")  # This will call the overridden method in MyModel
         """
         if self.idata is not None and "posterior" in self.idata:
             file = Path(str(fname))
@@ -432,7 +434,7 @@ class ModelBuilder:
         --------
         >>> class MyModel(ModelBuilder):
         >>>     ...
-        >>> name = './mymodel.nc'
+        >>> name = "./mymodel.nc"
         >>> imported_model = MyModel.load(name)
         """
         filepath = Path(str(fname))
@@ -444,6 +446,7 @@ class ModelBuilder:
             sampler_config=json.loads(idata.attrs["sampler_config"]),
         )
         model.idata = idata
+        model.is_fitted_ = True
         dataset = idata.fit_data.to_dataframe()
         X = dataset.drop(columns=[model.output_var])
         y = dataset[model.output_var]
@@ -524,6 +527,8 @@ class ModelBuilder:
             )
             self.idata.add_groups(fit_data=combined_data.to_xarray())  # type: ignore
 
+        self.is_fitted_ = True
+
         return self.idata  # type: ignore
 
     def predict(
@@ -554,7 +559,7 @@ class ModelBuilder:
         >>> model = MyModel()
         >>> idata = model.fit(data)
         >>> x_pred = []
-        >>> prediction_data = pd.DataFrame({'input':x_pred})
+        >>> prediction_data = pd.DataFrame({"input": x_pred})
         >>> pred_mean = model.predict(prediction_data)
         """
 

pymc_extras/prior.py

@@ -70,8 +70,10 @@ Create a prior with a custom transform function by registering it with
 
     from pymc_extras.prior import register_tensor_transform
 
+
     def custom_transform(x):
-        return x ** 2
+        return x**2
+
 
     register_tensor_transform("square", custom_transform)
 
@@ -138,9 +140,23 @@ def handle_dims(x: pt.TensorLike, dims: Dims, desired_dims: Dims) -> pt.TensorVa
 
     Doesn't check for validity of the dims
 
+    Parameters
+    ----------
+    x : pt.TensorLike
+        The tensor to align.
+    dims : Dims
+        The current dimensions of the tensor.
+    desired_dims : Dims
+        The desired dimensions of the tensor.
+
+    Returns
+    -------
+    pt.TensorVariable
+        The aligned tensor.
+
     Examples
     --------
-    1D to 2D with new dim
+    Handle transpose 1D to 2D with new dimension.
 
     .. code-block:: python
 
@@ -177,10 +193,40 @@ def handle_dims(x: pt.TensorLike, dims: Dims, desired_dims: Dims) -> pt.TensorVa
 
 
 DimHandler = Callable[[pt.TensorLike, Dims], pt.TensorLike]
+"""A function that takes a tensor and its current dims and makes it broadcastable to the desired dims."""
 
 
 def create_dim_handler(desired_dims: Dims) -> DimHandler:
-    """Wrap the `handle_dims` function to act like the previous `create_dim_handler` function."""
+    """Wrap the :func:`handle_dims` function to always use the same desired_dims.
+
+    Parameters
+    ----------
+    desired_dims : Dims
+        The desired dimensions to align to.
+
+    Returns
+    -------
+    DimHandler
+        A function that takes a tensor and its current dims and aligns it to
+        the desired dims.
+
+
+    Examples
+    --------
+    Create a dim handler to align to ("channel", "group").
+
+    .. code-block:: python
+
+        import numpy as np
+
+        from pymc_extras.prior import create_dim_handler
+
+        dim_handler = create_dim_handler(("channel", "group"))
+
+        result = dim_handler(np.array([1, 2, 3]), dims="channel")
+
+
+    """
 
     def func(x: pt.TensorLike, dims: Dims) -> pt.TensorVariable:
         return handle_dims(x, dims, desired_dims)
@@ -228,8 +274,10 @@ def register_tensor_transform(name: str, transform: Transform) -> None:
             register_tensor_transform,
         )
 
+
         def custom_transform(x):
-            return x ** 2
+            return x**2
+
 
         register_tensor_transform("square", custom_transform)
 
@@ -268,9 +316,50 @@ def _get_pymc_parameters(distribution: pm.Distribution) -> set[str]:
 
 @runtime_checkable
 class VariableFactory(Protocol):
-    """Protocol for something that works like a Prior class."""
+    '''Protocol for something that works like a Prior class.
+
+    Sample with :func:`sample_prior`.
+
+    Examples
+    --------
+    Create a custom variable factory.
+
+    .. code-block:: python
+
+        import pymc as pm
+
+        import pytensor.tensor as pt
+
+        from pymc_extras.prior import sample_prior, VariableFactory
+
+
+        class PowerSumDistribution:
+            """Create a distribution that is the sum of powers of a base distribution."""
+
+            def __init__(self, distribution: VariableFactory, n: int):
+                self.distribution = distribution
+                self.n = n
+
+            @property
+            def dims(self):
+                return self.distribution.dims
+
+            def create_variable(self, name: str) -> "TensorVariable":
+                raw = self.distribution.create_variable(f"{name}_raw")
+                return pm.Deterministic(
+                    name,
+                    pt.sum([raw**n for n in range(1, self.n + 1)], axis=0),
+                    dims=self.dims,
+                )
+
+
+        cubic = PowerSumDistribution(Prior("Normal"), n=3)
+        samples = sample_prior(cubic)
+
+    '''
 
     dims: tuple[str, ...]
+    """The dimensions of the variable to create."""
 
     def create_variable(self, name: str) -> pt.TensorVariable:
         """Create a TensorVariable."""
@@ -316,6 +405,7 @@ def sample_prior(
 
         from pymc_extras.prior import sample_prior
 
+
         class CustomVariableDefinition:
             def __init__(self, dims, n: int):
                 self.dims = dims
@@ -323,7 +413,8 @@ def sample_prior(
 
             def create_variable(self, name: str) -> "TensorVariable":
                 x = pm.Normal(f"{name}_x", mu=0, sigma=1, dims=self.dims)
-                return pt.sum([x ** n for n in range(1, self.n + 1)], axis=0)
+                return pt.sum([x**n for n in range(1, self.n + 1)], axis=0)
+
 
         cubic = CustomVariableDefinition(dims=("channel",), n=3)
         coords = {"channel": ["C1", "C2", "C3"]}
@@ -381,6 +472,82 @@ class Prior:
         be registered with `register_tensor_transform` function or
         be available in either `pytensor.tensor` or `pymc.math`.
 
+    Examples
+    --------
+    Create a normal prior.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import Prior
+
+        normal = Prior("Normal")
+
+    Create a hierarchical normal prior by using distributions for the parameters
+    and specifying the dims.
+
+    .. code-block:: python
+
+        hierarchical_normal = Prior(
+            "Normal",
+            mu=Prior("Normal"),
+            sigma=Prior("HalfNormal"),
+            dims="channel",
+        )
+
+    Create a non-centered hierarchical normal prior with the `centered` parameter.
+
+    .. code-block:: python
+
+        non_centered_hierarchical_normal = Prior(
+            "Normal",
+            mu=Prior("Normal"),
+            sigma=Prior("HalfNormal"),
+            dims="channel",
+            # Only change needed to make it non-centered
+            centered=False,
+        )
+
+    Create a hierarchical beta prior by using Beta distribution, distributions for
+    the parameters, and specifying the dims.
+
+    .. code-block:: python
+
+        hierarchical_beta = Prior(
+            "Beta",
+            alpha=Prior("HalfNormal"),
+            beta=Prior("HalfNormal"),
+            dims="channel",
+        )
+
+    Create a transformed hierarchical normal prior by using the `transform`
+    parameter. Here the "sigmoid" transformation comes from `pm.math`.
+
+    .. code-block:: python
+
+        transformed_hierarchical_normal = Prior(
+            "Normal",
+            mu=Prior("Normal"),
+            sigma=Prior("HalfNormal"),
+            transform="sigmoid",
+            dims="channel",
+        )
+
+    Create a prior with a custom transform function by registering it with
+    :func:`register_tensor_transform`.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import register_tensor_transform
+
+
+        def custom_transform(x):
+            return x**2
+
+
+        register_tensor_transform("square", custom_transform)
+
+        custom_distribution = Prior("Normal", transform="square")
+
     """
 
     # Taken from https://en.wikipedia.org/wiki/Location%E2%80%93scale_family
@@ -389,9 +556,13 @@ class Prior:
         "StudentT": {"mu": 0, "sigma": 1},
         "ZeroSumNormal": {"sigma": 1},
     }
+    """Available non-centered distributions and their default parameters."""
 
     pymc_distribution: type[pm.Distribution]
+    """The PyMC distribution class."""
+
     pytensor_transform: Callable[[pt.TensorLike], pt.TensorLike] | None
+    """The PyTensor transform function."""
 
     @validate_call
     def __init__(
@@ -1317,9 +1488,33 @@ class Censored:
 
 
 class Scaled:
-    """Scaled distribution for numerical stability."""
+    """Scaled distribution for numerical stability.
+
+    This is the same as multiplying the variable by a constant factor.
+
+    Parameters
+    ----------
+    dist : Prior
+        The prior distribution to scale.
+    factor : pt.TensorLike
+        The scaling factor. This will have to be broadcastable to the
+        dimensions of the distribution.
+
+    Examples
+    --------
+    Create a scaled normal distribution.
+
+    .. code-block:: python
+
+        from pymc_extras.prior import Prior, Scaled
+
+        normal = Prior("Normal", mu=0, sigma=1)
+        # Same as Normal(mu=0, sigma=10)
+        scaled_normal = Scaled(normal, factor=10)
+
+    """
 
-    def __init__(self, dist: Prior, factor: float | pt.TensorVariable) -> None:
+    def __init__(self, dist: Prior, factor: pt.TensorLike) -> None:
         self.dist = dist
         self.factor = factor
 

pymc_extras/statespace/core/compile.py

@@ -28,7 +28,7 @@ def compile_statespace(
         x0, P0, c, d, T, Z, R, H, Q, steps=steps, sequence_names=sequence_names
     )
 
-    inputs = list(pytensor.graph.basic.explicit_graph_inputs(outputs))
+    inputs = list(pytensor.graph.traversal.explicit_graph_inputs(outputs))
 
     _f = pm.compile(inputs, outputs, on_unused_input="ignore", **compile_kwargs)
 

pymc_extras/statespace/filters/kalman_filter.py

@@ -200,7 +200,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
         )
@@ -393,7 +393,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 +580,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 +630,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
 
@@ -658,14 +658,14 @@ class SquareRootFilter(BaseFilter):
         # 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 +677,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
@@ -800,7 +801,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,
@@ -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)