{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KBMOD Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from kbmod.configuration import SearchConfiguration\n",
    "from kbmod.fake_data.demo_helper import make_demo_data\n",
    "from kbmod.run_search import *\n",
    "from kbmod.search import *\n",
    "from kbmod.work_unit import WorkUnit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup file paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to run KBMOD you need to have the location of the input data and a `res_filepath` that provides a path to the directory where the output results will be stored. Input data can come from a variety of formats including Rubin’s Bulter, fits files, and `WorkUnit` files. In this demo we use the `WorkUnit` file which is an internal storage format used. For more information on generating a `WorkUnit` from the Butler or fits, see the standardizer notebooks.\n",
    "\n",
    "If you already have data files, you can use those. Below we create and use data in `data/demo_data.fits`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_filename = \"../data/demo_data.fits\"\n",
    "\n",
    "# Create the fake data usering a helper function.\n",
    "if not Path(input_filename).is_file():\n",
    "    make_demo_data(filename=input_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_filepath = \"./fake_results\"\n",
    "if not Path(res_filepath).is_dir():\n",
    "    print(f\"Directory {res_filepath} does not exist. Creating.\")\n",
    "    os.mkdir(res_filepath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run KBMOD\n",
    "\n",
    "The standard approach to running KBMOD is to perform a grid search over all starting pixels and a grid of velocities. The velocities are defined by steps in velocity space (in pixels per day) and angles. Let’s do a grid search with:\n",
    "* 21 different velocity steps from 0 pixels per day and 20 pixels per day\n",
    "* 11 different angles from 0.5 below the default angle (computed based on the ecliptic) to 0.5 above.\n",
    "\n",
    "KBMOD needs a series of configuration parameters to specify all the information about the search. In this notebook we explicitly provide the configuration parameters as a dictionary so users can see what is being specified. However most users will want to use the ``SearchConfiguration`` class. A ``SearchConfiguration`` object uses reasonable defaults when created:\n",
    "\n",
    "```\n",
    "config = SearchConfiguration()\n",
    "```\n",
    "\n",
    "Users can then override values one at a time or by passing a dictionary:\n",
    "\n",
    "```\n",
    "d = {\"result_filename\": \"Here\", \"encode_num_bytes\": 2}\n",
    "config.set_multiple(d)\n",
    "```\n",
    "\n",
    "More importantly ``SearchConfiguration`` can read from or written to a YAML file:\n",
    "\n",
    "```\n",
    "config = SearchConfiguration.from_file(file_path)\n",
    "```\n",
    "\n",
    "This allows users to define a per-task configuration and version control it.\n",
    "\n",
    "Most of the parameters you will not need to change. They are included to provide fine grained options for control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_suffix = \"DEMO\"\n",
    "\n",
    "# The demo data has an object moving at x_v=10 px/day\n",
    "# and y_v = 0 px/day. So we search velocities [0, 19].\n",
    "v_min = 0\n",
    "v_max = 20\n",
    "v_steps = 20\n",
    "v_arr = [v_min, v_max, v_steps]\n",
    "\n",
    "# and angles [-0.5, 0.5)\n",
    "ang_below = 0.5\n",
    "ang_above = 0.5\n",
    "ang_steps = 10\n",
    "ang_arr = [ang_below, ang_above, ang_steps]\n",
    "\n",
    "input_parameters = {\n",
    "    # Use search parameters (including a force ecliptic angle of 0.0)\n",
    "    # to match what we know is in the demo data.\n",
    "    \"generator_config\": {\n",
    "        \"name\": \"EclipticCenteredSearch\",\n",
    "        \"angles\": [-0.5, 0.5, 11],\n",
    "        \"velocities\": [0.0, 20.0, 21],\n",
    "        \"angle_units\": \"radian\",\n",
    "        \"given_ecliptic\": 0.0,\n",
    "    },\n",
    "    # Output parameters\n",
    "    \"result_filename\": \"./fake_results/results.ecsv\",\n",
    "    # Basic filtering (always applied)\n",
    "    \"num_obs\": 15,  # <-- Filter anything with fewer than 15 observations\n",
    "    \"lh_level\": 10.0,  # <-- Filter anything with a likelihood < 10.0\n",
    "    # SigmaG clipping parameters\n",
    "    \"sigmaG_lims\": [15, 60],  # <-- Clipping parameters (lower and upper percentile)\n",
    "    \"gpu_filter\": True,  # <-- Apply clipping and filtering on the GPU\n",
    "    \"clip_negative\": True,\n",
    "}\n",
    "config = SearchConfiguration.from_dict(input_parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To see the full list of parameters used and their values you can just print the configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load the data as a `WorkUnit` object. In general `WorkUnit`'s include a copy of their own configuration so they have all the information they need for a full run. We overwrite the stored configuration with the one we defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = WorkUnit.from_fits(input_filename)\n",
    "input_data.config = config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "KBMOD uses Python's logging library for output during the run. If you are interested in running with debug output (verbose mode), you can set the level of Python's logger to ``WARNING`` (for warning messages only), ``INFO`` (for moderate output), or ``DEBUG`` (for comprehensive output). By default logging is set to warning level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "\n",
    "# Or for a LOT of detail\n",
    "# logging.basicConfig(level=logging.DEBUG)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we have defined the search parameters, we can create a `SearchRunner` and use one of the run_search functions. In this case we use `run_search_from_work_unit` which uses the `WorkUnit` to define both the image data and the configuration information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = SearchRunner()\n",
    "results = rs.run_search_from_work_unit(input_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then check that results were written to an output directory. The configuration parameters above specify that KBMOD should write three types of output files:\n",
    "\n",
    "1. A combined serialized ``Results`` saved as a .ecsv file (``\"result_filename\": \"./fake_results/results.ecsv\")``.\n",
    "2. (Legacy format) A series of individual output files (``\"res_filepath\": res_filepath``). Currently this is just the results file (trajectory information) and a copy of the final configuration used. Recent versions of KBMOD has removed older files, such as the psi curves or phi curves, that were not being used. However we can easily add files that would be useful.\n",
    "\n",
    "Users can shut off these outputs but passing ``None`` to the configuration options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(res_filepath):\n",
    "    files = os.listdir(res_filepath)\n",
    "    print(files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyzing Results\n",
    "\n",
    "The run function we used returns a `Results` object containing the individual results of the run. We can perform basic actions on this data structure such as sorting it, extracting individual results, or performing additional filtering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Search found {len(results)} results.\")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot the different curves for the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axs = plt.subplots(2, 1)\n",
    "axs[0].plot(results[\"psi_curve\"][0])\n",
    "axs[0].set_title(\"Psi\")\n",
    "\n",
    "axs[1].plot(results[\"phi_curve\"][0])\n",
    "axs[1].set_title(\"Psi\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For additional analysis steps, including manual filtering and exploration, please refer to the `kbmod_analysis_demo` notebook which uses the data generated by this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Debugging Results\n",
    "\n",
    "What happens if there is a trajectory we expect to be in an image that we do not find? This can happen if we are using data with injected fakes or known objects. In this case, we can use KBMOD's `track_filtered` setting to do additional debugging. We start by making a few changes to the configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turn on filtered tracking\n",
    "input_data.config.set(\"track_filtered\", True)\n",
    "\n",
    "# Turn up filtering of stamp filtering. This will require 100% of the stamp's flux\n",
    "# to be at the center pixel and effectively filter every candidate trajectory.\n",
    "input_data.config.set(\"center_thresh\", 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we rerun the search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = SearchRunner()\n",
    "results = rs.run_search_from_work_unit(input_data)\n",
    "print(f\"Search found {len(results)} results.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected we found no results. Everything was filtered in at least one stage. Let's see where."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(results.filtered_stats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The stamp filter removed the majority of the results. Since we set `track_filtered` to true, we can look at the actual results removed during that stage.\n",
    "\n",
    "We revert all the filters and add the reason for the filtering in its own column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results.revert_filter(add_column=\"filtered_reason\")\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we search for one of the expected trajectories (starting at pixel (50, 40) at the first time step) by using the table's search functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = results.table[(results.table[\"x\"] == 50) & (results.table[\"y\"] == 40)]\n",
    "print(subset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see all of the potential trajectories were filtered by the stamp filter. We can use this information to help tune different filtering stages."
   ]
  },
  {
   "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
}