Merge pull request #427 from dirac-institute/error_checking

Improve error checking
dirac-institute · Jan 11, 2024 · 08317b7 · 08317b7
2 parents 83902da + 9566583
commit 08317b7
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git 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<RawImage>
         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<RawImage>&
         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<RawImage>&
 }
 
 void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector<RawImage>& psi_imgs,
-                        const std::vector<RawImage>& phi_imgs) {
+                        const std::vector<RawImage>& phi_imgs, bool debug) {
     if (result_data.get_cpu_array_ptr() != nullptr) {
         return;
     }
@@ -253,20 +259,36 @@ 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<uint8_t>(result_data, psi_imgs, phi_imgs);
         } else {
             set_encode_cpu_psi_phi_array<uint16_t>(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
@@ -26,7 +26,7 @@ namespace search {
 std::array<float, 3> compute_scale_params_from_image_vect(const std::vector<RawImage>& imgs, int num_bytes);
 
 void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vector<RawImage>& psi_imgs,
-                        const std::vector<RawImage>& phi_imgs);
+                        const std::vector<RawImage>& phi_imgs, bool debug=false);
 
 } /* namespace search */
 

diff --git 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
@@ -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)