{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AAS234 WWT Workshop Notebook (June 2019; St. Loui, MO)\n",
    "\n",
    "*Please note: these notebooks are preserved for posterity but not kept up-to-date, so you might run into issues with older notebooks.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data layers in pywwt.\n",
    "\n",
    "In addition to the two and three-dimensional layer data included out of the box, `pywwt` also allows users to upload their own `astropy` `Table` objects and FITS image data for quick and seamless visualizations.\n",
    "\n",
    "In this notebook, we'll use `astroquery` to download data in both formats and show how to plot them in multiple dimensions and against several all-sky, large-scale astronomical surveys.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and initialize the viewer.\n",
    "\n",
    "*(Select a cell and press `ctrl` + `Enter` to run it.)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from astropy import units as u\n",
    "from astropy.coordinates import SkyCoord\n",
    "from astropy.io import fits\n",
    "from astropy.table import Table, Column\n",
    "from astropy.time import Time, TimeDelta # note capitalization\n",
    "from astropy.wcs import WCS\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from astroquery.nasa_exoplanet_archive import NasaExoplanetArchive\n",
    "from astroquery.skyview import SkyView\n",
    "\n",
    "from pywwt.jupyter import WWTJupyterWidget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wwt = WWTJupyterWidget()\n",
    "wwt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*If you're using JupyterLab* and not just a plain Jupyter notebook, you can move the WWT view to a separate window pane. This is **extremely** useful since it lets you keep on typing code without scrolling WWT out of view. Here's how you do that:\n",
    "\n",
    "![Right click and select \"Create New View for Output\"](../data/separate-pane-instructions.jpg)\n",
    "\n",
    "If you don't get a menu when you right-click, or the menu doesn't look like the one pictured, you are using a plain Jupyter notebook and will have to scroll back and forth."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot K2’s footprint and exoplanet yield.\n",
    "\n",
    "_Here, we use `pywwt` to plot data from an `astropy` `Table` in both two and three dimensions._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**To begin, we use a class from the `astroquery` package to query NASA's Exoplanet Archive and return an `astropy` `Table` that contains information about every confirmed exoplanet so far. It's a large table, so we'll only keep the columns necessary to run this example.**\n",
    "\n",
    "*(If the `astroquery` call doesn't load, uncomment and run the cell following it to download the table from a `pywwt`-affiliated GitHub repository.)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#xarch = NasaExoplanetArchive.get_confirmed_planets_table()\n",
    "#xarch.keep_columns(['pl_hostname', 'pl_letter', 'st_dist', 'ra', 'dec'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xarch = Table.read('https://raw.githubusercontent.com/WorldWideTelescope/'\n",
    "                   'pywwt-notebooks/master/aas-tutorials/data/k2_table.ecsv',\n",
    "                   format='ascii.ecsv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*(Open a new cell above this one and see what the table looks like by entering `xarch` and running it.)*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Next, we slice the table twice so it only contains entries of 'b' planets observed by K2.** \n",
    "\n",
    "K2 planets and stars are named using a consistent system: \n",
    "\n",
    "```K2 + 'system number' + 'letter'```. \n",
    "\n",
    "The letter for a star is typically 'a', the first exoplanet confirmed in the system is 'b', and so on.\n",
    "\n",
    "We are plotting stars, but since the table only contains exoplanets (letters from 'b' onward), we will slice the table so only planets with the `b` suffix remain. That gives us one entry per system.\n",
    "\n",
    "**`k2_pl` searches the `xarch` column that lists star names and only keeps those with K2 at the beginning. Then, `k2_st` searches the suffix of the planet name and only keeps `b` planets.**\n",
    "\n",
    "*(If you aren't familiar with list comprehensions, see [Note 1](aas234_tutorial.ipynb#Note-1) at the end of the notebook.)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select only k2 planets\n",
    "k2_pl = xarch[ [row.startswith('K2') for row in xarch['pl_hostname']] ]\n",
    "\n",
    "# select only 'b' planets from k2\n",
    "k2_st = k2_pl[ ['b' in row for row in k2_pl['pl_letter']] ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**After that, we load the footprint of the K2 campaign fields on the sky using pywwt's `add_fov` method.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "field = wwt.add_fov(wwt.instruments.k2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Allow a few seconds for the footprint to load, then watch as `pywwt`draws it on the sky."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Once the footprint has loaded, we finally create the data layer that will plot our points in the viewer.**\n",
    "\n",
    "`pywwt` reads the locations from the keyword arguments `lon_att` and `lat_att`, and we feed it the names of the columns in `k2_st` that contain the stellar right ascensions and declinations. We can also set other attributes like color and size in the function call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lay = wwt.layers.add_table_layer(table=k2_st, frame='Sky',\n",
    "                                lon_att='ra', lat_att='dec',\n",
    "                                color='#c4d600', size_scale=40)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do you see the points? If not, try making them larger."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lay.size_scale *= 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to see the points and footprint without a sky background, `pywwt` also has a black layer available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wwt.foreground_opacity = 0\n",
    "wwt.background = wwt.imagery.other.black"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**It is also possible to visualize data layers in `pywwt`'s 3D solar system mode. To do so, we'll switch the view.**\n",
    "\n",
    "We will also remove the K2 footprint, as `add_fov` is designed to work primarily in 2D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "field.remove()\n",
    "wwt.set_view('solar system')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Zoom out, and you'll see K2's exoplanet-hosting stars in 3D. All of them are currently the same distance from Earth -- let's fix that.**\n",
    "\n",
    "Data layers also have attributes related to distance and altitude. We set `alt_type` to the correct type for our scenario, and let `pywwt` know which column of the table contains distance with `alt_att`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lay.alt_type = 'distance'\n",
    "lay.alt_att = 'st_dist'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Zoom out again and enjoy the view of K2's stars scattered across the Milky Way at their proper distances from Earth._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer FITS images over existing all-sky surveys.\n",
    "\n",
    "_Here, we use `pywwt` to view a FITS image against the backdrop of multiple all-sky surveys._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll begin by resetting the viewer to its originial state. Then, since we're working with an infrared image, we'll switch the background to one of the same wavelength and make the foreground layer invisible for simplicity's sake."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wwt.reset()\n",
    "wwt.background = wwt.imagery.ir.twomass\n",
    "wwt.foreground_opacity = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**We again use `astroquery` to get our data, this time harnessing the `SkyView` service to retrieve a FITS cutout of a portion of M101 that experienced a Type Ia supernova in 2011.**\n",
    "\n",
    "The M101 cutout comes from a pre-event image from the highest wavelength band of 2MASS, an infrared all-sky survey. `SkyView` can also access of ther bands of 2MASS and many other surveys, as shown at the end of this document in [Note 2](aas234_tutorial.ipynb#Note-2).\n",
    "\n",
    "*(Again, if the `astroquery` call doesn't load, uncomment and run the cell following it to download the table from a `pywwt`-affiliated GitHub repository.)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose the size of the image in pixels\n",
    "size = 500 # remember that larger sizes mean longer waits!\n",
    "\n",
    "img_list = SkyView.get_images(position='SN 2011FE',\n",
    "                              survey='2MASS-K', pixels=size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#img_list = [fits.open('https://github.com/WorldWideTelescope/'\n",
    "#                      'pywwt-notebooks/blob/live-image-beta/'\n",
    "#                      'tutorials/data/sn_2011fe_500.fits?raw=true')]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because `SkyView.get_images()` returns a list, we have to index `img_list` once the download has finished in order to get to the FITS image data itself. (In this example, `img_results` is a list of just one.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn11 = img_list[0]\n",
    "sn11.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**In the most basic cases, this is all that's needed to create an image layer.**\n",
    "\n",
    "If you have the `PrimaryHDU`, or \"Primary Header Data Unit\", `pywwt` only needs one command to plot it. Once uploaded, the viewer will take you to the proper location and scale to see your image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = wwt.layers.add_image_layer(sn11)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_(It's also possible to upload FITS images from a local file by using `image='the/path/to/your/file'` as the method's argument instead of an existing `Python` object name. `pywwt` will then automatically do the examination of the header and image data that we did in previous cells._)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also separate `PrimaryHDU` into its header and image data components and make edits to either before creating a layer. We'll explore that by splitting them into different variables and deleting our first image layer from the viewer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn11_data = sn11[0].data.copy()\n",
    "sn11_header = sn11[0].header.copy()\n",
    "img.remove()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pywwt` requires the image data to be structured as a 2-D array, which is what we see below.  The dimensions of our array are (`size` x `size`), just as we specified earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sn11_data.shape)\n",
    "sn11_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can perform operations on this array like scrubbing a column of pixels to, say, censor out an artifact that we don't want to display."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn11_data[0:size, 50:100] = np.median(sn11_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The header contains base-level data about the image such as its coordinates and the projection system it follows. We will need it to create an `astropy.wcs.WCS` object, which `pywwt` uses to convert the image to its preferred coordinate system and orient the data properly on its 2-D sky projection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn11_header"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Speaking of which, let's create the WCS (World Coordinate System) object now. As stated before, all we need is the header data; `astropy`'s internals do the rest of the heavy lifting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WCS(sn11_header)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then bring these separate header and data objects together in a tuple to run the `add_image_layer` method again. You should notice a vertical band that wasn't there before -- we created that by manipulating the image data array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = wwt.layers.add_image_layer(image=(sn11_data, WCS(sn11_header)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Once the image has loaded, we are able to edit a few of its attributes: `opacity`, `vmin`/`vmax`** (the minimum and maximum flux values covered by the colormap), **and `stretch`** (the scaling function used to display the data). There are more options coming as we continue development."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Changing opacity shows us that our image is indeed aligned with the background survey."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.opacity = .4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.opacity = .6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Shifting `vmin` and `vmax` changes the lower and upper bounds of our grayscale colormap.\n",
    "\n",
    "`pywwt` automatically sets these attributes such that ~99% of points in the image data array are covered, but we can adjust them depending on what exactly we'd like to emphasize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.percentile(sn11_data, .5), np.percentile(sn11_data, 99.5))\n",
    "print(img.vmin, img.vmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.vmin = 360"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.vmax = 370"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Valid `stretch` values are: `linear`, `log`, `power`, `sqrt`, and `histeq`. `pywwt` chooses a linear stretch by default, but again, we can adjust it to fit our individual needs -- aesthetic, scientific, or anywhere in between."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img.stretch = 'log'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The great advantage of being able to upload personal FITS images into `pywwt` is that we can easily compare them with multiple surveys at once.**\n",
    "\n",
    "Since our FITS image is in infrared, we can use a background and a foreground all-sky survey from `pywwt` to see M101 in three wavelengths at once. (Please give them a few seconds to load.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wwt.background = wwt.imagery.visible.sdss\n",
    "wwt.foreground = wwt.imagery.gamma.fermi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get the right balance between the all-sky layers, we can manually change the opacity of the foreground..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wwt.foreground_opacity = .7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "...or we can use a Jupyter-exclusive option to do the same -- and more -- using live sliders and menus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wwt.layer_controls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**We can use `pywwt`'s ability to add annotations to the viewer and `SkyCoord`'s location functionality to visually indicate the spot where our supernova took place.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sn11_coord = SkyCoord.from_name('Sn 2011FE')\n",
    "circle = wwt.add_circle(sn11_coord, radius=.003*u.deg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in the K2 example, we're able to change attributes of the annotation after its creation and remove it when we no longer need it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "circle.line_width = 6 * u.pixel\n",
    "circle.line_color = '#e56020'\n",
    "circle.radius = wwt.get_fov() / 20\n",
    "circle.set_center(wwt.get_center())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "circle.remove()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Finally, remember that `pywwt` is four-dimensional visualization engine -- this means we can also visualize the progression of time.**\n",
    "\n",
    "Using `astropy` `Time` objects to change the radius of our circle over time, we can make a basic animation of how SN 2011fe looked in the months immediately surrounding its detonation and view it in a matter of seconds.\n",
    "\n",
    "*(This animation is qualitative and is only meant to approximate how the supernova looked.)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "circ = wwt.add_circle(sn11_coord, fill=True, radius=.003*u.deg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save benchmark times for animation\n",
    "start = Time('2011-08-01', format='iso')\n",
    "peak = Time('2011-09-13', format='iso')\n",
    "finish = Time('2011-10-26', format='iso')\n",
    "span = peak - start\n",
    "\n",
    "now = start\n",
    "step = (finish - start) / 20  - TimeDelta(.01 * u.second)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop through a few months in a few seconds of real time\n",
    "max_rad = circ.radius\n",
    "\n",
    "while now <= finish:\n",
    "    wwt.set_current_time(now)\n",
    "    \n",
    "    circ.radius = max_rad * (1 - np.abs(now - peak) / span)\n",
    "    time.sleep(.1)\n",
    "    \n",
    "    now += step\n",
    "circ.remove()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***Is it working? If so, congratulations, enjoy the view, and stay in touch! Find us at the AAS booth, sign up for our [mailing list](https://bit.ly/wwt-signup), discuss with other users at the [WWT forum](https://wwt-forum.org/), and contribute suggestions and fixes at our [GitHub repository](https://github.com/WorldWideTelescope/pywwt).***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Note 1\n",
    "*The constructions used to build `k2_pl` and `k2_st` are called \"list comprehensions,\" and they are shortcuts for building lists. The regular, equivalent method to create `k2_st`, looks like this:*\n",
    "\n",
    "```\n",
    "index_list = []\n",
    "for i, row in enumerate(xarch, 0):\n",
    "    if row['pl_hostname'].startswith('K2'):\n",
    "        index_list.append(i)\n",
    "\n",
    "k2_pl = xarch[index_list]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Note 2\n",
    "\n",
    "You can see all of the layers available layers in `SkyView` by running `SkyView.list_surveys(`) or `SkyView.survey_dict()`, but be aware that these will take some time to load.\n",
    "\n",
    "To get links to image files from all bands of a particular survey, run:\n",
    "\n",
    "`SkyView.get_image_list('Sn 2011FE', survey=SkyView.survey_dict['IR:2MASS class='])`\n",
    "\n",
    "You can substitute the dictionary entry for whichever surveys are present in `SkyView.survey_dict()`. For example, if you wanted images from all bands of SDSS, the `survey` keyword argument in the call above would change to `SkyView.survey_dict['Optical:SDSS']`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Credits\n",
    "\n",
    "This notebook was prepared by O. Justin Otor."
   ]
  }
 ],
 "metadata": {
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}