notebooks/Kbmod_Reference.ipynb

{
 "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\n",
    "Before importing, make sure to run `source setup.bash` in the root directory.  \n",
    "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. In python this is something like:\n",
    "\n",
    "```\n",
    "import sys\n",
    "sys.path.insert(0, 'HOMEDIR/kbmod/search/pybinds')\n",
    "```\n",
    "\n",
    "where HOMEDIR is your home directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#everything we will need for this demo\n",
    "from kbmodpy import kbmod as kb\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "path = \"../data/demo/\""
   ]
  },
  {
   "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": 2,
   "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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.00367206, 0.01464826, 0.02320431, 0.01464826, 0.00367206],\n",
       "       [0.01464826, 0.05843356, 0.09256457, 0.05843356, 0.01464826],\n",
       "       [0.02320431, 0.09256457, 0.1466315 , 0.09256457, 0.02320431],\n",
       "       [0.01464826, 0.05843356, 0.09256457, 0.05843356, 0.01464826],\n",
       "       [0.00367206, 0.01464826, 0.02320431, 0.01464826, 0.00367206]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.      , 0.015625, 0.0625  ],\n",
       "       [0.140625, 0.25    , 0.390625],\n",
       "       [0.5625  , 0.765625, 1.      ]], dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.975315511226654"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p.get_dim() # dimension of kernel width and height\n",
    "p.get_radius() # distance from center of kernel to edge\n",
    "p.get_size() # total number of pixels in the kernel\n",
    "p.get_sum() # total sum of all pixels in the kernel, \n",
    "            #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.  \n",
    "It can be initialized 2 ways:  \n",
    "A. Load a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#im = kb.layered_image(path+\"example.fits\", p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "B. Generate a new image from scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Artificial objects can easily be added into a layered_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "im.add_object(20.0, 35.0, 2500.0)\n",
    "# x, y, flux, psf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The image pixels can be retrieved as a 2D numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-10.626595  ,  -6.433425  ,  -0.63092774, ...,   9.480564  ,\n",
       "          5.7158356 ,   3.428888  ],\n",
       "       [ -3.4042237 ,   5.1546435 ,  -9.349374  , ...,   4.2239113 ,\n",
       "         -1.8593729 ,  -0.2675743 ],\n",
       "       [  1.3247397 ,  -3.2443523 ,  -1.388663  , ...,   1.1003296 ,\n",
       "         -1.6527367 ,  -5.297088  ],\n",
       "       ...,\n",
       "       [ -0.2647652 ,   7.6661224 ,   2.7818446 , ...,  -5.5298233 ,\n",
       "         -1.4977067 ,  -3.243896  ],\n",
       "       [  0.4727264 ,  -1.4867105 ,   5.2125196 , ...,  -4.342271  ,\n",
       "          0.46082076,   1.0669568 ],\n",
       "       [ -6.926575  ,   0.5734095 ,   1.8180655 , ...,   0.05729845,\n",
       "         -4.188006  ,  -6.2619734 ]], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pixels = im.science()\n",
    "pixels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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": 10,
   "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": [
    "The image can be convolved with a psf kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "im.convolve_psf()\n",
    "# note: This function is called interally by stack_search and doesn't need to be\n",
    "# used directy. It is only exposed because it happens to be a fast \n",
    "# implementation of a generally useful function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The image at any point can be saved to a file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#im.save_layers(path+\"/out\") # file will use original name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A layered_image can have its layers set from any numpy array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw = kb.raw_image( np.ones_like(pixels) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 1., 1., ..., 1., 1., 1.],\n",
       "       [1., 1., 1., ..., 1., 1., 1.],\n",
       "       [1., 1., 1., ..., 1., 1., 1.],\n",
       "       ...,\n",
       "       [1., 1., 1., ..., 1., 1., 1.],\n",
       "       [1., 1., 1., ..., 1., 1., 1.],\n",
       "       [1., 1., 1., ..., 1., 1., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "im.set_science( raw )\n",
    "im.set_variance( raw )\n",
    "im.science()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Image width = 100\n",
      "Image height = 100\n",
      "Image time = 0\n",
      "Image ppi = 10000\n"
     ]
    }
   ],
   "source": [
    "print('Image width = %i' % im.get_width())\n",
    "print('Image height = %i' % im.get_height())\n",
    "print('Image time = %i' % im.get_time())\n",
    "print('Image ppi = %i' % im.get_ppi()) # pixels per image, width*height"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"stack\"></a>\n",
    "# image_stack\n",
    "A collection of layered_images. Used to apply operations to a group of images.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 )\n",
    "# this creates a stack with 10 50x50 images, and times ranging from 0 to 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manually set the times the images in the stack were taken "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "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": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cannot find data directory. Using fake images.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "if os.path.exists(path):\n",
    "    files = os.listdir(path)\n",
    "    files = [path+f for f in files if '.fits' in f]\n",
    "    files.sort()\n",
    "    print('Using loaded files:')\n",
    "    print(files)\n",
    "    stack = kb.image_stack(files)\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": 19,
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most features of the layered_image can be used on the whole stack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0, 2.0, 3.0, 4.5, 5.0, 6.0, 7.0, 10.0, 11.0, 14.0]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stack.apply_mask_flags(flags, flag_exceptions)\n",
    "stack.apply_global_mask(global_flags, 2)\n",
    "stack.convolve_psf()\n",
    "stack.get_width()\n",
    "stack.get_height()\n",
    "stack.get_ppi()\n",
    "stack.get_images() # retrieves list of layered_images back from the stack\n",
    "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": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "im_list = stack.get_images()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_im_list = []\n",
    "for im, time in zip(im_list, stack.get_times()):\n",
    "    im.add_object(20.0+(time*8.), 35.0+(time*0.), 25000.0)\n",
    "    new_im_list.append(im)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "stack = kb.image_stack(new_im_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"search\"></a>\n",
    "# stack_search\n",
    "Searches a stack of images for a given psf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data directory does not exist. Skipping file operations.\n"
     ]
    }
   ],
   "source": [
    "if os.path.exists(path):\n",
    "    if os.path.exists(os.path.join(path,'out/psi')) is False:\n",
    "        os.mkdir(os.path.join(path,'out/psi'))\n",
    "    \n",
    "    if os.path.exists(os.path.join(path,'out/phi')) is False:\n",
    "        os.mkdir(os.path.join(path,'out/phi'))\n",
    "\n",
    "    search.save_psi_phi(os.path.join(path, 'out'))\n",
    "else:\n",
    "    print('Data directory does not exist. Skipping file operations.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch a search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "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": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data directory does not exist. Skipping file operations.\n"
     ]
    }
   ],
   "source": [
    "if os.path.exists(path):\n",
    "    search.save_results(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": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_results = search.get_results(0, 100)\n",
    "# start, count"
   ]
  },
  {
   "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": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = top_results[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.1188055276870728"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# these numbers are wild because mask flags and search parameters above were chosen randomly\n",
    "best.flux \n",
    "best.lh\n",
    "best.x\n",
    "best.y\n",
    "best.x_v\n",
    "best.y_v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "tests/test_search.py shows a simple example of how to generate a set of images, add an artificial source, and recover it with search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[lh: 4704.088867 flux: 3326.292725 x: 52 y: 35 x_v: 13.955224 y_v: 1.118806 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 13.930058 y_v: -1.397668 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 13.997200 y_v: 0.279981 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 14.000000 y_v: -0.000000 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 13.997200 y_v: -0.279981 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 13.988802 y_v: -0.559851 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 13.974808 y_v: -0.839496 obs_count: 10,\n",
       " lh: 4703.314941 flux: 3325.745605 x: 52 y: 35 x_v: 13.955224 y_v: -1.118806 obs_count: 10,\n",
       " lh: 4702.495117 flux: 3325.166016 x: 52 y: 34 x_v: 13.955224 y_v: 1.118806 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 13.997200 y_v: -0.279981 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 13.988802 y_v: -0.559851 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 13.974808 y_v: -0.839496 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 14.000000 y_v: -0.000000 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 13.997200 y_v: 0.279981 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 13.930058 y_v: -1.397668 obs_count: 10,\n",
       " lh: 4701.964844 flux: 3324.791016 x: 52 y: 34 x_v: 13.955224 y_v: -1.118806 obs_count: 10,\n",
       " lh: 4700.060547 flux: 3323.444580 x: 53 y: 34 x_v: 13.930058 y_v: -1.397668 obs_count: 10,\n",
       " lh: 4700.060547 flux: 3323.444580 x: 53 y: 34 x_v: 13.955224 y_v: -1.118806 obs_count: 10,\n",
       " lh: 4700.060547 flux: 3323.444580 x: 53 y: 34 x_v: 13.974808 y_v: -0.839496 obs_count: 10,\n",
       " lh: 4700.060547 flux: 3323.444580 x: 53 y: 34 x_v: 13.988802 y_v: -0.559851 obs_count: 10]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}