{
 "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 matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "from kbmod.analysis.plotting import *\n",
    "from kbmod.core.image_stack_py import ImageStackPy\n",
    "from kbmod.core.stamp_utils import (\n",
    "    coadd_mean,\n",
    "    coadd_median,\n",
    "    coadd_sum,\n",
    "    create_stamps_from_image_stack,\n",
    ")\n",
    "from kbmod.util_functions import load_deccam_layered_image\n",
    "from kbmod.search import Trajectory\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)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = np.array([[1.0]])\n",
    "im = load_deccam_layered_image(im_path + \"000000.fits\", p)\n",
    "print(f\"Loaded a {im.width} by {im.height} image at time {im.time}\")"
   ]
  },
  {
   "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 or a `LayeredImagePy` 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.sci, 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 `ImageStackPy` or list of images. In the case of an `ImageStackPy` 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",
    "\n",
    "sci = []\n",
    "var = []\n",
    "times = []\n",
    "psfs = []\n",
    "for f in files:\n",
    "    im = load_deccam_layered_image(f, p)\n",
    "    sci.append(im.sci)\n",
    "    var.append(im.var)\n",
    "    times.append(im.time)\n",
    "    psfs.append(im.psf)\n",
    "\n",
    "# Load the images.\n",
    "stack = ImageStackPy(\n",
    "    times=times,\n",
    "    sci=sci,\n",
    "    var=var,\n",
    "    mask=None,\n",
    "    psfs=psfs,\n",
    ")\n",
    "\n",
    "num_images = stack.num_times\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 = create_stamps_from_image_stack(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 numpy arrays."
   ]
  },
  {
   "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 coadd from the stack of stamps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx-thumbnail": {
     "tooltip": "Sum, mean and median stacked postage stamp."
    }
   },
   "outputs": [],
   "source": [
    "stamps = np.asarray(stamps)\n",
    "\n",
    "plot_multiple_images(\n",
    "    [\n",
    "        coadd_sum(stamps),\n",
    "        coadd_mean(stamps),\n",
    "        coadd_median(stamps),\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
}