PixelHistogram (#317)

Multi threaded filling of per pixel histograms for example for detector calibration

1. PixelHistogram - Generic variant expects already pedestal subtracted
data
2. PedestalTrackingHistogram - Terrible name, useful class. Keeps it's
own pedestal and does conversion and pedestal tracking in the worker
threads.

---------

Co-authored-by: Lars Erik Fröjd <froejdh_e@pc-jungfrau-02.psi.ch>

python/aare/__init__.py

@@ -33,7 +33,7 @@ from .CtbRawFile import CtbRawFile
 from .RawFile import RawFile
 from .ScanParameters import ScanParameters
 
-from .utils import random_pixels, random_pixel, flat_list, add_colorbar
+from .utils import random_pixels, random_pixel, flat_list, add_colorbar, Timer
 
 
 #make functions available in the top level API
@@ -44,3 +44,13 @@ from ._aare import apply_calibration, count_switching_pixels
 from ._aare import calculate_pedestal, calculate_pedestal_float, calculate_pedestal_g0, calculate_pedestal_g0_float
 
 from ._aare import VarClusterFinder
+from ._aare import (
+    PedestalTrackingPixelHistogram,
+    PixelHistogram,
+    PixelHistogram_d,
+    PixelHistogram_f,
+    PixelHistogram_u8,
+    PixelHistogram_u16,
+    PixelHistogram_u32,
+    PixelHistogram_u64,
+)

python/aare/utils.py

@@ -2,6 +2,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import make_axes_locatable
+import time
 
 def random_pixels(n_pixels, xmin=0, xmax=512, ymin=0, ymax=1024):
     """Return a list of random pixels.
@@ -34,4 +35,18 @@ def add_colorbar(ax, im, size="5%", pad=0.05):
     divider = make_axes_locatable(ax)
     cax = divider.append_axes("right", size=size, pad=pad)
     plt.colorbar(im, cax=cax)
-    return ax, im, cax
+    return ax, im, cax
+
+
+class Timer:
+    def __init__(self, label="Elapsed time:", verbose=True):
+        self.label = label
+        self.verbose = verbose
+    def __enter__(self):
+        self.start = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.elapsed = time.perf_counter() - self.start
+        if self.verbose:
+            print(f"{self.label} {self.elapsed:.3f}s")

python/src/bind_PedestalTrackingPixelHistogram.hpp

@@ -0,0 +1,245 @@
+// SPDX-License-Identifier: MPL-2.0
+#include "aare/hist/PedestalTrackingPixelHistogram.hpp"
+#include "np_helper.hpp"
+
+#include <cstdint>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+using namespace ::aare;
+
+void define_pedestal_tracking_pixel_histogram_bindings(py::module &m) {
+    py::class_<PedestalTrackingPixelHistogram>(
+        m, "PedestalTrackingPixelHistogram",
+        "A pixel-wise histogram of frame - pedestal residuals, with a "
+        "per-pixel running pedestal estimate sharded across worker threads")
+        .def(
+            py::init<int, int, int, double, double, int, std::size_t, double>(),
+            R"(
+             Initialize a PedestalTrackingPixelHistogram.
+
+             Args:
+                 rows: Number of rows in the detector
+                 cols: Number of columns in the detector
+                 n_bins: Number of histogram bins along the residual axis
+                 xmin: Minimum residual value (inclusive)
+                 xmax: Maximum residual value (exclusive)
+                 n_threads: Number of worker threads (default: 1). Each
+                     worker owns a disjoint row slice of both the
+                     pedestal and the histogram, so the partition
+                     determines per-thread memory usage.
+                 max_pending: Maximum number of frames that can be
+                     queued for asynchronous filling before
+                     fill_async() applies backpressure
+                     on the caller (default: 16).
+                 n_sigma: Sigma multiplier used as the gate for the
+                     pedestal-update side effect of
+                     fill_async(): a pixel sample is
+                     pushed back into the pedestal estimate iff
+                     ``abs(residual) < n_sigma * cached_std``. Set to
+                     ``0.0`` to disable the pedestal update and get
+                     histogram-only async behaviour (default: 1.0).
+                     Also exposed live via the ``n_sigma`` property.
+             )",
+            py::kw_only(), py::arg("rows"), py::arg("cols"), py::arg("n_bins"),
+            py::arg("xmin"), py::arg("xmax"), py::arg("n_threads") = 1,
+            py::arg("max_pending") = std::size_t{16}, py::arg("n_sigma") = 1.0)
+
+        .def(
+            "push_pedestal_no_update",
+            [](PedestalTrackingPixelHistogram &self,
+               py::array_t<PedestalTrackingPixelHistogram::FrameType, 0>
+                   frame) {
+                auto view = make_view_2d(frame);
+                self.push_pedestal_no_update(view);
+            },
+            R"(
+             Accumulate `frame` into the per-pixel running pedestal
+             estimate without refreshing the cached mean.
+
+             Use repeatedly while bootstrapping the pedestal, then call
+             update_mean() once before starting to fill the histogram.
+
+             Args:
+                 frame: A 2D numpy array of raw pixel values (dtype: uint16)
+             )",
+            py::arg("frame").noconvert())
+
+        .def("update_mean", &PedestalTrackingPixelHistogram::update_mean,
+             R"(
+             Refresh each partial pedestal's cached per-pixel mean from
+             its running sums. Drains pending async fills first, then
+             dispatches the update to the worker pool so the writes to
+             each shard happen on the same thread that reads them in
+             fill_async().
+             )",
+             py::call_guard<py::gil_scoped_release>())
+
+        .def(
+            "pedestal_mean",
+            [](const PedestalTrackingPixelHistogram &self) {
+                // pedestal_mean() flushes + locks + memcpys; do all of
+                // that without the GIL, only reacquire to wrap into a
+                // numpy array.
+                NDArray<PedestalTrackingPixelHistogram::AxisType, 2> *ptr =
+                    nullptr;
+                {
+                    py::gil_scoped_release release;
+                    ptr = new NDArray<PedestalTrackingPixelHistogram::AxisType,
+                                      2>(self.pedestal_mean());
+                }
+                return return_image_data(ptr);
+            },
+            R"(
+             Snapshot the per-pixel pedestal mean stitched together
+             from all shards.
+
+             Returns:
+                 A 2D numpy array (rows x cols, dtype: float64)
+                 containing the current cached pedestal mean.
+             )")
+
+        .def(
+            "fill_async",
+            [](PedestalTrackingPixelHistogram &self,
+               py::array_t<PedestalTrackingPixelHistogram::FrameType, 0>
+                   image) {
+                // Copy the numpy buffer into an owned NDArray while we
+                // still hold the GIL so we don't depend on the array's
+                // backing storage outliving this call.
+                auto view = make_view_2d(image);
+                NDArray<PedestalTrackingPixelHistogram::FrameType, 2> owned(
+                    view);
+                // Release the GIL while enqueueing -
+                // fill_async can block on backpressure
+                // when the queue is full.
+                py::gil_scoped_release release;
+                self.fill_async(std::move(owned));
+            },
+            R"(
+             Submit an image for asynchronous filling with sigma-clipped
+             pedestal tracking.
+
+             For each pixel the worker pool:
+               * histograms the pedestal-subtracted residual when it
+                 falls in ``[xmin, xmax)``, and
+               * additionally pushes the raw pixel value back into the
+                 per-thread pedestal estimate when
+                 ``abs(residual) < n_sigma * cached_std`` (the
+                 sigma-clipped pedestal-update gate).
+
+             The cached std is populated by ``update_mean()``, so
+             ``push_pedestal_no_update()`` + ``update_mean()`` must have
+             run at least once for the pedestal-update side effect to
+             fire. Setting ``n_sigma = 0`` disables the side effect and
+             recovers plain histogram-only async filling.
+
+             The image is copied into an internal buffer before this call
+             returns, so the caller may mutate or free the numpy array
+             immediately. If the internal queue is full this call blocks
+             (with the GIL released) until a slot becomes available.
+
+             Args:
+                 image: A 2D numpy array of raw pixel values (dtype: uint16)
+             )",
+            py::arg("image").noconvert())
+
+        .def("fill_from_file", &PedestalTrackingPixelHistogram::fill_from_file,
+             R"(
+             Fill the histogram from a file.
+
+             Args:
+                 file_path: Path to the file to fill from
+                 max_frames: Maximum number of frames to fill from the file (default: -1)
+             )",
+             py::call_guard<py::gil_scoped_release>(), py::arg("fname"),
+             py::arg("max_frames") = -1, py::arg("verbose") = false)
+        .def("process_pedestal_file",
+             &PedestalTrackingPixelHistogram::process_pedestal_file,
+             R"(
+             Process a pedestal file.
+
+             Args:
+                 file_path: Path to the file to process
+                 max_frames: Maximum number of frames to process from the file (default: -1)
+             )",
+             py::call_guard<py::gil_scoped_release>(), py::arg("fname"),
+             py::arg("max_frames") = -1, py::arg("verbose") = false)
+        .def_property("n_sigma", &PedestalTrackingPixelHistogram::n_sigma,
+                      &PedestalTrackingPixelHistogram::set_n_sigma,
+                      R"(
+             Sigma multiplier used as the pedestal-update gate in
+             fill_async(). Atomic; safe to read or write
+             from any thread. Setting it to 0.0 disables the pedestal
+             update entirely. The new value takes effect on subsequent
+             per-pixel evaluations inside the worker pool.
+             )")
+
+        .def("flush", &PedestalTrackingPixelHistogram::flush,
+             R"(
+             Block until all images submitted via
+             fill_async() have been merged into the
+             accumulators. Cheap when nothing is pending.
+             )",
+             py::call_guard<py::gil_scoped_release>())
+
+        .def(
+            "values",
+            [](const PedestalTrackingPixelHistogram &self) {
+                // values() implicitly flushes - release the GIL while it
+                // does so. Allocation/copy into the NDArray runs without
+                // the GIL too; only the numpy wrapping needs it.
+                NDArray<PedestalTrackingPixelHistogram::StorageType, 3> *ptr =
+                    nullptr;
+                {
+                    py::gil_scoped_release release;
+                    ptr =
+                        new NDArray<PedestalTrackingPixelHistogram::StorageType,
+                                    3>(self.values());
+                }
+                return return_image_data(ptr);
+            },
+            R"(
+             Get the histogram data as a numpy array.
+
+             Implicitly flushes any pending asynchronous fills before
+             returning, so the snapshot is consistent with everything
+             submitted up to this call.
+
+             Returns:
+                 A 3D numpy array (rows x cols x n_bins, dtype: uint16)
+                 containing the histogram bins for each pixel.
+             )")
+
+        .def(
+            "bin_centers",
+            [](const PedestalTrackingPixelHistogram &self) {
+                auto ptr =
+                    new NDArray<PedestalTrackingPixelHistogram::AxisType, 1>(
+                        self.bin_centers());
+                return return_image_data(ptr);
+            },
+            R"(
+             Get the bin centers along the residual axis.
+
+             Returns:
+                 A 1D numpy array (dtype: float32) of bin center values.
+             )")
+
+        .def(
+            "bin_edges",
+            [](const PedestalTrackingPixelHistogram &self) {
+                auto ptr =
+                    new NDArray<PedestalTrackingPixelHistogram::AxisType, 1>(
+                        self.bin_edges());
+                return return_image_data(ptr);
+            },
+            R"(
+             Get the bin edges along the residual axis.
+
+             Returns:
+                 A 1D numpy array (dtype: float32) of bin edge values.
+             )");
+}

python/src/bind_PixelHistogram.hpp

@@ -0,0 +1,147 @@
+// SPDX-License-Identifier: MPL-2.0
+#include "aare/hist/PixelHistogram.hpp"
+#include "np_helper.hpp"
+
+#include <cstddef>
+#include <cstdint>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <string>
+
+namespace py = pybind11;
+using namespace ::aare;
+
+namespace {
+
+template <typename StorageType>
+void define_pixel_histogram_binding(py::module &m, const char *class_name,
+                                    const char *storage_dtype) {
+    using Hist = PixelHistogram<StorageType, double>;
+
+    const std::string doc =
+        std::string("A histogram for pixel-wise statistics with float64 input "
+                    "axis and ") +
+        storage_dtype + " bin storage";
+
+    py::class_<Hist>(m, class_name, doc.c_str())
+        .def(py::init<int, int, int, double, double, int, std::size_t>(),
+             R"(
+             Initialize a PixelHistogram.
+
+             Args:
+                 rows: Number of rows in the detector
+                 cols: Number of columns in the detector
+                 n_bins: Number of histogram bins
+                 xmin: Minimum value for histogram range
+                 xmax: Maximum value for histogram range
+                 n_threads: Number of threads for parallel filling (default: 1)
+                 max_pending: Maximum number of images that can be queued for
+                     asynchronous filling before fill_async() applies
+                     backpressure on the caller (default: 16)
+             )",
+             py::kw_only(), py::arg("rows"), py::arg("cols"), py::arg("n_bins"),
+             py::arg("xmin"), py::arg("xmax"), py::arg("n_threads") = 1,
+             py::arg("max_pending") = std::size_t{16})
+
+        .def(
+            "fill_async",
+            [](Hist &self, py::array_t<double, 0> image) {
+                // Copy the numpy buffer into an owned NDArray while we
+                // still hold the GIL so we don't depend on the array's
+                // backing storage outliving this call.
+                auto view = make_view_2d(image);
+                NDArray<double, 2> owned(view);
+                // Release the GIL while enqueueing - fill_async can block
+                // on backpressure when the queue is full.
+                py::gil_scoped_release release;
+                self.fill_async(std::move(owned));
+            },
+            R"(
+             Submit an image for asynchronous filling.
+
+             The image is copied into an internal buffer before this call
+             returns, so the caller may mutate or free the numpy array
+             immediately. The actual histogram update happens on a
+             background thread. If the internal queue is full this call
+             blocks (with the GIL released) until a slot becomes available.
+
+             Args:
+                 image: A 2D numpy array of pixel values (dtype: float64)
+             )",
+            py::arg("image").noconvert())
+
+        .def("flush", &Hist::flush,
+             R"(
+             Block until all images submitted via fill_async() have been
+             merged into the accumulators. Cheap when nothing is pending.
+             )",
+             py::call_guard<py::gil_scoped_release>())
+
+        .def(
+            "values",
+            [](const Hist &self) {
+                // values() implicitly flushes - release the GIL while it
+                // does so. Allocation/copy into the NDArray runs without
+                // the GIL too; only the numpy wrapping needs it.
+                NDArray<StorageType, 3> *ptr = nullptr;
+                {
+                    py::gil_scoped_release release;
+                    ptr = new NDArray<StorageType, 3>(self.values());
+                }
+                return return_image_data(ptr);
+            },
+            R"(
+             Get the histogram data as a numpy array.
+
+             Implicitly flushes any pending asynchronous fills before
+             returning, so the snapshot is consistent with everything
+             submitted up to this call.
+
+             Returns:
+                 A 3D numpy array containing the histogram bins for each pixel
+             )")
+
+        .def(
+            "bin_centers",
+            [](const Hist &self) {
+                auto ptr = new NDArray<double, 1>(self.bin_centers());
+                return return_image_data(ptr);
+            },
+            R"(
+             Get the bin centers along the value axis.
+
+             Returns:
+                 A 1D numpy array containing the center values for each histogram bin
+             )")
+        .def(
+            "bin_edges",
+            [](const Hist &self) {
+                auto ptr = new NDArray<double, 1>(self.bin_edges());
+                return return_image_data(ptr);
+            },
+            R"(
+             Get the bin edges along the value axis.
+
+             Returns:
+                 A 1D numpy array containing the edge values for the histogram bins
+             )");
+}
+
+} // namespace
+
+void define_pixel_histogram_bindings(py::module &m) {
+    define_pixel_histogram_binding<double>(m, "PixelHistogram_d", "float64");
+    define_pixel_histogram_binding<float>(m, "PixelHistogram_f", "float32");
+    define_pixel_histogram_binding<std::uint64_t>(m, "PixelHistogram_u64",
+                                                  "uint64");
+    define_pixel_histogram_binding<std::uint32_t>(m, "PixelHistogram_u32",
+                                                  "uint32");
+    define_pixel_histogram_binding<std::uint16_t>(m, "PixelHistogram_u16",
+                                                  "uint16");
+    define_pixel_histogram_binding<std::uint8_t>(m, "PixelHistogram_u8",
+                                                 "uint8");
+
+    // Backwards-compatible alias for the generic Python class name.
+    m.attr("PixelHistogram") = m.attr("PixelHistogram_d");
+}

python/src/module.cpp

@@ -12,6 +12,8 @@
 #include "bind_Defs.hpp"
 #include "bind_Eta.hpp"
 #include "bind_Interpolator.hpp"
+#include "bind_PedestalTrackingPixelHistogram.hpp"
+#include "bind_PixelHistogram.hpp"
 #include "bind_PixelMap.hpp"
 #include "bind_RawFile.hpp"
 #include "bind_calibration.hpp"
@@ -64,6 +66,8 @@ PYBIND11_MODULE(_aare, m) {
     define_raw_master_file_bindings(m);
     define_var_cluster_finder_bindings(m);
     define_pixel_map_bindings(m);
+    define_pixel_histogram_bindings(m);
+    define_pedestal_tracking_pixel_histogram_bindings(m);
     define_pedestal_bindings<double>(m, "Pedestal_d");
     define_pedestal_bindings<float>(m, "Pedestal_f");
     define_fit_bindings(m);

python/src/pedestal.hpp

@@ -21,6 +21,23 @@ void define_pedestal_bindings(py::module &m, const std::string &name) {
                  *mea = self.mean();
                  return return_image_data(mea);
              })
+        .def("view",
+             [](py::object self_py) {
+                 auto &self = self_py.cast<Pedestal<SUM_TYPE> &>();
+                 auto v = self.view();
+                 std::array<py::ssize_t, 2> shape{
+                     static_cast<py::ssize_t>(v.shape(0)),
+                     static_cast<py::ssize_t>(v.shape(1))};
+                 std::array<py::ssize_t, 2> byte_strides{
+                     static_cast<py::ssize_t>(v.strides()[0]) *
+                         static_cast<py::ssize_t>(sizeof(SUM_TYPE)),
+                     static_cast<py::ssize_t>(v.strides()[1]) *
+                         static_cast<py::ssize_t>(sizeof(SUM_TYPE))};
+                 auto arr = py::array_t<SUM_TYPE>(shape, byte_strides, v.data(),
+                                                  self_py);
+                 arr.attr("setflags")(py::arg("write") = false);
+                 return arr;
+             })
         .def("variance",
              [](Pedestal<SUM_TYPE> &self) {
                  auto var = new NDArray<SUM_TYPE, 2>{};
@@ -49,6 +66,16 @@ void define_pedestal_bindings(py::module &m, const std::string &name) {
                  auto v = make_view_2d(f);
                  pedestal.push(v);
              })
+        .def(
+            "push_with_threshold",
+            [](Pedestal<SUM_TYPE> &pedestal,
+               py::array_t<uint16_t, py::array::c_style> &f,
+               py::array_t<SUM_TYPE, py::array::c_style> &threshold) {
+                auto frame_view = make_view_2d(f);
+                auto threshold_view = make_view_2d(threshold);
+                pedestal.push_with_threshold(frame_view, threshold_view);
+            },
+            py::arg("frame").noconvert(), py::arg("threshold").noconvert())
         .def(
             "push_no_update",
             [](Pedestal<SUM_TYPE> &pedestal,