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
import kbmod.search as kb
import numpy as np
import matplotlib.pyplot as plt
import os

im_path = "../data/demo/"
res_path = "./results"

psf

2D Point Spread Function Array

### raw_image 2D Image array

layered_image

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

image_stack

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

stack_search

Searches an image_stack for a moving psf

trajectory

Stores an object’s position and motion through an image_stack

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 = kb.psf(1.0)

The psf can be cast into a numpy array

np.array(p)

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 = kb.psf(arr)  # initialized from array
arr = np.square(arr)
p2.set_array(arr)  # set from array
np.array(p2)

There are several methods that get information about its properties

print(f"dim = {p.get_dim()}")  # dimension of kernel width and height
print(f"radius = {p.get_radius()}")  # distance from center of kernel to edge
print(f"size = {p.get_size()}")  # total number of pixels in the kernel
print(
    f"sum = {p.get_sum()}"
)  # total sum of all pixels in the kernel. Should be close to 1.0 for a normalized kernel

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

A layered_image can be initialized 2 ways:

A. Load a file for kbmod reference:

The layered_image is loaded given the path and filename to the FITS file as well as the PSF for the image.

im = kb.layered_image(im_path + "000000.fits", p)
print(f"Loaded a {im.get_width()} by {im.get_height()} image at time {im.get_time()}")

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:

im = kb.layered_image("image2", 100, 100, 5.0, 25.0, 0.0, p)
# name, width, height, background_noise_sigma, variance, capture_time, PSF

You can access a variety of information from the layered_image object

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

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

pixels = im.science()
pixels

A layered_image can have its layers set from any numpy array

raw = kb.raw_image(np.ones_like(pixels))

im.set_science(raw)
im.set_variance(raw)
im.science()

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.

The image at any point can be saved to a file

# im.save_layers(im_path+"/out") # file will use original name

Inserting artificial objects

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.

im.add_object(20.0, 35.0, 2500.0)
# x, y, flux

Image masking

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)

flags = ~0
flag_exceptions = [32, 39]
# mask all of pixels with flags except those with specifed combiniations
im.apply_mask_flags(flags, flag_exceptions)

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

# image_stack A collection of layered_images (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 = [kb.layered_image("img" + str(n), 100, 100, 10.0, 5.0, n / count, p) for n in range(count)]
stack = kb.image_stack(imlist)

Manually set the times the images in the stack were taken

stack.set_times([0, 2, 3, 4.5, 5, 6, 7, 10, 11, 14])

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.

import os

if os.path.exists(im_path):
    files = os.listdir(im_path)
    files = [im_path + f for f in files if ".fits" in f]
    files.sort()
    print("Using loaded files:")
    print(files)

    # Create default PSFs for each image.
    all_psfs = [p for _ in range(len(files))]

    # Load the images.
    stack = kb.image_stack(files, all_psfs)
else:
    print("Cannot find data directory. Using fake images.")

A global mask can be generated and applied to the stack

flags = ~0  # mask pixels with any flags
flag_exceptions = [32, 39]  # unless it has one of these special combinations of flags
global_flags = int("100111", 2)  # mask any pixels which have any of
# these flags in more than two images

stack.apply_global_mask(global_flags, 2)

The global mask is saved for future reference

stack.get_global_mask()

Most features of the layered_image can be used on the whole stack

# Apply the masking flags to each image.
stack.apply_mask_flags(flags, flag_exceptions)

# Convolve each image with its saved PFS.
stack.convolve_psf()

# Get image statitics. These functions return the information for the first image in the case
# where the stack contains images of different sizes.
print(f"Width = {stack.get_width()}")
print(f"Height = {stack.get_height()}")
print(f"Pixels Per Image = {stack.get_ppi()}")

# Retrieve a list of layered_images back from the stack.
stack.get_images()

# Get the list of image time stamps.
print(f"Times = {stack.get_times()}")

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 layered_images.
im_list = stack.get_images()

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

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

# stack_search

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

search = kb.stack_search(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.

search.search(10, 10, -0.1, 0.1, 5, 15, 2)
# angle_steps, velocity_steps, min_angle, max_angle, min_velocity, max_velocity, min_observations

Save the results to a files

note: format is {x, y, xv, yv, likelihood, flux}

if os.path.exists(res_path):
    search.save_results(res_path + "results.txt", 0.05)
    # path, fraction of total results to save in file
else:
    print("Data directory does not exist. Skipping file operations.")

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.x_v}")
print(f"y_v = {best.y_v}")

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