{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KBMOD Reference  \n",
    "  \n",
    "This notebook demonstrates a gpu-accelerated image processing framework designed for image stack and time domain analysis, compatible with FITS and numpy.\n",
    "\n",
    "An example of the C++ interface can be found in search/src/kbmod.cpp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup for Kbmod Reference\n",
    "Before importing, make sure you have installed kbmod using `pip install .` in the root directory.  Also be sure you are running with python3.\n",
    "\n",
    "If you are running into trouble with importing `kbmod` and related libraries and get a `ModuleNotFoundError` or `ModelImportError`, make sure that: a) your notebook is using the correct kernel and b) the pybinds directory is in the python path. Part b should happen automatically with pip install, but if not the path can be manually added with:\n",
    "\n",
    "```\n",
    "import sys\n",
    "sys.path.insert(0, 'HOMEDIR/kbmod/src')\n",
    "```\n",
    "\n",
    "where HOMEDIR is the location of kbmod directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# everything we will need for this demo\n",
    "import kbmod.search as kb\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "im_path = \"../data/demo/\"\n",
    "res_path = \"./results\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "\n",
    "### [psf](#psf) \n",
    "2D Point Spread Function Array  \n",
    "### [raw_image](#raw)\n",
    "2D Image array  \n",
    "\n",
    "### [layered_image](#layered) \n",
    "A Complete image represented as 3 raw_image layers (science, mask, variance)   \n",
    "\n",
    "### [image_stack](#stack)  \n",
    "Stack of layered_images, intended to be the same frame captured at different times\n",
    "\n",
    "### [stack_search](#search)  \n",
    "Searches an image_stack for a moving psf\n",
    "\n",
    "### [trajectory](#traj)\n",
    "Stores an object's position and motion through an image_stack\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# psf\n",
    "A 2D psf kernel, for convolution and adding artificial sources to images  \n",
    "\n",
    "This simple constructor initializes a gaussian psf with a sigma of 1.0 pixels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = kb.psf(1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The psf can be cast into a numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A psf can also be initialized or set from a numpy array, but the array must be square and have odd dimensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = np.linspace(0.0, 1.0, 9).reshape(3, 3)\n",
    "p2 = kb.psf(arr)  # initialized from array\n",
    "arr = np.square(arr)\n",
    "p2.set_array(arr)  # set from array\n",
    "np.array(p2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are several methods that get information about its properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"dim = {p.get_dim()}\")  # dimension of kernel width and height\n",
    "print(f\"radius = {p.get_radius()}\")  # distance from center of kernel to edge\n",
    "print(f\"size = {p.get_size()}\")  # total number of pixels in the kernel\n",
    "print(\n",
    "    f\"sum = {p.get_sum()}\"\n",
    ")  # total sum of all pixels in the kernel. Should be close to 1.0 for a normalized kernel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"layered\"></a>\n",
    "# layered_image\n",
    "Stores the science, mask, and variance image for a single image. The \"layered\" means it contains all of them together.  The layered_image also stores auxiliary data, including the time of the image and the image’s PSF.\n",
    "\n",
    "A layered_image can be initialized 2 ways:  \n",
    "\n",
    "### A. Load a file for kbmod reference:\n",
    "The layered_image is loaded given the path and filename to the FITS file as well as the PSF for the image. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = kb.layered_image(im_path + \"000000.fits\", p)\n",
    "print(f\"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_time()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "KBMOD will attempt to read the timestamp from the FITS file (using the `MJD` field in the header). If no timestamp is provided then one can be set manually using `set_time`.\n",
    "\n",
    "### B. Generate a new image from scratch with random noise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = kb.layered_image(\"image2\", 100, 100, 5.0, 25.0, 0.0, p)\n",
    "# name, width, height, background_noise_sigma, variance, capture_time, PSF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can access a variety of information from the layered_image object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Width = {im.get_width()}\")\n",
    "print(f\"Height = {im.get_height()}\")\n",
    "print(f\"Pixels Per Image = {im.get_ppi()}\")\n",
    "print(f\"Time = {im.get_time()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The image pixels' values can be retrieved as a 2D numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixels = im.science()\n",
    "pixels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A layered_image can have its layers set from any numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw = kb.raw_image(np.ones_like(pixels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im.set_science(raw)\n",
    "im.set_variance(raw)\n",
    "im.science()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: Due to how pybind handles references, you will need to call the `set_` function for any layer you change. Simply getting a layer and changing the values directly will not propagate those changes back into the C++ object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The image at any point can be saved to a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# im.save_layers(im_path+\"/out\") # file will use original name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inserting artificial objects\n",
    "\n",
    "Artificial objects can easily be added into a layered_image. The layered_image generates a point observation at the given pixel and applies the image's PSF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im.add_object(20.0, 35.0, 2500.0)\n",
    "# x, y, flux"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image masking\n",
    "\n",
    "The image can mask itself by providing a bitmask of flags (note: masked pixels are set to -9999 so they can be distinguished later from 0.0 pixles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flags = ~0\n",
    "flag_exceptions = [32, 39]\n",
    "# mask all of pixels with flags except those with specifed combiniations\n",
    "im.apply_mask_flags(flags, flag_exceptions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convolution with PSF.\n",
    "\n",
    "The image can be convolved with a psf kernel using `convolve_psf`. Most users should not need to call this function explicitly since it is automatically called during relevant functions, such as inserting an artificial object or searching the image stack. The function is only exposed because it happens to be a fast implementation of a generally useful function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im.convolve_psf()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"stack\"></a>\n",
    "# image_stack\n",
    "A collection of layered_images (usually at different times). Used to apply operations to a group of images.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a stack with 10 50x50 images with random noise and times ranging from 0 to 1\n",
    "count = 10\n",
    "imlist = [kb.layered_image(\"img\" + str(n), 100, 100, 10.0, 5.0, n / count, p) for n in range(count)]\n",
    "stack = kb.image_stack(imlist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manually set the times the images in the stack were taken "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stack.set_times([0, 2, 3, 4.5, 5, 6, 7, 10, 11, 14])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A shortcut is provided to initialize a stack automatically from a list of files. If 'MJD' is in the header for each image, the stack will automatically load the times as well. If not, you can set them as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if os.path.exists(im_path):\n",
    "    files = os.listdir(im_path)\n",
    "    files = [im_path + f for f in files if \".fits\" in f]\n",
    "    files.sort()\n",
    "    print(\"Using loaded files:\")\n",
    "    print(files)\n",
    "\n",
    "    # Create default PSFs for each image.\n",
    "    all_psfs = [p for _ in range(len(files))]\n",
    "\n",
    "    # Load the images.\n",
    "    stack = kb.image_stack(files, all_psfs)\n",
    "else:\n",
    "    print(\"Cannot find data directory. Using fake images.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A global mask can be generated and applied to the stack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flags = ~0  # mask pixels with any flags\n",
    "flag_exceptions = [32, 39]  # unless it has one of these special combinations of flags\n",
    "global_flags = int(\"100111\", 2)  # mask any pixels which have any of\n",
    "# these flags in more than two images\n",
    "\n",
    "stack.apply_global_mask(global_flags, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The global mask is saved for future reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stack.get_global_mask()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most features of the layered_image can be used on the whole stack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply the masking flags to each image.\n",
    "stack.apply_mask_flags(flags, flag_exceptions)\n",
    "\n",
    "# Convolve each image with its saved PFS.\n",
    "stack.convolve_psf()\n",
    "\n",
    "# Get image statitics. These functions return the information for the first image in the case\n",
    "# where the stack contains images of different sizes.\n",
    "print(f\"Width = {stack.get_width()}\")\n",
    "print(f\"Height = {stack.get_height()}\")\n",
    "print(f\"Pixels Per Image = {stack.get_ppi()}\")\n",
    "\n",
    "# Retrieve a list of layered_images back from the stack.\n",
    "stack.get_images()\n",
    "\n",
    "# Get the list of image time stamps.\n",
    "print(f\"Times = {stack.get_times()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we will create a very bright object and add it to the images and create a new image stack with the new object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the individual layered_images.\n",
    "im_list = stack.get_images()\n",
    "\n",
    "# Create a new list of layered_images with the added object.\n",
    "new_im_list = []\n",
    "for im, time in zip(im_list, stack.get_times()):\n",
    "    im.add_object(20.0 + (time * 8.0), 35.0 + (time * 0.0), 25000.0)\n",
    "    new_im_list.append(im)\n",
    "\n",
    "# Save these images in a new image_stack.\n",
    "stack = kb.image_stack(new_im_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"search\"></a>\n",
    "# stack_search\n",
    "\n",
    "We can create a search object that will compute auxiliary data for the images and run the search algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search = kb.stack_search(stack)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To save psi and images, a directory with \"psi\" and \"phi\" folders must be specified. In general the psi and phi images are used for debugging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(res_path):\n",
    "    if os.path.exists(os.path.join(res_path, \"out/psi\")) is False:\n",
    "        os.mkdir(os.path.join(res_path, \"out/psi\"))\n",
    "\n",
    "    if os.path.exists(os.path.join(res_path, \"out/phi\")) is False:\n",
    "        os.mkdir(os.path.join(res_path, \"out/phi\"))\n",
    "\n",
    "    search.save_psi_phi(os.path.join(res_path, \"out\"))\n",
    "else:\n",
    "    print(\"Data directory does not exist. Skipping file operations.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch a basic search that uses the a grid of velocities and angles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search.search(10, 10, -0.1, 0.1, 5, 15, 2)\n",
    "# angle_steps, velocity_steps, min_angle, max_angle, min_velocity, max_velocity, min_observations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the results to a files  \n",
    "note: format is {x, y, xv, yv, likelihood, flux}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(res_path):\n",
    "    search.save_results(res_path + \"results.txt\", 0.05)\n",
    "    # path, fraction of total results to save in file\n",
    "else:\n",
    "    print(\"Data directory does not exist. Skipping file operations.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trajectories can be retrieved directly from search without writing and reading to file.  \n",
    "However, this is not recommended for a large number of trajectories, as it is not returned as a numpy array, but as a list of the trajectory objects described below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_results = search.get_results(0, 100)\n",
    "# start, count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The basic search does not do any filtering. You can enable basic GPU filtering using a clipped sigmaG filtering by calling `enable_gpu_sigmag_filter` before the search. The function takes a sigmaG coefficient that is derived from the percentiles and can be computed using `PostProcess._find_sigmaG_coeff()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search.enable_gpu_sigmag_filter([0.25, 0.75], 0.7413, 10.0)\n",
    "# sigmaG limits, sigmaG coefficient, the likelihood threshold"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note:* The sigmaG coefficient 0.7413 applies only to the percentile range of 25th and 75th. If you change the percentile range, then you will also need to update the coefficient."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"traj\"></a>\n",
    "# trajectory\n",
    "A simple container with properties representing an object and its path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = top_results[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# these numbers are wild because mask flags and search parameters above were chosen randomly\n",
    "print(f\"Flux = {best.flux}\")\n",
    "print(f\"Likelihood = {best.lh}\")\n",
    "print(f\"x = {best.x}\")\n",
    "print(f\"y = {best.y}\")\n",
    "print(f\"x_v = {best.x_v}\")\n",
    "print(f\"y_v = {best.y_v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These top_results are all be duplicating searches on the same bright object we added.\n",
    "top_results[:20]"
   ]
  },
  {
   "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": 2
}