Prototype multi threaded file reader (#343)

- Reading files from nfs shares tops out at ~1GB/s for single threaded
reads.
- MultiThreadedFileReader uses our File wrapper to read a generic file
in parallel
- Placed in aare::experimental to show that it's not production ready

python/tests/test_MultiThreadedFileReader.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: MPL-2.0
+import numpy as np
+import pytest
+
+import aare
+from aare.experimental import MultiThreadedFileReader
+
+
+@pytest.fixture
+def frame_file(tmp_path):
+    data = np.arange(10 * 2 * 3, dtype=np.uint16).reshape(10, 2, 3)
+    path = tmp_path / "frames.npy"
+    np.save(path, data)
+    return path, data
+
+
+def test_experimental_import_path():
+    assert aare.experimental.MultiThreadedFileReader is MultiThreadedFileReader
+    assert MultiThreadedFileReader.__module__ == "aare.experimental"
+    assert not hasattr(aare, "MultiThreadedFileReader")
+
+
+def test_reads_all_frames_in_order(frame_file):
+    path, expected = frame_file
+    reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=3)
+
+    first = reader.read()
+    second = reader.read()
+    exhausted = reader.read()
+
+    assert np.array_equal(first, expected[:6])
+    assert np.array_equal(second, expected[6:])
+    assert exhausted.shape == (0, 2, 3)
+    assert first.dtype == np.uint16
+    assert reader.n_threads == 2
+    assert reader.chunk_size == 3
+    assert reader.total_frames == 10
+    assert reader.source_total_frames == 10
+    assert reader.rows == 2
+    assert reader.cols == 3
+    assert reader.bitdepth == 16
+    assert reader.dtype == np.dtype(np.uint16)
+    assert reader.bytes_per_frame == 12
+    assert reader.total_bytes == expected.nbytes
+    assert len(reader) == 10
+    assert reader.tell() == 10
+    assert reader.remaining_frames == 0
+    assert reader.next_read_frames == 0
+    assert reader.next_read_bytes == 0
+
+
+def test_total_frame_limit(frame_file):
+    path, expected = frame_file
+    reader = MultiThreadedFileReader(
+        path, n_threads=2, chunk_size=2, total_frames=7
+    )
+
+    assert np.array_equal(reader.read(), expected[:4])
+    assert np.array_equal(reader.read_all(), expected[4:7])
+
+
+def test_iteration_yields_one_chunk_per_thread(frame_file):
+    path, expected = frame_file
+    reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=2)
+
+    batches = list(reader)
+
+    assert [len(batch) for batch in batches] == [4, 4, 2]
+    assert np.array_equal(np.concatenate(batches), expected)
+
+
+def test_seek_resets_iteration_position(frame_file):
+    path, expected = frame_file
+    reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=2)
+
+    reader.read()
+    assert reader.tell() == 4
+
+    reader.seek(1)
+    assert reader.tell() == 1
+    assert np.array_equal(reader.read(), expected[1:5])
+
+    with pytest.raises(IndexError):
+        reader.seek(11)
+
+
+def test_context_manager_closes_worker_files(frame_file):
+    path, expected = frame_file
+
+    with MultiThreadedFileReader(path, n_threads=2, chunk_size=2) as reader:
+        assert not reader.closed
+        assert np.array_equal(reader.read(), expected[:4])
+
+    assert reader.closed
+    reader.close()
+    with pytest.raises(RuntimeError):
+        reader.read()
+    with pytest.raises(RuntimeError):
+        reader.seek(0)
+    with pytest.raises(RuntimeError):
+        with reader:
+            pass
+
+
+def test_explicit_zero_frame_limit(frame_file):
+    path, expected = frame_file
+    reader = MultiThreadedFileReader(
+        path, n_threads=8, chunk_size=3, total_frames=0
+    )
+
+    actual = reader.read()
+    assert actual.shape == (0, *expected.shape[1:])
+    assert actual.dtype == expected.dtype
+
+
+@pytest.mark.parametrize(
+    "dtype", [np.int8, np.int32, np.uint64, np.float32, np.float64]
+)
+def test_preserves_numpy_dtype(tmp_path, dtype):
+    expected = np.arange(4 * 2 * 3, dtype=dtype).reshape(4, 2, 3)
+    path = tmp_path / "typed-frames.npy"
+    np.save(path, expected)
+
+    reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=3)
+    actual = reader.read()
+
+    assert actual.dtype == expected.dtype
+    assert reader.dtype == expected.dtype
+    assert np.array_equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    ("n_threads", "chunk_size", "total_frames"),
+    [(0, 1, None), (1, 0, None), (1, 1, 11)],
+)
+def test_invalid_configuration(
+    frame_file, n_threads, chunk_size, total_frames
+):
+    path, _ = frame_file
+
+    with pytest.raises(ValueError):
+        MultiThreadedFileReader(
+            path,
+            n_threads=n_threads,
+            chunk_size=chunk_size,
+            total_frames=total_frames,
+        )

python/tests/test_PixelHistogram.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: MPL-2.0
+import numpy as np
+import pytest
+
+from aare import PixelHistogram
+
+
+def _random_frames(rows, cols, n, xmin, xmax, seed=0):
+    rng = np.random.default_rng(seed)
+    return [rng.uniform(xmin - 0.25, xmax + 0.25, size=(rows, cols)).astype(np.float64)
+            for _ in range(n)]
+
+
+def _reference_hdata(frames, rows, cols, n_bins, xmin, xmax):
+    expected = np.zeros((rows, cols, n_bins), dtype=np.uint16)
+    inv_range = n_bins / (xmax - xmin)
+    for img in frames:
+        for r in range(rows):
+            for c in range(cols):
+                v = float(img[r, c])
+                if not (xmin <= v < xmax):
+                    continue
+                b = int((v - xmin) * inv_range)
+                if b >= n_bins:
+                    b = n_bins - 1
+                expected[r, c, b] += 1
+    return expected
+
+
+def test_async_fill_matches_reference():
+    rows, cols, n_bins = 5, 7, 8
+    xmin, xmax = 0.0, 2.0
+    frames = _random_frames(rows, cols, n=3, xmin=xmin, xmax=xmax, seed=1)
+
+    hist = PixelHistogram(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax)
+    for img in frames:
+        hist.fill_async(img)
+
+    np.testing.assert_array_equal(
+        hist.values(),
+        _reference_hdata(frames, rows, cols, n_bins, xmin, xmax),
+    )
+
+
+
+
+def test_fill_async_copies_buffer():
+    # After fill_async returns, the caller should be free to mutate the
+    # numpy array without affecting the pending fill.
+    rows, cols, n_bins = 4, 4, 4
+    xmin, xmax = 0.0, 1.0
+    hist = PixelHistogram(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax, n_threads=1, max_pending=8)
+
+    img = np.full((rows, cols), 0.1, dtype=np.float64)  # falls in bin 0
+    hist.fill_async(img)
+    # Mutate the original array immediately; this must not affect the
+    # value that was already enqueued.
+    img[:] = 0.9  # would be bin 3
+    hist.flush()
+
+    h = hist.values()
+    assert h.shape == (rows, cols, n_bins)
+    # Every pixel saw one value in bin 0, none elsewhere.
+    assert (h[:, :, 0] == 1).all()
+    assert (h[:, :, 1:] == 0).all()
+
+
+    def test_fill_async_rejects_wrong_shape():
+        hist = PixelHistogram(8, 8, 4, 0.0, 1.0)
+        bad = np.zeros((4, 4), dtype=np.float32)
+        with pytest.raises(ValueError):
+            hist.fill_async(bad)
+
+
+def test_hdata_flushes_pending():
+    # Submit several frames with a tiny queue and read hdata() without an
+    # explicit flush(); hdata() must drain everything first.
+    rows, cols, n_bins = 3, 3, 4
+    xmin, xmax = 0.0, 1.0
+    hist = PixelHistogram(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax,
+                          n_threads=1, max_pending=1)
+    frames = _random_frames(rows, cols, n=8, xmin=xmin, xmax=xmax, seed=3)
+    for img in frames:
+        hist.fill_async(img)
+
+    h = hist.values()  # no explicit flush()
+    np.testing.assert_array_equal(
+        h, _reference_hdata(frames, rows, cols, n_bins, xmin, xmax)
+    )
+
+
+def test_bin_centers_and_edges():
+    n_bins = 5
+    xmin, xmax = 0.0, 1.0
+    hist = PixelHistogram(rows=2, cols=2, n_bins=n_bins, xmin=xmin, xmax=xmax)
+    edges = hist.bin_edges()
+    centers = hist.bin_centers()
+    assert edges.shape == (n_bins + 1,)
+    assert centers.shape == (n_bins,)
+    np.testing.assert_allclose(edges, np.linspace(xmin, xmax, n_bins + 1), atol=1e-6)
+    np.testing.assert_allclose(centers, 0.5 * (edges[:-1] + edges[1:]), atol=1e-6)