pymc-extras 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

tests/model/marginal/test_distributions.py

@@ -1,132 +0,0 @@
-import numpy as np
-import pymc as pm
-import pytest
-
-from pymc.logprob.abstract import _logprob
-from pytensor import tensor as pt
-from scipy.stats import norm
-
-from pymc_extras import marginalize
-from pymc_extras.distributions import DiscreteMarkovChain
-from pymc_extras.model.marginal.distributions import MarginalFiniteDiscreteRV
-
-
-def test_marginalized_bernoulli_logp():
-    """Test logp of IR TestFiniteMarginalDiscreteRV directly"""
-    mu = pt.vector("mu")
-
-    idx = pm.Bernoulli.dist(0.7, name="idx")
-    y = pm.Normal.dist(mu=mu[idx], sigma=1.0, name="y")
-    marginal_rv_node = MarginalFiniteDiscreteRV(
-        [mu],
-        [idx, y],
-        dims_connections=(((),),),
-        dims=(),
-    )(mu)[0].owner
-
-    y_vv = y.clone()
-    (logp,) = _logprob(
-        marginal_rv_node.op,
-        (y_vv,),
-        *marginal_rv_node.inputs,
-    )
-
-    ref_logp = pm.logp(pm.NormalMixture.dist(w=[0.3, 0.7], mu=mu, sigma=1.0), y_vv)
-    np.testing.assert_almost_equal(
-        logp.eval({mu: [-1, 1], y_vv: 2}),
-        ref_logp.eval({mu: [-1, 1], y_vv: 2}),
-    )
-
-
-@pytest.mark.parametrize("batch_chain", (False, True), ids=lambda x: f"batch_chain={x}")
-@pytest.mark.parametrize("batch_emission", (False, True), ids=lambda x: f"batch_emission={x}")
-def test_marginalized_hmm_normal_emission(batch_chain, batch_emission):
-    if batch_chain and not batch_emission:
-        pytest.skip("Redundant implicit combination")
-
-    with pm.Model() as m:
-        P = [[0, 1], [1, 0]]
-        init_dist = pm.Categorical.dist(p=[1, 0])
-        chain = DiscreteMarkovChain(
-            "chain", P=P, init_dist=init_dist, steps=3, shape=(3, 4) if batch_chain else None
-        )
-        emission = pm.Normal(
-            "emission", mu=chain * 2 - 1, sigma=1e-1, shape=(3, 4) if batch_emission else None
-        )
-
-    marginal_m = marginalize(m, [chain])
-    logp_fn = marginal_m.compile_logp()
-
-    test_value = np.array([-1, 1, -1, 1])
-    expected_logp = pm.logp(pm.Normal.dist(0, 1e-1), np.zeros_like(test_value)).sum().eval()
-    if batch_emission:
-        test_value = np.broadcast_to(test_value, (3, 4))
-        expected_logp *= 3
-    np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp)
-
-
-@pytest.mark.parametrize(
-    "categorical_emission",
-    [False, True],
-)
-def test_marginalized_hmm_categorical_emission(categorical_emission):
-    """Example adapted from https://www.youtube.com/watch?v=9-sPm4CfcD0"""
-    with pm.Model() as m:
-        P = np.array([[0.5, 0.5], [0.3, 0.7]])
-        init_dist = pm.Categorical.dist(p=[0.375, 0.625])
-        chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, steps=2)
-        if categorical_emission:
-            emission = pm.Categorical("emission", p=pt.constant([[0.8, 0.2], [0.4, 0.6]])[chain])
-        else:
-            emission = pm.Bernoulli("emission", p=pt.where(pt.eq(chain, 0), 0.2, 0.6))
-    marginal_m = marginalize(m, [chain])
-
-    test_value = np.array([0, 0, 1])
-    expected_logp = np.log(0.1344)  # Shown at the 10m22s mark in the video
-    logp_fn = marginal_m.compile_logp()
-    np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp)
-
-
-@pytest.mark.parametrize("batch_chain", (False, True))
-@pytest.mark.parametrize("batch_emission1", (False, True))
-@pytest.mark.parametrize("batch_emission2", (False, True))
-def test_marginalized_hmm_multiple_emissions(batch_chain, batch_emission1, batch_emission2):
-    chain_shape = (3, 1, 4) if batch_chain else (4,)
-    emission1_shape = (
-        (2, *reversed(chain_shape)) if batch_emission1 else tuple(reversed(chain_shape))
-    )
-    emission2_shape = (*chain_shape, 2) if batch_emission2 else chain_shape
-    with pm.Model() as m:
-        P = [[0, 1], [1, 0]]
-        init_dist = pm.Categorical.dist(p=[1, 0])
-        chain = DiscreteMarkovChain("chain", P=P, init_dist=init_dist, shape=chain_shape)
-        emission_1 = pm.Normal(
-            "emission_1", mu=(chain * 2 - 1).T, sigma=1e-1, shape=emission1_shape
-        )
-
-        emission2_mu = (1 - chain) * 2 - 1
-        if batch_emission2:
-            emission2_mu = emission2_mu[..., None]
-        emission_2 = pm.Normal("emission_2", mu=emission2_mu, sigma=1e-1, shape=emission2_shape)
-
-    marginal_m = marginalize(m, [chain])
-
-    with pytest.warns(UserWarning, match="multiple dependent variables"):
-        logp_fn = marginal_m.compile_logp(sum=False)
-
-    test_value = np.array([-1, 1, -1, 1])
-    multiplier = 2 + batch_emission1 + batch_emission2
-    if batch_chain:
-        multiplier *= 3
-    expected_logp = norm.logpdf(np.zeros_like(test_value), 0, 1e-1).sum() * multiplier
-
-    test_value = np.broadcast_to(test_value, chain_shape)
-    test_value_emission1 = np.broadcast_to(test_value.T, emission1_shape)
-    if batch_emission2:
-        test_value_emission2 = np.broadcast_to(-test_value[..., None], emission2_shape)
-    else:
-        test_value_emission2 = np.broadcast_to(-test_value, emission2_shape)
-    test_point = {"emission_1": test_value_emission1, "emission_2": test_value_emission2}
-    res_logp, dummy_logp = logp_fn(test_point)
-    assert res_logp.shape == ((1, 3) if batch_chain else ())
-    np.testing.assert_allclose(res_logp.sum(), expected_logp)

tests/model/marginal/test_graph_analysis.py

@@ -1,182 +0,0 @@
-import pytensor.tensor as pt
-import pytest
-
-from pymc.distributions import CustomDist
-from pytensor.tensor.type_other import NoneTypeT
-
-from pymc_extras.model.marginal.graph_analysis import (
-    is_conditional_dependent,
-    subgraph_batch_dim_connection,
-)
-
-
-def test_is_conditional_dependent_static_shape():
-    """Test that we don't consider dependencies through "constant" shape Ops"""
-    x1 = pt.matrix("x1", shape=(None, 5))
-    y1 = pt.random.normal(size=pt.shape(x1))
-    assert is_conditional_dependent(y1, x1, [x1, y1])
-
-    x2 = pt.matrix("x2", shape=(9, 5))
-    y2 = pt.random.normal(size=pt.shape(x2))
-    assert not is_conditional_dependent(y2, x2, [x2, y2])
-
-
-class TestSubgraphBatchDimConnection:
-    def test_dimshuffle(self):
-        inp = pt.tensor(shape=(5, 1, 4, 3))
-        out1 = pt.matrix_transpose(inp)
-        out2 = pt.expand_dims(inp, 1)
-        out3 = pt.squeeze(inp)
-        [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [out1, out2, out3])
-        assert dims1 == (0, 1, 3, 2)
-        assert dims2 == (0, None, 1, 2, 3)
-        assert dims3 == (0, 2, 3)
-
-    def test_careduce(self):
-        inp = pt.tensor(shape=(4, 3, 2))
-
-        out = pt.sum(inp[:, None], axis=(1,))
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, 2)
-
-        invalid_out = pt.sum(inp, axis=(1,))
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-    def test_subtensor(self):
-        inp = pt.tensor(shape=(4, 3, 2))
-
-        invalid_out = inp[0, :1]
-        with pytest.raises(
-            ValueError,
-            match="Partial slicing or indexing of known dimensions not supported",
-        ):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-        # If we are selecting dummy / unknown dimensions that's fine
-        valid_out = pt.expand_dims(inp, (0, 1))[0, :1]
-        [dims] = subgraph_batch_dim_connection(inp, [valid_out])
-        assert dims == (None, 0, 1, 2)
-
-    def test_advanced_subtensor_value(self):
-        inp = pt.tensor(shape=(2, 4))
-        intermediate_out = inp[:, None, :, None] + pt.zeros((2, 3, 4, 5))
-
-        # Index on an unlabled dim introduced by broadcasting with zeros
-        out = intermediate_out[:, [0, 0, 1, 2]]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, None, 1, None)
-
-        # Indexing that introduces more dimensions
-        out = intermediate_out[:, [[0, 0], [1, 2]], :]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, None, None, 1, None)
-
-        # Special case where advanced dims are moved to the front of the output
-        out = intermediate_out[:, [0, 0, 1, 2], :, 0]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (None, 0, 1)
-
-        # Indexing on a labeled dim fails
-        out = intermediate_out[:, :, [0, 0, 1, 2]]
-        with pytest.raises(ValueError, match="Partial slicing or advanced integer indexing"):
-            subgraph_batch_dim_connection(inp, [out])
-
-    def test_advanced_subtensor_key(self):
-        inp = pt.tensor(shape=(5, 5), dtype=int)
-        base = pt.zeros((2, 3, 4))
-
-        out = base[inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, None, None)
-
-        out = base[:, :, inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (
-            None,
-            None,
-            0,
-            1,
-        )
-
-        out = base[1:, 0, inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (None, 0, 1)
-
-        # Special case where advanced dims are moved to the front of the output
-        out = base[0, :, inp]
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, None)
-
-        # Mix keys dimensions
-        out = base[:, inp, inp.T]
-        with pytest.raises(ValueError, match="Different known dimensions mixed via broadcasting"):
-            subgraph_batch_dim_connection(inp, [out])
-
-    def test_elemwise(self):
-        inp = pt.tensor(shape=(5, 5))
-
-        out = inp + inp
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1)
-
-        out = inp + inp.T
-        with pytest.raises(ValueError, match="Different known dimensions mixed via broadcasting"):
-            subgraph_batch_dim_connection(inp, [out])
-
-        out = inp[None, :, None, :] + inp[:, None, :, None]
-        with pytest.raises(
-            ValueError, match="Same known dimension used in different axis after broadcasting"
-        ):
-            subgraph_batch_dim_connection(inp, [out])
-
-    def test_blockwise(self):
-        inp = pt.tensor(shape=(5, 4))
-
-        invalid_out = inp @ pt.ones((4, 3))
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-        out = (inp[:, :, None, None] + pt.zeros((2, 3))) @ pt.ones((2, 3))
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, None, None)
-
-    def test_random_variable(self):
-        inp = pt.tensor(shape=(5, 4, 3))
-
-        out1 = pt.random.normal(loc=inp)
-        out2 = pt.random.categorical(p=inp[..., None])
-        out3 = pt.random.multivariate_normal(mean=inp[..., None], cov=pt.eye(1))
-        [dims1, dims2, dims3] = subgraph_batch_dim_connection(inp, [out1, out2, out3])
-        assert dims1 == (0, 1, 2)
-        assert dims2 == (0, 1, 2)
-        assert dims3 == (0, 1, 2, None)
-
-        invalid_out = pt.random.categorical(p=inp)
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-        invalid_out = pt.random.multivariate_normal(mean=inp, cov=pt.eye(3))
-        with pytest.raises(ValueError, match="Use of known dimensions"):
-            subgraph_batch_dim_connection(inp, [invalid_out])
-
-    def test_symbolic_random_variable(self):
-        inp = pt.tensor(shape=(4, 3, 2))
-
-        # Test univariate
-        out = CustomDist.dist(
-            inp,
-            dist=lambda mu, size: pt.random.normal(loc=mu, size=size),
-        )
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, 2)
-
-        # Test multivariate
-        def dist(mu, size):
-            if isinstance(size.type, NoneTypeT):
-                size = mu.shape
-            return pt.random.normal(loc=mu[..., None], size=(*size, 2))
-
-        out = CustomDist.dist(inp, dist=dist, size=(4, 3, 2), signature="()->(2)")
-        [dims] = subgraph_batch_dim_connection(inp, [out])
-        assert dims == (0, 1, 2, None)