Adding MCMC, ABC example

example/abc.py

@@ -0,0 +1,47 @@
+import sys
+sys.path.insert(0, '../')
+
+from retrieval import Planet
+
+import numpy as np
+import astropy.units as u
+import matplotlib.pyplot as plt
+
+example_spectrum = np.load('../retrieval/data/example_spectrum.npy')
+
+planet = Planet(1 * u.M_jup, 1 * u.R_jup, 1e-3 * u.bar, 2.2 * u.u)
+
+
+def distance(theta):
+    temperature = theta[0] * u.K
+    model = planet.transit_depth(temperature).flux
+    return np.sum((example_spectrum[:, 1] - model)**2 /
+                  example_spectrum[:, 2]**2)
+
+init_temp = 1480
+
+distance_chain = [distance([init_temp])]
+temperature_chain = [init_temp]
+
+n_steps = 1000
+
+threshold = 1.5
+i = 0
+total_steps = 1
+
+while len(temperature_chain) < n_steps:
+    total_steps += 1
+    trial_temp = temperature_chain[i] + 10 * np.random.randn()
+    trial_dist = distance([trial_temp])
+
+    if trial_dist < threshold:
+        i += 1
+        temperature_chain.append(trial_temp)
+        distance_chain.append(trial_dist)
+
+acceptance_rate = i / total_steps
+print("acceptance rate = ", acceptance_rate)
+
+plt.hist(temperature_chain)
+plt.xlabel('Temperature [K]')
+plt.show()

example/mcmc.py

@@ -25,13 +25,15 @@ def lnprior(theta):
 def lnlikelihood(theta):
     temperature = theta[0] * u.K
     model = planet.transit_depth(temperature).flux
-    return -0.5 * np.sum((example_spectrum[:, 1] - model)**2 /
-                         example_spectrum[:, 2]**2)
+    lp = lnprior(theta)
+    return lp + -0.5 * np.sum((example_spectrum[:, 1] - model)**2 /
+                               example_spectrum[:, 2]**2)
 
 nwalkers = 10
 ndim = 1
 
-p0 = [[1500 + 10 * np.random.randn()] for i in range(nwalkers)]
+p0 = [[1500 + 10 * np.random.randn()]
+      for i in range(nwalkers)]
 
 with Pool() as pool:
     sampler = EnsembleSampler(nwalkers, ndim, lnlikelihood, pool=pool)

retrieval/core.py

@@ -16,6 +16,18 @@ class Planet(object):
     Properties of an exoplanet.
     """
     def __init__(self, mass, radius, pressure, mu):
+        """
+        Parameters
+        ----------
+        mass : `~astropy.unit.Quantity`
+            Mass of the transiting planet
+        radius : `~astropy.unit.Quantity`
+            Radius of the transiting planet
+        pressure : `~astropy.unit.Quantity`
+            Pressure level probed in transmission
+        mu : `~astropy.unit.Quantity`
+            Mean molecular weight of the atmosphere
+        """
         self.mass = mass
         self.radius = radius
         self.pressure = pressure
@@ -28,6 +40,10 @@ def transit_depth(self, temperature, rstar=1 * u.R_sun):
         Parameters
         ----------
         temperature : `~astropy.units.Quantity`
+            Temperature of the atmosphere observed in transmission
+        rstar : `~astropy.units.Quantity`
+            Radius of the star, used to compute the ratio of the planet-to-star
+            radius. Default is one solar radius.
 
         Returns
         -------