Improved ClusterFile API (#358)

Improved API for ClusterFile

- read_frame now returns a std::optional<ClusterVector> to differentiate
between an empty frame and end of file
- f.frames() produces an iterator over frames
- f.chunks() produces an iterator reading chunks of clusters

python/aare/ClusterFinder.py

@@ -89,19 +89,55 @@ def ClusterFileSink(clusterfindermt, cluster_file, dtype=np.int32):
 
 
 def ClusterFile(fname, cluster_size=(3,3), dtype=np.int32, chunk_size = 1000, mode = "r"):
-    """
-    Factory function to create a ClusterFile object. Provides a cleaner syntax for
-    the templated ClusterFile in C++.
+    """Create a reader or writer for a binary cluster file.
+
+    Parameters
+    ----------
+    fname : path-like
+        Cluster file to open.
+    cluster_size : tuple[int, int], default=(3, 3)
+        Cluster dimensions stored in the file.
+    dtype : numpy dtype, default=numpy.int32
+        Data type of the cluster values stored in the file.
+    chunk_size : int, default=1000
+        Maximum number of selected clusters returned by ``chunks()`` and
+        default iteration. Must be positive when iterating over chunks.
+    mode : {"r", "w", "a"}, default="r"
+        Open for reading, truncate and write, or append, respectively.
+
+    Returns
+    -------
+    ClusterFile
+        The compiled ClusterFile specialization matching ``cluster_size`` and
+        ``dtype``.
+
+    Notes
+    -----
+    The file format contains no cluster shape or data-type metadata. Supplying
+    values that do not match the file causes its bytes to be interpreted
+    incorrectly. Use ``frames()`` to iterate over complete frames, including
+    empty or fully filtered frames with their stored frame numbers. Use
+    ``chunks()`` or ``chunks(chunk_size)`` to iterate over selected clusters
+    in batches. Chunks may combine frames, so their frame number is not
+    reliable per-cluster metadata.
+
+    Iterators consume the current file position without rewinding; use one
+    traversal at a time. Each result owns its storage and remains valid after
+    advancing the iterator or closing the file.
+
+    Examples
+    --------
 
     .. code-block:: python
 
         from aare import ClusterFile
-        
-        with ClusterFile("clusters.clust", cluster_size=(3,3), dtype=np.int32) as cf:
-            # cf is now a ClusterFile_Cluster3x3i object but you don't need to know that.
-            for clusters in cf:
-                # Loop over clusters in chunks of 1000 
-                # The type of clusters will be a ClusterVector_Cluster3x3i in this case
+
+        with ClusterFile(
+            "clusters.clust", cluster_size=(3, 3), dtype=np.int32
+        ) as cf:
+            for clusters in cf.chunks():
+                # Process clusters in chunks of at most 1000.
+                ...
 
     """
 

python/src/bind_ClusterFile.hpp

@@ -11,6 +11,7 @@
 #include <pybind11/stl.h>
 #include <pybind11/stl/filesystem.h>
 #include <string>
+#include <utility>
 
 // Disable warnings for unused parameters, as we ignore some
 // in the __exit__ method
@@ -20,18 +21,90 @@
 namespace py = pybind11;
 using namespace ::aare;
 
+template <typename Range> class ClusterFileIterator {
+    Range m_range;
+    std::optional<typename Range::Iterator> m_iterator;
+    bool m_done{false};
+
+  public:
+    explicit ClusterFileIterator(Range range) : m_range(std::move(range)) {}
+
+    typename Range::Iterator::value_type next() {
+        if (m_done) {
+            throw py::stop_iteration();
+        }
+        // Advance before returning the next result, never after yielding it.
+        if (m_iterator) {
+            ++*m_iterator;
+        } else {
+            m_iterator.emplace(m_range.begin());
+        }
+        if (*m_iterator == m_range.end()) {
+            m_done = true;
+            throw py::stop_iteration();
+        }
+        return std::move(**m_iterator);
+    }
+};
+
+template <typename Iterator>
+void define_cluster_file_iterator(py::module &m, const std::string &name) {
+    py::class_<Iterator>(m, name.c_str())
+        .def(
+            "__iter__", [](Iterator &self) -> Iterator & { return self; },
+            py::return_value_policy::reference_internal)
+        .def("__next__", &Iterator::next);
+}
+
 template <typename Type, uint8_t CoordSizeX, uint8_t CoordSizeY,
           typename CoordType = uint16_t>
 void define_ClusterFile(py::module &m, const std::string &typestr) {
 
     using ClusterType = Cluster<Type, CoordSizeX, CoordSizeY, CoordType>;
+    using File = ClusterFile<ClusterType>;
+    using FrameIterator = ClusterFileIterator<typename File::FrameRange>;
+    using ChunkIterator = ClusterFileIterator<typename File::ChunkRange>;
 
     auto class_name = fmt::format("ClusterFile_{}", typestr);
+    define_cluster_file_iterator<FrameIterator>(m,
+                                                class_name + "_FrameIterator");
+    define_cluster_file_iterator<ChunkIterator>(m,
+                                                class_name + "_ChunkIterator");
 
-    py::class_<ClusterFile<ClusterType>>(m, class_name.c_str())
+    py::class_<ClusterFile<ClusterType>>(
+        m, class_name.c_str(),
+        "Read and write legacy binary cluster files. The format contains no "
+        "cluster type or shape metadata, so this class must match the file.")
         .def(py::init<const std::filesystem::path &, size_t,
                       const std::string &>(),
-             py::arg(), py::arg("chunk_size") = 1000, py::arg("mode") = "r")
+             py::arg("fname"), py::arg("chunk_size") = 1000,
+             py::arg("mode") = "r",
+             "Open a cluster file. Mode must be 'r' to read, 'w' to truncate "
+             "and write, or 'a' to append.")
+        .def(
+            "frames", [](File &self) { return FrameIterator(self.frames()); },
+            py::keep_alive<0, 1>(),
+            "Iterate over complete frames from the current position, including "
+            "empty and fully filtered frames with their stored frame numbers. "
+            "Results own their storage. The iterator keeps the file alive but "
+            "shares its cursor; use only one traversal at a time.")
+        .def(
+            "chunks", [](File &self) { return ChunkIterator(self.chunks()); },
+            py::keep_alive<0, 1>(),
+            "Iterate from the current position using the constructor's chunk "
+            "size, which must be positive. Results own their storage. The "
+            "iterator keeps the file alive but shares its cursor; use only "
+            "one traversal at a time.")
+        .def(
+            "chunks",
+            [](File &self, size_t chunk_size) {
+                return ChunkIterator(self.chunks(chunk_size));
+            },
+            py::arg("chunk_size"), py::keep_alive<0, 1>(),
+            "Iterate over chunks of up to chunk_size selected clusters. The "
+            "size must be positive and does not change the constructor's "
+            "default. Chunks may span frames; their frame numbers are not "
+            "per-cluster metadata. Only the final chunk can be short.")
         .def(
             "read_clusters",
             [](ClusterFile<ClusterType> &self, size_t n_clusters) {
@@ -39,32 +112,60 @@ void define_ClusterFile(py::module &m, const std::string &typestr) {
                     self.read_clusters(n_clusters));
                 return v;
             },
-            py::return_value_policy::take_ownership, py::arg("n_clusters"))
-        .def("read_frame",
-             [](ClusterFile<ClusterType> &self) {
-                 auto v = new ClusterVector<ClusterType>(self.read_frame());
-                 return v;
-             })
-        .def("set_roi", &ClusterFile<ClusterType>::set_roi, py::arg("roi"))
-        .def("tell", &ClusterFile<ClusterType>::tell)
+            py::return_value_policy::take_ownership, py::arg("n_clusters"),
+            "Read up to n_clusters without preserving frame boundaries. The "
+            "result may combine frames, so its frame number is not reliable "
+            "per-cluster metadata.")
+        .def(
+            "read_frame",
+            [](ClusterFile<ClusterType> &self) -> py::object {
+                auto clusters = self.read_frame();
+                if (!clusters) {
+                    return py::none();
+                }
+                return py::cast(
+                    new ClusterVector<ClusterType>(std::move(*clusters)),
+                    py::return_value_policy::take_ownership);
+            },
+            "Read and return the next complete frame with its frame number, "
+            "or None at end of file.")
+        .def("set_roi", &ClusterFile<ClusterType>::set_roi, py::arg("roi"),
+             "Select clusters whose centers lie within the half-open ROI.")
+        .def("tell", &ClusterFile<ClusterType>::tell,
+             "Return the current byte position in the file.")
         .def("estimate_n_clusters",
-             &ClusterFile<ClusterType>::estimate_n_clusters)
+             &ClusterFile<ClusterType>::estimate_n_clusters,
+             "Estimate the number of clusters from the file size. Frame "
+             "headers can make this larger than the actual count.")
         .def(
             "set_noise_map",
-            [](ClusterFile<ClusterType> &self, py::array_t<int32_t> noise_map) {
+            [](ClusterFile<ClusterType> &self,
+               py::array_t<int32_t, py::array::c_style> noise_map) {
                 auto view = make_view_2d(noise_map);
                 self.set_noise_map(view);
             },
-            py::arg("noise_map"))
+            py::arg("noise_map"),
+            "Set a two-dimensional, C-contiguous int32 noise map indexed as "
+            "[y, x]. The map must cover every cluster center coordinate.")
 
-        .def("set_gain_map",
-             [](ClusterFile<ClusterType> &self, py::array_t<double> gain_map) {
-                 auto view = make_view_2d(gain_map);
-                 self.set_gain_map(view);
-             })
+        .def(
+            "set_gain_map",
+            [](ClusterFile<ClusterType> &self,
+               py::array_t<double, py::array::c_style> gain_map) {
+                auto view = make_view_2d(gain_map);
+                self.set_gain_map(view);
+            },
+            py::arg("gain_map"),
+            "Set a two-dimensional, C-contiguous float64 gain map in "
+            "ADU/energy, indexed as [y, x]. Clusters whose complete footprint "
+            "extends beyond the map are retained with all data values set to "
+            "zero.")
 
-        .def("close", &ClusterFile<ClusterType>::close)
-        .def("write_frame", &ClusterFile<ClusterType>::write_frame)
+        .def("close", &ClusterFile<ClusterType>::close,
+             "Close the file. Calling close more than once is safe.")
+        .def("write_frame", &ClusterFile<ClusterType>::write_frame,
+             py::arg("clusters"),
+             "Write one ClusterVector, including its frame number.")
         .def("__enter__", [](ClusterFile<ClusterType> &self) { return &self; })
         .def("__exit__",
              [](ClusterFile<ClusterType> &self,
@@ -74,14 +175,8 @@ void define_ClusterFile(py::module &m, const std::string &typestr) {
                  self.close();
              })
         .def("__iter__", [](ClusterFile<ClusterType> &self) { return &self; })
-        .def("__next__", [](ClusterFile<ClusterType> &self) {
-            auto v = new ClusterVector<ClusterType>(
-                self.read_clusters(self.chunk_size()));
-            if (v->size() == 0) {
-                throw py::stop_iteration();
-            }
-            return v;
-        });
+        .def("__next__",
+             [](File &self) { return ChunkIterator(self.chunks()).next(); });
 }
 
 #pragma GCC diagnostic pop

python/tests/test_ClusterFile.py

@@ -6,10 +6,279 @@ import boost_histogram as bh
 import time
 from pathlib import Path
 import pickle
+import struct
+import gc
+import weakref
 
-from aare import ClusterFile
+from aare import ClusterFile, ROI
 from conftest import test_data_path
 
+
+@pytest.fixture
+def iterator_file(tmp_path):
+    fname = tmp_path / "iteration.clust"
+    records = [
+        (-11, []),
+        (42, [(5, 6, 10), (7, 8, 20)]),
+        (103, []),
+        (104, [(5, 6, 30), (7, 8, 40)]),
+        (105, []),
+    ]
+    fname.write_bytes(b"".join(
+        struct.pack("@iI", number, len(clusters))
+        + b"".join(struct.pack("@HH9i", x, y, *([value] * 9))
+                   for x, y, value in clusters)
+        for number, clusters in records
+    ))
+    return fname
+
+
+def test_frames_preserve_empty_frames_and_frame_numbers(iterator_file):
+    with ClusterFile(iterator_file) as reader:
+        frames = reader.frames()
+        assert iter(frames) is frames
+        assert reader.tell() == 0
+        results = list(frames)
+        assert [frame.frame_number for frame in results] == [-11, 42, 103, 104, 105]
+        assert [frame.size for frame in results] == [0, 2, 0, 2, 0]
+        for _ in range(2):
+            with pytest.raises(StopIteration):
+                next(frames)
+        assert list(reader.frames()) == []
+
+
+@pytest.mark.parametrize("chunk_size, sizes", [(1, [1, 1, 1, 1]), (2, [2, 2]),
+                                             (3, [3, 1]), (10, [4])])
+@pytest.mark.parametrize("explicit", [False, True])
+def test_chunks_cross_frames(iterator_file, chunk_size, sizes, explicit):
+    with ClusterFile(iterator_file, chunk_size=chunk_size) as reader:
+        chunks = reader.chunks(chunk_size=chunk_size) if explicit else reader.chunks()
+        assert iter(chunks) is chunks
+        assert reader.tell() == 0
+        results = list(chunks)
+        assert [chunk.size for chunk in results] == sizes
+        data = np.concatenate([np.asarray(chunk) for chunk in results])
+        np.testing.assert_array_equal(data["x"], [5, 7, 5, 7])
+        np.testing.assert_array_equal(data["data"][:, 0, 0], [10, 20, 30, 40])
+        for _ in range(2):
+            with pytest.raises(StopIteration):
+                next(chunks)
+
+
+def test_chunk_override_preserves_default_iteration(iterator_file):
+    with ClusterFile(iterator_file, chunk_size=2) as reader:
+        assert iter(reader) is reader
+        assert next(reader.chunks(1)).size == 1
+        assert [chunk.size for chunk in reader] == [2, 1]
+
+
+@pytest.mark.parametrize("use_roi", [False, True])
+@pytest.mark.parametrize("use_noise", [False, True])
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+def test_iterators_apply_filters_and_gain(iterator_file, use_roi, use_noise, method):
+    with ClusterFile(iterator_file, chunk_size=3) as reader:
+        if use_roi:
+            reader.set_roi(ROI(0, 6, 0, 12))
+        if use_noise:
+            reader.set_noise_map(np.full((12, 12), 15, dtype=np.int32))
+        reader.set_gain_map(np.full((12, 12), 2.0))
+        results = list(getattr(reader, method)())
+        if method == "frames":
+            assert [frame.frame_number for frame in results] == [-11, 42, 103, 104, 105]
+        expected = [value / 2 for x, value in [(5, 10), (7, 20), (5, 30), (7, 40)]
+                    if (not use_roi or x == 5) and (not use_noise or value > 15)]
+        data = np.concatenate([np.asarray(result) for result in results])
+        np.testing.assert_array_equal(data["data"][:, 0, 0], expected)
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+def test_iterators_handle_all_rejected_clusters(iterator_file, method):
+    with ClusterFile(iterator_file) as reader:
+        reader.set_roi(ROI(0, 1, 0, 1))
+        results = list(getattr(reader, method)())
+        assert len(results) == (5 if method == "frames" else 0)
+        assert all(result.size == 0 for result in results)
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+def test_iterator_keeps_file_alive(iterator_file, method):
+    reader = ClusterFile(iterator_file)
+    owner = weakref.ref(reader)
+    iterator = getattr(reader, method)()
+    del reader
+    gc.collect()
+    assert owner() is not None
+    results = list(iterator)
+    assert sum(result.size for result in results) == 4
+    del iterator
+    gc.collect()
+    assert owner() is None
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+def test_retained_results_and_numpy_views_own_storage(iterator_file, method):
+    with ClusterFile(iterator_file, chunk_size=2) as reader:
+        reader.read_frame()  # Skip the initial empty frame.
+        iterator = getattr(reader, method)()
+        first = next(iterator)
+        array = np.asarray(first)
+        expected = array.copy()
+        results = list(iterator)
+        np.testing.assert_array_equal(np.asarray(first), expected)
+    del first, results, iterator, reader
+    gc.collect()
+    np.testing.assert_array_equal(array, expected)
+
+
+def test_frame_iteration_resumes_without_reading_ahead(iterator_file):
+    with ClusterFile(iterator_file) as reader:
+        reader.read_frame()
+        for frame in reader.frames():
+            assert frame.frame_number == 42
+            break
+        assert reader.read_frame().frame_number == 103
+        assert [frame.frame_number for frame in reader.frames()] == [104, 105]
+
+
+def test_partial_chunk_must_be_completed_before_frame_iteration(iterator_file):
+    with ClusterFile(iterator_file) as reader:
+        for chunk in reader.chunks(1):
+            assert chunk.size == 1
+            break
+        with pytest.raises(RuntimeError, match="clusters left"):
+            next(reader.frames())
+        assert reader.read_clusters(1).size == 1
+        assert [frame.frame_number for frame in reader.frames()] == [103, 104, 105]
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+def test_empty_iterators(tmp_path, method):
+    fname = tmp_path / "empty.clust"
+    fname.touch()
+    with ClusterFile(fname) as reader:
+        iterator = getattr(reader, method)()
+        for _ in range(2):
+            with pytest.raises(StopIteration):
+                next(iterator)
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+@pytest.mark.parametrize("started", [False, True])
+def test_iterators_reject_closed_files(iterator_file, method, started):
+    reader = ClusterFile(iterator_file, chunk_size=1)
+    iterator = getattr(reader, method)()
+    if started:
+        next(iterator)
+    reader.close()
+    with pytest.raises(RuntimeError, match="not opened for reading"):
+        next(iterator)
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+@pytest.mark.parametrize("mode", ["w", "a"])
+def test_iterators_reject_writing_modes(tmp_path, method, mode):
+    with ClusterFile(tmp_path / "output.clust", mode=mode) as reader:
+        with pytest.raises(RuntimeError, match="not opened for reading"):
+            next(getattr(reader, method)())
+
+
+def test_chunk_sizes_must_be_positive(iterator_file):
+    with ClusterFile(iterator_file, chunk_size=0) as reader:
+        with pytest.raises(ValueError, match="greater than zero"):
+            reader.chunks()
+        with pytest.raises(ValueError, match="greater than zero"):
+            reader.chunks(0)
+        with pytest.raises(ValueError, match="greater than zero"):
+            next(reader)
+        with pytest.raises(TypeError):
+            reader.chunks(-1)
+        assert reader.tell() == 0
+        assert reader.read_clusters(0).size == 0
+        assert len(list(reader.frames())) == 5
+
+
+@pytest.mark.parametrize("method", ["frames", "chunks"])
+def test_iterators_defer_read_errors_until_next_result(tmp_path, method):
+    fname = tmp_path / "truncated.clust"
+    cluster = struct.pack("@HH9i", 5, 6, *range(9))
+    frame = struct.pack("@iI", 42, 1) + cluster
+    fname.write_bytes(frame + frame[:-1])
+    with ClusterFile(fname, chunk_size=1) as reader:
+        iterator = getattr(reader, method)()
+        assert next(iterator).size == 1
+        assert reader.tell() == len(frame)
+        with pytest.raises(RuntimeError):
+            next(iterator)
+
+
+@pytest.mark.parametrize(
+    "shape, dtype",
+    [(shape, dtype)
+     for shape in [(2, 2), (3, 3), (5, 5), (7, 7), (9, 9)]
+     for dtype in [np.int32, np.float32, np.float64]]
+    + [((3, 3), np.int16)],
+)
+def test_iterators_are_bound_for_each_cluster_type(tmp_path, shape, dtype):
+    fname = tmp_path / "typed_frames.clust"
+    values = list(range(shape[0] * shape[1]))
+    record = struct.pack("@HH" + np.dtype(dtype).char * len(values), 5, 6, *values)
+    fname.write_bytes((struct.pack("@iI", -42, 1) + record) * 3)
+    with ClusterFile(fname, cluster_size=shape, dtype=dtype) as reader:
+        frame = next(reader.frames())
+        assert frame.frame_number == -42
+        np.testing.assert_array_equal(np.asarray(frame)["data"].reshape(-1), values)
+        chunks = list(reader.chunks())
+        assert len(chunks) == 1
+        assert chunks[0].size == 2
+        np.testing.assert_array_equal(np.asarray(chunks[0])["data"].reshape(-1), values * 2)
+
+
+def test_read_frame_returns_none_at_eof(tmp_path):
+    fname = tmp_path / "empty.clust"
+    fname.touch()
+
+    with ClusterFile(fname) as f:
+        assert f.read_frame() is None
+
+
+def test_read_frame_raises_for_malformed_file(tmp_path):
+    fname = tmp_path / "malformed.clust"
+    fname.write_bytes(b"\x00")
+
+    with ClusterFile(fname) as f, pytest.raises(RuntimeError):
+        f.read_frame()
+
+
+@pytest.mark.parametrize("use_roi", [False, True])
+@pytest.mark.parametrize("method", ["read_clusters", "default", "frames", "chunks"])
+@pytest.mark.parametrize(
+    "size",
+    [
+        pytest.param(3, id="partial-frame-number"),
+        pytest.param(4, id="missing-cluster-count"),
+        pytest.param(7, id="partial-cluster-count"),
+        pytest.param(48, id="missing-cluster-record"),
+        pytest.param(87, id="partial-cluster-record"),
+    ],
+)
+def test_chunk_reads_reject_incomplete_frames(tmp_path, use_roi, method, size):
+    fname = tmp_path / "incomplete.clust"
+    cluster = struct.pack("@HH9i", 5, 6, *range(9))
+    frame = struct.pack("@iI", 42, 2) + cluster * 2
+    fname.write_bytes(frame[:size])
+
+    with ClusterFile(fname) as reader:
+        if use_roi:
+            reader.set_roi(ROI(0, 10, 0, 10))
+        with pytest.raises(RuntimeError):
+            if method == "default":
+                next(reader)
+            elif method == "read_clusters":
+                reader.read_clusters(10)
+            else:
+                next(getattr(reader, method)())
+
+
 @pytest.mark.withdata
 def test_cluster_file(test_data_path): 
     """Test ClusterFile"""