Allow derived icdf for negative non-even power transforms

Author: ricardoV94

pymc/logprob/transforms.py

@@ -178,6 +178,7 @@ class MeasurableTransform(MeasurableElemwise):
         Erf,
         Erfc,
         Erfcx,
+        Sigmoid,
     )
 
     # Cannot use `transform` as name because it would clash with the property added by
@@ -227,7 +228,7 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
     return pt.switch(pt.isnan(jacobian), -np.inf, input_logprob + jacobian)
 
 
-MONOTONICALLY_INCREASING_OPS = (Exp, Log, Add, Sinh, Tanh, ArcSinh, ArcCosh, ArcTanh, Erf)
+MONOTONICALLY_INCREASING_OPS = (Exp, Log, Add, Sinh, Tanh, ArcSinh, ArcCosh, ArcTanh, Erf, Sigmoid)
 MONOTONICALLY_DECREASING_OPS = (Erfc, Erfcx)
 
 
@@ -300,7 +301,18 @@ def measurable_transform_icdf(op: MeasurableTransform, value, *inputs, **kwargs)
         value = pt.switch(pt.lt(scale, 0), 1 - value, value)
     elif isinstance(op.scalar_op, Pow):
         if op.transform_elemwise.power < 0:
-            raise NotImplementedError
+            # Note: Negative even powers will be rejected below when inverting the transform
+            # For the remaining negative powers the function is decreasing with a jump around 0
+            # We adjust the value with the mass below zero.
+            # For non-negative RVs with cdf(0)=0, it simplifies to 1 - value
+            cdf_zero = pt.exp(_logcdf_helper(measurable_input, 0))
+            # Use nan to not mask invalid values accidentally
+            value = pt.switch((value >= 0) & (value <= 1), value, np.nan)
+            value = pt.switch(
+                (cdf_zero > 0) & (value < cdf_zero),
+                cdf_zero - value,
+                1 + cdf_zero - value,
+            )
     else:
         raise NotImplementedError
 

tests/logprob/test_transforms.py

@@ -379,9 +379,7 @@ def test_reciprocal_rv_transform(self, numerator):
         x_vv = x_rv.clone()
         x_logp_fn = pytensor.function([x_vv], logp(x_rv, x_vv))
         x_logcdf_fn = pytensor.function([x_vv], logcdf(x_rv, x_vv))
-
-        with pytest.raises(NotImplementedError):
-            icdf(x_rv, x_vv)
+        x_icdf_fn = pytensor.function([x_vv], icdf(x_rv, x_vv))
 
         x_test_val = np.r_[-0.5, 1.5]
         np.testing.assert_allclose(
@@ -392,6 +390,10 @@ def test_reciprocal_rv_transform(self, numerator):
             x_logcdf_fn(x_test_val),
             sp.stats.invgamma(shape, scale=scale * numerator).logcdf(x_test_val),
         )
+        np.testing.assert_allclose(
+            x_icdf_fn(x_test_val),
+            sp.stats.invgamma(shape, scale=scale * numerator).ppf(x_test_val),
+        )
 
     def test_reciprocal_real_rv_transform(self):
         # 1 / Cauchy(mu, sigma) = Cauchy(mu / (mu^2 + sigma ^2), sigma / (mu ^ 2, sigma ^ 2))
@@ -406,8 +408,10 @@ def test_reciprocal_real_rv_transform(self):
             logcdf(test_rv, test_value).eval(),
             sp.stats.cauchy(1 / 5, 2 / 5).logcdf(test_value),
         )
-        with pytest.raises(NotImplementedError):
-            icdf(test_rv, test_value)
+        np.testing.assert_allclose(
+            icdf(test_rv, test_value).eval(),
+            sp.stats.cauchy(1 / 5, 2 / 5).ppf(test_value),
+        )
 
     def test_sqr_transform(self):
         # The square of a normal with unit variance is a noncentral chi-square with 1 df and nc = mean ** 2