fix bug with missed variable

Author: ferrine

pymc_experimental/distributions/multivariate/r2d2m2cp.py

@@ -71,18 +71,23 @@ def _R2D2M2CP_beta(
     if not centered:
         with pm.Model(name):
             if psi_mask is not None and psi_mask.any():
+                # limit case where some probs are not 1 or 0
+                # setsubtensor is required
                 r_idx = psi_mask.nonzero()
                 with pm.Model("raw"):
                     raw = pm.Normal("masked", shape=len(r_idx[0]))
                     raw = pt.set_subtensor(pt.zeros_like(mu_param)[r_idx], raw)
                 raw = pm.Deterministic("raw", raw, dims=dims)
             elif psi_mask is not None:
+                # all variables are deterministic
                 raw = pt.zeros_like(mu_param)
             else:
                 raw = pm.Normal("raw", dims=dims)
         beta = pm.Deterministic(name, (raw * std_param + mu_param) / input_sigma, dims=dims)
     else:
         if psi_mask is not None and psi_mask.any():
+            # limit case where some probs are not 1 or 0
+            # setsubtensor is required
             r_idx = psi_mask.nonzero()
             with pm.Model(name):
                 mean = (mu_param / input_sigma)[r_idx]
@@ -96,7 +101,8 @@ def _R2D2M2CP_beta(
                 beta = pt.set_subtensor(mean, masked)
             beta = pm.Deterministic(name, beta, dims=dims)
         elif psi_mask is not None:
-            beta = mean
+            # all variables are deterministic
+            beta = pm.Deterministic(name, (mu_param / input_sigma), dims=dims)
         else:
             beta = pm.Normal(name, mu_param / input_sigma, std_param / input_sigma, dims=dims)
     return beta

pymc_experimental/tests/distributions/test_multivariate.py

@@ -50,7 +50,7 @@ def r2(self):
     def r2_std(self, request):
         return request.param
 
-    @pytest.fixture(params=["true", "false", "limit-1", "limit-0"])
+    @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)
@@ -64,6 +64,8 @@ def positive_probs(self, input_std, request):
             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):
@@ -122,14 +124,15 @@ def test_init(
         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 ("beta::phi" in model.named_vars) == (
             "variables_importance" in phi_args or "importance_concentration" in phi_args
         ), set(model.named_vars)
-        assert ("beta::psi" in model.named_vars) == (positive_probs_std is not None), set(
-            model.named_vars
-        )
+        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())), model.point_logps()
 
     def test_failing_importance(self, dims, input_shape, output_std, input_std):