trying again

exoplanet/distributions_test.py

@@ -1,97 +1,104 @@
 # -*- coding: utf-8 -*-
 
-# from __future__ import division, print_function
-
-# import logging
-# 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 _sample(self, **kwargs):
-#         logger = logging.getLogger("pymc3")
-#         logger.propagate = False
-#         kwargs["draws"] = kwargs.get("draws", 1000)
-#         kwargs["progressbar"] = kwargs.get("progressbar", False)
-#         return pm.sample(**kwargs)
-
-#     def test_unit_vector(self):
-#         with pm.Model():
-#             UnitVector("x", shape=(2, 3))
-#             trace = self._sample()
-
-#         # 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 = self._sample()
-
-#         # 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 = self._sample()
-
-#         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 = self._sample()
-
-#         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))
+from __future__ import division, print_function
+
+import logging
+import numpy as np
+from scipy.stats import kstest
+
+import pymc3 as pm
+
+from .distributions import UnitVector, Angle, QuadLimbDark, RadiusImpact
+
+
+class TestDistributions(object):
+    random_seed = 20160911
+
+    @classmethod
+    def setup_class(cls):
+        np.random.seed(cls.random_seed)
+
+    def setup_method(self):
+        np.random.seed(self.random_seed)
+
+    def _sample(self, **kwargs):
+        logger = logging.getLogger("pymc3")
+        logger.propagate = False
+        kwargs["draws"] = kwargs.get("draws", 1000)
+        kwargs["progressbar"] = kwargs.get("progressbar", False)
+        return pm.sample(**kwargs)
+
+    def test_unit_vector(self):
+        with pm.Model():
+            UnitVector("x", shape=(2, 3))
+            trace = self._sample()
+
+        # 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 = self._sample()
+
+        # 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 = self._sample()
+
+        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 = self._sample()
+
+        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/orbits/keplerian.py

@@ -14,9 +14,6 @@
 from ..citations import add_citations_to_model
 from ..theano_ops.kepler.solver import KeplerOp
 
-gcc_to_sun = (constants.M_sun / constants.R_sun**3).to(u.g / u.cm**3).value
-G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value
-
 
 class KeplerianOrbit(object):
     """A system of bodies on Keplerian orbits around a common central
@@ -71,6 +68,10 @@ def __init__(self,
                  **kwargs):
         add_citations_to_model(self.__citations__, model=model)
 
+        self.gcc_to_sun = (
+            (constants.M_sun / constants.R_sun**3).to(u.g / u.cm**3).value)
+        self.G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value
+
         self.kepler_op = KeplerOp(**kwargs)
 
         # Parameters
@@ -143,7 +144,7 @@ def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star,
                                  "also define rho_star or m_star")
             if r_star is None:
                 r_star = 1.0
-            rho_star = 3*np.pi*(a / r_star)**3 / (G_grav*period**2)
+            rho_star = 3*np.pi*(a / r_star)**3 / (self.G_grav*period**2)
             rho_star -= 3*self.m_planet/(4*np.pi*r_star**3)
             rho_star_units = None
 
@@ -161,7 +162,7 @@ def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star,
             if rho_star_units is not None:
                 rho_star *= (1 * rho_star_units).to(u.M_sun / u.R_sun**3).value
             else:
-                rho_star /= gcc_to_sun
+                rho_star /= self.gcc_to_sun
         if r_star is not None:
             r_star = tt.as_tensor_variable(r_star)
         if m_star is not None:
@@ -177,11 +178,13 @@ def _get_consistent_inputs(self, a, period, rho_star, r_star, m_star,
 
         # Work out the planet parameters
         if a is None:
-            a = (G_grav*(m_star+self.m_planet)*period**2/(4*np.pi**2))**(1./3)
+            a = (self.G_grav*(m_star+self.m_planet)*period**2 /
+                 (4*np.pi**2))**(1./3)
         elif period is None:
-            period = 2*np.pi*a**(3/2)/(tt.sqrt(G_grav*(m_star+self.m_planet)))
+            period = 2*np.pi*a**(3/2)/(
+                tt.sqrt(self.G_grav*(m_star+self.m_planet)))
 
-        return a, period, rho_star * gcc_to_sun, r_star, m_star
+        return a, period, rho_star * self.gcc_to_sun, r_star, m_star
 
     def _rotate_vector(self, x, y):
         if self.ecc is None: