Merge pull request #342 from dirac-institute/cpu_versions2

Create a CPU version of Co-added Stamps

src/kbmod/analysis_utils.py

@@ -470,8 +470,11 @@ def apply_stamp_filter(
                 all_true = [True] * num_times
                 bool_slice = [all_true for _ in inds_to_use]
 
-            # Create and filter the results.
-            stamps_slice = search.gpu_coadded_stamps(trj_slice, bool_slice, params)
+            # Create and filter the results, using the GPU if there is one and enough
+            # trajectories to make it worthwhile.
+            stamps_slice = search.coadded_stamps(
+                trj_slice, bool_slice, params, kb.HAS_GPU and len(trj_slice) > 100
+            )
             for ind, stamp in enumerate(stamps_slice):
                 if stamp.get_width() > 1:
                     result_list.results[ind + start_idx].stamp = np.array(stamp)

src/kbmod/search/KBMOSearch.cpp

@@ -373,6 +373,56 @@ bool KBMOSearch::filterStamp(const RawImage& img, const stampParameters& params)
     return false;
 }
 
+std::vector<RawImage> KBMOSearch::coaddedScienceStamps(std::vector<trajectory>& t_array,
+                                                       std::vector<std::vector<bool> >& use_index_vect,
+                                                       const stampParameters& params,
+                                                       bool use_gpu) {
+    if (use_gpu) {
+        #ifdef HAVE_CUDA
+            return coaddedScienceStampsGPU(t_array, use_index_vect, params);
+        #else
+            print("WARNING: GPU is not enabled. Performing co-adds on the CPU.");
+        #endif
+    }
+    return coaddedScienceStampsCPU(t_array, use_index_vect, params);
+}
+
+std::vector<RawImage> KBMOSearch::coaddedScienceStampsCPU(std::vector<trajectory>& t_array,
+                                                          std::vector<std::vector<bool> >& use_index_vect,
+                                                          const stampParameters& params) {
+    const int num_trajectories = t_array.size();
+    std::vector<RawImage> results(num_trajectories);
+    std::vector<float> empty_pixels(1, NO_DATA);
+
+    for (int i = 0; i < num_trajectories; ++i) {
+        std::vector<RawImage> stamps = scienceStamps(t_array[i], params.radius, false, true, use_index_vect[i]);
+
+        RawImage coadd(1, 1);
+        switch (params.stamp_type) {
+            case STAMP_MEDIAN:
+                coadd = createMedianImage(stamps);
+                break;
+            case STAMP_MEAN:
+                coadd = createMeanImage(stamps);
+                break;
+            case STAMP_SUM:
+                coadd = createSummedImage(stamps);
+                break;
+            default:
+                throw std::runtime_error("Invalid stamp coadd type.");
+        }
+
+        // Do the filtering if needed.
+        if (params.do_filtering && filterStamp(coadd, params)) {
+            results[i] = RawImage(1, 1, empty_pixels);
+        } else {
+            results[i] = coadd;
+        }
+    }
+
+    return results;
+}
+
 std::vector<RawImage> KBMOSearch::coaddedScienceStampsGPU(std::vector<trajectory>& t_array,
                                                           std::vector<std::vector<bool> >& use_index_vect,
                                                           const stampParameters& params) {

src/kbmod/search/KBMOSearch.h

@@ -69,13 +69,14 @@ class KBMOSearch {
     RawImage meanScienceStamp(const trajectory& trj, int radius, const std::vector<bool>& use_index);
     RawImage summedScienceStamp(const trajectory& trj, int radius, const std::vector<bool>& use_index);
 
-    // Compute a mean or summed stamp for each trajectory on the GPU. This is slower than the
-    // above for small numbers of trajectories (< 500), but performs relatively better as the
-    // number of trajectories increases. If filtering is applied then will use a 1x1 image with
-    // NO_DATA to represent filtered images.
-    std::vector<RawImage> coaddedScienceStampsGPU(std::vector<trajectory>& t_array,
-                                                  std::vector<std::vector<bool> >& use_index_vect,
-                                                  const stampParameters& params);
+    // Compute a mean or summed stamp for each trajectory on the GPU or CPU.
+    // The GPU implementation is slower for small numbers of trajectories (< 500), but performs
+    // relatively better as the number of trajectories increases. If filtering is applied then
+    // the code will return a 1x1 image with NO_DATA to represent each filtered image.
+    std::vector<RawImage> coaddedScienceStamps(std::vector<trajectory>& t_array,
+                                               std::vector<std::vector<bool> >& use_index_vect,
+                                               const stampParameters& params,
+                                               bool use_cpu);
 
     // Function to do the actual stamp filtering.
     bool filterStamp(const RawImage& img, const stampParameters& params);
@@ -120,6 +121,13 @@ class KBMOSearch {
     void createSearchList(int angleSteps, int veloctiySteps, float minAngle, float maxAngle,
                           float minVelocity, float maxVelocity);
 
+    std::vector<RawImage> coaddedScienceStampsGPU(std::vector<trajectory>& t_array,
+                                                  std::vector<std::vector<bool> >& use_index_vect,
+                                                  const stampParameters& params);
+
+    std::vector<RawImage> coaddedScienceStampsCPU(std::vector<trajectory>& t_array,
+                                                  std::vector<std::vector<bool> >& use_index_vect,
+                                                  const stampParameters& params);
     // Helper functions for timing operations of the search.
     void startTimer(const std::string& message);
     void endTimer();

src/kbmod/search/bindings.cpp

@@ -199,10 +199,10 @@ PYBIND11_MODULE(search, m) {
             .def("median_sci_stamp", &ks::medianScienceStamp)
             .def("mean_sci_stamp", &ks::meanScienceStamp)
             .def("summed_sci_stamp", &ks::summedScienceStamp)
-            .def("gpu_coadded_stamps",
+            .def("coadded_stamps",
                  (std::vector<ri>(ks::*)(std::vector<tj> &, std::vector<std::vector<bool>> &,
-                                         const search::stampParameters &)) &
-                         ks::coaddedScienceStampsGPU)
+                                         const search::stampParameters &, bool)) &
+                         ks::coaddedScienceStamps)
             // For testing
             .def("filter_stamp", &ks::filterStamp)
             .def("get_traj_pos", &ks::getTrajPos)