Various fixes

dirac-institute · Dec 11, 2023 · 3e20b7f · 3e20b7f
1 parent 04c9ae6
commit 3e20b7f
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 46 deletions.
diff --git a/src/kbmod/search/kernels.cu b/src/kbmod/search/kernels.cu
@@ -54,7 +54,7 @@ __forceinline__ __device__ PsiPhi read_encoded_psi_phi(PsiPhiArrayMeta &params,
 
     // Compute the in-list index from the row, column, and time.
     int start_index = 2 * (params.pixels_per_image * time + row * params.width + col);
-    if (params.num_bytes == -1) {
+    if (params.num_bytes == 4) {
         // Short circuit the typical case of float encoding. No scaling or shifting done.
         return {reinterpret_cast<float *>(psi_phi_vect)[start_index],
                 reinterpret_cast<float *>(psi_phi_vect)[start_index + 1]};
@@ -278,6 +278,7 @@ extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, float *image_time
     // Check that the device psi_phi vector has been allocated.
     if (psi_phi_array.gpu_array_allocated() == false) {
         throw std::runtime_error("PsiPhi data has not been created.");
+    }
 
     // Copy trajectories to search
     if (params.debug) {
@@ -311,7 +312,7 @@ extern "C" void deviceSearchFilter(PsiPhiArray &psi_phi_array, float *image_time
     dim3 threads(THREAD_DIM_X, THREAD_DIM_Y);
 
     // Launch Search
-    searchFilterImages<<<blocks, threads>>>(psi_phi_array.get_meta_data(), psi_phi_array.get_gpu_array_ptr() device_img_times,
+    searchFilterImages<<<blocks, threads>>>(psi_phi_array.get_meta_data(), psi_phi_array.get_gpu_array_ptr(), device_img_times,
                                             params, num_trajectories, device_tests, device_search_results);
 
     // Read back results

diff --git a/src/kbmod/search/psi_phi_array.cpp b/src/kbmod/search/psi_phi_array.cpp
@@ -25,7 +25,7 @@ PsiPhiArray::~PsiPhiArray() {
     if (gpu_array_ptr != nullptr) {
         device_free_psi_phi_array(this);
     }
-#endif    
+#endif
 }
 
 void PsiPhiArray::clear() {
@@ -61,7 +61,7 @@ void PsiPhiArray::clear() {
 void PsiPhiArray::set_meta_data(int new_num_bytes, int new_num_times, int new_height, int new_width) {
     // Validity checking of parameters.
     if (new_num_bytes != -1 && new_num_bytes != 1 && new_num_bytes != 2 && new_num_bytes != 4) {
-        throw std::runtime_error("Invalid setting of num_bytes. Must be (-1, 1, 2, or 4).");
+        throw std::runtime_error("Invalid setting of num_bytes. Must be (-1 [use default], 1, 2, or 4).");
     }
     if (new_num_times <= 0) throw std::runtime_error("Invalid num_times passed to set_meta_data.\n");
     if (new_width <= 0) throw std::runtime_error("Invalid width passed to set_meta_data.\n");
@@ -78,7 +78,7 @@ void PsiPhiArray::set_meta_data(int new_num_bytes, int new_num_times, int new_he
     } else if (meta_data.num_bytes == 2) {
         meta_data.block_size = sizeof(uint16_t);
     } else {
-        meta_data.num_bytes = -1;
+        meta_data.num_bytes = 4;
         meta_data.block_size = sizeof(float);
     }
 
@@ -118,7 +118,7 @@ PsiPhi PsiPhiArray::read_psi_phi(int time, int row, int col) {
     // Compute the in-list index from the row, column, and time.
     int start_index = 2 * (meta_data.pixels_per_image * time + row * meta_data.width + col);
 
-    if (meta_data.num_bytes == -1) {
+    if (meta_data.num_bytes == 4) {
         // Short circuit the typical case of float encoding.
         // No scaling or shifting done.
         result.psi = reinterpret_cast<float*>(cpu_array_ptr)[start_index];
@@ -278,36 +278,45 @@ void fill_psi_phi_array(PsiPhiArray& result_data, int num_bytes, const std::vect
 static void psi_phi_array_binding(py::module& m) {
     using ppa = search::PsiPhiArray;
 
-    py::class_<search::PsiPhi>(m, "PsiPhi", pydocs::DOC_PsiPhi))
+    py::class_<search::PsiPhi>(m, "PsiPhi", pydocs::DOC_PsiPhi)
             .def(py::init<>())
             .def_readwrite("psi", &search::PsiPhi::psi)
             .def_readwrite("phi", &search::PsiPhi::phi);
 
-    py::class_<ppa>(m, "PsiPhiArray", DOC_PsiPhiArray)
+    py::class_<ppa>(m, "PsiPhiArray", pydocs::DOC_PsiPhiArray)
             .def(py::init<>())
-            .def_property_readonly("num_bytes", &ppa::get_num_bytes, DOC_PsiPhiArray_get_num_bytes)
-            .def_property_readonly("num_times", &ppa::get_num_times, DOC_PsiPhiArray_get_num_times)
-            .def_property_readonly("width", &ppa::get_width, DOC_PsiPhiArray_get_width)
-            .def_property_readonly("height", &ppa::get_height, DOC_PsiPhiArray_get_height)
-            .def_property_readonly("pixels_per_image", &ppa::get_pixels_per_image, DOC_PsiPhiArray_get_pixels_per_image)
-            .def_property_readonly("num_entries", &ppa::get_num_entries, DOC_PsiPhiArray_get_num_entries)
-            .def_property_readonly("total_array_size", &ppa::get_total_array_size, DOC_PsiPhiArray_get_total_array_size)
-            .def_property_readonly("block_size", &ppa::get_block_size, DOC_PsiPhiArray_get_block_size)
-            .def_property_readonly("psi_min_val", &ppa::get_psi_min_val, DOC_PsiPhiArray_get_psi_min_val)
-            .def_property_readonly("psi_max_val", &ppa::get_psi_max_val, DOC_PsiPhiArray_get_psi_max_val)
-            .def_property_readonly("psi_scale", &ppa::get_psi_scale, DOC_PsiPhiArray_get_psi_scale)
-            .def_property_readonly("phi_min_val", &ppa::get_phi_min_val, DOC_PsiPhiArray_get_phi_min_val)
-            .def_property_readonly("phi_max_val", &ppa::get_phi_max_val, DOC_PsiPhiArray_get_phi_max_val)
-            .def_property_readonly("phi_scale", &ppa::get_phi_scale, DOC_PsiPhiArray_get_phi_scale)
-            .def_property_readonly("cpu_array_allocated", &ppa::cpu_array_allocated, DOC_PsiPhiArray_get_cpu_array_allocated)
-            .def_property_readonly("gpu_array_allocated", &ppa::gpu_array_allocated, DOC_PsiPhiArray_get_gpu_array_allocated)
-            .def("set_meta_data", &ppa::set_meta_data, DOC_PsiPhiArray_set_meta_data)
-            .def("clear", &ppa::clear, DOC_PsiPhiArray_clear)
-            .def("read_psi_phi", &ppa::read_psi_phi, DOC_PsiPhiArray_read_psi_phi);
+            .def_property_readonly("num_bytes", &ppa::get_num_bytes, pydocs::DOC_PsiPhiArray_get_num_bytes)
+            .def_property_readonly("num_times", &ppa::get_num_times, pydocs::DOC_PsiPhiArray_get_num_times)
+            .def_property_readonly("width", &ppa::get_width, pydocs::DOC_PsiPhiArray_get_width)
+            .def_property_readonly("height", &ppa::get_height, pydocs::DOC_PsiPhiArray_get_height)
+            .def_property_readonly("pixels_per_image", &ppa::get_pixels_per_image,
+                                   pydocs::DOC_PsiPhiArray_get_pixels_per_image)
+            .def_property_readonly("num_entries", &ppa::get_num_entries,
+                                   pydocs::DOC_PsiPhiArray_get_num_entries)
+            .def_property_readonly("total_array_size", &ppa::get_total_array_size,
+                                   pydocs::DOC_PsiPhiArray_get_total_array_size)
+            .def_property_readonly("block_size", &ppa::get_block_size, pydocs::DOC_PsiPhiArray_get_block_size)
+            .def_property_readonly("psi_min_val", &ppa::get_psi_min_val,
+                                   pydocs::DOC_PsiPhiArray_get_psi_min_val)
+            .def_property_readonly("psi_max_val", &ppa::get_psi_max_val,
+                                   pydocs::DOC_PsiPhiArray_get_psi_max_val)
+            .def_property_readonly("psi_scale", &ppa::get_psi_scale, pydocs::DOC_PsiPhiArray_get_psi_scale)
+            .def_property_readonly("phi_min_val", &ppa::get_phi_min_val,
+                                   pydocs::DOC_PsiPhiArray_get_phi_min_val)
+            .def_property_readonly("phi_max_val", &ppa::get_phi_max_val,
+                                   pydocs::DOC_PsiPhiArray_get_phi_max_val)
+            .def_property_readonly("phi_scale", &ppa::get_phi_scale, pydocs::DOC_PsiPhiArray_get_phi_scale)
+            .def_property_readonly("cpu_array_allocated", &ppa::cpu_array_allocated,
+                                   pydocs::DOC_PsiPhiArray_get_cpu_array_allocated)
+            .def_property_readonly("gpu_array_allocated", &ppa::gpu_array_allocated,
+                                   pydocs::DOC_PsiPhiArray_get_gpu_array_allocated)
+            .def("set_meta_data", &ppa::set_meta_data, pydocs::DOC_PsiPhiArray_set_meta_data)
+            .def("clear", &ppa::clear, pydocs::DOC_PsiPhiArray_clear)
+            .def("read_psi_phi", &ppa::read_psi_phi, pydocs::DOC_PsiPhiArray_read_psi_phi);
     m.def("compute_scale_params_from_image_vect", &search::compute_scale_params_from_image_vect);
     m.def("decode_uint_scalar", &search::decode_uint_scalar);
     m.def("encode_uint_scalar", &search::encode_uint_scalar);
-    m.def("fill_psi_phi_array", &search::fill_psi_phi_array, DOC_PsiPhiArray_fill_psi_phi_array);
+    m.def("fill_psi_phi_array", &search::fill_psi_phi_array, pydocs::DOC_PsiPhiArray_fill_psi_phi_array);
 }
 #endif
 

diff --git a/src/kbmod/search/psi_phi_array_ds.h b/src/kbmod/search/psi_phi_array_ds.h
@@ -6,6 +6,12 @@
  * from the rest of the utility functions) to allow the CUDA files to import
  * only what they need.
  *
+ * The data structure allocates memory on both the CPU and GPU for the
+ * interleaved psi/phi array and maintains ownership of the pointers
+ * until clear() is called or the PsiPhiArray's destructor is called. This allows
+ * the object to be passed repeatedly to the on-device search without reallocating
+ * and copying the memory on the GPU.
+ *
  * Created on: Dec 5, 2023
  */
 
@@ -43,11 +49,11 @@ struct PsiPhiArrayMeta {
     int height = 0;
     int pixels_per_image = 0;
     int num_entries = 0;
-    int block_size = 0;
+    int block_size = 0;  // Actual memory used per entry.
     long unsigned total_array_size = 0;
 
     // Compression and scaling parameters of on GPU array.
-    int num_bytes = -1;  // -1 (float), 1 (unit8) or 2 (unit16)
+    int num_bytes = 4;  // 1 (unit8), 2 (unit16), or 4 (float)
 
     float psi_min_val = FLT_MAX;
     float psi_max_val = -FLT_MAX;

diff --git a/src/kbmod/search/pydocs/psi_phi_array_docs.h b/src/kbmod/search/pydocs/psi_phi_array_docs.h
@@ -13,13 +13,14 @@ static const auto DOC_PsiPhi = R"doc(
   phi : `float`
       The phi value at a pixel.
   )doc";
-  
-  static const auto DOC_PsiPhiArray = R"doc(
+
+static const auto DOC_PsiPhiArray = R"doc(
   An encoded array of Psi and Phi values along with their meta data.
   )doc";
 
 static const auto DOC_PsiPhiArray_get_num_bytes = R"doc(
-  The setting for encoding (-1 for float, 1 for uint8, and 2 for uint16)
+  The target number of bytes to use for encoding the data (1 for uint8, 2 for uint16,
+  or 4 for float32). Might differ from actual number of bytes (block_size).
   )doc";
 
 static const auto DOC_PsiPhiArray_get_num_times = R"doc(
@@ -74,10 +75,6 @@ static const auto DOC_PsiPhiArray_get_phi_scale = R"doc(
   The scaling parameter for phi.
   )doc";
 
-static const auto DOC_PsiPhiArray_get_phi_scale = R"doc(
-  The scaling parameter for phi.
-  )doc";
-
 static const auto DOC_PsiPhiArray_get_cpu_array_allocated = R"doc(
   A Boolean indicating whether the cpu array exists.
   )doc";
@@ -115,7 +112,7 @@ static const auto DOC_PsiPhiArray_set_meta_data = R"doc(
     Parameters
     ----------
     new_num_bytes : `int`
-        The type of encoding to use (-1, 1, or 2).
+        The type of encoding to use (1, 2, or 4).
     new_num_times : `int`
         The number of time steps in the data.
     new_height : `int`
@@ -132,7 +129,7 @@ static const auto DOC_PsiPhiArray_fill_psi_phi_array = R"doc(
     result_data : `PsiPhiArray`
         The location to store the data.
     num_bytes : `int`
-        The type of encoding to use (-1, 1, or 2).
+        The type of encoding to use (1, 2, or 4).
     psi_imgs : `list`
         A list of psi images.
     phi_imgs : `list`

diff --git a/tests/test_psi_phi_array.py b/tests/test_psi_phi_array.py
@@ -34,7 +34,7 @@ def setUp(self):
     def test_set_meta_data(self):
         arr = PsiPhiArray()
         self.assertEqual(arr.num_times, 0)
-        self.assertEqual(arr.num_bytes, -1)
+        self.assertEqual(arr.num_bytes, 4)
         self.assertEqual(arr.width, 0)
         self.assertEqual(arr.height, 0)
         self.assertEqual(arr.pixels_per_image, 0)
@@ -43,8 +43,8 @@ def test_set_meta_data(self):
         self.assertEqual(arr.total_array_size, 0)
 
         # Make float
-        arr.set_meta_data(-1, self.num_times, self.height, self.width)
-        self.assertEqual(arr.num_bytes, -1)
+        arr.set_meta_data(4, self.num_times, self.height, self.width)
+        self.assertEqual(arr.num_bytes, 4)
         self.assertEqual(arr.block_size, 4)
         self.assertEqual(arr.num_times, self.num_times)
         self.assertEqual(arr.width, self.width)
@@ -104,7 +104,7 @@ def test_compute_scale_params_from_image_vect(self):
         max_val = 2 * self.width * self.height - 1
 
         # Parameters for encoding to a float
-        result_float = compute_scale_params_from_image_vect([self.psi_1, self.psi_2], -1)
+        result_float = compute_scale_params_from_image_vect([self.psi_1, self.psi_2], 4)
         self.assertAlmostEqual(result_float[0], 0.0, delta=1e-5)
         self.assertAlmostEqual(result_float[1], max_val, delta=1e-5)
         self.assertAlmostEqual(result_float[2], 1.0, delta=1e-5)
@@ -122,7 +122,7 @@ def test_compute_scale_params_from_image_vect(self):
         self.assertAlmostEqual(result_uint16[2], max_val / 65535.0, delta=1e-5)
 
     def test_fill_psi_phi_array(self):
-        for num_bytes in [-1, 2]:
+        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])
 
@@ -133,13 +133,12 @@ def test_fill_psi_phi_array(self):
             self.assertEqual(arr.height, self.height)
             self.assertEqual(arr.pixels_per_image, self.width * self.height)
             self.assertEqual(arr.num_entries, 2 * arr.pixels_per_image * self.num_times)
-            if num_bytes == -1:
+            if num_bytes == 4:
                 self.assertEqual(arr.block_size, 4)
             else:
                 self.assertEqual(arr.block_size, num_bytes)
             self.assertEqual(arr.total_array_size, arr.num_entries * arr.block_size)
 
-
             # Check that we can correctly read the values from the CPU.
             self.assertTrue(arr.cpu_array_allocated)
             for time in range(self.num_times):