diff --git a/src/kbmod/search/psi_phi_array.cpp b/src/kbmod/search/psi_phi_array.cpp index 6b99e6523..5cc71fb11 100644 --- a/src/kbmod/search/psi_phi_array.cpp +++ b/src/kbmod/search/psi_phi_array.cpp @@ -177,6 +177,9 @@ void set_encode_cpu_psi_phi_array(PsiPhiArray& data, const std::vector throw std::runtime_error("CPU PsiPhi already allocated."); } T* encoded = (T*)malloc(data.get_total_array_size()); + if (encoded == nullptr) { + throw std::runtime_error("Unable to allocate space for CPU PsiPhi array."); + } // Create a safe maximum that is slightly less than the true max to avoid // rollover of the unsigned integer. @@ -214,6 +217,9 @@ void set_float_cpu_psi_phi_array(PsiPhiArray& data, const std::vector& throw std::runtime_error("CPU PsiPhi already allocated."); } float* encoded = (float*)malloc(data.get_total_array_size()); + if (encoded == nullptr) { + throw std::runtime_error("Unable to allocate space for CPU PsiPhi array."); + } int current_index = 0; for (int t = 0; t < data.get_num_times(); ++t) { @@ -229,7 +235,7 @@ void set_float_cpu_psi_phi_array(PsiPhiArray& data, const std::vector& } void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector& psi_imgs, - const std::vector& phi_imgs) { + const std::vector& phi_imgs, bool debug) { if (result_data.get_cpu_array_ptr() != nullptr) { return; } @@ -253,6 +259,15 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect compute_scale_params_from_image_vect(phi_imgs, result_data.get_num_bytes()); result_data.set_phi_scaling(phi_params[0], phi_params[1], phi_params[2]); + if (debug) { + printf("Encoding psi to %i bytes min=%f, max=%f, scale=%f\n", + result_data.get_num_bytes(), psi_params[0], psi_params[1], + psi_params[2]); + printf("Encoding phi to %i bytes min=%f, max=%f, scale=%f\n", + result_data.get_num_bytes(), phi_params[0], phi_params[1], + phi_params[2]); + } + // Do the local encoding. if (result_data.get_num_bytes() == 1) { set_encode_cpu_psi_phi_array(result_data, psi_imgs, phi_imgs); @@ -260,13 +275,20 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect set_encode_cpu_psi_phi_array(result_data, psi_imgs, phi_imgs); } } else { + if (debug) { printf("Encoding psi and phi as floats.\n"); } // Just interleave psi and phi images. set_float_cpu_psi_phi_array(result_data, psi_imgs, phi_imgs); } #ifdef HAVE_CUDA // Create a copy of the encoded data in GPU memory. + if (debug) { + printf("Allocating on device memory using %lu bytes.\n", result_data.get_total_array_size()); + } device_allocate_psi_phi_array(&result_data); + if (result_data.get_gpu_array_ptr() == nullptr) { + throw std::runtime_error("Unable to allocate GPU PsiPhi array."); + } #endif } diff --git a/src/kbmod/search/psi_phi_array_utils.h b/src/kbmod/search/psi_phi_array_utils.h index c69db11ef..d06a4483c 100644 --- a/src/kbmod/search/psi_phi_array_utils.h +++ b/src/kbmod/search/psi_phi_array_utils.h @@ -26,7 +26,7 @@ namespace search { std::array compute_scale_params_from_image_vect(const std::vector& imgs, int num_bytes); void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector& psi_imgs, - const std::vector& phi_imgs); + const std::vector& phi_imgs, bool debug=false); } /* namespace search */ diff --git a/src/kbmod/search/stack_search.cpp b/src/kbmod/search/stack_search.cpp index 3cf7c4040..0d88d719d 100644 --- a/src/kbmod/search/stack_search.cpp +++ b/src/kbmod/search/stack_search.cpp @@ -69,7 +69,6 @@ void StackSearch::set_start_bounds_y(int y_min, int y_max) { 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); - prepare_psi_phi(); create_search_list(ang_steps, vel_steps, min_ang, max_ang, min_vel, mavx); // Create a data stucture for the per-image data. @@ -78,7 +77,7 @@ void StackSearch::search(int ang_steps, int vel_steps, float min_ang, float max_ 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); + fill_psi_phi_array(psi_phi_data, params.encode_num_bytes, psi_images, phi_images, debug_info); psi_phi_timer.stop(); // Allocate a vector for the results. diff --git a/tests/test_psi_phi_array.py b/tests/test_psi_phi_array.py index a5abb7aaf..8260fffaf 100644 --- a/tests/test_psi_phi_array.py +++ b/tests/test_psi_phi_array.py @@ -124,7 +124,7 @@ def test_compute_scale_params_from_image_vect(self): def test_fill_psi_phi_array(self): for num_bytes in [2, 4]: arr = PsiPhiArray() - fill_psi_phi_array(arr, num_bytes, [self.psi_1, self.psi_2], [self.phi_1, self.phi_2]) + fill_psi_phi_array(arr, num_bytes, [self.psi_1, self.psi_2], [self.phi_1, self.phi_2], False) # Check the meta data. self.assertEqual(arr.num_times, self.num_times)