adding tests for distributions

exoplanet/distributions.py

@@ -2,7 +2,10 @@
 
 from __future__ import division, print_function
 
-__all__ = ["UnitVector", "Angle", "RadiusImpactParameter"]
+__all__ = [
+    "UnitVector", "Angle", "RadiusImpact", "QuadLimbDark",
+    "get_joint_radius_impact",
+]
 
 import numpy as np
 
@@ -60,7 +63,7 @@ def random(self, point=None, size=None):
                                 size=size)
 
 
-class Triangle(pm.Flat):
+class QuadLimbDark(pm.Flat):
     """An uninformative prior for quadratic limb darkening parameters
 
     This is an implementation of the `Kipping (2013)
@@ -85,9 +88,9 @@ def __init__(self, *args, **kwargs):
                 raise ValueError("the first dimension should be exactly 2")
 
         kwargs["shape"] = shape
-        kwargs["transform"] = tr.triangle
+        kwargs["transform"] = tr.quad_limb_dark
 
-        super(Triangle, self).__init__(*args, **kwargs)
+        super(QuadLimbDark, self).__init__(*args, **kwargs)
 
         # Work out some reasonable starting values for the parameters
         default = np.zeros(shape)
@@ -113,7 +116,7 @@ def random(self, point=None, size=None):
                                 size=size)
 
 
-class RadiusImpactParameter(pm.Flat):
+class RadiusImpact(pm.Flat):
     """The Espinoza (2018) distribution over radius and impact parameter
 
     This is an implementation of `Espinoza (2018)
@@ -122,6 +125,10 @@ class RadiusImpactParameter(pm.Flat):
     radius ratio will be in the zeroth entry in the first dimension and
     the impact parameter will be in the first.
 
+    Args:
+        min_radius: The minimum allowed radius.
+        max_radius: The maximum allowed radius.
+
     """
     __citations__ = ("espinoza18", )
 
@@ -143,7 +150,7 @@ def __init__(self, *args, **kwargs):
         kwargs["shape"] = shape
         kwargs["transform"] = transform
 
-        super(RadiusImpactParameter, self).__init__(*args, **kwargs)
+        super(RadiusImpact, self).__init__(*args, **kwargs)
 
         # Work out some reasonable starting values for the parameters
         default = np.zeros(shape)
@@ -188,10 +195,10 @@ def random(self, point=None, size=None):
                                 size=size)
 
 
-def get_joint_r_and_b_distribution(name="", N_planets=None,
-                                   min_radius=0, max_radius=1,
-                                   r_star=None, testval_r=None, testval_b=None,
-                                   **kwargs):
+def get_joint_radius_impact(name="", N_planets=None,
+                            min_radius=0, max_radius=1,
+                            testval_r=None, testval_b=None,
+                            **kwargs):
     """Get the joint distribution over radius and impact parameter
 
     This uses the Espinoza (2018) parameterization of the distribution (see
@@ -201,8 +208,8 @@ def get_joint_r_and_b_distribution(name="", N_planets=None,
         name (Optional[str]): A prefix that is added to all distribution names
             used in this parameterization. For example, if ``name`` is
             ``param_``, vars will be added to the PyMC3 model with names:
-            ``param_rb`` (for the joint distribution), ``param_b``,
-            ``param_r``, and optionally ``param_ror`` if ``r_star`` is given.
+            ``param_rb`` (for the joint distribution), ``param_b``, and
+            ``param_r``.
         N_planets (Optional[int]): The number of planets. If not provided, it
             will be inferred from the ``testval_*`` parameters or assumed to
             be 1.
@@ -241,7 +248,7 @@ def get_joint_r_and_b_distribution(name="", N_planets=None,
         rb_test[1, :] = testval_b
 
     # Construct the join distribution
-    rb = RadiusImpactParameter(
+    rb = RadiusImpact(
         name + "rb", min_radius=min_radius, max_radius=max_radius,
         shape=(2, N_planets), testval=rb_test, **kwargs)
 

exoplanet/distributions_test.py

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division, print_function
+
+import numpy as np
+from scipy.stats import kstest
+
+import pymc3 as pm
+from pymc3.tests.helpers import SeededTest
+
+from .distributions import UnitVector, Angle, QuadLimbDark, RadiusImpact
+
+
+class TestDistributions(SeededTest):
+
+    def test_unit_vector(self):
+        with pm.Model():
+            UnitVector("x", shape=(2, 3))
+            trace = pm.sample(draws=1000)
+
+        # Make sure that the unit vector constraint is satisfied
+        assert np.allclose(np.sum(trace["x"]**2, axis=-1), 1.0)
+
+        # Pull out the component and compute the angle
+        x = trace["x"][:, :, 0]
+        y = trace["x"][:, :, 1]
+        z = trace["x"][:, :, 2]
+        theta = np.arctan2(y, x)
+
+        # The angle should be uniformly distributed
+        cdf = lambda x: np.clip((x + np.pi) / (2 * np.pi), 0, 1)  # NOQA
+        for i in range(theta.shape[1]):
+            s, p = kstest(theta[:, i], cdf)
+            assert s < 0.05
+
+        # As should the vertical component
+        cdf = lambda x: np.clip((x + 1) / 2, 0, 1)  # NOQA
+        for i in range(z.shape[1]):
+            s, p = kstest(z[:, i], cdf)
+            assert s < 0.05
+
+    def test_angle(self):
+        with pm.Model():
+            Angle("theta", shape=(5, 2))
+            trace = pm.sample(draws=1000)
+
+        # The angle should be uniformly distributed
+        theta = trace["theta"]
+        theta = np.reshape(theta, (len(theta), -1))
+        cdf = lambda x: np.clip((x + np.pi) / (2 * np.pi), 0, 1)  # NOQA
+        for i in range(theta.shape[1]):
+            s, p = kstest(theta[:, i], cdf)
+            assert s < 0.05
+
+    def test_quad_limb_dark(self):
+        with pm.Model():
+            QuadLimbDark("u", shape=2)
+            trace = pm.sample(draws=1000)
+
+        u1 = trace["u"][:, 0]
+        u2 = trace["u"][:, 1]
+
+        # Make sure that the physical constraints are satisfied
+        assert np.all(u1 + u2 < 1)
+        assert np.all(u1 > 0)
+        assert np.all(u1 + 2 * u2 > 0)
+
+        # Make sure that the qs are uniform
+        q1 = (u1 + u2) ** 2
+        q2 = 0.5 * u1 / (u1 + u2)
+
+        cdf = lambda x: np.clip(x, 0, 1)  # NOQA
+        for q in (q1, q2):
+            s, p = kstest(q, cdf)
+            assert s < 0.05
+
+    def test_radius_impact(self):
+        min_radius = 0.01
+        max_radius = 0.1
+        with pm.Model():
+            RadiusImpact("rb", min_radius=min_radius, max_radius=max_radius)
+            trace = pm.sample(draws=1000)
+
+        r = trace["rb"][:, 0]
+        b = trace["rb"][:, 1]
+
+        # Make sure that the physical constraints are satisfied
+        assert np.all((r <= max_radius) & (min_radius <= r))
+        assert np.all((b >= 0) & (b <= 1 + r))

exoplanet/transforms.py

@@ -2,7 +2,7 @@
 
 from __future__ import division, print_function
 
-__all__ = ["unit_vector", "angle", "triangle", "radius_impact"]
+__all__ = ["unit_vector", "angle", "quad_limb_dark", "radius_impact"]
 
 import numpy as np
 
@@ -68,14 +68,14 @@ def jacobian_det(self, y):
 angle = AngleTransform()
 
 
-class TriangleTransform(tr.Transform):
+class QuadLimbDarkTransform(tr.Transform):
     """A triangle transformation for PyMC3
 
     Ref: https://arxiv.org/abs/1308.0009
 
     """
 
-    name = "triangle"
+    name = "quadlimbdark"
 
     def backward(self, y):
         q = tt.nnet.sigmoid(y)
@@ -108,7 +108,7 @@ def jacobian_det(self, y):
         return -2 * tt.nnet.softplus(-y) - y
 
 
-triangle = TriangleTransform()
+quad_limb_dark = QuadLimbDarkTransform()
 
 
 class RadiusImpactTransform(tr.Transform):