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

pymc_extras/inference/pathfinder/pathfinder.py

@@ -278,12 +278,13 @@ def alpha_recover(
     z = pt.diff(g, axis=0)
     alpha_l_init = pt.ones(N)
 
-    alpha, _ = pytensor.scan(
+    alpha = pytensor.scan(
         fn=compute_alpha_l,
         outputs_info=alpha_l_init,
         sequences=[s, z],
         n_steps=Lp1 - 1,
         allow_gc=False,
+        return_updates=False,
     )
 
     # assert np.all(alpha.eval() > 0), "alpha cannot be negative"
@@ -334,11 +335,12 @@ def inverse_hessian_factors(
             return pt.set_subtensor(chi_l[j_last], diff_l)
 
         chi_init = pt.zeros((J, N))
-        chi_mat, _ = pytensor.scan(
+        chi_mat = pytensor.scan(
             fn=chi_update,
             outputs_info=chi_init,
             sequences=[diff],
             allow_gc=False,
+            return_updates=False,
         )
 
         chi_mat = pt.matrix_transpose(chi_mat)
@@ -377,14 +379,14 @@ def inverse_hessian_factors(
     eta = pt.diagonal(E, axis1=-2, axis2=-1)
 
     # beta: (L, N, 2J)
-    alpha_diag, _ = pytensor.scan(lambda a: pt.diag(a), sequences=[alpha])
+    alpha_diag = pytensor.scan(lambda a: pt.diag(a), sequences=[alpha], return_updates=False)
     beta = pt.concatenate([alpha_diag @ Z, S], axis=-1)
 
     # more performant and numerically precise to use solve than inverse: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.inv.html
 
     # E_inv: (L, J, J)
-    E_inv = pt.slinalg.solve_triangular(E, Ij, check_finite=False)
-    eta_diag, _ = pytensor.scan(pt.diag, sequences=[eta])
+    E_inv = pt.linalg.solve_triangular(E, Ij, check_finite=False)
+    eta_diag = pytensor.scan(pt.diag, sequences=[eta], return_updates=False)
 
     # block_dd: (L, J, J)
     block_dd = (
@@ -530,7 +532,9 @@ def bfgs_sample_sparse(
 
     # qr_input: (L, N, 2J)
     qr_input = inv_sqrt_alpha_diag @ beta
-    (Q, R), _ = pytensor.scan(fn=pt.linalg.qr, sequences=[qr_input], allow_gc=False)
+    Q, R = pytensor.scan(
+        fn=pt.linalg.qr, sequences=[qr_input], allow_gc=False, return_updates=False
+    )
 
     IdN = pt.eye(R.shape[1])[None, ...]
     IdN += IdN * REGULARISATION_TERM
@@ -623,10 +627,11 @@ def bfgs_sample(
 
     L, N, JJ = beta.shape
 
-    (alpha_diag, inv_sqrt_alpha_diag, sqrt_alpha_diag), _ = pytensor.scan(
+    alpha_diag, inv_sqrt_alpha_diag, sqrt_alpha_diag = pytensor.scan(
         lambda a: [pt.diag(a), pt.diag(pt.sqrt(1.0 / a)), pt.diag(pt.sqrt(a))],
         sequences=[alpha],
         allow_gc=False,
+        return_updates=False,
     )
 
     u = pt.random.normal(size=(L, num_samples, N))

pymc_extras/inference/smc/sampling.py

@@ -238,7 +238,7 @@ class SMCDiagnostics(NamedTuple):
     def update_diagnosis(i, history, info, state):
         le, lli, ancestors, weights_evolution = history
         return SMCDiagnostics(
-            le.at[i].set(state.lmbda),
+            le.at[i].set(state.tempering_param),
             lli.at[i].set(info.log_likelihood_increment),
             ancestors.at[i].set(info.ancestors),
             weights_evolution.at[i].set(state.weights),
@@ -265,7 +265,7 @@ def inference_loop(rng_key, initial_state, kernel, iterations_to_diagnose, n_par
 
     def cond(carry):
         i, state, _, _ = carry
-        return state.lmbda < 1
+        return state.tempering_param < 1
 
     def one_step(carry):
         i, state, k, previous_info = carry

pymc_extras/model/marginal/distributions.py

@@ -282,11 +282,12 @@ def finite_discrete_marginal_rv_logp(op: MarginalFiniteDiscreteRV, values, *inpu
         def logp_fn(marginalized_rv_const, *non_sequences):
             return graph_replace(joint_logp, replace={marginalized_vv: marginalized_rv_const})
 
-        joint_logps, _ = scan_map(
+        joint_logps = scan_map(
             fn=logp_fn,
             sequences=marginalized_rv_domain_tensor,
             non_sequences=[*values, *inputs],
             mode=Mode().including("local_remove_check_parameter"),
+            return_updates=False,
         )
 
     joint_logp = pt.logsumexp(joint_logps, axis=0)
@@ -350,12 +351,13 @@ def marginal_hmm_logp(op, values, *inputs, **kwargs):
 
     P_bcast_dims = (len(chain_shape) - 1) - (P.type.ndim - 2)
     log_P = pt.shape_padright(pt.log(P), P_bcast_dims)
-    log_alpha_seq, _ = scan(
+    log_alpha_seq = scan(
         step_alpha,
         non_sequences=[log_P],
         outputs_info=[log_alpha_init],
         # Scan needs the time dimension first, and we already consumed the 1st logp computing the initial value
         sequences=pt.moveaxis(batch_logp_emissions[..., 1:], -1, 0),
+        return_updates=False,
     )
     # Final logp is just the sum of the last scan state
     joint_logp = pt.logsumexp(log_alpha_seq[-1], axis=0)

pymc_extras/model/marginal/marginal_model.py

@@ -11,7 +11,7 @@ from pymc.backends.arviz import coords_and_dims_for_inferencedata, dataset_to_po
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
 from pymc.distributions.transforms import Chain
 from pymc.logprob.transforms import IntervalTransform
-from pymc.model import Model
+from pymc.model import Model, modelcontext
 from pymc.model.fgraph import (
     ModelFreeRV,
     ModelValuedVar,
@@ -337,8 +337,9 @@ def transform_posterior_pts(model, posterior_pts):
 
 
 def recover_marginals(
-    model: Model,
     idata: InferenceData,
+    *,
+    model: Model | None = None,
     var_names: Sequence[str] | None = None,
     return_samples: bool = True,
     extend_inferencedata: bool = True,
@@ -389,6 +390,15 @@ def recover_marginals(
 
 
     """
+    # Temporary error message for helping with migration
+    # Will be removed in a future release
+    if isinstance(idata, Model):
+        raise TypeError(
+            "The order of arguments of `recover_marginals` changed. The first input must be an idata"
+        )
+
+    model = modelcontext(model)
+
     unmarginal_model = unmarginalize(model)
 
     # Find the names of the marginalized variables

pymc_extras/prior.py

@@ -1575,9 +1575,9 @@ def __getattr__(name: str):
         samples = dist.sample_prior(coords={"channel": ["C1", "C2", "C3"]})
 
     """
-    # Protect against doctest
-    if name == "__wrapped__":
-        return
+    # Ignore Python internal attributes needed for introspection
+    if name.startswith("__"):
+        raise AttributeError(name)
 
     _get_pymc_distribution(name)
     return partial(Prior, distribution=name)

pymc_extras/statespace/core/properties.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from copy import deepcopy
+from dataclasses import dataclass, fields
+from typing import Generic, Protocol, Self, TypeVar
+
+from pytensor.tensor.variable import TensorVariable
+
+from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_AUX_DIM,
+    ALL_STATE_DIM,
+    OBS_STATE_AUX_DIM,
+    OBS_STATE_DIM,
+    SHOCK_AUX_DIM,
+    SHOCK_DIM,
+)
+
+
+class StateSpaceLike(Protocol):
+    @property
+    def state_names(self) -> tuple[str, ...]: ...
+
+    @property
+    def observed_states(self) -> tuple[str, ...]: ...
+
+    @property
+    def shock_names(self) -> tuple[str, ...]: ...
+
+
+@dataclass(frozen=True)
+class Property:
+    def __str__(self) -> str:
+        return "\n".join(f"{f.name}: {getattr(self, f.name)}" for f in fields(self))
+
+
+T = TypeVar("T", bound=Property)
+
+
+@dataclass(frozen=True)
+class Info(Generic[T]):
+    items: tuple[T, ...] | None
+    key_field: str | tuple[str, ...] = "name"
+    _index: dict[str | tuple, T] | None = None
+
+    def __post_init__(self):
+        index = {}
+        if self.items is None:
+            object.__setattr__(self, "items", ())
+        else:
+            object.__setattr__(self, "items", tuple(self.items))
+
+        for item in self.items:
+            key = self._key(item)
+            if key in index:
+                raise ValueError(f"Duplicate {self.key_field} '{key}' detected.")
+            index[key] = item
+        object.__setattr__(self, "_index", index)
+
+    def _key(self, item: T) -> str | tuple:
+        if isinstance(self.key_field, tuple):
+            return tuple(getattr(item, f) for f in self.key_field)
+        return getattr(item, self.key_field)
+
+    def get(self, key: str | tuple, default=None) -> T | None:
+        return self._index.get(key, default)
+
+    def __getitem__(self, key: str | tuple) -> T:
+        try:
+            return self._index[key]
+        except KeyError as e:
+            available = ", ".join(str(k) for k in self._index.keys())
+            raise KeyError(f"No {self.key_field} '{key}'. Available: [{available}]") from e
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._index
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(self.items)
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __str__(self) -> str:
+        return f"{self.key_field}s: {tuple(self._index.keys())}"
+
+    def add(self, new_item: T) -> Self:
+        return type(self)((*self.items, new_item))
+
+    def merge(self, other: Self, overwrite_duplicates: bool = False) -> Self:
+        if not isinstance(other, type(self)):
+            raise TypeError(f"Cannot merge {type(other).__name__} with {type(self).__name__}")
+
+        overlapping = set(self._index.keys()) & set(other._index.keys())
+        if overlapping and overwrite_duplicates:
+            return type(self)(
+                (
+                    *self.items,
+                    *(item for item in other.items if self._key(item) not in overlapping),
+                )
+            )
+
+        return type(self)(self.items + other.items)
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        if isinstance(self.key_field, tuple):
+            return tuple(item.name for item in self.items)
+        return tuple(self._index.keys())
+
+    def copy(self) -> Info[T]:
+        return deepcopy(self)
+
+
+@dataclass(frozen=True)
+class Parameter(Property):
+    name: str
+    shape: tuple[int, ...] | None = None
+    dims: tuple[str, ...] | None = None
+    constraints: str | None = None
+
+
+@dataclass(frozen=True)
+class ParameterInfo(Info[Parameter]):
+    def __init__(self, parameters: tuple[Parameter, ...] | None):
+        super().__init__(items=parameters, key_field="name")
+
+    def to_dict(self):
+        return {
+            param.name: {"shape": param.shape, "constraints": param.constraints, "dims": param.dims}
+            for param in self.items
+        }
+
+
+@dataclass(frozen=True)
+class Data(Property):
+    name: str
+    shape: tuple[int | None, ...] | None = None
+    dims: tuple[str, ...] | None = None
+    is_exogenous: bool = False
+
+
+@dataclass(frozen=True)
+class DataInfo(Info[Data]):
+    def __init__(self, data: tuple[Data, ...] | None):
+        super().__init__(items=data, key_field="name")
+
+    @property
+    def needs_exogenous_data(self) -> bool:
+        return any(d.is_exogenous for d in self.items)
+
+    @property
+    def exogenous_names(self) -> tuple[str, ...]:
+        return tuple(d.name for d in self.items if d.is_exogenous)
+
+    def __str__(self) -> str:
+        return f"data: {[d.name for d in self.items]}\nneeds exogenous data: {self.needs_exogenous_data}"
+
+    def to_dict(self):
+        return {
+            data.name: {"shape": data.shape, "dims": data.dims, "exogenous": data.is_exogenous}
+            for data in self.items
+        }
+
+
+@dataclass(frozen=True)
+class Coord(Property):
+    dimension: str
+    labels: tuple[str | int, ...]
+
+
+@dataclass(frozen=True)
+class CoordInfo(Info[Coord]):
+    def __init__(self, coords: tuple[Coord, ...] | None = None):
+        super().__init__(items=coords, key_field="dimension")
+
+    def __str__(self) -> str:
+        base = "coordinates:"
+        for coord in self.items:
+            coord_str = str(coord)
+            indented = "\n".join("  " + line for line in coord_str.splitlines())
+            base += "\n" + indented + "\n"
+        return base
+
+    @classmethod
+    def default_coords_from_model(cls, model: StateSpaceLike) -> CoordInfo:
+        states = tuple(model.state_names)
+        obs_states = tuple(model.observed_states)
+        shocks = tuple(model.shock_names)
+
+        dim_to_labels = (
+            (ALL_STATE_DIM, states),
+            (ALL_STATE_AUX_DIM, states),
+            (OBS_STATE_DIM, obs_states),
+            (OBS_STATE_AUX_DIM, obs_states),
+            (SHOCK_DIM, shocks),
+            (SHOCK_AUX_DIM, shocks),
+        )
+
+        coords = tuple(Coord(dimension=dim, labels=labels) for dim, labels in dim_to_labels)
+        return cls(coords=coords)
+
+    def to_dict(self):
+        return {coord.dimension: tuple(coord.labels) for coord in self.items}
+
+
+@dataclass(frozen=True)
+class State(Property):
+    name: str
+    observed: bool
+    shared: bool = False
+
+
+@dataclass(frozen=True)
+class StateInfo(Info[State]):
+    def __init__(self, states: tuple[State, ...] | None):
+        super().__init__(items=states, key_field=("name", "observed"))
+
+    def __contains__(self, key: object) -> bool:
+        if isinstance(key, str):
+            return any(s.name == key for s in self.items)
+        return key in self._index
+
+    def __str__(self) -> str:
+        return (
+            f"states: {[s.name for s in self.items]}\nobserved: {[s.observed for s in self.items]}"
+        )
+
+    @property
+    def observed_state_names(self) -> tuple[str, ...]:
+        return tuple(s.name for s in self.items if s.observed)
+
+    @property
+    def unobserved_state_names(self) -> tuple[str, ...]:
+        return tuple(s.name for s in self.items if not s.observed)
+
+
+@dataclass(frozen=True)
+class Shock(Property):
+    name: str
+
+
+@dataclass(frozen=True)
+class ShockInfo(Info[Shock]):
+    def __init__(self, shocks: tuple[Shock, ...] | None):
+        super().__init__(items=shocks, key_field="name")
+
+
+@dataclass(frozen=True)
+class SymbolicVariable(Property):
+    name: str
+    symbolic_variable: TensorVariable
+
+
+@dataclass(frozen=True)
+class SymbolicVariableInfo(Info[SymbolicVariable]):
+    def __init__(self, symbolic_variables: tuple[SymbolicVariable, ...] | None = None):
+        super().__init__(items=symbolic_variables, key_field="name")
+
+    def to_dict(self):
+        return {variable.name: variable.symbolic_variable for variable in self.items}
+
+
+@dataclass(frozen=True)
+class SymbolicData(Property):
+    name: str
+    symbolic_data: TensorVariable
+
+
+@dataclass(frozen=True)
+class SymbolicDataInfo(Info[SymbolicData]):
+    def __init__(self, symbolic_data: tuple[SymbolicData, ...] | None = None):
+        super().__init__(items=symbolic_data, key_field="name")
+
+    def to_dict(self):
+        return {data.name: data.symbolic_data for data in self.items}