pytesprocess 0.1.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pytesprocess/__init__.py +9 -0
- pytesprocess/_version.py +2 -0
- pytesprocess/cli/__init__.py +1 -0
- pytesprocess/cli/commands/__init__.py +5 -0
- pytesprocess/cli/commands/event.py +66 -0
- pytesprocess/cli/commands/filter.py +17 -0
- pytesprocess/cli/commands/ivsweep.py +29 -0
- pytesprocess/cli/common.py +86 -0
- pytesprocess/cli/main.py +81 -0
- pytesprocess/config/__init__.py +4 -0
- pytesprocess/config/loader.py +94 -0
- pytesprocess/config/manager.py +297 -0
- pytesprocess/config/resolvers/__init__.py +5 -0
- pytesprocess/config/resolvers/common.py +56 -0
- pytesprocess/config/resolvers/feature.py +293 -0
- pytesprocess/config/resolvers/salting.py +86 -0
- pytesprocess/config/resolvers/trigger.py +84 -0
- pytesprocess/config/selectors.py +108 -0
- pytesprocess/config/validation.py +314 -0
- pytesprocess/config/warnings.py +2 -0
- pytesprocess/core/__init__.py +10 -0
- pytesprocess/core/algorithms.py +1455 -0
- pytesprocess/core/didv.py +1648 -0
- pytesprocess/core/eventbuilder.py +495 -0
- pytesprocess/core/filterbuilder.py +81 -0
- pytesprocess/core/filterdata.py +1849 -0
- pytesprocess/core/ivsweep.py +2072 -0
- pytesprocess/core/noise.py +923 -0
- pytesprocess/core/noisemodel.py +1408 -0
- pytesprocess/core/oftrigger.py +1035 -0
- pytesprocess/core/template.py +450 -0
- pytesprocess/process/__init__.py +6 -0
- pytesprocess/process/data_source.py +185 -0
- pytesprocess/process/event_context.py +35 -0
- pytesprocess/process/feature_plan.py +186 -0
- pytesprocess/process/feature_resources.py +267 -0
- pytesprocess/process/features.py +1024 -0
- pytesprocess/process/filterprocess.py +1176 -0
- pytesprocess/process/ivprocess.py +1380 -0
- pytesprocess/process/processing_data.py +967 -0
- pytesprocess/process/randoms.py +921 -0
- pytesprocess/process/triggers.py +1011 -0
- pytesprocess/salting/__init__.py +7 -0
- pytesprocess/salting/generator.py +364 -0
- pytesprocess/salting/injector.py +329 -0
- pytesprocess/salting/sampling.py +84 -0
- pytesprocess/utils/__init__.py +5 -0
- pytesprocess/utils/arg_utils.py +122 -0
- pytesprocess/utils/dataframe_output.py +120 -0
- pytesprocess/utils/filter_hdf5.py +594 -0
- pytesprocess/utils/utils.py +701 -0
- pytesprocess/workflows/__init__.py +3 -0
- pytesprocess/workflows/processing.py +317 -0
- pytesprocess/workflows/salting.py +133 -0
- pytesprocess-0.1.1.dist-info/METADATA +211 -0
- pytesprocess-0.1.1.dist-info/RECORD +60 -0
- pytesprocess-0.1.1.dist-info/WHEEL +5 -0
- pytesprocess-0.1.1.dist-info/entry_points.txt +2 -0
- pytesprocess-0.1.1.dist-info/licenses/LICENSE +21 -0
- pytesprocess-0.1.1.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,450 @@
|
|
|
1
|
+
import copy
|
|
2
|
+
|
|
3
|
+
import numpy as np
|
|
4
|
+
import qetpy as qp
|
|
5
|
+
from qetpy.core._template import Template as QETpyTemplate
|
|
6
|
+
from pytesdaqx.io import StreamReader
|
|
7
|
+
|
|
8
|
+
from pytesprocess.core.filterdata import FilterData
|
|
9
|
+
from pytesprocess.utils import convert_length_msec_to_samples
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class Template(FilterData):
|
|
13
|
+
"""
|
|
14
|
+
pytesprocess wrapper around qetpy.core._template.Template.
|
|
15
|
+
|
|
16
|
+
pytesprocess handles trace I/O and FilterData persistence.
|
|
17
|
+
QETpy handles the array-based template calculations and fitting.
|
|
18
|
+
"""
|
|
19
|
+
|
|
20
|
+
def __init__(self, verbose=True, filter_data=None):
|
|
21
|
+
super().__init__(verbose=verbose, filter_data=filter_data)
|
|
22
|
+
self._qp_template = QETpyTemplate(verbose=verbose)
|
|
23
|
+
|
|
24
|
+
def clear(self, channels=None, clear_filter_data=False, tag=None):
|
|
25
|
+
self._qp_template.clear(channels=channels)
|
|
26
|
+
if clear_filter_data:
|
|
27
|
+
self.clear_data(channels=channels, tag=tag)
|
|
28
|
+
|
|
29
|
+
def _normalize_channels(self, channels):
|
|
30
|
+
if isinstance(channels, str):
|
|
31
|
+
channels = [channels]
|
|
32
|
+
elif isinstance(channels, (list, tuple, np.ndarray)):
|
|
33
|
+
channels = list(channels)
|
|
34
|
+
else:
|
|
35
|
+
raise ValueError('ERROR: "channels" should be a string or a list of strings!')
|
|
36
|
+
if not channels:
|
|
37
|
+
raise ValueError('ERROR: "channels" cannot be empty!')
|
|
38
|
+
if not all(isinstance(chan, str) for chan in channels):
|
|
39
|
+
raise ValueError('ERROR: all channel names should be strings!')
|
|
40
|
+
return channels
|
|
41
|
+
|
|
42
|
+
def _get_trace_axis_info(self, metadata, trace_length_samples=None,
|
|
43
|
+
pretrigger_length_msec=None,
|
|
44
|
+
pretrigger_length_samples=None):
|
|
45
|
+
fs = float(metadata[0]['sample_rate_hz'])
|
|
46
|
+
nsamples = int(trace_length_samples) if trace_length_samples is not None else None
|
|
47
|
+
if nsamples is None:
|
|
48
|
+
if 'nb_samples' in metadata[0]:
|
|
49
|
+
nsamples = int(metadata[0]['nb_samples'])
|
|
50
|
+
else:
|
|
51
|
+
raise ValueError('ERROR: unable to determine trace length in samples.')
|
|
52
|
+
|
|
53
|
+
trigger_index = nsamples // 2
|
|
54
|
+
if pretrigger_length_msec is not None:
|
|
55
|
+
trigger_index = convert_length_msec_to_samples(pretrigger_length_msec, fs) - 1
|
|
56
|
+
elif pretrigger_length_samples is not None:
|
|
57
|
+
trigger_index = int(pretrigger_length_samples) - 1
|
|
58
|
+
|
|
59
|
+
t = np.arange(nsamples, dtype=np.float64) / fs
|
|
60
|
+
freqs = np.fft.fftfreq(nsamples, d=1.0 / fs)
|
|
61
|
+
return fs, nsamples, trigger_index, t, freqs
|
|
62
|
+
|
|
63
|
+
def _store_template(self, channels, template, sample_rate,
|
|
64
|
+
pretrigger_length_samples, tag='default', metadata=None):
|
|
65
|
+
template = np.asarray(template)
|
|
66
|
+
if template.ndim == 1:
|
|
67
|
+
if len(channels) != 1:
|
|
68
|
+
raise ValueError('ERROR: 1D template provided for multiple channels.')
|
|
69
|
+
self.set_template(
|
|
70
|
+
channels=channels[0],
|
|
71
|
+
template=template,
|
|
72
|
+
sample_rate=sample_rate,
|
|
73
|
+
pretrigger_length_samples=pretrigger_length_samples,
|
|
74
|
+
metadata=metadata,
|
|
75
|
+
tag=tag,
|
|
76
|
+
)
|
|
77
|
+
elif template.ndim == 2:
|
|
78
|
+
if template.shape[0] != len(channels):
|
|
79
|
+
raise ValueError('ERROR: template channel dimension is inconsistent with channels.')
|
|
80
|
+
for ichan, chan in enumerate(channels):
|
|
81
|
+
chan_metadata = copy.deepcopy(metadata) if metadata is not None else {}
|
|
82
|
+
self.set_template(
|
|
83
|
+
channels=chan,
|
|
84
|
+
template=template[ichan],
|
|
85
|
+
sample_rate=sample_rate,
|
|
86
|
+
pretrigger_length_samples=pretrigger_length_samples,
|
|
87
|
+
metadata=chan_metadata,
|
|
88
|
+
tag=tag,
|
|
89
|
+
)
|
|
90
|
+
else:
|
|
91
|
+
raise ValueError('ERROR: template should be 1D or 2D [nchans, nsamples].')
|
|
92
|
+
|
|
93
|
+
def _load_dpdi_from_store(self, channels, poles, dpdi_tag):
|
|
94
|
+
if poles is None:
|
|
95
|
+
raise ValueError(
|
|
96
|
+
'ERROR: dpdi=None and no internal dpdi available. '
|
|
97
|
+
'Need "poles" (and optionally "dpdi_tag") to load stored dpdi.'
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
dpdi_dict = {}
|
|
101
|
+
dpdi_freqs_dict = {}
|
|
102
|
+
metadata_dict = {}
|
|
103
|
+
for chan in channels:
|
|
104
|
+
dpdi_vals, dpdi_freqs = self.get_dpdi(chan, poles=poles, tag=dpdi_tag)
|
|
105
|
+
dpdi_dict[chan] = np.asarray(dpdi_vals, dtype=np.complex128)
|
|
106
|
+
dpdi_freqs_dict[chan] = np.asarray(dpdi_freqs, dtype=np.float64)
|
|
107
|
+
metadata_dict[chan] = {'poles': poles, 'tag': dpdi_tag}
|
|
108
|
+
return dpdi_dict, dpdi_freqs_dict, metadata_dict
|
|
109
|
+
|
|
110
|
+
def calc_average_pulses(self, channels, file_path, event_list,
|
|
111
|
+
trace_length_msec=None,
|
|
112
|
+
pretrigger_length_msec=None,
|
|
113
|
+
trace_length_samples=None,
|
|
114
|
+
pretrigger_length_samples=None,
|
|
115
|
+
nevents=2000,
|
|
116
|
+
lgc_plot=False,
|
|
117
|
+
lgc_filter_freq=True,
|
|
118
|
+
filter_freq=50e3,
|
|
119
|
+
time_lims=None,
|
|
120
|
+
store_filterdata=True,
|
|
121
|
+
tag='default_mean_current'):
|
|
122
|
+
channels = self._normalize_channels(channels)
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
# get traces
|
|
126
|
+
reader = StreamReader(file_path)
|
|
127
|
+
traces, metadata = reader.read_records(
|
|
128
|
+
record_list=event_list,
|
|
129
|
+
n_records=nevents,
|
|
130
|
+
channels=channels,
|
|
131
|
+
trace_length_msec=trace_length_msec,
|
|
132
|
+
trace_length_samples=trace_length_samples,
|
|
133
|
+
pretrigger_length_msec=pretrigger_length_msec,
|
|
134
|
+
pretrigger_length_samples=pretrigger_length_samples,
|
|
135
|
+
include_metadata=True,
|
|
136
|
+
units='amps',
|
|
137
|
+
stack=True)
|
|
138
|
+
|
|
139
|
+
detector_config = reader.get_detector_settings()
|
|
140
|
+
|
|
141
|
+
if traces.ndim != 3:
|
|
142
|
+
raise ValueError(
|
|
143
|
+
'ERROR: expected traces array with shape [nevents, nchans, nsamples]. '
|
|
144
|
+
f'Got shape {traces.shape}.'
|
|
145
|
+
)
|
|
146
|
+
if traces.shape[1] != len(channels):
|
|
147
|
+
raise ValueError(
|
|
148
|
+
'ERROR: traces shape is inconsistent with the requested channels. '
|
|
149
|
+
f'Expected {len(channels)} channels, got {traces.shape[1]}.'
|
|
150
|
+
)
|
|
151
|
+
|
|
152
|
+
fs, _, trigger_index, _, _ = self._get_trace_axis_info(
|
|
153
|
+
metadata,
|
|
154
|
+
trace_length_samples=traces.shape[-1],
|
|
155
|
+
pretrigger_length_msec=pretrigger_length_msec,
|
|
156
|
+
pretrigger_length_samples=pretrigger_length_samples,
|
|
157
|
+
)
|
|
158
|
+
|
|
159
|
+
self._qp_template.calc_average_pulses(
|
|
160
|
+
traces=traces,
|
|
161
|
+
channels=channels,
|
|
162
|
+
sample_rate=fs,
|
|
163
|
+
trigger_index=trigger_index,
|
|
164
|
+
lgc_plot=lgc_plot,
|
|
165
|
+
lgc_filter_freq=lgc_filter_freq,
|
|
166
|
+
filter_freq=filter_freq,
|
|
167
|
+
time_lims=time_lims,
|
|
168
|
+
)
|
|
169
|
+
|
|
170
|
+
if store_filterdata:
|
|
171
|
+
mean_i_t, _ = self._qp_template.get_template_in_current(
|
|
172
|
+
channels, use_fit=False, domain='time'
|
|
173
|
+
)
|
|
174
|
+
metadata_out = {
|
|
175
|
+
'template_source': 'average_current_pulse',
|
|
176
|
+
'template_domain': 'current',
|
|
177
|
+
'template_kind': 'mean',
|
|
178
|
+
'nb_events': traces.shape[0],
|
|
179
|
+
}
|
|
180
|
+
self._store_template(
|
|
181
|
+
channels=channels,
|
|
182
|
+
template=mean_i_t,
|
|
183
|
+
sample_rate=fs,
|
|
184
|
+
pretrigger_length_samples=trigger_index,
|
|
185
|
+
tag=tag,
|
|
186
|
+
metadata=metadata_out,
|
|
187
|
+
)
|
|
188
|
+
|
|
189
|
+
def calc_power_template(self, channels, dpdi=None, dpdi_err=None,
|
|
190
|
+
dpdi_freqs=None, poles=None, dpdi_tag='default',
|
|
191
|
+
dpdi_metadata=None, lgc_plot=False,
|
|
192
|
+
filter_freq=50e3, time_lims=None,
|
|
193
|
+
store_filterdata=True, tag='default_mean_power'):
|
|
194
|
+
channels = self._normalize_channels(channels)
|
|
195
|
+
|
|
196
|
+
if dpdi is None:
|
|
197
|
+
missing = [chan for chan in channels if chan not in self._qp_template._dpdi]
|
|
198
|
+
if missing:
|
|
199
|
+
loaded_dpdi, loaded_freqs, loaded_meta = self._load_dpdi_from_store(channels, poles, dpdi_tag)
|
|
200
|
+
dpdi_dict = loaded_dpdi
|
|
201
|
+
dpdi_freqs_dict = loaded_freqs
|
|
202
|
+
if dpdi_metadata is None:
|
|
203
|
+
dpdi_metadata = loaded_meta
|
|
204
|
+
else:
|
|
205
|
+
dpdi_dict = {chan: self._qp_template._dpdi[chan] for chan in channels}
|
|
206
|
+
dpdi_freqs_dict = {chan: self._qp_template._dpdi_freqs.get(chan, self._qp_template._freqs[chan])
|
|
207
|
+
for chan in channels}
|
|
208
|
+
else:
|
|
209
|
+
dpdi_dict = self._qp_template._prepare_channel_dict(dpdi, channels, dtype=np.complex128, name='dpdi')
|
|
210
|
+
dpdi_freqs_dict = None if dpdi_freqs is None else self._qp_template._prepare_channel_dict(
|
|
211
|
+
dpdi_freqs, channels, dtype=np.float64, name='dpdi_freqs'
|
|
212
|
+
)
|
|
213
|
+
|
|
214
|
+
if dpdi_err is None:
|
|
215
|
+
dpdi_err_dict = {chan: self._qp_template._dpdi_err.get(
|
|
216
|
+
chan, np.zeros_like(dpdi_dict[chan], dtype=np.complex128)
|
|
217
|
+
) for chan in channels}
|
|
218
|
+
else:
|
|
219
|
+
dpdi_err_dict = self._qp_template._prepare_channel_dict(dpdi_err, channels, dtype=np.complex128, name='dpdi_err')
|
|
220
|
+
|
|
221
|
+
self._qp_template.calc_power_template(
|
|
222
|
+
channels=channels,
|
|
223
|
+
dpdi=dpdi_dict,
|
|
224
|
+
dpdi_err=dpdi_err_dict,
|
|
225
|
+
dpdi_freqs=dpdi_freqs_dict,
|
|
226
|
+
dpdi_metadata=dpdi_metadata,
|
|
227
|
+
lgc_plot=lgc_plot,
|
|
228
|
+
filter_freq=filter_freq,
|
|
229
|
+
time_lims=time_lims,
|
|
230
|
+
)
|
|
231
|
+
|
|
232
|
+
if store_filterdata:
|
|
233
|
+
mean_p_t, _ = self._qp_template.get_template_in_power(
|
|
234
|
+
channels, use_fit=False, domain='time'
|
|
235
|
+
)
|
|
236
|
+
fs = self._qp_template._sample_rate[channels[0]]
|
|
237
|
+
pretrigger_samples = self._qp_template._pretrigger_samples[channels[0]]
|
|
238
|
+
metadata_out = {
|
|
239
|
+
'template_source': 'average_power_pulse',
|
|
240
|
+
'template_domain': 'power',
|
|
241
|
+
'template_kind': 'mean',
|
|
242
|
+
}
|
|
243
|
+
self._store_template(
|
|
244
|
+
channels=channels,
|
|
245
|
+
template=mean_p_t,
|
|
246
|
+
sample_rate=fs,
|
|
247
|
+
pretrigger_length_samples=pretrigger_samples,
|
|
248
|
+
tag=tag,
|
|
249
|
+
metadata=metadata_out,
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
def fit_templates(self, channels, template_model='twopole',
|
|
253
|
+
guess=None, bounds=None, f_fit_cutoff=50e3,
|
|
254
|
+
max_nfev=800, dt=0.0, lgc_diagnostics=False,
|
|
255
|
+
lgc_plot=True, filter_freq=50e3,
|
|
256
|
+
time_lims=None, store_filterdata=True,
|
|
257
|
+
tag='default'):
|
|
258
|
+
channels = self._normalize_channels(channels)
|
|
259
|
+
fit_result = self._qp_template.fit_templates(
|
|
260
|
+
channels=channels,
|
|
261
|
+
template_model=template_model,
|
|
262
|
+
guess=guess,
|
|
263
|
+
bounds=bounds,
|
|
264
|
+
dt=dt,
|
|
265
|
+
f_fit_cutoff=f_fit_cutoff,
|
|
266
|
+
max_nfev=max_nfev,
|
|
267
|
+
lgc_diagnostics=lgc_diagnostics,
|
|
268
|
+
lgc_plot=lgc_plot,
|
|
269
|
+
filter_freq=filter_freq,
|
|
270
|
+
time_lims=time_lims,
|
|
271
|
+
)
|
|
272
|
+
|
|
273
|
+
if store_filterdata:
|
|
274
|
+
fit_i_t, _ = self._qp_template.get_template_in_current(
|
|
275
|
+
channels, use_fit=True, domain='time'
|
|
276
|
+
)
|
|
277
|
+
fs = self._qp_template._sample_rate[channels[0]]
|
|
278
|
+
pretrigger_samples = self._qp_template._pretrigger_samples[channels[0]]
|
|
279
|
+
metadata_out = {
|
|
280
|
+
'template_source': 'fitted_current_template',
|
|
281
|
+
'template_domain': 'current',
|
|
282
|
+
'template_kind': 'fit',
|
|
283
|
+
'template_model': template_model,
|
|
284
|
+
}
|
|
285
|
+
self._store_template(
|
|
286
|
+
channels=channels,
|
|
287
|
+
template=fit_i_t,
|
|
288
|
+
sample_rate=fs,
|
|
289
|
+
pretrigger_length_samples=pretrigger_samples,
|
|
290
|
+
tag=tag,
|
|
291
|
+
metadata=metadata_out,
|
|
292
|
+
)
|
|
293
|
+
|
|
294
|
+
return fit_result
|
|
295
|
+
|
|
296
|
+
def get_template_in_current(self, channels, use_fit=True, domain='time', return_metadata=False):
|
|
297
|
+
channels = self._normalize_channels(channels)
|
|
298
|
+
return self._qp_template.get_template_in_current(
|
|
299
|
+
channels=channels,
|
|
300
|
+
use_fit=use_fit,
|
|
301
|
+
domain=domain,
|
|
302
|
+
return_metadata=return_metadata,
|
|
303
|
+
)
|
|
304
|
+
|
|
305
|
+
def get_template_in_power(self, channels, dpdi=None, dpdi_freqs=None,
|
|
306
|
+
use_fit=True, domain='time', return_metadata=False):
|
|
307
|
+
channels = self._normalize_channels(channels)
|
|
308
|
+
return self._qp_template.get_template_in_power(
|
|
309
|
+
channels=channels,
|
|
310
|
+
dpdi=dpdi,
|
|
311
|
+
dpdi_freqs=dpdi_freqs,
|
|
312
|
+
use_fit=use_fit,
|
|
313
|
+
domain=domain,
|
|
314
|
+
return_metadata=return_metadata
|
|
315
|
+
)
|
|
316
|
+
|
|
317
|
+
@property
|
|
318
|
+
def qp_template(self):
|
|
319
|
+
return self._qp_template
|
|
320
|
+
|
|
321
|
+
# ------------------------------------------------------------------
|
|
322
|
+
# Original template creation helpers retained
|
|
323
|
+
# ------------------------------------------------------------------
|
|
324
|
+
def create_template(self, channels,
|
|
325
|
+
sample_rate=None,
|
|
326
|
+
trace_length_msec=None,
|
|
327
|
+
trace_length_samples=None,
|
|
328
|
+
pretrigger_length_msec=None,
|
|
329
|
+
pretrigger_length_samples=None,
|
|
330
|
+
A=1, B=None, C=None,
|
|
331
|
+
tau_r=None,
|
|
332
|
+
tau_f1=None, tau_f2=None, tau_f3=None,
|
|
333
|
+
tag='default'):
|
|
334
|
+
if sample_rate is None:
|
|
335
|
+
raise ValueError('ERROR: "sample_rate" argument required')
|
|
336
|
+
if trace_length_msec is None and trace_length_samples is None:
|
|
337
|
+
raise ValueError('ERROR: Trace length required ("trace_length_msec" or "trace_length_samples")!')
|
|
338
|
+
if pretrigger_length_msec is None and pretrigger_length_samples is None:
|
|
339
|
+
raise ValueError('ERROR: Pretrigger length required ("pretrigger_length_msec" or "pretrigger_length_samples")!')
|
|
340
|
+
if tau_r is None:
|
|
341
|
+
raise ValueError('ERROR: "tau_r" argument required')
|
|
342
|
+
if A is None and B is None and C is None:
|
|
343
|
+
raise ValueError('ERROR: "A" and/or "B" and/or "C" argument(s) required!')
|
|
344
|
+
if tau_f1 is None and tau_f2 is None and tau_f3 is None:
|
|
345
|
+
raise ValueError('ERROR: "tau_f1" and/or "tau_f2" and/or "tau_f3" argument(s) required!')
|
|
346
|
+
|
|
347
|
+
if trace_length_samples is None:
|
|
348
|
+
trace_length_samples = int(round(1e-3 * trace_length_msec * sample_rate))
|
|
349
|
+
if pretrigger_length_msec is None:
|
|
350
|
+
pretrigger_length_msec = 1e3 * pretrigger_length_samples / sample_rate
|
|
351
|
+
else:
|
|
352
|
+
pretrigger_length_samples = int(round(1e-3 * pretrigger_length_msec * sample_rate))
|
|
353
|
+
|
|
354
|
+
dt = 1 / sample_rate
|
|
355
|
+
t0 = pretrigger_length_msec * 1e-3
|
|
356
|
+
t = np.asarray(list(range(trace_length_samples))) * dt
|
|
357
|
+
|
|
358
|
+
if A is not None and B is not None and C is not None:
|
|
359
|
+
poles = 4
|
|
360
|
+
template = qp.utils.make_template_fourpole(
|
|
361
|
+
t, A, B, C, tau_r, tau_f1, tau_f2, tau_f3,
|
|
362
|
+
t0=t0, fs=sample_rate, normalize=True
|
|
363
|
+
)
|
|
364
|
+
elif A is not None and B is not None:
|
|
365
|
+
poles = 3
|
|
366
|
+
template = qp.utils.make_template_threepole(
|
|
367
|
+
t, A, B, tau_r, tau_f1, tau_f2,
|
|
368
|
+
t0=t0, fs=sample_rate, normalize=True
|
|
369
|
+
)
|
|
370
|
+
else:
|
|
371
|
+
poles = 2
|
|
372
|
+
template = qp.utils.make_template_twopole(
|
|
373
|
+
t, A, tau_r, tau_f1,
|
|
374
|
+
t0=t0, fs=sample_rate, normalize=True
|
|
375
|
+
)
|
|
376
|
+
|
|
377
|
+
metadata = {
|
|
378
|
+
'sample_rate': sample_rate,
|
|
379
|
+
'sample_rate_hz': sample_rate,
|
|
380
|
+
'nb_samples': trace_length_samples,
|
|
381
|
+
'nb_pretrigger_samples': pretrigger_length_samples,
|
|
382
|
+
'nb_poles': poles,
|
|
383
|
+
'A': A, 'tau_r': tau_r, 'tau_f1': tau_f1,
|
|
384
|
+
}
|
|
385
|
+
if B is not None:
|
|
386
|
+
metadata['B'] = B
|
|
387
|
+
metadata['tau_f2'] = tau_f2
|
|
388
|
+
if C is not None:
|
|
389
|
+
metadata['C'] = C
|
|
390
|
+
metadata['tau_f3'] = tau_f3
|
|
391
|
+
|
|
392
|
+
channels = self._normalize_channels(channels)
|
|
393
|
+
for chan in channels:
|
|
394
|
+
self.set_template(
|
|
395
|
+
chan,
|
|
396
|
+
template,
|
|
397
|
+
sample_rate=sample_rate,
|
|
398
|
+
pretrigger_length_samples=pretrigger_length_samples,
|
|
399
|
+
metadata=metadata,
|
|
400
|
+
tag=tag,
|
|
401
|
+
)
|
|
402
|
+
|
|
403
|
+
def create_template_sum_twopoles(self, channels,
|
|
404
|
+
amplitudes,
|
|
405
|
+
rise_times,
|
|
406
|
+
fall_times,
|
|
407
|
+
sample_rate=None,
|
|
408
|
+
trace_length_msec=None,
|
|
409
|
+
trace_length_samples=None,
|
|
410
|
+
pretrigger_length_msec=None,
|
|
411
|
+
pretrigger_length_samples=None,
|
|
412
|
+
tag='default'):
|
|
413
|
+
if sample_rate is None:
|
|
414
|
+
raise ValueError('ERROR: "sample_rate" argument required')
|
|
415
|
+
if trace_length_msec is None and trace_length_samples is None:
|
|
416
|
+
raise ValueError('ERROR: Trace length required ("trace_length_msec" or "trace_length_samples")!')
|
|
417
|
+
if pretrigger_length_msec is None and pretrigger_length_samples is None:
|
|
418
|
+
raise ValueError('ERROR: Pretrigger length required ("pretrigger_length_msec" or "pretrigger_length_samples")!')
|
|
419
|
+
|
|
420
|
+
if trace_length_samples is None:
|
|
421
|
+
trace_length_samples = int(round(1e-3 * trace_length_msec * sample_rate))
|
|
422
|
+
if pretrigger_length_msec is None:
|
|
423
|
+
pretrigger_length_msec = 1e3 * pretrigger_length_samples / sample_rate
|
|
424
|
+
else:
|
|
425
|
+
pretrigger_length_samples = int(round(1e-3 * pretrigger_length_msec * sample_rate))
|
|
426
|
+
|
|
427
|
+
dt = 1 / sample_rate
|
|
428
|
+
time_array = np.asarray(list(range(trace_length_samples))) * dt
|
|
429
|
+
template = qp.utils.make_template_sum_twopoles(
|
|
430
|
+
time_array, amplitudes, rise_times, fall_times, normalize=True
|
|
431
|
+
)
|
|
432
|
+
|
|
433
|
+
metadata = {
|
|
434
|
+
'sample_rate': sample_rate,
|
|
435
|
+
'sample_rate_hz': sample_rate,
|
|
436
|
+
'nb_samples': trace_length_samples,
|
|
437
|
+
'nb_pretrigger_samples': pretrigger_length_samples,
|
|
438
|
+
'nb_sum_twopoles': len(amplitudes),
|
|
439
|
+
}
|
|
440
|
+
|
|
441
|
+
channels = self._normalize_channels(channels)
|
|
442
|
+
for chan in channels:
|
|
443
|
+
self.set_template(
|
|
444
|
+
chan,
|
|
445
|
+
template,
|
|
446
|
+
sample_rate=sample_rate,
|
|
447
|
+
pretrigger_length_samples=pretrigger_length_samples,
|
|
448
|
+
metadata=metadata,
|
|
449
|
+
tag=tag,
|
|
450
|
+
)
|
|
@@ -0,0 +1,185 @@
|
|
|
1
|
+
"""Raw-stream and trigger-dataframe sources used by ProcessingData."""
|
|
2
|
+
|
|
3
|
+
import copy
|
|
4
|
+
|
|
5
|
+
import vaex as vx
|
|
6
|
+
|
|
7
|
+
from pytesdaqx.io import StreamReader
|
|
8
|
+
from pytesprocess.utils import extract_stream_id
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class RawDataSource:
|
|
12
|
+
"""Storage-independent access to one raw stream at a time."""
|
|
13
|
+
|
|
14
|
+
def __init__(self, raw_files, catalog=None):
|
|
15
|
+
self._raw_files = raw_files
|
|
16
|
+
self._catalog = catalog
|
|
17
|
+
self._reader = None
|
|
18
|
+
self._stream_id = None
|
|
19
|
+
|
|
20
|
+
self._partition_target_duration_s = 10.0
|
|
21
|
+
self._partition_overlap_left_samples = 0
|
|
22
|
+
self._partition_overlap_right_samples = 0
|
|
23
|
+
self._partition_next_core_start_index = 0
|
|
24
|
+
self._partition_core_start_index = None
|
|
25
|
+
self._partition_core_length_samples = None
|
|
26
|
+
self._partition_read_start_index = None
|
|
27
|
+
self._partition_read_length_samples = None
|
|
28
|
+
self._stream_length_samples = None
|
|
29
|
+
self._use_continuous_partitions = False
|
|
30
|
+
|
|
31
|
+
@property
|
|
32
|
+
def stream_id(self):
|
|
33
|
+
return self._stream_id
|
|
34
|
+
|
|
35
|
+
@property
|
|
36
|
+
def reader(self):
|
|
37
|
+
return self._reader
|
|
38
|
+
|
|
39
|
+
@property
|
|
40
|
+
def uses_continuous_partitions(self):
|
|
41
|
+
return self._use_continuous_partitions
|
|
42
|
+
|
|
43
|
+
def configure_partitions(self, partition_target_duration_s=10.0,
|
|
44
|
+
overlap_left_samples=0,
|
|
45
|
+
overlap_right_samples=0):
|
|
46
|
+
if float(partition_target_duration_s) <= 0:
|
|
47
|
+
raise ValueError('ERROR: partition_target_duration_s must be positive!')
|
|
48
|
+
if int(overlap_left_samples) < 0 or int(overlap_right_samples) < 0:
|
|
49
|
+
raise ValueError('ERROR: partition overlap must be non-negative!')
|
|
50
|
+
self._partition_target_duration_s = float(partition_target_duration_s)
|
|
51
|
+
self._partition_overlap_left_samples = int(overlap_left_samples)
|
|
52
|
+
self._partition_overlap_right_samples = int(overlap_right_samples)
|
|
53
|
+
|
|
54
|
+
def select_stream(self, stream, force=False):
|
|
55
|
+
stream_id = extract_stream_id(stream)
|
|
56
|
+
if stream_id not in self._raw_files:
|
|
57
|
+
raise ValueError(f'ERROR: raw stream {stream_id} not found!')
|
|
58
|
+
if (not force and self._reader is not None
|
|
59
|
+
and self._stream_id == stream_id):
|
|
60
|
+
return
|
|
61
|
+
|
|
62
|
+
self.close()
|
|
63
|
+
self._stream_id = stream_id
|
|
64
|
+
self._reader = StreamReader(copy.deepcopy(self._raw_files[stream_id]))
|
|
65
|
+
self._reset_partition_state()
|
|
66
|
+
|
|
67
|
+
if self._catalog is not None:
|
|
68
|
+
view = self._catalog.filter(streams=stream_id)
|
|
69
|
+
if (view.storage_format == 'zarr'
|
|
70
|
+
and self._reader.is_continuous_stream):
|
|
71
|
+
if view.n_samples is None:
|
|
72
|
+
raise ValueError(
|
|
73
|
+
f'ERROR: Unable to determine stream length for {stream_id}!'
|
|
74
|
+
)
|
|
75
|
+
self._stream_length_samples = int(view.n_samples)
|
|
76
|
+
self._use_continuous_partitions = True
|
|
77
|
+
|
|
78
|
+
def read_next_native(self, channels=None):
|
|
79
|
+
return self._reader.read_next(
|
|
80
|
+
channels=channels, units='amps', include_metadata=True
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
def read_record(self, record, channels=None, **kwargs):
|
|
84
|
+
return self._reader.read_record(
|
|
85
|
+
record, channels=channels, units='amps', include_metadata=True,
|
|
86
|
+
**kwargs,
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
def read_next_partition(self, channels=None, sample_rate_hz=None):
|
|
90
|
+
core_start = int(self._partition_next_core_start_index)
|
|
91
|
+
if (self._stream_length_samples is None
|
|
92
|
+
or core_start >= self._stream_length_samples):
|
|
93
|
+
return None
|
|
94
|
+
if sample_rate_hz is None:
|
|
95
|
+
raise ValueError('ERROR: sample_rate_hz is required for partitions!')
|
|
96
|
+
|
|
97
|
+
target_samples = max(
|
|
98
|
+
1, int(round(self._partition_target_duration_s * float(sample_rate_hz)))
|
|
99
|
+
)
|
|
100
|
+
core_stop = min(
|
|
101
|
+
core_start + target_samples, self._stream_length_samples
|
|
102
|
+
)
|
|
103
|
+
read_start = max(
|
|
104
|
+
0, core_start - self._partition_overlap_left_samples
|
|
105
|
+
)
|
|
106
|
+
read_stop = min(
|
|
107
|
+
self._stream_length_samples,
|
|
108
|
+
core_stop + self._partition_overlap_right_samples,
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
traces, info = self._reader.read_partition(
|
|
112
|
+
partition_start_index=read_start,
|
|
113
|
+
partition_length_samples=read_stop-read_start,
|
|
114
|
+
channels=channels, units='amps', include_metadata=True,
|
|
115
|
+
)
|
|
116
|
+
|
|
117
|
+
self._partition_core_start_index = core_start
|
|
118
|
+
self._partition_core_length_samples = core_stop-core_start
|
|
119
|
+
self._partition_read_start_index = read_start
|
|
120
|
+
self._partition_read_length_samples = read_stop-read_start
|
|
121
|
+
self._partition_next_core_start_index = core_stop
|
|
122
|
+
|
|
123
|
+
return traces, info, read_start
|
|
124
|
+
|
|
125
|
+
def partition_context(self):
|
|
126
|
+
if not self._use_continuous_partitions:
|
|
127
|
+
return None
|
|
128
|
+
return {
|
|
129
|
+
'partition_start_index': self._partition_core_start_index,
|
|
130
|
+
'partition_length_samples': self._partition_core_length_samples,
|
|
131
|
+
'partition_read_start_index': self._partition_read_start_index,
|
|
132
|
+
'partition_read_length_samples': self._partition_read_length_samples,
|
|
133
|
+
'stream_length_samples': self._stream_length_samples,
|
|
134
|
+
'stream_id': self._stream_id,
|
|
135
|
+
}
|
|
136
|
+
|
|
137
|
+
def get_detector_settings(self):
|
|
138
|
+
return self._reader.get_detector_settings()
|
|
139
|
+
|
|
140
|
+
def close(self):
|
|
141
|
+
if self._reader is not None:
|
|
142
|
+
try:
|
|
143
|
+
self._reader.close()
|
|
144
|
+
except Exception:
|
|
145
|
+
pass
|
|
146
|
+
self._reader = None
|
|
147
|
+
|
|
148
|
+
def _reset_partition_state(self):
|
|
149
|
+
self._partition_next_core_start_index = 0
|
|
150
|
+
self._partition_core_start_index = None
|
|
151
|
+
self._partition_core_length_samples = None
|
|
152
|
+
self._partition_read_start_index = None
|
|
153
|
+
self._partition_read_length_samples = None
|
|
154
|
+
self._stream_length_samples = None
|
|
155
|
+
self._use_continuous_partitions = False
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
class TriggerDataframeSource:
|
|
159
|
+
"""Iterate independent trigger-dataframe tasks."""
|
|
160
|
+
|
|
161
|
+
def __init__(self, trigger_files):
|
|
162
|
+
self._trigger_files = trigger_files or {}
|
|
163
|
+
self._dataframe = None
|
|
164
|
+
self._index = -1
|
|
165
|
+
self._task_id = None
|
|
166
|
+
|
|
167
|
+
@property
|
|
168
|
+
def task_id(self):
|
|
169
|
+
return self._task_id
|
|
170
|
+
|
|
171
|
+
def select_task(self, task_id):
|
|
172
|
+
task_id = str(task_id)
|
|
173
|
+
if task_id not in self._trigger_files:
|
|
174
|
+
raise ValueError(
|
|
175
|
+
f'ERROR: trigger dataframe task {task_id} not found!'
|
|
176
|
+
)
|
|
177
|
+
self._task_id = task_id
|
|
178
|
+
self._dataframe = vx.open_many(self._trigger_files[task_id])
|
|
179
|
+
self._index = -1
|
|
180
|
+
|
|
181
|
+
def next_row(self):
|
|
182
|
+
self._index += 1
|
|
183
|
+
if self._dataframe is None or self._index >= len(self._dataframe):
|
|
184
|
+
return None
|
|
185
|
+
return self._dataframe.to_records(index=self._index)
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
"""Canonical state for the event/window currently being processed."""
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass, field
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@dataclass
|
|
7
|
+
class EventContext:
|
|
8
|
+
"""Mutable per-event state shared by ProcessingData accessors.
|
|
9
|
+
|
|
10
|
+
Raw-data provenance is canonicalized around streams. Legacy HDF5
|
|
11
|
+
coordinates are retained only where they are physically meaningful.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
full_traces: object = None
|
|
15
|
+
admin_info: dict | None = None
|
|
16
|
+
truncated_traces: dict = field(default_factory=dict)
|
|
17
|
+
dataframe_info: object = None
|
|
18
|
+
stream_id: str | None = None
|
|
19
|
+
stream_number: int | None = None
|
|
20
|
+
stream_trigger_index: int | None = None
|
|
21
|
+
global_segment_number: int | None = None
|
|
22
|
+
segment_trigger_index: int | None = None
|
|
23
|
+
salting_info: dict | None = None
|
|
24
|
+
|
|
25
|
+
def reset(self, *, stream_id=None):
|
|
26
|
+
self.full_traces = None
|
|
27
|
+
self.admin_info = None
|
|
28
|
+
self.truncated_traces = {}
|
|
29
|
+
self.dataframe_info = None
|
|
30
|
+
self.stream_id = stream_id
|
|
31
|
+
self.stream_number = None
|
|
32
|
+
self.stream_trigger_index = None
|
|
33
|
+
self.global_segment_number = None
|
|
34
|
+
self.segment_trigger_index = None
|
|
35
|
+
self.salting_info = None
|