Optimize VariablePsfStack using KdTree

.gitignore

@@ -11,3 +11,4 @@ __pycache__
 /packages
 /doc/build/
 /build/
+.vscode

SEBenchmarks/CMakeLists.txt

@@ -56,6 +56,8 @@ elements_add_executable(BenchRendering src/program/BenchRendering.cpp
         LINK_LIBRARIES SEFramework SEImplementation ${Boost_LIBRARIES})
 elements_add_executable(BenchBackgroundModel src/program/BenchBackgroundModel.cpp
         LINK_LIBRARIES SEFramework SEImplementation ${Boost_LIBRARIES})
+elements_add_executable(BenchVariablePsfStack src/program/BenchVariablePsfStack.cpp
+        LINK_LIBRARIES SEFramework ${Boost_LIBRARIES})
 
 #===============================================================================
 # Declare the Boost tests here

SEBenchmarks/src/program/BenchVariablePsfStack.cpp

@@ -0,0 +1,136 @@
+/** Copyright © 2019-2025 Université de Genève, LMU Munich - Faculty of Physics, IAP-CNRS/Sorbonne Université
+ *
+ * This library is free software; you can redistribute it and/or modify it under
+ * the terms of the GNU Lesser General Public License as published by the Free
+ * Software Foundation; either version 3.0 of the License, or (at your option)
+ * any later version.
+ *
+ * This library is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+ * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
+ * details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this library; if not, write to the Free Software Foundation, Inc.,
+ * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+/**
+ * @file src/program/BenchVariablePsfStack.cpp
+ * @date 06/27/25
+ * @author marc schefer
+ */
+
+#include <map>
+#include <string>
+#include <random>
+
+#include <boost/program_options.hpp>
+#include <boost/timer/timer.hpp>
+#include "ElementsKernel/ProgramHeaders.h"
+#include "ElementsKernel/Real.h"
+#include "SEFramework/Psf/VariablePsfStack.h"
+
+namespace po = boost::program_options;
+namespace timer = boost::timer;
+using namespace SourceXtractor;
+
+static Elements::Logging logger = Elements::Logging::getLogger("BenchVariablePsfStack");
+
+class BenchVariablePsfStack : public Elements::Program {
+private:
+  std::default_random_engine random_generator;
+  std::uniform_real_distribution<double> random_dist{0.0, 1000.0};
+
+public:
+
+  po::options_description defineSpecificProgramOptions() override {
+    po::options_description options{};
+    options.add_options()
+      ("iterations", po::value<int>()->default_value(100000), "Number of getPsf calls to benchmark")
+      ("measures", po::value<int>()->default_value(3), "Number of timing measurements to take")
+      ("fits-file", po::value<std::string>()->default_value(""), "FITS file containing PSF stack");
+    return options;
+  }
+
+  Elements::ExitCode mainMethod(std::map<std::string, po::variable_value> &args) override {
+
+    auto iterations = args["iterations"].as<int>();
+    auto measures = args["measures"].as<int>();
+    auto fits_file = args["fits-file"].as<std::string>();
+
+    logger.info() << "Benchmarking VariablePsfStack::getPsf() with " << iterations << " iterations";
+    logger.info() << "Taking " << measures << " timing measurements";
+
+    // Initialize VariablePsfStack with FITS file if provided, otherwise nullptr
+    std::shared_ptr<CCfits::FITS> fitsPtr = nullptr;
+    if (!fits_file.empty()) {
+      try {
+        fitsPtr = std::make_shared<CCfits::FITS>(fits_file);
+        logger.info() << "Using FITS file: " << fits_file;
+      } catch (const std::exception& e) {
+        logger.error() << "Failed to load FITS file '" << fits_file << "': " << e.what();
+        return Elements::ExitCode::DATAERR;
+      }
+    } else {
+      logger.error() << "No FITS file provided";
+      return Elements::ExitCode::USAGE;
+    }
+
+    try {
+      VariablePsfStack psfStack(fitsPtr);
+
+      logger.info() << "VariablePsfStack loaded successfully with " << psfStack.getNumberOfPsfs() << " PSFs";
+      logger.info() << "PSF size: " << psfStack.getWidth() << "x" << psfStack.getHeight();
+      logger.info() << "Pixel sampling: " << psfStack.getPixelSampling();
+
+      std::cout << "Iterations,Measurement,Time_nanoseconds" << std::endl;
+
+      for (int m = 0; m < measures; ++m) {
+        logger.info() << "Measurement " << (m + 1) << "/" << measures;
+
+        timer::cpu_timer timer;
+        timer.stop();
+
+        // Prepare test values for getPsf calls
+        std::vector<std::vector<double>> testValues;
+        testValues.reserve(iterations);
+
+        for (int i = 0; i < iterations; ++i) {
+          testValues.push_back({random_dist(random_generator), random_dist(random_generator)});
+        }
+
+        // Start timing
+        timer.start();
+
+        for (int i = 0; i < iterations; ++i) {
+          try {
+            auto psf = psfStack.getPsf(testValues[i]);
+            // Prevent compiler optimization by using the result
+            volatile auto width = psf->getWidth();
+            (void)width; // Suppress unused variable warning
+          } catch (const std::exception& e) {
+            // Expected to fail with nullptr, but we still measure the timing
+            // until the exception is thrown
+          }
+        }
+
+        timer.stop();
+
+        auto elapsed_ns = timer.elapsed().wall;
+        std::cout << iterations << "," << (m + 1) << "," << elapsed_ns << std::endl;
+
+        logger.info() << "Time for " << iterations << " calls: " << (elapsed_ns / 1e9) << " seconds";
+        logger.info() << "Average time per call: " << (elapsed_ns / iterations) << " nanoseconds";
+      }
+
+    } catch (const std::exception& e) {
+      logger.error() << "Error initializing VariablePsfStack: " << e.what();
+      return Elements::ExitCode::DATAERR;
+    }
+
+    return Elements::ExitCode::OK;
+  }
+};
+
+MAIN_FOR(BenchVariablePsfStack)

SEFramework/SEFramework/Psf/VariablePsfStack.h

@@ -25,8 +25,10 @@
 #define _SEIMPLEMENTATION_PSF_VARIABLEPSFSTACK_H_
 
 #include <CCfits/CCfits>
+#include <memory>
 #include <SEFramework/Image/VectorImage.h>
 #include <SEFramework/Psf/Psf.h>
+#include <SEUtils/KdTree.h>
 
 namespace SourceXtractor {
 
@@ -42,6 +44,18 @@ namespace SourceXtractor {
  */
 class VariablePsfStack final : public Psf {
 public:
+  /**
+   * @brief Structure to hold PSF position data
+   */
+  struct PsfPosition {
+    SeFloat ra;
+    SeFloat dec;
+    SeFloat x;
+    SeFloat y;
+    double gridx;
+    double gridy;
+  };
+
   /**
    * Constructor
    */
@@ -83,6 +97,13 @@ class VariablePsfStack final : public Psf {
     return m_components;
   };
 
+  /**
+   * @return The number of PSFs loaded in the stack
+   */
+  long getNumberOfPsfs() const {
+    return m_nrows;
+  };
+
   /**
    *
    */
@@ -99,12 +120,8 @@ class VariablePsfStack final : public Psf {
 
   long m_nrows;
 
-  std::vector<SeFloat> m_ra_values;
-  std::vector<SeFloat> m_dec_values;
-  std::vector<SeFloat> m_x_values;
-  std::vector<SeFloat> m_y_values;
-  std::vector<double>  m_gridx_values;
-  std::vector<double>  m_gridy_values;
+  std::vector<PsfPosition> m_positions;
+  std::unique_ptr<KdTree<PsfPosition>> m_kdtree;
 
   std::vector<std::string> m_components = {"X_IMAGE", "Y_IMAGE"};
 
@@ -119,6 +136,16 @@ class VariablePsfStack final : public Psf {
   void selfTest();
 };
 
+/**
+ * @brief KdTree traits specialization for PsfPosition
+ */
+template <>
+struct KdTreeTraits<VariablePsfStack::PsfPosition> {
+  static double getCoord(const VariablePsfStack::PsfPosition& pos, size_t index) {
+    return (index == 0) ? static_cast<double>(pos.x) : static_cast<double>(pos.y);
+  }
+};
+
 }  // namespace SourceXtractor
 
 #endif  //_SEIMPLEMENTATION_PSF_VARIABLEPSFSTACK_H_

SEFramework/src/lib/Psf/VariablePsfStack.cpp

@@ -21,6 +21,7 @@
  *      Author: Martin Kuemmel
  */
 #include <algorithm>
+#include <memory>
 #include <ElementsKernel/Logging.h>
 #include <ElementsKernel/Exception.h>
 #include "SEFramework/Psf/VariablePsfStack.h"
@@ -73,18 +74,33 @@ void VariablePsfStack::setup(std::shared_ptr<CCfits::FITS> pFits) {
     // read the nrows value
     m_nrows = position_data.rows();
 
+    // Temporary vectors for reading data
+    std::vector<SeFloat> ra_values, dec_values, x_values, y_values;
+    std::vector<double> gridx_values, gridy_values;
+
     try {
       // read in all the EXT specific columns
-      position_data.column("GRIDX", false).read(m_gridx_values, 0, m_nrows);
-      position_data.column("GRIDY", false).read(m_gridy_values, 0, m_nrows);
+      position_data.column("GRIDX", false).read(gridx_values, 0, m_nrows);
+      position_data.column("GRIDY", false).read(gridy_values, 0, m_nrows);
     } catch (CCfits::Table::NoSuchColumn) {
-      position_data.column("X_CENTER", false).read(m_gridx_values, 0, m_nrows);
-      position_data.column("Y_CENTER", false).read(m_gridy_values, 0, m_nrows);
+      position_data.column("X_CENTER", false).read(gridx_values, 0, m_nrows);
+      position_data.column("Y_CENTER", false).read(gridy_values, 0, m_nrows);
+    }
+    position_data.column("RA", false).read(ra_values, 0, m_nrows);
+    position_data.column("DEC", false).read(dec_values, 0, m_nrows);
+    position_data.column("X", false).read(x_values, 0, m_nrows);
+    position_data.column("Y", false).read(y_values, 0, m_nrows);
+
+    // Populate the positions vector
+    m_positions.reserve(m_nrows);
+    for (long i = 0; i < m_nrows; ++i) {
+      m_positions.push_back({ra_values[i], dec_values[i], x_values[i], y_values[i], gridx_values[i], gridy_values[i]});
+    }
+
+    // Build KdTree for fast nearest neighbor searches
+    if (!m_positions.empty()) {
+      m_kdtree = std::make_unique<KdTree<PsfPosition>>(m_positions);
     }
-    position_data.column("RA", false).read(m_ra_values, 0, m_nrows);
-    position_data.column("DEC", false).read(m_dec_values, 0, m_nrows);
-    position_data.column("X", false).read(m_x_values, 0, m_nrows);
-    position_data.column("Y", false).read(m_y_values, 0, m_nrows);
 
   } catch (CCfits::FitsException& e) {
     throw Elements::Exception() << "Error loading stacked PSF file: " << e.message();
@@ -95,54 +111,54 @@ void VariablePsfStack::selfTest() {
   int naxis1, naxis2;
 
   // read in the min/max grid values in x/y
-  const auto x_grid_minmax = std::minmax_element(begin(m_gridx_values), end(m_gridx_values));
-  const auto y_grid_minmax = std::minmax_element(begin(m_gridy_values), end(m_gridy_values));
+  auto x_grid_minmax = std::minmax_element(m_positions.begin(), m_positions.end(),
+    [](const PsfPosition& a, const PsfPosition& b) { return a.gridx < b.gridx; });
+  auto y_grid_minmax = std::minmax_element(m_positions.begin(), m_positions.end(),
+    [](const PsfPosition& a, const PsfPosition& b) { return a.gridy < b.gridy; });
 
   // read the image size
   m_pFits->extension(1).readKey("NAXIS1", naxis1);
   m_pFits->extension(1).readKey("NAXIS2", naxis2);
 
   // make sure all PSF in the grid are there
-  if (*x_grid_minmax.first - m_grid_offset < 1)
-    throw Elements::Exception() << "The PSF at the smallest x-grid starts at: " << *x_grid_minmax.first - m_grid_offset;
-  if (*y_grid_minmax.first - m_grid_offset < 1)
-    throw Elements::Exception() << "The PSF at the smallest y-grid starts at: " << *y_grid_minmax.first - m_grid_offset;
-  if (*x_grid_minmax.second + m_grid_offset > naxis1)
-    throw Elements::Exception() << "The PSF at the largest x-grid is too large: " << *x_grid_minmax.second + m_grid_offset
+  if (x_grid_minmax.first->gridx - m_grid_offset < 1)
+    throw Elements::Exception() << "The PSF at the smallest x-grid starts at: " << x_grid_minmax.first->gridx - m_grid_offset;
+  if (y_grid_minmax.first->gridy - m_grid_offset < 1)
+    throw Elements::Exception() << "The PSF at the smallest y-grid starts at: " << y_grid_minmax.first->gridy - m_grid_offset;
+  if (x_grid_minmax.second->gridx + m_grid_offset > naxis1)
+    throw Elements::Exception() << "The PSF at the largest x-grid is too large: " << x_grid_minmax.second->gridx + m_grid_offset
                                 << " NAXIS1: " << naxis1;
-  if (*y_grid_minmax.second + m_grid_offset > naxis2)
-    throw Elements::Exception() << "The PSF at the largest y-grid is too large: " << *y_grid_minmax.second + m_grid_offset
-                                << " NAXIS2: " << naxis1;
+  if (y_grid_minmax.second->gridy + m_grid_offset > naxis2)
+    throw Elements::Exception() << "The PSF at the largest y-grid is too large: " << y_grid_minmax.second->gridy + m_grid_offset
+                                << " NAXIS2: " << naxis2;
 }
 
 std::shared_ptr<VectorImage<SeFloat>> VariablePsfStack::getPsf(const std::vector<double>& values) const {
-  long   index_min_distance = 0;
-  double min_distance       = 1.0e+32;
-
   // make sure there are only two positions
   if (values.size() != 2)
     throw Elements::Exception() << "There can be only two positional value for the stacked PSF!";
 
-  // find the position of minimal distance
-  for (int act_index = 0; act_index < m_nrows; act_index++) {
-    double act_distance = (values[0] - m_x_values[act_index]) * (values[0] - m_x_values[act_index]) +
-                          (values[1] - m_y_values[act_index]) * (values[1] - m_y_values[act_index]);
-    if (act_distance < min_distance) {
-      index_min_distance = act_index;
-      min_distance       = act_distance;
-    }
-  }
+  // Use KdTree to find the nearest PSF position
+  KdTree<PsfPosition>::Coord coord;
+  coord.coord[0] = values[0];  // x coordinate
+  coord.coord[1] = values[1];  // y coordinate
+
+  PsfPosition nearest_position = m_kdtree->findNearest(coord);
+
+  // Calculate distance for logging
+  double min_distance = (values[0] - nearest_position.x) * (values[0] - nearest_position.x) +
+                        (values[1] - nearest_position.y) * (values[1] - nearest_position.y);
+
   // give some feedback
   stack_logger.debug() << "Distance: " << sqrt(min_distance) << " (" << values[0] << "," << values[1] << ")<-->("
-                       << m_x_values[index_min_distance] << "," << m_y_values[index_min_distance]
-                       << ") index: " << index_min_distance;
+                       << nearest_position.x << "," << nearest_position.y << ")";
 
   // get the first and last pixels for the PSF to be extracted
   // NOTE: CCfits has 1-based indices, also the last index is *included* in the reading
   // NOTE: the +0.5 forces a correct cast/ceiling
-  std::vector<long> first_vertex{long(m_gridx_values[index_min_distance]+.5) - long(m_grid_offset),  long(m_gridy_values[index_min_distance]+.5) - long(m_grid_offset)};
+  std::vector<long> first_vertex{long(nearest_position.gridx+.5) - long(m_grid_offset),  long(nearest_position.gridy+.5) - long(m_grid_offset)};
   stack_logger.debug() << "First vertex: ( " << first_vertex[0] << ", " << first_vertex[1] << ") First vertex alternative: " <<
-      m_gridx_values[index_min_distance]-m_grid_offset << " " << m_gridy_values[index_min_distance]-m_grid_offset <<
+      nearest_position.gridx-m_grid_offset << " " << nearest_position.gridy-m_grid_offset <<
       " grid offset:" << m_grid_offset;
 
   std::vector<long> last_vertex{first_vertex[0] + long(m_psf_size) - 1, first_vertex[1] +long( m_psf_size) - 1};
@@ -155,8 +171,6 @@ std::shared_ptr<VectorImage<SeFloat>> VariablePsfStack::getPsf(const std::vector
     m_pFits->extension(1).read(stamp_data, first_vertex, last_vertex, stride);
   }
 
-  //stack_logger.info() << "DDD ( " << first_vertex[0] << ", " << first_vertex[1] << ") --> ( " << last_vertex[0] << ", " << last_vertex[1] << "): " << stamp_data.size();
-
   // create and return the psf image
   return VectorImage<SeFloat>::create(m_psf_size, m_psf_size, std::begin(stamp_data), std::end(stamp_data));
 }

SEUtils/CMakeLists.txt

@@ -66,6 +66,9 @@ elements_add_unit_test(Misc_test tests/src/Misc_test.cpp
 elements_add_unit_test(QuadTree_test tests/src/QuadTree_test.cpp
                      LINK_LIBRARIES SEUtils
                      TYPE Boost)
+elements_add_unit_test(KdTree_test tests/src/KdTree_test.cpp
+                     LINK_LIBRARIES SEUtils
+                     TYPE Boost)
 
 if(GMOCK_FOUND)
 elements_add_unit_test(Observable_test tests/src/Observable_test.cpp

SEUtils/SEUtils/KdTree.h

@@ -21,6 +21,7 @@
 #include <vector>
 #include <memory>
 #include <algorithm>
+#include <limits>
 
 namespace SourceXtractor {
 
@@ -49,6 +50,7 @@ class KdTree {
 
   explicit KdTree(const std::vector<T>& data);
   std::vector<T> findPointsWithinRadius(Coord coord, double radius) const;
+  T findNearest(Coord coord) const;
 
 private:
   class Node;