TrajectoryExplorer Demo#

[ ]:

import math
import matplotlib.pyplot as plt

from kbmod.analysis.plotting import plot_multiple_images
from kbmod.trajectory_explorer import TrajectoryExplorer
from kbmod.search import Trajectory

# We can turn on verbose (debug) logging for this notebook by uncommenting the lines below
# import logging
# logging.basicConfig(level=logging.DEBUG)

TrajectoryExplorer#

TrajectoryExplorer is a class that allows us to evaluate individual trajectories for debugging or scientific analysis. It handles all the loading of data onto GPU, configuration of the parameters, and other set up.

Create fake data#

Let’s start by setting up a fake data set with a single known object.

[ ]:

from kbmod.fake_data.fake_data_creator import create_fake_times, FakeDataSet

# Create a list of fake times from MJD=57130.2 onward
num_images = 12
fake_times = create_fake_times(num_images, t0=57130.2)

# Create fake data set.
fake_ds = FakeDataSet(
    400,  # Width
    500,  # Height
    fake_times,
    noise_level=2.0,
    psf_val=1.0,
)

# Create a bright fake object starting at pixel x=50, y=60 and moving with velocity vx=5.0, vy=-2.0.
fake_trj = Trajectory(50, 60, 5.0, -2.0, flux=100.0)
fake_ds.insert_object(fake_trj)

# Display the image from the first time step.
img = fake_ds.stack.get_single_image(0).get_science().image
plt.imshow(img, cmap="gray")

To make things more interesting, let’s block out the pixels for timestep 6, so we cannot see the object then. We need to do this manually (there is no helper function to do this).

[ ]:

# Remove at least observation from the trajectory.
dt = fake_times[6] - fake_times[0]
pred_x = int(50 + 5.0 * dt)
pred_y = int(60 - 2.0 * dt)

sci_t6 = fake_ds.stack.get_single_image(6).get_science()
for dy in [-2, -1, 0, 1, 2]:
    for dx in [-2, -1, 0, 1, 2]:
        sci_t6.set_pixel(pred_y + dy, pred_x + dx, 0.00001)

img = fake_ds.stack.get_single_image(6).get_science().image
plt.imshow(img, cmap="gray")

Using TrajectoryExplorer#

TrajectoryExplorer is constructed from the ImageStack and an optional configuration. Once it is instantiated, we can use the evaluate_linear_trajectory function to query specific trajectories. The function returns a Results object with a single row and a variety of columns, including the raw statistics from the GPU search.

[ ]:

explorer = TrajectoryExplorer(fake_ds.stack)

result = explorer.evaluate_linear_trajectory(50, 60, 5.0, -2.0)
print(result)

We can access individual statistics about the results by using the column name and row index of 0.

[ ]:

best_row = result[0]

print(
    f"The trajectory starts at ({best_row['x']}, {best_row['y']}) and moves as ({best_row['vx']}, {best_row['vy']})"
)
print(f"It has {best_row['obs_count']} observations and a likelihood of {best_row['likelihood']}.")

We also have the row psi-curves, phi-curves, and likelihood curves:

[ ]:

fig, axs = plt.subplots(2, 1)
axs[0].plot(fake_times, best_row["psi_curve"])
axs[0].set_title("Psi")

axs[1].plot(fake_times, best_row["phi_curve"])
axs[1].set_title("Psi")

The TrajectoryExplorer also extracts and saves the different coadds (mean, sum, and median) as well as the individual stamp for each time step. We can visualize them both:

[ ]:

coadds = [best_row["coadd_sum"], best_row["coadd_mean"], best_row["coadd_median"]]
plot_multiple_images(coadds, labels=["Sum", "Mean", "Median"])

[ ]:

plot_multiple_images(best_row["all_stamps"], columns=4)

Using (RA, dec)#

You can also use the TrajectoryExplorer with a WCS and (RA, dec) coordinates. Both RA and dec are specified in degrees and the corresponding velocities are expressed as degrees per day.

We start by creating a fake WCS to match our fake images. Then we evaluate the trajectory based on this WCS.

[ ]:

from astropy.wcs import WCS

my_wcs = WCS(naxis=2)
my_wcs.wcs.crpix = [201.0, 251.0]  # Reference point on the image (center of the image 1-indexed)
my_wcs.wcs.crval = [45.0, -15.0]  # Reference pointing on the sky
my_wcs.wcs.cdelt = [0.1, 0.1]  # Pixel step size
my_wcs.wcs.ctype = ["RA---TAN-SIP", "DEC--TAN-SIP"]

# Use the WCS to compute the angular coordinates of the inserted object at two times.
sky_pos0 = my_wcs.pixel_to_world(50, 60)
sky_pos1 = my_wcs.pixel_to_world(55, 58)

ra0 = sky_pos0.ra.deg
dec0 = sky_pos0.dec.deg
v_ra = sky_pos1.ra.deg - ra0
v_dec = sky_pos1.dec.deg - dec0
print(f"Object starts at ({ra0}, {dec0}) with velocity ({v_ra}, {v_dec}).")

result = explorer.evaluate_angle_trajectory(ra0, dec0, v_ra, v_dec, my_wcs)
print(result)

Filtering#

A key component of KBMOD is the clipped Sigma-G filtering. TrajectoryExplorer does not do any filtering by default so that users can see all the information for the given trajectory. However, it provides the ability to apply the filtering manually. The mask of time step validity is stored in the column obs_valid.

[ ]:

explorer.apply_sigma_g(result)

for i in range(len(fake_times)):
    print(f"Time {i} is valid = {result['obs_valid'][0][i]}.")

Neighborhood Searches#

The TrajectoryExplorer can also be used to perform a hyper-localized search. This search effectively uses a small neighborhood around a given trajectory (both in terms of starting pixel and velocities) and returns all results from this neighborhood. This localized set can be used to:

  1. refine trajectories by searching a finer parameter space around the best results found by the initial search, or

  2. collect a distribution of trajectories and their likelihoods around a single result. For this search, the TrajectoryExplorer does not perform any filtering, so it will return all trajectories and their likelihoods (even one <= -1.0)

Only basic statistics, such as likelihood and flux, are returned from the search. Stamps and lightcurves are not computed.

[ ]:

samples = explorer.evaluate_around_linear_trajectory(50, 60, 5.0, -2.0, pixel_radius=5)
print(samples[0:10])