mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-10-17 15:35:10 +02:00
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
This commit is contained in:
Erik Fröjdh
+123
-28
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user