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.
Files changed (60) hide show
  1. pytesprocess/__init__.py +9 -0
  2. pytesprocess/_version.py +2 -0
  3. pytesprocess/cli/__init__.py +1 -0
  4. pytesprocess/cli/commands/__init__.py +5 -0
  5. pytesprocess/cli/commands/event.py +66 -0
  6. pytesprocess/cli/commands/filter.py +17 -0
  7. pytesprocess/cli/commands/ivsweep.py +29 -0
  8. pytesprocess/cli/common.py +86 -0
  9. pytesprocess/cli/main.py +81 -0
  10. pytesprocess/config/__init__.py +4 -0
  11. pytesprocess/config/loader.py +94 -0
  12. pytesprocess/config/manager.py +297 -0
  13. pytesprocess/config/resolvers/__init__.py +5 -0
  14. pytesprocess/config/resolvers/common.py +56 -0
  15. pytesprocess/config/resolvers/feature.py +293 -0
  16. pytesprocess/config/resolvers/salting.py +86 -0
  17. pytesprocess/config/resolvers/trigger.py +84 -0
  18. pytesprocess/config/selectors.py +108 -0
  19. pytesprocess/config/validation.py +314 -0
  20. pytesprocess/config/warnings.py +2 -0
  21. pytesprocess/core/__init__.py +10 -0
  22. pytesprocess/core/algorithms.py +1455 -0
  23. pytesprocess/core/didv.py +1648 -0
  24. pytesprocess/core/eventbuilder.py +495 -0
  25. pytesprocess/core/filterbuilder.py +81 -0
  26. pytesprocess/core/filterdata.py +1849 -0
  27. pytesprocess/core/ivsweep.py +2072 -0
  28. pytesprocess/core/noise.py +923 -0
  29. pytesprocess/core/noisemodel.py +1408 -0
  30. pytesprocess/core/oftrigger.py +1035 -0
  31. pytesprocess/core/template.py +450 -0
  32. pytesprocess/process/__init__.py +6 -0
  33. pytesprocess/process/data_source.py +185 -0
  34. pytesprocess/process/event_context.py +35 -0
  35. pytesprocess/process/feature_plan.py +186 -0
  36. pytesprocess/process/feature_resources.py +267 -0
  37. pytesprocess/process/features.py +1024 -0
  38. pytesprocess/process/filterprocess.py +1176 -0
  39. pytesprocess/process/ivprocess.py +1380 -0
  40. pytesprocess/process/processing_data.py +967 -0
  41. pytesprocess/process/randoms.py +921 -0
  42. pytesprocess/process/triggers.py +1011 -0
  43. pytesprocess/salting/__init__.py +7 -0
  44. pytesprocess/salting/generator.py +364 -0
  45. pytesprocess/salting/injector.py +329 -0
  46. pytesprocess/salting/sampling.py +84 -0
  47. pytesprocess/utils/__init__.py +5 -0
  48. pytesprocess/utils/arg_utils.py +122 -0
  49. pytesprocess/utils/dataframe_output.py +120 -0
  50. pytesprocess/utils/filter_hdf5.py +594 -0
  51. pytesprocess/utils/utils.py +701 -0
  52. pytesprocess/workflows/__init__.py +3 -0
  53. pytesprocess/workflows/processing.py +317 -0
  54. pytesprocess/workflows/salting.py +133 -0
  55. pytesprocess-0.1.1.dist-info/METADATA +211 -0
  56. pytesprocess-0.1.1.dist-info/RECORD +60 -0
  57. pytesprocess-0.1.1.dist-info/WHEEL +5 -0
  58. pytesprocess-0.1.1.dist-info/entry_points.txt +2 -0
  59. pytesprocess-0.1.1.dist-info/licenses/LICENSE +21 -0
  60. 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,6 @@
1
+ from .features import *
2
+ from .triggers import *
3
+ from .processing_data import *
4
+ from .randoms import *
5
+ from .ivprocess import *
6
+ from .filterprocess import *
@@ -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