adding periodic and unituniform distributions

docs/notebooks/tess.ipynb

@@ -340,7 +340,7 @@
     "        # Orbital parameters for the planets\n",
     "        logP = pm.Normal(\"logP\", mu=np.log(bls_period), sd=1)\n",
     "        t0 = pm.Normal(\"t0\", mu=bls_t0, sd=1)\n",
-    "        b = pm.Uniform(\"b\", lower=0, upper=1, testval=0.5)\n",
+    "        b = xo.distributions.UnitUniform(\"b\")\n",
     "        logr = pm.Normal(\"logr\", sd=1.0,\n",
     "                         mu=0.5*np.log(1e-3*np.array(bls_depth))+np.log(R_star_huang[0]))\n",
     "        r_pl = pm.Deterministic(\"r_pl\", tt.exp(logr))\n",
@@ -375,7 +375,7 @@
     "            ecc=ecc, omega=omega)\n",
     "\n",
     "        # Compute the model light curve using starry\n",
-    "        light_curves = xo.IntegratedLimbDarkLightCurve(u_star).get_light_curve(\n",
+    "        light_curves = xo.LimbDarkLightCurve(u_star).get_light_curve(\n",
     "            orbit=orbit, r=r_pl, t=x[mask], texp=texp)*1e3\n",
     "        light_curve = pm.math.sum(light_curves, axis=-1) + mean\n",
     "        pm.Deterministic(\"light_curves\", light_curves)\n",

docs/notebooks/together.ipynb

@@ -315,9 +315,7 @@
     "                         sd=1.0, shape=2)\n",
     "        r_pl = pm.Deterministic(\"r_pl\", tt.exp(logr))\n",
     "        ror = pm.Deterministic(\"ror\", r_pl / r_star)\n",
-    "        b = pm.Flat(\"b\", shape=2,\n",
-    "                    transform=pm.distributions.transforms.logodds,\n",
-    "                    testval=0.5+np.zeros(2))\n",
+    "        b = xo.distributions.UnitUniform(\"b\", shape=2)\n",
     "        \n",
     "        # This is the eccentricity prior from Kipping (2013):\n",
     "        # https://arxiv.org/abs/1306.4982\n",

docs/user/api.rst

@@ -54,7 +54,9 @@ Distributions
 -------------
 
 .. autoclass:: exoplanet.distributions.UnitVector
+.. autoclass:: exoplanet.distributions.UnitUniform
 .. autoclass:: exoplanet.distributions.Angle
+.. autoclass:: exoplanet.distributions.Periodic
 .. autoclass:: exoplanet.distributions.QuadLimbDark
 .. autoclass:: exoplanet.distributions.RadiusImpact
 .. autofunction:: exoplanet.distributions.get_joint_radius_impact

exoplanet/distributions.py

@@ -5,6 +5,7 @@
 __all__ = [
     "UnitVector",
     "Angle",
+    "UnitUniform",
     "RadiusImpact",
     "QuadLimbDark",
     "get_joint_radius_impact",
@@ -56,7 +57,16 @@ class Angle(pm.Continuous):
     """
 
     def __init__(self, *args, **kwargs):
-        kwargs["transform"] = kwargs.pop("transform", tr.angle)
+        transform = kwargs.pop("transform", None)
+        if transform is None:
+            if "regularized" in kwargs:
+                transform = tr.AngleTransform(
+                    regularized=kwargs.pop("regularized")
+                )
+            else:
+                transform = tr.angle
+        kwargs["transform"] = transform
+
         shape = kwargs.get("shape", None)
         if shape is None:
             testval = 0.0
@@ -80,6 +90,88 @@ def logp(self, value):
         return tt.zeros_like(value)
 
 
+class Periodic(pm.Continuous):
+    """An periodic parameter in a given range
+
+    Like the :class:`Angle` distribution, the actual sampling is performed in
+    a two dimensional vector space ``(sin(theta), cos(theta))`` and then
+    transformed into the range ``[lower, upper)``.
+
+    Args:
+        lower: The lower bound on the range.
+        upper: The upper bound on the range.
+
+    """
+
+    def __init__(self, lower=0, upper=1, **kwargs):
+        self.lower = lower
+        self.upper = upper
+
+        transform = kwargs.pop("transform", None)
+        if transform is None:
+            transform = tr.PeriodicTransform(
+                lower=lower,
+                upper=upper,
+                regularized=kwargs.pop("regularized", 10.0),
+            )
+        kwargs["transform"] = transform
+
+        shape = kwargs.get("shape", None)
+        if shape is None:
+            testval = 0.5 * (lower + upper)
+        else:
+            testval = 0.5 * (lower + upper) + np.zeros(shape)
+        kwargs["testval"] = kwargs.pop("testval", testval)
+        super(Periodic, self).__init__(**kwargs)
+
+    def _random(self, size=None):
+        return np.random.uniform(self.lower, self.upper, size)
+
+    def random(self, point=None, size=None):
+        return generate_samples(
+            self._random,
+            dist_shape=self.shape,
+            broadcast_shape=self.shape,
+            size=size,
+        )
+
+    def logp(self, value):
+        return tt.zeros_like(value)
+
+
+class UnitUniform(pm.Flat):
+    """A uniform distribution between zero and one
+
+    This can sometimes get better performance than ``pm.Uniform.dist(0, 1)``.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        kwargs["transform"] = kwargs.pop(
+            "transform", pm.distributions.transforms.logodds
+        )
+
+        shape = kwargs.get("shape", None)
+        if shape is None:
+            testval = 0.5
+        else:
+            testval = 0.5 + np.zeros(shape)
+        kwargs["testval"] = kwargs.pop("testval", testval)
+
+        super(UnitUniform, self).__init__(*args, **kwargs)
+
+    def _random(self, size=None):
+        return np.random.uniform(0, 1, size)
+
+    def random(self, point=None, size=None):
+        return generate_samples(
+            self._random,
+            dist_shape=self.shape,
+            broadcast_shape=self.shape,
+            size=size,
+        )
+
+
 class QuadLimbDark(pm.Flat):
     """An uninformative prior for quadratic limb darkening parameters
 

exoplanet/distributions_test.py

@@ -8,7 +8,14 @@
 
 import pymc3 as pm
 
-from .distributions import UnitVector, Angle, QuadLimbDark, RadiusImpact
+from .distributions import (
+    UnitVector,
+    UnitUniform,
+    Angle,
+    QuadLimbDark,
+    RadiusImpact,
+    Periodic,
+)
 
 
 class TestDistributions(object):
@@ -78,6 +85,32 @@ def test_angle(self):
             s, p = kstest(theta[:, i], cdf)
             assert s < 0.05
 
+    def test_periodic(self):
+        lower = -3.245
+        upper = 5.123
+        with self._model():
+            Periodic("p", lower=lower, upper=upper, shape=(5, 2))
+            trace = self._sample()
+
+        p = trace["p"]
+        p = np.reshape(p, (len(p), -1))
+        cdf = lambda x: np.clip((x - lower) / (upper - lower), 0, 1)  # NOQA
+        for i in range(p.shape[1]):
+            s, _ = kstest(p[:, i], cdf)
+            assert s < 0.05
+
+    def test_unit_uniform(self):
+        with self._model():
+            UnitUniform("u", shape=(5, 2))
+            trace = self._sample()
+
+        u = trace["u"]
+        u = np.reshape(u, (len(u), -1))
+        cdf = lambda x: np.clip(x, 0, 1)  # NOQA
+        for i in range(u.shape[1]):
+            s, p = kstest(u[:, i], cdf)
+            assert s < 0.05
+
     def test_quad_limb_dark(self):
         with self._model():
             QuadLimbDark("u", shape=2)

exoplanet/transforms.py

@@ -64,7 +64,7 @@ def jacobian_det(self, y):
 
 
 class AngleTransform(tr.Transform):
-    """An angle transformation for PyMC3
+    """An angle transformation
 
     The variable is augmented to sample an isotropic 2D normal and the angle
     is given by the arctan of the ratio of the two coordinates. This will have
@@ -105,6 +105,52 @@ def jacobian_det(self, y):
 angle = AngleTransform()
 
 
+class PeriodicTransform(tr.Transform):
+    """An periodic transformation
+
+    This extends the :class:`Angle` transform to have a uniform distribution
+    between ``lower`` and ``upper``.
+
+    Args:
+        lower: The lower bound of the range.
+        upper: The upper bound of the range.
+        regularized: The amplitude of the regularization term. If ``None``,
+            no regularization is applied. This has no effect on the
+            distribution over the transformed parameter, but it can make
+            sampling more efficient in some cases.
+
+    """
+
+    name = "periodic"
+
+    def __init__(self, lower=0, upper=1, **kwargs):
+        self.mid = tt.as_tensor_variable(0.5 * (lower + upper))
+        self.delta = tt.as_tensor_variable(0.5 * (upper - lower) / np.pi)
+        self.mid_ = 0.5 * (lower + upper)
+        self.delta_ = 0.5 * (upper - lower) / np.pi
+        self.regularized = kwargs.pop("regularized", 10.0)
+        super(PeriodicTransform, self).__init__(**kwargs)
+
+    def backward(self, y):
+        return self.mid + self.delta * tt.arctan2(y[0], y[1])
+
+    def forward(self, x):
+        a = (x - self.mid) / self.delta
+        return tt.concatenate(
+            (tt.shape_padleft(tt.sin(a)), tt.shape_padleft(tt.cos(a))), axis=0
+        )
+
+    def forward_val(self, x, point=None):
+        a = (x - self.mid_) / self.delta_
+        return np.array([np.sin(a), np.cos(a)])
+
+    def jacobian_det(self, y):
+        sm = tt.sum(tt.square(y), axis=0)
+        if self.regularized is not None:
+            return self.regularized * tt.log(sm) - 0.5 * sm
+        return -0.5 * sm
+
+
 class QuadLimbDarkTransform(tr.Transform):
     """A triangle transformation for PyMC3