Adding first working example

docs/index.rst

@@ -1,14 +1,31 @@
-Documentation
-=============
+retrieval
+=========
+
+This is the documentation for ``retrieval``.
+
+Here's a simple example:
+
+.. plot::
+
+    from retrieval import transit_depth
+    import astropy.units as u
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    temperatures = np.arange(1000, 3000, 500) * u.K
+
+    for temperature in temperatures:
+        sp = transit_depth(temperature)
+
+        ax = sp.plot(label=temperature)
+
+    ax.set_xlabel('Wavelength [$\mu$m]')
+    ax.set_ylabel('Transit depth')
+    ax.legend()
+    plt.show()
 
-This is the documentation for retrieval.
 
 .. toctree::
   :maxdepth: 2
 
   retrieval/index.rst
-
-.. note:: The layout of this directory is simply a suggestion.  To follow
-          traditional practice, do *not* edit this page, but instead place
-          all documentation for the package inside ``retrieval/``.
-          You can follow this practice or choose your own layout.

retrieval/core.py

@@ -0,0 +1,45 @@
+import numpy as np
+import astropy.units as u
+from astropy.constants import G, k_B, R_jup, M_jup, R_sun
+
+from .opacity import water_opacity
+from .spectrum import Spectrum
+
+
+__all__ = ['transit_depth']
+
+gamma = 0.57721
+
+
+def transit_depth(temperature):
+    """
+    Compute the transit depth with wavelength at ``temperature``.
+
+    Parameters
+    ----------
+    temperature : `~astropy.units.Quantity`
+
+    Returns
+    -------
+    sp : `~retrieval.Spectrum`
+        Transit depth spectrum
+    """
+    wavenumber, kappa = water_opacity(temperature)
+
+
+    g = G * M_jup / R_jup**2
+    rstar = 1 * R_sun
+
+    R0 = R_jup
+    P0 = 1e-3 * u.bar
+
+    mu = 2 * u.u
+
+    scale_height = k_B * temperature / mu / g
+    tau = P0 * kappa / g * np.sqrt(2.0 * np.pi * R0 / scale_height)
+    r = R0 + scale_height * (gamma + np.log(tau))
+
+    depth = (r / rstar) ** 2
+    wavelength = wavenumber.to(u.um, u.spectral())
+
+    return Spectrum(wavelength, depth)

retrieval/opacity.py

@@ -0,0 +1,57 @@
+import os
+
+import numpy as np
+import astropy.units as u
+from scipy.interpolate import RegularGridInterpolator
+
+__all__ = ['water_opacity']
+
+water_opacity_path = os.path.join(os.path.dirname(__file__), 'data',
+                                  'water_opacity_grid.npy')
+
+wavenumber = np.arange(2000, 13000, 5) / u.cm
+
+
+class Interpolator(object):
+    """
+    Lazy loader for the water opacity grid interpolator
+    """
+    def __init__(self):
+        self._interpolator = None
+
+    @property
+    def interp(self):
+        # Lazy load the interpolator
+        if self._interpolator is None:
+            opacity_grid = np.load(water_opacity_path)
+            temperature_grid = np.concatenate([np.arange(50, 700, 50),
+                                               np.arange(700, 1500, 100),
+                                               np.arange(1500, 3100,
+                                                         200)]) * u.K
+            interp = RegularGridInterpolator((wavenumber.value,
+                                              temperature_grid.value),
+                                             opacity_grid)
+            self._interpolator = lambda temp: (interp((wavenumber.value,
+                                                       temp.value))
+                                               * u.cm ** 2 / u.g)
+        return self._interpolator
+
+
+interpolator = Interpolator()
+
+
+def water_opacity(temperature):
+    """
+
+    Parameters
+    ----------
+    temperature : `~astropy.units.Quantity`
+        Temperature of the exoplanet atmosphere
+
+    Returns
+    -------
+    wavenumber : `~astropy.units.Quantity`
+
+    opacity : `~astropy.units.Quantity`
+    """
+    return wavenumber, interpolator.interp(temperature)

retrieval/spectrum.py

@@ -0,0 +1,27 @@
+import matplotlib.pyplot as plt
+
+__all__ = ['Spectrum']
+
+
+class Spectrum(object):
+    """
+    A simple spectrum-like object.
+
+    Consists of a wavelength axis plus a "flux" axis, where "flux" is a
+    stand-in for anything that's measured as a function of wavelength, including
+    flux, transit depth, etc.
+    """
+    def __init__(self, wavelength, flux):
+        self.wavelength = wavelength
+        self.flux = flux
+
+    def plot(self, ax=None, **kwargs):
+        """
+        Plot the wavelength axis against the "flux" axis.
+        """
+        if ax is None:
+            ax = plt.gca()
+
+        ax.plot(self.wavelength, self.flux, **kwargs)
+
+        return ax

tox.ini

@@ -43,6 +43,9 @@ description =
 # The following provides some specific pinnings for key packages
 deps =
 
+    scipy
+    matplotlib
+
     numpy116: numpy==1.16.*
     numpy117: numpy==1.17.*
     numpy118: numpy==1.18.*