Translate entire one pass of common-mode correction into cython

src/cal_tools/agipdalgs.pyx

@@ -147,69 +147,79 @@ def gain_choose(cnp.ndarray[cnp.uint8_t, ndim=3] a, cnp.ndarray[choices_t, ndim=
 
 @boundscheck(False)
 @wraparound(False)
-def sum_and_count_in_range_asic(float[:, :, :] arr, long[:] ix,
-                                float lower, float upper):
-    """
-    Return the sum & count of values where lower <= x <= upper,
-    across axes 2 & 3 (pixels within an ASIC, as reshaped by AGIPD correction code).
-    Specialised function for performance.
-    """
-    cdef int i, j, k, l, m
+def cm_correction_1pass(float[:, :, :] arr, unsigned short[:] cellid,
+                        float lower, float upper, float fraction):
+    """Apply 1 pass of common mode correction"""
+    cdef long long nfrm = arr.shape[0]
+    cdef long long nx = arr.shape[1]
+    cdef long long ny = arr.shape[2]
+
+    cdef float asic_thr = fraction * nx * ny
+
+    cdef double[:, :, :] crng_sum = np.zeros([11, nx, ny], dtype=np.float64)
+    cdef long long[:, :, :] crng_cnt = np.zeros([11, nx, ny], dtype=np.int64)
+    cdef long long[:, :, :] crng_tot = np.zeros([11, nx, ny], dtype=np.int64)
+
+    cdef long long asic_cnt, cnt, tot, i, l, m
+    cdef int row
+    cdef double asic_sum
     cdef float value
-    
-    cdef int nfrm = ix.size
-    cdef int nx = arr.shape[1]
-    cdef int ny = arr.shape[2]
-    
-    # Drop axes -2 & -1 (pixel dimensions within each ASIC)
-    count_arr = np.zeros(nfrm, dtype=np.uint64)
-    sum_arr = np.zeros(nfrm, dtype=np.float64)
-
-    cdef unsigned long long[:] count = count_arr
-    cdef double[:] sum_ = sum_arr
+    cdef bint z
+    cdef bint used_row[11]
+
     with nogil:
+        for row in range(11):
+            used_row[row] = False
+        # sum and count intensities over cell rows and trains
+        # the result has dimentionality [11, 64, 64]
         for i in range(nfrm):
-            k = ix[i]
+            row = cellid[i] // 32
+            used_row[row] = True
             for l in range(nx):
                 for m in range(ny):
-                    value = arr[k, l, m]
-                    if lower <= value <= upper:
-                        sum_[i] += value
-                        count[i] += 1
-
-    return sum_arr, count_arr
-
+                    value = arr[i, l, m]
+                    z = (lower <= value) and (value <= upper)
+                    crng_sum[row, l, m] += z * value
+                    crng_cnt[row, l, m] += z
+                    crng_tot[row, l, m] += 1
+
+        # find average values if there are more values in the interval
+        # than `fraction`
+        for row in range(11):
+            if used_row[row]:
+                for l in range(nx):
+                    for m in range(ny):
+                        tot = crng_tot[row, l, m]
+                        cnt = crng_cnt[row, l, m]
+                        z = cnt < (fraction * tot)
+                        crng_sum[row, l, m] = (
+                            .0 if z else (crng_sum[row, l, m] / cnt))
+
+        # substract mean value from the intensities
+        for i in range(nfrm):
+            row = cellid[i] // 32
+            for l in range(nx):
+                for m in range(ny):
+                    arr[i, l, m] -= crng_sum[row, l, m]
 
-@boundscheck(False)
-@wraparound(False)
-def sum_and_count_in_range_cell(float[:, :, :] arr, long[:] ix,
-                                float lower, float upper):
-    """
-    Return the sum & count of values where lower <= x <= upper,
-    across axes 0 & 1 (memory cells in the same row, as reshaped by AGIPD correction code).
-    Specialised function for performance.
-    """
-    cdef int i, j, k, l, m
-    cdef float value
-    
-    cdef int nfrm = ix.size
-    cdef int nx = arr.shape[1]
-    cdef int ny = arr.shape[2]
-    
-    # Drop axes 0
-    count_arr = np.zeros([nx, ny], dtype=np.uint64)
-    sum_arr = np.zeros([nx, ny], dtype=np.float64)
-    
-    cdef unsigned long long[:, :] count = count_arr
-    cdef double[:, :] sum_ = sum_arr
-    with nogil:
+        # over ASIC pixels
         for i in range(nfrm):
-            k = ix[i]
+            asic_sum = .0
+            asic_cnt = 0
+            # sum and count
             for l in range(nx):
                 for m in range(ny):
-                    value = arr[k, l, m]
-                    if lower <= value <= upper:
-                        sum_[l, m] += value
-                        count[l, m] += 1
+                    value = arr[i, l, m]
+                    z = (lower <= value) and (value <= upper)
+                    asic_sum += z * value
+                    asic_cnt += z
+
+            # find average values if there are more values in the interval
+            # than `fraction`
+            z = asic_cnt < asic_thr
+            asic_sum = .0 if z else (asic_sum / asic_cnt)
 
-    return sum_arr, count_arr
+            # substract mean value from the intensities
+            for l in range(nx):
+                for m in range(ny):
+                    arr[i, l, m] -= asic_sum

src/cal_tools/agipdlib.py

@@ -636,7 +636,6 @@ class AgipdCorrections:
         :param i_proc: Index of shared memory array to process
         :param asic: Asic number to process
         """
-
         if not self.corr_bools.get("common_mode"):
             return
         dark_min = self.cm_dark_min
@@ -647,35 +646,13 @@ class AgipdCorrections:
         if n_img == 0:
             return
         cell_id = self.shared_dict[i_proc]['cellId'][:n_img]
-        row_id = cell_id // 32
         data = self.shared_dict[i_proc]['data'][:n_img]
         data = data.reshape(-1, 8, 64, 2, 64)
 
         asic_data = data[:, asic % 8, :, asic // 8, :]
-        # Loop over rows of cells
-        for cell_row in range(11):
-            irow = np.flatnonzero(row_id == cell_row)
-            if not irow.size:
-                continue
-            # Loop over iterations
-            for _ in range(n_itr):
-                # Cell common mode
-                cell_cm_sum, cell_cm_count = \
-                    calgs.sum_and_count_in_range_cell(
-                        asic_data, irow, dark_min, dark_max)
-                cell_cm = cell_cm_sum / cell_cm_count
-
-                cell_cm[cell_cm_count < fraction * irow.size] = 0
-                asic_data[irow] -= cell_cm[None, :, :]
-
-                # Asics common mode
-                asic_cm_sum, asic_cm_count = \
-                    calgs.sum_and_count_in_range_asic(
-                        asic_data, irow, dark_min, dark_max)
-                asic_cm = asic_cm_sum / asic_cm_count
-
-                asic_cm[asic_cm_count < fraction * 4096] = 0
-                asic_data[irow] -= asic_cm[:, None, None]
+        for _ in range(n_itr):
+            calgs.cm_correction_1pass(
+                asic_data, cell_id, dark_min, dark_max, fraction)
 
     def mask_zero_std(self, i_proc, cells):
         """