{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Rossiter-McLaughlin effect"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we'll show how to use `starry` to model the Rossiter-McLaughlin effect (and radial velocity measurements in general). Check out the notebook on the **Derivation of the radial velocity field** to understand how `starry` models radial velocity under the hood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "%run notebook_setup.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import starry\n",
    "import numpy as np\n",
    "import astropy.units as u\n",
    "\n",
    "starry.config.lazy = False\n",
    "starry.config.quiet = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The basics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To define a radial velocity map, simply pass the `rv=True` keyword when instantiating a `starry` `Map` object. We'll add quadratic limb darkening just for fun. Note that the spherical harmonic degree `ydeg` is implicitly set to zero, since we're not modeling any brightness structure on the surface of the star other than limb darkening (though we can; see below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map = starry.Map(udeg=2, rv=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's set the properties that affect the projected radial velocity field. We'll set the inclination (in degrees), obliquity (also in degrees), equatorial velocity (in meters per second), and the differential rotation shear (unitless):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map.inc = 60\n",
    "map.obl = 30\n",
    "map.veq = 1.0e4\n",
    "map.alpha = 0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also set the limb darkening coefficients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map[1] = 0.5\n",
    "map[2] = 0.25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see what the map currently looks like. We can choose to either view the brightness map (using `rv=False`) or the velocity-weighted brightness map (using `rv=True`). The former is uninteresting, so let's view the latter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map.show(rv=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the map is inclined toward the observer and rotated on the plane of the sky. The left hemisphere is blueshifted (negative radial velocities) and the right hemisphere is redshifted (positive radial velocities). Limb darkening and differential rotation add some additional structure, causing deviations from a perfect dipolar field."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example, we can plot what happens when the star is transited by a planet. As usual, we can call `map.flux()` to get a light curve, but for RV maps we can also call `map.rv()` to get the radial velocity anomaly one would measure from the object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Occultor properties\n",
    "xo = np.linspace(-1.5, 1.5, 1000)\n",
    "yo = -0.25\n",
    "ro = 0.1\n",
    "\n",
    "# Plot the flux\n",
    "plt.figure(figsize=(12, 5))\n",
    "plt.plot(xo, map.flux(xo=xo, yo=yo, ro=ro))\n",
    "plt.xlabel(\"Occultor x position [stellar radii]\", fontsize=24)\n",
    "plt.ylabel(\"Flux [normalized]\", fontsize=24)\n",
    "\n",
    "# Plot the radial velocity\n",
    "plt.figure(figsize=(12, 5))\n",
    "plt.plot(xo, map.rv(xo=xo, yo=yo, ro=ro))\n",
    "plt.xlabel(\"Occultor x position [stellar radii]\", fontsize=24)\n",
    "plt.ylabel(\"Radial velocity [m/s]\", fontsize=24);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first plot is the usual transit light curve, and the second plot is the Rossiter-McLaughlin effect. Note that the units of the RV in the second plot are given by `map.velocity_units`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Accounting for surface features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The effect of a planet occulting a star on the observed radial velocity is similar to that of a spot rotating in and out of view. With `starry`, we can model both using the same formalism. Let's define a map of a higher spherical harmonic degree to model the effect of a spot on the RV measurements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map = starry.Map(ydeg=10, udeg=2, rv=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll give it the same properties as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map.inc = 60\n",
    "map.obl = 30\n",
    "map.veq = 1.0e4\n",
    "map.alpha = 0.3\n",
    "map[1] = 0.5\n",
    "map[2] = 0.25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And this time we'll also add a large spot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map.add_spot(-0.015, sigma=0.03, lat=30, lon=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the map looks like in white light:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map.show(rv=False, theta=np.linspace(0, 360, 50))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what it looks like in velocity space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map.show(rv=True, theta=np.linspace(0, 360, 50))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the light curve and the radial velocity anomaly for this map over one rotation period:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = np.linspace(-180, 180, 1000)\n",
    "\n",
    "# Plot the flux\n",
    "plt.figure(figsize=(12, 5))\n",
    "plt.plot(theta, map.flux(theta=theta))\n",
    "plt.xlabel(\"Angle of rotation [degrees]\", fontsize=24)\n",
    "plt.ylabel(\"Flux [normalized]\", fontsize=24)\n",
    "\n",
    "# Plot the radial velocity\n",
    "plt.figure(figsize=(12, 5))\n",
    "plt.plot(theta, map.rv(theta=theta))\n",
    "plt.xlabel(\"Angle of rotation [degrees]\", fontsize=24)\n",
    "plt.ylabel(\"Radial velocity [m/s]\", fontsize=24);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the signal is similar to the Rossiter-McLaughlin effect! Note that this is a *very* simple version of Doppler imaging, as we are modeling only the observed radial velocity --- not any details of the line shape changes due to the spot rotating in and out of view. This feature of `starry` is therefore not meant to enable the mapping of stellar surfaces from RV data, but it can be used to model the RV effects of spots on the surface for, say, de-trending radial velocity observations when modeling planetary signals."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The full RV model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far we've shown how to manually specify the occultor properties and position to get the Rossiter-McLaughlin signal, but we can use the formalism of `starry.System` to model a full system of Keplerian bodies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example of a misaligned hot Jupiter on an eccentric orbit, where we compute both the flux and the radial velocity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the star\n",
    "A = starry.Primary(\n",
    "    starry.Map(udeg=2, rv=True, amp=1, veq=5e4, alpha=0, obl=30),\n",
    "    r=1.0,\n",
    "    m=1.0,\n",
    "    length_unit=u.Rsun,\n",
    "    mass_unit=u.Msun,\n",
    ")\n",
    "A.map[1] = 0.5\n",
    "A.map[2] = 0.25\n",
    "\n",
    "# Define the planet\n",
    "b = starry.Secondary(\n",
    "    starry.Map(rv=True, amp=0, veq=0),\n",
    "    r=0.1,\n",
    "    porb=1.0,\n",
    "    m=0.01,\n",
    "    t0=0.0,\n",
    "    inc=80.0,\n",
    "    ecc=0.3,\n",
    "    w=60,\n",
    "    length_unit=u.Rsun,\n",
    "    mass_unit=u.Msun,\n",
    "    angle_unit=u.degree,\n",
    "    time_unit=u.day,\n",
    ")\n",
    "\n",
    "# Define the system\n",
    "sys = starry.System(A, b)\n",
    "\n",
    "# Compute the flux & RV signal\n",
    "time = np.linspace(-0.5, 0.5, 1000)\n",
    "flux = sys.flux(time)\n",
    "rv = sys.rv(time)\n",
    "\n",
    "# Plot it\n",
    "fig, ax = plt.subplots(2, figsize=(12, 8))\n",
    "ax[0].plot(time, flux)\n",
    "ax[1].plot(time, rv)\n",
    "ax[1].set_xlabel(\"time [days]\", fontsize=24)\n",
    "ax[0].set_ylabel(\"flux [normalized]\", fontsize=24)\n",
    "ax[1].set_ylabel(\"radial velocity [m / s]\", fontsize=24);"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}