From a24bbd9cf93961546451094b4a269a0b93289ae4 Mon Sep 17 00:00:00 2001
From: Mazzoleni Alice Francesca <mazzol_a@pc17378.psi.ch>
Date: Thu, 3 Apr 2025 11:56:25 +0200
Subject: [PATCH] started to do python refactoring

---
 python/src/cluster.hpp       | 185 ++++++++++++++++++++---------------
 python/src/cluster_file.hpp  |  89 ++++++++++-------
 python/src/interpolation.hpp |  51 ++++++----
 python/src/module.cpp        |  19 ++--
 4 files changed, 198 insertions(+), 146 deletions(-)

diff --git a/python/src/cluster.hpp b/python/src/cluster.hpp
index 3db816a..5657288 100644
--- a/python/src/cluster.hpp
+++ b/python/src/cluster.hpp
@@ -16,149 +16,176 @@
 namespace py = pybind11;
 using pd_type = double;
 
-template <typename T>
+using namespace aare;
+
+template <typename ClusterType>
 void define_cluster_vector(py::module &m, const std::string &typestr) {
+
+    using T = typename extract_template_arguments<ClusterType>::value_type;
+
     auto class_name = fmt::format("ClusterVector_{}", typestr);
-    py::class_<ClusterVector<T>>(m, class_name.c_str(), py::buffer_protocol())
-        .def(py::init<int, int>(),
-    py::arg("cluster_size_x") = 3, py::arg("cluster_size_y") = 3)
+    py::class_<ClusterVector<ClusterType>>(m, class_name.c_str(),
+                                           py::buffer_protocol())
+        .def(py::init<int, int>(), py::arg("cluster_size_x") = 3,
+             py::arg("cluster_size_y") = 3) // TODO change!!!
         .def("push_back",
-             [](ClusterVector<T> &self, int x, int y, py::array_t<T> data) {
-                //  auto view = make_view_2d(data);
-                 self.push_back(x, y, reinterpret_cast<const std::byte*>(data.data()));
+             [](ClusterVector<ClusterType> &self, ClusterType &cl) {
+                 //  auto view = make_view_2d(data);
+                 self.push_back(cl);
              })
-        .def_property_readonly("size", &ClusterVector<T>::size)
-        .def("item_size", &ClusterVector<T>::item_size)
+        .def_property_readonly("size", &ClusterVector<ClusterType>::size)
+        .def("item_size", &ClusterVector<ClusterType>::item_size)
         .def_property_readonly("fmt",
-                               [typestr](ClusterVector<T> &self) {
+                               [typestr](ClusterVector<ClusterType> &self) {
                                    return fmt::format(
                                        self.fmt_base(), self.cluster_size_x(),
                                        self.cluster_size_y(), typestr);
                                })
         .def("sum",
-             [](ClusterVector<T> &self) {
+             [](ClusterVector<ClusterType> &self) {
                  auto *vec = new std::vector<T>(self.sum());
                  return return_vector(vec);
              })
-        .def("sum_2x2", [](ClusterVector<T> &self) {
-            auto *vec = new std::vector<T>(self.sum_2x2());
-            return return_vector(vec);
-        })
-        .def_property_readonly("cluster_size_x", &ClusterVector<T>::cluster_size_x)
-        .def_property_readonly("cluster_size_y", &ClusterVector<T>::cluster_size_y)
-        .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 */
-                self.item_size(), /* Size of one scalar */
-                fmt::format(self.fmt_base(), self.cluster_size_x(),
-                            self.cluster_size_y(),
-                            typestr), /* Format descriptor */
-                1,                    /* Number of dimensions */
-                {self.size()},        /* Buffer dimensions */
-                {self.item_size()}    /* Strides (in bytes) for each index */
-            );
-        });
+        .def("sum_2x2",
+             [](ClusterVector<ClusterType> &self) {
+                 auto *vec = new std::vector<T>(self.sum_2x2());
+                 return return_vector(vec);
+             })
+        .def_property_readonly("cluster_size_x",
+                               &ClusterVector<ClusterType>::cluster_size_x)
+        .def_property_readonly("cluster_size_y",
+                               &ClusterVector<ClusterType>::cluster_size_y)
+        .def_property_readonly("capacity",
+                               &ClusterVector<ClusterType>::capacity)
+        .def_property("frame_number", &ClusterVector<ClusterType>::frame_number,
+                      &ClusterVector<ClusterType>::set_frame_number)
+        .def_buffer(
+            [typestr](ClusterVector<ClusterType> &self) -> py::buffer_info {
+                return py::buffer_info(
+                    self.data(),      /* Pointer to buffer */
+                    self.item_size(), /* Size of one scalar */
+                    fmt::format(self.fmt_base(), self.cluster_size_x(),
+                                self.cluster_size_y(),
+                                typestr), /* Format descriptor */
+                    1,                    /* Number of dimensions */
+                    {self.size()},        /* Buffer dimensions */
+                    {self.item_size()} /* Strides (in bytes) for each index */
+                );
+            });
 }
 
+template <typename ClusterType>
 void define_cluster_finder_mt_bindings(py::module &m) {
-    py::class_<ClusterFinderMT<uint16_t, pd_type>>(m, "ClusterFinderMT")
+    py::class_<ClusterFinderMT<ClusterType, 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") = 2048,
              py::arg("n_threads") = 3)
         .def("push_pedestal_frame",
-             [](ClusterFinderMT<uint16_t, pd_type> &self,
+             [](ClusterFinderMT<ClusterType, 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,
+            [](ClusterFinderMT<ClusterType, 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, frame_number);
                 return;
             },
             py::arg(), py::arg("frame_number") = 0)
-        .def("clear_pedestal", &ClusterFinderMT<uint16_t, pd_type>::clear_pedestal)
-        .def("sync", &ClusterFinderMT<uint16_t, pd_type>::sync)
-        .def("stop", &ClusterFinderMT<uint16_t, pd_type>::stop)
-        .def("start", &ClusterFinderMT<uint16_t, pd_type>::start)
-        .def("pedestal",
-             [](ClusterFinderMT<uint16_t, pd_type> &self, size_t thread_index) {
-                 auto pd = new NDArray<pd_type, 2>{};
-                 *pd = self.pedestal(thread_index);
-                 return return_image_data(pd);
-             },py::arg("thread_index") = 0)
-        .def("noise",
-             [](ClusterFinderMT<uint16_t, pd_type> &self, size_t thread_index) {
-                 auto arr = new NDArray<pd_type, 2>{};
-                 *arr = self.noise(thread_index);
-                 return return_image_data(arr);
-             },py::arg("thread_index") = 0);
+        .def("clear_pedestal",
+             &ClusterFinderMT<ClusterType, uint16_t, pd_type>::clear_pedestal)
+        .def("sync", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::sync)
+        .def("stop", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::stop)
+        .def("start", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::start)
+        .def(
+            "pedestal",
+            [](ClusterFinderMT<ClusterType, uint16_t, pd_type> &self,
+               size_t thread_index) {
+                auto pd = new NDArray<pd_type, 2>{};
+                *pd = self.pedestal(thread_index);
+                return return_image_data(pd);
+            },
+            py::arg("thread_index") = 0)
+        .def(
+            "noise",
+            [](ClusterFinderMT<ClusterType, uint16_t, pd_type> &self,
+               size_t thread_index) {
+                auto arr = new NDArray<pd_type, 2>{};
+                *arr = self.noise(thread_index);
+                return return_image_data(arr);
+            },
+            py::arg("thread_index") = 0);
 }
 
+template <typename ClusterType>
 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)
+    py::class_<ClusterCollector<ClusterType>>(m, "ClusterCollector")
+        .def(py::init<ClusterFinderMT<ClusterType, uint16_t, double> *>())
+        .def("stop", &ClusterCollector<ClusterType>::stop)
         .def(
             "steal_clusters",
-            [](ClusterCollector &self) {
-                auto v =
-                    new std::vector<ClusterVector<int>>(self.steal_clusters());
+            [](ClusterCollector<ClusterType> &self) {
+                auto v = new std::vector<ClusterVector<ClusterType>>(
+                    self.steal_clusters());
                 return v;
             },
             py::return_value_policy::take_ownership);
 }
 
+template <typename ClusterType>
 void define_cluster_file_sink_bindings(py::module &m) {
-    py::class_<ClusterFileSink>(m, "ClusterFileSink")
-        .def(py::init<ClusterFinderMT<uint16_t, double, int32_t> *,
+    py::class_<ClusterFileSink<ClusterType>>(m, "ClusterFileSink")
+        .def(py::init<ClusterFinderMT<ClusterType, uint16_t, double> *,
                       const std::filesystem::path &>())
-        .def("stop", &ClusterFileSink::stop);
+        .def("stop", &ClusterFileSink<ClusterType>::stop);
 }
 
+template <typename ClusterType>
 void define_cluster_finder_bindings(py::module &m) {
-    py::class_<ClusterFinder<uint16_t, pd_type>>(m, "ClusterFinder")
+    py::class_<ClusterFinder<ClusterType, uint16_t, pd_type>>(m,
+                                                              "ClusterFinder")
         .def(py::init<Shape<2>, Shape<2>, pd_type, size_t>(),
              py::arg("image_size"), py::arg("cluster_size"),
              py::arg("n_sigma") = 5.0, py::arg("capacity") = 1'000'000)
         .def("push_pedestal_frame",
-             [](ClusterFinder<uint16_t, pd_type> &self,
+             [](ClusterFinder<ClusterType, uint16_t, pd_type> &self,
                 py::array_t<uint16_t> frame) {
                  auto view = make_view_2d(frame);
                  self.push_pedestal_frame(view);
              })
-        .def("clear_pedestal", &ClusterFinder<uint16_t, pd_type>::clear_pedestal)
-        .def_property_readonly("pedestal",
-                               [](ClusterFinder<uint16_t, pd_type> &self) {
-                                   auto pd = new NDArray<pd_type, 2>{};
-                                   *pd = self.pedestal();
-                                   return return_image_data(pd);
-                               })
-        .def_property_readonly("noise",
-                               [](ClusterFinder<uint16_t, pd_type> &self) {
-                                   auto arr = new NDArray<pd_type, 2>{};
-                                   *arr = self.noise();
-                                   return return_image_data(arr);
-                               })
+        .def("clear_pedestal",
+             &ClusterFinder<ClusterType, uint16_t, pd_type>::clear_pedestal)
+        .def_property_readonly(
+            "pedestal",
+            [](ClusterFinder<ClusterType, uint16_t, pd_type> &self) {
+                auto pd = new NDArray<pd_type, 2>{};
+                *pd = self.pedestal();
+                return return_image_data(pd);
+            })
+        .def_property_readonly(
+            "noise",
+            [](ClusterFinder<ClusterType, uint16_t, pd_type> &self) {
+                auto arr = new NDArray<pd_type, 2>{};
+                *arr = self.noise();
+                return return_image_data(arr);
+            })
         .def(
             "steal_clusters",
-            [](ClusterFinder<uint16_t, pd_type> &self,
+            [](ClusterFinder<ClusterType, uint16_t, pd_type> &self,
                bool realloc_same_capacity) {
-                auto v = new ClusterVector<int>(
+                auto v = new ClusterVector<ClusterType>(
                     self.steal_clusters(realloc_same_capacity));
                 return v;
             },
             py::arg("realloc_same_capacity") = false)
         .def(
             "find_clusters",
-            [](ClusterFinder<uint16_t, pd_type> &self,
+            [](ClusterFinder<ClusterType, 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, frame_number);
@@ -167,7 +194,7 @@ void define_cluster_finder_bindings(py::module &m) {
             py::arg(), py::arg("frame_number") = 0);
 
     m.def("hitmap",
-          [](std::array<size_t, 2> image_size, ClusterVector<int32_t> &cv) {
+          [](std::array<size_t, 2> image_size, ClusterVector<ClusterType> &cv) {
               py::array_t<int32_t> hitmap(image_size);
               auto r = hitmap.mutable_unchecked<2>();
 
diff --git a/python/src/cluster_file.hpp b/python/src/cluster_file.hpp
index b807712..576c3bb 100644
--- a/python/src/cluster_file.hpp
+++ b/python/src/cluster_file.hpp
@@ -10,69 +10,86 @@
 #include <pybind11/stl/filesystem.h>
 #include <string>
 
-//Disable warnings for unused parameters, as we ignore some
-//in the __exit__ method
+// Disable warnings for unused parameters, as we ignore some
+// in the __exit__ method
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wunused-parameter"
 
-
 namespace py = pybind11;
 using namespace ::aare;
 
+template <typename ClusterType>
 void define_cluster_file_io_bindings(py::module &m) {
     PYBIND11_NUMPY_DTYPE(Cluster3x3, x, y, data);
 
-    py::class_<ClusterFile>(m, "ClusterFile")
+    py::class_<ClusterFile<ClusterType>>(m, "ClusterFile")
         .def(py::init<const std::filesystem::path &, size_t,
                       const std::string &>(),
              py::arg(), py::arg("chunk_size") = 1000, py::arg("mode") = "r")
-        .def("read_clusters",
-             [](ClusterFile &self, size_t n_clusters) {
-                auto v = new ClusterVector<int32_t>(self.read_clusters(n_clusters));
+        .def(
+            "read_clusters",
+            [](ClusterFile<ClusterType> &self, size_t n_clusters) {
+                auto v = new ClusterVector<ClusterType>(
+                    self.read_clusters(n_clusters));
                 return v;
-             },py::return_value_policy::take_ownership)
-        .def("read_clusters",
-                [](ClusterFile &self, size_t n_clusters, ROI roi) {
-                   auto v = new ClusterVector<int32_t>(self.read_clusters(n_clusters, roi));
-                   return v;
-                },py::return_value_policy::take_ownership)
+            },
+            py::return_value_policy::take_ownership)
+        .def(
+            "read_clusters",
+            [](ClusterFile<ClusterType> &self, size_t n_clusters, ROI roi) {
+                auto v = new ClusterVector<ClusterType>(
+                    self.read_clusters(n_clusters, roi));
+                return v;
+            },
+            py::return_value_policy::take_ownership)
         .def("read_frame",
-             [](ClusterFile &self) {
-                auto v = new ClusterVector<int32_t>(self.read_frame());
-                return v;
+             [](ClusterFile<ClusterType> &self) {
+                 auto v = new ClusterVector<ClusterType>(self.read_frame());
+                 return v;
              })
-        .def("set_roi", &ClusterFile::set_roi)
-        .def("set_noise_map", [](ClusterFile &self, py::array_t<int32_t> noise_map) {
-            auto view = make_view_2d(noise_map);
-            self.set_noise_map(view);
-        })
-        .def("set_gain_map", [](ClusterFile &self, py::array_t<double> gain_map) {
-            auto view = make_view_2d(gain_map);
-            self.set_gain_map(view);
-        })
-        .def("close", &ClusterFile::close)
-        .def("write_frame", &ClusterFile::write_frame)
-        .def("__enter__", [](ClusterFile &self) { return &self; })
+        .def("set_roi", &ClusterFile<ClusterType>::set_roi)
+        .def(
+            "set_noise_map",
+            [](ClusterFile<ClusterType> &self, py::array_t<int32_t> noise_map) {
+                auto view = make_view_2d(noise_map);
+                self.set_noise_map(view);
+            })
+
+        .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);
+             })
+
+        // void set_gain_map(const GainMap &gain_map); //TODO do i need a
+        // gainmap constructor?
+
+        .def("close", &ClusterFile<ClusterType>::close)
+        .def("write_frame", &ClusterFile<ClusterType>::write_frame)
+        .def("__enter__", [](ClusterFile<ClusterType> &self) { return &self; })
         .def("__exit__",
-             [](ClusterFile &self,
+             [](ClusterFile<ClusterType> &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__", [](ClusterFile &self) { return &self; })
-        .def("__next__", [](ClusterFile &self) {
-            auto v = new ClusterVector<int32_t>(self.read_clusters(self.chunk_size()));
+        .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;
         });
 
-        m.def("calculate_eta2", []( aare::ClusterVector<int32_t> &clusters) {
-            auto eta2 = new NDArray<double, 2>(calculate_eta2(clusters));
-            return return_image_data(eta2);
-        });
+    /*
+    m.def("calculate_eta2", []( aare::ClusterVector<ClusterType> &clusters) {
+        auto eta2 = new NDArray<double, 2>(calculate_eta2(clusters));
+        return return_image_data(eta2);
+    });
+    */ //add in different file
 }
 
 #pragma GCC diagnostic pop
\ No newline at end of file
diff --git a/python/src/interpolation.hpp b/python/src/interpolation.hpp
index 02742e1..47b5203 100644
--- a/python/src/interpolation.hpp
+++ b/python/src/interpolation.hpp
@@ -8,31 +8,40 @@
 #include <pybind11/stl.h>
 
 namespace py = pybind11;
+
 void define_interpolation_bindings(py::module &m) {
 
-    PYBIND11_NUMPY_DTYPE(aare::Photon, x,y,energy);
+    PYBIND11_NUMPY_DTYPE(aare::Photon, x, y, energy);
 
     py::class_<aare::Interpolator>(m, "Interpolator")
-        .def(py::init([](py::array_t<double, py::array::c_style | py::array::forcecast> etacube, py::array_t<double> xbins,
-                         py::array_t<double> ybins, py::array_t<double> ebins) {
-            return Interpolator(make_view_3d(etacube), make_view_1d(xbins),
-                                make_view_1d(ybins), make_view_1d(ebins));
-        }))
-        .def("get_ietax", [](Interpolator& self){
-            auto*ptr = new NDArray<double,3>{};
-            *ptr = self.get_ietax();
-            return return_image_data(ptr);
-        })
-        .def("get_ietay", [](Interpolator& self){
-            auto*ptr = new NDArray<double,3>{};
-            *ptr = self.get_ietay();
-            return return_image_data(ptr);
-        })
-        .def("interpolate", [](Interpolator& self, const ClusterVector<int32_t>& clusters){
-            auto photons = self.interpolate(clusters);
-            auto* ptr = new std::vector<Photon>{photons};
-            return return_vector(ptr);
-        });
+        .def(py::init(
+            [](py::array_t<double, py::array::c_style | py::array::forcecast>
+                   etacube,
+               py::array_t<double> xbins, py::array_t<double> ybins,
+               py::array_t<double> ebins) {
+                return Interpolator(make_view_3d(etacube), make_view_1d(xbins),
+                                    make_view_1d(ybins), make_view_1d(ebins));
+            }))
+        .def("get_ietax",
+             [](Interpolator &self) {
+                 auto *ptr = new NDArray<double, 3>{};
+                 *ptr = self.get_ietax();
+                 return return_image_data(ptr);
+             })
+        .def("get_ietay",
+             [](Interpolator &self) {
+                 auto *ptr = new NDArray<double, 3>{};
+                 *ptr = self.get_ietay();
+                 return return_image_data(ptr);
+             })
+
+        // TODO take care of clustertype template
+        .def("interpolate",
+             [](Interpolator &self, const ClusterVector<auto> &clusters) {
+                 auto photons = self.interpolate<ClusterType>(clusters);
+                 auto *ptr = new std::vector<Photon>{photons};
+                 return return_vector(ptr);
+             });
 
     // TODO! Evaluate without converting to double
     m.def(
diff --git a/python/src/module.cpp b/python/src/module.cpp
index 43f48ba..ba681be 100644
--- a/python/src/module.cpp
+++ b/python/src/module.cpp
@@ -1,17 +1,17 @@
-//Files with bindings to the different classes
-#include "file.hpp"
-#include "raw_file.hpp"
-#include "ctb_raw_file.hpp"
-#include "raw_master_file.hpp"
-#include "var_cluster.hpp"
-#include "pixel_map.hpp"
-#include "pedestal.hpp"
+// Files with bindings to the different classes
 #include "cluster.hpp"
 #include "cluster_file.hpp"
+#include "ctb_raw_file.hpp"
+#include "file.hpp"
 #include "fit.hpp"
 #include "interpolation.hpp"
+#include "pedestal.hpp"
+#include "pixel_map.hpp"
+#include "raw_file.hpp"
+#include "raw_master_file.hpp"
+#include "var_cluster.hpp"
 
-//Pybind stuff
+// Pybind stuff
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -33,5 +33,4 @@ PYBIND11_MODULE(_aare, m) {
     define_cluster_file_sink_bindings(m);
     define_fit_bindings(m);
     define_interpolation_bindings(m);
-
 }
\ No newline at end of file