WIP

python/aare/__init__.py

@@ -5,7 +5,14 @@ from . import _aare
 from . import transform
 
 from ._aare import File, RawMasterFile, RawSubFile, JungfrauDataFile
-from ._aare import Pedestal_d, Pedestal_f, ClusterFinder_Cluster3x3i, VarClusterFinder
+from ._aare import (
+    FastPedestal_d,
+    FastPedestal_f,
+    Pedestal_d,
+    Pedestal_f,
+    ClusterFinder_Cluster3x3i,
+    VarClusterFinder,
+)
 from ._aare import DetectorType, ReadoutMode 
 from ._aare import hitmap
 from ._aare import ROI

python/src/fast_pedestal.hpp

@@ -0,0 +1,117 @@
+// 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(), py::buffer_protocol())
+        .def(py::init<int, int, int>())
+        .def(py::init<int, int>())
+        .def("mean",
+             [](FastPedestal<SUM_TYPE> &self) {
+                 auto mean = new NDArray<SUM_TYPE, 2>{};
+                 *mean = self.mean();
+                 return return_image_data(mean);
+             })
+        .def("view",
+             [](py::object self_py) {
+                 auto &self = self_py.cast<FastPedestal<SUM_TYPE> &>();
+                 auto view = self.view();
+                 std::array<py::ssize_t, 2> shape{
+                     static_cast<py::ssize_t>(view.shape(0)),
+                     static_cast<py::ssize_t>(view.shape(1))};
+                 std::array<py::ssize_t, 2> byte_strides{
+                     static_cast<py::ssize_t>(view.strides()[0]) *
+                         static_cast<py::ssize_t>(sizeof(SUM_TYPE)),
+                     static_cast<py::ssize_t>(view.strides()[1]) *
+                         static_cast<py::ssize_t>(sizeof(SUM_TYPE))};
+                 auto array = py::array_t<SUM_TYPE>(shape, byte_strides,
+                                                    view.data(), self_py);
+                 array.attr("setflags")(py::arg("write") = false);
+                 return array;
+             })
+        .def("variance",
+             [](FastPedestal<SUM_TYPE> &self) {
+                 auto variance = new NDArray<SUM_TYPE, 2>{};
+                 *variance = self.variance();
+                 return return_image_data(variance);
+             })
+        .def("std",
+             [](FastPedestal<SUM_TYPE> &self) {
+                 auto standard_deviation = new NDArray<SUM_TYPE, 2>{};
+                 *standard_deviation = self.std();
+                 return return_image_data(standard_deviation);
+             })
+        .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))
+        .def_property_readonly("rows", &FastPedestal<SUM_TYPE>::rows)
+        .def_property_readonly("cols", &FastPedestal<SUM_TYPE>::cols)
+        .def_property_readonly("cur_samples",
+                               &FastPedestal<SUM_TYPE>::cur_samples)
+        .def_property_readonly("ready", &FastPedestal<SUM_TYPE>::ready)
+        .def_property_readonly("n_samples", &FastPedestal<SUM_TYPE>::n_samples)
+        // .def_property_readonly("sum", &FastPedestal<SUM_TYPE>::get_sum)
+        .def("clone",
+             [](FastPedestal<SUM_TYPE> &pedestal) {
+                 return FastPedestal<SUM_TYPE>(pedestal);
+             })
+        .def(
+            "push",
+            [](FastPedestal<SUM_TYPE> &pedestal,
+               py::array_t<uint16_t, py::array::c_style> &frame) {
+                pedestal.push(make_view_2d(frame));
+            },
+            py::arg("frame").noconvert())
+        .def(
+            "push_init",
+            [](FastPedestal<SUM_TYPE> &pedestal,
+               py::array_t<uint16_t, py::array::c_style> &frame) {
+                pedestal.push_init(make_view_2d(frame));
+            },
+            py::arg("frame").noconvert())
+        // .def(
+        //     "push_no_update",
+        //     [](FastPedestal<SUM_TYPE> &pedestal,
+        //        py::array_t<uint16_t, py::array::c_style> &frame) {
+        //         pedestal.push_no_update(make_view_2d(frame));
+        //     },
+        //     py::arg("frame").noconvert())
+        .def("update_mean", &FastPedestal<SUM_TYPE>::update_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);
+        });
+}

python/src/module.cpp

@@ -20,6 +20,7 @@
 
 // TODO! migrate the other names
 #include "ctb_raw_file.hpp"
+#include "fast_pedestal.hpp"
 #include "file.hpp"
 #include "fit.hpp"
 #include "jungfrau_data_file.hpp"
@@ -70,6 +71,8 @@ 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_fast_pedestal_bindings<double>(m, "FastPedestal_d");
+    define_fast_pedestal_bindings<float>(m, "FastPedestal_f");
     define_fit_bindings(m);
     define_interpolation_bindings(m);
     define_jungfrau_data_file_io_bindings(m);

python/tests/test_FastPedestal.py

@@ -0,0 +1,56 @@
+import numpy as np
+import pytest
+
+from aare import FastPedestal_d, FastPedestal_f
+
+
+@pytest.mark.parametrize(
+    ("pedestal_type", "expected_dtype"),
+    [(FastPedestal_d, np.float64), (FastPedestal_f, np.float32)],
+)
+def test_fast_pedestal_initialization(pedestal_type, expected_dtype):
+    pedestal = pedestal_type(2, 3, 2)
+    first = np.array([[2, 4, 6], [8, 10, 12]], dtype=np.uint16)
+    second = np.array([[4, 6, 8], [10, 12, 14]], dtype=np.uint16)
+
+    pedestal.push_init(first)
+    pedestal.push_init(second)
+    pedestal.update_mean()
+
+    expected_mean = np.array(
+        [[3, 5, 7], [9, 11, 13]], dtype=expected_dtype
+    )
+    np.testing.assert_array_equal(pedestal.mean(), expected_mean)
+    np.testing.assert_array_equal(pedestal.std(), np.ones((2, 3)))
+
+
+def test_fast_pedestal_steady_state_push():
+    pedestal = FastPedestal_d(1, 2, 2)
+    pedestal.push_init(np.array([[2, 4]], dtype=np.uint16))
+    pedestal.push_init(np.array([[4, 6]], dtype=np.uint16))
+    pedestal.update_mean()
+
+    pedestal.push(np.array([[6, 8]], dtype=np.uint16))
+
+    np.testing.assert_array_equal(pedestal.mean(), [[4.5, 6.5]])
+
+
+def test_fast_pedestal_exposes_read_only_buffer_and_subtraction():
+    pedestal = FastPedestal_d(1, 2, 1)
+    pedestal.push_init(np.array([[2, 4]], dtype=np.uint16))
+    pedestal.update_mean()
+
+    view = np.asarray(pedestal)
+    result = np.array([[12, 14]], dtype=np.uint16) - pedestal
+
+    np.testing.assert_array_equal(view, [[2, 4]])
+    np.testing.assert_array_equal(result, [[10, 10]])
+    assert np.shares_memory(view, pedestal.view())
+    assert not view.flags.writeable
+
+
+def test_fast_pedestal_rejects_wrong_shape():
+    pedestal = FastPedestal_d(2, 3)
+
+    with pytest.raises(RuntimeError, match="shape"):
+        pedestal.push_init(np.zeros((2, 2), dtype=np.uint16))