Make extract_digitizer_peaks() compatible with irregular patterns

src/toolbox_scs/detectors/digitizers.py

@@ -40,7 +40,7 @@ log = logging.getLogger(__name__)
 
 
 def peaks_from_raw_trace(traces, pulseStart, pulseStop, baseStart, baseStop,
-                         period, npulses, extra_dim):
+                         period=None, npulses=None, extra_dim=None):
     """
     Computes peaks from raw digitizer traces by trapezoidal integration.
 
@@ -62,7 +62,7 @@ def peaks_from_raw_trace(traces, pulseStart, pulseStop, baseStart, baseStop,
     npulses: int
         number of pulses. Ignored if intstart is a 1D-array.
     extra_dim: str
-        Name given to the dimension along the peaks.
+        Name given to the dimension along the peaks. Defaults to 'pulseId'.
 
     Returns
     -------
@@ -87,6 +87,8 @@ def peaks_from_raw_trace(traces, pulseStart, pulseStop, baseStart, baseStop,
     else:
         pulses = range(0, npulses*period, period)
 
+    if extra_dim is None:
+        extra_dim = 'pulseId'
     results = xr.DataArray(np.empty((traces.shape[0], len(pulses))),
                            coords=coords,
                            dims=['trainId', extra_dim])
@@ -1259,7 +1261,7 @@ def timFactorFromTable(voltage, photonEnergy, mcp=1):
 #############################################################################################
 
 def extract_digitizer_peaks(proposal, runNB, mnemonic, bunchPattern=None,
-                            integParams=None, autoFind=True, save=True,
+                            integParams=None, autoFind=True, save=False,
                             subdir='usr/processed_runs'):
     """
     Extract the peaks from digitizer raw traces and saves them into a file.
@@ -1321,7 +1323,6 @@ def extract_digitizer_peaks(proposal, runNB, mnemonic, bunchPattern=None,
         pulse_period = 440
 
     pattern = None
-    regular = True
     try:
         if bunchPattern == 'sase3':
             pattern = XrayPulses(run)
@@ -1331,58 +1332,56 @@ def extract_digitizer_peaks(proposal, runNB, mnemonic, bunchPattern=None,
         log.warning(e)
         bunchPattern = None
 
+    # Extract pulse_ids, npulses and period from bunch pattern
+    pulse_ids, npulses, period = None, None, None
+    regular = True
+    if pattern is not None:
+        if pattern.is_constant_pattern() is False:
+            log.info('The number of pulses changed during the run.')
+            pulse_ids = np.unique(pattern.pulse_ids(labelled=False, copy=False))
+            npulses, period = None, None
+            regular = False
+        else:
+            pulse_ids = pattern.peek_pulse_ids(labelled=False)
+            npulses = len(pulse_ids)
+            if npulses > 1:
+                periods = np.diff(pulse_ids)
+                if len(np.unique(periods)) > 1:
+                    regular = False
+                else:
+                    period = np.unique(periods)[0] * pulse_period
+
+    # Use integration parameters, adjust them to match bunch pattern if necessary
     if integParams is not None:
         required_keys = ['pulseStart', 'pulseStop', 'baseStart',
-                         'baseStop', 'period', 'npulses']
+                         'baseStop']
+        if bunchPattern is None:
+            required_keys += ['period', 'npulses']
         if not all(name in integParams for name in required_keys):
             raise TypeError('All keys of integParams argument '
                             f'{required_keys} are required.')
         params = integParams.copy()
         autoFind = False
         if pattern is not None:
-            # use period and npulses from pattern
-            pulse_ids = pattern.peek_pulse_ids(labelled=False)
-            npulses_from_bp = len(pulse_ids)
-            period_from_bp = 0
-            if npulses_from_bp > 1:
-                period_from_bp = min(np.diff(pulse_ids)) * pulse_period
-            if (npulses_from_bp != params['npulses'] or
-                period_from_bp != params['period']):
+            if (npulses != params.get('npulses') or
+                period != params.get('period')):
                 log.warning(f'Integration parameters '
-                    f'(npulses={params["npulses"]}, '
-                    f'period={params["period"]}) do not match the '
-                    f'the bunch pattern (npulses={npulses_from_bp}, '
-                    f'period={period_from_bp}). Using bunch pattern parameters.')
-                params['npulses'] = npulses_from_bp
-                params['period'] = period_from_bp
-        else:
-            period = params['period']
-            npulses = params['npulses']
-            pulse_ids = np.arange(params['npulses'], dtype=np.uint64)
-
-    elif pattern is not None:
-        if pattern.is_constant_pattern() is False:
-            print('The number of pulses changed during the run.')
-            pulse_ids = np.unique(pattern.pulse_ids(labelled=False, copy=False))
-            regular = False
+                    f'(npulses={params.get("npulses")}, '
+                    f'period={params.get("period")}) do not match the '
+                    f'the bunch pattern (npulses={npulses}, '
+                    f'period={period}). Using bunch pattern parameters.')
+            params['npulses'] = npulses
+            params['period'] = period
         else:
-            pulse_ids = pattern.peek_pulse_ids(labelled=False)
-        npulses = len(pulse_ids)
-        period = 0
-        if npulses > 1:
-            periods = np.diff(pulse_ids)
-            if len(np.unique(periods)) > 1:
-                regular = False
-            period = min(periods)
-            npulses = int((max(pulse_ids) - min(pulse_ids)) / period) + 1
-            period *= pulse_period
-    else:
-        pulse_ids = npulses = period = None
+            pulse_ids = np.arange(0, params['npulses'], params['period'],
+                                  dtype=np.uint64)
+        start = params['pulseStart']
+        params['pulseStart'] = [start + (pid - pulse_ids[0]) * pulse_period
+                                for pid in pulse_ids]
 
-    # generate average trace
-    traces = run[source, key].xarray(name=mnemonic.replace('raw', 'avg'),
-                                     extra_dims=mnemo['dim'])
-    trace = traces.mean('trainId')
+    # load traces and generate average
+    traces = run[source, key].xarray(name=mnemonic, extra_dims=mnemo['dim'])
+    trace = traces.mean('trainId').rename(mnemonic.replace('raw', 'avg'))
 
     # find peak integration parameters
     if autoFind == True:
@@ -1393,31 +1392,34 @@ def extract_digitizer_peaks(proposal, runNB, mnemonic, bunchPattern=None,
                   f'found integration parameters (npulses={params["npulses"]}, '
                   f'period={params["period"]}) do not match. Using bunch '
                   'pattern parameters.')
-            params['period'] = period
-            params['npulses'] = min(len(trace) // period, npulses)
+            params["npulses"] = npulses
+            params["period"] = period
         if pulse_ids is None:
             pulse_ids = np.arange(params['npulses'], dtype=np.uint64)
-
-        if params is None:
-            log.warning('Could not find peak integration parameters.')
-            return xr.DataArray()
+        start = params['pulseStart']
+        params['pulseStart'] = [start + (pid - pulse_ids[0]) * pulse_period
+                                for pid in pulse_ids]
 
     # extract peaks
     data = peaks_from_raw_trace(traces, **params, extra_dim=extra_dim)
     data = data.rename(mnemonic.replace('raw', 'peaks'))
-    data = data.assign_coords({extra_dim: pulse_ids})
 
-    if regular is False:
-        period = int(period / pulse_period)
+    if regular is True:
+        data = data.assign_coords({extra_dim: pulse_ids})
+    else:
         mask = pattern.pulse_mask(labelled=False)
         mask = xr.DataArray(mask, dims=['trainId', extra_dim],
                             coords={'trainId': run[source, key].train_ids,
                                     extra_dim: np.arange(mask.shape[1])})
-        mask = mask.sel({extra_dim: slice(pulse_ids[0],
-                                          pulse_ids[0] + npulses*period,period)})
+        mask = mask.sel({extra_dim: pulse_ids})
         data = data.where(mask, drop=True)
 
     data.attrs['params_keys'] = list(params.keys())
+    if regular:
+        params['pulseStart'] = params['pulseStart'][0]
+    for p in params:
+        if params[p] is None:
+            params[p] = 'None'
     data.attrs[f'params_{data.name}'] = list(params.values())
     if save:
         save_peaks(proposal, runNB, data, trace, subdir)