1#!/usr/bin/env python 

2import numpy as np 

3import h5py 

4 

5 

6class WavesetReader: 

7 def __init__(self, filename): 

8 self.filename = filename 

9 self._wavesets = None 

10 

11 @property 

12 def wavesets(self): 

13 if self._wavesets is None: 

14 with h5py.File(self.filename, "r") as f: 

15 self._wavesets = list(f.keys()) 

16 return self._wavesets 

17 

18 def __getitem__(self, key): 

19 with h5py.File(self.filename, "r") as f: 

20 raw_waveforms = f[f"{key}/waveforms"][:] 

21 v_gain = f[f"{key}/waveform_info/v_gain"][()] 

22 h_int = f[f"{key}/waveform_info/h_int"][()] 

23 return Waveset(raw_waveforms, v_gain, h_int) 

24 

25 

26class Waveset: 

27 def __init__(self, raw_waveforms, v_gain, h_int): 

28 self.raw_waveforms = raw_waveforms 

29 self.v_gain = v_gain 

30 self.h_int = h_int 

31 self._waveforms = None 

32 

33 @property 

34 def waveforms(self): 

35 if self._waveforms is None: 

36 self._waveforms = self.raw_waveforms * self.v_gain 

37 return self._waveforms 

38 

39 def zeroed_waveforms(self, baseline_min, baseline_max): 

40 return ( 

41 self.waveforms.T 

42 - np.mean(self.waveforms[:, baseline_min:baseline_max], axis=1) 

43 ).T