Add gap pixel insertion for images in memory (#423)

Second part of splitting up dev/gap. Only the in-memory transform; no
reader changes.

- `GapPixels` and `ModuleGaps` configure 2 pixel gaps at the chip
boundaries inside each Jungfrau or Eiger module and, only when module
gaps are set, gaps between modules. Gap pixels get a fill value or, with
`split_counts`, half of the neighbouring double size pixel.
- `insert_gap_pixels` and `gapped_shape` apply the transform to an image
in memory. Gap positions follow detector coordinates, so a ROI receives
exactly the gaps that fall inside it.
- Python bindings, tests and docs for the above.

Chip and module sizes are fixed at 256 and 1024x512, so there is no
detector type or quad argument: an Eiger quad is 512x512 and its only
chip boundaries are at pixel 256.

Reading raw files with gaps comes in the next PR.

python/aare/__init__.py

@@ -25,6 +25,7 @@ from ._aare import Pedestal_d, Pedestal_f, ClusterFinder_Cluster3x3i, VarCluster
 from ._aare import DetectorType, FrameDiscardPolicy, ReadoutMode
 from ._aare import hitmap
 from ._aare import ROI
+from ._aare import GapPixels, ModuleGaps, insert_gap_pixels, gapped_shape
 from ._aare import xy
 from ._aare import corner 
 

python/src/bind_GapPixels.hpp

@@ -0,0 +1,174 @@
+// SPDX-License-Identifier: MPL-2.0
+#pragma once
+#include "aare/GapPixels.hpp"
+#include "aare/ROI.hpp"
+#include "aare/defs.hpp"
+#include "np_helper.hpp"
+
+#include <fmt/format.h>
+#include <optional>
+#include <pybind11/numpy.h>
+#include <pybind11/operators.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+using namespace ::aare;
+
+namespace {
+
+template <typename T>
+py::array insert_gap_pixels_typed(const py::array &image, const ROI &roi,
+                                  const GapPixels &gaps) {
+    py::array_t<T, py::array::c_style | py::array::forcecast> typed(image);
+    auto source = make_const_view_2d(typed);
+    auto *result = new NDArray<T, 2>(insert_gap_pixels<T>(source, roi, gaps));
+    return return_image_data(result);
+}
+
+} // namespace
+
+void define_gap_pixels_bindings(py::module &m) {
+    py::class_<ModuleGaps>(m, "ModuleGaps", R"(
+        Pixels inserted at each boundary between modules.
+        )")
+        .def(py::init<>())
+        .def(py::init<ssize_t, ssize_t>(), py::arg("x"), py::arg("y"))
+        .def_readwrite("x", &ModuleGaps::x)
+        .def_readwrite("y", &ModuleGaps::y)
+        .def(py::self == py::self)
+        .def(py::self != py::self)
+        .def("__repr__", [](const ModuleGaps &self) {
+            return fmt::format("ModuleGaps(x={}, y={})", self.x, self.y);
+        });
+
+    py::class_<GapPixels>(m, "GapPixels", R"(
+        Configuration for inserting gap pixels into Jungfrau and Eiger
+        images. Chip and module sizes are fixed: 256x256 pixel chips in
+        1024x512 pixel modules.
+
+        Parameters
+        ----------
+        chip_gap : int
+            Pixels inserted at each chip boundary inside a module, at most
+            1024. Source column 255 expands to the right and column 256 to
+            the left.
+        module_gaps : ModuleGaps, optional
+            Pixels inserted at module boundaries, at most 1024 each. None
+            means no pixels are inserted between modules.
+        fill_value : float
+            Value written to gap pixels, converted to the pixel type. Finite
+            values wrap for integer images; NaN, infinity and magnitudes of
+            2**63 or more raise ValueError for integer images. Module gaps
+            always receive this value.
+        split_counts : bool
+            Split the counts of the double-size pixels at chip boundaries
+            between the pixel and its gap pixel. Integer remainders go to
+            one side at random. Requires a chip gap of 2.
+        seed : int, optional
+            Seed for the random assignment of remainders.
+        )")
+        .def(
+            py::init([](ssize_t chip_gap, std::optional<ModuleGaps> module_gaps,
+                        double fill_value, bool split_counts,
+                        std::optional<uint64_t> seed) {
+                GapPixels gaps;
+                gaps.chip_gap = chip_gap;
+                gaps.module_gaps = module_gaps;
+                gaps.fill_value = fill_value;
+                gaps.split_counts = split_counts;
+                gaps.seed = seed;
+                return gaps;
+            }),
+            py::arg("chip_gap") = 2, py::arg("module_gaps") = py::none(),
+            py::arg("fill_value") = 0.0, py::arg("split_counts") = false,
+            py::arg("seed") = py::none())
+        .def_readwrite("chip_gap", &GapPixels::chip_gap)
+        .def_readwrite("module_gaps", &GapPixels::module_gaps)
+        .def_readwrite("fill_value", &GapPixels::fill_value)
+        .def_readwrite("split_counts", &GapPixels::split_counts)
+        .def_readwrite("seed", &GapPixels::seed)
+        .def(py::self == py::self)
+        .def(py::self != py::self)
+        .def("__repr__", [](const GapPixels &self) {
+            return fmt::format(
+                "GapPixels(chip_gap={}, module_gaps={}, fill_value={}, "
+                "split_counts={}, seed={})",
+                self.chip_gap,
+                self.module_gaps
+                    ? fmt::format("ModuleGaps(x={}, y={})", self.module_gaps->x,
+                                  self.module_gaps->y)
+                    : "None",
+                self.fill_value, self.split_counts ? "True" : "False",
+                self.seed ? fmt::format("{}", *self.seed) : "None");
+        });
+
+    m.def(
+        "gapped_shape",
+        [](const ROI &roi, const GapPixels &gaps) {
+            const auto shape = gapped_shape(roi, gaps);
+            return py::make_tuple(shape[0], shape[1]);
+        },
+        py::arg("roi"), py::arg("gaps") = GapPixels{}, R"(
+        Shape (rows, cols) of the gapped image of a detector coordinate ROI.
+
+        Parameters
+        ----------
+        roi : ROI
+            Detector coordinates of the image.
+        gaps : GapPixels
+            Gap configuration.
+        )");
+
+    m.def(
+        "insert_gap_pixels",
+        [](const py::array &image, const GapPixels &gaps,
+           std::optional<ROI> roi) -> py::array {
+            if (image.ndim() != 2) {
+                throw py::value_error(
+                    fmt::format("Expected a 2D image, got {}D", image.ndim()));
+            }
+            const ROI rect =
+                roi.value_or(ROI{0, image.shape(1), 0, image.shape(0)});
+            if (py::isinstance<py::array_t<uint8_t>>(image)) {
+                return insert_gap_pixels_typed<uint8_t>(image, rect, gaps);
+            }
+            if (py::isinstance<py::array_t<uint16_t>>(image)) {
+                return insert_gap_pixels_typed<uint16_t>(image, rect, gaps);
+            }
+            if (py::isinstance<py::array_t<uint32_t>>(image)) {
+                return insert_gap_pixels_typed<uint32_t>(image, rect, gaps);
+            }
+            if (py::isinstance<py::array_t<int32_t>>(image)) {
+                return insert_gap_pixels_typed<int32_t>(image, rect, gaps);
+            }
+            if (py::isinstance<py::array_t<float>>(image)) {
+                return insert_gap_pixels_typed<float>(image, rect, gaps);
+            }
+            if (py::isinstance<py::array_t<double>>(image)) {
+                return insert_gap_pixels_typed<double>(image, rect, gaps);
+            }
+            throw py::type_error(
+                "Unsupported dtype, expected uint8, uint16, uint32, int32, "
+                "float32 or float64");
+        },
+        py::arg("image"), py::arg("gaps") = GapPixels{},
+        py::arg("roi") = py::none(), R"(
+        Insert gap pixels into a 2D image.
+
+        Parameters
+        ----------
+        image : numpy.ndarray
+            Ungapped image of ``roi``.
+        gaps : GapPixels
+            Gap configuration.
+        roi : ROI, optional
+            Detector coordinates of the image. Defaults to the full image
+            starting at pixel (0, 0).
+
+        Returns
+        -------
+        numpy.ndarray
+            New array with the same dtype as ``image``.
+        )");
+}

python/src/module.cpp

@@ -14,6 +14,7 @@
 #include "bind_Defs.hpp"
 #include "bind_Eta.hpp"
 #include "bind_FastPedestal.hpp"
+#include "bind_GapPixels.hpp"
 #include "bind_Interpolator.hpp"
 #include "bind_MultiThreadedFileReader.hpp"
 #include "bind_Pedestal.hpp"
@@ -156,6 +157,7 @@ PYBIND11_MODULE(_aare, m) {
     register_calculate_3x3eta<int16_t, 3, 3, uint16_t>(m);
 
     define_defs_bindings(m);
+    define_gap_pixels_bindings(m);
 
     using Sum_index_pair_d = Sum_index_pair<double, corner>;
     PYBIND11_NUMPY_DTYPE(Sum_index_pair_d, sum, index);

python/tests/test_GapPixels.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: MPL-2.0
+import numpy as np
+import pytest
+
+from aare import GapPixels, ModuleGaps, ROI, gapped_shape, insert_gap_pixels
+
+
+def reference(image, boundaries_x, boundaries_y, fill=0):
+    """Insert gap columns and rows with numpy: boundaries map a source index to
+    the number of pixels inserted before it."""
+    cols = [index for index, count in boundaries_x for _ in range(count)]
+    rows = [index for index, count in boundaries_y for _ in range(count)]
+    result = np.insert(image, cols, fill, axis=1)
+    return np.insert(result, rows, fill, axis=0)
+
+
+def test_configuration_defaults_and_equality():
+    gaps = GapPixels()
+    assert gaps.chip_gap == 2
+    assert gaps.module_gaps is None
+    assert gaps.fill_value == 0
+    assert gaps.split_counts is False
+    assert gaps.seed is None
+    assert gaps == GapPixels()
+    assert gaps != GapPixels(fill_value=1)
+
+    gaps = GapPixels(chip_gap=4, module_gaps=ModuleGaps(8, 36), fill_value=-1, split_counts=False, seed=5)
+    assert gaps.module_gaps == ModuleGaps(x=8, y=36)
+    assert gaps.seed == 5
+    assert "ModuleGaps(x=8, y=36)" in repr(gaps)
+    gaps.module_gaps = None
+    assert gaps.module_gaps is None
+
+
+def test_gapped_shape():
+    assert gapped_shape(ROI(0, 1024, 0, 512)) == (514, 1030)
+    assert gapped_shape(ROI(0, 1024, 0, 1024)) == (1028, 1030)
+    assert gapped_shape(ROI(0, 1024, 0, 1024), GapPixels(module_gaps=ModuleGaps(8, 36))) == (1064, 1030)
+    assert gapped_shape(ROI(0, 512, 0, 512)) == (514, 514)  # Eiger quad
+    assert gapped_shape(ROI(256, 512, 0, 256)) == (256, 256)
+    assert gapped_shape(ROI(255, 257, 0, 1)) == (1, 4)
+
+
+@pytest.mark.parametrize("dtype", [np.uint8, np.uint16, np.uint32, np.int32, np.float32, np.float64])
+def test_insert_gap_pixels_matches_numpy(dtype):
+    rng = np.random.default_rng(1)
+    image = rng.integers(0, 200, size=(512, 1024)).astype(dtype)
+    result = insert_gap_pixels(image, GapPixels(fill_value=3))
+    expected = reference(image, [(256, 2), (512, 2), (768, 2)], [(256, 2)], fill=3)
+    assert result.dtype == image.dtype
+    assert result.shape == (514, 1030)
+    assert (result == expected).all()
+
+
+def test_insert_gap_pixels_without_module_gaps():
+    # two stacked modules: chip gaps at rows 256 and 768, none at the module
+    # boundary at row 512
+    image = np.arange(1024 * 1024, dtype=np.uint32).reshape(1024, 1024)
+    result = insert_gap_pixels(image)
+    expected = reference(image, [(256, 2), (512, 2), (768, 2)], [(256, 2), (768, 2)])
+    assert result.shape == (1028, 1030)
+    assert (result == expected).all()
+
+
+def test_insert_gap_pixels_with_module_gaps():
+    image = np.arange(1024 * 1024, dtype=np.uint32).reshape(1024, 1024)
+    gaps = GapPixels(module_gaps=ModuleGaps(8, 36), fill_value=9)
+    result = insert_gap_pixels(image, gaps)
+    expected = reference(image, [(256, 2), (512, 2), (768, 2)], [(256, 2), (512, 36), (768, 2)], fill=9)
+    assert result.shape == (1064, 1030)
+    assert (result == expected).all()
+
+
+def test_insert_gap_pixels_in_roi():
+    image = np.array([[11, 22]], dtype=np.uint16)
+    result = insert_gap_pixels(image, roi=ROI(255, 257, 10, 11))
+    assert result.tolist() == [[11, 0, 0, 22]]
+
+    # a gap on the ROI border is not inserted
+    image = np.arange(8, dtype=np.uint16).reshape(2, 4)
+    result = insert_gap_pixels(image, roi=ROI(256, 260, 0, 2))
+    assert (result == image).all()
+
+
+def test_insert_gap_pixels_split_counts():
+    rng = np.random.default_rng(2)
+    image = rng.integers(0, 1000, size=(512, 1024)).astype(np.uint32)
+    gaps = GapPixels(split_counts=True, seed=3)
+    first = insert_gap_pixels(image, gaps)
+    second = insert_gap_pixels(image, gaps)
+    assert first.sum() == image.sum()
+    assert (first == second).all()
+    # edge pixel and its gap pixel hold floor and ceil of the half; row 255
+    # is excluded because it is split again vertically
+    pair = np.stack([first[:255, 255], first[:255, 256]])
+    assert (pair.sum(axis=0) == image[:255, 255]).all()
+    assert (pair.max(axis=0) == (image[:255, 255] + 1) // 2).all()
+
+    halves = insert_gap_pixels(np.full((2, 2), 3.0), gaps, roi=ROI(255, 257, 255, 257))
+    assert (halves == 0.75).all()
+
+
+def test_insert_gap_pixels_rejects_bad_input():
+    image = np.zeros((400, 400), dtype=np.uint16)
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, roi=ROI(0, 401, 0, 400))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(np.zeros((2, 2, 2), dtype=np.uint16))
+    with pytest.raises(TypeError):
+        insert_gap_pixels(np.zeros((2, 2), dtype=np.int64))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, GapPixels(split_counts=True, chip_gap=4))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, GapPixels(fill_value=float("nan")))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, GapPixels(fill_value=1e20))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, GapPixels(chip_gap=6148914691236517205))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, GapPixels(module_gaps=ModuleGaps(8, 2000)))
+    with pytest.raises(ValueError):
+        insert_gap_pixels(image, roi=ROI(2**40, 2**40 + 400, 0, 400))
+
+
+def test_insert_gap_pixels_nan_fill_for_float_images():
+    image = np.ones((2, 2), dtype=np.float32)
+    result = insert_gap_pixels(image, GapPixels(fill_value=float("nan")), roi=ROI(255, 257, 0, 2))
+    assert result.shape == (2, 4)
+    assert np.isnan(result[:, 1:3]).all()
+    assert (result[:, [0, 3]] == 1).all()