diff --git a/src/kbmod/search/kernels.cu b/src/kbmod/search/kernels.cu index 50c497754..a2eb45950 100644 --- a/src/kbmod/search/kernels.cu +++ b/src/kbmod/search/kernels.cu @@ -24,6 +24,9 @@ namespace search { extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray *data) { + if (data == nullptr) { + throw std::runtime_error("No data given."); + } if (!data->cpu_array_allocated() || !data->cpu_time_array_allocated()) { throw std::runtime_error("CPU data is not allocated."); } @@ -48,6 +51,9 @@ extern "C" void device_allocate_psi_phi_arrays(PsiPhiArray *data) { } extern "C" void device_free_psi_phi_arrays(PsiPhiArray *data) { + if (data == nullptr) { + throw std::runtime_error("No data given."); + } if (data->gpu_array_allocated()) { checkCudaErrors(cudaFree(data->get_gpu_array_ptr())); data->set_gpu_array_ptr(nullptr); @@ -61,7 +67,8 @@ extern "C" void device_free_psi_phi_arrays(PsiPhiArray *data) { __host__ __device__ PsiPhi read_encoded_psi_phi(PsiPhiArrayMeta ¶ms, void *psi_phi_vect, int time, int row, int col) { // Bounds checking. - if ((row < 0) || (col < 0) || (row >= params.height) || (col >= params.width)) { + if ((row < 0) || (col < 0) || (row >= params.height) || (col >= params.width) || + (psi_phi_vect == nullptr)) { return {NO_DATA, NO_DATA}; } @@ -92,6 +99,9 @@ extern "C" __device__ __host__ void SigmaGFilteredIndicesCU(float *values, int n float sgl1, float sigmag_coeff, float width, int *idx_array, int *min_keep_idx, int *max_keep_idx) { + // Basic data checking. + assert(idx_array != nullptr && min_keep_idx != nullptr && max_keep_idx != nullptr); + // Clip the percentiles to [0.01, 99.99] to avoid invalid array accesses. if (sgl0 < 0.0001) sgl0 = 0.0001; if (sgl1 > 0.9999) sgl1 = 0.9999; @@ -147,6 +157,9 @@ extern "C" __device__ __host__ void SigmaGFilteredIndicesCU(float *values, int n extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_vect, float *image_times, SearchParameters params, Trajectory *candidate) { + // Basic data validity check. + assert(psi_phi_vect != nullptr && image_times != nullptr && candidate != nullptr); + // Data structures used for filtering. We fill in only what we need. float psi_array[MAX_NUM_IMAGES]; float phi_array[MAX_NUM_IMAGES]; @@ -226,6 +239,10 @@ extern "C" __device__ __host__ void evaluateTrajectory(PsiPhiArrayMeta psi_phi_m __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_vect, float *image_times, SearchParameters params, int num_trajectories, Trajectory *trajectories, Trajectory *results) { + // Basic data validity check. + assert(psi_phi_vect != nullptr && image_times != nullptr && trajectories != nullptr && + results != nullptr); + // Get the x and y coordinates within the search space. const int x_i = blockIdx.x * THREAD_DIM_X + threadIdx.x; const int y_i = blockIdx.y * THREAD_DIM_Y + threadIdx.y; @@ -294,6 +311,9 @@ __global__ void searchFilterImages(PsiPhiArrayMeta psi_phi_meta, void *psi_phi_v extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, SearchParameters params, int num_trajectories, Trajectory *trj_to_search, int num_results, Trajectory *best_results) { + // Basic data validity check. + assert(trj_to_search != nullptr && best_results != nullptr); + // Allocate Device memory Trajectory *device_tests; Trajectory *device_search_results; diff --git a/src/kbmod/search/pydocs/stack_search_docs.h b/src/kbmod/search/pydocs/stack_search_docs.h index f52e2aece..1af2936d1 100644 --- a/src/kbmod/search/pydocs/stack_search_docs.h +++ b/src/kbmod/search/pydocs/stack_search_docs.h @@ -10,6 +10,25 @@ static const auto DOC_StackSearch_search = R"doc( todo )doc"; +static const auto DOC_StackSearch_set_min_obs = R"doc( + Sets the minimum number of observations for valid result. + + Parameters + ---------- + new_value : `int` + The minimum number of observations for a trajectory to be returned. + )doc"; + +static const auto DOC_StackSearch_set_min_lh = R"doc( + Sets the minimum likelihood for valid result. + + Parameters + ---------- + new_value : `float` + The minimum likelihood value for a trajectory to be returned. + )doc"; + + static const auto DOC_StackSearch_enable_gpu_sigmag_filter = R"doc( todo )doc"; @@ -111,6 +130,10 @@ static const auto DOC_StackSearch_get_phi_curves = R"doc( The phi values at each time step with NO_DATA replaced by 0.0. )doc"; +static const auto DOC_StackSearch_clear_psi_phi = R"doc( + Clear the pre-computed psi and phi data. + )doc"; + static const auto DOC_StackSearch_prepare_psi_phi = R"doc( Compute the cached psi and phi data. )doc"; @@ -123,5 +146,43 @@ static const auto DOC_StackSearch_set_results = R"doc( todo )doc"; +static const auto DOC_StackSearch_evaluate_single_trajectory = R"doc( + Performs the evaluation of a single Trajectory object. Modifies the object + in-place. + + Note + ---- + Runs on the CPU, but requires CUDA compiler. + + Parameters + ---------- + trj : `kb.Trajectory` + The trjactory to evaluate. + )doc"; + +static const auto DOC_StackSearch_search_linear_trajectory = R"doc( + Performs the evaluation of a linear trajectory in pixel space. + + Note + ---- + Runs on the CPU, but requires CUDA compiler. + + Parameters + ---------- + x : `short` + The starting x pixel of the trajectory. + y : `short` + The starting y pixel of the trajectory. + vx : `float` + The x velocity of the trajectory in pixels per day. + vy : `float` + The y velocity of the trajectory in pixels per day. + + Returns + ------- + result : `kb.Trajectory` + The trajectory object with statistics set. + )doc"; + } // namespace pydocs #endif /* STACKSEARCH_DOCS */ diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index cf06968a6..3e81e107f 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -4,6 +4,9 @@ namespace search { #ifdef HAVE_CUDA extern "C" void deviceSearchFilter(PsiPhiArray& psi_phi_array, SearchParameters params, int num_trajectories, Trajectory* trj_to_search, int num_results, Trajectory* best_results); + +extern "C" void evaluateTrajectory(PsiPhiArrayMeta psi_phi_meta, void* psi_phi_vect, float* image_times, + SearchParameters params, Trajectory* candidate); #endif StackSearch::StackSearch(ImageStack& imstack) : stack(imstack) { @@ -32,11 +35,19 @@ StackSearch::StackSearch(ImageStack& imstack) : stack(imstack) { params.debug = false; } +// -------------------------------------------- +// Configuration functions +// -------------------------------------------- + void StackSearch::set_debug(bool d) { debug_info = d; params.debug = d; } +void StackSearch::set_min_obs(int new_value) { params.min_observations = new_value; } + +void StackSearch::set_min_lh(float new_value) { params.min_lh = new_value; } + void StackSearch::enable_gpu_sigmag_filter(std::vector percentiles, float sigmag_coeff, float min_lh) { params.do_sigmag_filter = true; params.sgl_L = percentiles[0]; @@ -46,6 +57,11 @@ void StackSearch::enable_gpu_sigmag_filter(std::vector percentiles, float } void StackSearch::enable_gpu_encoding(int encode_num_bytes) { + // If changing a setting that would impact the search data encoding, clear the cached values. + if (params.encode_num_bytes != encode_num_bytes) { + clear_psi_phi(); + } + // Make sure the encoding is one of the supported options. // Otherwise use default float (aka no encoding). if (encode_num_bytes == 1 || encode_num_bytes == 2) { @@ -65,19 +81,67 @@ void StackSearch::set_start_bounds_y(int y_min, int y_max) { params.y_start_max = y_max; } +// -------------------------------------------- +// Data precomputation functions +// -------------------------------------------- + +void StackSearch::prepare_psi_phi() { + if (!psi_phi_generated) { + DebugTimer timer = DebugTimer("Preparing Psi and Phi images", debug_info); + fill_psi_phi_array_from_image_stack(psi_phi_array, stack, params.encode_num_bytes, debug_info); + timer.stop(); + psi_phi_generated = true; + } + + // Perform additional error checking that the arrays are allocated (checked even if + // using the cached values). + if (!psi_phi_array.cpu_array_allocated() || !psi_phi_array.cpu_time_array_allocated()) { + throw std::runtime_error("PsiPhiArray arrays unallocated after prepare_psi_phi_array."); + } +} + +void StackSearch::clear_psi_phi() { + if (psi_phi_generated) { + psi_phi_array.clear(); + psi_phi_generated = false; + } +} + +// -------------------------------------------- +// Core search functions +// -------------------------------------------- + +void StackSearch::evaluate_single_trajectory(Trajectory& trj) { + prepare_psi_phi(); + if (!psi_phi_array.cpu_array_allocated()) std::runtime_error("Data not allocated."); + +#ifdef HAVE_CUDA + evaluateTrajectory(psi_phi_array.get_meta_data(), psi_phi_array.get_cpu_array_ptr(), + psi_phi_array.get_cpu_time_array_ptr(), params, &trj); +#else + throw std::runtime_error("CUDA installation is needed for single trajectory search."); +#endif +} + +Trajectory StackSearch::search_linear_trajectory(short x, short y, float vx, float vy) { + Trajectory result; + result.x = x; + result.y = y; + result.vx = vx; + result.vy = vy; + + evaluate_single_trajectory(result); + + return result; +} + void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ang, float min_vel, float mavx, int min_observations) { DebugTimer core_timer = DebugTimer("Running core search", debug_info); create_search_list(ang_steps, vel_steps, min_ang, max_ang, min_vel, mavx); - // Create a data stucture for the per-image data. - std::vector image_times = stack.build_zeroed_times(); - DebugTimer psi_phi_timer = DebugTimer("Creating psi/phi buffers", debug_info); prepare_psi_phi(); - PsiPhiArray psi_phi_data; - fill_psi_phi_array(psi_phi_data, params.encode_num_bytes, psi_images, phi_images, image_times, - debug_info); psi_phi_timer.stop(); // Allocate a vector for the results. @@ -98,7 +162,7 @@ void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ // Do the actual search on the GPU. DebugTimer search_timer = DebugTimer("Running search", debug_info); #ifdef HAVE_CUDA - deviceSearchFilter(psi_phi_data, params, search_list.size(), search_list.data(), max_results, + deviceSearchFilter(psi_phi_array, params, search_list.size(), search_list.data(), max_results, results.data()); #else throw std::runtime_error("Non-GPU search is not implemented."); @@ -111,27 +175,6 @@ void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ core_timer.stop(); } -void StackSearch::prepare_psi_phi() { - if (!psi_phi_generated) { - DebugTimer timer = DebugTimer("Preparing Psi and Phi images", debug_info); - psi_images.clear(); - phi_images.clear(); - - // Compute Phi and Psi from convolved images - // while leaving masked pixels alone - // Reinsert 0s for NO_DATA? - const int num_images = stack.img_count(); - for (int i = 0; i < num_images; ++i) { - LayeredImage& img = stack.get_single_image(i); - psi_images.push_back(img.generate_psi_image()); - phi_images.push_back(img.generate_phi_image()); - } - - psi_phi_generated = true; - timer.stop(); - } -} - void StackSearch::create_search_list(int angle_steps, int velocity_steps, float min_ang, float max_ang, float min_vel, float mavx) { DebugTimer timer = DebugTimer("Creating search candidate list", debug_info); @@ -159,54 +202,34 @@ void StackSearch::create_search_list(int angle_steps, int velocity_steps, float timer.stop(); } -std::vector StackSearch::create_curves(Trajectory t, const std::vector& imgs) { - /*Create a lightcurve from an image along a trajectory - * - * INPUT- - * Trajectory t - The trajectory along which to compute the lightcurve - * std::vector imgs - The image from which to compute the - * trajectory. Most likely a psiImage or a phiImage. - * Output- - * std::vector lightcurve - The computed trajectory - */ - - int img_size = imgs.size(); - std::vector lightcurve; - lightcurve.reserve(img_size); - for (int i = 0; i < img_size; ++i) { - /* Do not use get_pixel_interp(), because results from create_curves must - * be able to recover the same likelihoods as the ones reported by the - * gpu search. Shift by 0.5 pixels to center as done on GPU. */ - float time = stack.get_zeroed_time(i); - Point p{t.get_x_pos(time) + 0.5f, t.get_y_pos(time) + 0.5f}; - - float pix_val = imgs[i].get_pixel(p.to_index()); - if (pix_val == NO_DATA) pix_val = 0.0; - lightcurve.push_back(pix_val); +std::vector StackSearch::extract_psi_or_phi_curve(Trajectory& trj, bool extract_psi) { + prepare_psi_phi(); + + const int num_times = stack.img_count(); + std::vector result(num_times, 0.0); + + for (int i = 0; i < num_times; ++i) { + float time = psi_phi_array.read_time(i); + + // Query the center of the predicted location's pixel. + Point pred_pt = {trj.get_x_pos(time) + 0.5f, trj.get_y_pos(time) + 0.5f}; + Index pred_idx = pred_pt.to_index(); + PsiPhi psi_phi_val = psi_phi_array.read_psi_phi(i, pred_idx.i, pred_idx.j); + + float value = (extract_psi) ? psi_phi_val.psi : psi_phi_val.phi; + if (value != NO_DATA) { + result[i] = value; + } } - return lightcurve; + return result; } -std::vector StackSearch::get_psi_curves(Trajectory& t) { - /*Generate a psi lightcurve for further analysis - * INPUT- - * Trajectory& t - The trajectory along which to find the lightcurve - * OUTPUT- - * std::vector - A vector of the lightcurve values - */ - prepare_psi_phi(); - return create_curves(t, psi_images); +std::vector StackSearch::get_psi_curves(Trajectory& trj) { + return extract_psi_or_phi_curve(trj, true); } -std::vector StackSearch::get_phi_curves(Trajectory& t) { - /*Generate a phi lightcurve for further analysis - * INPUT- - * Trajectory& t - The trajectory along which to find the lightcurve - * OUTPUT- - * std::vector - A vector of the lightcurve values - */ - prepare_psi_phi(); - return create_curves(t, phi_images); +std::vector StackSearch::get_phi_curves(Trajectory& trj) { + return extract_psi_or_phi_curve(trj, false); } void StackSearch::sort_results() { @@ -250,6 +273,12 @@ static void stack_search_bindings(py::module& m) { py::class_(m, "StackSearch", pydocs::DOC_StackSearch) .def(py::init()) .def("search", &ks::search, pydocs::DOC_StackSearch_search) + .def("evaluate_single_trajectory", &ks::evaluate_single_trajectory, + pydocs::DOC_StackSearch_evaluate_single_trajectory) + .def("search_linear_trajectory", &ks::search_linear_trajectory, + pydocs::DOC_StackSearch_search_linear_trajectory) + .def("set_min_obs", &ks::set_min_obs, pydocs::DOC_StackSearch_set_min_obs) + .def("set_min_lh", &ks::set_min_lh, pydocs::DOC_StackSearch_set_min_lh) .def("enable_gpu_sigmag_filter", &ks::enable_gpu_sigmag_filter, pydocs::DOC_StackSearch_enable_gpu_sigmag_filter) .def("enable_gpu_encoding", &ks::enable_gpu_encoding, pydocs::DOC_StackSearch_enable_gpu_encoding) @@ -269,6 +298,7 @@ static void stack_search_bindings(py::module& m) { .def("get_phi_curves", (std::vector(ks::*)(tj&)) & ks::get_phi_curves, pydocs::DOC_StackSearch_get_phi_curves) .def("prepare_psi_phi", &ks::prepare_psi_phi, pydocs::DOC_StackSearch_prepare_psi_phi) + .def("clear_psi_phi", &ks::clear_psi_phi, pydocs::DOC_StackSearch_clear_psi_phi) .def("get_results", &ks::get_results, pydocs::DOC_StackSearch_get_results) .def("set_results", &ks::set_results, pydocs::DOC_StackSearch_set_results); } diff --git a/src/kbmod/search/stack_search.h b/src/kbmod/search/stack_search.h index 48f3386d9..f7d2c7771 100644 --- a/src/kbmod/search/stack_search.h +++ b/src/kbmod/search/stack_search.h @@ -35,55 +35,62 @@ class StackSearch { int get_image_npixels() const { return stack.get_npixels(); } const ImageStack& get_imagestack() const { return stack; } + // Parameter setters used to control the searches. void set_debug(bool d); - - // The primary search functions. + void set_min_obs(int new_value); + void set_min_lh(float new_value); void enable_gpu_sigmag_filter(std::vector percentiles, float sigmag_coeff, float min_lh); void enable_gpu_encoding(int num_bytes); - void set_start_bounds_x(int x_min, int x_max); void set_start_bounds_y(int y_min, int y_max); + // The primary search functions + void evaluate_single_trajectory(Trajectory& trj); + Trajectory search_linear_trajectory(short x, short y, float vx, float vy); void search(int a_steps, int v_steps, float min_angle, float max_angle, float min_velocity, float max_velocity, int min_observations); - // Gets the vector of result trajectories. + // Gets the vector of result trajectories from the grid search. std::vector get_results(int start, int end); - // Filters the results based on various parameters. + // Filters the results of the grid search based on various parameters void filter_results(int min_observations); void filter_results_lh(float min_lh); - // Getters for the Psi and Phi data. + // Getters for the Psi and Phi data std::vector get_psi_curves(Trajectory& t); std::vector get_phi_curves(Trajectory& t); - // Helper functions for computing Psi and Phi. + // Helper functions for computing Psi and Phi void prepare_psi_phi(); + void clear_psi_phi(); - // Helper functions for testing. + // Helper functions for testing void set_results(const std::vector& new_results); virtual ~StackSearch(){}; protected: void sort_results(); - std::vector create_curves(Trajectory t, const std::vector& imgs); // Creates list of trajectories to search. void create_search_list(int angle_steps, int velocity_steps, float min_ang, float max_ang, float min_vel, float max_vel); - bool psi_phi_generated; - bool debug_info; + std::vector extract_psi_or_phi_curve(Trajectory& trj, bool extract_psi); + + // Core data and search parameters ImageStack stack; + SearchParameters params; + bool debug_info; + + // Precomputed and cached search data + bool psi_phi_generated; + PsiPhiArray psi_phi_array; + + // Cached data for grid search. TODO: see if we can remove this. std::vector search_list; - std::vector psi_images; - std::vector phi_images; std::vector results; - - // Parameters for the GPU search. - SearchParameters params; }; } /* namespace search */ diff --git a/tests/test_search.py b/tests/test_search.py index 33c06cc82..48fbb87ae 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -4,6 +4,7 @@ from kbmod.fake_data_creator import add_fake_object from kbmod.search import * +from kbmod.trajectory_utils import make_trajectory class test_search(unittest.TestCase): @@ -29,11 +30,12 @@ def setUp(self): self.vyel = 16.0 # create a Trajectory for the object - self.trj = Trajectory() - self.trj.x = self.start_x - self.trj.y = self.start_y - self.trj.vx = self.vxel - self.trj.vy = self.vyel + self.trj = make_trajectory( + x=self.start_x, + y=self.start_y, + vx=self.vxel, + vy=self.vyel, + ) # Add convenience array of all true bools for the stamp tests. self.all_valid = [True] * self.imCount @@ -87,6 +89,35 @@ def setUp(self): self.params.m02_limit = 35.5 self.params.m20_limit = 35.5 + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") + def test_evaluate_single_trajectory(self): + test_trj = make_trajectory( + x=self.start_x, + y=self.start_y, + vx=self.vxel, + vy=self.vyel, + ) + self.search.evaluate_single_trajectory(test_trj) + + # We found a valid result. + self.assertGreater(test_trj.obs_count, 0) + self.assertGreater(test_trj.flux, 0.0) + self.assertGreater(test_trj.lh, 0.0) + + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") + def test_search_linear_trajectory(self): + test_trj = self.search.search_linear_trajectory( + self.start_x, + self.start_y, + self.vxel, + self.vyel, + ) + + # We found a valid result. + self.assertGreater(test_trj.obs_count, 0) + self.assertGreater(test_trj.flux, 0.0) + self.assertGreater(test_trj.lh, 0.0) + @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_results(self): self.search.search( @@ -341,11 +372,7 @@ def test_mean_stamps_trj(self): def test_mean_stamps_no_data(self): # Create a Trajectory that goes through the masked pixels. - trj = Trajectory() - trj.x = self.masked_x - trj.y = self.masked_y - trj.vx = 0 - trj.vy = 0 + trj = make_trajectory(x=self.masked_x, y=self.masked_y, vx=0.0, vy=0.0) # Compute the stacked science from a single Trajectory meanStamp = StampCreator.get_mean_stamp(self.search.get_imagestack(), trj, 2, self.all_valid) @@ -747,25 +774,23 @@ def test_coadd_gpu_use_inds(self): def test_coadd_filter_cpu(self): # Create a second Trajectory that isn't any good. - trj2 = Trajectory() - trj2.x = 1 - trj2.y = 1 - trj2.vx = 0 - trj2.vy = 0 + trj2 = make_trajectory(x=1, y=1, vx=0.0, vy=0.0) # Create a third Trajectory that is close to good, but offset. - trj3 = Trajectory() - trj3.x = self.trj.x + 2 - trj3.y = self.trj.y + 2 - trj3.vx = self.trj.vx - trj3.vy = self.trj.vy + trj3 = make_trajectory( + x=self.trj.x + 2, + y=self.trj.y + 2, + vx=self.trj.vx, + vy=self.trj.vy, + ) # Create a fourth Trajectory that is close enough - trj4 = Trajectory() - trj4.x = self.trj.x + 1 - trj4.y = self.trj.y + 1 - trj4.vx = self.trj.vx - trj4.vy = self.trj.vy + trj4 = make_trajectory( + x=self.trj.x + 1, + y=self.trj.y + 1, + vx=self.trj.vx, + vy=self.trj.vy, + ) # Compute the stacked science from a single Trajectory. all_valid_vect = [(self.all_valid) for i in range(4)] @@ -788,25 +813,23 @@ def test_coadd_filter_cpu(self): @unittest.skipIf(not HAS_GPU, "Skipping test (no GPU detected)") def test_coadd_filter_gpu(self): # Create a second Trajectory that isn't any good. - trj2 = Trajectory() - trj2.x = 1 - trj2.y = 1 - trj2.vx = 0 - trj2.vy = 0 + trj2 = make_trajectory(x=1, y=1, vx=0.0, vy=0.0) # Create a third Trajectory that is close to good, but offset. - trj3 = Trajectory() - trj3.x = self.trj.x + 2 - trj3.y = self.trj.y + 2 - trj3.vx = self.trj.vx - trj3.vy = self.trj.vy + trj3 = make_trajectory( + x=self.trj.x + 2, + y=self.trj.y + 2, + vx=self.trj.vx, + vy=self.trj.vy, + ) # Create a fourth Trajectory that is close enough - trj4 = Trajectory() - trj4.x = self.trj.x + 1 - trj4.y = self.trj.y + 1 - trj4.vx = self.trj.vx - trj4.vy = self.trj.vy + trj4 = make_trajectory( + x=self.trj.x + 1, + y=self.trj.y + 1, + vx=self.trj.vx, + vy=self.trj.vy, + ) # Compute the stacked science from a single Trajectory. all_valid_vect = [(self.all_valid) for i in range(4)]