Cluster cuts (#146)

Co-authored-by: Patrick <patrick.sieberer@psi.ch>
Co-authored-by: JulianHeymes <julian.heymes@psi.ch>
Co-authored-by: Dhanya Thattil <dhanya.thattil@psi.ch>
Co-authored-by: Xiangyu Xie <45243914+xiangyuxie@users.noreply.github.com>
Co-authored-by: xiangyu.xie <xiangyu.xie@psi.ch>

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,25 +1,17 @@
 #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];
-};
 
+//TODO! Legacy enums, migrate to enum class
 typedef enum {
     cBottomLeft = 0,
     cBottomRight = 1,
@@ -53,15 +45,7 @@ struct ClusterAnalysis {
     double etay;
 };
 
-/*
-Binary cluster file. Expects data to be layed out as:
-int32_t frame_number
-uint32_t number_of_clusters
-int16_t x, int16_t y, int32_t data[9] x number_of_clusters
-int32_t frame_number
-uint32_t number_of_clusters
-....
-*/
+
 
 /**
  * @brief Class to read and write cluster files
@@ -70,16 +54,19 @@ uint32_t number_of_clusters
  *
  *       int32_t frame_number
  *       uint32_t number_of_clusters
- *       int16_t x, int16_t y, int32_t data[9] x number_of_clusters
+ *       int16_t x, int16_t y, int32_t data[9] * number_of_clusters
  *       int32_t frame_number
  *       uint32_t number_of_clusters
  *       etc.
  */
 class ClusterFile {
     FILE *fp{};
-    uint32_t m_num_left{};
-    size_t m_chunk_size{};
-    const std::string m_mode;
+    uint32_t m_num_left{};            /*Number of photons left in frame*/
+    size_t m_chunk_size{};            /*Number of clusters to read at a time*/
+    const std::string m_mode;         /*Mode to open the file in*/
+    std::optional<ROI> m_roi;         /*Region of interest, will be applied if set*/
+    std::optional<NDArray<int32_t, 2>> m_noise_map; /*Noise map to cut photons, will be applied if set*/
+    std::optional<NDArray<double, 2>> m_gain_map; /*Gain map to apply to the clusters, will be applied if set*/
 
   public:
     /**
@@ -117,29 +104,49 @@ class ClusterFile {
 
     void write_frame(const ClusterVector<int32_t> &clusters);
     
-    // Need to be migrated to support NDArray and return a ClusterVector
-    // std::vector<Cluster3x3>
-    // read_cluster_with_cut(size_t n_clusters, double *noise_map, int nx, int ny);
-
     /**
      * @brief Return the chunk size
      */
     size_t chunk_size() const { return m_chunk_size; }
+
+    /**
+     * @brief Set the region of interest to use when reading clusters. If set only clusters within
+     * the ROI will be read.
+     */
+    void set_roi(ROI roi);
+
+    /**
+     * @brief Set the noise map to use when reading clusters. If set clusters below the noise
+     * level will be discarded. Selection criteria one of: Central pixel above noise, highest
+     * 2x2 sum above 2 * noise, total sum above 3 * noise.
+     */
+    void set_noise_map(const NDView<int32_t, 2> noise_map);
+
+    /**
+     * @brief Set the gain map to use when reading clusters. If set the gain map will be applied
+     * to the clusters that pass ROI and noise_map selection.
+     */
+    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);
+        ClusterVector<int32_t> read_frame_with_cut();
+        ClusterVector<int32_t> read_frame_without_cut();
+        bool is_selected(Cluster3x3 &cl);
+        Cluster3x3 read_one_cluster();
 };
 
-int analyze_data(int32_t *data, int32_t *t2, int32_t *t3, char *quad,
-                 double *eta2x, double *eta2y, double *eta3x, double *eta3y);
-int analyze_cluster(Cluster3x3 &cl, int32_t *t2, int32_t *t3, char *quad,
-                    double *eta2x, double *eta2y, double *eta3x, double *eta3y);
-
+//TODO! helper functions that doesn't really belong here
 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

@@ -1,11 +1,9 @@
 #pragma once
 
 #include "aare/Dtype.hpp"
-// #include "aare/utils/logger.hpp"
 
 #include <array>
 #include <stdexcept>
-
 #include <cassert>
 #include <cstdint>
 #include <cstring>
@@ -43,6 +41,7 @@ inline constexpr size_t bits_per_byte = 8;
 void assert_failed(const std::string &msg);
 
 
+
 class DynamicCluster {
   public:
     int cluster_sizeX;
@@ -215,6 +214,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;
+    }
   };