From 9a7713e98a86445eaa7953b825a7bd7baf3d5fee Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Tue, 22 Jul 2025 16:42:09 +0200
Subject: [PATCH 01/13] added g0 calibration, pedestal and pixel counting

---
 CMakeLists.txt                    |   1 +
 include/aare/NDView.hpp           |  29 +++-
 include/aare/VarClusterFinder.hpp |   6 +-
 include/aare/calibration.hpp      | 211 +++++++++++++++++++++++++++++-
 python/aare/__init__.py           |   3 +-
 python/src/bind_calibration.hpp   |  71 +++++++++-
 src/NDArray.test.cpp              |   4 +-
 src/calibration.cpp               |  78 +++++++++++
 8 files changed, 379 insertions(+), 24 deletions(-)
 create mode 100644 src/calibration.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index d6a22d7..3cfc6a7 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -368,6 +368,7 @@ set(PUBLICHEADERS
 
 
 set(SourceFiles
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/calibration.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/CtbRawFile.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/decode.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/src/defs.cpp
diff --git a/include/aare/NDView.hpp b/include/aare/NDView.hpp
index 28f5371..4c16742 100644
--- a/include/aare/NDView.hpp
+++ b/include/aare/NDView.hpp
@@ -66,17 +66,28 @@ class NDView : public ArrayExpr<NDView<T, Ndim>, Ndim> {
         : buffer_(buffer), strides_(c_strides<Ndim>(shape)), shape_(shape),
           size_(std::accumulate(std::begin(shape), std::end(shape), 1,
                                 std::multiplies<>())) {}
-
+                   
     template <typename... Ix>
     std::enable_if_t<sizeof...(Ix) == Ndim, T &> operator()(Ix... index) {
         return buffer_[element_offset(strides_, index...)];
     }
 
     template <typename... Ix>
-    const std::enable_if_t<sizeof...(Ix) == Ndim, T &> operator()(Ix... index) const {
+    std::enable_if_t<sizeof...(Ix) == 1 && (Ndim > 1), NDView<T, Ndim - 1>> operator()(Ix... index) {
+        // return a view of the next dimension
+        std::array<ssize_t, Ndim - 1> new_shape{};
+        std::copy_n(shape_.begin() + 1, Ndim - 1, new_shape.begin());
+        return NDView<T, Ndim - 1>(&buffer_[element_offset(strides_, index...)],
+                                   new_shape);
+        
+    }
+
+    template <typename... Ix>
+    std::enable_if_t<sizeof...(Ix) == Ndim, const T &> operator()(Ix... index) const {
         return buffer_[element_offset(strides_, index...)];
     }
 
+
     ssize_t size() const { return static_cast<ssize_t>(size_); }
     size_t total_bytes() const { return size_ * sizeof(T); }
     std::array<ssize_t, Ndim> strides() const noexcept { return strides_; }
@@ -85,9 +96,19 @@ class NDView : public ArrayExpr<NDView<T, Ndim>, Ndim> {
     T *end() { return buffer_ + size_; }
     T const *begin() const { return buffer_; }
     T const *end() const { return buffer_ + size_; }
-    T &operator()(ssize_t i)  { return buffer_[i]; }
+    
+    
+
+
+
+    /**
+     * @brief Access element at index i.
+     */
     T &operator[](ssize_t i)  { return buffer_[i]; }
-    const T &operator()(ssize_t i) const { return buffer_[i]; }
+
+    /**
+     * @brief Access element at index i.
+     */
     const T &operator[](ssize_t i) const { return buffer_[i]; }
 
     bool operator==(const NDView &other) const {
diff --git a/include/aare/VarClusterFinder.hpp b/include/aare/VarClusterFinder.hpp
index 3679d39..a002f9b 100644
--- a/include/aare/VarClusterFinder.hpp
+++ b/include/aare/VarClusterFinder.hpp
@@ -240,14 +240,14 @@ template <typename T> void VarClusterFinder<T>::first_pass() {
 
     for (ssize_t i = 0; i < original_.size(); ++i) {
         if (use_noise_map)
-            threshold_ = 5 * noiseMap(i);
-        binary_(i) = (original_(i) > threshold_);
+            threshold_ = 5 * noiseMap[i];
+        binary_[i] = (original_[i] > threshold_);
     }
 
     for (int i = 0; i < shape_[0]; ++i) {
         for (int j = 0; j < shape_[1]; ++j) {
 
-            // do we have someting to process?
+            // do we have something to process?
             if (binary_(i, j)) {
                 auto tmp = check_neighbours(i, j);
                 if (tmp != 0) {
diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index aab1f89..8d672b4 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -1,5 +1,7 @@
 #pragma once
 
+#include "aare/NDArray.hpp"
+#include "aare/NDView.hpp"
 #include "aare/defs.hpp"
 #include "aare/utils/task.hpp"
 #include <cstdint>
@@ -55,32 +57,227 @@ ALWAYS_INLINE std::pair<uint16_t, int16_t> get_value_and_gain(uint16_t raw) {
 
 template <class T>
 void apply_calibration_impl(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
-                       NDView<T, 3> ped, NDView<T, 3> cal, int start,
-                       int stop) {
+                            NDView<T, 3> ped, NDView<T, 3> cal, int start,
+                            int stop) {
 
     for (int frame_nr = start; frame_nr != stop; ++frame_nr) {
         for (int row = 0; row != raw_data.shape(1); ++row) {
             for (int col = 0; col != raw_data.shape(2); ++col) {
-                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
+                auto [value, gain] =
+                    get_value_and_gain(raw_data(frame_nr, row, col));
+
+                // Using multiplication does not seem to speed up the code here
+                // ADU/keV is the standard unit for the calibration which
+                // means rewriting the formula is not worth it.
                 res(frame_nr, row, col) =
-                    (value - ped(gain, row, col)) / cal(gain, row, col); //TODO! use multiplication
+                    (value - ped(gain, row, col)) / cal(gain, row, col);
             }
         }
     }
 }
 
 template <class T>
+void apply_calibration_impl(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
+                            NDView<T, 2> ped, NDView<T, 2> cal, int start,
+                            int stop) {
+
+    for (int frame_nr = start; frame_nr != stop; ++frame_nr) {
+        for (int row = 0; row != raw_data.shape(1); ++row) {
+            for (int col = 0; col != raw_data.shape(2); ++col) {
+                auto [value, gain] =
+                    get_value_and_gain(raw_data(frame_nr, row, col));
+
+                // Using multiplication does not seem to speed up the code here
+                // ADU/keV is the standard unit for the calibration which
+                // means rewriting the formula is not worth it.
+
+                // ignore anything that is not gain 0
+                // TODO! deal with fixed gain?
+                if (gain == 0)
+                    res(frame_nr, row, col) =
+                        (value - ped(row, col)) / cal(row, col);
+            }
+        }
+    }
+}
+
+template <class T, ssize_t Ndim=3>
 void apply_calibration(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
-                       NDView<T, 3> ped, NDView<T, 3> cal,
+                       NDView<T, Ndim> ped, NDView<T, Ndim> cal,
                        ssize_t n_threads = 4) {
     std::vector<std::future<void>> futures;
     futures.reserve(n_threads);
     auto limits = split_task(0, raw_data.shape(0), n_threads);
     for (const auto &lim : limits)
-        futures.push_back(std::async(&apply_calibration_impl<T>, res, raw_data, ped, cal,
-                                     lim.first, lim.second));
+        futures.push_back(std::async(
+            static_cast<void (*)(NDView<T, 3>, NDView<uint16_t, 3>,
+                                 NDView<T, Ndim>, NDView<T, Ndim>, int, int)>(
+                apply_calibration_impl),
+            res, raw_data, ped, cal, lim.first, lim.second));
     for (auto &f : futures)
         f.get();
 }
 
+std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>
+sum_and_count_per_gain(NDView<uint16_t, 3> raw_data);
+
+std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>
+sum_and_count_g0(NDView<uint16_t, 3> raw_data);
+
+template <typename T>
+NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data) {
+    auto [accumulator, count] = sum_and_count_per_gain(raw_data);
+
+    NDArray<T, 3> pedestal(
+        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (int gain = 0; gain < 3; ++gain) {
+        for (int row = 0; row < raw_data.shape(1); ++row) {
+            for (int col = 0; col < raw_data.shape(2); ++col) {
+                if (count(gain, row, col) != 0) {
+                    pedestal(gain, row, col) =
+                        static_cast<T>(accumulator(gain, row, col)) /
+                        static_cast<T>(count(gain, row, col));
+                }
+            }
+        }
+    }
+    return pedestal;
+}
+
+template <typename T>
+NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data) {
+    auto [accumulator, count] = sum_and_count_g0(raw_data);
+
+    NDArray<T, 2> pedestal(
+        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (int row = 0; row < raw_data.shape(1); ++row) {
+        for (int col = 0; col < raw_data.shape(2); ++col) {
+            if (count(row, col) != 0) {
+                pedestal(row, col) =
+                    static_cast<T>(accumulator(row, col)) /
+                    static_cast<T>(count(row, col));
+            }
+        }
+    }
+    return pedestal;
+}
+template <typename T>
+NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
+                                 ssize_t n_threads) {
+    std::vector<std::future<std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>>>
+        futures;
+    futures.reserve(n_threads);
+    auto limits = split_task(0, raw_data.shape(0), n_threads);  
+    // make subviews for each thread
+    std::vector<NDView<uint16_t, 3>> subviews;
+    for (const auto &lim : limits) {
+        subviews.emplace_back(
+            raw_data.data() + lim.first * raw_data.strides()[0],
+            std::array<ssize_t, 3>{lim.second - lim.first, raw_data.shape(1),
+                                   raw_data.shape(2)});
+    }
+    for (auto view : subviews) {
+        futures.push_back(std::async(
+            static_cast<std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>> (*)(
+                NDView<uint16_t, 3>)>(&sum_and_count_g0),
+            view));
+    }
+    NDArray<size_t, 2> accumulator(
+        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    NDArray<size_t, 2> count(
+        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (auto &f : futures) {
+        auto [acc, cnt] = f.get();
+        accumulator += acc;
+        count += cnt;
+    }
+    NDArray<T, 2> pedestal(
+        std::array<ssize_t,2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (int gain = 0; gain < 3; ++gain) {
+        for (int row = 0; row < raw_data.shape(1); ++row) {
+            for (int col = 0; col < raw_data.shape(2); ++col) {
+                if (count(row, col) != 0) {
+                    pedestal(row, col) =
+                        static_cast<T>(accumulator(row, col)) /
+                        static_cast<T>(count(row, col));
+                }
+            }
+        }
+    }
+    return pedestal;
+
+}
+
+
+template <typename T>
+NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data,
+                                 ssize_t n_threads) {
+    NDArray<int, 2> switched(
+        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    std::vector<std::future<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>>>
+        futures;
+    futures.reserve(n_threads);
+    auto limits = split_task(0, raw_data.shape(0), n_threads);
+
+    // make subviews for each thread
+    std::vector<NDView<uint16_t, 3>> subviews;
+    for (const auto &lim : limits) {
+        subviews.emplace_back(
+            raw_data.data() + lim.first * raw_data.strides()[0],
+            std::array<ssize_t, 3>{lim.second - lim.first, raw_data.shape(1),
+                                   raw_data.shape(2)});
+    }
+    for (auto view : subviews) {
+        futures.push_back(std::async(
+            static_cast<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>> (*)(
+                NDView<uint16_t, 3>)>(&sum_and_count_per_gain),
+            view));
+    }
+
+    NDArray<size_t, 3> accumulator(
+        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+    NDArray<size_t, 3> count(
+        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (auto &f : futures) {
+        auto [acc, cnt] = f.get();
+        accumulator += acc;
+        count += cnt;
+    }
+
+    NDArray<T, 3> pedestal(
+        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (int gain = 0; gain < 3; ++gain) {
+        for (int row = 0; row < raw_data.shape(1); ++row) {
+            for (int col = 0; col < raw_data.shape(2); ++col) {
+                if (count(gain, row, col) != 0) {
+                    pedestal(gain, row, col) =
+                        static_cast<T>(accumulator(gain, row, col)) /
+                        static_cast<T>(count(gain, row, col));
+                }
+            }
+        }
+    }
+    return pedestal;
+}
+
+/**
+ * @brief Count the number of switching pixels in the raw data.
+ * This function counts the number of pixels that switch between G1 and G2 gain.
+ * It returns an NDArray with the number of switching pixels per pixel.
+ * @param raw_data The NDView containing the raw data
+ * @return An NDArray with the number of switching pixels per pixel
+ */
+NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data);
+
+/**
+ * @brief Count the number of switching pixels in the raw data.
+ * This function counts the number of pixels that switch between G1 and G2 gain.
+ * It returns an NDArray with the number of switching pixels per pixel.
+ * @param raw_data The NDView containing the raw data
+ * @param n_threads The number of threads to use for parallel processing
+ * @return An NDArray with the number of switching pixels per pixel
+ */
+NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data,
+                                       ssize_t n_threads);
+
 } // namespace aare
\ No newline at end of file
diff --git a/python/aare/__init__.py b/python/aare/__init__.py
index c48da49..1d34a85 100644
--- a/python/aare/__init__.py
+++ b/python/aare/__init__.py
@@ -32,6 +32,7 @@ from .utils import random_pixels, random_pixel, flat_list, add_colorbar
 from .func import *
 
 from .calibration import *
-from ._aare import apply_calibration
+from ._aare import apply_calibration, count_switching_pixels
+from ._aare import calculate_pedestal, calculate_pedestal_float, calculate_pedestal_g0, calculate_pedestal_g0_float
 
 from ._aare import VarClusterFinder
diff --git a/python/src/bind_calibration.hpp b/python/src/bind_calibration.hpp
index df4b207..d2e43f1 100644
--- a/python/src/bind_calibration.hpp
+++ b/python/src/bind_calibration.hpp
@@ -17,19 +17,60 @@ py::array_t<DataType> pybind_apply_calibration(
         calibration,
     int n_threads = 4) {
 
-    auto data_span = make_view_3d(data);
-    auto ped = make_view_3d(pedestal);
-    auto cal = make_view_3d(calibration);
-
+    auto data_span = make_view_3d(data); // data is always 3D
     /* No pointer is passed, so NumPy will allocate the buffer */
     auto result = py::array_t<DataType>(data_span.shape());
     auto res = make_view_3d(result);
-
-    aare::apply_calibration<DataType>(res, data_span, ped, cal, n_threads);
-
+    if (data.ndim() == 3 && pedestal.ndim() == 3 && calibration.ndim() == 3) {
+        auto ped = make_view_3d(pedestal);
+        auto cal = make_view_3d(calibration);
+        aare::apply_calibration<DataType, 3>(res, data_span, ped, cal,
+                                             n_threads);
+    } else if (data.ndim() == 3 && pedestal.ndim() == 2 &&
+               calibration.ndim() == 2) {
+        auto ped = make_view_2d(pedestal);
+        auto cal = make_view_2d(calibration);
+        aare::apply_calibration<DataType, 2>(res, data_span, ped, cal,
+                                             n_threads);
+    } else {
+        throw std::runtime_error(
+            "Invalid number of dimensions for data, pedestal or calibration");
+    }
     return result;
 }
 
+py::array_t<int> pybind_count_switching_pixels(
+    py::array_t<uint16_t, py::array::c_style | py::array::forcecast> data,
+    ssize_t n_threads = 4) {
+
+    auto data_span = make_view_3d(data);
+    auto arr = new NDArray<int, 2>{};
+    *arr = aare::count_switching_pixels(data_span, n_threads);
+    return return_image_data(arr);
+}
+
+template <typename T>
+py::array_t<T> pybind_calculate_pedestal(
+    py::array_t<uint16_t, py::array::c_style | py::array::forcecast> data,
+    ssize_t n_threads) {
+
+    auto data_span = make_view_3d(data);
+    auto arr = new NDArray<T, 3>{};
+    *arr = aare::calculate_pedestal<T>(data_span, n_threads);
+    return return_image_data(arr);
+}
+
+template <typename T>
+py::array_t<T> pybind_calculate_pedestal_g0(
+    py::array_t<uint16_t, py::array::c_style | py::array::forcecast> data,
+    ssize_t n_threads) {
+
+    auto data_span = make_view_3d(data);
+    auto arr = new NDArray<T, 2>{};
+    *arr = aare::calculate_pedestal_g0<T>(data_span, n_threads);
+    return return_image_data(arr);
+}
+
 void bind_calibration(py::module &m) {
     m.def("apply_calibration", &pybind_apply_calibration<float>,
           py::arg("raw_data").noconvert(), py::kw_only(),
@@ -40,4 +81,20 @@ void bind_calibration(py::module &m) {
           py::arg("raw_data").noconvert(), py::kw_only(),
           py::arg("pd").noconvert(), py::arg("cal").noconvert(),
           py::arg("n_threads") = 4);
+
+    m.def("count_switching_pixels", &pybind_count_switching_pixels,
+          py::arg("raw_data").noconvert(), py::kw_only(),
+          py::arg("n_threads") = 4);
+
+    m.def("calculate_pedestal_float", &pybind_calculate_pedestal<float>,
+          py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
+
+    m.def("calculate_pedestal", &pybind_calculate_pedestal<double>,
+          py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
+
+    m.def("calculate_pedestal_g0", &pybind_calculate_pedestal_g0<double>,
+          py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
+
+    m.def("calculate_pedestal_g0_float", &pybind_calculate_pedestal_g0<float>,
+          py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
 }
\ No newline at end of file
diff --git a/src/NDArray.test.cpp b/src/NDArray.test.cpp
index 91b5933..9af8379 100644
--- a/src/NDArray.test.cpp
+++ b/src/NDArray.test.cpp
@@ -25,13 +25,13 @@ TEST_CASE("Construct from an NDView") {
     REQUIRE(image.data() != view.data());
 
     for (uint32_t i = 0; i < image.size(); ++i) {
-        REQUIRE(image(i) == view(i));
+        REQUIRE(image[i] == view[i]);
     }
 
     // Changing the image doesn't change the view
     image = 43;
     for (uint32_t i = 0; i < image.size(); ++i) {
-        REQUIRE(image(i) != view(i));
+        REQUIRE(image[i] != view[i]);
     }
 }
 
diff --git a/src/calibration.cpp b/src/calibration.cpp
new file mode 100644
index 0000000..c5578e3
--- /dev/null
+++ b/src/calibration.cpp
@@ -0,0 +1,78 @@
+#include "aare/calibration.hpp"
+
+namespace aare {
+
+std::pair<NDArray<size_t, 3>, NDArray<size_t,3>> sum_and_count_per_gain(NDView<uint16_t,3> raw_data){
+    NDArray<size_t, 3> accumulator(std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+    NDArray<size_t, 3> count(std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
+        for (int row = 0; row != raw_data.shape(1); ++row) {
+            for (int col = 0; col != raw_data.shape(2); ++col) {
+                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
+                accumulator(gain, row, col) += value;
+                count(gain, row, col) += 1;
+            }
+        }
+    }
+
+    return {std::move(accumulator), std::move(count)};
+}
+
+std::pair<NDArray<size_t, 2>, NDArray<size_t,2>> sum_and_count_g0(NDView<uint16_t, 3> raw_data){
+    NDArray<size_t, 2> accumulator(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    NDArray<size_t, 2> count(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
+    for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
+        for (int row = 0; row != raw_data.shape(1); ++row) {
+            for (int col = 0; col != raw_data.shape(2); ++col) {
+                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
+                if (gain != 0)
+                    continue; // we only care about gain 0
+                accumulator(row, col) += value;
+                count(row, col) += 1;
+            }
+        }
+    }
+
+    return {std::move(accumulator), std::move(count)};
+}
+
+NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data) {
+    NDArray<int, 2> switched(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)},0);
+    for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
+        for (int row = 0; row != raw_data.shape(1); ++row) {
+            for (int col = 0; col != raw_data.shape(2); ++col) {
+                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
+                if (gain != 0) {
+                    switched(row, col) += 1;
+                }
+            }
+        }
+    }
+    return switched;
+}
+
+
+NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data, ssize_t n_threads){
+    NDArray<int, 2> switched(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)},0);
+    std::vector<std::future<NDArray<int, 2>>> futures;
+    futures.reserve(n_threads);
+    auto limits = split_task(0, raw_data.shape(0), n_threads);
+
+    // make subviews for each thread
+    std::vector<NDView<uint16_t, 3>> subviews;
+    for (const auto &lim : limits) {
+        subviews.emplace_back(raw_data.data() + lim.first * raw_data.strides()[0],
+                              std::array<ssize_t, 3>{lim.second-lim.first, raw_data.shape(1), raw_data.shape(2)});
+    }
+
+    for (auto view : subviews) {
+        futures.push_back(std::async(static_cast<NDArray<int, 2>(*)(NDView<uint16_t, 3>)>(&count_switching_pixels), view));
+    }
+    
+    for (auto &f : futures) {
+        switched += f.get();
+    }
+    return switched;
+}
+
+}// namespace aare
\ No newline at end of file

From 5de402f91bc4d489d33866323dba1c27c7e1dcb7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Fr=C3=B6jdh?= <erik.frojdh@gmail.com>
Date: Wed, 23 Jul 2025 11:05:44 +0200
Subject: [PATCH 02/13] added docs

---
 RELEASE.md                      |  9 +++++
 docs/src/pycalibration.rst      | 16 ++++++++
 python/src/bind_calibration.hpp | 67 +++++++++++++++++++++++++++++----
 3 files changed, 85 insertions(+), 7 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index dc10406..a0c0704 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,6 +1,15 @@
 # Release notes
 
 
+### head
+
+Features:
+
+- Apply calibration works in G0 if passes a 2D calibration and pedestal
+- count pixels that switch
+- calculate pedestal (also g0 version)
+
+
 ### 2025.07.18
 
 Features:
diff --git a/docs/src/pycalibration.rst b/docs/src/pycalibration.rst
index 0ea3fba..6083162 100644
--- a/docs/src/pycalibration.rst
+++ b/docs/src/pycalibration.rst
@@ -17,8 +17,24 @@ Functions for applying calibration to data.
     # Apply calibration to raw data to convert from raw ADC values to keV
     data = aare.apply_calibration(raw_data, pd=pedestal, cal=calibration)
 
+    # If you pass a 2D pedestal and calibration only G0 will be used for the conversion
+    # Pixels that switched to G1 or G2 will be set to 0
+    data = aare.apply_calibration(raw_data, pd=pedestal[0], cal=calibration[0])
+
+    
+
 .. py:currentmodule:: aare
 
 .. autofunction:: apply_calibration
 
 .. autofunction:: load_calibration
+
+.. autofunction:: calculate_pedestal
+
+.. autofunction:: calculate_pedestal_float
+
+.. autofunction:: calculate_pedestal_g0
+
+.. autofunction:: calculate_pedestal_g0_float
+
+.. autofunction:: count_switching_pixels
diff --git a/python/src/bind_calibration.hpp b/python/src/bind_calibration.hpp
index d2e43f1..b95fa34 100644
--- a/python/src/bind_calibration.hpp
+++ b/python/src/bind_calibration.hpp
@@ -72,29 +72,82 @@ py::array_t<T> pybind_calculate_pedestal_g0(
 }
 
 void bind_calibration(py::module &m) {
-    m.def("apply_calibration", &pybind_apply_calibration<float>,
-          py::arg("raw_data").noconvert(), py::kw_only(),
-          py::arg("pd").noconvert(), py::arg("cal").noconvert(),
-          py::arg("n_threads") = 4);
-
     m.def("apply_calibration", &pybind_apply_calibration<double>,
           py::arg("raw_data").noconvert(), py::kw_only(),
           py::arg("pd").noconvert(), py::arg("cal").noconvert(),
           py::arg("n_threads") = 4);
 
+    m.def("apply_calibration", &pybind_apply_calibration<float>,
+          py::arg("raw_data").noconvert(), py::kw_only(),
+          py::arg("pd").noconvert(), py::arg("cal").noconvert(),
+          py::arg("n_threads") = 4);
+
     m.def("count_switching_pixels", &pybind_count_switching_pixels,
+          R"(
+        Count the number of time each pixel switches to G1 or G2.
+
+        Parameters
+        ----------
+        raw_data : array_like
+            3D array of shape (frames, rows, cols) to count the switching pixels from.
+        n_threads : int 
+            The number of threads to use for the calculation.
+        )",
           py::arg("raw_data").noconvert(), py::kw_only(),
           py::arg("n_threads") = 4);
 
-    m.def("calculate_pedestal_float", &pybind_calculate_pedestal<float>,
+    m.def("calculate_pedestal", &pybind_calculate_pedestal<double>,
+          R"(
+        Calculate the pedestal for all three gains and return the result as a 3D array of doubles.
+
+        Parameters
+        ----------
+        raw_data : array_like
+            3D array of shape (frames, rows, cols) to calculate the pedestal from.
+            Needs to contain data for all three gains (G0, G1, G2).
+        n_threads : int 
+            The number of threads to use for the calculation.
+        )",
           py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
 
-    m.def("calculate_pedestal", &pybind_calculate_pedestal<double>,
+    m.def("calculate_pedestal_float", &pybind_calculate_pedestal<float>,
+          R"(
+        Same as `calculate_pedestal` but returns a 3D array of floats.
+
+        Parameters
+        ----------
+        raw_data : array_like
+            3D array of shape (frames, rows, cols) to calculate the pedestal from.
+            Needs to contain data for all three gains (G0, G1, G2).
+        n_threads : int 
+            The number of threads to use for the calculation.
+        )",
           py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
 
     m.def("calculate_pedestal_g0", &pybind_calculate_pedestal_g0<double>,
+          R"(
+        Calculate the pedestal for G0 and return the result as a 2D array of doubles.
+        Pixels in G1 and G2 are ignored.
+
+        Parameters
+        ----------
+        raw_data : array_like
+            3D array of shape (frames, rows, cols) to calculate the pedestal from.
+        n_threads : int 
+            The number of threads to use for the calculation.
+        )",
           py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
 
     m.def("calculate_pedestal_g0_float", &pybind_calculate_pedestal_g0<float>,
+          R"(
+        Same as `calculate_pedestal_g0` but returns a 2D array of floats.
+
+        Parameters
+        ----------
+        raw_data : array_like
+            3D array of shape (frames, rows, cols) to calculate the pedestal from.
+        n_threads : int 
+            The number of threads to use for the calculation.
+        )",
           py::arg("raw_data").noconvert(), py::arg("n_threads") = 4);
 }
\ No newline at end of file

From cb439efb4839e13a8da76b088145c288646cba08 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Fr=C3=B6jdh?= <erik.frojdh@gmail.com>
Date: Wed, 23 Jul 2025 11:34:47 +0200
Subject: [PATCH 03/13] added tests

---
 python/tests/test_calibration.py | 96 +++++++++++++++++++++++++++++++-
 1 file changed, 94 insertions(+), 2 deletions(-)

diff --git a/python/tests/test_calibration.py b/python/tests/test_calibration.py
index cce4344..bbd980c 100644
--- a/python/tests/test_calibration.py
+++ b/python/tests/test_calibration.py
@@ -1,6 +1,7 @@
 import pytest 
 import numpy as np
-from aare import apply_calibration
+
+import aare
 
 def test_apply_calibration_small_data():
     # The raw data consists of 10 4x5 images
@@ -27,7 +28,7 @@ def test_apply_calibration_small_data():
 
     
 
-    data = apply_calibration(raw, pd = pedestal, cal = calibration)
+    data = aare.apply_calibration(raw, pd = pedestal, cal = calibration)
 
 
     # The formula that is applied is:
@@ -41,3 +42,94 @@ def test_apply_calibration_small_data():
     assert data[2,2,2] == 0
     assert data[0,1,1] == 0
     assert data[1,3,0] == 0
+
+
+@pytest.fixture
+def raw_data_3x2x2():
+    raw = np.zeros((3, 2, 2), dtype=np.uint16)
+    raw[0, 0, 0] = 100
+    raw[1,0, 0] =  200
+    raw[2, 0, 0] = 300
+
+    raw[0, 0, 1] = (1<<14) + 100
+    raw[1, 0, 1] = (1<<14) + 200
+    raw[2, 0, 1] = (1<<14) + 300
+
+    raw[0, 1, 0] = (1<<14) + 37
+    raw[1, 1, 0] =  38
+    raw[2, 1, 0] = (3<<14) + 39
+
+    raw[0, 1, 1] = (3<<14) + 100
+    raw[1, 1, 1] = (3<<14) + 200
+    raw[2, 1, 1] = (3<<14) + 300
+    return raw
+
+def test_calculate_pedestal(raw_data_3x2x2):
+    # Calculate the pedestal
+    pd = aare.calculate_pedestal(raw_data_3x2x2)
+    assert pd.shape == (3, 2, 2)
+    assert pd.dtype == np.float64
+    assert pd[0, 0, 0] == 200
+    assert pd[1, 0, 0] == 0
+    assert pd[2, 0, 0] == 0
+
+    assert pd[0, 0, 1] == 0
+    assert pd[1, 0, 1] == 200
+    assert pd[2, 0, 1] == 0
+
+    assert pd[0, 1, 0] == 38
+    assert pd[1, 1, 0] == 37
+    assert pd[2, 1, 0] == 39
+
+    assert pd[0, 1, 1] == 0
+    assert pd[1, 1, 1] == 0
+    assert pd[2, 1, 1] == 200
+
+def test_calculate_pedestal_float(raw_data_3x2x2):
+    #results should be the same for float
+    pd2 = aare.calculate_pedestal_float(raw_data_3x2x2)
+    assert pd2.shape == (3, 2, 2)
+    assert pd2.dtype == np.float32
+    assert pd2[0, 0, 0] == 200
+    assert pd2[1, 0, 0] == 0
+    assert pd2[2, 0, 0] == 0
+
+    assert pd2[0, 0, 1] == 0
+    assert pd2[1, 0, 1] == 200
+    assert pd2[2, 0, 1] == 0
+
+    assert pd2[0, 1, 0] == 38
+    assert pd2[1, 1, 0] == 37
+    assert pd2[2, 1, 0] == 39
+
+    assert pd2[0, 1, 1] == 0
+    assert pd2[1, 1, 1] == 0
+    assert pd2[2, 1, 1] == 200
+
+def test_calculate_pedestal_g0(raw_data_3x2x2):
+    pd = aare.calculate_pedestal_g0(raw_data_3x2x2)
+    assert pd.shape == (2, 2)
+    assert pd.dtype == np.float64
+    assert pd[0, 0] == 200
+    assert pd[1, 0] == 38
+    assert pd[0, 1] == 0
+    assert pd[1, 1] == 0
+
+def test_calculate_pedestal_g0_float(raw_data_3x2x2):
+    pd = aare.calculate_pedestal_g0_float(raw_data_3x2x2)
+    assert pd.shape == (2, 2)
+    assert pd.dtype == np.float32
+    assert pd[0, 0] == 200
+    assert pd[1, 0] == 38
+    assert pd[0, 1] == 0
+    assert pd[1, 1] == 0
+
+def test_count_switching_pixels(raw_data_3x2x2):
+    # Count the number of pixels that switched gain
+    count = aare.count_switching_pixels(raw_data_3x2x2)
+    assert count.shape == (2, 2)
+    assert count.sum() == 8
+    assert count[0, 0] == 0 
+    assert count[1, 0] == 2 
+    assert count[0, 1] == 3 
+    assert count[1, 1] == 3
\ No newline at end of file

From 0fea0f5b0e403bea2bbb98c6144bb09b72dec9d8 Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Thu, 24 Jul 2025 09:40:38 +0200
Subject: [PATCH 04/13] added safe_divide to NDArray and used it for pedestal

---
 include/aare/NDArray.hpp     | 27 +++++++++++++++++++++------
 include/aare/calibration.hpp | 36 ++++++------------------------------
 2 files changed, 27 insertions(+), 36 deletions(-)

diff --git a/include/aare/NDArray.hpp b/include/aare/NDArray.hpp
index 1a501eb..f73bfd6 100644
--- a/include/aare/NDArray.hpp
+++ b/include/aare/NDArray.hpp
@@ -217,6 +217,7 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
         std::fill(shape_.begin(), shape_.end(), 0);
         std::fill(strides_.begin(), strides_.end(), 0);
     }
+
 };
 
 // Move assign
@@ -380,12 +381,8 @@ NDArray<T, Ndim> NDArray<T, Ndim>::operator*(const T &value) {
     result *= value;
     return result;
 }
-// template <typename T, ssize_t Ndim> void NDArray<T, Ndim>::Print() {
-//     if (shape_[0] < 20 && shape_[1] < 20)
-//         Print_all();
-//     else
-//         Print_some();
-// }
+
+
 
 template <typename T, ssize_t Ndim>
 std::ostream &operator<<(std::ostream &os, const NDArray<T, Ndim> &arr) {
@@ -437,4 +434,22 @@ NDArray<T, Ndim> load(const std::string &pathname,
     return img;
 }
 
+
+template <typename RT, typename NT, typename DT, ssize_t Ndim>
+NDArray<RT, Ndim> safe_divide(const NDArray<NT, Ndim> &numerator,
+                                 const NDArray<DT, Ndim> &denominator) {
+    if (numerator.shape() != denominator.shape()) {
+        throw std::runtime_error("Shapes of numerator and denominator must match");
+    }
+    NDArray<RT, Ndim> result(numerator.shape());
+    for (ssize_t i = 0; i < numerator.size(); ++i) {
+        if (denominator[i] != 0) {
+            result[i] = static_cast<RT>(numerator[i]) / static_cast<RT>(denominator[i]);
+        } else {
+            result[i] = RT{0}; // or handle division by zero as needed
+        }
+    }
+    return result;
+}
+
 } // namespace aare
\ No newline at end of file
diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index 8d672b4..d61f706 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -161,6 +161,9 @@ NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data) {
     }
     return pedestal;
 }
+
+
+
 template <typename T>
 NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
                                  ssize_t n_threads) {
@@ -191,20 +194,8 @@ NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
         accumulator += acc;
         count += cnt;
     }
-    NDArray<T, 2> pedestal(
-        std::array<ssize_t,2>{raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (int gain = 0; gain < 3; ++gain) {
-        for (int row = 0; row < raw_data.shape(1); ++row) {
-            for (int col = 0; col < raw_data.shape(2); ++col) {
-                if (count(row, col) != 0) {
-                    pedestal(row, col) =
-                        static_cast<T>(accumulator(row, col)) /
-                        static_cast<T>(count(row, col));
-                }
-            }
-        }
-    }
-    return pedestal;
+
+    return safe_divide<T>(accumulator, count);
 
 }
 
@@ -212,8 +203,6 @@ NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
 template <typename T>
 NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data,
                                  ssize_t n_threads) {
-    NDArray<int, 2> switched(
-        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
     std::vector<std::future<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>>>
         futures;
     futures.reserve(n_threads);
@@ -244,20 +233,7 @@ NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data,
         count += cnt;
     }
 
-    NDArray<T, 3> pedestal(
-        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (int gain = 0; gain < 3; ++gain) {
-        for (int row = 0; row < raw_data.shape(1); ++row) {
-            for (int col = 0; col < raw_data.shape(2); ++col) {
-                if (count(gain, row, col) != 0) {
-                    pedestal(gain, row, col) =
-                        static_cast<T>(accumulator(gain, row, col)) /
-                        static_cast<T>(count(gain, row, col));
-                }
-            }
-        }
-    }
-    return pedestal;
+    return safe_divide<T>(accumulator, count);
 }
 
 /**

From 348fd0f937da19cadcc27be531813d1e79cb5c0b Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Thu, 24 Jul 2025 10:14:29 +0200
Subject: [PATCH 05/13] removed unused code

---
 include/aare/calibration.hpp | 39 ------------------------------------
 1 file changed, 39 deletions(-)

diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index d61f706..ce69e95 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -124,45 +124,6 @@ sum_and_count_per_gain(NDView<uint16_t, 3> raw_data);
 std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>
 sum_and_count_g0(NDView<uint16_t, 3> raw_data);
 
-template <typename T>
-NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data) {
-    auto [accumulator, count] = sum_and_count_per_gain(raw_data);
-
-    NDArray<T, 3> pedestal(
-        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (int gain = 0; gain < 3; ++gain) {
-        for (int row = 0; row < raw_data.shape(1); ++row) {
-            for (int col = 0; col < raw_data.shape(2); ++col) {
-                if (count(gain, row, col) != 0) {
-                    pedestal(gain, row, col) =
-                        static_cast<T>(accumulator(gain, row, col)) /
-                        static_cast<T>(count(gain, row, col));
-                }
-            }
-        }
-    }
-    return pedestal;
-}
-
-template <typename T>
-NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data) {
-    auto [accumulator, count] = sum_and_count_g0(raw_data);
-
-    NDArray<T, 2> pedestal(
-        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (int row = 0; row < raw_data.shape(1); ++row) {
-        for (int col = 0; col < raw_data.shape(2); ++col) {
-            if (count(row, col) != 0) {
-                pedestal(row, col) =
-                    static_cast<T>(accumulator(row, col)) /
-                    static_cast<T>(count(row, col));
-            }
-        }
-    }
-    return pedestal;
-}
-
-
 
 template <typename T>
 NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,

From 46876bfa731b954df7faf409572117b9e0f01b7e Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Thu, 24 Jul 2025 10:57:02 +0200
Subject: [PATCH 06/13] reduced duplicate code

---
 include/aare/NDView.hpp      | 16 ++++++++++++++++
 include/aare/calibration.hpp | 33 ++++++++++-----------------------
 include/aare/utils/par.hpp   | 16 ++++++++++++++++
 include/aare/utils/task.hpp  |  2 +-
 4 files changed, 43 insertions(+), 24 deletions(-)

diff --git a/include/aare/NDView.hpp b/include/aare/NDView.hpp
index 4c16742..56bca2f 100644
--- a/include/aare/NDView.hpp
+++ b/include/aare/NDView.hpp
@@ -178,6 +178,22 @@ class NDView : public ArrayExpr<NDView<T, Ndim>, Ndim> {
     const T *data() const { return buffer_; }
     void print_all() const;
 
+    /**
+     * @brief Create a subview of a range of the first dimension. 
+     * This is useful for splitting a batches of frames in parallel processing.
+     * @param first The first index of the subview (inclusive).
+     * @param last The last index of the subview (exclusive).
+     * @return A new NDView that is a subview of the current view.
+     * @throws std::runtime_error if the range is invalid.
+     */
+    NDView sub_view(ssize_t first, ssize_t last) const {
+        if (first < 0 || last > shape_[0] || first >= last)
+            throw std::runtime_error(LOCATION + "Invalid sub_view range");
+        auto new_shape = shape_;
+        new_shape[0] = last - first;
+        return NDView(buffer_ + first * strides_[0], new_shape);
+    }
+
   private:
     T *buffer_{nullptr};
     std::array<ssize_t, Ndim> strides_{};
diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index ce69e95..0677127 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -3,6 +3,7 @@
 #include "aare/NDArray.hpp"
 #include "aare/NDView.hpp"
 #include "aare/defs.hpp"
+#include "aare/utils/par.hpp"
 #include "aare/utils/task.hpp"
 #include <cstdint>
 #include <future>
@@ -111,35 +112,30 @@ void apply_calibration(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
     for (const auto &lim : limits)
         futures.push_back(std::async(
             static_cast<void (*)(NDView<T, 3>, NDView<uint16_t, 3>,
-                                 NDView<T, Ndim>, NDView<T, Ndim>, int, int)>(
+                                 NDView<T, Ndim>, NDView<T, Ndim>, int,
+                                 int)>(
                 apply_calibration_impl),
             res, raw_data, ped, cal, lim.first, lim.second));
     for (auto &f : futures)
         f.get();
 }
 
+
 std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>
 sum_and_count_per_gain(NDView<uint16_t, 3> raw_data);
 
 std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>
 sum_and_count_g0(NDView<uint16_t, 3> raw_data);
 
-
 template <typename T>
 NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
-                                 ssize_t n_threads) {
+                                    ssize_t n_threads) {
     std::vector<std::future<std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>>>
         futures;
     futures.reserve(n_threads);
-    auto limits = split_task(0, raw_data.shape(0), n_threads);  
-    // make subviews for each thread
-    std::vector<NDView<uint16_t, 3>> subviews;
-    for (const auto &lim : limits) {
-        subviews.emplace_back(
-            raw_data.data() + lim.first * raw_data.strides()[0],
-            std::array<ssize_t, 3>{lim.second - lim.first, raw_data.shape(1),
-                                   raw_data.shape(2)});
-    }
+
+    auto subviews = make_subviews(raw_data, n_threads);
+
     for (auto view : subviews) {
         futures.push_back(std::async(
             static_cast<std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>> (*)(
@@ -157,26 +153,17 @@ NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
     }
 
     return safe_divide<T>(accumulator, count);
-
 }
 
-
 template <typename T>
 NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data,
                                  ssize_t n_threads) {
     std::vector<std::future<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>>>
         futures;
     futures.reserve(n_threads);
-    auto limits = split_task(0, raw_data.shape(0), n_threads);
 
-    // make subviews for each thread
-    std::vector<NDView<uint16_t, 3>> subviews;
-    for (const auto &lim : limits) {
-        subviews.emplace_back(
-            raw_data.data() + lim.first * raw_data.strides()[0],
-            std::array<ssize_t, 3>{lim.second - lim.first, raw_data.shape(1),
-                                   raw_data.shape(2)});
-    }
+    auto subviews = make_subviews(raw_data, n_threads);
+
     for (auto view : subviews) {
         futures.push_back(std::async(
             static_cast<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>> (*)(
diff --git a/include/aare/utils/par.hpp b/include/aare/utils/par.hpp
index e52c897..55f60f4 100644
--- a/include/aare/utils/par.hpp
+++ b/include/aare/utils/par.hpp
@@ -1,7 +1,10 @@
+#pragma once
 #include <thread>
 #include <utility>
 #include <vector>
 
+#include "aare/utils/task.hpp"
+
 namespace aare {
 
 template <typename F>
@@ -15,4 +18,17 @@ void RunInParallel(F func, const std::vector<std::pair<int, int>> &tasks) {
         thread.join();
     }
 }
+
+
+template <typename T>
+std::vector<NDView<T,3>> make_subviews(NDView<T, 3> &data, ssize_t n_threads) {
+    std::vector<NDView<T, 3>> subviews;
+    subviews.reserve(n_threads);
+    auto limits = split_task(0, data.shape(0), n_threads);
+    for (const auto &lim : limits) {
+        subviews.push_back(data.sub_view(lim.first, lim.second));
+    }
+    return subviews;
+}
+
 } // namespace aare
\ No newline at end of file
diff --git a/include/aare/utils/task.hpp b/include/aare/utils/task.hpp
index a6ee142..63fe691 100644
--- a/include/aare/utils/task.hpp
+++ b/include/aare/utils/task.hpp
@@ -1,4 +1,4 @@
-
+#pragma once
 #include <utility>
 #include <vector>
 

From b8e91d0282419d5eed0d30b050d5b7c030b091da Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Thu, 24 Jul 2025 12:10:13 +0200
Subject: [PATCH 07/13] zero out switching pixels if 2D calibration is used

---
 include/aare/calibration.hpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index 0677127..bf543aa 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -92,11 +92,12 @@ void apply_calibration_impl(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
                 // ADU/keV is the standard unit for the calibration which
                 // means rewriting the formula is not worth it.
 
-                // ignore anything that is not gain 0
-                // TODO! deal with fixed gain?
+                // Set the value to 0 if the gain is not 0
                 if (gain == 0)
                     res(frame_nr, row, col) =
                         (value - ped(row, col)) / cal(row, col);
+                else
+                    res(frame_nr, row, col) = 0;
             }
         }
     }

From 8c4d8b687eeee37d24bf0359298f825bdacc3ea5 Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Thu, 24 Jul 2025 12:16:08 +0200
Subject: [PATCH 08/13] using make_subview

---
 src/calibration.cpp | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/src/calibration.cpp b/src/calibration.cpp
index c5578e3..eeaae84 100644
--- a/src/calibration.cpp
+++ b/src/calibration.cpp
@@ -56,14 +56,8 @@ NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data, ssize_t n_t
     NDArray<int, 2> switched(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)},0);
     std::vector<std::future<NDArray<int, 2>>> futures;
     futures.reserve(n_threads);
-    auto limits = split_task(0, raw_data.shape(0), n_threads);
 
-    // make subviews for each thread
-    std::vector<NDView<uint16_t, 3>> subviews;
-    for (const auto &lim : limits) {
-        subviews.emplace_back(raw_data.data() + lim.first * raw_data.strides()[0],
-                              std::array<ssize_t, 3>{lim.second-lim.first, raw_data.shape(1), raw_data.shape(2)});
-    }
+    auto subviews = make_subviews(raw_data, n_threads);
 
     for (auto view : subviews) {
         futures.push_back(std::async(static_cast<NDArray<int, 2>(*)(NDView<uint16_t, 3>)>(&count_switching_pixels), view));

From 13471582354bdf235bd65188f00c6f9cf19e0f67 Mon Sep 17 00:00:00 2001
From: Alice <alice.mazzoleni@psi.ch>
Date: Thu, 24 Jul 2025 15:40:05 +0200
Subject: [PATCH 09/13] templated calculate_pedestal with boolean template
 argument only_gain0, added drop_dimension to NDArray and reference pointer to
 data

---
 CMakeLists.txt                  |  1 +
 include/aare/NDArray.hpp        | 33 +++++++++----
 include/aare/calibration.hpp    | 83 +++++++++++++++++----------------
 python/src/bind_calibration.hpp |  4 +-
 src/calibration.cpp             | 56 ++++++----------------
 5 files changed, 85 insertions(+), 92 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3cfc6a7..6338f93 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -438,6 +438,7 @@ endif()
 if(AARE_TESTS)
     set(TestSources
             ${CMAKE_CURRENT_SOURCE_DIR}/src/algorithm.test.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/src/calibration.test.cpp
             ${CMAKE_CURRENT_SOURCE_DIR}/src/defs.test.cpp
             ${CMAKE_CURRENT_SOURCE_DIR}/src/decode.test.cpp
             ${CMAKE_CURRENT_SOURCE_DIR}/src/Dtype.test.cpp
diff --git a/include/aare/NDArray.hpp b/include/aare/NDArray.hpp
index f73bfd6..c3edbfb 100644
--- a/include/aare/NDArray.hpp
+++ b/include/aare/NDArray.hpp
@@ -33,7 +33,7 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
      * @brief Default constructor. Will construct an empty NDArray.
      *
      */
-    NDArray() : shape_(), strides_(c_strides<Ndim>(shape_)), data_(nullptr){};
+    NDArray() : shape_(), strides_(c_strides<Ndim>(shape_)), data_(nullptr) {};
 
     /**
      * @brief Construct a new NDArray object with a given shape.
@@ -185,6 +185,7 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
     const T &operator[](ssize_t i) const { return data_[i]; }
 
     T *data() { return data_; }
+    T *&data_ref() { return data_; }
     std::byte *buffer() { return reinterpret_cast<std::byte *>(data_); }
     ssize_t size() const { return static_cast<ssize_t>(size_); }
     size_t total_bytes() const { return size_ * sizeof(T); }
@@ -211,13 +212,30 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
     void Print_all();
     void Print_some();
 
+    template <ssize_t M = Ndim, typename = std::enable_if_t<(M > 2)>>
+    NDArray<T, Ndim - 1> drop_dimension() && {
+
+        std::array<ssize_t, Ndim - 1> new_shape;
+
+        std::copy(shape_.begin() + 1, shape_.begin() + Ndim, new_shape.begin());
+
+        NDArray<T, Ndim - 1> new_array(new_shape);
+
+        delete new_array.data();
+
+        new_array.data_ref() = data_;
+
+        this->reset();
+
+        return new_array;
+    }
+
     void reset() {
         data_ = nullptr;
         size_ = 0;
         std::fill(shape_.begin(), shape_.end(), 0);
         std::fill(strides_.begin(), strides_.end(), 0);
     }
-
 };
 
 // Move assign
@@ -382,8 +400,6 @@ NDArray<T, Ndim> NDArray<T, Ndim>::operator*(const T &value) {
     return result;
 }
 
-
-
 template <typename T, ssize_t Ndim>
 std::ostream &operator<<(std::ostream &os, const NDArray<T, Ndim> &arr) {
     for (auto row = 0; row < arr.shape(0); ++row) {
@@ -434,17 +450,18 @@ NDArray<T, Ndim> load(const std::string &pathname,
     return img;
 }
 
-
 template <typename RT, typename NT, typename DT, ssize_t Ndim>
 NDArray<RT, Ndim> safe_divide(const NDArray<NT, Ndim> &numerator,
-                                 const NDArray<DT, Ndim> &denominator) {
+                              const NDArray<DT, Ndim> &denominator) {
     if (numerator.shape() != denominator.shape()) {
-        throw std::runtime_error("Shapes of numerator and denominator must match");
+        throw std::runtime_error(
+            "Shapes of numerator and denominator must match");
     }
     NDArray<RT, Ndim> result(numerator.shape());
     for (ssize_t i = 0; i < numerator.size(); ++i) {
         if (denominator[i] != 0) {
-            result[i] = static_cast<RT>(numerator[i]) / static_cast<RT>(denominator[i]);
+            result[i] =
+                static_cast<RT>(numerator[i]) / static_cast<RT>(denominator[i]);
         } else {
             result[i] = RT{0}; // or handle division by zero as needed
         }
diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index bf543aa..a94705f 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -103,7 +103,7 @@ void apply_calibration_impl(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
     }
 }
 
-template <class T, ssize_t Ndim=3>
+template <class T, ssize_t Ndim = 3>
 void apply_calibration(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
                        NDView<T, Ndim> ped, NDView<T, Ndim> cal,
                        ssize_t n_threads = 4) {
@@ -113,52 +113,44 @@ void apply_calibration(NDView<T, 3> res, NDView<uint16_t, 3> raw_data,
     for (const auto &lim : limits)
         futures.push_back(std::async(
             static_cast<void (*)(NDView<T, 3>, NDView<uint16_t, 3>,
-                                 NDView<T, Ndim>, NDView<T, Ndim>, int,
-                                 int)>(
+                                 NDView<T, Ndim>, NDView<T, Ndim>, int, int)>(
                 apply_calibration_impl),
             res, raw_data, ped, cal, lim.first, lim.second));
     for (auto &f : futures)
         f.get();
 }
 
-
+template <bool only_gain0>
 std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>
-sum_and_count_per_gain(NDView<uint16_t, 3> raw_data);
-
-std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>
-sum_and_count_g0(NDView<uint16_t, 3> raw_data);
-
-template <typename T>
-NDArray<T, 2> calculate_pedestal_g0(NDView<uint16_t, 3> raw_data,
-                                    ssize_t n_threads) {
-    std::vector<std::future<std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>>>>
-        futures;
-    futures.reserve(n_threads);
-
-    auto subviews = make_subviews(raw_data, n_threads);
-
-    for (auto view : subviews) {
-        futures.push_back(std::async(
-            static_cast<std::pair<NDArray<size_t, 2>, NDArray<size_t, 2>> (*)(
-                NDView<uint16_t, 3>)>(&sum_and_count_g0),
-            view));
-    }
-    NDArray<size_t, 2> accumulator(
-        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
-    NDArray<size_t, 2> count(
-        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (auto &f : futures) {
-        auto [acc, cnt] = f.get();
-        accumulator += acc;
-        count += cnt;
+sum_and_count_per_gain(NDView<uint16_t, 3> raw_data) {
+    constexpr ssize_t num_gains = only_gain0 ? 1 : 3;
+    NDArray<size_t, 3> accumulator(
+        std::array<ssize_t, 3>{num_gains, raw_data.shape(1), raw_data.shape(2)},
+        0);
+    NDArray<size_t, 3> count(
+        std::array<ssize_t, 3>{num_gains, raw_data.shape(1), raw_data.shape(2)},
+        0);
+    for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
+        for (int row = 0; row != raw_data.shape(1); ++row) {
+            for (int col = 0; col != raw_data.shape(2); ++col) {
+                auto [value, gain] =
+                    get_value_and_gain(raw_data(frame_nr, row, col));
+                if (gain != 0 && only_gain0)
+                    continue;
+                accumulator(gain, row, col) += value;
+                count(gain, row, col) += 1;
+            }
+        }
     }
 
-    return safe_divide<T>(accumulator, count);
+    return {std::move(accumulator), std::move(count)};
 }
 
-template <typename T>
-NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data,
-                                 ssize_t n_threads) {
+template <typename T, bool only_gain0 = false>
+NDArray<T, 3 - static_cast<ssize_t>(only_gain0)>
+calculate_pedestal(NDView<uint16_t, 3> raw_data, ssize_t n_threads) {
+
+    constexpr ssize_t num_gains = only_gain0 ? 1 : 3;
     std::vector<std::future<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>>>>
         futures;
     futures.reserve(n_threads);
@@ -168,21 +160,27 @@ NDArray<T, 3> calculate_pedestal(NDView<uint16_t, 3> raw_data,
     for (auto view : subviews) {
         futures.push_back(std::async(
             static_cast<std::pair<NDArray<size_t, 3>, NDArray<size_t, 3>> (*)(
-                NDView<uint16_t, 3>)>(&sum_and_count_per_gain),
+                NDView<uint16_t, 3>)>(&sum_and_count_per_gain<only_gain0>),
             view));
     }
 
     NDArray<size_t, 3> accumulator(
-        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+        std::array<ssize_t, 3>{num_gains, raw_data.shape(1), raw_data.shape(2)},
+        0);
     NDArray<size_t, 3> count(
-        std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
+        std::array<ssize_t, 3>{num_gains, raw_data.shape(1), raw_data.shape(2)},
+        0);
     for (auto &f : futures) {
         auto [acc, cnt] = f.get();
         accumulator += acc;
         count += cnt;
     }
 
-    return safe_divide<T>(accumulator, count);
+    if constexpr (only_gain0) {
+        return safe_divide<T>(accumulator, count).drop_dimension();
+    } else {
+        return safe_divide<T>(accumulator, count);
+    }
 }
 
 /**
@@ -205,4 +203,9 @@ NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data);
 NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data,
                                        ssize_t n_threads);
 
+template <typename T>
+auto calculate_pedestal_g0(NDView<uint16_t, 3> raw_data, ssize_t n_threads) {
+    return calculate_pedestal<T, true>(raw_data, n_threads);
+}
+
 } // namespace aare
\ No newline at end of file
diff --git a/python/src/bind_calibration.hpp b/python/src/bind_calibration.hpp
index b95fa34..836a6de 100644
--- a/python/src/bind_calibration.hpp
+++ b/python/src/bind_calibration.hpp
@@ -56,7 +56,7 @@ py::array_t<T> pybind_calculate_pedestal(
 
     auto data_span = make_view_3d(data);
     auto arr = new NDArray<T, 3>{};
-    *arr = aare::calculate_pedestal<T>(data_span, n_threads);
+    *arr = aare::calculate_pedestal<T, false>(data_span, n_threads);
     return return_image_data(arr);
 }
 
@@ -67,7 +67,7 @@ py::array_t<T> pybind_calculate_pedestal_g0(
 
     auto data_span = make_view_3d(data);
     auto arr = new NDArray<T, 2>{};
-    *arr = aare::calculate_pedestal_g0<T>(data_span, n_threads);
+    *arr = aare::calculate_pedestal<T, true>(data_span, n_threads);
     return return_image_data(arr);
 }
 
diff --git a/src/calibration.cpp b/src/calibration.cpp
index eeaae84..0cb3b99 100644
--- a/src/calibration.cpp
+++ b/src/calibration.cpp
@@ -2,46 +2,14 @@
 
 namespace aare {
 
-std::pair<NDArray<size_t, 3>, NDArray<size_t,3>> sum_and_count_per_gain(NDView<uint16_t,3> raw_data){
-    NDArray<size_t, 3> accumulator(std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
-    NDArray<size_t, 3> count(std::array<ssize_t, 3>{3, raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
-        for (int row = 0; row != raw_data.shape(1); ++row) {
-            for (int col = 0; col != raw_data.shape(2); ++col) {
-                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
-                accumulator(gain, row, col) += value;
-                count(gain, row, col) += 1;
-            }
-        }
-    }
-
-    return {std::move(accumulator), std::move(count)};
-}
-
-std::pair<NDArray<size_t, 2>, NDArray<size_t,2>> sum_and_count_g0(NDView<uint16_t, 3> raw_data){
-    NDArray<size_t, 2> accumulator(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
-    NDArray<size_t, 2> count(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
-    for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
-        for (int row = 0; row != raw_data.shape(1); ++row) {
-            for (int col = 0; col != raw_data.shape(2); ++col) {
-                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
-                if (gain != 0)
-                    continue; // we only care about gain 0
-                accumulator(row, col) += value;
-                count(row, col) += 1;
-            }
-        }
-    }
-
-    return {std::move(accumulator), std::move(count)};
-}
-
 NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data) {
-    NDArray<int, 2> switched(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)},0);
+    NDArray<int, 2> switched(
+        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
     for (int frame_nr = 0; frame_nr != raw_data.shape(0); ++frame_nr) {
         for (int row = 0; row != raw_data.shape(1); ++row) {
             for (int col = 0; col != raw_data.shape(2); ++col) {
-                auto [value, gain] = get_value_and_gain(raw_data(frame_nr, row, col));
+                auto [value, gain] =
+                    get_value_and_gain(raw_data(frame_nr, row, col));
                 if (gain != 0) {
                     switched(row, col) += 1;
                 }
@@ -51,22 +19,26 @@ NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data) {
     return switched;
 }
 
-
-NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data, ssize_t n_threads){
-    NDArray<int, 2> switched(std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)},0);
+NDArray<int, 2> count_switching_pixels(NDView<uint16_t, 3> raw_data,
+                                       ssize_t n_threads) {
+    NDArray<int, 2> switched(
+        std::array<ssize_t, 2>{raw_data.shape(1), raw_data.shape(2)}, 0);
     std::vector<std::future<NDArray<int, 2>>> futures;
     futures.reserve(n_threads);
 
     auto subviews = make_subviews(raw_data, n_threads);
 
     for (auto view : subviews) {
-        futures.push_back(std::async(static_cast<NDArray<int, 2>(*)(NDView<uint16_t, 3>)>(&count_switching_pixels), view));
+        futures.push_back(
+            std::async(static_cast<NDArray<int, 2> (*)(NDView<uint16_t, 3>)>(
+                           &count_switching_pixels),
+                       view));
     }
-    
+
     for (auto &f : futures) {
         switched += f.get();
     }
     return switched;
 }
 
-}// namespace aare
\ No newline at end of file
+} // namespace aare
\ No newline at end of file

From 1195a5e100b355c9b6353e2bc0c5299547e6ab58 Mon Sep 17 00:00:00 2001
From: Alice <alice.mazzoleni@psi.ch>
Date: Fri, 25 Jul 2025 10:18:55 +0200
Subject: [PATCH 10/13] added drop dimension test, added file
 calibration.test.cpp

---
 src/NDArray.test.cpp     | 19 +++++++++++++++-
 src/calibration.test.cpp | 49 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 src/calibration.test.cpp

diff --git a/src/NDArray.test.cpp b/src/NDArray.test.cpp
index 9af8379..97a7f09 100644
--- a/src/NDArray.test.cpp
+++ b/src/NDArray.test.cpp
@@ -427,4 +427,21 @@ TEST_CASE("Construct an NDArray from an std::array") {
     for (uint32_t i = 0; i < a.size(); ++i) {
         REQUIRE(a(i) == b[i]);
     }
-}
\ No newline at end of file
+}
+
+TEST_CASE("Drop dimension") {
+    NDArray<int, 3> array(std::array<ssize_t, 3>{2, 2, 2});
+
+    std::fill(array.begin(), array.begin() + 4, 0);
+    std::fill(array.begin() + 4, array.end(), 1);
+
+    auto new_array = std::move(array).drop_dimension();
+
+    CHECK(new_array.shape() == std::array<ssize_t, 2>{2, 2});
+    CHECK(new_array.size() == 4);
+
+    std::for_each(new_array.begin(), new_array.end(),
+                  [](int &element) { CHECK(element == 0); });
+
+    CHECK(array.size() == 0); // array was moved
+}
diff --git a/src/calibration.test.cpp b/src/calibration.test.cpp
new file mode 100644
index 0000000..712dcd8
--- /dev/null
+++ b/src/calibration.test.cpp
@@ -0,0 +1,49 @@
+/************************************************
+ * @file test-Cluster.cpp
+ * @short test case for generic Cluster, ClusterVector, and calculate_eta2
+ ***********************************************/
+
+#include "aare/calibration.hpp"
+
+// #include "catch.hpp"
+#include <array>
+#include <catch2/catch_all.hpp>
+#include <catch2/catch_test_macros.hpp>
+
+using namespace aare;
+
+TEST_CASE("Test Pedestal Generation", "[.calibration]") {
+    NDArray<uint16_t, 3> raw(std::array<ssize_t, 3>{3, 2, 2}, 0);
+
+    // gain 0
+    raw(0, 0, 0) = 100;
+    raw(1, 0, 0) = 200;
+    raw(2, 0, 0) = 300;
+
+    // gain 1
+    raw(0, 0, 1) = (1 << 14) + 100;
+    raw(1, 0, 1) = (1 << 14) + 200;
+    raw(2, 0, 1) = (1 << 14) + 300;
+
+    raw(0, 1, 0) = (1 << 14) + 37;
+    raw(1, 1, 0) = 38;
+    raw(2, 1, 0) = (3 << 14) + 39;
+
+    // gain 2
+    raw(0, 1, 1) = (3 << 14) + 100;
+    raw(1, 1, 1) = (3 << 14) + 200;
+    raw(2, 1, 1) = (3 << 14) + 300;
+
+    auto pedestal = calculate_pedestal<double>(raw.view(), 4);
+
+    REQUIRE(pedestal.size() == raw.size());
+    CHECK(pedestal(0, 0, 0) == 200);
+    CHECK(pedestal(1, 0, 0) == 0);
+    CHECK(pedestal(1, 0, 1) == 200);
+
+    auto pedestal_gain0 = calculate_pedestal_g0<double>(raw.view(), 4);
+
+    REQUIRE(pedestal_gain0.size() == 4);
+    CHECK(pedestal_gain0(0, 0) == 200);
+    CHECK(pedestal_gain0(1, 0) == 38);
+}
\ No newline at end of file

From d6222027d07fe7bc6f89661f856301ec970de633 Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Fri, 25 Jul 2025 10:40:32 +0200
Subject: [PATCH 11/13] move constructor for Ndim-1

---
 include/aare/NDArray.hpp     | 41 ++++++++++++++++++------------------
 include/aare/NDView.hpp      | 26 +++++++++++++++++++++++
 include/aare/calibration.hpp | 20 ++++++++----------
 src/NDArray.test.cpp         | 24 +++++++++++++++++++++
 4 files changed, 79 insertions(+), 32 deletions(-)

diff --git a/include/aare/NDArray.hpp b/include/aare/NDArray.hpp
index c3edbfb..3495276 100644
--- a/include/aare/NDArray.hpp
+++ b/include/aare/NDArray.hpp
@@ -25,7 +25,7 @@ template <typename T, ssize_t Ndim = 2>
 class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
     std::array<ssize_t, Ndim> shape_;
     std::array<ssize_t, Ndim> strides_;
-    size_t size_{};
+    size_t size_{}; //TODO! do we need to store size when we have shape?
     T *data_;
 
   public:
@@ -43,8 +43,7 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
      */
     explicit NDArray(std::array<ssize_t, Ndim> shape)
         : shape_(shape), strides_(c_strides<Ndim>(shape_)),
-          size_(std::accumulate(shape_.begin(), shape_.end(), 1,
-                                std::multiplies<>())),
+          size_(num_elements(shape_)),
           data_(new T[size_]) {}
 
     /**
@@ -79,6 +78,24 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
         other.reset(); // TODO! is this necessary?
     }
 
+
+    //Move constructor from an an array with Ndim + 1
+    template <ssize_t M, typename = std::enable_if_t<(M == Ndim + 1)>>
+    NDArray(NDArray<T, M> &&other) 
+        : shape_(drop_first_dim(other.shape())),
+          strides_(c_strides<Ndim>(shape_)), size_(num_elements(shape_)),
+          data_(other.data()) {
+
+            // For now only allow move if the size matches, to avoid unreachable data
+            // if the use case arises we can remove this check
+            if(size() != other.size()) {
+                data_ = nullptr; // avoid double free, other will clean up the memory in it's destructor
+                throw std::runtime_error(LOCATION +
+                                         "Size mismatch in move constructor of NDArray<T, Ndim-1>");
+            }
+        other.reset();
+    }
+
     // Copy constructor
     NDArray(const NDArray &other)
         : shape_(other.shape_), strides_(c_strides<Ndim>(shape_)),
@@ -212,24 +229,6 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
     void Print_all();
     void Print_some();
 
-    template <ssize_t M = Ndim, typename = std::enable_if_t<(M > 2)>>
-    NDArray<T, Ndim - 1> drop_dimension() && {
-
-        std::array<ssize_t, Ndim - 1> new_shape;
-
-        std::copy(shape_.begin() + 1, shape_.begin() + Ndim, new_shape.begin());
-
-        NDArray<T, Ndim - 1> new_array(new_shape);
-
-        delete new_array.data();
-
-        new_array.data_ref() = data_;
-
-        this->reset();
-
-        return new_array;
-    }
-
     void reset() {
         data_ = nullptr;
         size_ = 0;
diff --git a/include/aare/NDView.hpp b/include/aare/NDView.hpp
index 56bca2f..c946c26 100644
--- a/include/aare/NDView.hpp
+++ b/include/aare/NDView.hpp
@@ -26,6 +26,32 @@ Shape<Ndim> make_shape(const std::vector<size_t> &shape) {
     return arr;
 }
 
+
+/**
+ * @brief Helper function to drop the first dimension of a shape.
+ * This is useful when you want to create a 2D view from a 3D array.
+ * @param shape The shape to drop the first dimension from.
+ * @return A new shape with the first dimension dropped.
+ */
+template<size_t Ndim>
+Shape<Ndim-1> drop_first_dim(const Shape<Ndim> &shape) {
+    Shape<Ndim - 1> new_shape;
+    std::copy(shape.begin() + 1, shape.end(), new_shape.begin());
+    return new_shape;
+}
+
+/**
+ * @brief Helper function when constructing NDArray/NDView. Calculates the number
+ * of elements in the resulting array from a shape.
+ * @param shape The shape to calculate the number of elements for.
+ * @return The number of elements in and NDArray/NDView of that shape.
+ */
+template <size_t Ndim>
+size_t num_elements(const Shape<Ndim> &shape) {
+    return std::accumulate(shape.begin(), shape.end(), 1,
+                           std::multiplies<size_t>());
+}
+
 template <ssize_t Dim = 0, typename Strides>
 ssize_t element_offset(const Strides & /*unused*/) {
     return 0;
diff --git a/include/aare/calibration.hpp b/include/aare/calibration.hpp
index a94705f..3b18da6 100644
--- a/include/aare/calibration.hpp
+++ b/include/aare/calibration.hpp
@@ -163,24 +163,22 @@ calculate_pedestal(NDView<uint16_t, 3> raw_data, ssize_t n_threads) {
                 NDView<uint16_t, 3>)>(&sum_and_count_per_gain<only_gain0>),
             view));
     }
+    Shape<3> shape{num_gains, raw_data.shape(1), raw_data.shape(2)};
+    NDArray<size_t, 3> accumulator(shape, 0);
+    NDArray<size_t, 3> count(shape, 0);
 
-    NDArray<size_t, 3> accumulator(
-        std::array<ssize_t, 3>{num_gains, raw_data.shape(1), raw_data.shape(2)},
-        0);
-    NDArray<size_t, 3> count(
-        std::array<ssize_t, 3>{num_gains, raw_data.shape(1), raw_data.shape(2)},
-        0);
+    // Combine the results from the futures
     for (auto &f : futures) {
         auto [acc, cnt] = f.get();
         accumulator += acc;
         count += cnt;
     }
 
-    if constexpr (only_gain0) {
-        return safe_divide<T>(accumulator, count).drop_dimension();
-    } else {
-        return safe_divide<T>(accumulator, count);
-    }
+
+    // Will move to a NDArray<T, 3 - static_cast<ssize_t>(only_gain0)>
+    // if only_gain0 is true
+    return safe_divide<T>(accumulator, count);
+
 }
 
 /**
diff --git a/src/NDArray.test.cpp b/src/NDArray.test.cpp
index 9af8379..a4d687a 100644
--- a/src/NDArray.test.cpp
+++ b/src/NDArray.test.cpp
@@ -427,4 +427,28 @@ TEST_CASE("Construct an NDArray from an std::array") {
     for (uint32_t i = 0; i < a.size(); ++i) {
         REQUIRE(a(i) == b[i]);
     }
+}
+
+
+TEST_CASE("Move construct from an array with Ndim + 1") {
+    NDArray<int, 3> a({{1,2,2}}, 0);
+    a(0, 0, 0) = 1;
+    a(0, 0, 1) = 2;
+    a(0, 1, 0) = 3;
+    a(0, 1, 1) = 4;
+
+
+    NDArray<int, 2> b(std::move(a));
+    REQUIRE(b.shape() == Shape<2>{2,2});
+    REQUIRE(b.size() == 4);
+    REQUIRE(b(0, 0) == 1);
+    REQUIRE(b(0, 1) == 2);
+    REQUIRE(b(1, 0) == 3);
+    REQUIRE(b(1, 1) == 4);
+
+}
+
+TEST_CASE("Move construct from an array with Ndim + 1 throws on size mismatch") {
+    NDArray<int, 3> a({{2,2,2}}, 0);
+    REQUIRE_THROWS(NDArray<int, 2>(std::move(a)));
 }
\ No newline at end of file

From 3d6858ad3382a6ee11abc6357697382eac7f94f9 Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Fri, 25 Jul 2025 10:42:47 +0200
Subject: [PATCH 12/13] removed data_ref

---
 include/aare/NDArray.hpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/include/aare/NDArray.hpp b/include/aare/NDArray.hpp
index 3495276..731ba2f 100644
--- a/include/aare/NDArray.hpp
+++ b/include/aare/NDArray.hpp
@@ -202,7 +202,6 @@ class NDArray : public ArrayExpr<NDArray<T, Ndim>, Ndim> {
     const T &operator[](ssize_t i) const { return data_[i]; }
 
     T *data() { return data_; }
-    T *&data_ref() { return data_; }
     std::byte *buffer() { return reinterpret_cast<std::byte *>(data_); }
     ssize_t size() const { return static_cast<ssize_t>(size_); }
     size_t total_bytes() const { return size_ * sizeof(T); }

From 89bb8776ea8ff1699ec340bf279e05a78b503687 Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Fri, 25 Jul 2025 11:44:27 +0200
Subject: [PATCH 13/13] check Ndim on drop_first_dim

---
 include/aare/NDView.hpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/include/aare/NDView.hpp b/include/aare/NDView.hpp
index c946c26..0aa4f78 100644
--- a/include/aare/NDView.hpp
+++ b/include/aare/NDView.hpp
@@ -35,6 +35,7 @@ Shape<Ndim> make_shape(const std::vector<size_t> &shape) {
  */
 template<size_t Ndim>
 Shape<Ndim-1> drop_first_dim(const Shape<Ndim> &shape) {
+    static_assert(Ndim > 1, "Cannot drop first dimension from a 1D shape");
     Shape<Ndim - 1> new_shape;
     std::copy(shape.begin() + 1, shape.end(), new_shape.begin());
     return new_shape;