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

Before importing, make sure to run source setup.bash in the root directory.
Also be sure you are running with python3.


In [4]:
#everything we will need for this demo
from kbmodpy import kbmod as kb
import numpy as np
import matplotlib.pyplot as plt
import os
path = "../data/demo/"

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

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


In [5]:
p = kb.psf(1.0)

The psf can be cast into a numpy array


In [6]:
np.array(p)


Out[6]:
array([[ 0.00367206,  0.01464826,  0.02320431,  0.01464826,  0.00367206],
       [ 0.01464826,  0.05843356,  0.09256457,  0.05843356,  0.01464826],
       [ 0.02320431,  0.09256457,  0.14663149,  0.09256457,  0.02320431],
       [ 0.01464826,  0.05843356,  0.09256457,  0.05843356,  0.01464826],
       [ 0.00367206,  0.01464826,  0.02320431,  0.01464826,  0.00367206]], dtype=float32)

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


In [7]:
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)


Out[7]:
array([[ 0.      ,  0.015625,  0.0625  ],
       [ 0.140625,  0.25    ,  0.390625],
       [ 0.5625  ,  0.765625,  1.      ]], dtype=float32)

There are several methods that get information about its properties


In [8]:
p.get_dim() # dimension of kernel width and height
p.get_radius() # distance from center of kernel to edge
p.get_size() # total number of pixels in the kernel
p.get_sum() # total sum of all pixels in the kernel, 
            #should be close to 1.0 for a normalized kernel


Out[8]:
0.975315511226654

layered_image

Stores the science, mask, and variance image for a single image. The "layered" means it contains all of them together.
It can be initialized 2 ways:
A. Load a file


In [9]:
#im = kb.layered_image(path+"example.fits")

B. Generate a new image from scratch


In [10]:
im = kb.layered_image("image2", 100, 100, 5.0, 25.0, 0.0)
# name, width, height, background_noise_sigma, variance, capture_time

Artificial objects can easily be added into a layered_image


In [11]:
im.add_object(20.0, 35.0, 2500.0, p)
# x, y, flux, psf

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


In [17]:
pixels = im.science()
pixels


Out[17]:
array([[ 1.41452205,  1.20091569,  1.08562613, ...,  0.16255999,
        -0.71415883, -0.77766204],
       [ 0.2462137 , -0.14571723, -0.17865041, ...,  0.91385883,
         0.11997563, -0.12612684],
       [ 0.19908643, -0.26934448, -0.297548  , ...,  1.01619554,
         0.30717701,  0.08602032],
       ..., 
       [-0.41059819, -0.40178138, -0.25474659, ...,  0.87045324,
         1.94920647,  2.66952872],
       [-0.86084074, -0.22506261, -0.19994019, ...,  0.89172727,
         1.37131238,  1.7821281 ],
       [-1.99527967, -0.99903578, -1.02412295, ..., -0.13386457,
         0.44371662,  1.71280479]], dtype=float32)

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)


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

The image can be convolved with a psf kernel


In [19]:
im.convolve(p)
# note: This function is called interally by stack_search and doesn't need to be
# used directy. It is only exposed because it happens to be a fast 
# implementation of a generally useful function

The image at any point can be saved to a file


In [20]:
#im.save_layers(path+"/out") # file will use original name

A layered_image can have its layers set from any numpy array


In [33]:
raw = kb.raw_image( np.ones_like(pixels) )

In [35]:
im.set_science( raw )
im.set_variance( raw )
im.science()


Out[35]:
array([[ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       ..., 
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.]], dtype=float32)

Get properties


In [16]:
im.get_width()
im.get_height()
im.get_time()
im.get_ppi() # pixels per image, width*height


Out[16]:
10000

image_stack

A collection of layered_images. Used to apply operations to a group of images.


In [14]:
count = 10
imlist = [ kb.layered_image("img"+str(n), 50, 50, 10.0, 5.0, n/count) for n in range(count) ]
stack = kb.image_stack( imlist )
# this creates a stack with 10 50x50 images, and times ranging from 0 to 1

Manually set the times the images in the stack were taken


In [15]:
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.


In [16]:
import os
files = os.listdir(path)
files = [path+f for f in files if '.fits' in f]
files.sort()
files


Out[16]:
['../data/demo/CORR40535777.fits',
 '../data/demo/CORR40535787.fits',
 '../data/demo/CORR40535797.fits',
 '../data/demo/CORR40535807.fits',
 '../data/demo/CORR40535817.fits']

In [17]:
stack = kb.image_stack(files)

A master mask can be generated and applied to the stack


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

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


In [19]:
stack.apply_mask_flags(flags, flag_exceptions)
stack.apply_master_mask(master_flags, 2)
stack.convolve(p)
stack.get_width()
stack.get_height()
stack.get_ppi()
stack.get_images() # retrieves list of layered_images back from the stack
stack.get_times()


Out[19]:
[0.0,
 0.0030539999715983868,
 0.006087000016123056,
 0.00913199968636036,
 0.012175999581813812]

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


In [20]:
im_list = stack.get_images()

In [21]:
new_im_list = []
for im, time in zip(im_list, stack.get_times()):
    im.add_object(20.0+(time*800.), 35.0+(time*0.), 25000.0, p)
    new_im_list.append(im)

In [22]:
stack = kb.image_stack(new_im_list)

stack_search

Searches a stack of images for a given psf


In [23]:
search = kb.stack_search( stack, p )

To save psi and images, a directory with "psi" and "phi" folders must be specified.


In [24]:
if os.path.exists(os.path.join(path,'out/psi')) is False:
    os.mkdir(os.path.join(path,'out/psi'))
    
if os.path.exists(os.path.join(path,'out/phi')) is False:
    os.mkdir(os.path.join(path,'out/phi'))

search.save_psi_phi(os.path.join(path, 'out'))

Launch a search


In [25]:
#search.gpu(100, 100, 0.2, 0.4, 750, 850, 2)
search.gpu(10, 10, -0.1, 0.1, 750, 850, 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}


In [26]:
search.save_results(path+"results.txt", 0.05) 
# path, fraction of total results to save in file

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


In [27]:
top_results = search.get_results(0, 100)
# start, count

trajectory

A simple container with properties representing an object and its path


In [28]:
best = top_results[0]

In [29]:
# these numbers are wild because mask flags and search parameters above were chosen randomly
best.flux 
best.lh
best.x
best.y
best.x_v
best.y_v


Out[29]:
-82.86174011230469

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


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


Out[30]:
[lh: 941.987000 flux: 21225.396484 x: 19 y: 35 x_v: 825.853455 y_v: -82.861740 obs_count: 5,
 lh: 941.987000 flux: 21225.396484 x: 19 y: 35 x_v: 835.803528 y_v: -83.860069 obs_count: 5,
 lh: 941.987000 flux: 21225.396484 x: 19 y: 35 x_v: 827.345398 y_v: -66.329193 obs_count: 5,
 lh: 941.987000 flux: 21225.396484 x: 19 y: 35 x_v: 837.313416 y_v: -67.128342 obs_count: 5,
 lh: 941.960388 flux: 21226.257812 x: 19 y: 34 x_v: 827.345398 y_v: 66.329185 obs_count: 5,
 lh: 941.960388 flux: 21226.257812 x: 19 y: 34 x_v: 837.313416 y_v: 67.128334 obs_count: 5,
 lh: 941.794128 flux: 21220.515625 x: 19 y: 34 x_v: 838.488464 y_v: 50.369762 obs_count: 5,
 lh: 941.794128 flux: 21220.515625 x: 19 y: 34 x_v: 828.506470 y_v: 49.770123 obs_count: 5,
 lh: 931.414917 flux: 20985.867188 x: 20 y: 35 x_v: 746.253113 y_v: -74.875061 obs_count: 5,
 lh: 931.414917 flux: 20985.867188 x: 20 y: 35 x_v: 756.203186 y_v: -75.873398 obs_count: 5,
 lh: 931.414917 flux: 20985.867188 x: 20 y: 35 x_v: 766.153198 y_v: -76.871735 obs_count: 5,
 lh: 931.414917 flux: 20985.867188 x: 20 y: 35 x_v: 776.103271 y_v: -77.870064 obs_count: 5,
 lh: 931.380676 flux: 20988.455078 x: 20 y: 34 x_v: 757.569275 y_v: 60.735161 obs_count: 5,
 lh: 931.380676 flux: 20988.455078 x: 20 y: 34 x_v: 767.537292 y_v: 61.534306 obs_count: 5,
 lh: 931.380676 flux: 20988.455078 x: 20 y: 34 x_v: 777.505310 y_v: 62.333454 obs_count: 5,
 lh: 931.380676 flux: 20988.455078 x: 20 y: 34 x_v: 747.601257 y_v: 59.936012 obs_count: 5,
 lh: 754.681458 flux: 17003.689453 x: 19 y: 33 x_v: 837.313416 y_v: 67.128334 obs_count: 5,
 lh: 754.681458 flux: 17003.689453 x: 19 y: 33 x_v: 827.345398 y_v: 66.329185 obs_count: 5,
 lh: 747.403381 flux: 16838.832031 x: 20 y: 33 x_v: 747.601257 y_v: 59.936012 obs_count: 5,
 lh: 747.403381 flux: 16838.832031 x: 20 y: 33 x_v: 757.569275 y_v: 60.735161 obs_count: 5]

Demo


In [ ]: