{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Orbit visualization\n",
    "\n",
    "In this notebook, we'll visualize a simple star-planet system and its light curve."
   ]
  },
  {
   "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": {
    "tags": [
     "hide_input"
    ]
   },
   "outputs": [],
   "source": [
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we import stuff. Note that we disable lazy evaluation in `starry`: all map attributes and method return values will be actual numpy floats and arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import starry\n",
    "\n",
    "starry.config.lazy = False\n",
    "starry.config.quiet = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's instantiate a star. We'll give it a bit of quadratic limb darkening."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = starry.Primary(starry.Map(udeg=2, amp=1.0), r=1.0, m=1.0)\n",
    "A.map[1:] = [0.5, 0.25]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we instantiate a planet and give it the 10th degree spherical harmonic expansion of the map of the Earth. The values below aren't at all realistic, but they'll make for a cool visualization below. By default, mass and radius units are solar, angles are in degrees, and times are in days. Finally, we set the map inclination and the orbital inclination to the same value, and likewise the map obliquity and longitude of ascending node. This causes the axis of rotation to be perpendicular to the orbital plane (i.e., the orbital and rotational angular momentum vectors are parallel to each other, as we'd expect for a tidally-locked planet)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = starry.Secondary(\n",
    "    starry.Map(ydeg=10, inc=80.0, obl=30.0, amp=0.1),\n",
    "    r=0.5,\n",
    "    m=0.5,\n",
    "    porb=1.0,\n",
    "    prot=1.0,\n",
    "    t0=0.0,\n",
    "    inc=80.0,\n",
    "    Omega=30.0,\n",
    ")\n",
    "b.map.load(\"earth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we instantiate a Keplerian system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys = starry.System(A, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the positions of the two bodies over the course of one orbit of the planet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "npts = 200\n",
    "t = np.linspace(-0.5, 0.5, npts)\n",
    "x, y, z = sys.position(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Render the maps over the same time period:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": false
   },
   "outputs": [],
   "source": [
    "res = 300\n",
    "theta_sec = [360.0 / sec.prot * (t - sec.t0) - sec.theta0 for sec in sys.secondaries]\n",
    "img = np.array(\n",
    "    [np.tile(sys.primary.map.render(res=300), (npts, 1, 1))]\n",
    "    + [\n",
    "        sec.map.render(theta=theta_sec[i], res=res)\n",
    "        for i, sec in enumerate(sys.secondaries)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the full system light curve:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flux = sys.flux(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize everything. We can normally visualize the orbit by calling\n",
    "\n",
    "```python\n",
    "sys.show(t)\n",
    "```\n",
    "\n",
    "but here's a more souped-up version just for fun:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the plot\n",
    "fig, ax = plt.subplots(1, figsize=(6.5, 7))\n",
    "ax_xz = fig.add_axes([0.275, 0.8, 0.2, 0.2])\n",
    "ax_xz.annotate(\n",
    "    \"Top\", fontsize=12, xy=(0, 0), xycoords=\"axes fraction\", ha=\"left\", va=\"bottom\"\n",
    ")\n",
    "ax_zy = fig.add_axes([0.525, 0.8, 0.2, 0.2])\n",
    "ax_zy.annotate(\n",
    "    \"Side\", fontsize=12, xy=(0, 0), xycoords=\"axes fraction\", ha=\"left\", va=\"bottom\"\n",
    ")\n",
    "ax_lc = fig.add_axes([0.125, 0.05, 0.775, 0.2])\n",
    "\n",
    "xz = [None] + [None for sec in sys.secondaries]\n",
    "xy = [None] + [None for sec in sys.secondaries]\n",
    "zy = [None] + [None for sec in sys.secondaries]\n",
    "circ = [None] + [None for sec in sys.secondaries]\n",
    "maps = [sys.primary.map] + [sec.map for sec in sys.secondaries]\n",
    "radii = np.array([sys.primary.r] + [sec.r for sec in sys.secondaries])\n",
    "\n",
    "for axis, arrs in zip([ax, ax_xz, ax_zy], [(x, y), (x, z), (z, y)]):\n",
    "    axis.axis(\"off\")\n",
    "    R = 1.2 * max(-np.min(arrs), np.max(arrs))\n",
    "    axis.set_xlim(-R, R)\n",
    "    axis.set_ylim(-R, R)\n",
    "\n",
    "# Plot the light curve\n",
    "ax_lc.plot(t, flux, \"k-\")\n",
    "(lc,) = ax_lc.plot(t[0], flux[0], \"o\", color=\"k\")\n",
    "ax_lc.axis(\"off\")\n",
    "\n",
    "# Plot the first frame\n",
    "for i, xi, yi, zi, map, r in zip(range(1 + len(sys.secondaries)), x, y, z, maps, radii):\n",
    "\n",
    "    # Orbit outlines\n",
    "    ax_xz.plot(xi, zi)\n",
    "    ax_zy.plot(zi, yi)\n",
    "\n",
    "    # Body positions\n",
    "    xz[i] = ax_xz.scatter(xi[0], zi[0])\n",
    "    zy[i] = ax_zy.scatter(zi[0], yi[0])\n",
    "\n",
    "    # Maps\n",
    "    extent = np.array([xi[0], xi[0], yi[0], yi[0]]) + np.array([-1, 1, -1, 1]) * r\n",
    "    xy[i] = ax.imshow(\n",
    "        img[i, 0],\n",
    "        origin=\"lower\",\n",
    "        cmap=\"plasma\",\n",
    "        extent=extent,\n",
    "        clip_on=False,\n",
    "        zorder=zi[0],\n",
    "    )\n",
    "    circ[i] = plt.Circle(\n",
    "        (xi[0], yi[0]), r, color=\"k\", fill=False, zorder=zi[0] + 1e-3, lw=3\n",
    "    )\n",
    "    ax.add_artist(circ[i])\n",
    "\n",
    "# Animation\n",
    "def updatefig(k):\n",
    "    for i, xi, yi, zi, map, r in zip(\n",
    "        range(1 + len(sys.secondaries)), x, y, z, maps, radii\n",
    "    ):\n",
    "        xz[i].set_offsets((xi[k], zi[k]))\n",
    "        zy[i].set_offsets((zi[k], yi[k]))\n",
    "        xy[i].set_extent(\n",
    "            np.array([xi[k], xi[k], yi[k], yi[k]]) + np.array([-1, 1, -1, 1]) * r\n",
    "        )\n",
    "        xy[i].set_zorder(zi[k])\n",
    "        xy[i].set_data(img[i, k])\n",
    "        circ[i].center = (xi[k], yi[k])\n",
    "        circ[i].set_zorder(zi[k] + 1e-3)\n",
    "    lc.set_xdata(t[k])\n",
    "    lc.set_ydata(flux[k])\n",
    "    return xz + xy + zy + circ + [lc]\n",
    "\n",
    "\n",
    "ani = FuncAnimation(fig, updatefig, interval=30, blit=False, frames=len(t))\n",
    "plt.close()\n",
    "display(HTML(ani.to_html5_video()))"
   ]
  }
 ],
 "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
}