Adds pes_get_voltages(), updates pes_get_params()

src/toolbox_scs/detectors/pes.py

@@ -10,6 +10,7 @@ import logging
 import numpy as np
 import xarray as xr
 import extra_data as ed
+import re
 
 from ..misc.bunch_pattern_external import is_sase_3
 from ..mnemonics_machinery import (mnemonics_to_process,
@@ -119,10 +120,10 @@ def get_pes_tof(run, mnemonics=None, merge_with=None,
         else:
             arr = run.get_array(*run_mnemonics[m].values(), name=m)
         if arr.sizes['PESsampleId'] < npulses*period*440 + start + width:
-            log.warning('Not all pulses were recorded. The number of samples on '
-                        f'the digitizer {arr.sizes["PESsampleId"]} is not enough '
-                        f'to cover the {npulses} spectra. Missing spectra will be '
-                       'filled with NaNs.')
+            log.warning('Not all pulses were recorded. The number of samples '
+                        f'on the digitizer {arr.sizes["PESsampleId"]} is not '
+                        f'enough to cover the {npulses} spectra. Missing '
+                        'spectra will be filled with NaNs.')
         spectra = []
         for p in range(npulses):
             begin = p*period*440 + start
@@ -134,14 +135,15 @@ def get_pes_tof(run, mnemonics=None, merge_with=None,
                 baseBegin = p*period*440 + baseStart
                 baseEnd = baseBegin + baseWidth
                 bl = arr.isel(
-                    PESsampleId=slice(baseBegin, baseEnd)).mean(dim='PESsampleId')
+                 PESsampleId=slice(baseBegin, baseEnd)).mean(dim='PESsampleId')
                 pes = pes - bl
             spectra.append(pes)
         spectra = xr.concat(spectra,
                             dim='sa3_pId').rename(m.replace('raw', 'tof'))
         ds = ds.merge(spectra)
     if len(ds.variables) > 0:
-        ds = ds.assign_coords({'sa3_pId': mask_on['pulse_slot'][:ds.sizes['sa3_pId']].values})
+        ds = ds.assign_coords(
+             {'sa3_pId': mask_on['pulse_slot'][:ds.sizes['sa3_pId']].values})
         ds = ds.rename({'PESsampleId': 'time_ns'})
         ds = ds.assign_coords({'time_ns': time_ns})
     if bool(merge_with):
@@ -154,7 +156,7 @@ def get_pes_tof(run, mnemonics=None, merge_with=None,
 def get_pes_params(run):
     """
     Extract PES parameters for a given extra_data DataCollection.
-    Parameters are gas, retardation voltage.
+    Parameters are gas, binding energy, voltages of the MPOD.
 
     Parameters
     ----------
@@ -168,7 +170,6 @@ def get_pes_params(run):
     """
     params = {}
     sel = run.select_trains(ed.by_index[:20])
-    mnemonics = mnemonics_for_run(run)
     gas_dict = {'N2': 409.9, 'Ne': 870.2, 'Kr': 1921, 'Xe': 1148.7}
     for gas in gas_dict.keys():
         mnemo = _mnemonics[f'PES_{gas}'][0]
@@ -180,6 +181,31 @@ def get_pes_params(run):
     if 'gas' not in params:
         params['gas'] = 'unknown'
         log.warning('Could not find which PES gas was used.')
-    arr = sel.get_array(*mnemonics['PES_RV'].values())
-    params['ret_voltage'] = float(arr[0].values)
+    voltages = get_pes_voltages(run)
+    params.update(voltages)
     return params
+
+
+def get_pes_voltages(run, device='SA3_XTD10_PES/MDL/DAQ_MPOD'):
+    """
+    Extract PES voltages read by the MDL watchdog of the MPOD device.
+
+    Parameters
+    ----------
+    run: extra_data.DataCollection
+        DataCollection containing PES data.
+    device: string
+        Name of the device containing the voltage data.
+
+    Returns
+    -------
+    voltages: dict
+        dictionnary of voltages
+    """
+    a = re.compile('[u]\d{3}.value')
+    tid, da = run.train_from_index(0, devices=device)
+    voltages = {}
+    for k in da[device]:
+        if len(a.findall(k)) == 1:
+            voltages[k.split('.')[0]] = da[device][k]
+    return voltages