{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KBMOD Visualization and Analysis\n",
    "  \n",
    "This notebook demonstrates the basic functionality for visualizing both the input data and the results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup kbmod visualization demo\n",
    "\n",
    "Before importing, make sure you have installed kbmod using `pip install .` in the root directory.  Also be sure you are running with python3 and using the correct notebook kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "from kbmod.analysis.plotting import *\n",
    "from kbmod.util_functions import load_deccam_layered_image\n",
    "from kbmod.search import (\n",
    "    ImageStack,\n",
    "    PSF,\n",
    "    Trajectory,\n",
    "    get_stamps,\n",
    "    get_median_stamp,\n",
    "    get_mean_stamp,\n",
    "    get_summed_stamp,\n",
    ")\n",
    "from kbmod.results import Results\n",
    "\n",
    "# Data paths\n",
    "im_path = \"../data/small/\"\n",
    "res_path = \"../data/fake_results/results.ecsv\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the Input\n",
    "\n",
    "We can visual input files with the `kbmod.plotting.plot_img()` function which will display the impage and optionally normalize the image (with the `norm` parameter) and add a title (with the `title` parameter).\n",
    "\n",
    "We can load a `LayeredImage` given the path and filename to the FITS file as well as the `PSF` for the image. We use a default psf here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = PSF(1.0)\n",
    "im = load_deccam_layered_image(im_path + \"000000.fits\", p)\n",
    "print(f\"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_obstime()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the image the plotting library's `plot_img()` function which is a wrapper around matplotlib's `imshow()` function. The `plot_img()` function can take images in a variety of formats including a numpy array of the pixel values, a `RawImage` object, or a `LayeredImage` object.\n",
    "\n",
    "*Note*: The data/demo images contain a single bright object, so the majority of the image should be empty with a single bright spot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_image(im.get_science().image, norm=True, title=\"Image 000000.fits\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also load a collection of images and plot them with `plot_multiple_images()`. This function handles the arrangement of subplots given a single `columns` parameter that defines the total number of columns to use.\n",
    "\n",
    "The function can take an `ImageStack` or list of images. In the case of an `ImageStack` it automatically labels each image with its time stamp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = [im_path + f for f in os.listdir(im_path) if \".fits\" in f]\n",
    "files.sort()\n",
    "imgs = [load_deccam_layered_image(f, p) for f in files]\n",
    "\n",
    "# Load the images.\n",
    "stack = ImageStack(imgs)\n",
    "\n",
    "num_images = stack.img_count()\n",
    "print(f\"Loaded {num_images} images.\")\n",
    "\n",
    "plot_multiple_images(stack, columns=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create and Visualize Stamps\n",
    "\n",
    "Stamps are a critical tool for analyzing and debugging proposed detections. They can be created automatically using the stamp creation utilities. It requires a few pieces of data:\n",
    "* search_stack - provides the machinery for making predictions on the image (needed to handle the various corrections).\n",
    "* trajectory - Contains the information about where to place the stamps (the underlying trajectory).\n",
    "* stamp_radius - The radius in pixels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the trajectory with a given parameters and then the trajectory result.\n",
    "trj = Trajectory()\n",
    "trj.x = 11\n",
    "trj.y = 27\n",
    "trj.vx = 16.0\n",
    "trj.vy = 3.3\n",
    "\n",
    "# Create the stamps around this trajectory.\n",
    "stamps = get_stamps(stack, trj, 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can display the stamps around each predicted object position using the plotting library's `plot_multiple_images()` function. Note that this time the function is taking a list of `RawImage` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_multiple_images(stamps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also compute the sum, mean, or median stamp. The stamp functions take a list of bools corresponding to whether each time is valid. These can be extracted from the result data. You can use an empty array (as we do below) to build the stamp out of all valid times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx-thumbnail": {
     "tooltip": "Sum, mean and median stacked postage stamp."
    }
   },
   "outputs": [],
   "source": [
    "# The coadds requires a vector of which indices to use.\n",
    "inds = [True] * stack.img_count()\n",
    "\n",
    "plot_multiple_images(\n",
    "    [\n",
    "        get_summed_stamp(stack, trj, 10, inds),\n",
    "        get_mean_stamp(stack, trj, 10, inds),\n",
    "        get_median_stamp(stack, trj, 10, inds),\n",
    "    ],\n",
    "    labels=[\"Summed\", \"Mean\", \"Median\"],\n",
    "    norm=True,\n",
    "    columns=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting Results\n",
    "\n",
    "We can plot the results stored in the `Results` object from a KBMOD run. In this example, you need results in the form of either a `Results` object or saved files. For the purposes of this notebook you can use that fake result data (corresponding to the `data/demo` images) included in `data/fake_results/results.ecsv`. These results were generated by the KBMOD_Demo notebook.\n",
    "\n",
    "We load all of the result data as a `Results` object using the class's `read_table` function. We extract the first row to use for later examples. \n",
    "\n",
    "Note that in previous versions (<=1.1) KBMOD saved output was also saved in a series of individual files such as `results_DEMO.txt` for the a file of the Trajectory objects. You can still load these using the `from_trajectory_file` function:\n",
    "\n",
    "```\n",
    "results = Results.from_trajectory_file(\"../data/fake_results/results_DEMO.txt\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = Results.read_table(\"../data/fake_results/results.ecsv\")\n",
    "print(f\"Loaded {len(results)} results.\")\n",
    "\n",
    "row0 = results[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We plot the results using the same helper functions as above and accessing the \"stamp\" column from the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_image(row0[\"stamp\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or the \"all_stamps\" column:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_multiple_images(row0[\"all_stamps\"], columns=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the time series such as: `psi_curve`, `phi_curve`, or `light_curve` using the `plot_time_series` function. If we provide timestamps (such as with the top and bottom figure), they are used to both label the x-axis and set spacing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_times = [\n",
    "    57130.19921875,\n",
    "    57130.2109375,\n",
    "    57130.21875,\n",
    "    57131.19921875,\n",
    "    57131.2109375,\n",
    "    57131.21875,\n",
    "    57132.19921875,\n",
    "    57132.2109375,\n",
    "    57132.21875,\n",
    "    57133.19921875,\n",
    "]\n",
    "\n",
    "fig1 = plt.figure(figsize=(10, 15))\n",
    "(ax1, ax2, ax3) = fig1.subplots(3, 1)\n",
    "\n",
    "plot_time_series(row0[\"psi_curve\"], times=all_times, ax=ax1, figure=fig1, title=\"Psi curve with time spacing\")\n",
    "\n",
    "plot_time_series(row0[\"psi_curve\"], ax=ax2, figure=fig1, title=\"Psi curve with equal spacing\")\n",
    "\n",
    "plot_time_series(\n",
    "    row0[\"psi_curve\"],\n",
    "    times=all_times,\n",
    "    indices=[True, True, True, False, True, True, False, True, True, False],\n",
    "    ax=ax3,\n",
    "    figure=fig1,\n",
    "    title=\"Psi curve with time spacing and invalid indices\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "The plotting library also contains a helper function to display all this information from a single row of the `Results` data structure. The `plot_result_row()` function will read data out of the \"stamp\", \"psi_curve\", \"phi_curve\", \"obs_valid\", and \"all_stamps\" columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_result_row(row0)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "If a column does not exist, the function will simply indicate this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.table.remove_column(\"all_stamps\")\n",
    "plot_result_row(results[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Jeremy's KBMOD",
   "language": "python",
   "name": "kbmod_jk"
  },
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}