allow passing mask to clustervector (#304)

- passing mask to ClusterVector 
- creates a copy of the ClusterVector

Co-authored-by: Erik Fröjdh <erik.frojdh@psi.ch>

RELEASE.md

@@ -7,6 +7,7 @@
 
 - Added a new Minuit2-based fitting framework for ``Gaussian``, ``RisingScurve``, ``FallingScurve``, ``Pol1`` and ``Pol2`` models.
 - setter and getter for nSigma for ClusterFinder ``aare.ClusterFinder().nSigma = 2``, ``aare.ClusterFinderMT().set_nSigma(2)`` 
+- mask opeartor for ClusterVector ``masked_clustervector = aare.ClusterVector()(mask)`` 
 
 ### Bugfixes:
 

docs/src/python/cluster/pyClusterVector.rst

@@ -38,7 +38,7 @@ C++ functions that support the ClusterVector or to view it as a numpy array.
 Below is the API of the ClusterVector_Cluster3x3i but all variants share the same API.
 
 .. autoclass::  aare._aare.ClusterVector_Cluster3x3i
-    :special-members: __init__
+    :special-members: __init__, __call__ 
     :members:
     :undoc-members:
     :show-inheritance:

include/aare/ClusterVector.hpp

@@ -59,6 +59,27 @@ class ClusterVector<Cluster<T, ClusterSizeX, ClusterSizeY, CoordType>> {
         other.m_data.clear();
     }
 
+    /**
+     * @brief Create a copy of the clustervector by filtering clusters in the
+     * ClusterVector using a boolean mask.
+     * @param mask boolean 1d mask
+     * @return ClusterVector containing only the clusters where the mask is
+     * true
+     */
+    ClusterVector operator()(NDView<bool, 1> mask) {
+        if (static_cast<size_t>(mask.size()) != m_data.size()) {
+            throw std::runtime_error(
+                LOCATION + "Mask size does not match number of clusters");
+        }
+        ClusterVector result(capacity(), frame_number());
+        for (size_t i = 0; i < m_data.size(); ++i) {
+            if (mask(i)) {
+                result.push_back(m_data[i]);
+            }
+        }
+        return result;
+    }
+
     // Move assignment operator
     ClusterVector &operator=(ClusterVector &&other) noexcept {
         if (this != &other) {

python/src/bind_ClusterVector.hpp

@@ -35,6 +35,22 @@ void define_ClusterVector(py::module &m, const std::string &typestr) {
 
         .def(py::init()) // TODO change!!!
 
+        .def(
+            "__call__",
+            [](ClusterVector<ClusterType> &self, py::array_t<bool> mask) {
+                return self(make_view_1d(mask));
+            },
+            py::arg("mask"), R"(
+            Create a copy of the clustervector and apply a boolean mask to the ClusterVector.
+
+            Parameters
+            ----------
+
+            mask : 1d boolean numpy array
+                Mask to apply to the ClusterVector. Must be the same length as the number of clusters in the ClusterVector. 
+
+            )")
+
         .def("push_back",
              [](ClusterVector<ClusterType> &self, const ClusterType &cluster) {
                  self.push_back(cluster);

python/tests/test_ClusterVector.py

@@ -109,4 +109,22 @@ def test_3x3_reduction():
     assert reduced_cv.size == 2
     assert reduced_cv[0]["x"] == 5
     assert reduced_cv[0]["y"] == 5
-    assert (reduced_cv[0]["data"] == np.array([[2.0, 1.0, 1.0], [2.0, 3.0, 1.0], [2.0, 1.0, 1.0]], dtype=np.double)).all()
+    assert (reduced_cv[0]["data"] == np.array([[2.0, 1.0, 1.0], [2.0, 3.0, 1.0], [2.0, 1.0, 1.0]], dtype=np.double)).all()
+
+
+def test_masking(): 
+
+    cv = _aare.ClusterVector_Cluster3x3i()
+    cv.push_back(_aare.Cluster3x3i(19, 22, np.array([0,1,0,2,3,0,2,1,0], dtype=np.int32)))
+    cv.push_back(_aare.Cluster3x3i(1, 2, np.ones(9, dtype=np.int32)))
+    assert cv.size == 2
+
+    mask = np.array([False, True], dtype=bool)
+    cv_masked = cv(mask)
+    assert cv_masked.size == 1
+
+    cv_masked_array = np.array(cv_masked, copy=False)
+    
+    assert cv_masked_array[0]["x"] == 1
+    assert cv_masked_array[0]["y"] == 2
+    assert (cv_masked_array[0]["data"] == np.ones((3,3),dtype=np.int32)).all()