added roi, noise and gain

include/aare/Cluster.hpp

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <algorithm>
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <numeric>
+
+namespace aare {
+
+//TODO! Template this? 
+struct Cluster3x3 {
+    int16_t x;
+    int16_t y;
+    int32_t data[9];
+
+    int32_t sum_2x2() const{
+        std::array<int32_t, 4> total;
+        total[0] = data[0] + data[1] + data[3] + data[4];
+        total[1] = data[1] + data[2] + data[4] + data[5];
+        total[2] = data[3] + data[4] + data[6] + data[7];
+        total[3] = data[4] + data[5] + data[7] + data[8];
+        return *std::max_element(total.begin(), total.end());
+    }
+
+    int32_t sum() const{
+        return std::accumulate(data, data + 9, 0);
+    }
+};
+struct Cluster2x2 {
+    int16_t x;
+    int16_t y;
+    int32_t data[4];
+};
+
+} // namespace aare

include/aare/ClusterFile.hpp

@@ -1,24 +1,15 @@
 #pragma once
 
+#include "aare/Cluster.hpp"
 #include "aare/ClusterVector.hpp"
 #include "aare/NDArray.hpp"
 #include "aare/defs.hpp"
 #include <filesystem>
 #include <fstream>
+#include <optional>
 
 namespace aare {
 
-//TODO! Template this? 
-struct Cluster3x3 {
-    int16_t x;
-    int16_t y;
-    int32_t data[9];
-};
-struct Cluster2x2 {
-    int16_t x;
-    int16_t y;
-    int32_t data[4];
-};
 
 typedef enum {
     cBottomLeft = 0,
@@ -80,6 +71,9 @@ class ClusterFile {
     uint32_t m_num_left{};
     size_t m_chunk_size{};
     const std::string m_mode;
+    std::optional<ROI> m_roi;
+    std::optional<NDArray<int32_t, 2>> m_noise_map;
+    std::optional<NDArray<double, 2>> m_gain_map;
 
   public:
     /**
@@ -104,7 +98,9 @@ class ClusterFile {
      */
     ClusterVector<int32_t> read_clusters(size_t n_clusters);
 
-    ClusterVector<int32_t> read_clusters(size_t n_clusters, ROI roi);
+    
+
+
 
     /**
      * @brief Read a single frame from the file and return the clusters. The
@@ -126,11 +122,23 @@ class ClusterFile {
      */
     size_t chunk_size() const { return m_chunk_size; }
 
+    void set_roi(ROI roi);
+
+    void set_noise_map(const NDView<int32_t, 2> noise_map);
+
+    void set_gain_map(const NDView<double, 2> gain_map);
+    
     
     /**
      * @brief Close the file. If not closed the file will be closed in the destructor
      */
     void close();
+
+    private:
+        ClusterVector<int32_t> read_clusters_with_cut(size_t n_clusters);
+        ClusterVector<int32_t> read_clusters_without_cut(size_t n_clusters);
+        bool is_selected(Cluster3x3 &cl);
+        Cluster3x3 read_one_cluster();
 };
 
 int analyze_data(int32_t *data, int32_t *t2, int32_t *t3, char *quad,
@@ -142,4 +150,6 @@ NDArray<double, 2> calculate_eta2(ClusterVector<int> &clusters);
 Eta2 calculate_eta2(Cluster3x3 &cl);
 Eta2 calculate_eta2(Cluster2x2 &cl);
 
+
+
 } // namespace aare

include/aare/ClusterVector.hpp

@@ -8,6 +8,9 @@
 
 #include <fmt/core.h>
 
+#include "aare/Cluster.hpp"
+#include "aare/NDView.hpp"
+
 namespace aare {
 
 /**
@@ -265,6 +268,28 @@ template <typename T, typename CoordType = int16_t> class ClusterVector {
         m_size = new_size;
     }
 
+    void apply_gain_map(const NDView<double> gain_map){
+        //in principle we need to know the size of the image for this lookup
+        //TODO! check orientations
+        std::array<int64_t, 9> xcorr = {-1, 0, 1, -1, 0, 1, -1, 0, 1};
+        std::array<int64_t, 9> ycorr = {-1, -1, -1, 0, 0, 0, 1, 1, 1};
+        for (size_t i=0; i<m_size; i++){
+            auto& cl = at<Cluster3x3>(i);
+
+            if (cl.x > 0 && cl.y > 0 && cl.x < gain_map.shape(1)-1 && cl.y < gain_map.shape(0)-1){
+                for (size_t j=0; j<9; j++){
+                    size_t x = cl.x + xcorr[j];
+                    size_t y = cl.y + ycorr[j];
+                    cl.data[j] = static_cast<T>(cl.data[j] * gain_map(y, x));
+                }
+            }else{
+                memset(cl.data, 0, 9*sizeof(T)); //clear edge clusters
+            }
+
+            
+        }
+    }
+
   private:
     void allocate_buffer(size_t new_capacity) {
         size_t num_bytes = item_size() * new_capacity;

include/aare/defs.hpp

@@ -6,6 +6,9 @@
 #include <array>
 #include <stdexcept>
 
+// #include <algorithm>
+// #include <array>
+// #include <numeric>
 #include <cassert>
 #include <cstdint>
 #include <cstring>
@@ -43,6 +46,7 @@ inline constexpr size_t bits_per_byte = 8;
 void assert_failed(const std::string &msg);
 
 
+
 class DynamicCluster {
   public:
     int cluster_sizeX;
@@ -215,6 +219,9 @@ struct ROI{
   
     int64_t height() const { return ymax - ymin; }
     int64_t width() const { return xmax - xmin; }
+    bool contains(int64_t x, int64_t y) const {
+        return x >= xmin && x < xmax && y >= ymin && y < ymax;
+    }
   };
 
 

python/src/cluster_file.hpp

@@ -31,26 +31,22 @@ void define_cluster_file_io_bindings(py::module &m) {
                 auto v = new ClusterVector<int32_t>(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)
         .def("read_frame",
              [](ClusterFile &self) {
                 auto v = new ClusterVector<int32_t>(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("read_cluster_with_cut",
-        //      [](ClusterFile &self, size_t n_clusters,
-        //         py::array_t<double> noise_map, int nx, int ny) {
-        //          auto view = make_view_2d(noise_map);
-        //          auto *vec =
-        //              new std::vector<Cluster3x3>(self.read_cluster_with_cut(
-        //                  n_clusters, view.data(), nx, ny));
-        //          return return_vector(vec);
-        //      })
         .def("__enter__", [](ClusterFile &self) { return &self; })
         .def("__exit__",
              [](ClusterFile &self,

src/ClusterFile.cpp

@@ -31,6 +31,18 @@ ClusterFile::ClusterFile(const std::filesystem::path &fname, size_t chunk_size,
     }
 }
 
+void ClusterFile::set_roi(ROI roi){
+    m_roi = roi;
+}
+
+void ClusterFile::set_noise_map(const NDView<int32_t, 2> noise_map){
+    m_noise_map = NDArray<int32_t, 2>(noise_map);
+}
+
+void ClusterFile::set_gain_map(const NDView<double, 2> gain_map){
+    m_gain_map = NDArray<double, 2>(gain_map);
+}
+
 ClusterFile::~ClusterFile() { close(); }
 
 void ClusterFile::close() {
@@ -55,7 +67,19 @@ void ClusterFile::write_frame(const ClusterVector<int32_t> &clusters) {
     fwrite(clusters.data(), clusters.item_size(), clusters.size(), fp);
 }
 
-ClusterVector<int32_t> ClusterFile::read_clusters(size_t n_clusters) {
+
+ClusterVector<int32_t> ClusterFile::read_clusters(size_t n_clusters){
+    if (m_mode != "r") {
+        throw std::runtime_error("File not opened for reading");
+    }
+    if (m_noise_map || m_roi){
+        return read_clusters_with_cut(n_clusters);
+    }else{
+        return read_clusters_without_cut(n_clusters);
+    }
+}
+
+ClusterVector<int32_t> ClusterFile::read_clusters_without_cut(size_t n_clusters) {
     if (m_mode != "r") {
         throw std::runtime_error("File not opened for reading");
     }
@@ -105,70 +129,69 @@ ClusterVector<int32_t> ClusterFile::read_clusters(size_t n_clusters) {
     // Resize the vector to the number of clusters.
     // No new allocation, only change bounds.
     clusters.resize(nph_read);
+    if(m_gain_map)
+        clusters.apply_gain_map(m_gain_map->view());
     return clusters;
 }
 
-ClusterVector<int32_t> ClusterFile::read_clusters(size_t n_clusters, ROI roi) {
-    if (m_mode != "r") {
-        throw std::runtime_error("File not opened for reading");
-    }
 
+
+ClusterVector<int32_t> ClusterFile::read_clusters_with_cut(size_t n_clusters) {
     ClusterVector<int32_t> clusters(3,3);
     clusters.reserve(n_clusters);
 
-    Cluster3x3 tmp; //this would break if the cluster size changes
-
-    
     // if there are photons left from previous frame read them first
     if (m_num_left) {
-        size_t nph_read = 0;
+        while(m_num_left && clusters.size() < n_clusters){
-        while(nph_read < m_num_left && clusters.size() < n_clusters){
+            Cluster3x3 c = read_one_cluster();
-            fread(&tmp, sizeof(tmp), 1, fp);
+            if(is_selected(c)){
-            nph_read++;
+                clusters.push_back(c.x, c.y, reinterpret_cast<std::byte*>(c.data));
-            if(tmp.x >= roi.xmin && tmp.x <= roi.xmax && tmp.y >= roi.ymin && tmp.y <= roi.ymax){
-                clusters.push_back(tmp.x, tmp.y, reinterpret_cast<std::byte*>(tmp.data));
             }
         }
-        m_num_left -= nph_read;
     }
 
-
+    // we did not have enough clusters left in the previous frame
+    // keep on reading frames until reaching n_clusters
     if (clusters.size() < n_clusters) {
+        // sanity check
         if (m_num_left) {
             throw std::runtime_error(LOCATION + "Entered second loop with clusters left\n");
         }
-        // we did not have enough clusters left in the previous frame
-        // keep on reading frames until reaching n_clusters
         
         int32_t frame_number = 0; // frame number needs to be 4 bytes!
         while (fread(&frame_number, sizeof(frame_number), 1, fp)) {
-            uint32_t nph_in_frame = 0; //number of photons we can read until next frame number
+            if (fread(&m_num_left, sizeof(m_num_left), 1, fp)) {
-            size_t nph_read = 0;       //number of photons read in this frame
+                clusters.set_frame_number(frame_number); //cluster vector will hold the last frame number
-
+                while(m_num_left && clusters.size() < n_clusters){
-            if (fread(&nph_in_frame, sizeof(nph_in_frame), 1, fp)) {
+                    Cluster3x3 c = read_one_cluster();
-                if(frame_number != 1){
+                    if(is_selected(c)){
-                    throw std::runtime_error("Frame number is not 1");
+                        clusters.push_back(c.x, c.y, reinterpret_cast<std::byte*>(c.data));
+                    }
+                }
             }
 
-                while(nph_read < nph_in_frame && clusters.size() < n_clusters){
+            // we have enough clusters, break out of the outer while loop
-                    fread(&tmp, sizeof(tmp), 1, fp);
+            if (clusters.size() >= n_clusters)
-                    nph_read++;
-                    if(tmp.x >= roi.xmin && tmp.x <= roi.xmax && tmp.y >= roi.ymin && tmp.y <= roi.ymax){
-                        clusters.push_back(tmp.x, tmp.y, reinterpret_cast<std::byte*>(tmp.data));
-                    }
-                }
-                m_num_left = nph_in_frame - nph_read;
-            }
-
-            if (clusters.size() >= n_clusters){
                 break;
         }
-        }
 
     }
+    if(m_gain_map)
+        clusters.apply_gain_map(m_gain_map->view());
+
     return clusters;
 }
 
+Cluster3x3 ClusterFile::read_one_cluster(){
+    Cluster3x3 c;
+    auto rc = fread(&c, sizeof(c), 1, fp);
+    if (rc != 1) {
+        throw std::runtime_error(LOCATION + "Could not read cluster");
+    }
+    --m_num_left;
+    return c;
+}
+
 ClusterVector<int32_t> ClusterFile::read_frame() {
     if (m_mode != "r") {
         throw std::runtime_error("File not opened for reading");
@@ -198,6 +221,26 @@ ClusterVector<int32_t> ClusterFile::read_frame() {
     return clusters;
 }
 
+bool ClusterFile::is_selected(Cluster3x3 &cl) {
+    //Should fail fast
+    if (m_roi) {
+        if (!(m_roi->contains(cl.x, cl.y))) {
+            return false;
+        }
+    }
+    if (m_noise_map){
+        int32_t sum_1x1 = cl.data[4]; // central pixel
+        int32_t sum_2x2 = cl.sum_2x2(); // highest sum of 2x2 subclusters
+        int32_t sum_3x3 = cl.sum(); // sum of all pixels
+
+        auto noise = (*m_noise_map)(cl.y, cl.x); //TODO! check if this is correct
+        if (sum_1x1 <= noise || sum_2x2 <= 2 * noise || sum_3x3 <= 3 * noise) {
+            return false;
+        }
+    }
+    //we passed all checks
+    return true;
+}
 
 // std::vector<Cluster3x3> ClusterFile::read_cluster_with_cut(size_t n_clusters,
 //                                                            double *noise_map,