From 595c2583ff233f3106a53a0f430869da3598241d Mon Sep 17 00:00:00 2001
From: Jeremy Kubica <104161096+jeremykubica@users.noreply.github.com>
Date: Fri, 11 Oct 2024 12:48:33 -0400
Subject: [PATCH] Add an alternative way to plot the result data

---
 notebooks/results_viz_tool.ipynb | 97 ++++++++++++++++++++++++++++++++
 src/kbmod/analysis/plotting.py   | 58 ++++++++++++++++++-
 2 files changed, 154 insertions(+), 1 deletion(-)
 create mode 100644 notebooks/results_viz_tool.ipynb

diff --git a/notebooks/results_viz_tool.ipynb b/notebooks/results_viz_tool.ipynb
new file mode 100644
index 00000000..3bb8a584
--- /dev/null
+++ b/notebooks/results_viz_tool.ipynb
@@ -0,0 +1,97 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# KBMOD Visualization Tool\n",
+    "  \n",
+    "This notebook is a tool for loading and visualizing results data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "from kbmod.analysis.plotting import *\n",
+    "from kbmod.results import Results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Load the results data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Clean Data:\n",
+    "res_file = \"/epyc/users/dinob/public_html/kbmod_runs/results/20190402_A0b_001.results.ecsv\"\n",
+    "\n",
+    "# Noisy data:\n",
+    "# res_file = \"/epyc/users/dinob/public_html/kbmod_runs/results/20190601_A1b_041.results.ecsv\"\n",
+    "\n",
+    "results = Results.read_table(res_file)\n",
+    "print(f\"Loaded {len(results)} results.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot each result."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "num_results = len(results)\n",
+    "\n",
+    "for idx in range(num_results):\n",
+    "    row = results[idx]\n",
+    "    plot_result_row_summary(row)"
+   ]
+  },
+  {
+   "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
+}
diff --git a/src/kbmod/analysis/plotting.py b/src/kbmod/analysis/plotting.py
index 5e4b2ca0..96d1d337 100644
--- a/src/kbmod/analysis/plotting.py
+++ b/src/kbmod/analysis/plotting.py
@@ -29,6 +29,7 @@
     "plot_multiple_images",
     "plot_time_series",
     "plot_result_row",
+    "plot_result_row_summary",
     "plot_search_trajectories",
 ]
 
@@ -527,7 +528,7 @@ def plot_result_row(row, times=None, figure=None):
     # In the top subfigure plot the coadded stamp on the left and
     # the light curve on the right.
     (ax_stamp, ax_lc) = fig_top.subplots(1, 2)
-    if row["stamp"] is not None:
+    if "stamp" in row and row["stamp"] is not None:
         plot_image(row["stamp"], ax=ax_stamp, figure=fig_top, norm=True, title="Coadded Stamp")
     else:
         ax_stamp.text(0.5, 0.5, "No Stamp")
@@ -558,6 +559,61 @@ def plot_result_row(row, times=None, figure=None):
         ax.text(0.5, 0.5, "No Individual Stamps")
 
 
+def compute_lightcurve_histogram(row, min_val=0.0, max_val=1000.0, bins=20):
+    psi = row["psi_curve"]
+    phi = row["phi_curve"]
+    valid = (phi != 0) & np.isfinite(psi) & np.isfinite(phi)
+
+    lc = psi[valid] / phi[valid]
+    lc[lc < min_val] = min_val
+    lc[lc > max_val] = max_val
+
+    return np.histogram(lc, bins=bins)
+
+
+def plot_result_row_summary(row, times=None, figure=None):
+    """Plot a single row of the results table.
+
+    Parameters
+    ----------
+    row : `astropy.table.row.Row`
+        The information from the results to plot.
+    times : a `list` or `numpy.ndarray` of floats
+        The array of the time stamps. If ``None`` then uses equally
+        spaced points. `None` by default.
+    figure : `matplotlib.pyplot.Figure` or `None`
+        Figure, `None` by default.
+    """
+    if figure is None:
+        figure = plt.figure(layout="constrained")
+
+    figure = plt.figure()
+    (fig_top, fig_bot) = figure.subfigures(2, 1)
+
+    # Plot the light curves on the top
+    ax_curves = fig_top.subplots(1, 2)
+    if "psi_curve" in row.colnames and "psi_curve" in row.colnames:
+        psi = row["psi_curve"]
+        phi = row["phi_curve"]
+
+        valid = (phi != 0) & np.isfinite(psi) & np.isfinite(phi)
+        if "obs_valid" in row.colnames:
+            valid = valid & row["obs_valid"]
+
+        lc = np.full(psi.shape, 0.0)
+        lc[valid] = psi[valid] / phi[valid]
+        plot_time_series(lc, times, indices=valid, ax=ax_curves[0], figure=fig_top, title=f"Psi/Phi")
+
+        counts, bins = compute_lightcurve_histogram(row)
+        ax_curves[1].stairs(counts, bins)
+
+    # Plot the stamps along the bottom.
+    ax_stamps = fig_bot.subplots(1, 4)
+    for col, name in enumerate(["coadd_sum", "coadd_mean", "coadd_median", "coadd_weighted"]):
+        if name in row.colnames and row[name] is not None:
+            plot_image(row[name], ax=ax_stamps[col], figure=fig_top, norm=True, title=name, show_counts=False)
+
+
 def plot_search_trajectories(gen, figure=None):
     """Plot the search trajectorys as created by a TrajectoryGenerator.