pymc-extras 0.3.1__tar.gz → 0.4.0__tar.gz

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.3.1
+Version: 0.4.0
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Project-URL: Documentation, https://pymc-extras.readthedocs.io/
 Project-URL: Repository, https://github.com/pymc-devs/pymc-extras.git
@@ -232,9 +232,11 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Requires-Python: >=3.11
-Requires-Dist: better-optimize>=0.1.4
+Requires-Dist: better-optimize>=0.1.5
+Requires-Dist: preliz>=0.20.0
 Requires-Dist: pydantic>=2.0.0
-Requires-Dist: pymc>=5.21.1
+Requires-Dist: pymc>=5.24.1
+Requires-Dist: pytensor>=2.31.4
 Requires-Dist: scikit-learn
 Provides-Extra: complete
 Requires-Dist: dask[complete]<2025.1.1; extra == 'complete'
@@ -245,7 +247,6 @@ Requires-Dist: xhistogram; extra == 'dask-histogram'
 Provides-Extra: dev
 Requires-Dist: blackjax; extra == 'dev'
 Requires-Dist: dask[all]<2025.1.1; extra == 'dev'
-Requires-Dist: preliz>=0.5.0; extra == 'dev'
 Requires-Dist: pytest-mock; extra == 'dev'
 Requires-Dist: pytest>=6.0; extra == 'dev'
 Requires-Dist: statsmodels; extra == 'dev'

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3.1'
-__version_tuple__ = version_tuple = (0, 3, 1)
+__version__ = version = '0.4.0'
+__version_tuple__ = version_tuple = (0, 4, 0)

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/conda-envs/environment-test.yml

@@ -3,9 +3,10 @@ channels:
 - conda-forge
 - nodefaults
 dependencies:
-  - pymc>=5.21.1
+  - pymc>=5.24.1
+  - pytensor>=2.31.4
   - scikit-learn
-  - better-optimize>=0.1.2
+  - better-optimize>=0.1.5
   - dask<2025.1.1
   - xhistogram
   - statsmodels
@@ -13,7 +14,7 @@ dependencies:
   - pytest
   - pytest-cov
   - pydantic>=2.0.0
-  - preliz>=0.5.0
+  - preliz>=0.20.0
   - pip
   - pip:
       - jax

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/pymc_extras/inference/laplace_approx/find_map.py

@@ -326,7 +326,7 @@ def find_MAP(
     )
 
     raveled_optimized = RaveledVars(optimizer_result.x, initial_params.point_map_info)
-    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed)
+    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     unobserved_vars_values = model.compile_fn(unobserved_vars, mode="FAST_COMPILE")(
         DictToArrayBijection.rmap(raveled_optimized)
     )
@@ -335,13 +335,20 @@ def find_MAP(
         var.name: value for var, value in zip(unobserved_vars, unobserved_vars_values)
     }
 
-    idata = map_results_to_inference_data(optimized_point, frozen_model)
-    idata = add_fit_to_inference_data(idata, raveled_optimized, H_inv)
+    idata = map_results_to_inference_data(
+        map_point=optimized_point, model=frozen_model, include_transformed=include_transformed
+    )
+
+    idata = add_fit_to_inference_data(
+        idata=idata, mu=raveled_optimized, H_inv=H_inv, model=frozen_model
+    )
+
     idata = add_optimizer_result_to_inference_data(
-        idata, optimizer_result, method, raveled_optimized, model
+        idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model
     )
+
     idata = add_data_to_inference_data(
-        idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
+        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
     )
 
     return idata

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/pymc_extras/inference/laplace_approx/idata.py

@@ -59,6 +59,7 @@ def make_unpacked_variable_names(names: list[str], model: pm.Model) -> list[str]
 def map_results_to_inference_data(
     map_point: dict[str, float | int | np.ndarray],
     model: pm.Model | None = None,
+    include_transformed: bool = True,
 ):
     """
     Add the MAP point to an InferenceData object in the posterior group.
@@ -68,13 +69,13 @@ def map_results_to_inference_data(
 
     Parameters
     ----------
-    idata: az.InferenceData
-        An InferenceData object to which the MAP point will be added.
     map_point: dict
         A dictionary containing the MAP point estimates for each variable. The keys should be the variable names, and
         the values should be the corresponding MAP estimates.
     model: Model, optional
         A PyMC model. If None, the model is taken from the current model context.
+    include_transformed: bool
+        Whether to return transformed (unconstrained) variables in the constrained_posterior group. Default is True.
 
     Returns
     -------
@@ -118,7 +119,7 @@ def map_results_to_inference_data(
         dims=dims,
     )
 
-    if unconstrained_names:
+    if unconstrained_names and include_transformed:
         unconstrained_posterior = az.from_dict(
             posterior={
                 k: np.expand_dims(v, (0, 1))

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/pymc_extras/inference/laplace_approx/laplace.py

@@ -302,7 +302,7 @@ def fit_laplace(
     ----------
     model : pm.Model
         The PyMC model to be fit. If None, the current model context is used.
-    method : str
+    optimize_method : str
         The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
         trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.
 
@@ -441,9 +441,11 @@ def fit_laplace(
             .rename({"temp_chain": "chain", "temp_draw": "draw"})
         )
 
-        idata.unconstrained_posterior = unstack_laplace_draws(
-            new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
-        )
+        if include_transformed:
+            idata.unconstrained_posterior = unstack_laplace_draws(
+                new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
+            )
+
         idata.posterior = new_posterior.drop_vars(
             ["laplace_approximation", "unpacked_variable_names"]
         )

{pymc_extras-0.3.1 → pymc_extras-0.4.0}/pymc_extras/inference/pathfinder/pathfinder.py

@@ -38,16 +38,15 @@ from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import modelcontext
 from pymc.model.core import Point
+from pymc.progress_bar import CustomProgress, default_progress_theme
 from pymc.pytensorf import (
     compile,
     find_rng_nodes,
     reseed_rngs,
 )
 from pymc.util import (
-    CustomProgress,
     RandomSeed,
     _get_seeds_per_chain,
-    default_progress_theme,
     get_default_varnames,
 )
 from pytensor.compile.function.types import Function

pymc_extras-0.4.0/pymc_extras/statespace/models/structural/__init__.py

@@ -0,0 +1,21 @@
+from pymc_extras.statespace.models.structural.components.autoregressive import (
+    AutoregressiveComponent,
+)
+from pymc_extras.statespace.models.structural.components.cycle import CycleComponent
+from pymc_extras.statespace.models.structural.components.level_trend import LevelTrendComponent
+from pymc_extras.statespace.models.structural.components.measurement_error import MeasurementError
+from pymc_extras.statespace.models.structural.components.regression import RegressionComponent
+from pymc_extras.statespace.models.structural.components.seasonality import (
+    FrequencySeasonality,
+    TimeSeasonality,
+)
+
+__all__ = [
+    "LevelTrendComponent",
+    "MeasurementError",
+    "AutoregressiveComponent",
+    "TimeSeasonality",
+    "FrequencySeasonality",
+    "RegressionComponent",
+    "CycleComponent",
+]

pymc_extras-0.4.0/pymc_extras/statespace/models/structural/components/autoregressive.py

@@ -0,0 +1,188 @@
+import numpy as np
+import pytensor.tensor as pt
+
+from pymc_extras.statespace.models.structural.core import Component
+from pymc_extras.statespace.models.structural.utils import order_to_mask
+from pymc_extras.statespace.utils.constants import AR_PARAM_DIM
+
+
+class AutoregressiveComponent(Component):
+    r"""
+    Autoregressive timeseries component
+
+    Parameters
+    ----------
+    order: int or sequence of int
+
+        If int, the number of lags to include in the model.
+        If a sequence, an array-like of zeros and ones indicating which lags to include in the model.
+
+    name: str, default "auto_regressive"
+        A name for this autoregressive component. Used to label dimensions and coordinates.
+
+    observed_state_names: list[str] | None, default None
+        List of strings for observed state labels. If None, defaults to ["data"].
+
+    Notes
+    -----
+    An autoregressive component can be thought of as a way o introducing serially correlated errors into the model.
+    The process is modeled:
+
+    .. math::
+        x_t = \sum_{i=1}^p \rho_i x_{t-i}
+
+    Where ``p``, the number of autoregressive terms to model, is the order of the process. By default, all lags up to
+    ``p`` are included in the model. To disable lags, pass a list of zeros and ones to the ``order`` argumnet. For
+    example, ``order=[1, 1, 0, 1]`` would become:
+
+    .. math::
+        x_t = \rho_1 x_{t-1} + \rho_2 x_{t-1} + \rho_4 x_{t-1}
+
+    The coefficient :math:`\rho_3` has been constrained to zero.
+
+    .. warning:: This class is meant to be used as a component in a structural time series model. For modeling of
+              stationary processes with ARIMA, use ``statespace.BayesianSARIMA``.
+
+    Examples
+    --------
+    Model a timeseries as an AR(2) process with non-zero mean:
+
+    .. code:: python
+
+        from pymc_extras.statespace import structural as st
+        import pymc as pm
+        import pytensor.tensor as pt
+
+        trend = st.LevelTrendComponent(order=1, innovations_order=0)
+        ar = st.AutoregressiveComponent(2)
+        ss_mod = (trend + ar).build()
+
+        with pm.Model(coords=ss_mod.coords) as model:
+            P0 = pm.Deterministic('P0', pt.eye(ss_mod.k_states) * 10, dims=ss_mod.param_dims['P0'])
+            intitial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
+            ar_params = pm.Normal('ar_params', dims=ss_mod.param_dims['ar_params'])
+            sigma_ar = pm.Exponential('sigma_ar', 1, dims=ss_mod.param_dims['sigma_ar'])
+
+            ss_mod.build_statespace_graph(data)
+            idata = pm.sample(nuts_sampler='numpyro')
+
+    """
+
+    def __init__(
+        self,
+        order: int = 1,
+        name: str = "auto_regressive",
+        observed_state_names: list[str] | None = None,
+    ):
+        if observed_state_names is None:
+            observed_state_names = ["data"]
+
+        k_posdef = k_endog = len(observed_state_names)
+
+        order = order_to_mask(order)
+        ar_lags = np.flatnonzero(order).ravel().astype(int) + 1
+        k_states = len(order)
+
+        self.order = order
+        self.ar_lags = ar_lags
+
+        super().__init__(
+            name=name,
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef,
+            measurement_error=True,
+            combine_hidden_states=True,
+            observed_state_names=observed_state_names,
+            obs_state_idxs=np.tile(np.r_[[1.0], np.zeros(k_states - 1)], k_endog),
+        )
+
+    def populate_component_properties(self):
+        k_states = self.k_states // self.k_endog  # this is also the number of AR lags
+
+        self.state_names = [
+            f"L{i + 1}[{state_name}]"
+            for state_name in self.observed_state_names
+            for i in range(k_states)
+        ]
+
+        self.shock_names = [f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names]
+        self.param_names = [f"params_{self.name}", f"sigma_{self.name}"]
+        self.param_dims = {f"params_{self.name}": (f"lag_{self.name}",)}
+        self.coords = {f"lag_{self.name}": self.ar_lags.tolist()}
+
+        if self.k_endog > 1:
+            self.param_dims[f"params_{self.name}"] = (
+                f"endog_{self.name}",
+                AR_PARAM_DIM,
+            )
+            self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
+
+            self.coords[f"endog_{self.name}"] = self.observed_state_names
+
+        self.param_info = {
+            f"params_{self.name}": {
+                "shape": (k_states,) if self.k_endog == 1 else (self.k_endog, k_states),
+                "constraints": None,
+                "dims": (AR_PARAM_DIM,)
+                if self.k_endog == 1
+                else (
+                    f"endog_{self.name}",
+                    f"lag_{self.name}",
+                ),
+            },
+            f"sigma_{self.name}": {
+                "shape": () if self.k_endog == 1 else (self.k_endog,),
+                "constraints": "Positive",
+                "dims": None if self.k_endog == 1 else (f"endog_{self.name}",),
+            },
+        }
+
+    def make_symbolic_graph(self) -> None:
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        k_posdef = self.k_posdef
+
+        k_nonzero = int(sum(self.order))
+        ar_params = self.make_and_register_variable(
+            f"params_{self.name}", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
+        )
+        sigma_ar = self.make_and_register_variable(
+            f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
+        )
+
+        if k_endog == 1:
+            T = pt.eye(k_states, k=-1)
+            ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+            T = T[ar_idx].set(ar_params)
+
+        else:
+            transition_matrices = []
+
+            for i in range(k_endog):
+                T = pt.eye(k_states, k=-1)
+                ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
+                T = T[ar_idx].set(ar_params[i])
+                transition_matrices.append(T)
+            T = pt.specify_shape(
+                pt.linalg.block_diag(*transition_matrices), (self.k_states, self.k_states)
+            )
+
+        self.ssm["transition", :, :] = T
+
+        R = np.eye(k_states)
+        R_mask = np.full((k_states), False)
+        R_mask[0] = True
+        R = R[:, R_mask]
+
+        self.ssm["selection", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+        )
+
+        Z = pt.zeros((1, k_states))[0, 0].set(1.0)
+        self.ssm["design", :, :] = pt.specify_shape(
+            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+        )
+
+        cov_idx = ("state_cov", *np.diag_indices(k_posdef))
+        self.ssm[cov_idx] = sigma_ar**2