Merge branch 'fix/cm-for-variable-pattern' into 'master'

[AGIPD] [CORRECT] Fix array reshaping in common mode correction

See merge request detectors/pycalibration!724

src/cal_tools/agipdalgs.pyx

@@ -147,63 +147,108 @@ 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(cnp.ndarray[float, ndim=4] arr, float lower, float upper):
+def cm_correction(float[:, :, :] arr, unsigned short[:] cellid,
+                  float lower, float upper, float fraction):
+    """Apply one iteration of common-mode correction
+
+    The common-mode correction shifts the position of the noise peak for
+    a slice of pixels to zero. The position is estimated as mean signal of
+    pixels in a range between `lower` and `upper` boundaries. If the noise
+    peak located in this boundaries then the mean value works as a robust
+    estimator, which converges to the true noise peak position with
+    iterations. The correction is applied only if the number of pixels in
+    the given range above the `fraction` of the total number of the pixels
+    in the slice.
+
+    This function performs consequently cells common-mode correction and
+    ASIC common-mode correction.
+    Cells correction is calculated across trains and groups of 32 cells.
+    ASIC correction is calculated across the pixes in ASIC on a single image.
+
+    To converge this function should be iterated.
+    
+    This function performs correction in-place, altering `arr`.
+
+    :param arr: array of images cropped to a single ASIC and
+        stacked together over trains and pulses in the first dimension
+    :param cellid: array of cell IDs, must have the length equal to
+        the number of images
+    :param lower: the lower signal value in ADU to consider pixel as a dark
+    :param upper: the upper signal value in ADU to consider pixel as a dark
+    :param fraction: the fraction of the dark pixels in the slice which is
+        considered to be enough to apply the common-mode correction to this
+        slice
     """
-    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
-    cdef float value
-    cdef cnp.ndarray[unsigned long long, ndim=2] count
-    cdef cnp.ndarray[double, ndim=2] sum_
-
-    # Drop axes -2 & -1 (pixel dimensions within each ASIC)
-    out_shape = arr[:, :, 0, 0].shape
-    count = np.zeros(out_shape, dtype=np.uint64)
-    sum_ = np.zeros(out_shape, dtype=np.float64)
-
-    with nogil:
-        for i in range(arr.shape[0]):
-            for k in range(arr.shape[1]):
-                for l in range(arr.shape[2]):
-                    for m in range(arr.shape[3]):
-                        value = arr[i, k, l, m]
-                        if lower <= value <= upper:
-                            sum_[i, k] += value
-                            count[i, k] += 1
-    return sum_, count
+    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
 
-@boundscheck(False)
-@wraparound(False)
-def sum_and_count_in_range_cell(cnp.ndarray[float, ndim=4] arr, 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 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 int i, j, k, l, m,
+    cdef long long asic_cnt, cnt, tot, i, l, m
+    cdef int row
+    cdef double asic_sum
     cdef float value
-    cdef cnp.ndarray[unsigned long long, ndim=2] count
-    cdef cnp.ndarray[double, ndim=2] sum_
-
-    # Drop axes 0 & 1
-    out_shape = arr[0, 0, :, :].shape
-    count = np.zeros(out_shape, dtype=np.uint64)
-    sum_ = np.zeros(out_shape, dtype=np.float64)
-
+    cdef bint z
+    cdef bint used_row[11]
 
     with nogil:
-        for i in range(arr.shape[0]):
-            for k in range(arr.shape[1]):
-                for l in range(arr.shape[2]):
-                    for m in range(arr.shape[3]):
-                        value = arr[i, k, l, m]
-                        if lower <= value <= upper:
-                            sum_[l, m] += value
-                            count[l, m] += 1
-
-    return sum_, count
+        for row in range(11):
+            used_row[row] = False
+        # sum and count intensities over cell rows and trains
+        # the result has dimensionality [11, 64, 64]
+        for i in range(nfrm):
+            row = cellid[i] // 32
+            used_row[row] = True
+            for l in range(nx):
+                for m in range(ny):
+                    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))
+
+        # subtract 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]
+
+        # over ASIC pixels
+        for i in range(nfrm):
+            asic_sum = .0
+            asic_cnt = 0
+            # sum and count
+            for l in range(nx):
+                for m in range(ny):
+                    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)
+
+            # subtract 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,44 +646,13 @@ class AgipdCorrections:
         if n_img == 0:
             return
         cell_id = self.shared_dict[i_proc]['cellId'][:n_img]
-        train_id = self.shared_dict[i_proc]['trainId'][:n_img]
-        cell_ids = cell_id[train_id == train_id[0]]
-        n_cells = cell_ids.size
-        data = self.shared_dict[i_proc]['data'][:n_img].reshape(-1, n_cells,
-                                                                8, 64, 2, 64)
+        data = self.shared_dict[i_proc]['data'][:n_img]
+        data = data.reshape(-1, 8, 64, 2, 64)
 
-        # Loop over iterations
+        asic_data = data[:, asic % 8, :, asic // 8, :]
         for _ in range(n_itr):
-            # Loop over rows of cells
-            # TODO: check what occurs in case of 64 memory cells
-            #       as it will have less than 11 iterations
-            first = 0
-            for cell_row in range(11):
-                last = first + cell_ids[(cell_ids // 32) == cell_row].size
-                if first == last:
-                    continue
-                asic_data = data[:, first:last, asic % 8, :, asic // 8, :]
-
-                # Cell common mode
-                cell_cm_sum, cell_cm_count = \
-                    calgs.sum_and_count_in_range_cell(asic_data, dark_min,
-                                                      dark_max)
-                cell_cm = cell_cm_sum / cell_cm_count
-
-                # TODO: check files with less 256 trains
-                cell_cm[cell_cm_count < fraction * 32 * 256] = 0
-                asic_data[...] -= cell_cm[None, None, :, :]
-
-                # Asics common mode
-                asic_cm_sum, asic_cm_count = \
-                    calgs.sum_and_count_in_range_asic(asic_data, dark_min,
-                                                      dark_max)
-                asic_cm = asic_cm_sum / asic_cm_count
-
-                asic_cm[asic_cm_count < fraction * 64 * 64] = 0
-                asic_data[...] -= asic_cm[:, :, None, None]
-
-                first = last
+            calgs.cm_correction(
+                asic_data, cell_id, dark_min, dark_max, fraction)
 
     def mask_zero_std(self, i_proc, cells):
         """