Extend Rice distribution

RELEASE-NOTES.md

@@ -20,6 +20,7 @@
   context manager instance. If they do not, the conditional relations between
   the distribution's parameters could be broken, and `random` could return
   values drawn from an incorrect distribution.
+- `Rice` distribution is now defined with either the noncentrality parameter or the shape parameter (#3287).
 
 ### Maintenance
 
@@ -29,6 +30,9 @@
 - Removed use of deprecated `ymin` keyword in matplotlib's `Axes.set_ylim` (#3279)
 - Fix for #3210. Now `distribution.draw_values(params)`, will draw the `params` values from their joint probability distribution and not from combinations of their marginals (Refer to PR #3273).
 - Rewrote `Multinomial._random` method to better handle shape broadcasting (#3271)
+- Fixed `Rice` distribution, which inconsistently mixed two parametrizations (#3286).
+- `Rice` distribution now accepts multiple parameters and observations and is usable with NUTS (#3289).
+
 
 ### Deprecations
 

pymc3/distributions/continuous.py

@@ -12,8 +12,6 @@
 from scipy.interpolate import InterpolatedUnivariateSpline
 import warnings
 
-from theano.scalar import i1, i0
-
 from pymc3.theanof import floatX
 from . import transforms
 from pymc3.util import get_variable_name
@@ -3505,28 +3503,57 @@ class Rice(PositiveContinuous):
     ========  ==============================================================
     Support   :math:`x \in (0, \infty)`
     Mean      :math:`\sigma {\sqrt  {\pi /2}}\,\,L_{{1/2}}(-\nu ^{2}/2\sigma ^{2})`
-    Variance  :math:`2\sigma ^{2}+\nu ^{2}-{\frac  {\pi \sigma ^{2}}{2}}L_{{1/2}}^{2}
-                        \left({\frac  {-\nu ^{2}}{2\sigma ^{2}}}\right)`
+    Variance  :math:`2\sigma ^{2}+\nu ^{2}-{\frac  {\pi \sigma ^{2}}{2}}L_{{1/2}}^{2}\left({\frac  {-\nu ^{2}}{2\sigma ^{2}}}\right)`
     ========  ==============================================================
 
 
     Parameters
     ----------
     nu : float
-        shape parameter.
+        noncentrality parameter.
     sd : float
-        standard deviation.
+        scale parameter.
+    b : float
+        shape parameter (alternative to nu).
+
+    Notes
+    -----
+    The distribution :math:`\mathrm{Rice}\left(|\nu|,\sigma\right)` is the
+    distribution of :math:`R=\sqrt{X^2+Y^2}` where :math:`X\sim N(\nu \cos{\theta}, \sigma^2)`,
+    :math:`Y\sim N(\nu \sin{\theta}, \sigma^2)` are independent and for any
+    real :math:`\theta`.
+
+    The distribution is defined with either nu or b. 
+    The link between the two parametrizations is given by
+
+    .. math::
+
+       b = \dfrac{\nu}{\sigma}
 
     """
 
-    def __init__(self, nu=None, sd=None, *args, **kwargs):
+    def __init__(self, nu=None, sd=None, b=None, *args, **kwargs):
         super(Rice, self).__init__(*args, **kwargs)
+        nu, b, sd = self.get_nu_b(nu, b, sd)
         self.nu = nu = tt.as_tensor_variable(nu)
         self.sd = sd = tt.as_tensor_variable(sd)
+        self.b = b = tt.as_tensor_variable(b)
         self.mean = sd * np.sqrt(np.pi / 2) * tt.exp((-nu**2 / (2 * sd**2)) / 2) * ((1 - (-nu**2 / (2 * sd**2)))
-                                 * i0(-(-nu**2 / (2 * sd**2)) / 2) - (-nu**2 / (2 * sd**2)) * i1(-(-nu**2 / (2 * sd**2)) / 2))
+                                 * tt.i0(-(-nu**2 / (2 * sd**2)) / 2) - (-nu**2 / (2 * sd**2)) * tt.i1(-(-nu**2 / (2 * sd**2)) / 2))
         self.variance = 2 * sd**2 + nu**2 - (np.pi * sd**2 / 2) * (tt.exp((-nu**2 / (2 * sd**2)) / 2) * ((1 - (-nu**2 / (
-            2 * sd**2))) * i0(-(-nu**2 / (2 * sd**2)) / 2) - (-nu**2 / (2 * sd**2)) * i1(-(-nu**2 / (2 * sd**2)) / 2)))**2
+            2 * sd**2))) * tt.i0(-(-nu**2 / (2 * sd**2)) / 2) - (-nu**2 / (2 * sd**2)) * tt.i1(-(-nu**2 / (2 * sd**2)) / 2)))**2
+
+    def get_nu_b(self, nu, b, sd):
+        if sd is None:
+            sd = 1.
+        if nu is None and b is not None:
+            nu = b * sd
+            return nu, b, sd
+        elif nu is not None and b is None:
+            b = nu / sd
+            return nu, b, sd
+        raise ValueError('Rice distribution must specify either nu'
+                         ' or b.')
 
     def random(self, point=None, size=None):
         """
@@ -3547,7 +3574,7 @@ def random(self, point=None, size=None):
         """
         nu, sd = draw_values([self.nu, self.sd],
                              point=point, size=size)
-        return generate_samples(stats.rice.rvs, b=nu, scale=sd, loc=0,
+        return generate_samples(stats.rice.rvs, b=nu / sd, scale=sd, loc=0,
                                 dist_shape=self.shape, size=size)
 
     def logp(self, value):
@@ -3566,8 +3593,9 @@ def logp(self, value):
         """
         nu = self.nu
         sd = self.sd
+        b = self.b
         x = value / sd
-        return bound(tt.log(x * tt.exp((-(x - nu) * (x - nu)) / 2) * i0e(x * nu) / sd),
+        return bound(tt.log(x * tt.exp((-(x - b) * (x - b)) / 2) * i0e(x * b) / sd),
                      sd >= 0,
                      nu >= 0,
                      value > 0,

pymc3/distributions/dist_math.py

@@ -9,7 +9,7 @@
 import scipy.linalg
 import theano.tensor as tt
 import theano
-from theano.scalar import UnaryScalarOp, upgrade_to_float
+from theano.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from theano.tensor.slinalg import Cholesky
 from theano.scan_module import until
 from theano import scan
@@ -270,6 +270,19 @@ def grad(self, inputs, grads):
         return [x_grad * self.grad_op(x)]
 
 
+class I1e(UnaryScalarOp):
+    """
+    Modified Bessel function of the first kind of order 1, exponentially scaled.
+    """
+    nfunc_spec = ('scipy.special.i1e', 1, 1)
+
+    def impl(self, x):
+        return scipy.special.i1e(x)
+
+
+i1e_scalar = I1e(upgrade_to_float_no_complex, name="i1e")
+i1e = tt.Elemwise(i1e_scalar, name="Elemwise{i1e,no_inplace}")
+
 
 class I0e(UnaryScalarOp):
     """
@@ -280,8 +293,14 @@ class I0e(UnaryScalarOp):
     def impl(self, x):
         return scipy.special.i0e(x)
 
+    def grad(self, inp, grads):
+        x, = inp
+        gz, = grads
+        return gz * (i1e_scalar(x) - theano.scalar.sgn(x) * i0e_scalar(x)),
+
 
-i0e = I0e(upgrade_to_float, name='i0e')
+i0e_scalar = I0e(upgrade_to_float_no_complex, name="i0e")
+i0e = tt.Elemwise(i0e_scalar, name="Elemwise{i0e,no_inplace}")
 
 
 def random_choice(*args, **kwargs):

pymc3/tests/test_dist_math.py

@@ -10,7 +10,7 @@
 from ..theanof import floatX
 from ..distributions import Discrete
 from ..distributions.dist_math import (
-    bound, factln, alltrue_scalar, MvNormalLogp, SplineWrapper)
+    bound, factln, alltrue_scalar, MvNormalLogp, SplineWrapper, i0e)
 
 
 def test_bound():
@@ -193,3 +193,12 @@ def test_hessian(self):
         g_x, = tt.grad(spline(x_var), [x_var])
         with pytest.raises(NotImplementedError):
             tt.grad(g_x, [x_var])
+
+
+class TestI0e(object):
+    @theano.configparser.change_flags(compute_test_value="ignore")
+    def test_grad(self):
+        utt.verify_grad(i0e, [0.5])
+        utt.verify_grad(i0e, [-2.])
+        utt.verify_grad(i0e, [[0.5, -2.]])
+        utt.verify_grad(i0e, [[[0.5, -2.]]])

pymc3/tests/test_distributions.py

@@ -1178,7 +1178,9 @@ def test_multidimensional_beta_construction(self):
 
     def test_rice(self):
         self.pymc3_matches_scipy(Rice, Rplus, {'nu': Rplus, 'sd': Rplusbig},
-                                 lambda value, nu, sd: sp.rice.logpdf(value, b=nu, loc=0, scale=sd))
+                                 lambda value, nu, sd: sp.rice.logpdf(value, b=nu / sd, loc=0, scale=sd))
+        self.pymc3_matches_scipy(Rice, Rplus, {'b': Rplus, 'sd': Rplusbig},
+                                 lambda value, b, sd: sp.rice.logpdf(value, b=b, loc=0, scale=sd))
 
     @pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on float32")
     def test_interpolated(self):