MultiThreaded Cluster finder

This commit is contained in:
froejdh_e
2025-01-09 16:53:22 +01:00
parent dc9e10016d
commit cc95561eda
16 changed files with 268 additions and 111 deletions

View File

@ -8,6 +8,8 @@ from ._aare import DetectorType
from ._aare import ClusterFile
from ._aare import hitmap
from ._aare import ClusterFinderMT, ClusterCollector, ClusterFileSink
from .CtbRawFile import CtbRawFile
from .RawFile import RawFile
from .ScanParameters import ScanParameters

View File

@ -8,18 +8,19 @@ import numpy as np
import boost_histogram as bh
import time
from aare import File, ClusterFinder, VarClusterFinder
from aare import File, ClusterFinder, VarClusterFinder, ClusterFile
base = Path('/mnt/sls_det_storage/matterhorn_data/aare_test_data/')
f = File(base/'Moench03new/cu_half_speed_master_4.json')
from aare._aare import ClusterFinderMT, ClusterCollector
from aare._aare import ClusterFinderMT, ClusterCollector, ClusterFileSink
cf = ClusterFinderMT((400,400), (3,3), n_threads = 3)
collector = ClusterCollector(cf)
# collector = ClusterCollector(cf)
out_file = ClusterFileSink(cf, "test.clust")
for i in range(1000):
img = f.read_frame()
@ -34,13 +35,13 @@ for i in range(100):
# time.sleep(1)
cf.stop()
collector.stop()
cv = collector.steal_clusters()
print(f'Processed {len(cv)} frames')
out_file.stop()
print('Done')
cfile = ClusterFile("test.clust")
# cf = ClusterFinder((400,400), (3,3))

View File

@ -1,7 +1,8 @@
#include "aare/ClusterCollector.hpp"
#include "aare/ClusterFileSink.hpp"
#include "aare/ClusterFinder.hpp"
#include "aare/ClusterFinderMT.hpp"
#include "aare/ClusterVector.hpp"
#include "aare/ClusterCollector.hpp"
#include "aare/NDView.hpp"
#include "aare/Pedestal.hpp"
#include "np_helper.hpp"
@ -9,13 +10,12 @@
#include <cstdint>
#include <filesystem>
#include <pybind11/pybind11.h>
#include <pybind11/stl_bind.h>
#include <pybind11/stl.h>
#include <pybind11/stl_bind.h>
namespace py = pybind11;
using pd_type = double;
template <typename T>
void define_cluster_vector(py::module &m, const std::string &typestr) {
auto class_name = fmt::format("ClusterVector_{}", typestr);
@ -64,42 +64,51 @@ void define_cluster_finder_mt_bindings(py::module &m) {
auto view = make_view_2d(frame);
self.push_pedestal_frame(view);
})
.def("find_clusters",
[](ClusterFinderMT<uint16_t, pd_type> &self,
py::array_t<uint16_t> frame) {
auto view = make_view_2d(frame);
self.find_clusters(view);
return;
})
.def("sync", &ClusterFinderMT<uint16_t, pd_type>::sync)
.def(
"steal_clusters",
"find_clusters",
[](ClusterFinderMT<uint16_t, pd_type> &self,
bool realloc_same_capacity) {
auto v = new ClusterVector<int>(
self.steal_clusters(realloc_same_capacity));
return v;
py::array_t<uint16_t> frame, uint64_t frame_number) {
auto view = make_view_2d(frame);
self.find_clusters(view, frame_number);
return;
},
py::arg("realloc_same_capacity") = false)
.def("stop", &ClusterFinderMT<uint16_t, pd_type>::stop);
py::arg(), py::arg("frame_number") = 0)
.def("sync", &ClusterFinderMT<uint16_t, pd_type>::sync)
.def("stop", &ClusterFinderMT<uint16_t, pd_type>::stop)
.def_property_readonly("pedestal",
[](ClusterFinderMT<uint16_t, pd_type> &self) {
auto pd = new NDArray<pd_type, 2>{};
*pd = self.pedestal();
return return_image_data(pd);
})
.def_property_readonly("noise",
[](ClusterFinderMT<uint16_t, pd_type> &self) {
auto arr = new NDArray<pd_type, 2>{};
*arr = self.noise();
return return_image_data(arr);
});
}
void define_cluster_collector_bindings(py::module &m) {
py::class_<ClusterCollector>(m, "ClusterCollector")
.def(py::init<ClusterFinderMT<uint16_t, double, int32_t>*>())
.def(py::init<ClusterFinderMT<uint16_t, double, int32_t> *>())
.def("stop", &ClusterCollector::stop)
.def("steal_clusters",
[](ClusterCollector &self) {
auto v = new std::vector<ClusterVector<int>>(
self.steal_clusters());
return v;
}, py::return_value_policy::take_ownership);
.def(
"steal_clusters",
[](ClusterCollector &self) {
auto v =
new std::vector<ClusterVector<int>>(self.steal_clusters());
return v;
},
py::return_value_policy::take_ownership);
}
void define_cluster_file_sink_bindings(py::module &m) {
py::class_<ClusterFileSink>(m, "ClusterFileSink")
.def(py::init<ClusterFinderMT<uint16_t, double, int32_t> *,
const std::filesystem::path &>())
.def("stop", &ClusterFileSink::stop);
}
void define_cluster_finder_bindings(py::module &m) {
py::class_<ClusterFinder<uint16_t, pd_type>>(m, "ClusterFinder")
@ -133,12 +142,15 @@ void define_cluster_finder_bindings(py::module &m) {
return v;
},
py::arg("realloc_same_capacity") = false)
.def("find_clusters", [](ClusterFinder<uint16_t, pd_type> &self,
py::array_t<uint16_t> frame) {
auto view = make_view_2d(frame);
self.find_clusters(view);
return;
});
.def(
"find_clusters",
[](ClusterFinder<uint16_t, pd_type> &self,
py::array_t<uint16_t> frame, uint64_t frame_number) {
auto view = make_view_2d(frame);
self.find_clusters(view);
return;
},
py::arg(), py::arg("frame_number") = 0);
m.def("hitmap",
[](std::array<size_t, 2> image_size, ClusterVector<int32_t> &cv) {

View File

@ -28,16 +28,13 @@ void define_cluster_file_io_bindings(py::module &m) {
py::arg(), py::arg("chunk_size") = 1000, py::arg("mode") = "r")
.def("read_clusters",
[](ClusterFile &self, size_t n_clusters) {
auto *vec =
new std::vector<Cluster3x3>(self.read_clusters(n_clusters));
return return_vector(vec);
auto v = new ClusterVector<int32_t>(self.read_clusters(n_clusters));
return v;
})
.def("read_frame",
[](ClusterFile &self) {
int32_t frame_number;
auto *vec =
new std::vector<Cluster3x3>(self.read_frame(frame_number));
return py::make_tuple(frame_number, return_vector(vec));
auto v = new ClusterVector<int32_t>(self.read_frame());
return v;
})
.def("write_frame", &ClusterFile::write_frame)
.def("read_cluster_with_cut",
@ -59,12 +56,11 @@ void define_cluster_file_io_bindings(py::module &m) {
})
.def("__iter__", [](ClusterFile &self) { return &self; })
.def("__next__", [](ClusterFile &self) {
auto vec =
new std::vector<Cluster3x3>(self.read_clusters(self.chunk_size()));
if (vec->size() == 0) {
auto v = new ClusterVector<int32_t>(self.read_clusters(self.chunk_size()));
if (v->size() == 0) {
throw py::stop_iteration();
}
return return_vector(vec);
return v;
});
m.def("calculate_eta2", []( aare::ClusterVector<int32_t> &clusters) {

View File

@ -124,8 +124,41 @@ void define_file_io_bindings(py::module &m) {
self.read_into(reinterpret_cast<std::byte *>(image.mutable_data()),
n_frames);
return image;
})
.def("__enter__", [](File &self) { return &self; })
.def("__exit__",
[](File &self,
const std::optional<pybind11::type> &exc_type,
const std::optional<pybind11::object> &exc_value,
const std::optional<pybind11::object> &traceback) {
// self.close();
})
.def("__iter__", [](File &self) { return &self; })
.def("__next__", [](File &self) {
try{
const uint8_t item_size = self.bytes_per_pixel();
py::array image;
std::vector<ssize_t> shape;
shape.reserve(2);
shape.push_back(self.rows());
shape.push_back(self.cols());
if (item_size == 1) {
image = py::array_t<uint8_t>(shape);
} else if (item_size == 2) {
image = py::array_t<uint16_t>(shape);
} else if (item_size == 4) {
image = py::array_t<uint32_t>(shape);
}
self.read_into(
reinterpret_cast<std::byte *>(image.mutable_data()));
return image;
}catch(std::runtime_error &e){
throw py::stop_iteration();
}
});
py::class_<FileConfig>(m, "FileConfig")
.def(py::init<>())
.def_readwrite("rows", &FileConfig::rows)

View File

@ -28,5 +28,6 @@ PYBIND11_MODULE(_aare, m) {
define_cluster_finder_mt_bindings(m);
define_cluster_file_io_bindings(m);
define_cluster_collector_bindings(m);
define_cluster_file_sink_bindings(m);
}