diff --git a/src/cal_tools/agipdalgs.pyx b/src/cal_tools/agipdalgs.pyx
index 649511df6ecdf15b3674590d27c29c6b9e56e34b..a0dc4912c489de0811e7062f2e2d25a4d0c1e4a0 100644
--- a/src/cal_tools/agipdalgs.pyx
+++ b/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
diff --git a/src/cal_tools/agipdlib.py b/src/cal_tools/agipdlib.py
index 03f9680d70ff9f2f113ade8e92c83776ea1910cb..862864080e8b84a169bad003710c11b45b35f491 100644
--- a/src/cal_tools/agipdlib.py
+++ b/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):
"""