shone-pr-11.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecb858af-9ef0-4643-8453-0a66195cd76e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from jax import jit, vmap, numpy as jnp\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import astropy.units as u\n",
    "from astropy.constants import G\n",
    "\n",
    "from shone.chemistry import FastchemWrapper\n",
    "from shone.chemistry.translate import species_name_to_fastchem_name\n",
    "from shone.opacity import Opacity\n",
    "from shone.spectrum import bin_spectrum\n",
    "from shone.transmission import heng_kitzmann_2017, de_wit_seager_2013"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8d27b7e-442a-4c8f-a2de-391f0e1b2fa8",
   "metadata": {},
   "source": [
    "Shop for opacities to load:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ff215e1-7664-4c58-ae47-9e14d6f55793",
   "metadata": {},
   "outputs": [],
   "source": [
    "Opacity.get_available_species()['size_gb'].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac11882-80d3-4171-800c-8b94301a09aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "available_species = Opacity.get_available_species()\n",
    "\n",
    "# available_species"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea63dda1-3e11-40aa-b8ef-61e9e726dadb",
   "metadata": {},
   "source": [
    "Select the species you'd like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d406fbbf-3983-42c4-a5b0-856bdb069c8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "species_indices = [2, 11, 16, 21]\n",
    "\n",
    "available_species[species_indices]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9acf544a-3feb-42d0-896b-1363ecf218e1",
   "metadata": {},
   "source": [
    "Downselect portions of the opacity grid to compile:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fbc0fd3-066e-4f8b-9715-2f94264fa839",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_grid(op):\n",
    "    # set wavelength limits\n",
    "    wavelength_mask = (\n",
    "        (0.5 < op.grid.wavelength) & (op.grid.wavelength < 5.5)\n",
    "    )\n",
    "    wavelength_mask &= (np.arange(len(wavelength_mask)) % 100 == 0)\n",
    "\n",
    "    # set temperature limits\n",
    "    temperature_mask = (\n",
    "        (op.grid.temperature <= op.grid.temperature.max()) & \n",
    "        (op.grid.temperature >= op.grid.temperature.min())\n",
    "    )\n",
    "\n",
    "    op.grid = op.grid.isel(\n",
    "        wavelength=wavelength_mask, \n",
    "        temperature=temperature_mask\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e7dca2-6cb5-4b41-ac41-8a5254f4f6b0",
   "metadata": {},
   "source": [
    "Compile opacity interpolators for each selected species:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "903593e7-9fa6-4fba-827a-9606ecaa846c",
   "metadata": {},
   "outputs": [],
   "source": [
    "opacities = []\n",
    "\n",
    "for idx in tqdm(species_indices):\n",
    "    op = Opacity.load_species_from_index(idx)\n",
    "\n",
    "    # discard some wavelengths:\n",
    "    crop_grid(op)\n",
    "\n",
    "    # get compiled interpolator:\n",
    "    opacities.append(\n",
    "        op.get_interpolator()\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dcc5aa4-b2aa-4093-9449-3f7a5608f05e",
   "metadata": {},
   "source": [
    "We define a p-T grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "718e2cde-1876-452a-b6e1-15c2d0ab505f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pressure = np.geomspace(1e-6, 1.5)\n",
    "temperature = 1500 * (pressure / 0.01) ** 0.08\n",
    "\n",
    "ax = plt.gca()\n",
    "ax.semilogy(temperature, pressure)\n",
    "ax.invert_yaxis()\n",
    "ax.set(\n",
    "    xlabel='Temperature [K]',\n",
    "    ylabel='Pressure [bar]',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ac3831-d26a-4423-8b96-6777088538fc",
   "metadata": {},
   "source": [
    "Compute volume mixing ratios:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d5691e5-a26a-4f44-95ac-3f2bb9a476de",
   "metadata": {},
   "outputs": [],
   "source": [
    "chem = FastchemWrapper(temperature, pressure)\n",
    "\n",
    "from shone.chemistry import get_fastchem_interpolator\n",
    "\n",
    "interp_chem = get_fastchem_interpolator()\n",
    "\n",
    "fastchem_names = [\n",
    "    species_name_to_fastchem_name(name)\n",
    "    for name in available_species.iloc[species_indices]['name']\n",
    "]\n",
    "\n",
    "fastchem_indices = [\n",
    "    # returns 9999999 if not found, so take min:\n",
    "    min(chem.fastchem.getGasSpeciesIndex(name),\n",
    "        chem.fastchem.getElementIndex(name)\n",
    "    ) for name in fastchem_names\n",
    "]\n",
    "\n",
    "weights_amu = np.array([\n",
    "    # returns 9999999 if not found, so take min:\n",
    "    min(chem.fastchem.getGasSpeciesWeight(idx),\n",
    "        chem.fastchem.getElementWeight(idx)\n",
    "    ) for idx in fastchem_indices\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c658e99d-fae7-4567-bac0-125429e81324",
   "metadata": {},
   "source": [
    "Compile the opacity interpolator for all species:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f31cdb-58c2-459c-a6f5-794559a289fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit\n",
    "def interp_opacity(wavelength, temperature, pressure):\n",
    "\n",
    "    interp_op = []\n",
    "\n",
    "    for op in opacities:\n",
    "        interp_op.append(\n",
    "            op(wavelength, temperature, pressure)\n",
    "        )\n",
    "        \n",
    "    total_opacity = jnp.array(interp_op)\n",
    "    return total_opacity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1496d27e-154c-4b01-870b-92d398304d19",
   "metadata": {},
   "source": [
    "Define the wavelength grid to sample in the interpolator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e050eca-a4db-4357-9396-f842d65b8df1",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength = jnp.geomspace(0.5, 5, 10_000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8e4c2d2-0e3f-4bdd-a7e7-5f7905cd7373",
   "metadata": {},
   "source": [
    "Set the planetary mass and radius:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ced243e-d0de-41e5-a85e-ee6c8040392b",
   "metadata": {},
   "outputs": [],
   "source": [
    "R_p0 = (1 * u.R_jup).cgs.value\n",
    "mass = (1 * u.M_jup).cgs.value\n",
    "g = (G * mass / R_p0**2).cgs.value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72bed1dc-a36c-4a38-be97-b43e42a2f960",
   "metadata": {},
   "source": [
    "Define a compiled model for the radius spectrum, and bin the model spectrum to the observed wavelength grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c03e7b8-c3d2-4274-b757-8047d9d8c656",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "obs_wavelength = np.linspace(0.55, 4.5, 200)\n",
    "\n",
    "@partial(jit, static_argnames=\"rayleigh_scattering absorption\".split())\n",
    "def compute_transmission_spectrum(\n",
    "    temperature=temperature, pressure=pressure, \n",
    "    m_to_h=1, c_to_o=0.58, \n",
    "    rayleigh_scattering=True, absorption=True\n",
    "):\n",
    "    vmr = interp_chem(\n",
    "        temperature, pressure, \n",
    "        log_m_to_h=jnp.log10(m_to_h), log_c_to_o=jnp.log10(c_to_o)\n",
    "    )\n",
    "\n",
    "    opacity = interp_opacity(wavelength, temperature, pressure)\n",
    "    \n",
    "    R_p = de_wit_seager_2013.transmission_radius(\n",
    "        wavelength, temperature, pressure, g, R_p0, \n",
    "        opacity, vmr, np.array(fastchem_indices),\n",
    "        weights_amu=weights_amu, \n",
    "        rayleigh_scattering=rayleigh_scattering, \n",
    "        absorption=absorption\n",
    "    )\n",
    "    R_p = R_p / u.R_jup.to(u.cm) # normalize to jupiter radii\n",
    "\n",
    "    R_p_binned = bin_spectrum(obs_wavelength, wavelength, R_p)\n",
    "    return R_p, R_p_binned"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d63f264-45fb-4110-9330-4a62f391a9e3",
   "metadata": {},
   "source": [
    "Compute the transmission spectrum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a76b3f53-4bd9-4744-af64-1855ca23913c",
   "metadata": {},
   "outputs": [],
   "source": [
    "R_p, R_p_binned = compute_transmission_spectrum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff3faad6-52f2-4791-aa5c-f1e10a254f36",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(12, 6))\n",
    "ax.semilogx(wavelength, R_p, alpha=0.8, lw=0.5, color='silver')\n",
    "ax.scatter(obs_wavelength, R_p_binned, marker='o', color='k', zorder=5)\n",
    "\n",
    "xticks = [0.5, 1, 2, 3, 4, 5]\n",
    "ax.set(\n",
    "    xlabel='Wavelength [µm]',\n",
    "    ylabel='Radius [R$_{\\\\rm Jup}$]',\n",
    "    xticks=xticks,\n",
    "    xticklabels=xticks,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5077efeb-30d9-4382-98dd-291a808f794f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_interactive(\n",
    "    T_0=1500, alpha=0.08, min_temp=600, \n",
    "    m_to_h=1, c_to_o=0.58, \n",
    "    rayleigh_scattering=True, absorption=True\n",
    "):\n",
    "    pressure = np.geomspace(1e-6, 1.5)\n",
    "    temperature = T_0 * (pressure / 0.01) ** alpha\n",
    "    temperature = np.clip(temperature, min_temp, None)\n",
    "\n",
    "    R_p, R_p_binned = compute_transmission_spectrum(\n",
    "        temperature, pressure, m_to_h=m_to_h, c_to_o=c_to_o,\n",
    "        rayleigh_scattering=rayleigh_scattering, absorption=absorption\n",
    "    )\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 2, figsize=(15, 6))\n",
    "\n",
    "    # I'm downsampling the wavelengths that get plotted here only because \n",
    "    # they slow down the interactivity\n",
    "    ax[0].semilogx(wavelength[::5], R_p[::5], alpha=0.8, lw=0.5, color='silver')\n",
    "    ax[0].scatter(obs_wavelength, R_p_binned, marker='o', color='k', zorder=5)\n",
    "    xticks = [0.5, 1, 2, 3, 4, 5]\n",
    "    ax[0].set(\n",
    "        xlabel='Wavelength [µm]',\n",
    "        ylabel='Radius [R$_{\\\\rm Jup}$]',\n",
    "        xticks=xticks,\n",
    "        xticklabels=xticks,\n",
    "    )\n",
    "\n",
    "\n",
    "    title = (\n",
    "        f'$T(p) = {{\\\\rm min}} \\\\left( {min_temp:.0f}, {T_0:.0f} \\\\left '\n",
    "        f'( \\\\frac{{p}}{{0.01~{{\\\\rm bar}}}} \\\\right)^{{{alpha:.2f}}} \\\\right)$ K'\n",
    "    )\n",
    "    ax[1].semilogy(temperature, pressure)\n",
    "    ax[1].invert_yaxis()\n",
    "    ax[1].set(\n",
    "        xlabel='Temperature [K]',\n",
    "        ylabel='Pressure [bar]',\n",
    "        title=title\n",
    "    )\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dcba74c-b278-4290-b31c-4dab7a4c97cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interactive\n",
    "\n",
    "interactive(\n",
    "    plot_interactive,\n",
    "    T_0=(900, 2300, 50),\n",
    "    alpha=(0.05, 0.1, 0.001),\n",
    "    min_temp=(500, 2000, 50),\n",
    "    m_to_h=(0.1, 10, 0.1),\n",
    "    c_to_o=(0.1, 1, 0.05),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8cdb8b8-e748-407a-a81e-ebfb502780e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "compute_transmission_spectrum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3d11742-7cf5-45de-8525-9e92c0e95c57",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}