Compute BOZ pixel position in a single function

src/toolbox_scs/routines/boz.py

@@ -267,6 +267,34 @@ def _get_pixel_pos(module):
     # keeping only module 15 pixel X,Y position
     return g.get_pixel_positions()[module][:, :, :2]
 
+def get_roi_pixel_pos(roi, params):
+    """Compute fake or real pixel position of an roi from roi center.
+
+    Inputs:
+    -------
+    roi: dictionnary
+    params: parameters
+
+    Returns:
+    --------
+    X, Y: 1-d array of pixel position.
+    """
+    if params.use_hex:
+        # DSSC pixel position on hexagonal lattice
+        X = params.pixel_pos[roi['yl']:roi['yh'], roi['xl']:roi['xh'], 0]
+        Y = params.pixel_pos[roi['yl']:roi['yh'], roi['xl']:roi['xh'], 1]
+    else:
+        nY, nX = roi['yh'] - roi['yl'], roi['xh'] - roi['xl']
+        X = np.arange(nX)/100
+        Y = np.arange(nY)[:, np.newaxis]/100
+    
+    # center of ROI is put to 0,0
+    X -= np.mean(X)
+    Y -= np.mean(Y)
+
+    return X, Y
+
+
 def _get_pixel_corners(module):
     """Compute the pixel corners of DSSC module."""
     # module pixel position
@@ -671,7 +699,7 @@ def inspect_rois(data_mean, rois, threshold=None, allrois=False):
 
 # Flat field related functions
 
-def _plane_flat_field(p, roi, pixel_pos, use_hex=False):
+def _plane_flat_field(p, roi, params):
     """Compute the p plane over the given roi.
 
     Given the plane parameters p, compute the plane over the roi
@@ -682,9 +710,7 @@ def _plane_flat_field(p, roi, pixel_pos, use_hex=False):
     p: a vector of a, b, c, d plane parameter with the
        plane given by ax+ by + cz + d = 0
     roi: a dictionnary roi['yh', 'yl', 'xh', 'xl']
-    pixel_pos: array of DSSC pixel position on hexagonal lattice
-    use_hex: boolean, use actual DSSC pixel position from pixel_pos
-        or fake cartesian pixel position
+    params: parameters
 
     Returns
     -------
@@ -693,19 +719,7 @@ def _plane_flat_field(p, roi, pixel_pos, use_hex=False):
     """
     a, b, c, d = p
 
-    if use_hex:
-        # DSSC pixel position on hexagonal lattice
-        pos = pixel_pos[roi['yl']:roi['yh'], roi['xl']:roi['xh'], :]
-        X = pos[:, :, 0]
-        Y = pos[:, :, 1]
-    else:
-        nY, nX = roi['yh'] - roi['yl'], roi['xh'] - roi['xl']
-        X = np.arange(nX)/100
-        Y = np.arange(nY)[:, np.newaxis]/100
-
-    # center of ROI is put to 0,0
-    X -= np.mean(X)
-    Y -= np.mean(Y)
+    X, Y = get_roi_pixel_pos(roi, params)
 
     Z = -(a*X + b*Y + d)/c
 
@@ -729,22 +743,20 @@ def compute_flat_field_correction(rois, params, plot=False):
     flat_field = np.ones((128, 512))
 
     plane = params.get_flat_field()
-    use_hex = params.use_hex
-    pixel_pos = params.pixel_pos
     force_mirror = params.force_mirror
 
     r = rois['n']
     flat_field[r['yl']:r['yh'], r['xl']:r['xh']] = \
-                       _plane_flat_field(plane[:4], r, pixel_pos, use_hex)
+                       _plane_flat_field(plane[:4], r, params)
     
     r = rois['p']
     if force_mirror:
         a, b, c, d = plane[:4]
         flat_field[r['yl']:r['yh'], r['xl']:r['xh']] = \
-                       _plane_flat_field([-a, b, c, d], r, pixel_pos, use_hex)
+                       _plane_flat_field([-a, b, c, d], r, params)
     else:
         flat_field[r['yl']:r['yh'], r['xl']:r['xh']] = \
-                       _plane_flat_field(plane[4:], r, pixel_pos, use_hex)
+                       _plane_flat_field(plane[4:], r, params)
 
     if plot:
         f, ax = plt.subplots(1, 1, figsize=(6, 2))
@@ -916,15 +928,17 @@ def plane_fitting_domain(avg, rois, prod_th, ratio_th):
     """
     centers = {}
 
-    for k, r in enumerate(['n', '0', 'p']):
-        centers[r] = np.array([(rois[r]['yl'] + rois[r]['yh'])//2,
-                      (rois[r]['xl'] + rois[r]['xh'])//2])
+    for k in ['n', '0', 'p']:
+        r = rois[k]
+        centers[k] = np.array([(r['yl'] + r['yh'])//2,
+                      (r['xl'] + r['xh'])//2])
 
     k = 'n'
-    num = avg[rois[k]['yl']:rois[k]['yh'], rois[k]['xl']:rois[k]['xh']]
+    r = rois[k]
+    num = avg[r['yl']:r['yh'], r['xl']:r['xh']]
     d = '0'
     denom = np.roll(avg, tuple(centers[k] - centers[d]))[
-        rois[k]['yl']:rois[k]['yh'], rois[k]['xl']:rois[k]['xh']]
+        r['yl']:r['yh'], r['xl']:r['xh']]
     n = num/denom
     prod = num*denom
     n_m = ((prod > prod_th[0]) * (prod < prod_th[1]) *
@@ -933,10 +947,11 @@ def plane_fitting_domain(avg, rois, prod_th, ratio_th):
     n[~np.isfinite(n)] = 0
 
     k = 'p'
-    num = avg[rois[k]['yl']:rois[k]['yh'], rois[k]['xl']:rois[k]['xh']]
+    r = rois[k]
+    num = avg[r['yl']:r['yh'], r['xl']:r['xh']]
     d = '0'
     denom = np.roll(avg, tuple(centers[k] - centers[d]))[
-        rois[k]['yl']:rois[k]['yh'], rois[k]['xl']:rois[k]['xh']]
+        r['yl']:r['yh'], r['xl']:r['xh']]
     p = num/denom
     prod = num*denom
     p_m = ((prod > prod_th[0]) * (prod < prod_th[1]) *
@@ -984,39 +999,13 @@ def plane_fitting(params):
         num_n = a_n**2 + b_n**2 + c_n**2
 
         roi = params.rois['n']
-        if params.use_hex:
-            # DSSC pixel position on hexagonal lattice
-            pos = params.pixel_pos[roi['yl']:roi['yh'], roi['xl']:roi['xh'], :]
-            X = pos[:, :, 0]
-            Y = pos[:, :, 1]
-        else:
-            nY, nX = n.shape
-            X = np.arange(nX)/100
-            Y = np.arange(nY)[:, np.newaxis]/100
-
-        # center of ROI is put to 0,0
-        X -= np.mean(X)
-        Y -= np.mean(Y)
-
+        X, Y = get_roi_pixel_pos(roi, params)
         d0_2 = np.sum(n_m*(a_n*X + b_n*Y + c_n*n + d_n)**2)/num_n
 
         num_p = a_p**2 + b_p**2 + c_p**2
 
         roi = params.rois['p']
-        if params.use_hex:
-            # DSSC pixel position on hexagonal lattice
-            pos = params.pixel_pos[roi['yl']:roi['yh'], roi['xl']:roi['xh'], :]
-            X = pos[:, :, 0]
-            Y = pos[:, :, 1]
-        else:
-            nY, nX = p.shape
-            X = np.arange(nX)/100
-            Y = np.arange(nY)[:, np.newaxis]/100
-
-        # center of ROI is put to 0,0
-        X -= np.mean(X)
-        Y -= np.mean(Y)
-
+        X, Y = get_roi_pixel_pos(roi, params)
         if params.force_mirror:
             d2_2 = np.sum(p_m*(-a_n*X + b_n*Y + c_n*p + d_n)**2)/num_n
         else: