{
"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",
"from kbmod.core.psf import PSF\n",
"import kbmod.search as kb\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import os\n",
"\n",
"im_file = \"../data/demo_image.fits\"\n",
"res_path = \"./results\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"### [psf](#psf) \n",
"2D Point Spread Function Array \n",
"### [RawImage](#raw)\n",
"2D Image array \n",
"\n",
"### [LayeredImage](#layered) \n",
"A Complete image represented as 3 RawImage layers (science, mask, variance) \n",
"\n",
"### [ImageStack](#stack) \n",
"Stack of LayeredImages, intended to be the same frame captured at different times\n",
"\n",
"### [StackSearch](#search) \n",
"Searches an ImageStack for a moving psf\n",
"\n",
"### [trajectory](#traj)\n",
"Stores an object's position and motion through an ImageStack\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 = PSF(1.0)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The psf contains a numpy array of its kernel"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"p.kernel"
]
},
{
"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 = PSF(arr) # initialized from array"
]
},
{
"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.width}\") # dimension of kernel width and height\n",
"print(f\"radius = {p.radius}\") # distance from center of kernel to edge"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"# LayeredImage\n",
"Stores the science, mask, and variance image for a single image. The \"layered\" means it contains all of them together. The LayeredImage also stores auxiliary data, including the time of the image and the image’s PSF.\n",
"\n",
"A LayeredImage can be initialized 2 ways: \n",
"\n",
"### A. Load a file for kbmod reference:\n",
"The LayeredImage is loaded given the path and filename to the FITS file as well as the PSF's kernel for the image. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from kbmod.util_functions import load_deccam_layered_image\n",
"\n",
"im = load_deccam_layered_image(im_file, p.kernel)\n",
"print(f\"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_obstime()}\")"
]
},
{
"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": [
"from kbmod.fake_data.fake_data_creator import make_fake_layered_image\n",
"\n",
"im = make_fake_layered_image(100, 100, 5.0, 25.0, 0.0, p.kernel)\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 LayeredImage 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_npixels()}\")\n",
"print(f\"Time = {im.get_obstime()}\")"
]
},
{
"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.get_science()\n",
"pixels"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"A LayeredImage can have its layers set from any numpy array with type `float32` and observation time."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"all_ones = np.ones((im.get_height(), im.get_width()))\n",
"raw = kb.RawImage(all_ones.astype(np.float32))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"im.set_science(raw)\n",
"im.set_variance(raw)\n",
"im.get_science()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Inserting artificial objects\n",
"\n",
"Artificial objects can easily be added into a LayeredImage. The LayeredImage generates a point observation at the given pixel and applies the image's PSF."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from kbmod.fake_data.fake_data_creator import add_fake_object\n",
"\n",
"add_fake_object(im, 20.0, 35.0, 2500.0)\n",
"# x, y, flux"
]
},
{
"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": [
"\n",
"# ImageStack\n",
"A collection of LayeredImages (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 = [make_fake_layered_image(100, 100, 5.0, 25.0, n / count, p.kernel) for n in range(count)]\n",
"stack = kb.ImageStack(imlist)"
]
},
{
"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 LayeredImages.\n",
"im_list = stack.get_images()\n",
"\n",
"# Create a new list of LayeredImages with the added object.\n",
"new_im_list = []\n",
"for im, time in zip(im_list, stack.build_zeroed_times()):\n",
" add_fake_object(im, 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 ImageStack.\n",
"stack = kb.ImageStack(new_im_list)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"# StackSearch\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.StackSearch(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. To do this we need to first create a generator object to generate the trajectories. Those trajectories get feed into the search function."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from kbmod.trajectory_generator import KBMODV1Search\n",
"\n",
"gen = KBMODV1Search(\n",
" 10, 5, 15, 10, -0.1, 0.1\n",
") # velocity_steps, min_vel, max_vel, angle_steps, min_ang, max_ang,\n",
"candidates = [trj for trj in gen]\n",
"print(f\"Created {len(candidates)} candidate trajectories per pixel.\")\n",
"\n",
"search.search_all(candidates, 2)"
]
},
{
"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": [
"\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.vx}\")\n",
"print(f\"y_v = {best.vy}\")"
]
},
{
"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": "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
}