{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KBMOD Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from kbmod.fake_data_creator import *\n",
    "from kbmod.run_search import *\n",
    "from kbmod.search import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup file paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are at least two file paths you need to setup in order to run kbmod:\n",
    "1. The im_filepath provides a path to the input images.\n",
    "1. The res_filepath provides a path to the directory where the output results will be stored.\n",
    "\n",
    "A time and psf file can optionally be specified.\n",
    "\n",
    "If you already have data files, you can use those. Below we use the data in `data/demo`. You can also create your own fake data using `fake_data_creator.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['000000.fits', '000001.fits', '000002.fits', '000003.fits', '000004.fits', '000005.fits', '000006.fits', '000007.fits', '000008.fits', '000009.fits']\n"
     ]
    }
   ],
   "source": [
    "# Define the path for the data.\n",
    "im_filepath = \"../data/demo\"\n",
    "print(os.listdir(im_filepath))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------------------------\n",
      "Loading Images\n",
      "---------------------------------------\n",
      "Loaded 10 images\n",
      "Times set\n",
      "Starting Search\n",
      "---------------------------------------\n",
      "Ecliptic Angle = 1.1901\n",
      "Min. Search Angle = -0.5000\n",
      "Max Search Angle = 0.5000\n",
      "Min Velocity = 0.0000\n",
      "Max Velocity = 20.0000\n",
      "Using in-line GPU sigmaG filtering methods\n",
      "sigmaG limits: [15,60]\n",
      "sigmaG coeff: 0.7753\n",
      "Search finished in 0.057s\n",
      "---------------------------------------\n",
      "Retrieving Results\n",
      "---------------------------------------\n",
      "Getting results...\n",
      "---------------------------------------\n",
      "Chunk Start = 0\n",
      "Chunk Max Likelihood = 243.35\n",
      "Chunk Min. Likelihood = -1.00\n",
      "---------------------------------------\n",
      "Extracted batch of 1343 results for total of 1343\n",
      "Applying Clipped-sigmaG Filtering\n",
      "sigmaG limits: [15,60]\n",
      "sigmaG coeff: 0.7753\n",
      "0.18s elapsed\n",
      "Completed filtering.\n",
      "---------------------------------------\n",
      "---------------------------------------\n",
      "Applying Stamp Filtering\n",
      "---------------------------------------\n",
      "Stamp filtering 1342 results\n",
      "Keeping 588 results\n",
      "0.02s elapsed\n",
      "Clustering 588 results\n",
      "---------------------------------------\n",
      "Saving Results\n",
      "---------------------------------------\n",
      "Time taken for patch:  0.9508676528930664\n"
     ]
    }
   ],
   "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, 20].\n",
    "v_min = 0\n",
    "v_max = 20\n",
    "v_steps = 21\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 = 11\n",
    "ang_arr = [ang_below, ang_above, ang_steps]\n",
    "\n",
    "# There are 10 images in the demo data. Make sure we see\n",
    "# the object in at least 7.\n",
    "num_obs = 7\n",
    "\n",
    "input_parameters = {\n",
    "    # Required\n",
    "    \"im_filepath\": im_filepath,\n",
    "    \"res_filepath\": res_filepath,\n",
    "    \"time_file\": None,\n",
    "    \"output_suffix\": results_suffix,\n",
    "    \"v_arr\": v_arr,\n",
    "    \"ang_arr\": ang_arr,\n",
    "    # Important\n",
    "    \"num_obs\": num_obs,\n",
    "    \"do_mask\": True,\n",
    "    \"lh_level\": 10.0,\n",
    "    \"gpu_filter\": True,\n",
    "    # Fine tuning\n",
    "    \"sigmaG_lims\": [15, 60],\n",
    "    \"mom_lims\": [37.5, 37.5, 1.5, 1.0, 1.0],\n",
    "    \"peak_offset\": [3.0, 3.0],\n",
    "    \"chunk_size\": 1000000,\n",
    "    \"stamp_type\": \"cpp_median\",\n",
    "    \"eps\": 0.03,\n",
    "    \"clip_negative\": True,\n",
    "    \"mask_num_images\": 10,\n",
    "    \"cluster_type\": \"position\",\n",
    "    # Override the ecliptic angle for the demo data since we\n",
    "    # know the true angle in pixel space.\n",
    "    \"average_angle\": 0.0,\n",
    "}\n",
    "\n",
    "rs = run_search(input_parameters)\n",
    "rs.run_search()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that results were written."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['results_DEMO.txt', 'lc_DEMO.txt', 'psi_DEMO.txt', 'phi_DEMO.txt', 'lc_index_DEMO.txt', 'times_DEMO.txt', 'filtered_likes_DEMO.txt', 'ps_DEMO.txt', 'all_ps_DEMO.npy']\n"
     ]
    }
   ],
   "source": [
    "if os.path.exists(res_filepath):\n",
    "    files = os.listdir(res_filepath)\n",
    "    print(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (.conda-kbmod)",
   "language": "python",
   "name": "conda-env-.conda-kbmod-py"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}