From 9400e046be97ffccd5eeb4b86847cbb428e3b526 Mon Sep 17 00:00:00 2001
From: Saianeesh Keshav Haridas <accounts@saianeesh.com>
Date: Thu, 26 Dec 2024 18:49:32 +0000
Subject: [PATCH] feat: add basic jump fixing

---
 src/array_ops.cxx | 65 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/src/array_ops.cxx b/src/array_ops.cxx
index b32bee8f..aec403c0 100644
--- a/src/array_ops.cxx
+++ b/src/array_ops.cxx
@@ -882,6 +882,52 @@ void find_quantized_jumps(const bp::object & tod, const bp::object & out, const
     }
 }
 
+template <typename T>
+void subtract_jump_heights(const bp::object & tod, const bp::object & out, const bp::object & heights, const bp::object & jumps) {
+    BufferWrapper<T> tod_buf  ("tod",  tod,  false, std::vector<int>{-1, -1});
+    int ndet = tod_buf->shape[0];
+    int nsamp = tod_buf->shape[1];
+    if (tod_buf->strides[1] != tod_buf->itemsize || tod_buf->strides[0] != tod_buf->itemsize*nsamp)
+        throw buffer_exception("tod must be C-contiguous along last axis");
+    T* tod_data = (T*)tod_buf->buf;
+    BufferWrapper<T> out_buf  ("out",  out,  false, std::vector<int>{ndet, nsamp});
+    if (out_buf->strides[1] != out_buf->itemsize || out_buf->strides[0] != out_buf->itemsize*nsamp)
+        throw buffer_exception("out must be C-contiguous along last axis");
+    T* output = (T*)out_buf->buf;
+    BufferWrapper<T> h_buf  ("heights",  heights,  false, std::vector<int>{ndet, nsamp});
+    if (h_buf->strides[1] != h_buf->itemsize || h_buf->strides[0] != h_buf->itemsize*nsamp)
+        throw buffer_exception("heights must be C-contiguous along last axis");
+    T* jump_heights = (T*)h_buf->buf;
+    auto ranges = extract_ranges<int32_t>(jumps);
+
+    #pragma omp parallel for
+    for(int di = 0; di < ranges.size(); di++) {
+        int start = 0;
+        int stop = 0;
+        T min_h;
+        T max_h;
+        T height;
+        T to_sub = 0;
+        for (auto const &r: ranges[di].segments) {
+            start = di*nsamp + r.first;
+            for(int j = stop; j < start && to_sub != 0; j++) {
+                output[j] = tod_data[j] - to_sub;
+            }
+            stop = di*nsamp + r.second;
+            min_h = *(std::min_element(jump_heights+start, jump_heights+stop));
+            max_h = *(std::max_element(jump_heights+start, jump_heights+stop));
+            height = (max(abs(min_h), abs(max_h)) == abs(min_h)) ? min_h : max_h;
+            to_sub = to_sub + height;
+            for(int j = start; j < stop && to_sub != 0; j++) {
+                output[j] = tod_data[j] - to_sub;
+            }
+        }
+        for(int j = stop; j < di*nsamp + nsamp && to_sub != 0; j++) {
+            output[j] = tod_data[j] - to_sub;
+        }
+    }
+}
+
 template <typename T>
 using _interp_func_pointer = void (*)(const double* x, const double* y,
                                       const double* x_interp, T* y_interp,
@@ -1150,6 +1196,25 @@ PYBINDINGS("so3g")
             "Output will be 0 where jumps are not found and the assumed jump height where jumps are found.\n"
             "\n"
             "See details in docstring for find_quantized_jumps.\n");
+    bp::def("subtract_jump_heights", subtract_jump_heights<float>,
+            "subtract_jump_heights(tod, out, heights, jumps)"
+            "\n"
+            "Subtract fit jump heights from known jump locatations in a float32 array."
+            "If multiple samples in a jump have different heights, the largest height is used.\n"
+            "\n"
+            "Args:\n"
+            "  tod: data array (float32) with shape (ndet, nsamp)\n"
+            "  out: output array (float32) with shape (ndet, nsamp)\n"
+            "  heights: the height of the jump at each samples\n"
+            "           should be an array (float32) with shape (ndet, nsamp)\n"
+            "  jumps: RangesMatrix with the jump locations and shape (ndet, nsamp).\n");
+    bp::def("subtract_jump_heights64", subtract_jump_heights<double>,
+            "subtract_jump_heights64(tod, out, heights, jumps)"
+            "\n"
+            "Subtract fit jump heights from known jump locatations in a float64 array."
+            "If multiple samples in a jump have different heights, the largest height is used.\n"
+            "\n"
+            "See details in docstring for subtract_jump_heights.\n");
     bp::def("clean_flag", clean_flag,
             "clean_flag(flag, width)"
             "\n"