S K-edge normalization

This change is to accommodate the case (encountered at the sulfur K-edge) where only the sum of 'n' and 'p' ROI can be normalized against '0'.

Changes

• depreciate inspect_plane_fitting and replace it with inspect_ff_fitting
• framework to work with different flat field normalization function
• implement 2D polyline normalization
• compute np_mean directly in process

src/toolbox_scs/routines/boz.py

@@ -874,6 +874,8 @@ def _plane_flat_field(p, roi, params):
 def compute_flat_field_correction(rois, params, plot=False):
     if params.ff_type == 'plane':
         return compute_plane_flat_field_correction(rois, params, plot)
+    elif params.ff_type == 'polyline':
+        return compute_polyline_flat_field_correction(rois, params, plot)
     else:
         raise ValueError(f'Uknown flat field type {params.ff_type}')
 
@@ -920,6 +922,54 @@ def compute_plane_flat_field_correction(rois, params, plot=False):
 
     return flat_field
 
+def compute_polyline_flat_field_correction(rois, params, plot=False):
+    """Compute the 1D polyline field correction on beam rois.
+
+    Inputs
+    ------
+    rois: dictionnary of beam rois['n', '0', 'p']
+    params: parameters
+    plot: boolean, True by default, diagnostic plot
+
+    Returns
+    -------
+    numpy 2D array of the flat-field correction evaluated over one DSSC ladder
+    (2 sensors)
+    """
+    flat_field = np.ones((128, 512))
+
+    z = params.get_flat_field()
+    coeffs = np.logspace(-(z.shape[0]-1), 0, z.shape[0])
+    poly = np.poly1d(z*coeffs)
+    n = rois['n']
+    p = rois['p']
+    wn = n['xh']-n['xl']
+    wp = p['xh']-p['xl']
+    assert wn == wp, (\
+        f"For polyline flat field normalization, both 'n' and 'p' ROIs "
+        f"must have the same width {wn} and {wp}px"
+    )
+    x = np.arange(wn)
+    norm = poly(x)
+    
+    n_int = flat_field[n['yl']:n['yh'], n['xl']:n['xh']]
+    flat_field[n['yl']:n['yh'], n['xl']:n['xh']]  = \
+                        norm*n_int
+    
+    p_int = flat_field[p['yl']:p['yh'], p['xl']:p['xh']]
+    flat_field[p['yl']:p['yh'], p['xl']:p['xh']] = \
+                        norm*p_int  # not the mirror
+
+    if plot:
+        f, ax = plt.subplots(1, 1, figsize=(6, 2))
+        img = ax.pcolormesh(
+            np.flipud(flat_field[:, :256]), cmap='Greys_r')
+        f.colorbar(img, ax=[ax], label='amplitude')
+        ax.set_xlabel('px')
+        ax.set_ylabel('px')
+        ax.set_aspect('equal')
+
+    return flat_field
 
 def inspect_flat_field_domain(avg, rois, prod_th, ratio_th, vmin=None, vmax=None):
     """Extract beams roi from average image and compute the ratio.
@@ -1056,6 +1106,87 @@ def inspect_plane_fitting(avg, rois, domain=None, vmin=None, vmax=None):
 
     return fig
 
+def inspect_ff_fitting_sk(avg, rois, ff, domain=None, vmin=None, vmax=None):
+    """Extract beams roi from average image and compute the ratio.
+
+    Inputs
+    ------
+    avg: module average image with no saturated shots for the flat-field
+         determination
+    rois: dictionnary of rois
+    ff: 2D array, flat field normalization
+    domain: list of domain mask for the -1st and +1st order
+    vmin: imshow vmin level, default None will use 5 percentile value
+    vmax: imshow vmax level, default None will use 99.8 percentile value
+
+    Returns
+    -------
+    fig: matplotlib figure plotted
+    """
+    if vmin is None:
+        vmin = np.percentile(avg, 5)
+    if vmax is None:
+        vmax = np.percentile(avg, 99.8)
+
+    refn = avg[rois['n']['yl']:rois['n']['yh'],
+               rois['n']['xl']:rois['n']['xh']]
+    refp = avg[rois['p']['yl']:rois['p']['yh'],
+               rois['p']['xl']:rois['p']['xh']]
+    mid = avg[rois['0']['yl']:rois['0']['yh'],
+              rois['0']['xl']:rois['0']['xh']]
+    mref = 0.5*(refn + refp)
+
+    ffn = ff[rois['n']['yl']:rois['n']['yh'],
+             rois['n']['xl']:rois['n']['xh']]
+    ffp = ff[rois['p']['yl']:rois['p']['yh'],
+             rois['p']['xl']:rois['p']['xh']]
+    ffmid = ff[rois['0']['yl']:rois['0']['yh'],
+               rois['0']['xl']:rois['0']['xh']]
+    norm = (ffn+ffp) 
+
+    fig, axs = plt.subplots(3, 3, sharex=True, sharey=True,
+        figsize=(8, 4))
+    im = axs[0, 0].imshow(mref)
+    axs[0, 0].set_title('(n+p)/2')
+    fig.colorbar(im, ax=axs[0, 0])
+
+    im = axs[1, 0].imshow(mid)
+    axs[1, 0].set_title('0')
+    fig.colorbar(im, ax=axs[1, 0])
+
+    im = axs[2, 0].imshow(mid/mref-1, cmap='RdBu_r', vmin=-1, vmax=1)
+    axs[2, 0].set_title('02/(n+p) - 1')
+    fig.colorbar(im, ax=axs[2, 0])
+
+    im = axs[0, 1].imshow(norm)
+    axs[0, 1].set_title('n+p norm')
+    fig.colorbar(im, ax=axs[0, 1])
+
+    im = axs[1, 1].imshow(ffmid)
+    axs[1, 1].set_title('0 norm')
+    fig.colorbar(im, ax=axs[1, 1])
+
+    im = axs[2, 1].imshow(ffmid/norm-1, cmap='RdBu_r', vmin=-1, vmax=1)
+    axs[2, 1].set_title('0/(n+p)norm-1')
+    fig.colorbar(im, ax=axs[2, 1])
+
+
+    im = axs[0, 2].imshow(mref/norm)
+    axs[0, 2].set_title('(n+p)/norm')
+    fig.colorbar(im, ax=axs[0, 2])
+
+    im = axs[1, 2].imshow(mid)
+    axs[1, 2].set_title('0')
+    fig.colorbar(im, ax=axs[1, 2])
+
+    im = axs[2, 2].imshow(mid/(mref/norm)-1, cmap='RdBu_r', vmin=-1, vmax=1)
+    axs[2, 2].set_title('0/(n+p) norm')
+    fig.colorbar(im, ax=axs[2, 2])
+
+    # fig.suptitle(f'{proposalNB}-run{runNB}-dark{darkrunNB} sat={sat_level}')
+
+    return fig
+
 
 def plane_fitting_domain(avg, rois, prod_th, ratio_th):
     """Extract beams roi, compute their ratio and the domain.
@@ -1216,8 +1347,47 @@ def ff_refine_crit(p, alpha, params, arr_dark, arr, tid, rois,
 
     return bad + 1e3*(alpha*err_sigma + (1-alpha)*err_mean)
 
+def ff_refine_crit_sk(p, alpha, params, arr_dark, arr, tid, rois,
+    mask, sat_level=511):
+    """Criteria for the ff_refine_fit, combining 'n' and 'p' as reference.
+
+    Inputs
+    ------
+    p: ff plane
+    params: parameters
+    arr_dark: dark data
+    arr: data
+    tid: train id of arr data
+    rois: ['n', '0', 'p', 'sat'] rois
+    mask: mask fo good pixels
+    sat_level: integer, default 511, at which level pixel begin to saturate
+
+    Returns
+    -------
+    sum of standard deviation on binned 0th order intensity
+    """
+    params.set_flat_field(p, params.ff_type)
+    ff = compute_flat_field_correction(rois, params)
+    if np.any(ff < 0.0):
+        bad = 1e6
+    else:
+        bad = 0.0
+
+    data = process(None, arr_dark, arr, tid, rois, mask, ff,
+        sat_level, params._using_gpu)
+
+    # drop saturated shots
+    d = data.where(data['sat_sat'] == False, drop=True)
+
+    r = xas(d, 40, Iokey='np_mean', Itkey='0', nrjkey='0', fluorescence=True)
+
+    err_sigma = np.nansum(r['sigmaA'])
+    err_mean = (1.0 - np.nanmean(r['muA']))**2
+
+    return bad + 1e3*(alpha*err_sigma + (1-alpha)*err_mean)
+
 
-def ff_refine_fit(params):
+def ff_refine_fit(params, crit=ff_refine_crit):
     """Refine the flat-field fit by minimizing data spread.
 
     Inputs
@@ -1258,11 +1428,11 @@ def ff_refine_fit(params):
         fit_callback.counter += 1
 
         temp = list(fixed_p)
-        Jalpha = ff_refine_crit(x, *temp)
+        Jalpha = crit(x, *temp)
         temp[0] = 0
-        J0 = ff_refine_crit(x, *temp)
+        J0 = crit(x, *temp)
         temp[0] = 1
-        J1 = ff_refine_crit(x, *temp)
+        J1 = crit(x, *temp)
         fit_callback.res.append([J0, Jalpha, J1])
         print(f'{fit_callback.counter-1}: {time_delta} '
                 f'({J0}, {Jalpha}, {J1}), {x}')
@@ -1270,7 +1440,7 @@ def ff_refine_fit(params):
         return False
 
     fit_callback(p0)
-    res = minimize(ff_refine_crit, p0, fixed_p,
+    res = minimize(crit, p0, fixed_p,
         options={'disp': True, 'maxiter': params.ff_max_iter},
         callback=fit_callback)
 
@@ -1400,8 +1570,6 @@ def nl_crit_sk(p, domain, alpha, arr_dark, arr, tid, rois, mask, flat_field,
     # drop saturated shots
     d = data.where(data['sat_sat'] == False, drop=True)
 
-    d['np_mean'] = 0.5*(d['n'] + d['p'])
-
     v = snr(d['np_mean'].values.flatten(), d['0'].values.flatten(),
             methods=['weighted'])
     err = 1e8*v['weighted']['s']**2
@@ -1596,7 +1764,7 @@ def inspect_Fnl(Fnl):
 
 
 def inspect_correction(params, gain=None):
-    """Criteria for the non linear correction.
+    """Comparison plot of the different corrections.
 
     Inputs
     ------
@@ -1716,7 +1884,7 @@ def inspect_correction(params, gain=None):
     return f
 
 def inspect_correction_sk(params, ff, gain=None):
-    """Criteria for the non linear correction.
+    """Comparison plot of the different corrections, combining 'n' and 'p'.
 
     Inputs
     ------
@@ -2018,6 +2186,8 @@ def process_module(arr, tid, dark, rois, mask=None, sat_level=511,
 
     # compute dark corrected ROI values
     v = {}
+    r_ff = {}
+    ff = {}
     for n in rois.keys():
 
         r[n] = r[n] - rd[n]
@@ -2026,25 +2196,36 @@ def process_module(arr, tid, dark, rois, mask=None, sat_level=511,
             # TODO:  flat-field should not be applied on intra darks
             # ff = flat_field[:, rois[n]['yl']:rois[n]['yh'],
             #                 rois[n]['xl']:rois[n]['xh']]
-            ff = flat_field[rois[n]['yl']:rois[n]['yh'],
-                            rois[n]['xl']:rois[n]['xh']]
-            r[n] = r[n]/ff
+            ff[n] = flat_field[rois[n]['yl']:rois[n]['yh'],
+                        rois[n]['xl']:rois[n]['xh']]
+            r_ff[n] = r[n]/ff[n]
+        else:
+            ff[n] = 1.0
+            r_ff[n] = r[n]
 
-        v[n] = r[n].sum(axis=(2, 3))
+        v[n] = r_ff[n].sum(axis=(2, 3))
+    
+    # np_mean roi where we normalize the sum of flat_field
+    np_mean  = 0.5*(r['n'] + r['p'])/(ff['n'] + ff['p'])
+    v['np_mean'] = np_mean.sum(axis=(2,3))
 
     res = xr.Dataset()
 
     dims = ['trainId', 'pulseId']
     r_coords = {'trainId': tid, 'pulseId': np.arange(0, narr.shape[1])}
     for n in rois.keys():
+        res[n] = xr.DataArray(ensure_on_host(v[n]), coords=r_coords, dims=dims)
         res[n + '_sat'] = xr.DataArray(ensure_on_host(r_sat[n][:, :]),
                                        coords=r_coords, dims=dims)
-        res[n] = xr.DataArray(ensure_on_host(v[n]), coords=r_coords, dims=dims)
+    res['np_mean'] = xr.DataArray(ensure_on_host(v['np_mean']),
+                                  coords=r_coords, dims=dims)
+    res['np_mean_sat'] = res['n_sat'] + res['p_sat']
 
     for n in rois.keys():
         roi = rois[n]
         res[n + '_area'] = xr.DataArray(np.array([
             (roi['yh'] - roi['yl'])*(roi['xh'] - roi['xl'])]))
+    res['np_mean_area'] = res['n_area'] + res['p_area']
 
     return res