FastPedestal + optimizations in the cluster finder (#341)

Optimized pedestal tracking with FastPedestal and performance improvements to the cluster finder. 
- ~2x faster hitting 14k FPS on benchmark dataset on 16 core threadripper pro


**Changes**
- No bounds check on interior of frame. (only within half of a cluster)
(~15% improvement)
- Cache threshold (std*n_rms) Biggest improvement comes from not
computing std for each access
- Cache pedestal subtracted frame
- Switched to FastPedestal which assumes we reached steady state

**Things tried but rejected:**
- Always store cluster values, commit on val==max (~15% drop in frame
rate)

**Options**
- using 16 bit clusters and 16 bit pedestal. (slows down the single
threaded case but allow us to reach 12k with multi threading)

**Other notes on performance**
- Passing in memory frames from python and doing nothing: 0.8us/frame
- Passing + pedestal subtraction: 28us/frame
- Passing + pedestal + cluster finding (no store): 652us/frame
- Passing + pedestal + cluster finding: 886 us/frame

python/src/bind_Cluster.hpp

@@ -10,7 +10,6 @@
 #include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
-using pd_type = double;
 
 using namespace aare;
 

python/src/bind_ClusterCollector.hpp

@@ -6,6 +6,8 @@
 #include "aare/ClusterVector.hpp"
 #include "aare/NDView.hpp"
 #include "aare/Pedestal.hpp"
+
+#include "module_config.hpp"
 #include "np_helper.hpp"
 
 #include <cstdint>
@@ -15,7 +17,6 @@
 #include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
-using pd_type = double;
 
 using namespace aare;
 
@@ -30,7 +31,7 @@ void define_ClusterCollector(py::module &m, const std::string &typestr) {
     using ClusterType = Cluster<T, ClusterSizeX, ClusterSizeY, CoordType>;
 
     py::class_<ClusterCollector<ClusterType>>(m, class_name.c_str())
-        .def(py::init<ClusterFinderMT<ClusterType, uint16_t, double> *>())
+        .def(py::init<ClusterFinderMT<ClusterType, uint16_t, pd_type> *>())
         .def("stop", &ClusterCollector<ClusterType>::stop)
         .def(
             "steal_clusters",

python/src/bind_ClusterFileSink.hpp

@@ -15,8 +15,6 @@
 #include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
-using pd_type = double;
-
 using namespace aare;
 
 #pragma GCC diagnostic push
@@ -29,6 +27,8 @@ void define_ClusterFileSink(py::module &m, const std::string &typestr) {
 
     using ClusterType = Cluster<T, ClusterSizeX, ClusterSizeY, CoordType>;
 
+    // TODO! adapt to set pedestal type (needs templating of ClusterFileSink)
+    // or maybe access through base class?
     py::class_<ClusterFileSink<ClusterType>>(m, class_name.c_str())
         .def(py::init<ClusterFinderMT<ClusterType, uint16_t, double> *,
                       const std::filesystem::path &>())

python/src/bind_ClusterFinder.hpp

@@ -6,6 +6,8 @@
 #include "aare/ClusterVector.hpp"
 #include "aare/NDView.hpp"
 #include "aare/Pedestal.hpp"
+
+#include "module_config.hpp"
 #include "np_helper.hpp"
 
 #include <cstdint>
@@ -15,7 +17,6 @@
 #include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
-using pd_type = double;
 
 using namespace aare;
 
@@ -48,6 +49,8 @@ void define_ClusterFinder(py::module &m, const std::string &typestr) {
              })
         .def("clear_pedestal",
              &ClusterFinder<ClusterType, uint16_t, pd_type>::clear_pedestal)
+        .def("update_threshold",
+             &ClusterFinder<ClusterType, uint16_t, pd_type>::update_threshold)
         .def_property_readonly(
             "pedestal",
             [](ClusterFinder<ClusterType, uint16_t, pd_type> &self) {

python/src/bind_ClusterFinderMT.hpp

@@ -6,6 +6,8 @@
 #include "aare/ClusterVector.hpp"
 #include "aare/NDView.hpp"
 #include "aare/Pedestal.hpp"
+
+#include "module_config.hpp"
 #include "np_helper.hpp"
 
 #include <cstdint>
@@ -15,7 +17,6 @@
 #include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
-using pd_type = double;
 
 using namespace aare;
 
@@ -31,9 +32,10 @@ void define_ClusterFinderMT(py::module &m, const std::string &typestr) {
 
     py::class_<ClusterFinderMT<ClusterType, uint16_t, pd_type>>(
         m, class_name.c_str())
-        .def(py::init<Shape<2>, pd_type, size_t, size_t>(),
+        .def(py::init<Shape<2>, pd_type, size_t, size_t, size_t>(),
              py::arg("image_size"), py::arg("n_sigma") = 5.0,
-             py::arg("capacity") = 2048, py::arg("n_threads") = 3)
+             py::arg("capacity") = 2048, py::arg("n_threads") = 3,
+             py::arg("queue_depth") = 16)
         .def("push_pedestal_frame",
              [](ClusterFinderMT<ClusterType, uint16_t, pd_type> &self,
                 py::array_t<uint16_t> frame) {
@@ -46,7 +48,6 @@ void define_ClusterFinderMT(py::module &m, const std::string &typestr) {
                py::array_t<uint16_t> frame, uint64_t frame_number) {
                 auto view = make_view_2d(frame);
                 self.find_clusters(view, frame_number);
-                return;
             },
             py::arg(), py::arg("frame_number") = 0)
         .def_property_readonly(
@@ -56,6 +57,8 @@ void define_ClusterFinderMT(py::module &m, const std::string &typestr) {
             })
         .def("clear_pedestal",
              &ClusterFinderMT<ClusterType, uint16_t, pd_type>::clear_pedestal)
+        .def("update_threshold",
+             &ClusterFinderMT<ClusterType, uint16_t, pd_type>::update_threshold)
         .def("sync", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::sync)
         .def("stop", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::stop)
         .def("start", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::start)

python/src/bind_ClusterVector.hpp

@@ -15,7 +15,6 @@
 #include <pybind11/stl_bind.h>
 
 namespace py = pybind11;
-using pd_type = double;
 
 using namespace aare;
 

python/src/bind_FastPedestal.hpp

@@ -0,0 +1,147 @@
+// SPDX-License-Identifier: MPL-2.0
+
+#include "aare/FastPedestal.hpp"
+#include "np_helper.hpp"
+
+#include <cstdint>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+template <typename SUM_TYPE>
+void define_fast_pedestal_bindings(py::module &m, const std::string &name) {
+
+    py::class_<FastPedestal<SUM_TYPE>>(
+        m, name.c_str(),
+        "Maintain a per-pixel running mean, population variance, and "
+        "standard deviation.",
+        py::buffer_protocol())
+        .def(py::init<uint32_t, uint32_t, uint32_t>(), py::arg("rows"),
+             py::arg("cols"), py::arg("n_samples"),
+             "Construct an empty pedestal. It becomes ready after n_samples "
+             "calls to add_init_frame().")
+        .def(py::init<uint32_t, uint32_t>(), py::arg("rows"), py::arg("cols"),
+             "Construct an empty pedestal with n_samples=1000.")
+
+        .def(
+            "mean",
+            [](FastPedestal<SUM_TYPE> &self) {
+                auto mean = new NDArray<SUM_TYPE, 2>{};
+                *mean = self.mean();
+                return return_image_data(mean);
+            },
+            "Return a copy of the cached mean. The pedestal must be ready.")
+        .def(
+            "var",
+            [](FastPedestal<SUM_TYPE> &self) {
+                auto variance = new NDArray<SUM_TYPE, 2>{};
+                *variance = self.variance();
+                return return_image_data(variance);
+            },
+            "Return the population variance, normalized by n_samples, as a "
+            "NumPy array. The pedestal must be ready.")
+        .def(
+            "std",
+            [](FastPedestal<SUM_TYPE> &self) {
+                auto standard_deviation = new NDArray<SUM_TYPE, 2>{};
+                *standard_deviation = self.std();
+                return return_image_data(standard_deviation);
+            },
+            "Return the population standard deviation as a NumPy array. The "
+            "pedestal must be ready.")
+        .def(
+            "view",
+            [](py::object self_py) {
+                return py::module_::import("numpy").attr("asarray")(self_py);
+            },
+            "Return a non-owning, non-writable NumPy view of the cached mean. "
+            "The pedestal must be ready.")
+
+        // We need to buffer protocol to allow for numpy operations using the
+        // pedestal mean
+        .def_buffer([](FastPedestal<SUM_TYPE> &self) {
+            auto mean = self.view();
+            return py::buffer_info(
+                const_cast<SUM_TYPE *>(mean.data()), sizeof(SUM_TYPE),
+                py::format_descriptor<SUM_TYPE>::format(), 2,
+                {static_cast<py::ssize_t>(mean.shape(0)),
+                 static_cast<py::ssize_t>(mean.shape(1))},
+                {static_cast<py::ssize_t>(mean.strides()[0] * sizeof(SUM_TYPE)),
+                 static_cast<py::ssize_t>(mean.strides()[1] *
+                                          sizeof(SUM_TYPE))},
+                true);
+        })
+        // Subtracting a FastPedestal from a NumPy array
+        .def(
+            "__array_ufunc__",
+            [](py::object self, py::object ufunc, const std::string &method,
+               py::args inputs, py::kwargs kwargs) -> py::object {
+                if (method != "__call__" || inputs.size() != 2 ||
+                    inputs[1].ptr() != self.ptr() ||
+                    py::cast<std::string>(ufunc.attr("__name__")) !=
+                        "subtract") {
+                    return py::reinterpret_borrow<py::object>(
+                        Py_NotImplemented);
+                }
+
+                auto mean =
+                    py::module_::import("builtins").attr("memoryview")(self);
+                return ufunc(inputs[0], mean, **kwargs);
+            },
+            "Support subtracting a FastPedestal from a NumPy array.")
+        .def("clear", py::overload_cast<>(&FastPedestal<SUM_TYPE>::clear),
+             "Reset all statistics and initialization state to zero.")
+        .def_property_readonly("rows", &FastPedestal<SUM_TYPE>::rows,
+                               "Number of image rows.")
+        .def_property_readonly("cols", &FastPedestal<SUM_TYPE>::cols,
+                               "Number of image columns.")
+        .def_property_readonly("cur_samples",
+                               &FastPedestal<SUM_TYPE>::cur_samples,
+                               "Number of initialization frames accumulated. "
+                               "Steady-state pushes do not change it.")
+        .def_property_readonly(
+            "ready", &FastPedestal<SUM_TYPE>::ready,
+            "Whether n_samples initialization frames have been accumulated.")
+        .def_property_readonly(
+            "n_samples", &FastPedestal<SUM_TYPE>::n_samples,
+            "Initialization frame count and steady-state update-weight "
+            "denominator.")
+        .def(
+            "clone",
+            [](FastPedestal<SUM_TYPE> &pedestal) {
+                return FastPedestal<SUM_TYPE>(pedestal);
+            },
+            "Return an independent copy of the pedestal and its state.")
+        .def(
+            "push_ema",
+            [](FastPedestal<SUM_TYPE> &pedestal,
+               py::array_t<uint16_t, py::array::c_style> &frame) {
+                pedestal.push_ema(make_view_2d(frame));
+            },
+            py::arg("frame").noconvert(),
+            "Update exponential moving average. The pedstal must  "
+            "already be ready for this update.")
+        .def(
+            "add_init_frame",
+            [](FastPedestal<SUM_TYPE> &pedestal,
+               py::array_t<uint16_t, py::array::c_style> &frame) {
+                pedestal.add_init_frame(make_view_2d(frame));
+            },
+            py::arg("frame").noconvert(),
+            "Accumulate one uint16 initialization frame. Call exactly "
+            "n_samples times to make the pedestal ready.")
+        .def_static(
+            "from_file",
+            [](const std::filesystem::path &filename, uint32_t n_samples,
+               uint32_t skip_first) {
+                return FastPedestal<SUM_TYPE>::template from_file<uint16_t>(
+                    filename, n_samples, skip_first);
+            },
+            py::arg("filename"), py::arg("n_samples") = 1000,
+            py::arg("skip_first") = 0,
+            "Create a pedestal from a uint16 file. Skip skip_first frames, use "
+            "the next n_samples for initialization, then apply every remaining "
+            "frame as a steady-state update.");
+}

python/src/bind_PedestalTrackingPixelHistogram.hpp

@@ -153,9 +153,13 @@ void define_pedestal_tracking_pixel_histogram_bindings(py::module &m) {
              Args:
                  file_path: Path to the file to fill from
                  max_frames: Maximum number of frames to fill from the file (default: -1)
+                 reader_threads: Number of parallel file reader workers (default: 2)
+                 reader_chunk_size: Frames claimed by each reader worker per batch (default: 4)
              )",
              py::call_guard<py::gil_scoped_release>(), py::arg("fname"),
-             py::arg("max_frames") = -1, py::arg("verbose") = false)
+             py::arg("max_frames") = -1, py::arg("verbose") = false,
+             py::arg("reader_threads") = std::size_t{2},
+             py::arg("reader_chunk_size") = std::size_t{4})
         .def("process_pedestal_file",
              &PedestalTrackingPixelHistogram::process_pedestal_file,
              R"(

python/src/module.cpp

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: MPL-2.0
 // Files with bindings to the different classes
 
+#include "module_config.hpp"
+
 // New style file naming
 #include "bind_Cluster.hpp"
 #include "bind_ClusterCollector.hpp"
@@ -11,6 +13,7 @@
 #include "bind_ClusterVector.hpp"
 #include "bind_Defs.hpp"
 #include "bind_Eta.hpp"
+#include "bind_FastPedestal.hpp"
 #include "bind_Interpolator.hpp"
 #include "bind_MultiThreadedFileReader.hpp"
 #include "bind_PedestalTrackingPixelHistogram.hpp"
@@ -30,6 +33,7 @@
 #include "var_cluster.hpp"
 
 // Pybind stuff
+#include <cstdint>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -75,6 +79,10 @@ PYBIND11_MODULE(_aare, m) {
     define_pedestal_tracking_pixel_histogram_bindings(m);
     define_pedestal_bindings<double>(m, "Pedestal_d");
     define_pedestal_bindings<float>(m, "Pedestal_f");
+    define_pedestal_bindings<int16_t>(m, "Pedestal_i16");
+    define_fast_pedestal_bindings<double>(m, "FastPedestal_d");
+    define_fast_pedestal_bindings<float>(m, "FastPedestal_f");
+    define_fast_pedestal_bindings<int16_t>(m, "FastPedestal_i16");
     define_fit_bindings(m);
     define_interpolation_bindings(m);
     define_jungfrau_data_file_io_bindings(m);
@@ -106,6 +114,7 @@ PYBIND11_MODULE(_aare, m) {
     DEFINE_BINDINGS_CLUSTERFINDER(int, 3, 3, uint16_t, i);
     DEFINE_BINDINGS_CLUSTERFINDER(double, 3, 3, uint16_t, d);
     DEFINE_BINDINGS_CLUSTERFINDER(float, 3, 3, uint16_t, f);
+    DEFINE_BINDINGS_CLUSTERFINDER(int16_t, 3, 3, uint16_t, i16);
 
     DEFINE_BINDINGS_CLUSTERFINDER(int, 5, 5, uint16_t, i);
     DEFINE_BINDINGS_CLUSTERFINDER(double, 5, 5, uint16_t, d);

python/src/module_config.hpp

@@ -0,0 +1,5 @@
+#pragma once
+#include <cstdint>
+
+// Configure module wide pedestal type for cluster finding
+using pd_type = double;

python/src/pedestal.hpp

@@ -52,6 +52,12 @@ void define_pedestal_bindings(py::module &m, const std::string &name) {
                  *std = self.std();
                  return return_image_data(std);
              })
+        .def("cached_std",
+             [](Pedestal<SUM_TYPE> &self) {
+                 auto standard_deviation = new NDArray<SUM_TYPE, 2>{};
+                 *standard_deviation = self.cached_std();
+                 return return_image_data(standard_deviation);
+             })
         .def(
             "__array_ufunc__",
             [](py::object self, py::object ufunc, const std::string &method,