{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import wobble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "mpl.rcParams['xtick.labelsize'] = 12\n",
    "mpl.rcParams['ytick.labelsize'] = 12\n",
    "mpl.rcParams['axes.labelsize'] = 14\n",
    "mpl.rcParams['figure.figsize'] = (12.0, 4.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### First, you'll need some data to load up.\n",
    "\n",
    "You can download example HARPS data files (and results files) to play around with linked in the [documentation](https://wobble.readthedocs.io/en/latest/quickstart.html#input-data).\n",
    "\n",
    "Here we'll assume that you have the data `51peg_e2ds.hdf5` saved in the `wobble/data` directory.\n",
    "\n",
    "By default, loading the data will load all echelle orders and all epochs in the data file; you can change this with the optional `orders` and `epochs` kwargs, which each take lists (or 1-d numpy arrays) of indices for the desired orders/epochs to load."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = wobble.Data('../data/51peg_e2ds.hdf5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data we just loaded are assumed to be continuum normalized, with regions of bad data (negative flux values or very low SNR) \"masked out\" by setting their uncertainties to be infinite. In this example, the data are also in units of log(wavelength) vs. log(flux)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 0 # index of echelle order to plot\n",
    "n = 0 # index of epoch to plot\n",
    "plt.plot(data.xs[r][n], data.ys[r][n], 'k.', ms=6)\n",
    "mask = data.ivars[r][n] <= 1.e-8 # masked-out bad data\n",
    "plt.plot(data.xs[r][n][mask], data.ys[r][n][mask], 'w.', ms=4)\n",
    "plt.ylabel('ln(flux)')\n",
    "plt.xlabel('ln(wave)');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Now let's create a results object in which to store the outputs of wobble:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = wobble.Results(data=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This object is not currently populated with useful information (because we haven't optimized anything yet!), but once it is we'll be able to save it with the `results.write('filename.hdf5')` function.\n",
    "\n",
    "A saved results file can be loaded as:\n",
    "```results = wobble.Results(filename='filename.hdf5')```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Here's a minimal example of optimizing a model consisting of a star and tellurics for a single order:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = 67 # index into data.orders for the desired order\n",
    "model = wobble.Model(data, results, r)\n",
    "model.add_star('star')\n",
    "model.add_telluric('tellurics')\n",
    "wobble.optimize_order(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results have been automatically saved and we can now view them by generating a plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 40 # epoch to plot\n",
    "results.plot_spectrum(r, n, data, 'demo1.png')\n",
    "from IPython.display import Image\n",
    "Image(filename='demo1.png') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The residuals look good for the star but not great around the tellurics. Let's try running with variable tellurics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results2 = wobble.Results(data=data)\n",
    "model = wobble.Model(data, results2, r)\n",
    "model.add_star('star')\n",
    "model.add_telluric('tellurics', variable_bases=2)\n",
    "wobble.optimize_order(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results2.plot_spectrum(r, n, data, 'demo2.png')\n",
    "Image(filename='demo2.png') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks better! Here are the RVs for this single order. Once we run on all orders, we can combine the order-by-order velocities using `results.combine_orders('star')` and access the final (non-barycentric-corrected RVs as `results.star_time_rvs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.errorbar(results2.dates, results2.star_rvs[r] + results2.bervs, \n",
    "             1./np.sqrt(results2.star_ivars_rvs[r]), \n",
    "             fmt='o', ms=5, elinewidth=1)\n",
    "plt.xlabel('JD')\n",
    "plt.ylabel(r'RV (m s$^{-1}$)')\n",
    "plt.xlim([2456505, 2456570]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now let's generalize this to multiple orders and get RVs for the full* spectrum:\n",
    "\n",
    "##### * not actually the complete spectrum in this tutorial because that would take a long time to run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we'll overwrite the Data object with one that contains only a subset of spectral orders. The following commands could (and should) be used on the entire object containing all orders, but that can take a long time (up to an hour) so for the sake of the tutorial we'll use a smaller subsample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = wobble.Data('../data/51peg_e2ds.hdf5', orders=np.arange(65,70))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following loop, we'll continually be overwriting the \"model\" variable. That's ok! All optimized results will be copied over to the \"results\" object automatically, and as long as the star component is given the same name in the model for every order, they'll be associated in the results object as we'd expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = wobble.Results(data=data)\n",
    "for r in range(len(data.orders)):\n",
    "        print('starting order {0} of {1}'.format(r+1, len(data.orders)))\n",
    "        model = wobble.Model(data, results, r)\n",
    "        model.add_star('star')\n",
    "        model.add_telluric('tellurics', variable_bases=2)\n",
    "        wobble.optimize_order(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have RVs, let's do some post-processing on them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.combine_orders('star')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.apply_drifts('star') # instrumental drift corrections\n",
    "results.apply_bervs('star') # barycentric corrections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can look at the resulting RVs and save them to a text file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.errorbar(data.dates, results.star_time_rvs - np.mean(results.star_time_rvs), \n",
    "                     results.star_time_sigmas,\n",
    "                    fmt='o', ms=5, elinewidth=1)\n",
    "plt.xlabel('JD')\n",
    "plt.ylabel(r'RV (m s$^{-1}$)');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.write_rvs('star', 'demo_rvs.txt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above command saved the RVs; we probably also want to save the spectral fits and other diagnostic information for future reference. We can do that with the following command, which preserves the entire Results object (including RVs and spectra):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.write('demo_results.hdf5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Further exercises:\n",
    "- [optimize the regularization amplitudes](https://github.com/megbedell/wobble/blob/master/scripts/regularization.py) and rerun\n",
    "- check [the API](https://wobble.readthedocs.io/en/latest/api.html) and figure out what keywords to use if you want to:\n",
    "    - calculate uncertainties on the stellar template (warning: this will run slowly)\n",
    "    - give your own custom starting RV guess to a component"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "wobble-install",
   "language": "python",
   "name": "wobble-install"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}