Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found

Target

Select target project
  • SCS/ToolBox
  • kluyvert/ToolBox
2 results
Show changes
Commits on Source (20)
Expand all Hide whitespace changes
Inline Side-by-side
doc/PES_spectra_extraction.ipynb
View file @ dc092fc6
This diff is collapsed.
doc/changelog.rst
View file @ dc092fc6
......@@ -23,6 +23,8 @@ unreleased
- drop MaranaX non-existing mnemonics and add attributes on dataset :mr:`300`
- streamline digitizer functions :mr:`304`
- add BAM average and BAM feedbacks mnemonics, Viking mnemonics :mr:`305`, :mr:`306`
- improved function to load PES spectra :mr:`309`
- Gotthard-II mnemonics and pulse alignment :mr:`310`, :mr:`311`
- **New Features**
......@@ -31,6 +33,7 @@ unreleased
- BOZ GPU notebook :mr:`294`
- update hRIXS class for MaranaX :mr:`287`
- introduce MaranaX class for parallelized centroiding :mr:`297`
- New PES functions to save / load average traces :mr:`309`
1.7.0
-----
......
src/toolbox_scs/constants.py
View file @ dc092fc6
......@@ -158,54 +158,54 @@ mnemonics = {
'dim': None},),
# PES
"PES_N_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_A.raw.samples',
'dim': ['PESsampleId']},),
"PES_NNE_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_B.raw.samples',
'dim': ['PESsampleId']},),
"PES_NE_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_C.raw.samples',
'dim': ['PESsampleId']},),
"PES_ENE_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_D.raw.samples',
'dim': ['PESsampleId']},),
"PES_E_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_A.raw.samples',
'dim': ['PESsampleId']},),
"PES_ESE_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_B.raw.samples',
'dim': ['PESsampleId']},),
"PES_SE_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_C.raw.samples',
'dim': ['PESsampleId']},),
"PES_SSE_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_D.raw.samples',
'dim': ['PESsampleId']},),
"PES_S_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_1Araw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_1_A.raw.samples',
'dim': ['PESsampleId']},),
"PES_SSW_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_1Braw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_1_B.raw.samples',
'dim': ['PESsampleId']},),
"PES_SW_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_1Craw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_1_C.raw.samples',
'dim': ['PESsampleId']},),
"PES_WSW_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_1Draw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_1_D.raw.samples',
'dim': ['PESsampleId']},),
"PES_W_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_2Araw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_2_A.raw.samples',
'dim': ['PESsampleId']},),
"PES_WNW_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_2Braw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_2_B.raw.samples',
'dim': ['PESsampleId']},),
"PES_NW_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_2Craw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_2_C.raw.samples',
'dim': ['PESsampleId']},),
"PES_NNW_raw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
"PES_2Draw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_2_D.raw.samples',
'dim': ['PESsampleId']},),
"PES_4Araw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_A.raw.samples',
'dim': ['PESsampleId']},),
"PES_4Braw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_B.raw.samples',
'dim': ['PESsampleId']},),
"PES_4Craw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_C.raw.samples',
'dim': ['PESsampleId']},),
"PES_4Draw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_4_D.raw.samples',
'dim': ['PESsampleId']},),
"PES_3Aaw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_A.raw.samples',
'dim': ['PESsampleId']},),
"PES_3Braw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_B.raw.samples',
'dim': ['PESsampleId']},),
"PES_3Craw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_C.raw.samples',
'dim': ['PESsampleId']},),
"PES_3Draw": ({'source': 'SA3_XTD10_PES/ADC/1:network',
'key': 'digitizers.channel_3_D.raw.samples',
'dim': ['PESsampleId']},),
"PES_pressure": ({'source': 'SA3_XTD10_PES/GAUGE/G30310F',
'key': 'value.value',
'dim': None},),
......@@ -417,6 +417,9 @@ mnemonics = {
{'source': 'SCS_ILH_LAS/MOTOR/LT3',
'key': 'AActualPosition.value',
'dim': None},),
"PP800_T0_mm": ({'source': 'SCS_ILH_LAS/MDL/OPTICALDELAY_PP800',
'key': 'motorOffset.value',
'dim': None},),
"PP800_HalfWP": ({'source': 'SCS_ILH_LAS/MOTOR/ROT8WP1',
'key': 'actualPosition.value',
'dim': None},),
......@@ -1077,6 +1080,14 @@ mnemonics = {
'key': 'data.rawData',
'dim': ['mcpbig_samplesId'],
'extract': 'peaks'},),
"XRD_PD2raw": ({'source': 'SCS_XRD_DET/ADC/DIODE2:output',
'key': 'data.rawData',
'dim': ['pd2_samplesId'],
'extract': 'peaks'},),
"XRD_PD6raw": ({'source': 'SCS_XRD_DET/ADC/DIODE6:output',
'key': 'data.rawData',
'dim': ['pd6_samplesId'],
'extract': 'peaks'},),
# KARABACON
"KARABACON": ({'source': 'SCS_DAQ_SCAN/MDL/KARABACON',
......@@ -1090,16 +1101,22 @@ mnemonics = {
"Gotthard2": ({'source': 'SCS_PAM_XOX/DET/GOTTHARD_RECEIVER2:daqOutput',
'key': 'data.adc',
'dim': ['gott_pId', 'pixelId']},),
"GH21": ({'source': 'SCS_XOX_GH21/CORR/GOTTHARD2_RECEIVER1:daqOutput',
"GH21": ({'source': 'SCS_XOX_GH21/CORR/RECEIVER:daqOutput',
'key': 'data.adc',
'dim': ['gh2_pId', 'pixelId']},
{'source': 'SCS_XOX_GH21/CORR/GOTTHARD2_RECEIVER1:daqOutput',
'key': 'data.adc',
'dim': ['gh2_pId', 'pixelId']},),
"GH22": ({'source': 'SCS_XOX_GH22/CORR/GOTTHARD2_RECEIVER1:daqOutput',
"GH22": ({'source': 'SCS_XOX_GH22/CORR/RECEIVER:daqOutput',
'key': 'data.adc',
'dim': ['gh2_pId', 'pixelId']},
{'source': 'SCS_XOX_GH22/CORR/GOTTHARD2_RECEIVER2:daqOutput',
'key': 'data.adc',
'dim': ['gh2_pId', 'pixelId']},),
"GH21_raw": ({'source': 'SCS_XOX_GH21/DET/GOTTHARD2_RECEIVER1:daqOutput',
'key': 'data.adc',
'dim': ['gh2_pId', 'pixelId']},),
"GH22_raw": ({'source': 'SCS_XOX_GH22/DET/GOTTHARD2_RECEIVER1:daqOutput',
"GH22_raw": ({'source': 'SCS_XOX_GH22/DET/GOTTHARD2_RECEIVER2:daqOutput',
'key': 'data.adc',
'dim': ['gh2_pId', 'pixelId']},),
......
src/toolbox_scs/detectors/__init__.py
View file @ dc092fc6
......@@ -5,6 +5,7 @@ from .dssc import *
from .dssc_data import *
from .dssc_misc import *
from .dssc_processing import *
from .gotthard2 import *
from .hrixs import *
from .pes import *
from .viking import *
......@@ -18,6 +19,7 @@ __all__ = (
+ dssc_data.__all__
+ dssc_misc.__all__
+ dssc_processing.__all__
+ gotthard2.__all__
+ hrixs.__all__
+ pes.__all__
+ viking.__all__
......
src/toolbox_scs/detectors/gotthard2.py 0 → 100644
View file @ dc092fc6
""" Gotthard-II detector related sub-routines
Copyright (2024) SCS Team.
(contributions preferrably comply with pep8 code structure
guidelines.)
"""
from extra.components import OpticalLaserPulses, XrayPulses
import numpy as np
import xarray as xr
import logging
__all__ = [
'extract_GH2',
]
log = logging.getLogger(__name__)
def extract_GH2(ds, run, firstFrame=0, bunchPattern='scs_ppl',
gh2_dim='gh2_pId'):
'''
Select and align the frames of the Gotthard-II that have been exposed
to light.
Parameters
------
ds: xarray.Dataset
The dataset containing GH2 data
run: extra_data.DataCollection
The run containing the bunch pattern source
firstFrame: int
The GH2 frame number corresponding to the first pulse of the train.
bunchPattern: str in ['scs_ppl', 'sase3']
the bunch pattern used to align data. For 'scs_ppl', the gh2_pId
dimension in renamed 'ol_pId', and for 'sase3' gh2_pId is renamed
'sa3_pId'.
gh2_dim: str
The name of the dimension that corresponds to the Gotthard-II frames.
Returns
-------
nds: xarray Dataset
The aligned and reduced dataset with only-data-containing GH2
variables.
'''
if gh2_dim not in ds.dims:
log.warning(f'gh2_dim "{gh2_dim}" not in dataset. Skipping.')
return ds
if bunchPattern == 'scs_ppl':
pattern = OpticalLaserPulses(run)
dim = 'ol_pId'
else:
pattern = XrayPulses(run)
dim = 'sa3_pId'
others = [var for var in ds if dim in ds[var].coords]
nds = ds.drop_dims(dim, errors='ignore')
if pattern.is_constant_pattern():
pulse_ids = pattern.peek_pulse_ids(labelled=False)
nds = nds.isel({gh2_dim: pulse_ids + firstFrame})
nds = nds.assign_coords({gh2_dim: pulse_ids})
nds = nds.rename({gh2_dim: dim})
else:
log.warning('The number of pulses has changed during the run.')
pulse_ids = np.unique(pattern.pulse_ids(labelled=False, copy=False))
nds = nds.isel({gh2_dim: pulse_ids + firstFrame})
nds = nds.assign_coords({gh2_dim: pulse_ids})
nds = nds.rename({gh2_dim: dim})
mask = pattern.pulse_mask(labelled=False)
mask = xr.DataArray(mask, dims=['trainId', dim],
coords={'trainId': run.train_ids,
dim: np.arange(mask.shape[1])})
mask = mask.sel({dim: pulse_ids})
nds = nds.where(mask, drop=True)
ret = ds[others].merge(nds, join='inner')
return ret
src/toolbox_scs/detectors/pes.py
View file @ dc092fc6
......@@ -11,54 +11,54 @@ import numpy as np
import xarray as xr
import extra_data as ed
import re
import os
from pathlib import Path
from multiprocessing import Pool
from extra.components import XrayPulses, AdqRawChannel
from ..misc.bunch_pattern_external import is_sase_3
from ..mnemonics_machinery import (mnemonics_to_process,
mnemonics_for_run)
from extra_data.read_machinery import find_proposal
from ..constants import mnemonics as _mnemonics
__all__ = [
'get_pes_params',
'get_pes_tof',
'save_pes_avg_traces',
'load_pes_avg_traces'
]
log = logging.getLogger(__name__)
def get_pes_tof(run, mnemonics=None, merge_with=None,
start=31390, width=300, origin=None, width_ns=None,
subtract_baseline=True,
baseStart=None, baseWidth=80,
sample_rate=2e9):
def get_pes_tof(proposal, runNB, mnemonic, start=0, origin=None,
width=None, subtract_baseline=False,
baseStart=None, baseWidth=40,merge_with=None):
"""
Extracts time-of-flight spectra from raw digitizer traces. The
tracesvare either loaded via ToolBox mnemonics or those in the
optionally provided merge_with dataset. The spectra are aligned
by pulse Id using the SASE 3 bunch pattern, and have time coordinates
in nanoseconds.
Extracts time-of-flight spectra from raw digitizer traces. The spectra
are aligned by pulse Id using the SASE 3 bunch pattern. If origin is
not None, a time coordinate in nanoseconds 'time_ns' is computed and
added to the DataArray.
Parameters
----------
run: extra_data.DataCollection
DataCollection containing the digitizer data
mnemonics: str or list of str
mnemonics for PES, e.g. "PES_W_raw" or ["PES_W_raw", "PES_ENE_raw"].
If None and no merge_with dataset is provided, defaults to "PES_W_raw".
merge_with: xarray Dataset
If provided, the resulting Dataset will be merged with this
one. The PES variables of merge_with (if any) will also be
computed and merged.
proposal:int
The proposal number.
runNB: int
The run number.
mnemonic: str
mnemonic for PES, e.g. "PES_2Araw".
start: int
starting sample of the first spectrum in the raw trace.
width: int
number of samples per spectra.
origin: int
sample of the raw trace that corresponds to time-of-flight origin.
If None, origin is equal to start.
width_ns: float
time window for one spectrum. If None, the time window is defined by
width / sample rate.
sample of the trace that corresponds to time-of-flight origin,
also called prompt. Used to compute the 'time_ns'
coordinates. If None, computation of 'time_ns' is skipped.
width: int
number of samples per spectra. If None, the number of samples
for 4.5 MHz repetition rate is used.
subtract_baseline: bool
If True, subtract baseline defined by baseStart and baseWidth to each
spectrum.
......@@ -67,102 +67,107 @@ def get_pes_tof(run, mnemonics=None, merge_with=None,
baseWidth: int
number of samples to average (starting from baseStart) for baseline
calculation.
sample_rate: float
sample rate of the digitizer.
merge_with: xarray Dataset
If provided, the resulting Dataset will be merged with this
one.
Returns
-------
pes: xarray Dataset
Dataset containing the PES time-of-flight spectra (e.g. "PES_W_tof"),
merged with optionally provided merg_with dataset.
pes: xarray DataArray
DataArray containing the PES time-of-flight spectra.
Example
-------
>>> import toolbox_scs as tb
>>> import toolbox_scs.detectors as tbdet
>>> run, ds = tb.load(2927, 100, "PES_W_raw")
>>> pes = tbdet.get_pes_tof(run, merge_with=ds)
>>> proposal, runNB = 900447, 12
>>> pes = tbdet.get_pes_tof(proposal, runNB, 'PES_2Araw',
>>> start=2557, origin=76)
"""
def to_processed_name(name):
return name.replace('raw', 'tof')
to_process = mnemonics_to_process(mnemonics, merge_with,
'PES', to_processed_name)
run_mnemonics = mnemonics_for_run(run)
# check if bunch pattern table exists
if bool(merge_with) and 'bunchPatternTable' in merge_with:
bpt = merge_with['bunchPatternTable']
elif 'bunchPatternTable' in run_mnemonics:
bpt = run.get_array(*run_mnemonics['bunchPatternTable'].values())
elif 'bunchPatternTable_SA3' in run_mnemonics:
bpt = run.get_array(*run_mnemonics['bunchPatternTable_SA3'].values())
run = ed.open_run(proposal, runNB)
all_mnemonics = mnemonics_for_run(run)
mnemo = all_mnemonics.get(mnemonic)
if mnemo is None:
print(f'Mnemonic {mnemonic} not found in run. Skipping.')
return None
PULSE_PERIOD = 440
pattern = XrayPulses(run)
regular = True
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
else:
bpt = None
mask = is_sase_3(bpt).assign_coords({'pulse_slot': np.arange(2700)})
mask_on = mask.where(mask, drop=True)
npulses = mask.sum(dim='pulse_slot')[0].values
pulse_ids = pattern.peek_pulse_ids(labelled=False)
npulses = len(pulse_ids)
period = 1
if npulses > 1:
period = mask_on['pulse_slot'].diff(dim='pulse_slot')[0].values
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
kd = run[mnemo['source'], mnemo['key']]
nsamples = kd.shape[1]
npulses_trace = int(np.floor((nsamples - start) / period))
if npulses_trace < npulses:
log.warning(f'The digitizer only recorded {npulses_trace} pulses '
f'out of the {npulses} pulses in the train.')
else:
period = 0
if origin is None:
origin = start
if baseStart is None:
baseStart = start
if width_ns is not None:
width = int(sample_rate * width_ns * 1e-9)
time_ns = 1e9 * (np.arange(start, start + width) - origin) / sample_rate
npulses_trace = npulses
indices = np.array([start + i*period for i in range(npulses_trace + 1)])
ds = xr.Dataset()
for m in to_process:
if bool(merge_with) and m in merge_with:
arr = merge_with[m]
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 '
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
end = begin + width
if end > arr.sizes['PESsampleId']:
break
pes = arr.isel(PESsampleId=slice(begin, end))
if subtract_baseline:
baseBegin = p*period*440 + baseStart
baseEnd = baseBegin + baseWidth
bl = arr.isel(
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.rename({'PESsampleId': 'time_ns'})
ds = ds.assign_coords({'time_ns': time_ns})
if bool(merge_with):
ds = merge_with.drop(to_process,
errors='ignore').merge(ds, join='left')
return ds
def get_pes_params(run):
outp = kd.ndarray()
outp = outp[:, indices[0]:indices[-1]].reshape([outp.shape[0],
npulses_trace, period])
spectra = xr.DataArray(outp, dims=['trainId', 'sa3_pId', 'sampleId'],
coords={'trainId': kd.train_id_coordinates(),
'sa3_pId': pulse_ids[:npulses_trace],
'sampleId': np.arange(period)})
if subtract_baseline:
if baseStart is None:
baseStart = 0
spectra = spectra - spectra.isel(
sampleId=slice(baseStart, baseStart + baseWidth)).mean('sampleId')
if width is None:
width = PULSE_PERIOD
spectra = spectra.isel(sampleId=slice(0, width), drop=True)
spectra.attrs['start'] = start
if origin is not None:
try:
adq = AdqRawChannel(run, mnemonic.split('_')[1].split('raw')[0])
sample_rate = adq.sampling_rate
except Exception as e:
log.warning(e)
sample_rate = 1986111111.1111112
origin -= start
time_ns = (np.arange(spectra.sizes['sampleId'])
- origin)/sample_rate * 1e9
spectra = spectra.assign_coords(time_ns=('sampleId', time_ns))
spectra.attrs['origin'] = origin
spectra = spectra.rename(mnemonic.replace('raw', 'spectrum'))
if merge_with is not None:
return merge_with.merge(spectra.to_dataset(promote_attrs=True),
join='inner')
return spectra.rename(mnemonic.replace('raw', 'spectrum'))
def get_pes_params(run, channel=None):
"""
Extract PES parameters for a given extra_data DataCollection.
Parameters are gas, binding energy, voltages of the MPOD.
Parameters are gas, binding energy, retardation voltages or all
voltages of the MPOD.
Parameters
----------
run: extra_data.DataCollection
DataCollection containing the digitizer data
channel: str
Channel name or PES mnemonic, e.g. '2A' or 'PES_1Craw'.
If None, or if the channel is not found in the data,
the retardation voltage for all channels is retrieved.
Returns
-------
params: dict
......@@ -181,11 +186,41 @@ def get_pes_params(run):
if 'gas' not in params:
params['gas'] = 'unknown'
log.warning('Could not find which PES gas was used.')
voltages = get_pes_voltages(run)
params.update(voltages)
mpod_mapper = 'SA3_XTD10_PES/MDL/MPOD_MAPPER'
if mpod_mapper in run.all_sources:
if channel is None:
channel = [f'{a//4 + 1}{b}' for a, b in enumerate(
['A', 'B', 'C', 'D']*4)]
else:
if 'raw' in channel:
channel = channel.split('raw')[0].split('_')[1]
channel = [channel]
for c in channel:
rv = get_pes_rv(run, c, mpod_mapper)
if rv is not None:
params[f'{c}_RV'] = rv
elif 'SA3_XTD10_PES/MDL/DAQ_MPOD' in run.all_sources:
voltages = get_pes_voltages(run)
if voltages is not None:
params.update(voltages)
return params
def get_pes_rv(run, channel, mpod_mapper='SA3_XTD10_PES/MDL/MPOD_MAPPER'):
channel_to_group = {'4D': 1, '4B': 1, '4C': 1, '4A': 1,
'2A': 2, '1C': 2, '2C': 2, '2D': 2,
'3D': 3, '3B': 3, '2B': 3, '3A': 3,
'1A': 4, '3C': 4, '1B': 4, '1D': 4}
group = channel_to_group.get(channel)
if group is None:
log.warning(f'Could not find group for channel {channel}.')
return None
key = f'Group{group}D.channelMeasurementSenseVoltage.value'
voltage = run[mpod_mapper, key].ndarray().mean()
return voltage
def get_pes_voltages(run, device='SA3_XTD10_PES/MDL/DAQ_MPOD'):
"""
Extract PES voltages read by the MDL watchdog of the MPOD device.
......@@ -200,12 +235,120 @@ def get_pes_voltages(run, device='SA3_XTD10_PES/MDL/DAQ_MPOD'):
Returns
-------
voltages: dict
dictionnary of voltages
dictionnary of voltages, empty if device not found.
"""
a = re.compile('[u]\d{3}.value')
if device not in run.all_sources:
log.warning(f'Could not find {device} in run. Skipping.')
return None
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
def calculate_average(run, name, mnemo):
return run[mnemo['source'], mnemo['key']].xarray(
name=name, extra_dims=mnemo['dim']).mean('trainId')
def save_pes_avg_traces(proposal, runNB, channels=None,
subdir='usr/processed_runs'):
'''
Save average traces of PES into an h5 file.
Parameters
----------
proposal:int
The proposal number.
runNB: int
The run number.
channels: str or list
The PES channels or mnemonics, e.g. '2A', ['2A', '3C'],
['PES_1Araw', 'PES_4Draw', '3B']
subdir: str
subdirectory. The data is stored in
<proposal path>/<subdir>/r{runNB:04d}/f'r{runNB:04d}-pes-data.h5'
Output
------
xarray Dataset saved in a h5 file containing the PES average traces.
'''
root = find_proposal(f'p{proposal:06d}')
path = os.path.join(root, subdir + f'/r{runNB:04d}/')
Path(path).mkdir(parents=True, exist_ok=True)
fname = path + f'r{runNB:04d}-pes-data.h5'
if channels is None:
channels = [f'{a//4 + 1}{b}' for a, b in enumerate(
['A', 'B', 'C', 'D']*4)]
if isinstance(channels, str):
channels = [channels]
run = ed.open_run(proposal, runNB, parallelize=False)
all_mnemos = mnemonics_for_run(run)
# use multiprocessing.Pool
args = []
for c in channels:
m = c
if 'raw' not in c:
m = f'PES_{c}raw'
if m not in all_mnemos:
continue
args.append([run, m.replace('raw', 'avg'), all_mnemos[m]])
if len(args) == 0:
log.warning('No pes average trace to save. Skipping')
return
with Pool(len(args)) as pool:
avg_traces = pool.starmap(calculate_average, args)
avg_traces = xr.merge(avg_traces)
ds = xr.Dataset()
if os.path.isfile(fname):
ds = xr.load_dataset(fname)
ds = ds.drop_vars(channels, errors='ignore')
ds = ds.merge(avg_traces)
ds.to_netcdf(fname, format='NETCDF4')
return
def load_pes_avg_traces(proposal, runNB, channels=None,
subdir='usr/processed_runs'):
'''
Load existing PES average traces.
Parameters
----------
proposal:int
The proposal number.
runNB: int
The run number.
channels: str or list
The PES channels or mnemonics, e.g. '2A', ['2A', '3C'],
['PES_1Araw', 'PES_4Draw', '3B']
subdir: str
subdirectory. The data is stored in
<proposal path>/<subdir>/r{runNB:04d}/f'r{runNB:04d}-pes-data.h5'
Output
------
ds: xarray Dataset
dataset containing the PES average traces.
'''
root = find_proposal(f'p{proposal:06d}')
path = os.path.join(root, subdir + f'/r{runNB:04d}/')
fname = path + f'r{runNB:04d}-pes-data.h5'
if channels is None:
channels = [f'PES_{a//4 + 1}{b}avg' for a, b in
enumerate(['A', 'B', 'C', 'D']*4)]
if isinstance(channels, str):
channels = [channels]
for i, c in enumerate(channels):
if 'PES_' not in c and 'avg' not in c:
channels[i] = f'PES_{c}avg'
if os.path.isfile(fname):
ds = xr.load_dataset(fname)
channels = [c for c in channels if c in ds]
ds = ds[channels]
return ds
else:
log.warning(f'{fname} is not a valid file.')

Legal Notice | Contacts