From 6ff664f812478ebf40bc83c8a4d7d80a89f927f6 Mon Sep 17 00:00:00 2001
From: AliceMazzoleni99 <alice.mazzoleni@psi.ch>
Date: Fri, 17 Apr 2026 17:13:02 +0200
Subject: [PATCH] allow passing mask to clustervector (#304)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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

Co-authored-by: Erik Fröjdh <erik.frojdh@psi.ch>
---
 RELEASE.md                                  |  1 +
 docs/src/python/cluster/pyClusterVector.rst |  2 +-
 include/aare/ClusterVector.hpp              | 21 +++++++++++++++++++++
 python/src/bind_ClusterVector.hpp           | 16 ++++++++++++++++
 python/tests/test_ClusterVector.py          | 20 +++++++++++++++++++-
 5 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index 3057de7..9e50145 100644
--- a/RELEASE.md
+++ b/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:
 
diff --git a/docs/src/python/cluster/pyClusterVector.rst b/docs/src/python/cluster/pyClusterVector.rst
index 3dfe1e0..1aaf1b7 100644
--- a/docs/src/python/cluster/pyClusterVector.rst
+++ b/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:
diff --git a/include/aare/ClusterVector.hpp b/include/aare/ClusterVector.hpp
index a7f6dc4..0adcb5e 100644
--- a/include/aare/ClusterVector.hpp
+++ b/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) {
diff --git a/python/src/bind_ClusterVector.hpp b/python/src/bind_ClusterVector.hpp
index a9a71ac..254f96a 100644
--- a/python/src/bind_ClusterVector.hpp
+++ b/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);
diff --git a/python/tests/test_ClusterVector.py b/python/tests/test_ClusterVector.py
index f5b1e54..0ae96c3 100644
--- a/python/tests/test_ClusterVector.py
+++ b/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()
\ No newline at end of file
+    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()
\ No newline at end of file