mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-09-22 00:52:13 +02:00
Erik Fröjdh
b0875106e3
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
337 lines
12 KiB
Python
337 lines
12 KiB
Python
# SPDX-License-Identifier: MPL-2.0
|
|
|
|
import pytest
|
|
import numpy as np
|
|
import boost_histogram as bh
|
|
import time
|
|
from pathlib import Path
|
|
import pickle
|
|
import struct
|
|
import gc
|
|
import weakref
|
|
|
|
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"""
|
|
f = ClusterFile(test_data_path / "clust/single_frame_97_clustrers.clust")
|
|
assert f.estimate_n_clusters() == 97
|
|
assert f.tell() == 0
|
|
cv = f.read_clusters(10) #conversion does not work
|
|
|
|
|
|
assert cv.frame_number == 135
|
|
assert cv.size == 10
|
|
|
|
#Known data
|
|
#frame_number, num_clusters [135] 97
|
|
#[ 1 200] [0 1 2 3 4 5 6 7 8]
|
|
#[ 2 201] [ 9 10 11 12 13 14 15 16 17]
|
|
#[ 3 202] [18 19 20 21 22 23 24 25 26]
|
|
#[ 4 203] [27 28 29 30 31 32 33 34 35]
|
|
#[ 5 204] [36 37 38 39 40 41 42 43 44]
|
|
#[ 6 205] [45 46 47 48 49 50 51 52 53]
|
|
#[ 7 206] [54 55 56 57 58 59 60 61 62]
|
|
#[ 8 207] [63 64 65 66 67 68 69 70 71]
|
|
#[ 9 208] [72 73 74 75 76 77 78 79 80]
|
|
#[ 10 209] [81 82 83 84 85 86 87 88 89]
|
|
|
|
#conversion to numpy array
|
|
arr = np.array(cv, copy = False)
|
|
|
|
assert arr.size == 10
|
|
for i in range(10):
|
|
assert arr[i]['x'] == i+1
|
|
|
|
@pytest.mark.withdata
|
|
def test_read_clusters_and_fill_histogram(test_data_path):
|
|
# Create the histogram
|
|
n_bins = 100
|
|
xmin = -100
|
|
xmax = 1e4
|
|
hist_aare = bh.Histogram(bh.axis.Regular(n_bins, xmin, xmax))
|
|
|
|
fname = test_data_path / "clust/beam_En700eV_-40deg_300V_10us_d0_f0_100.clust"
|
|
|
|
#Read clusters and fill the histogram with pixel values
|
|
with ClusterFile(fname, chunk_size = 10000) as f:
|
|
for clusters in f:
|
|
arr = np.array(clusters, copy = False)
|
|
hist_aare.fill(arr['data'].flat)
|
|
|
|
|
|
#Load the histogram from the pickle file
|
|
with open(fname.with_suffix('.pkl'), 'rb') as f:
|
|
hist_py = pickle.load(f)
|
|
|
|
#Compare the two histograms
|
|
assert hist_aare == hist_py
|