def setup_search(path, tfile):
with open(path, 'r') as fp:
params = json.load(fp)
print("Loading trajectories from " + tfile + "... ", end="")
t_list = kb.load_trajectories(tfile)#params['trajectories_file'])
print("Done.")
psf = kb.psf(params['psf_sigma'])
print("Generating images... ", end="", flush=True)
imgs = []
for i in range(params['img_count']):
time = i/(params['img_count']-1)
im = kb.layered_image('img'+str(i),
params['x_dim'],
params['y_dim'],
params['noise'],
params['noise']*params['noise'],
time)
for t in t_list:
im.add_object(t.x+time*t.x_v, t.y+time*t.y_v, t.flux, psf)
imgs.append(im)
stack = kb.image_stack(imgs)
del imgs
search = kb.stack_search(stack, psf)
search.set_debug(True)
del stack
print("Done.")
return search, t_list, params
def run_search(par, run_number):
par = Bunch(par)
files = os.listdir(par.path)
files.sort()
files = [par.path + f for f in files]
files = files[:par.max_images]
images = [kb.layered_image(f) for f in files]
p = kb.psf(par.psf_width)
angle_range = par.angle_range
velocity_range = par.velocity_range
results_key = []
for _ in range(par.object_count):
traj = kb.trajectory()
traj.x = int(rd.uniform(*par.x_range))
traj.y = int(rd.uniform(*par.y_range))
ang = rd.uniform(*angle_range)
vel = rd.uniform(*velocity_range)
traj.x_v = vel * math.cos(ang)
traj.y_v = vel * math.sin(ang)
traj.flux = rd.uniform(*par.flux_range)
results_key.append(traj)
results_key.extend(par.real_results)
for t in results_key:
add_trajectory(images, t, p)
stack = kb.image_stack(images)
stack.apply_mask_flags(par.flags, par.flag_exceptions)
stack.apply_master_mask(par.master_flags, 2)
search = kb.stack_search(stack, p)
search_angle_r = (angle_range[0] / par.search_margin,
angle_range[1] * par.search_margin)
search_velocity_r = (velocity_range[0] / par.search_margin,
velocity_range[1] * par.search_margin)
search.gpu(par.angle_steps, par.velocity_steps, *search_angle_r,
*search_velocity_r, par.min_observations)
#results = search.get_results(0, par.results_count)
search.save_results(
par.results_file_path + 'run{0:03d}.txt'.format(run_number + 1), 0.03)
if par.save_science:
images = stack.sciences()
for i in range(len(images)):
np.save(
par.img_save_path + 'R{0:03d}'.format(run_number + 1) +
'SCI{0:03d}.npy'.format(i + 1), images[i])
return results_key
def setUp(self):
self.im_count = 5
p = kb.psf(0.05)
self.images = []
for c in range(self.im_count):
im = kb.layered_image(str(c), 10, 10, 0.0, 1.0, c)
im.add_object(2 + c * 0.5 + 0.5, 2 + c * 0.5 + 0.5, 1, p)
self.images.append(im)
def setUp(self):
# test pass thresholds
im_count = 10
self.ix = 136
self.iy = 103
self.xv = 34.0
self.yv = 21.0
self.flux = 350.0
p = kb.psf(1.0)
imgs = []
for i in range(im_count):
im = kb.layered_image("im" + str(i + 1), 500, 500, 0.0, 100.0,
i * 0.1)
im.add_object(self.ix + 0.1 * i * self.xv,
self.iy + 0.1 * i * self.yv, self.flux, p)
imgs.append(im)
stack = kb.image_stack(imgs)
self.search = kb.stack_search(stack, p)
def setUp(self):
# test pass thresholds
im = kb.layered_image("", 171, 111, 5.0, 25.0, 0)
stack = kb.image_stack([im])
p = kb.psf(1.0)
self.search = kb.stack_search(stack, p)
max_img = im.get_science()
max_img.set_pixel(38, 39, 117)
max_img.set_pixel(24, 63, 1000)
max_img.set_pixel(50, 27, 1000)
max_img.set_pixel(80, 82, 1000)
self.pooled_max = []
while (max_img.get_ppi() > 1):
self.pooled_max.append(max_img)
max_img = max_img.pool_max()
min_img = im.get_science()
self.pooled_min = []
while (min_img.get_ppi() > 1):
self.pooled_min.append(min_img)
min_img = min_img.pool_min()
def setUp(self):
p = kb.psf(1.0)
img = kb.layered_image("test", 4, 4, 0.0, 0.0, 0.0)
stack = kb.image_stack([img])
self.search = kb.stack_search(stack, p)
def run_search(self,
im_filepath,
res_filepath,
out_suffix,
time_file,
likelihood_level=10.,
mjd_lims=None,
num_fakes=25,
rand_seed=42):
visit_nums, visit_times = np.genfromtxt(time_file, unpack=True)
image_time_dict = OrderedDict()
for visit_num, visit_time in zip(visit_nums, visit_times):
image_time_dict[str(int(visit_num))] = visit_time
chunk_size = 100000
start = time.time()
patch_visits = sorted(os.listdir(im_filepath))
patch_visit_ids = np.array(
[int(visit_name[1:7]) for visit_name in patch_visits])
patch_visit_times = np.array(
[image_time_dict[str(visit_id)] for visit_id in patch_visit_ids])
if mjd_lims is None:
use_images = patch_visit_ids
else:
visit_only = np.where(((patch_visit_times > mjd_lims[0])
& (patch_visit_times < mjd_lims[1])))[0]
print(visit_only)
use_images = patch_visit_ids[visit_only]
image_mjd = np.array(
[image_time_dict[str(visit_id)] for visit_id in use_images])
times = image_mjd - image_mjd[0]
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
hdulist = fits.open('%s/v%i-fg.fits' % (im_filepath, use_images[0]))
f0 = hdulist[0].header['FLUXMAG0']
w = WCS(hdulist[1].header)
ec_angle = self.calc_ecliptic_angle(w)
del (hdulist)
images = [
kb.layered_image('%s/v%i-fg.fits' % (im_filepath, f))
for f in np.sort(use_images)
]
print('Images Loaded')
p = kb.psf(1.4)
# Add fakes steps
print('Adding fake objects')
x_fake_range = (5, 3650)
y_fake_range = (5, 3650)
angle_range = (ec_angle - (np.pi / 15.), ec_angle + (np.pi / 15.))
velocity_range = (100, 500)
mag_range = (20, 26)
fake_results = []
fake_output = []
np.random.seed(rand_seed)
for val in range(num_fakes):
traj = kb.trajectory()
traj.x = int(np.random.uniform(*x_fake_range))
traj.y = int(np.random.uniform(*y_fake_range))
ang = np.random.uniform(*angle_range)
vel = np.random.uniform(*velocity_range)
traj.x_v = vel * np.cos(ang)
traj.y_v = vel * np.sin(ang)
mag_val = np.random.uniform(*mag_range)
traj.flux = f0 * np.power(10, -0.4 * mag_val)
fake_results.append(traj)
fake_output.append(
[traj.x, traj.y, traj.x_v, traj.y_v, traj.flux, mag_val])
for fake_obj in fake_results:
tf.add_trajectory(images, fake_obj, p, times)
stack = kb.image_stack(images)
del (images)
stack.apply_mask_flags(flags, flag_exceptions)
stack.apply_master_mask(master_flags, 2)
stack.grow_mask()
stack.grow_mask()
stack.apply_mask_threshold(120.)
stack.set_times(times)
print("Times set")
x_size = stack.get_width()
y_size = stack.get_width()
search = kb.stack_search(stack, p)
del (stack)
ang_min = ec_angle - self.ang_arr[0]
ang_max = ec_angle + self.ang_arr[1]
vel_min = self.v_arr[0]
vel_max = self.v_arr[1]
print("Starting Search")
print('---------------------------------------')
param_headers = ("Ecliptic Angle", "Min. Search Angle",
"Max Search Angle", "Min Velocity", "Max Velocity")
param_values = (ec_angle, ang_min, ang_max, vel_min, vel_max)
for header, val in zip(param_headers, param_values):
print('%s = %.4f' % (header, val))
search.gpu(int(self.ang_arr[2]), int(self.v_arr[2]), ang_min, ang_max,
vel_min, vel_max, int(self.num_obs))
keep_stamps = []
keep_new_lh = []
keep_results = []
keep_times = []
memory_error = False
keep_lc = []
likelihood_limit = False
res_num = 0
chunk_size = 500000
print('---------------------------------------')
print("Processing Results")
print('---------------------------------------')
while likelihood_limit is False:
pool = mp.Pool(processes=16)
results = search.get_results(res_num, chunk_size)
chunk_headers = ("Chunk Start", "Chunk Size",
"Chunk Max Likelihood", "Chunk Min. Likelihood")
chunk_values = (res_num, len(keep_results), results[0].lh,
results[-1].lh)
for header, val, in zip(chunk_headers, chunk_values):
if type(val) == np.int:
print('%s = %i' % (header, val))
else:
print('%s = %.2f' % (header, val))
print('---------------------------------------')
psi_curves = []
phi_curves = []
for line in results:
psi_curve, phi_curve = search.lightcurve(line)
psi_curves.append(np.array(psi_curve).flatten())
phi_curve = np.array(phi_curve).flatten()
phi_curve[phi_curve == 0.] = 99999999.
phi_curves.append(phi_curve)
if line.lh < likelihood_level:
likelihood_limit = True
break
keep_idx_results = pool.starmap_async(
return_indices,
zip(psi_curves, phi_curves,
[j for j in range(len(psi_curves))]))
pool.close()
pool.join()
keep_idx_results = keep_idx_results.get()
if len(keep_idx_results[0]) < 3:
keep_idx_results = [(0, [-1], 0.)]
for result_on in range(len(psi_curves)):
if keep_idx_results[result_on][1][0] == -1:
continue
elif len(keep_idx_results[result_on][1]) < 3:
continue
elif keep_idx_results[result_on][2] < likelihood_level:
continue
else:
keep_idx = keep_idx_results[result_on][1]
new_likelihood = keep_idx_results[result_on][2]
keep_results.append(results[result_on])
keep_new_lh.append(new_likelihood)
stamps = search.sci_stamps(results[result_on], 10)
stamp_arr = np.array(
[np.array(stamps[s_idx]) for s_idx in keep_idx])
keep_stamps.append(np.sum(stamp_arr, axis=0))
keep_lc.append((psi_curves[result_on] /
phi_curves[result_on])[keep_idx])
#keep_times.append(image_mjd[keep_idx])
keep_times.append(keep_idx)
# if len(keep_results) > 800000:
# with open('%s/memory_error_tr_%s.txt' %
# (res_filepath, out_suffix), 'w') as f:
# f.write('In %i total results, %i were kept. Needs manual look.' %
# (res_num + chunk_size, len(keep_results)))
# memory_error = True
# likelihood_limit = True
# if res_num+chunk_size >= 8000000:
# likelihood_level = 20.
# with open('%s/overload_error_tr_%s.txt' %
# (res_filepath, out_suffix), 'w') as f:
# f.write('In %i total results, %i were kept. Likelihood level down to %f.' %
# (res_num + chunk_size, len(keep_results), line.lh))
res_num += chunk_size
del (search)
lh_sorted_idx = np.argsort(np.array(keep_new_lh))[::-1]
if len(lh_sorted_idx) > 0:
print("Stamp filtering %i results" % len(lh_sorted_idx))
pool = mp.Pool(processes=16)
stamp_filt_pool = pool.map_async(
stamp_filter_parallel,
np.array(keep_stamps)[lh_sorted_idx])
pool.close()
pool.join()
stamp_filt_results = stamp_filt_pool.get()
stamp_filt_idx = lh_sorted_idx[np.where(
np.array(stamp_filt_results) == 1)]
if len(stamp_filt_idx) > 0:
print("Clustering %i results" % len(stamp_filt_idx))
cluster_idx = self.cluster_results(
np.array(keep_results)[stamp_filt_idx], x_size, y_size,
[vel_min, vel_max], [ang_min, ang_max])
final_results = stamp_filt_idx[cluster_idx]
else:
cluster_idx = []
final_results = []
del (cluster_idx)
del (stamp_filt_results)
del (stamp_filt_idx)
del (stamp_filt_pool)
else:
final_results = lh_sorted_idx
print('Keeping %i results' % len(final_results))
np.savetxt('%s/results_%s.txt' % (res_filepath, out_suffix),
np.array(keep_results)[final_results],
fmt='%s')
np.savetxt('%s/results_fakes_%s.txt' % (res_filepath, out_suffix),
np.array(fake_output),
header='x,y,xv,yv,flux,mag')
# np.savetxt('%s/lc_%s.txt' % (res_filepath, out_suffix),
# np.array(keep_lc)[final_results], fmt='%s')
with open('%s/lc_%s.txt' % (res_filepath, out_suffix), 'w') as f:
writer = csv.writer(f)
writer.writerows(np.array(keep_lc)[final_results])
# np.savetxt('%s/times_%s.txt' % (res_filepath, out_suffix),
# np.array(keep_times)[final_results], fmt='%s')
with open('%s/times_%s.txt' % (res_filepath, out_suffix), 'w') as f:
writer = csv.writer(f)
writer.writerows(np.array(keep_times)[final_results])
np.savetxt('%s/filtered_likes_%s.txt' % (res_filepath, out_suffix),
np.array(keep_new_lh)[final_results],
fmt='%.4f')
np.savetxt('%s/ps_%s.txt' % (res_filepath, out_suffix),
np.array(keep_stamps).reshape(len(keep_stamps),
441)[final_results],
fmt='%.4f')
end = time.time()
del (keep_stamps)
del (keep_times)
del (keep_results)
del (keep_new_lh)
del (keep_lc)
print("Time taken for patch: ", end - start)
def run_search(self):
"""
This function serves as the highest-level python interface for starting
a KBMOD search.
INPUT-
im_filepath : string
Path to the folder containing the images to be ingested into
KBMOD and searched over.
res_filepath : string
Path to the folder that will contain the results from the
search.
out_suffix : string
Suffix to append to the output files. Used to differentiate
between different searches over the same stack of images.
time_file : string
Path to the file containing the image times.
lh_level : float
Minimum acceptable likelihood level for a trajectory.
Trajectories with likelihoods below this value will be
discarded.
psf_val : float
Determines the size of the psf generated by the kbmod stack.
mjd_lims : numpy array
Limits the search to images taken within the limits input by
mjd_lims.
average_angle : float
Overrides the ecliptic angle calculation and instead centers
the average search around average_angle.
"""
start = time.time()
kb_interface = Interface()
kb_post_process = PostProcess(self.config)
# Load images to search
stack, image_params = kb_interface.load_images(
self.config['im_filepath'], self.config['time_file'],
self.config['mjd_lims'], self.config['visit_in_filename'],
self.config['file_format'])
# Save values in image_params for later use
if self.config['do_mask']:
stack = kb_post_process.apply_mask(
stack,
mask_num_images=self.config['mask_num_images'],
mask_threshold=self.config['mask_threshold'])
psf = kb.psf(self.config['psf_val'])
search = kb.stack_search(stack, psf)
search, image_params = self.do_gpu_search(search, image_params)
# Load the KBMOD results into Python and apply a filter based on
# 'filter_type'
keep = kb_post_process.load_results(
search,
image_params,
self.config['lh_level'],
chunk_size=self.config['chunk_size'],
filter_type=self.config['filter_type'],
max_lh=self.config['max_lh'])
keep = kb_post_process.get_coadd_stamps(keep, search)
keep = kb_post_process.apply_stamp_filter(
keep,
center_thresh=self.config['center_thresh'],
peak_offset=self.config['peak_offset'],
mom_lims=self.config['mom_lims'])
keep = kb_post_process.apply_clustering(keep, image_params)
keep = kb_post_process.get_all_stamps(keep, search)
del (search)
# Save the results
kb_interface.save_results(self.config['res_filepath'],
self.config['output_suffix'], keep)
end = time.time()
del (keep)
print("Time taken for patch: ", end - start)
def run_search(par):
par = Bunch(par)
files = os.listdir(par.path)
files.sort()
files = [par.path+f for f in files]
files = files[:par.max_images]
images = [kb.layered_image(f) for f in files]
p = kb.psf(par.psf_width)
angle_range = par.angle_range
velocity_range = par.velocity_range
results_key = []
for _ in range(par.object_count):
traj = kb.trajectory()
traj.x = int(rd.uniform(*par.x_range))
traj.y = int(rd.uniform(*par.y_range))
ang = rd.uniform(*angle_range)
vel = rd.uniform(*velocity_range)
traj.x_v = vel*math.cos(ang)
traj.y_v = vel*math.sin(ang)
traj.flux = rd.uniform(*par.flux_range)
results_key.append(traj)
results_key.extend(par.real_results)
for t in results_key:
add_trajectory(images, t, p)
stack = kb.image_stack(images)
stack.apply_mask_flags(par.flags, par.flag_exceptions)
stack.apply_master_mask(par.master_flags, 2)
search = kb.stack_search(stack, p)
search_angle_r = (angle_range[0]/par.search_margin,
angle_range[1]*par.search_margin)
search_velocity_r = (velocity_range[0]/par.search_margin,
velocity_range[1]*par.search_margin)
search.gpu(par.angle_steps, par.velocity_steps,
*search_angle_r, *search_velocity_r, par.min_observations)
results = search.get_results(0, par.results_count)
results_clustered = [ results[i] for i in
cluster_trajectories(results,
dbscan_args=dict(eps=par.cluster_eps,
n_jobs=-1, min_samples=1))[1] ]
results_matched, results_unmatched = \
match_trajectories(results_clustered,
results_key, par.match_v, par.match_coord)
results_to_plot = results_unmatched
images = [i.science() for i in stack.get_images()]
stamps = [create_postage_stamp(images, t, stack.get_times(), [21,21])[0] \
for t in results_to_plot]
return results_matched, results_unmatched, stamps
def load_images(self, im_filepath, time_file, psf_val=1.4, mjd_lims=None):
"""
This function loads images and ingests them into a search object
Input
---------
im_filepath : string
Image file path from which to load images
time_file : string
File name containing image times
Output
---------
search : kb.stack_search object
image_params : dictionary
Contains image parameters such as ecliptic angle and mean Julian day
"""
# Empty for now. Will contain x_size, y_size, ec_angle, and mjd before being returned.
image_params = {}
visit_nums, visit_times = np.genfromtxt(time_file, unpack=True)
image_time_dict = OrderedDict()
for visit_num, visit_time in zip(visit_nums, visit_times):
image_time_dict[str(int(visit_num))] = visit_time
chunk_size = 100000
patch_visits = sorted(os.listdir(im_filepath))
patch_visit_ids = self.get_folder_visits(patch_visits)
patch_visit_times = np.array(
[image_time_dict[str(visit_id)] for visit_id in patch_visit_ids])
if mjd_lims is None:
use_images = patch_visit_ids
else:
visit_only = np.where(((patch_visit_times > mjd_lims[0])
& (patch_visit_times < mjd_lims[1])))[0]
print(visit_only)
use_images = patch_visit_ids[visit_only]
image_params['mjd'] = np.array(
[image_time_dict[str(visit_id)] for visit_id in use_images])
times = image_params['mjd'] - image_params['mjd'][0]
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
hdulist = fits.open('%s/%s' %
(im_filepath, self.return_filename(use_images[0])))
w = WCS(hdulist[1].header)
image_params['ec_angle'] = self.calc_ecliptic_angle(w)
del (hdulist)
images = [
kb.layered_image('%s/%s' % (im_filepath, self.return_filename(f)))
for f in np.sort(use_images)
]
print('Loaded %i images' % (len(images)))
p = kb.psf(psf_val)
stack = kb.image_stack(images)
# Apply masks
stack.apply_mask_flags(flags, flag_exceptions)
stack.apply_master_mask(master_flags, 2)
stack.grow_mask()
stack.grow_mask()
stack.apply_mask_threshold(120.)
stack.set_times(times)
print("Times set")
image_params['x_size'] = stack.get_width()
image_params['y_size'] = stack.get_width()
search = kb.stack_search(stack, p)
return (search, image_params)
image_folder = sys.argv[1]
results_file = sys.argv[2]
# Following sets up ability to create psi/phi and is from
# HITS_Main_Belt_Example.ipynb
flags = ~0
flag_exceptions = [32, 39]
master_flags = int('100111', 2)
image_file_list = [
str(image_folder + '/' + filename)
for filename in os.listdir(image_folder)
]
image_file_list.sort()
images = [kb.layered_image(f) for f in image_file_list]
p = kb.psf(1.4)
stack = kb.image_stack(images)
stack.apply_mask_flags(flags, flag_exceptions)
stack.apply_master_mask(master_flags, 2)
image_array = stack.get_images()
search = kb.stack_search(stack, p)
search.gpu(1, 2, -0.0442959674533, 0.741102195944, 1920.0, 4032.0, 3)
psi = search.get_psi()
phi = search.get_phi()
image_times = np.array(stack.get_times())
file_len, header_len = file_and_header_length(results_file)
