{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tidal effects\n",
    "\n",
    "In this notebook we explore the Roche lobe effects in one-dimensional Parker wind models for exospheres of hot exoplanets. These effects were notably discussed in [Erkaev et al. (2007)](https://ui.adsabs.harvard.edu/abs/2007A%26A...472..329E/abstract), and were implemented in `p-winds` [version 1.3.0](https://github.com/ladsantos/p-winds/releases/tag/v1.3.0).\n",
    "\n",
    "The process of including the Roche lobe (or tidal) effects in `p-winds` models is simple, and it involves changing how the radius at the sonic point and outflow structure are calculated.\n",
    "\n",
    "We start by importing all the necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.pylab as pylab\n",
    "import astropy.units as u\n",
    "import astropy.constants as c\n",
    "from p_winds import tools, parker, hydrogen, helium\n",
    "\n",
    "# Uncomment the next line if you have a MacBook with retina screen\n",
    "# %config InlineBackend.figure_format = 'retina'\n",
    "pylab.rcParams['figure.figsize'] = 9.0,6.5\n",
    "pylab.rcParams['font.size'] = 18"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to replicate the [quickstart example for HD 209458 b](https://p-winds.readthedocs.io/en/latest/quickstart.html) and include the tidal effects. We will assume that the planet has an isothermal upper atmosphere with temperature of $9\\,100$ K and a total mass loss rate of $2 \\times 10^{10}$ g s$^{-1}$ based on the results from [Salz et al. 2016](https://ui.adsabs.harvard.edu/abs/2016A%26A...586A..75S/abstract). We will also assume: \n",
    "* The atmosphere is made up of only H and He\n",
    "* The H number fraction is $0.9$\n",
    "* Initially a fully neutral atmosphere (this is going to be self-consistently calculated later).\n",
    "\n",
    "We will also need to know other parameters, namely: the stellar mass and radius, and the semi-major axis of the orbit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HD 209458 b planetary parameters, measured\n",
    "R_pl = 1.39  # Planetary radius in Jupiter radii\n",
    "M_pl = 0.73  # Planetary mass in Jupiter masses\n",
    "impact_parameter = 0.499  # Transit impact parameter\n",
    "a_pl = 0.04634  # Orbital semi-major axis in astronomical units\n",
    "\n",
    "# HD 209458 stellar parameters\n",
    "R_star = 1.20  # Stellar radius in solar radii\n",
    "M_star = 1.07  # Stellar mass in solar masses\n",
    "\n",
    "# A few assumptions about the planet's atmosphere\n",
    "m_dot = 10 ** 10.27  # Total atmospheric escape rate in g / s\n",
    "T_0 = 9100  # Wind temperature in K\n",
    "h_fraction = 0.90  # H number fraction\n",
    "he_fraction = 1 - h_fraction  # He number fraction\n",
    "he_h_fraction = he_fraction / h_fraction\n",
    "mean_f_ion = 0.0  # Mean ionization fraction (will be self-consistently calculated later)\n",
    "mu_0 = (1 + 4 * he_h_fraction) / (1 + he_h_fraction + mean_f_ion)  \n",
    "# mu_0 is the constant mean molecular weight (assumed for now, will be updated later)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we retrieve the high-energy spectrum of the host star with fluxes at the planet. For this example, we use the solar spectrum for convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "units = {'wavelength': u.angstrom, 'flux': u.erg / u.s / u.cm ** 2 / u.angstrom}\n",
    "spectrum = tools.make_spectrum_from_file('../../data/solar_spectrum_scaled_lambda.dat',\n",
    "                                    units)\n",
    "plt.loglog(spectrum['wavelength'], spectrum['flux_lambda'])\n",
    "plt.ylim(1E-5, 1E4)\n",
    "plt.xlabel(r'Wavelength (${\\rm \\AA}$)')\n",
    "plt.ylabel(r'Flux density (erg s$^{-1}$ cm$^{-2}$ ${\\rm \\AA}^{-1}$)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can calculate the distribution of ionized/neutral hydrogen. And here is where things differ a bit to include the tidal effects. We need to specify what is the stellar mass and the semi-major axis of the planet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_f_ion = 0.0\n",
    "r = np.logspace(0, np.log10(20), 100)  # Radial distance profile in unit of planetary radii\n",
    "\n",
    "f_r, mu_bar = hydrogen.ion_fraction(r, R_pl, T_0, h_fraction, \n",
    "                            m_dot, M_pl, mu_0, star_mass=M_star, \n",
    "                            semimajor_axis=a_pl,\n",
    "                            spectrum_at_planet=spectrum, exact_phi=True,\n",
    "                            initial_f_ion=initial_f_ion, relax_solution=True,\n",
    "                            return_mu=True)\n",
    "\n",
    "f_ion = f_r\n",
    "f_neutral = 1 - f_r\n",
    "\n",
    "plt.plot(r, f_neutral, color='C0', label='f$_\\mathrm{neutral}$')\n",
    "plt.plot(r, f_ion, color='C1', label='f$_\\mathrm{ion}$')\n",
    "plt.xlabel(r'Radius (R$_\\mathrm{pl}$)')\n",
    "plt.ylabel('Number fraction')\n",
    "plt.xlim(1, 10)\n",
    "plt.ylim(0, 1)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of using the functions `parker.radius_sonic_point()` and `parker.structure()`, we use `parker.radius_sonic_point_tidal()` and `parker.structure_tidal()`. Their inputs are a bit different, but the outputs are the same. As before, the velocities and densities calculated by `parker.structure_tidal()` are measured in units of sound speed and density at the sonic point, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vs = parker.sound_speed(T_0, mu_bar)  # Speed of sound (km/s, assumed to be constant)\n",
    "rs = parker.radius_sonic_point_tidal(M_pl, vs, M_star, a_pl)  # Radius at the sonic point (jupiterRad)\n",
    "rhos = parker.density_sonic_point(m_dot, rs, vs)  # Density at the sonic point (g/cm^3)\n",
    "\n",
    "r_array = r * R_pl / rs\n",
    "v_array, rho_array = parker.structure_tidal(r_array, vs, rs, M_pl, M_star, a_pl)\n",
    "\n",
    "# Convenience arrays for the plots\n",
    "r_plot = r_array * rs / R_pl\n",
    "v_plot = v_array * vs\n",
    "rho_plot = rho_array * rhos\n",
    "\n",
    "# Finally plot the structure of the upper atmosphere\n",
    "# The circles mark the velocity and density at the sonic point\n",
    "ax1 = plt.subplot()\n",
    "ax2 = ax1.twinx()\n",
    "ax1.semilogy(r_plot, v_plot, color='C0')\n",
    "ax1.plot(rs / R_pl, vs, marker='o', markeredgecolor='w', color='C0')\n",
    "ax2.semilogy(r_plot, rho_plot, color='C1')\n",
    "ax2.plot(rs / R_pl, rhos, marker='o', markeredgecolor='w', color='C1')\n",
    "ax1.set_xlabel(r'Radius (R$_{\\rm pl}$)')\n",
    "ax1.set_ylabel(r'Velocity (km s$^{-1}$)', color='C0')\n",
    "ax2.set_ylabel(r'Density (g cm$^{-3}$)', color='C1')\n",
    "ax1.set_xlim(1, 10)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare the structure with and without tidal effects to see how they differ. We will plot the profiles with the effects implemented in full lines, and those without in dashed lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_r_nte, mu_bar_nte = hydrogen.ion_fraction(r, R_pl, T_0, h_fraction, \n",
    "                            m_dot, M_pl, mu_0,\n",
    "                            spectrum_at_planet=spectrum, exact_phi=True,\n",
    "                            initial_f_ion=initial_f_ion, relax_solution=True,\n",
    "                            return_mu=True)\n",
    "\n",
    "vs_nte = parker.sound_speed(T_0, mu_bar_nte)  # Speed of sound (km/s, assumed to be constant)\n",
    "rs_nte = parker.radius_sonic_point(M_pl, vs_nte)  # Radius at the sonic point (jupiterRad)\n",
    "rhos_nte = parker.density_sonic_point(m_dot, rs_nte, vs_nte)  # Density at the sonic point (g/cm^3)\n",
    "\n",
    "r_array_nte = r * R_pl / rs_nte\n",
    "v_array_nte, rho_array_nte = parker.structure(r_array_nte)\n",
    "\n",
    "# Convenience arrays for the plots\n",
    "v_plot_nte = v_array_nte * vs_nte\n",
    "rho_plot_nte = rho_array_nte * rhos_nte\n",
    "\n",
    "# Finally plot the structure of the upper atmosphere\n",
    "# The circles mark the velocity and density at the sonic point\n",
    "ax1 = plt.subplot()\n",
    "ax2 = ax1.twinx()\n",
    "\n",
    "# With tidal effects\n",
    "ax1.semilogy(r_plot, v_plot, color='C0')\n",
    "ax1.plot(rs / R_pl, vs, marker='o', markeredgecolor='w', color='C0')\n",
    "ax2.semilogy(r_plot, rho_plot, color='C1')\n",
    "ax2.plot(rs / R_pl, rhos, marker='o', markeredgecolor='w', color='C1')\n",
    "\n",
    "# No tidal effects\n",
    "ax1.semilogy(r_plot, v_plot_nte, color='C0', ls='--')\n",
    "ax1.plot(rs_nte / R_pl, vs_nte, marker='o', markeredgecolor='w', color='C0')\n",
    "ax2.semilogy(r_plot, rho_plot_nte, color='C1', ls='--')\n",
    "ax2.plot(rs_nte / R_pl, rhos_nte, marker='o', markeredgecolor='w', color='C1')\n",
    "\n",
    "ax1.set_xlabel(r'Radius (R$_{\\rm pl}$)')\n",
    "ax1.set_ylabel(r'Velocity (km s$^{-1}$)', color='C0')\n",
    "ax2.set_ylabel(r'Density (g cm$^{-3}$)', color='C1')\n",
    "ax1.set_xlim(1, 10)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next step, we calculate the population of helium in singlet, triplet and ionized states using `helium.population_fraction()`. Similar to the hydrogen module, we integrate starting from the inner layer of the atmosphere."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In the initial state, the fraction of singlet and triplet helium \n",
    "# are, respectively, 1.0 and 0.0\n",
    "initial_state = np.array([1.0, 0.0])\n",
    "f_he_1, f_he_3 = helium.population_fraction(\n",
    "    r, v_array, rho_array, f_ion,\n",
    "    R_pl, T_0, h_fraction, vs, rs, rhos, spectrum,\n",
    "    initial_state=initial_state, relax_solution=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we plot the number densities of ionized helium, and helium in singlet and triplet states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hydrogen atom mass\n",
    "m_h = c.m_p.to(u.g).value\n",
    "\n",
    "# Number density of helium nuclei \n",
    "he_fraction = 1 - h_fraction\n",
    "n_he = (rho_array * rhos * he_fraction / (h_fraction + 4 * he_fraction) / m_h)\n",
    "\n",
    "n_he_1 = f_he_1 * n_he\n",
    "n_he_3 = f_he_3 * n_he\n",
    "n_he_ion = (1 - f_he_1 - f_he_3) * n_he\n",
    "\n",
    "plt.semilogy(r, n_he_1, color='C0', label='He singlet')\n",
    "plt.semilogy(r, n_he_3, color='C1', label='He triplet')\n",
    "plt.semilogy(r, n_he_ion, color='C2', label='He ionized')\n",
    "plt.xlabel(r'Radius (R$_\\mathrm{pl}$)')\n",
    "plt.ylabel('Number density (cm$^{-3}$)')\n",
    "plt.xlim(1, 10)\n",
    "plt.ylim(1E-2, 1E10)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the resulting structure and He number density profiles are significantly different than the ones from the [quickstart example](https://p-winds.readthedocs.io/en/latest/quickstart.html), so for HD 209458 b the tidal effects are very important. You can find more details about these effects in [Erkaev et al. (2007)](https://ui.adsabs.harvard.edu/abs/2007A%26A...472..329E/abstract), [Murray-Clay et al. (2009)](https://ui.adsabs.harvard.edu/abs/2009ApJ...693...23M/abstract) and in Appendix E in [Vissapragada et al. (2022)](https://ui.adsabs.harvard.edu/abs/2022arXiv220411865V/abstract)."
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}