Merge pull request #337 from dirac-institute/cpu_versions

Create CPU versions of several functions

src/kbmod/search/KBMOSearch.cpp

@@ -333,6 +333,46 @@ RawImage KBMOSearch::summedScienceStamp(const trajectory& trj, int radius,
             scienceStamps(trj, radius, false /*=interpolate*/, false /*=keep_no_data*/, use_index));
 }
 
+bool KBMOSearch::filterStamp(const RawImage& img, const stampParameters& params) {
+    // Allocate space for the coadd information and initialize to zero.
+    const int stamp_width = 2 * params.radius + 1;
+    const int stamp_ppi = stamp_width * stamp_width;
+    const std::vector<float>& pixels = img.getPixels();
+
+    // Filter on the peak's position.
+    pixelPos pos = img.findPeak(true);
+    if ((abs(pos.x - params.radius) >= params.peak_offset_x) ||
+        (abs(pos.y - params.radius) >= params.peak_offset_y)) {
+        return true;
+    }
+
+    // Filter on the percentage of flux in the central pixel.
+    if (params.center_thresh > 0.0) {
+        const std::vector<float>& pixels = img.getPixels();
+        float center_val = pixels[(int)pos.y * stamp_width + (int)pos.x];
+        float pixel_sum = 0.0;
+        for (int p = 0; p < stamp_ppi; ++p) {
+            pixel_sum += pixels[p];
+        }
+
+        if (center_val / pixel_sum < params.center_thresh) {
+            return true;
+        }
+    }
+
+    // Filter on the image moments.
+    imageMoments moments = img.findCentralMoments();
+    if ((fabs(moments.m01) >= params.m01_limit) ||
+        (fabs(moments.m10) >= params.m10_limit) ||
+        (fabs(moments.m11) >= params.m11_limit) ||
+        (moments.m02 >= params.m02_limit) ||
+        (moments.m20 >= params.m20_limit)) {
+        return true;
+    }
+
+    return false;
+}
+
 std::vector<RawImage> KBMOSearch::coaddedScienceStampsGPU(std::vector<trajectory>& t_array,
                                                           std::vector<std::vector<bool> >& use_index_vect,
                                                           const stampParameters& params) {
@@ -364,19 +404,19 @@ std::vector<RawImage> KBMOSearch::coaddedScienceStampsGPU(std::vector<trajectory
         throw std::runtime_error("Non-GPU co-adds is not implemented.");
     #endif
 
-    // Copy the stamps into RawImages
+    // Copy the stamps into RawImages and do the filtering.
     std::vector<RawImage> results(num_trajectories);
     std::vector<float> current_pixels(stamp_ppi, 0.0);
     std::vector<float> empty_pixels(1, NO_DATA);
     for (int t = 0; t < num_trajectories; ++t) {
-        bool all_no_data = true;
+        // Copy the data into a single RawImage.
         int offset = t * stamp_ppi;
         for (unsigned p = 0; p < stamp_ppi; ++p) {
             current_pixels[p] = stamp_data[offset + p];
-            all_no_data = all_no_data && (stamp_data[offset + p] == NO_DATA);
         }
+        RawImage current_image = RawImage(stamp_width, stamp_width, current_pixels);
 
-        if (all_no_data && params.do_filtering) {
+        if (params.do_filtering && filterStamp(current_image, params)) {
             results[t] = RawImage(1, 1, empty_pixels);
         } else {
             results[t] = RawImage(stamp_width, stamp_width, current_pixels);

src/kbmod/search/KBMOSearch.h

@@ -77,6 +77,9 @@ class KBMOSearch {
                                                   std::vector<std::vector<bool> >& use_index_vect,
                                                   const stampParameters& params);
 
+    // Function to do the actual stamp filtering.
+    bool filterStamp(const RawImage& img, const stampParameters& params);
+
     // Getters for the Psi and Phi data.
     std::vector<RawImage>& getPsiImages();
     std::vector<RawImage>& getPhiImages();

src/kbmod/search/PointSpreadFunc.h

@@ -39,9 +39,10 @@ class PointSpreadFunc {
     PointSpreadFunc& operator=(const PointSpreadFunc& other);  // Copy assignment
     PointSpreadFunc& operator=(PointSpreadFunc&& other);       // Move assignment
 
-    // Getter functions.
+    // Getter functions (inlined)
     float getStdev() const { return width; }
     float getSum() const { return sum; }
+    float getValue(int x, int y) const { return kernel[y * dim + x]; }
     int getDim() const { return dim; }
     int getRadius() const { return radius; }
     int getSize() const { return kernel.size(); }

src/kbmod/search/RawImage.cpp

@@ -14,10 +14,6 @@ namespace search {
     // and a PSF on a GPU device.
     extern "C" void deviceConvolve(float* sourceImg, float* resultImg, int width, int height, float* psfKernel,
                                    int psfSize, int psfDim, int psfRadius, float psfSum);
-
-    extern "C" pixelPos findPeakImageVect(int width, int height, float* img, bool furthest_from_center);
-
-    extern "C" imageMoments findCentralMomentsImageVect(int width, int height, float* img);
 #endif
 
 RawImage::RawImage() : width(0), height(0) { pixels = std::vector<float>(); }
@@ -142,12 +138,50 @@ RawImage RawImage::createStamp(float x, float y, int radius, bool interpolate, b
     return stamp;
 }
 
+void RawImage::convolve_cpu(const PointSpreadFunc& psf) {
+    std::vector<float> result(width * height, 0.0);
+    const int psfRad = psf.getRadius();
+    const float psfTotal = psf.getSum();
+
+    for (int y = 0; y < height; ++y) {
+        for (int x = 0; x < width; ++x) {
+            // Pixels with NO_DATA remain NO_DATA.
+            if (pixels[y * width + x] == NO_DATA) {
+                result[y * width + x] = NO_DATA;
+                continue;
+            }
+
+            float sum = 0.0;
+            float psfPortion = 0.0;
+            for (int j = -psfRad; j <= psfRad; j++) {
+                for (int i = -psfRad; i <= psfRad; i++) {
+                    if ((x + i >= 0) && (x + i < width) && (y + j >= 0) && (y + j < height)) {
+                        float currentPixel = pixels[(y + j) * width + (x + i)];
+                        if (currentPixel != NO_DATA) {
+                            float currentPSF = psf.getValue(i + psfRad, j + psfRad);
+                            psfPortion += currentPSF;
+                            sum += currentPixel * currentPSF;
+                        }
+                    }
+                }
+            }
+            result[y * width + x] = (sum * psfTotal) / psfPortion;
+        }
+    }
+
+    // Copy the data into the pixels vector.
+    const int ppi = width * height;
+    for(int i = 0; i < ppi; ++i) {
+        pixels[i] = result[i];
+    }
+}
+
 void RawImage::convolve(PointSpreadFunc psf) {
     #ifdef HAVE_CUDA
         deviceConvolve(pixels.data(), pixels.data(), getWidth(), getHeight(), psf.kernelData(),
                        psf.getSize(), psf.getDim(), psf.getRadius(), psf.getSum());
     #else
-        throw std::runtime_error("Non-GPU convolution is not implemented.");
+        convolve_cpu(psf);
     #endif
 }
 
@@ -325,20 +359,81 @@ std::array<float, 2> RawImage::computeBounds() const {
 }
 
 // The maximum value of the image and return the coordinates.
-pixelPos RawImage::findPeak(bool furthest_from_center) {
-    #ifdef HAVE_CUDA
-        return findPeakImageVect(width, height, pixels.data(), furthest_from_center);
-    #else
-        throw std::runtime_error("Non-GPU findPeak is not implemented.");
-    #endif
+pixelPos RawImage::findPeak(bool furthest_from_center) const {
+    int c_x = width / 2;
+    int c_y = height / 2;
+
+    // Initialize the variables for tracking the peak's location.
+    pixelPos result = {0, 0};
+    float max_val = pixels[0];
+    float dist2 = c_x * c_x + c_y * c_y;
+
+    // Search each pixel for the peak.
+    for (int y = 0; y < height; ++y) {
+        for (int x = 0; x < width; ++x) {
+            float pix_val = pixels[y * width + x];
+            if (pix_val > max_val) {
+                max_val = pix_val;
+                result.x = x;
+                result.y = y;
+                dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y);
+            } else if (pix_val == max_val) {
+                int new_dist2 = (c_x - x) * (c_x - x) + (c_y - y) * (c_y - y);
+                if ((furthest_from_center && (new_dist2 > dist2)) ||
+                    (!furthest_from_center && (new_dist2 < dist2))) {
+                    max_val = pix_val;
+                    result.x = x;
+                    result.y = y;
+                    dist2 = new_dist2;
+                }
+            }
+        }
+    }
+
+    return result;
 }
 
-imageMoments RawImage::findCentralMoments() {
-    #ifdef HAVE_CUDA
-        return findCentralMomentsImageVect(width, height, pixels.data());
-    #else
-        throw std::runtime_error("Non-GPU findCentralMoments is not implemented.");
-    #endif
+// Find the basic image moments in order to test if stamps have a gaussian shape.
+// It computes the moments on the "normalized" image where the minimum
+// value has been shifted to zero and the sum of all elements is 1.0.
+// Elements with NO_DATA are treated as zero.
+imageMoments RawImage::findCentralMoments() const {
+    const int num_pixels = width * height;
+    const int c_x = width / 2;
+    const int c_y = height / 2;
+
+    // Set all the moments to zero initially.
+    imageMoments res = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
+
+    // Find the min (non-NO_DATA) value to subtract off.
+    float min_val = FLT_MAX;
+    for (int p = 0; p < num_pixels; ++p) {
+        min_val = ((pixels[p] != NO_DATA) && (pixels[p] < min_val)) ? pixels[p] : min_val;
+    }
+
+    // Find the sum of the zero-shifted (non-NO_DATA) pixels.
+    double sum = 0.0;
+    for (int p = 0; p < num_pixels; ++p) {
+        sum += (pixels[p] != NO_DATA) ? (pixels[p] - min_val) : 0.0;
+    }
+    if (sum == 0.0) return res;
+
+    // Compute the rest of the moments.
+    for (int y = 0; y < height; ++y) {
+        for (int x = 0; x < width; ++x) {
+            int ind = y * width + x;
+            float pix_val = (pixels[ind] != NO_DATA) ? (pixels[ind] - min_val) / sum : 0.0;
+
+            res.m00 += pix_val;
+            res.m10 += (x - c_x) * pix_val;
+            res.m20 += (x - c_x) * (x - c_x) * pix_val;
+            res.m01 += (y - c_y) * pix_val;
+            res.m02 += (y - c_y) * (y - c_y) * pix_val;
+            res.m11 += (x - c_x) * (y - c_y) * pix_val;
+        }
+    }
+
+    return res;
 }
 
 RawImage createMedianImage(const std::vector<RawImage>& images) {

src/kbmod/search/RawImage.h

@@ -95,6 +95,7 @@ class RawImage {
 
     // Convolve the image with a point spread function.
     void convolve(PointSpreadFunc psf);
+    void convolve_cpu(const PointSpreadFunc& psf);
 
     // Create a "stamp" image of a give radius (width=2*radius+1)
     // about the given point.
@@ -104,13 +105,13 @@ class RawImage {
     // The maximum value of the image and return the coordinates. The parameter
     // furthest_from_center indicates whether to break ties using the peak further
     // or closer to the center of the image.
-    pixelPos findPeak(bool furthest_from_center);
+    pixelPos findPeak(bool furthest_from_center) const;
 
     // Find the basic image moments in order to test if stamps have a gaussian shape.
     // It computes the moments on the "normalized" image where the minimum
     // value has been shifted to zero and the sum of all elements is 1.0.
     // Elements with NO_DATA are treated as zero.
-    imageMoments findCentralMoments();
+    imageMoments findCentralMoments() const;
 
     virtual ~RawImage(){};
 

src/kbmod/search/bindings.cpp

@@ -24,6 +24,7 @@ using std::to_string;
 
 PYBIND11_MODULE(search, m) {
     m.attr("KB_NO_DATA") = pybind11::float_(search::NO_DATA);
+    m.attr("HAS_GPU") = pybind11::bool_(search::HAVE_GPU);
     py::enum_<search::StampType>(m, "StampType")
             .value("STAMP_SUM", search::StampType::STAMP_SUM)
             .value("STAMP_MEAN", search::StampType::STAMP_MEAN)
@@ -71,6 +72,7 @@ PYBIND11_MODULE(search, m) {
             .def("get_radius", &pf::getRadius, "Returns the radius of the PSF")
             .def("get_size", &pf::getSize, "Returns the number of elements in the PSFs kernel.")
             .def("get_kernel", &pf::getKernel, "Returns the PSF kernel.")
+            .def("get_value", &pf::getValue, "Returns the PSF kernel value at a specific point.")
             .def("square_psf", &pf::squarePSF,
                  "Squares, raises to the power of two, the elements of the PSF kernel.")
             .def("print_psf", &pf::printPSF, "Pretty-prints the PSF.");
@@ -104,6 +106,7 @@ PYBIND11_MODULE(search, m) {
             .def("get_pixel", &ri::getPixel, "Returns the value of a pixel.")
             .def("get_pixel_interp", &ri::getPixelInterp, "Get the interoplated value of a pixel.")
             .def("convolve", &ri::convolve, "Convolve the image with a PSF.")
+            .def("convolve_cpu", &ri::convolve_cpu, "Convolve the image with a PSF.")
             .def("save_fits", &ri::saveToFile, "Save the image to a FITS file.");
     m.def("create_median_image", &search::createMedianImage);
     m.def("create_summed_image", &search::createSummedImage);
@@ -202,6 +205,7 @@ PYBIND11_MODULE(search, m) {
                                          const search::stampParameters &)) &
                          ks::coaddedScienceStampsGPU)
             // For testing
+            .def("filter_stamp", &ks::filterStamp)
             .def("get_traj_pos", &ks::getTrajPos)
             .def("get_mult_traj_pos", &ks::getMultTrajPos)
             .def("psi_curves", (std::vector<float>(ks::*)(tj &)) & ks::psiCurves)

src/kbmod/search/common.h

@@ -10,6 +10,12 @@
 
 namespace search {
 
+#ifdef HAVE_CUDA
+    constexpr bool HAVE_GPU = true;
+#else
+    constexpr bool HAVE_GPU = false;
+#endif
+
 constexpr unsigned int MAX_KERNEL_RADIUS = 15;
 constexpr unsigned short MAX_STAMP_EDGE = 64;
 constexpr unsigned short CONV_THREAD_DIM = 32;