updating notebooks to new distributions

exoplanet-dev · Sep 19, 2019 · bcc5a51 · bcc5a51
1 parent 91fc861
commit bcc5a51
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 71 deletions.
diff --git a/docs/notebooks/citation.ipynb b/docs/notebooks/citation.ipynb
@@ -39,7 +39,7 @@
     "with pm.Model() as model:\n",
     "    u = xo.distributions.QuadLimbDark(\"u\")\n",
     "    orbit = xo.orbits.KeplerianOrbit(period=10.0)\n",
-    "    light_curve = xo.StarryLightCurve(u)\n",
+    "    light_curve = xo.LimbDarkLightCurve(u)\n",
     "    transit = light_curve.get_light_curve(r=0.1, orbit=orbit, t=[0.0, 0.1])\n",
     "    \n",
     "    txt, bib = xo.citations.get_citations_for_model()"

diff --git a/docs/notebooks/tess.ipynb b/docs/notebooks/tess.ipynb
@@ -185,7 +185,7 @@
     "\n",
     "## Transit search\n",
     "\n",
-    "Now, let's use [the box least squares periodogram from AstroPy](http://docs.astropy.org/en/latest/stats/bls.html)\n",
+    "Now, let's use [the box least squares periodogram from AstroPy](http://docs.astropy.org/en/latest/timeseries/bls.html)\n",
     "(Note: you'll need AstroPy v3.1 or more recent to use this feature) to estimate the period, phase, and depth of the transit."
    ]
   },
@@ -195,7 +195,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from astropy.stats import BoxLeastSquares\n",
+    "from astropy.timeseries import BoxLeastSquares\n",
     "\n",
     "period_grid = np.exp(np.linspace(np.log(1), np.log(15), 50000))\n",
     "\n",
@@ -308,7 +308,7 @@
    "source": [
     "## The transit model in PyMC3\n",
     "\n",
-    "The transit model, initialization, and sampling are all nearly the same as the one in :ref:`together`, but we'll use a [more informative prior on eccentricity](https://arxiv.org/abs/1306.4982)."
+    "The transit model, initialization, and sampling are all nearly the same as the one in :ref:`together`."
    ]
   },
   {
@@ -340,30 +340,21 @@
     "        # 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 = 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",
     "        ror = pm.Deterministic(\"ror\", r_pl / r_star)\n",
+    "        b = xo.distributions.ImpactParameter(\"b\", ror=ror)\n",
     "        \n",
     "        # This is the eccentricity prior from Kipping (2013):\n",
     "        # https://arxiv.org/abs/1306.4982\n",
-    "        BoundedBeta = pm.Bound(pm.Beta, lower=0, upper=1-1e-5)\n",
-    "        ecc = BoundedBeta(\"ecc\", alpha=0.867, beta=3.03, testval=0.1)\n",
+    "        ecc = xo.distributions.eccentricity.kipping13(\"ecc\", testval=0.1)\n",
     "        omega = xo.distributions.Angle(\"omega\")\n",
     "        \n",
     "        # Transit jitter & GP parameters\n",
     "        logs2 = pm.Normal(\"logs2\", mu=np.log(np.var(y[mask])), sd=10)\n",
-    "        logw0_guess = np.log(2*np.pi/10)\n",
-    "        logw0 = pm.Normal(\"logw0\", mu=logw0_guess, sd=10)\n",
-    "        \n",
-    "        # We'll parameterize using the maximum power (S_0 * w_0^4) instead of\n",
-    "        # S_0 directly because this removes some of the degeneracies between\n",
-    "        # S_0 and omega_0\n",
-    "        logpower = pm.Normal(\"logpower\",\n",
-    "                             mu=np.log(np.var(y[mask]))+4*logw0_guess,\n",
-    "                             sd=10)\n",
-    "        logS0 = pm.Deterministic(\"logS0\", logpower - 4 * logw0)\n",
+    "        logw0 = pm.Normal(\"logw0\", mu=0, sd=10)\n",
+    "        logSw4 = pm.Normal(\"logSw4\", mu=np.log(np.var(y[mask])), sd=10)\n",
     "\n",
     "        # Tracking planet parameters\n",
     "        period = pm.Deterministic(\"period\", tt.exp(logP))\n",
@@ -381,7 +372,7 @@
     "        pm.Deterministic(\"light_curves\", light_curves)\n",
     "\n",
     "        # GP model for the light curve\n",
-    "        kernel = xo.gp.terms.SHOTerm(log_S0=logS0, log_w0=logw0, Q=1/np.sqrt(2))\n",
+    "        kernel = xo.gp.terms.SHOTerm(log_Sw4=logSw4, log_w0=logw0, Q=1/np.sqrt(2))\n",
     "        gp = xo.gp.GP(kernel, x[mask], tt.exp(logs2) + tt.zeros(mask.sum()), J=2)\n",
     "        pm.Potential(\"transit_obs\", gp.log_likelihood(y[mask] - light_curve))\n",
     "        pm.Deterministic(\"gp_pred\", gp.predict())\n",
@@ -390,7 +381,7 @@
     "        # a better solution by trying different combinations of parameters in turn\n",
     "        if start is None:\n",
     "            start = model.test_point\n",
-    "        map_soln = xo.optimize(start=start, vars=[logs2, logpower, logw0])\n",
+    "        map_soln = xo.optimize(start=start, vars=[logs2, logSw4, logw0])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[logr])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[b])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[logP, t0])\n",
@@ -399,7 +390,7 @@
     "        map_soln = xo.optimize(start=map_soln, vars=[b])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[ecc, omega])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[mean])\n",
-    "        map_soln = xo.optimize(start=map_soln, vars=[logs2, logpower, logw0])\n",
+    "        map_soln = xo.optimize(start=map_soln, vars=[logs2, logSw4, logw0])\n",
     "        map_soln = xo.optimize(start=map_soln)\n",
     "\n",
     "    return model, map_soln\n",
@@ -503,7 +494,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now that we have the model, we can sample it using a :class:`exoplanet.PyMC3Sampler`:"
+    "Now that we have the model, we can sample:"
    ]
   },
   {

diff --git a/docs/notebooks/together.ipynb b/docs/notebooks/together.ipynb
@@ -141,7 +141,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We can initialize the transit parameters using [the box least squares periodogram from AstroPy](http://docs.astropy.org/en/latest/stats/bls.html).\n",
+    "We can initialize the transit parameters using [the box least squares periodogram from AstroPy](http://docs.astropy.org/en/latest/timeseries/bls.html).\n",
     "(Note: you'll need AstroPy v3.1 or more recent to use this feature.)\n",
     "A full discussion of transit detection and vetting is beyond the scope of this tutorial so let's assume that we know that there are two periodic transiting planets in this dataset."
    ]
@@ -152,7 +152,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from astropy.stats import BoxLeastSquares\n",
+    "from astropy.timeseries import BoxLeastSquares\n",
     "\n",
     "m = np.zeros(len(x), dtype=bool)\n",
     "period_grid = np.exp(np.linspace(np.log(5), np.log(50), 50000))\n",
@@ -315,10 +315,10 @@
     "                         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 = xo.distributions.UnitUniform(\"b\", shape=2)\n",
+    "        b = xo.distributions.ImpactParameter(\"b\", ror=ror, shape=2)\n",
     "\n",
-    "        ecc = xo.distributions.UnitUniform(\n",
-    "            \"ecc\", shape=2, testval=np.array([0.01, 0.01]))\n",
+    "        ecc = xo.distributions.eccentricity.vaneylen19(\n",
+    "            \"ecc\", multi=True, shape=2, testval=np.array([0.01, 0.01]))\n",
     "        omega = xo.distributions.Angle(\"omega\", shape=2)\n",
     "\n",
     "        # RV jitter & a quadratic RV trend\n",
@@ -327,16 +327,8 @@
     "\n",
     "        # Transit jitter & GP parameters\n",
     "        logs2 = pm.Normal(\"logs2\", mu=np.log(np.var(y[mask])), sd=10)\n",
-    "        logw0_guess = np.log(2*np.pi/10)\n",
-    "        logw0 = pm.Normal(\"logw0\", mu=logw0_guess, sd=10)\n",
-    "        \n",
-    "        # We'll parameterize using the maximum power (S_0 * w_0^4) instead of\n",
-    "        # S_0 directly because this removes some of the degeneracies between\n",
-    "        # S_0 and omega_0\n",
-    "        logpower = pm.Normal(\"logpower\",\n",
-    "                             mu=np.log(np.var(y[mask]))+4*logw0_guess,\n",
-    "                             sd=10)\n",
-    "        logS0 = pm.Deterministic(\"logS0\", logpower - 4 * logw0)\n",
+    "        logw0 = pm.Normal(\"logw0\", mu=0.0, sd=10)\n",
+    "        logSw4 = pm.Normal(\"logSw4\", mu=np.log(np.var(y[mask])), sd=10)\n",
     "\n",
     "        # Tracking planet parameters\n",
     "        period = pm.Deterministic(\"period\", tt.exp(logP))\n",
@@ -345,9 +337,9 @@
     "        # Orbit model\n",
     "        orbit = xo.orbits.KeplerianOrbit(\n",
     "            r_star=r_star, m_star=m_star,\n",
-    "            period=period, t0=t0, b=b, m_planet=m_pl,\n",
-    "            ecc=ecc, omega=omega,\n",
-    "            m_planet_units=msini.unit)\n",
+    "            period=period, t0=t0, b=b,\n",
+    "            m_planet=xo.units.with_unit(m_pl, msini.unit),\n",
+    "            ecc=ecc, omega=omega)\n",
     "\n",
     "        # Compute the model light curve using starry\n",
     "        light_curves = xo.LimbDarkLightCurve(u_star).get_light_curve(\n",
@@ -356,8 +348,8 @@
     "        pm.Deterministic(\"light_curves\", light_curves)\n",
     "\n",
     "        # GP model for the light curve\n",
-    "        kernel = xo.gp.terms.SHOTerm(log_S0=logS0, log_w0=logw0, Q=1/np.sqrt(2))\n",
-    "        gp = xo.gp.GP(kernel, x[mask], tt.exp(logs2) + tt.zeros(mask.sum()), J=2)\n",
+    "        kernel = xo.gp.terms.SHOTerm(log_Sw4=logSw4, log_w0=logw0, Q=1/np.sqrt(2))\n",
+    "        gp = xo.gp.GP(kernel, x[mask], tt.exp(logs2) + tt.zeros(mask.sum()))\n",
     "        pm.Potential(\"transit_obs\", gp.log_likelihood(y[mask] - light_curve))\n",
     "        pm.Deterministic(\"gp_pred\", gp.predict())\n",
     "\n",
@@ -389,19 +381,12 @@
     "        map_soln = xo.optimize(start=map_soln, vars=[logs2])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[logr, b])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[logP, t0])\n",
-    "        map_soln = xo.optimize(start=map_soln, vars=[logs2, logpower])\n",
+    "        map_soln = xo.optimize(start=map_soln, vars=[logs2, logSw4])\n",
     "        map_soln = xo.optimize(start=map_soln, vars=[logw0])\n",
     "        map_soln = xo.optimize(start=map_soln)\n",
     "\n",
-    "    return model, map_soln"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "    return model, map_soln\n",
+    "\n",
     "model0, map_soln0 = build_model()"
    ]
   },

diff --git a/docs/notebooks/transit.ipynb b/docs/notebooks/transit.ipynb
@@ -92,8 +92,8 @@
    "source": [
     "Then, define the parameters.\n",
     "In this simple model, we'll just fit for the limb darkening parameters of the star, and the period, phase, impact parameter, and radius ratio of the planets (note: this is already 10 parameters and running MCMC to convergence using [emcee](https://emcee.readthedocs.io) would probably take at least an hour).\n",
-    "For the limb darkening, we'll use a quadratic law as parameterized by [Kipping (2013)](https://arxiv.org/abs/1308.0009) and for the joint radius ratio and impact parameter distribution we'll use the parameterization from [Espinoza (2018)](https://arxiv.org/abs/1811.04859).\n",
-    "Both of these reparameterizations are implemented in *exoplanet* as custom *PyMC3* distributions (:class:`exoplanet.distributions.QuadLimbDark` and :class:`exoplanet.distributions.RadiusImpact` respectively)."
+    "For the limb darkening, we'll use a quadratic law as parameterized by [Kipping (2013)](https://arxiv.org/abs/1308.0009).\n",
+    "This reparameterizations is implemented in *exoplanet* as custom *PyMC3* distribution :class:`exoplanet.distributions.QuadLimbDark`."
    ]
   },
   {
@@ -119,19 +119,9 @@
     "    # The Kipping (2013) parameterization for quadratic limb darkening paramters\n",
     "    u = xo.distributions.QuadLimbDark(\"u\", testval=np.array([0.3, 0.2]))\n",
     "    \n",
-    "    # The Espinoza (2018) parameterization for the joint radius ratio and\n",
-    "    # impact parameter distribution\n",
-    "    r, b = xo.distributions.get_joint_radius_impact(\n",
-    "        min_radius=0.01, max_radius=0.1,\n",
-    "        testval_r=np.array([0.04, 0.06]),\n",
-    "        testval_b=np.random.rand(2)\n",
-    "    )\n",
-    "    \n",
-    "    # This shouldn't make a huge difference, but I like to put a uniform\n",
-    "    # prior on the *log* of the radius ratio instead of the value. This\n",
-    "    # can be implemented by adding a custom \"potential\" (log probability).\n",
-    "    pm.Potential(\"r_prior\", -pm.math.log(r))\n",
-    "    \n",
+    "    r = pm.Uniform(\"r\", lower=0.01, upper=0.1, shape=2, testval=np.array([0.04, 0.06]))\n",
+    "    b = xo.distributions.ImpactParameter(\"b\", ror=r, shape=2, testval=np.random.rand(2))\n",
+    "        \n",
     "    # Set up a Keplerian orbit for the planets\n",
     "    orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b)\n",
     "    \n",

diff --git a/exoplanet/distributions/transforms.py b/exoplanet/distributions/transforms.py
@@ -302,7 +302,7 @@ def forward(self, x):
         )
 
     def forward_val(self, x, point=None):
-        opror = draw_values(self.one_plus_ror - 0.0, point=point)
+        opror, = draw_values([self.one_plus_ror - 0.0], point=point)
         return super(ImpactParameterTransform, self).forward_val(
             x / opror, point=point
         )

diff --git a/exoplanet/utils.py b/exoplanet/utils.py
@@ -23,7 +23,12 @@
 from pymc3.blocking import ArrayOrdering, DictToArrayBijection
 from pymc3.model import Point
 from pymc3.theanof import inputvars
-from pymc3.util import get_default_varnames, update_start_vals
+from pymc3.util import (
+    get_default_varnames,
+    update_start_vals,
+    get_untransformed_name,
+    is_transformed_name,
+)
 
 logger = logging.getLogger("exoplanet")
 
@@ -133,10 +138,14 @@ def optimize(
     func = get_theano_function_for_var([nlp] + grad, model=model)
 
     if verbose:
+        names = [
+            get_untransformed_name(v.name)
+            if is_transformed_name(v.name)
+            else v.name
+            for v in vars
+        ]
         sys.stderr.write(
-            "optimizing logp for variables: {0}\n".format(
-                [v.name for v in vars]
-            )
+            "optimizing logp for variables: [{0}]\n".format(",".join(names))
         )
         bar = tqdm.tqdm()