add init test

Author: ferrine

pymc_experimental/distributions/__init__.py

@@ -20,9 +20,11 @@
 from pymc_experimental.distributions.continuous import GenExtreme
 from pymc_experimental.distributions.discrete import GeneralizedPoisson
 from pymc_experimental.distributions.timeseries import DiscreteMarkovChain
+from pymc_experimental.distributions.multivatiate import R2D2M2CP
 
 __all__ = [
     "DiscreteMarkovChain",
     "GeneralizedPoisson",
     "GenExtreme",
+    "R2D2M2CP",
 ]

pymc_experimental/distributions/multivatiate.py

@@ -1,13 +1,28 @@
+#   Copyright 2022 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+
 from typing import Sequence, Union
 
 import pymc as pm
 import pytensor.tensor as pt
 
 
-def _psivar2musigma(psi: pt.TensorVariable, var: pt.TensorVariable):
+def _psivar2musigma(psi: pt.TensorVariable, explained_var: pt.TensorVariable):
     pi = pt.erfinv(2 * psi - 1)
     f = (1 / (2 * pi**2 + 1)) ** 0.5
-    sigma = pt.expand_dims(var, -1) ** 0.5 * f
+    sigma = explained_var**0.5 * f
     mu = sigma * pi * 2**0.5
     return mu, sigma
 
@@ -38,7 +53,7 @@ def _R2D2M2CP_beta(
         probability of a coefficients to be positive
     """
     tau2 = r2 / (1 - r2)
-    explained_variance = phi * tau2 * pt.expand_dims(output_sigma**2, -1)
+    explained_variance = phi * pt.expand_dims(tau2 * output_sigma**2, -1)
     mu_param, std_param = _psivar2musigma(psi, explained_variance)
     if not centered:
         with pm.Model(name):
@@ -103,6 +118,8 @@ def R2D2M2CP(
 
     Notes
     -----
+    The R2D2M2CP prior is a modification of R2D2M2 prior.
+
     - ``(R2D2M2)``CP is taken from https://arxiv.org/abs/2208.07132
     - R2D2M2``(CP)``, (Correlation Probability) is proposed and implemented by Max Kochurov (@ferrine)
     """
@@ -116,7 +133,7 @@ def R2D2M2CP(
             if variance_explained is not None:
                 raise TypeError("Can't use variable importance with variance explained")
             phi = pm.Dirichlet("phi", pt.as_tensor(variables_importance), dims=hierarchy + [dim])
-        elif variance_explained:
+        elif variance_explained is not None:
             phi = pt.as_tensor(variance_explained)
         else:
             phi = 1 / len(model.coords[dim])

pymc_experimental/tests/distributions/test_multivariate.py

@@ -0,0 +1,94 @@
+import numpy as np
+import pymc as pm
+import pytest
+
+import pymc_experimental as pmx
+
+
+class TestR2D2M2CP:
+    @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])
+    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_std(self, dims, model):
+        input_shape = [int(model.dim_lengths[d].eval()) for d in dims]
+        return np.ones(input_shape)
+
+    @pytest.fixture
+    def output_std(self, dims, model):
+        *hierarchy, _ = dims
+        output_shape = [int(model.dim_lengths[d].eval()) for d in hierarchy]
+        return np.ones(output_shape)
+
+    @pytest.fixture
+    def r2(self):
+        return 0.8
+
+    @pytest.fixture(params=[None, 0.1])
+    def r2_std(self, request):
+        return request.param
+
+    @pytest.fixture(params=[True, False])
+    def positive_probs(self, input_std, request):
+        if request.param:
+            return np.full_like(input_std, 0.5)
+        else:
+            return 0.5
+
+    @pytest.fixture(params=[True, False])
+    def positive_probs_std(self, positive_probs, request):
+        if request.param:
+            return np.full_like(positive_probs, 0.1)
+        else:
+            return None
+
+    @pytest.fixture(params=["importance", "explained"])
+    def phi_args(self, request, input_std):
+        if request.param == "importance":
+            return {"variables_importance": np.full_like(input_std, 2)}
+        else:
+            val = np.full_like(input_std, 2)
+            return {"variance_explained": val / val.sum(-1, keepdims=True)}
+
+    def test_init(
+        self,
+        dims,
+        centered,
+        input_std,
+        output_std,
+        r2,
+        r2_std,
+        positive_probs,
+        positive_probs_std,
+        phi_args,
+    ):
+        eps, beta = pmx.distributions.R2D2M2CP(
+            "beta",
+            output_std,
+            input_std,
+            dims=dims,
+            r2=r2,
+            r2_std=r2_std,
+            centered=centered,
+            positive_probs_std=positive_probs_std,
+            positive_probs=positive_probs,
+            **phi_args
+        )
+        assert eps.eval().shape == output_std.shape
+        assert beta.eval().shape == input_std.shape