pymc-extras 0.2.0__py3-none-any.whl

tests/distributions/test_multivariate.py

@@ -0,0 +1,304 @@
+import numpy as np
+import pymc as pm
+import pytensor
+import pytest
+
+import pymc_extras as pmx
+
+
+class TestR2D2M2CP:
+    @pytest.fixture(autouse=True)
+    def fast_compile(self):
+        with pytensor.config.change_flags(mode="FAST_COMPILE", exception_verbosity="high"):
+            yield
+
+    @pytest.fixture(autouse=True)
+    def model(self):
+        # every method is within a model
+        with pm.Model() as model:
+            yield model
+
+    @pytest.fixture(params=[True, False], ids=["centered", "non-centered"])
+    def centered(self, request):
+        return request.param
+
+    @pytest.fixture(params=[["a"], ["a", "b"], ["one"]])
+    def dims(self, model: pm.Model, request):
+        for i, c in enumerate(request.param):
+            if c == "one":
+                model.add_coord(c, range(1))
+            else:
+                model.add_coord(c, range((i + 2) ** 2))
+        return request.param
+
+    @pytest.fixture
+    def input_shape(self, dims, model):
+        return [int(model.dim_lengths[d].eval()) for d in dims]
+
+    @pytest.fixture
+    def output_shape(self, dims, model):
+        *hierarchy, _ = dims
+        return [int(model.dim_lengths[d].eval()) for d in hierarchy]
+
+    @pytest.fixture
+    def input_std(self, input_shape):
+        return np.ones(input_shape)
+
+    @pytest.fixture
+    def output_std(self, output_shape):
+        return np.ones(output_shape)
+
+    @pytest.fixture
+    def r2(self):
+        return 0.8
+
+    @pytest.fixture(params=[None, 0.1], ids=["r2-std", "no-r2-std"])
+    def r2_std(self, request):
+        return request.param
+
+    @pytest.fixture(params=["true", "false", "limit-1", "limit-0", "limit-all"])
+    def positive_probs(self, input_std, request):
+        if request.param == "true":
+            return np.full_like(input_std, 0.5)
+        elif request.param == "false":
+            return 0.5
+        elif request.param == "limit-1":
+            ret = np.full_like(input_std, 0.5)
+            ret[..., 0] = 1
+            return ret
+        elif request.param == "limit-0":
+            ret = np.full_like(input_std, 0.5)
+            ret[..., 0] = 0
+            return ret
+        elif request.param == "limit-all":
+            return np.full_like(input_std, 0)
+
+    @pytest.fixture(params=[True, False], ids=["probs-std", "no-probs-std"])
+    def positive_probs_std(self, positive_probs, request):
+        if request.param:
+            std = np.full_like(positive_probs, 0.1)
+            std[positive_probs == 0] = 0
+            std[positive_probs == 1] = 0
+            return std
+        else:
+            return None
+
+    @pytest.fixture(params=[None, "importance", "explained"])
+    def phi_args_base(self, request, input_shape):
+        if input_shape[-1] < 2 and request.param is not None:
+            pytest.skip("not compatible")
+        elif request.param is None:
+            return {}
+        elif request.param == "importance":
+            return {"variables_importance": np.full(input_shape, 2)}
+        else:
+            val = np.full(input_shape, 2)
+            return {"variance_explained": val / val.sum(-1, keepdims=True)}
+
+    @pytest.fixture(params=["concentration", "no-concentration"])
+    def phi_args(self, request, phi_args_base):
+        if request.param == "concentration":
+            phi_args_base["importance_concentration"] = 10
+        return phi_args_base
+
+    def test_init_r2(
+        self,
+        dims,
+        input_std,
+        output_std,
+        r2,
+        r2_std,
+        model: pm.Model,
+    ):
+        eps, beta = pmx.distributions.R2D2M2CP(
+            "beta",
+            output_std,
+            input_std,
+            dims=dims,
+            r2=r2,
+            r2_std=r2_std,
+        )
+        assert not np.isnan(beta.eval()).any()
+        assert eps.eval().shape == output_std.shape
+        assert beta.eval().shape == input_std.shape
+        # r2 rv is only created if r2 std is not None
+        assert "beta" in model.named_vars
+        assert ("beta::r2" in model.named_vars) == (r2_std is not None), set(model.named_vars)
+        # phi is only created if variable importance is not None and there is more than one var
+        assert np.isfinite(model.compile_logp()(model.initial_point()))
+
+    def test_init_importance(
+        self,
+        dims,
+        centered,
+        input_std,
+        output_std,
+        phi_args,
+        model: pm.Model,
+    ):
+        eps, beta = pmx.distributions.R2D2M2CP(
+            "beta",
+            output_std,
+            input_std,
+            dims=dims,
+            r2=1,
+            centered=centered,
+            **phi_args,
+        )
+        assert not np.isnan(beta.eval()).any()
+        assert eps.eval().shape == output_std.shape
+        assert beta.eval().shape == input_std.shape
+        # r2 rv is only created if r2 std is not None
+        assert "beta" in model.named_vars
+        # phi is only created if variable importance is not None and there is more than one var
+        assert ("beta::phi" in model.named_vars) == (
+            "variables_importance" in phi_args or "importance_concentration" in phi_args
+        ), set(model.named_vars)
+        assert np.isfinite(model.compile_logp()(model.initial_point()))
+
+    def test_init_positive_probs(
+        self,
+        dims,
+        centered,
+        input_std,
+        output_std,
+        positive_probs,
+        positive_probs_std,
+        model: pm.Model,
+    ):
+        eps, beta = pmx.distributions.R2D2M2CP(
+            "beta",
+            output_std,
+            input_std,
+            dims=dims,
+            r2=1.0,
+            centered=centered,
+            positive_probs_std=positive_probs_std,
+            positive_probs=positive_probs,
+        )
+        assert not np.isnan(beta.eval()).any()
+        assert eps.eval().shape == output_std.shape
+        assert beta.eval().shape == input_std.shape
+        # r2 rv is only created if r2 std is not None
+        assert "beta" in model.named_vars
+        # phi is only created if variable importance is not None and there is more than one var
+        assert ("beta::psi" in model.named_vars) == (
+            positive_probs_std is not None and positive_probs_std.any()
+        ), set(model.named_vars)
+        assert np.isfinite(sum(model.point_logps().values()))
+
+    def test_failing_importance(self, dims, input_shape, output_std, input_std):
+        if input_shape[-1] < 2:
+            with pytest.raises(TypeError, match="less than two variables"):
+                pmx.distributions.R2D2M2CP(
+                    "beta",
+                    output_std,
+                    input_std,
+                    dims=dims,
+                    r2=0.8,
+                    variables_importance=abs(input_std),
+                )
+        else:
+            pmx.distributions.R2D2M2CP(
+                "beta",
+                output_std,
+                input_std,
+                dims=dims,
+                r2=0.8,
+                variables_importance=abs(input_std),
+            )
+
+    def test_failing_variance_explained(self, dims, input_shape, output_std, input_std):
+        if input_shape[-1] < 2:
+            with pytest.raises(TypeError, match="less than two variables"):
+                pmx.distributions.R2D2M2CP(
+                    "beta",
+                    output_std,
+                    input_std,
+                    dims=dims,
+                    r2=0.8,
+                    variance_explained=abs(input_std),
+                )
+        else:
+            pmx.distributions.R2D2M2CP(
+                "beta", output_std, input_std, dims=dims, r2=0.8, variance_explained=abs(input_std)
+            )
+
+    def test_failing_mutual_exclusive(self, model: pm.Model):
+        with pytest.raises(TypeError, match="variable importance with variance explained"):
+            with model:
+                model.add_coord("a", range(2))
+            pmx.distributions.R2D2M2CP(
+                "beta",
+                1,
+                [1, 1],
+                dims="a",
+                r2=0.8,
+                variance_explained=[0.5, 0.5],
+                variables_importance=[1, 1],
+            )
+
+    def test_limit_case_requires_std_0(self, model: pm.Model):
+        model.add_coord("a", range(2))
+        with pytest.raises(ValueError, match="Can't have both positive_probs"):
+            pmx.distributions.R2D2M2CP(
+                "beta",
+                1,
+                [1, 1],
+                dims="a",
+                r2=0.8,
+                positive_probs=[0.5, 0],
+                positive_probs_std=[0.3, 0.1],
+            )
+        with pytest.raises(ValueError, match="Can't have both positive_probs"):
+            pmx.distributions.R2D2M2CP(
+                "beta",
+                1,
+                [1, 1],
+                dims="a",
+                r2=0.8,
+                positive_probs=[0.5, 1],
+                positive_probs_std=[0.3, 0.1],
+            )
+
+    def test_limit_case_creates_masked_vars(self, model: pm.Model, centered: bool):
+        model.add_coord("a", range(2))
+        pmx.distributions.R2D2M2CP(
+            "beta0",
+            1,
+            [1, 1],
+            dims="a",
+            r2=0.8,
+            positive_probs=[0.5, 1],
+            positive_probs_std=[0.3, 0],
+            centered=centered,
+        )
+        pmx.distributions.R2D2M2CP(
+            "beta1",
+            1,
+            [1, 1],
+            dims="a",
+            r2=0.8,
+            positive_probs=[0.5, 0],
+            positive_probs_std=[0.3, 0],
+            centered=centered,
+        )
+        if not centered:
+            assert "beta0::raw::masked" in model.named_vars, model.named_vars
+            assert "beta1::raw::masked" in model.named_vars, model.named_vars
+        else:
+            assert "beta0::masked" in model.named_vars, model.named_vars
+            assert "beta1::masked" in model.named_vars, model.named_vars
+        assert "beta1::psi::masked" in model.named_vars
+        assert "beta0::psi::masked" in model.named_vars
+
+    def test_zero_length_rvs_not_created(self, model: pm.Model):
+        model.add_coord("a", range(2))
+        # deterministic case which should not have any new variables
+        b = pmx.distributions.R2D2M2CP("b1", 1, [1, 1], r2=0.5, positive_probs=[1, 1], dims="a")
+        assert not model.free_RVs, model.free_RVs
+
+        b = pmx.distributions.R2D2M2CP(
+            "b2", 1, [1, 1], r2=0.5, positive_probs=[1, 1], positive_probs_std=[0, 0], dims="a"
+        )
+        assert not model.free_RVs, model.free_RVs

tests/model/marginal/test_distributions.py

@@ -0,0 +1,131 @@
+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 MarginalModel
+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=(((),),),
+    )(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 MarginalModel() 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
+        )
+
+    m.marginalize([chain])
+    logp_fn = 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 MarginalModel() 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))
+    m.marginalize([chain])
+
+    test_value = np.array([0, 0, 1])
+    expected_logp = np.log(0.1344)  # Shown at the 10m22s mark in the video
+    logp_fn = 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 MarginalModel() 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)
+
+    with pytest.warns(UserWarning, match="multiple dependent variables"):
+        m.marginalize([chain])
+
+    logp_fn = 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

@@ -0,0 +1,182 @@
+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)