pymc-extras 0.2.0__tar.gz → 0.2.1__tar.gz

{pymc_extras-0.2.0 → pymc_extras-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pymc-extras
-Version: 0.2.0
+Version: 0.2.1
 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.
 Home-page: http://github.com/pymc-devs/pymc-extras
 Maintainer: PyMC Developers
@@ -63,10 +63,9 @@ import pymc as pm
 import pymc_extras as pmx
 
 with pm.Model():
+  alpha = pmx.ParabolicFractal('alpha', b=1, c=1)
 
-    alpha = pmx.ParabolicFractal('alpha', b=1, c=1)
-
-    ...
+  ...
 
 ```
 

{pymc_extras-0.2.0 → pymc_extras-0.2.1}/README.md

@@ -26,10 +26,9 @@ import pymc as pm
 import pymc_extras as pmx
 
 with pm.Model():
+  alpha = pmx.ParabolicFractal('alpha', b=1, c=1)
 
-    alpha = pmx.ParabolicFractal('alpha', b=1, c=1)
-
-    ...
+  ...
 
 ```
 

{pymc_extras-0.2.0 → pymc_extras-0.2.1}/pymc_extras/__init__.py

@@ -16,7 +16,11 @@ import logging
 from pymc_extras import gp, statespace, utils
 from pymc_extras.distributions import *
 from pymc_extras.inference.fit import fit
-from pymc_extras.model.marginal.marginal_model import MarginalModel, marginalize
+from pymc_extras.model.marginal.marginal_model import (
+    MarginalModel,
+    marginalize,
+    recover_marginals,
+)
 from pymc_extras.model.model_api import as_model
 from pymc_extras.version import __version__
 

{pymc_extras-0.2.0 → pymc_extras-0.2.1}/pymc_extras/distributions/timeseries.py

@@ -214,8 +214,8 @@ class DiscreteMarkovChain(Distribution):
         discrete_mc_op = DiscreteMarkovChainRV(
             inputs=[P_, steps_, init_dist_, state_rng],
             outputs=[state_next_rng, discrete_mc_],
-            ndim_supp=1,
             n_lags=n_lags,
+            extended_signature="(p,p),(),(p),[rng]->[rng],(t)",
         )
 
         discrete_mc = discrete_mc_op(P, steps, init_dist, state_rng)

{pymc_extras-0.2.0 → pymc_extras-0.2.1}/pymc_extras/model/marginal/distributions.py

@@ -1,20 +1,25 @@
+import warnings
+
 from collections.abc import Sequence
 
 import numpy as np
 import pytensor.tensor as pt
 
 from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
+from pymc.distributions.distribution import _support_point, support_point
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
 from pytensor import Variable
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.mode import Mode
-from pytensor.graph import Op, vectorize_graph
+from pytensor.graph import FunctionGraph, Op, vectorize_graph
+from pytensor.graph.basic import equal_computations
 from pytensor.graph.replace import clone_replace, graph_replace
 from pytensor.scan import map as scan_map
 from pytensor.scan import scan
 from pytensor.tensor import TensorVariable
+from pytensor.tensor.random.type import RandomType
 
 from pymc_extras.distributions import DiscreteMarkovChain
 
@@ -22,8 +27,15 @@ from pymc_extras.distributions import DiscreteMarkovChain
 class MarginalRV(OpFromGraph, MeasurableOp):
     """Base class for Marginalized RVs"""
 
-    def __init__(self, *args, dims_connections: tuple[tuple[int | None]], **kwargs) -> None:
+    def __init__(
+        self,
+        *args,
+        dims_connections: tuple[tuple[int | None], ...],
+        dims: tuple[Variable, ...],
+        **kwargs,
+    ) -> None:
         self.dims_connections = dims_connections
+        self.dims = dims
         super().__init__(*args, **kwargs)
 
     @property
@@ -43,6 +55,74 @@ class MarginalRV(OpFromGraph, MeasurableOp):
             )
         return tuple(support_axes_vars)
 
+    def __eq__(self, other):
+        # Just to allow easy testing of equivalent models,
+        # This can be removed once https://github.com/pymc-devs/pytensor/issues/1114 is fixed
+        if type(self) is not type(other):
+            return False
+
+        return equal_computations(
+            self.inner_outputs,
+            other.inner_outputs,
+            self.inner_inputs,
+            other.inner_inputs,
+        )
+
+    def __hash__(self):
+        # Just to allow easy testing of equivalent models,
+        # This can be removed once https://github.com/pymc-devs/pytensor/issues/1114 is fixed
+        return hash((type(self), len(self.inner_inputs), len(self.inner_outputs)))
+
+
+@_support_point.register
+def support_point_marginal_rv(op: MarginalRV, rv, *inputs):
+    """Support point for a marginalized RV.
+
+    The support point of a marginalized RV is the support point of the inner RV,
+    conditioned on the marginalized RV taking its support point.
+    """
+    outputs = rv.owner.outputs
+
+    inner_rv = op.inner_outputs[outputs.index(rv)]
+    marginalized_inner_rv, *other_dependent_inner_rvs = (
+        out
+        for out in op.inner_outputs
+        if out is not inner_rv and not isinstance(out.type, RandomType)
+    )
+
+    # Replace references to inner rvs by the dummy variables (including the marginalized RV)
+    # This is necessary because the inner RVs may depend on each other
+    marginalized_inner_rv_dummy = marginalized_inner_rv.clone()
+    other_dependent_inner_rv_to_dummies = {
+        inner_rv: inner_rv.clone() for inner_rv in other_dependent_inner_rvs
+    }
+    inner_rv = clone_replace(
+        inner_rv,
+        replace={marginalized_inner_rv: marginalized_inner_rv_dummy}
+        | other_dependent_inner_rv_to_dummies,
+    )
+
+    # Get support point of inner RV and marginalized RV
+    inner_rv_support_point = support_point(inner_rv)
+    marginalized_inner_rv_support_point = support_point(marginalized_inner_rv)
+
+    replacements = [
+        # Replace the marginalized RV dummy by its support point
+        (marginalized_inner_rv_dummy, marginalized_inner_rv_support_point),
+        # Replace other dependent RVs dummies by the respective outer outputs.
+        # PyMC will replace them by their support points later
+        *(
+            (v, outputs[op.inner_outputs.index(k)])
+            for k, v in other_dependent_inner_rv_to_dummies.items()
+        ),
+        # Replace outer input RVs
+        *zip(op.inner_inputs, inputs),
+    ]
+    fgraph = FunctionGraph(outputs=[inner_rv_support_point], clone=False)
+    fgraph.replace_all(replacements, import_missing=True)
+    [rv_support_point] = fgraph.outputs
+    return rv_support_point
+
 
 class MarginalFiniteDiscreteRV(MarginalRV):
     """Base class for Marginalized Finite Discrete RVs"""
@@ -132,12 +212,27 @@ def inline_ofg_outputs(op: OpFromGraph, inputs: Sequence[Variable]) -> tuple[Var
     Whereas `OpFromGraph` "wraps" a graph inside a single Op, this function "unwraps"
     the inner graph.
     """
-    return clone_replace(
+    return graph_replace(
         op.inner_outputs,
         replace=tuple(zip(op.inner_inputs, inputs)),
+        strict=False,
     )
 
 
+class NonSeparableLogpWarning(UserWarning):
+    pass
+
+
+def warn_non_separable_logp(values):
+    if len(values) > 1:
+        warnings.warn(
+            "There are multiple dependent variables in a FiniteDiscreteMarginalRV. "
+            f"Their joint logp terms will be assigned to the first value: {values[0]}.",
+            NonSeparableLogpWarning,
+            stacklevel=2,
+        )
+
+
 DUMMY_ZERO = pt.constant(0, name="dummy_zero")
 
 
@@ -199,6 +294,7 @@ def finite_discrete_marginal_rv_logp(op: MarginalFiniteDiscreteRV, values, *inpu
     # Align logp with non-collapsed batch dimensions of first RV
     joint_logp = align_logp_dims(dims=op.dims_connections[0], logp=joint_logp)
 
+    warn_non_separable_logp(values)
     # We have to add dummy logps for the remaining value variables, otherwise PyMC will raise
     dummy_logps = (DUMMY_ZERO,) * (len(values) - 1)
     return joint_logp, *dummy_logps
@@ -272,5 +368,6 @@ def marginal_hmm_logp(op, values, *inputs, **kwargs):
 
     # If there are multiple emission streams, we have to add dummy logps for the remaining value variables. The first
     # return is the joint probability of everything together, but PyMC still expects one logp for each emission stream.
+    warn_non_separable_logp(values)
     dummy_logps = (DUMMY_ZERO,) * (len(values) - 1)
     return joint_logp, *dummy_logps

{pymc_extras-0.2.0 → pymc_extras-0.2.1}/pymc_extras/model/marginal/graph_analysis.py

@@ -4,8 +4,8 @@ from collections.abc import Sequence
 from itertools import zip_longest
 
 from pymc import SymbolicRandomVariable
-from pytensor.compile import SharedVariable
-from pytensor.graph import Constant, Variable, ancestors
+from pymc.model.fgraph import ModelVar
+from pytensor.graph import Variable, ancestors
 from pytensor.graph.basic import io_toposort
 from pytensor.tensor import TensorType, TensorVariable
 from pytensor.tensor.blockwise import Blockwise
@@ -35,13 +35,9 @@ def static_shape_ancestors(vars):
 
 def find_conditional_input_rvs(output_rvs, all_rvs):
     """Find conditionally indepedent input RVs."""
-    blockers = [other_rv for other_rv in all_rvs if other_rv not in output_rvs]
-    blockers += static_shape_ancestors(tuple(all_rvs) + tuple(output_rvs))
-    return [
-        var
-        for var in ancestors(output_rvs, blockers=blockers)
-        if var in blockers or (var.owner is None and not isinstance(var, Constant | SharedVariable))
-    ]
+    other_rvs = [other_rv for other_rv in all_rvs if other_rv not in output_rvs]
+    blockers = other_rvs + static_shape_ancestors(tuple(all_rvs) + tuple(output_rvs))
+    return [var for var in ancestors(output_rvs, blockers=blockers) if var in other_rvs]
 
 
 def is_conditional_dependent(
@@ -141,6 +137,9 @@ def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars)
             # None of the inputs are related to the batch_axes of the input_vars
             continue
 
+        elif isinstance(node.op, ModelVar):
+            var_dims[node.outputs[0]] = inputs_dims[0]
+
         elif isinstance(node.op, DimShuffle):
             [input_dims] = inputs_dims
             output_dims = tuple(None if i == "x" else input_dims[i] for i in node.op.new_order)