WIP

python/examples/play.py

@@ -13,36 +13,66 @@ from aare import File, ClusterFinder, VarClusterFinder
 base = Path('/mnt/sls_det_storage/matterhorn_data/aare_test_data/')
 
 f = File(base/'Moench03new/cu_half_speed_master_4.json')
-cf = ClusterFinder((400,400), (3,3))
+
+
+from aare._aare import ClusterFinderMT, ClusterCollector
+
+
+cf = ClusterFinderMT((400,400), (3,3), n_threads = 3)
+collector = ClusterCollector(cf)
+
 for i in range(1000):
-    cf.push_pedestal_frame(f.read_frame())
+    img = f.read_frame()
+    cf.push_pedestal_frame(img)
+print('Pedestal done')
+cf.sync()
 
-fig, ax = plt.subplots()
-im = ax.imshow(cf.pedestal())
-cf.pedestal()
-cf.noise()
+for i in range(100):
+    img = f.read_frame()
+    cf.find_clusters(img)
+
+
+# time.sleep(1)
+cf.stop()  
+collector.stop()
+cv = collector.steal_clusters()
+print(f'Processed {len(cv)} frames')
+
+print('Done')
 
 
 
-N = 500
-t0 = time.perf_counter()
-hist1 = bh.Histogram(bh.axis.Regular(40, -2, 4000))
-f.seek(0)
 
-t0 = time.perf_counter()
-data = f.read_n(N)
-t_elapsed = time.perf_counter()-t0
+# cf = ClusterFinder((400,400), (3,3))
+# for i in range(1000):
+#     cf.push_pedestal_frame(f.read_frame())
+
+# fig, ax = plt.subplots()
+# im = ax.imshow(cf.pedestal())
+# cf.pedestal()
+# cf.noise()
 
 
-n_bytes = data.itemsize*data.size
 
-print(f'Reading {N} frames took {t_elapsed:.3f}s {N/t_elapsed:.0f} FPS, {n_bytes/1024**2:.4f} GB/s')
+# N = 500
+# t0 = time.perf_counter()
+# hist1 = bh.Histogram(bh.axis.Regular(40, -2, 4000))
+# f.seek(0)
+
+# t0 = time.perf_counter()
+# data = f.read_n(N)
+# t_elapsed = time.perf_counter()-t0
 
 
-for frame in data:
-    a = cf.find_clusters(frame)
+# n_bytes = data.itemsize*data.size
 
-clusters = cf.steal_clusters()
+# print(f'Reading {N} frames took {t_elapsed:.3f}s {N/t_elapsed:.0f} FPS, {n_bytes/1024**2:.4f} GB/s')
+
+
+# for frame in data:
+#     a = cf.find_clusters(frame)
+
+# clusters = cf.steal_clusters()
 
 # t_elapsed = time.perf_counter()-t0
 # print(f'Clustering {N} frames took {t_elapsed:.2f}s  {N/t_elapsed:.0f} FPS')

python/src/cluster.hpp

@@ -1,5 +1,7 @@
 #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"
@@ -7,11 +9,13 @@
 #include <cstdint>
 #include <filesystem>
 #include <pybind11/pybind11.h>
+#include <pybind11/stl_bind.h>
 #include <pybind11/stl.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);
@@ -31,6 +35,9 @@ void define_cluster_vector(py::module &m, const std::string &typestr) {
                  auto *vec = new std::vector<T>(self.sum());
                  return return_vector(vec);
              })
+        .def_property_readonly("capacity", &ClusterVector<T>::capacity)
+        .def_property("frame_number", &ClusterVector<T>::frame_number,
+                      &ClusterVector<T>::set_frame_number)
         .def_buffer([typestr](ClusterVector<T> &self) -> py::buffer_info {
             return py::buffer_info(
                 self.data(),           /* Pointer to buffer */
@@ -45,6 +52,55 @@ void define_cluster_vector(py::module &m, const std::string &typestr) {
         });
 }
 
+void define_cluster_finder_mt_bindings(py::module &m) {
+    py::class_<ClusterFinderMT<uint16_t, pd_type>>(m, "ClusterFinderMT")
+        .def(py::init<Shape<2>, Shape<2>, pd_type, size_t, size_t>(),
+             py::arg("image_size"), py::arg("cluster_size"),
+             py::arg("n_sigma") = 5.0, py::arg("capacity") = 1000,
+             py::arg("n_threads") = 3)
+        .def("push_pedestal_frame",
+             [](ClusterFinderMT<uint16_t, pd_type> &self,
+                py::array_t<uint16_t> frame) {
+                 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",
+            [](ClusterFinderMT<uint16_t, pd_type> &self,
+               bool realloc_same_capacity) {
+                auto v = new ClusterVector<int>(
+                    self.steal_clusters(realloc_same_capacity));
+                return v;
+            },
+            py::arg("realloc_same_capacity") = false)
+        .def("stop", &ClusterFinderMT<uint16_t, pd_type>::stop);
+}
+
+
+void define_cluster_collector_bindings(py::module &m) {
+    py::class_<ClusterCollector>(m, "ClusterCollector")
+        .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);
+
+
+
+}
+
+
 void define_cluster_finder_bindings(py::module &m) {
     py::class_<ClusterFinder<uint16_t, pd_type>>(m, "ClusterFinder")
         .def(py::init<Shape<2>, Shape<2>, pd_type, size_t>(),

python/src/module.cpp

@@ -25,5 +25,8 @@ PYBIND11_MODULE(_aare, m) {
     define_pedestal_bindings<double>(m, "Pedestal_d");
     define_pedestal_bindings<float>(m, "Pedestal_f");
     define_cluster_finder_bindings(m);
+    define_cluster_finder_mt_bindings(m);
     define_cluster_file_io_bindings(m);
+    define_cluster_collector_bindings(m);
+
 }