KBMOD Reference

This notebook demonstrates a gpu-accelerated image processing framework designed for image stack and time domain analysis, compatible with FITS and numpy.

An example of the C++ interface can be found in search/src/kbmod.cpp

Setup for Kbmod Reference

Before importing, make sure you have installed kbmod using pip install . in the root directory. Also be sure you are running with python3.

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:

import sys
sys.path.insert(0, 'HOMEDIR/kbmod/src')

where HOMEDIR is the location of kbmod directory.

# everything we will need for this demo
from kbmod.core.psf import PSF
import kbmod.search as kb
import numpy as np
import matplotlib.pyplot as plt
import os

im_file = "../data/demo_image.fits"
res_path = "./results"

psf

2D Point Spread Function Array

### RawImage 2D Image array

LayeredImage

A Complete image represented as 3 RawImage layers (science, mask, variance)

ImageStack

Stack of LayeredImages, intended to be the same frame captured at different times

StackSearch

Searches an ImageStack for a moving psf

trajectory

Stores an object’s position and motion through an ImageStack

psf

A 2D psf kernel, for convolution and adding artificial sources to images

This simple constructor initializes a gaussian psf with a sigma of 1.0 pixels

p = PSF(1.0)

The psf contains a numpy array of its kernel

p.kernel

A psf can also be initialized or set from a numpy array, but the array must be square and have odd dimensions

arr = np.linspace(0.0, 1.0, 9).reshape(3, 3)
p2 = PSF(arr)  # initialized from array

There are several methods that get information about its properties

print(f"dim = {p.width}")  # dimension of kernel width and height
print(f"radius = {p.radius}")  # distance from center of kernel to edge

# LayeredImage 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.

A LayeredImage can be initialized 2 ways:

A. Load a file for kbmod reference:

The LayeredImage is loaded given the path and filename to the FITS file as well as the PSF’s kernel for the image.

from kbmod.util_functions import load_deccam_layered_image

im = load_deccam_layered_image(im_file, p.kernel)
print(f"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_obstime()}")

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.

B. Generate a new image from scratch with random noise:

from kbmod.fake_data.fake_data_creator import make_fake_layered_image

im = make_fake_layered_image(100, 100, 5.0, 25.0, 0.0, p.kernel)
# name, width, height, background_noise_sigma, variance, capture_time, PSF

You can access a variety of information from the LayeredImage object

print(f"Width = {im.get_width()}")
print(f"Height = {im.get_height()}")
print(f"Pixels Per Image = {im.get_npixels()}")
print(f"Time = {im.get_obstime()}")

The image pixels’ values can be retrieved as a 2D numpy array

pixels = im.get_science()
pixels

A LayeredImage can have its layers set from any numpy array with type float32 and observation time.

all_ones = np.ones((im.get_height(), im.get_width()))
raw = kb.RawImage(all_ones.astype(np.float32))

im.set_science(raw)
im.set_variance(raw)
im.get_science()

Inserting artificial objects

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.

from kbmod.fake_data.fake_data_creator import add_fake_object

add_fake_object(im, 20.0, 35.0, 2500.0)
# x, y, flux

Convolution with PSF.

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.

im.convolve_psf()

# ImageStack A collection of LayeredImages (usually at different times). Used to apply operations to a group of images.

# Create a stack with 10 50x50 images with random noise and times ranging from 0 to 1
count = 10
imlist = [make_fake_layered_image(100, 100, 5.0, 25.0, n / count, p.kernel) for n in range(count)]
stack = kb.ImageStack(imlist)

Here, we will create a very bright object and add it to the images and create a new image stack with the new object.

# Get the individual LayeredImages.
im_list = stack.get_images()

# Create a new list of LayeredImages with the added object.
new_im_list = []
for im, time in zip(im_list, stack.build_zeroed_times()):
    add_fake_object(im, 20.0 + (time * 8.0), 35.0 + (time * 0.0), 25000.0)
    new_im_list.append(im)

# Save these images in a new ImageStack.
stack = kb.ImageStack(new_im_list)

# StackSearch

We can create a search object that will compute auxiliary data for the images and run the search algorithms.

search = kb.StackSearch(stack)

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.

if os.path.exists(res_path):
    if os.path.exists(os.path.join(res_path, "out/psi")) is False:
        os.mkdir(os.path.join(res_path, "out/psi"))

    if os.path.exists(os.path.join(res_path, "out/phi")) is False:
        os.mkdir(os.path.join(res_path, "out/phi"))

    search.save_psi_phi(os.path.join(res_path, "out"))
else:
    print("Data directory does not exist. Skipping file operations.")

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.

from kbmod.trajectory_generator import KBMODV1Search

gen = KBMODV1Search(
    10, 5, 15, 10, -0.1, 0.1
)  # velocity_steps, min_vel, max_vel, angle_steps, min_ang, max_ang,
candidates = [trj for trj in gen]
print(f"Created {len(candidates)} candidate trajectories per pixel.")

search.search_all(candidates, 2)

Trajectories can be retrieved directly from search without writing and reading to file.

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

top_results = search.get_results(0, 100)
# start, count

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().

search.enable_gpu_sigmag_filter([0.25, 0.75], 0.7413, 10.0)
# sigmaG limits, sigmaG coefficient, the likelihood threshold

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.

# trajectory A simple container with properties representing an object and its path

best = top_results[0]

# these numbers are wild because mask flags and search parameters above were chosen randomly
print(f"Flux = {best.flux}")
print(f"Likelihood = {best.lh}")
print(f"x = {best.x}")
print(f"y = {best.y}")
print(f"x_v = {best.vx}")
print(f"y_v = {best.vy}")

# These top_results are all be duplicating searches on the same bright object we added.
top_results[:20]