pytesprocess 0.1.1__py3-none-any.whl

pytesprocess/core/noise.py

@@ -0,0 +1,923 @@
+import copy
+from glob import glob
+from pathlib import Path
+import os
+
+import numpy as np
+import qetpy as qp
+import vaex as vx
+from pytesdaqx.io import AcquisitionCatalog, StreamReader
+
+from pytesprocess.core.filterdata import FilterData
+from pytesprocess.process.randoms import Randoms
+from pytesprocess.utils import (
+    convert_length_msec_to_samples,
+    extract_stream_id,
+    extract_stream_num,
+    utils,
+)
+
+
+class Noise(FilterData):
+    """
+    Class to manage noise calculation from 
+    randoms for multiple channels
+    """
+
+    def __init__(self, verbose=True, filter_data=None):
+        super().__init__(verbose=verbose, filter_data=filter_data)
+
+        self._raw_data_files = None
+        self._acquisition_name = None
+        self._raw_base_path = None
+        self._stream_list = None
+        self._detector_config = None
+        self._dataframe = None
+        self._record_list = None
+        self._fs = None
+        self._offset = dict()
+        self._offset_err = dict()
+        self._available_channels = None
+        self._catalog = None
+
+        
+    def get_detector_config(self, channel):
+        """
+        get detector config
+        """
+        if self._detector_config is None:
+            print('WARNING: No data has been set yet! '
+                  'Returning None ')
+            return None
+        elif channel not in self._detector_config.keys():
+            print(f'WARNING: No channel {channel}  found! '
+                  f'Returning None ')
+            return None
+        return self._detector_config[channel]
+                
+    def get_sample_rate(self):
+        """
+        Get sample rate in Hz ("calc_psd" needs to be 
+        called before)
+        """
+
+        return self._fs
+
+
+    def get_offset(self, channel):
+        """
+        Get offset, return None if no offset
+        available
+        """
+        offset = None
+        if channel in self._offset:
+            offset = self._offset[channel]
+        else:
+            print(f'WARNING: No offset available for channel '
+                  f'{channel}. You need to calculate psd first! '
+                  f'Returning None. ')
+            
+        return offset
+
+    def get_offset_error(self, channel):
+        """Get the statistical uncertainty on the stored baseline offset."""
+        if channel in self._offset_err:
+            return self._offset_err[channel]
+        if self._verbose:
+            print(
+                f"WARNING: No offset uncertainty available for channel {channel}. "
+                "Calculate an offset or PSD first. Returning None."
+            )
+        return None
+
+    def clear_randoms(self):
+        """
+        Clear internal data, however
+        keep self._filter_data
+        """
+
+        # clear data
+        self._dataframe = None
+        self._record_list = None
+        self._raw_data_files = None
+        self._acquisition_name = None
+        self._raw_base_path = None
+        self._stream_list = None
+        self._detector_config = None
+        self._fs = None
+        self._offset = dict()
+        self._offset_err = dict()
+        self._available_channels = None
+        self._catalog = None
+        
+    def set_randoms(self, data_path,
+                    streams=None,
+                    dataframe=None,
+                    record_list=None,
+                    restricted=False,
+                    data_type=None):
+        """
+        Set randoms events
+           -  Select randoms using datafame (if not None) or record_list
+           -  OR use full trace (if datafame is None)
+        """
+
+        # initialize data
+        self.clear_randoms()
+  
+        # check arguments
+        if (dataframe is not None
+            and record_list is not None):
+            raise ValueError('ERROR: choose between "dataframe" and '
+                             '"record_list", not both')
+        
+        # data catalog
+        catalog_full = AcquisitionCatalog(data_path, verbose=True)
+        self._catalog = catalog_full.filter(
+            measurement_types=data_type,
+            streams=streams,
+            restricted=restricted
+        )
+
+        # get file dictionary  
+        data_dict = self._catalog.select_files_by_stream(
+            stream_key='stream_id'
+        )
+
+        if not data_dict:
+            raise ValueError(f'No files with data type "{data_type}" '
+                             f'were found! Check raw data path...')
+
+        # storage format
+        self._storage_format = self._catalog.storage_format
+   
+        
+        # save info
+        self._raw_data_files = copy.deepcopy(data_dict)
+        acquisition_path = self._catalog.acquisition_path
+        path_obj = Path(acquisition_path)
+        self._acquisition_name = str(path_obj.name)
+        self._raw_base_path = str(path_obj.parent)
+        self._stream_list = list(self._raw_data_files.keys())
+        self._detector_config = dict()
+        self._available_channels = self._catalog.entries[0]['record_channels']
+        
+
+        # check dataframe
+        if dataframe is not None:
+            
+            if isinstance(dataframe, vx.dataframe.DataFrame):
+                if len(dataframe)<1:
+                    raise ValueError('ERROR: No event found in the datafame!')
+            elif isinstance(dataframe, str):    
+                dataframe = self._load_dataframe(dataframe)
+
+            self._dataframe = dataframe
+
+        elif record_list is not None:
+            self._record_list = record_list
+
+        # check filter data
+        if self._filter_data:
+            print('WARNING: Some noise data have been previously saved. '
+                  'Use "describe()" to check. If needed clear data '
+                  'using "clear_data(channels=None, tag=None)" function!')
+            
+            
+    def generate_randoms(self, data_path,
+                         streams=None,
+                         random_rate=None, nrandoms=None,
+                         min_separation_msec=100,
+                         edge_exclusion_msec=50,
+                         partition_target_duration_s=None,
+                         random_seed=None,
+                         restricted=False,
+                         data_type='background'):
+        """
+        Generate randoms from continuous data
+        """
+        # initialize data
+        self.clear_randoms()
+        
+        # data catalog
+        catalog_full = AcquisitionCatalog(data_path, verbose=True)
+        self._catalog = catalog_full.filter(
+            measurement_types=data_type,
+            streams=streams,
+            restricted=restricted
+        )
+
+        # get file dictionary  
+        data_dict = self._catalog.select_files_by_stream(
+            stream_key='stream_id'
+        )
+        
+        if not data_dict:
+            raise ValueError(f'No files with data type "{data_type}" '
+                             f'were found! Check raw data path...')
+
+        # raw data format
+        self._storage_format = self._catalog.storage_format
+      
+        # save info
+        self._raw_data_files = copy.deepcopy(data_dict)
+        acquisition_path = self._catalog.acquisition_path
+        path_obj = Path(acquisition_path)
+        self._acquisition_name = str(path_obj.name)
+        self._raw_base_path = str(path_obj.parent)
+        self._stream_list = list(self._raw_data_files.keys())
+        self._detector_config = dict()
+        self._available_channels = self._catalog.entries[0]['record_channels']
+                       
+        # generate randoms
+        rand_inst = Randoms(data_path, streams=streams,
+                            verbose=self._verbose,
+                            restricted=restricted,
+                            data_type=data_type)
+
+        
+        self._dataframe = rand_inst.process(
+            random_rate=random_rate,
+            nrandoms=nrandoms,
+            min_separation_msec=min_separation_msec,
+            edge_exclusion_msec=edge_exclusion_msec,
+            partition_target_duration_s=partition_target_duration_s,
+            random_seed=random_seed,
+            lgc_save=False,
+            lgc_output=True
+        )
+
+        # check filter data
+        if self._filter_data:
+            print('WARNING: Some noise data have been previously saved. '
+                  'Use "describe()" to check. If needed clear data '
+                  'using "clear_data(channels=None, tag=None)" function!')
+
+            
+    def calc_offset(self, channels=None,
+                    stream=None,
+                    trace_length_msec=None,
+                    trace_length_samples=None,
+                    pretrigger_length_msec=None,
+                    pretrigger_length_samples=None,
+                    nrandoms=None):
+        """Calculate baseline offsets from the currently selected randoms.
+
+        The same random-event selection used by :meth:`calc_psd` is used here,
+        but no spectrum is calculated.  This is useful for per-cycle operating
+        point calibration where the background baseline is needed before the
+        dIdV/noise processing chain.
+        """
+        if nrandoms is None:
+            raise ValueError('ERROR: Maximum number of randoms required! Add "nrandoms" argument.')
+        if self._raw_data_files is None:
+            raise ValueError('ERROR: No raw data available. Use set_randoms() or generate_randoms() first!')
+
+        if channels is None:
+            channels = self._available_channels
+        if isinstance(channels, str):
+            channels = [channels]
+
+        results = {}
+        for channel in channels:
+            channel_list, separator = utils.split_channel_name(
+                channel, self._available_channels
+            )
+            if separator is not None or len(channel_list) != 1:
+                raise ValueError(
+                    'ERROR: calc_offset() requires individual channels; '
+                    f'got {channel!r}.'
+                )
+
+            traces, metadata = self._get_traces(
+                channel_list,
+                nrandoms=nrandoms,
+                trace_length_msec=trace_length_msec,
+                trace_length_samples=trace_length_samples,
+                pretrigger_length_msec=pretrigger_length_msec,
+                pretrigger_length_samples=pretrigger_length_samples,
+                stream=stream,
+            )
+            fs = float(metadata['sample_rate_hz'])
+            traces = traces[:, 0, :] if traces.ndim == 3 else traces
+            cut = qp.autocuts_noise(traces, fs=fs)
+            if np.sum(cut) == 0:
+                raise ValueError(
+                    f'ERROR: No events selected after noise autocut for channel {channel}!'
+                )
+
+            offset, offset_err = qp.utils.calc_offset(traces[cut], fs=fs)
+            self._fs = fs
+            self._offset[channel] = offset
+            self._offset_err[channel] = offset_err
+            results[channel] = {
+                'offset': offset,
+                'offset_err': offset_err,
+                'cut_efficiency': float(np.sum(cut)) / len(cut) * 100.0,
+                'sample_rate_hz': fs,
+            }
+
+        return results
+
+
+    def calc_psd(self, channels=None, 
+                 stream=None,
+                 trace_length_msec=None,
+                 trace_length_samples=None,
+                 pretrigger_length_msec=None,
+                 pretrigger_length_samples=None,
+                 nrandoms=None,
+                 weights=None,
+                 tag='default'):
+        """
+        Calculate two-sided and folded-over PSD in Amps^2/Hz
+        A specific stream can be specified
+        
+        """
+
+        # --------------------------------
+        # Check arguments
+        # --------------------------------
+
+        if channels is None:
+            channels = self._available_channels
+        if isinstance(channels, str):
+            channels = [channels]
+            
+        if nrandoms is None:
+            raise ValueError('ERROR: Maximum number of randoms required!'
+                             ' Add "nrandoms" argument.')
+        
+        # check raw data has been loaded
+        if self._raw_data_files is None:
+            raise ValueError('ERROR: No raw data available. Use '
+                             + '"set_randoms()" or "generate_randoms() '
+                             + 'function first!')
+
+        # check trace length
+        if (trace_length_msec is not None
+            and trace_length_samples is not None):
+            raise ValueError('ERROR: Trace length need to be '
+                             'in msec OR samples, nto both')
+
+        if (pretrigger_length_msec is not None
+            and pretrigger_length_samples is not None):
+            raise ValueError('ERROR: Pretrigger length need to be '
+                             'in msec OR samples, nto both')
+
+        
+        # --------------------------------
+        # Loop channels and calculate PSD
+        # --------------------------------
+        for chan in channels:
+            
+            if self._verbose:
+                if stream is None:
+                    print('INFO: Processing PSD for channel '
+                          +  chan)
+                else:
+                    print('INFO: Processing PSD for channel '
+                          +  chan + ' using stream '
+                          + str(stream))
+
+
+            # let's check if sum of pulses 
+            separator = None
+            chan_list, separator = utils.split_channel_name(
+                chan, self._available_channels
+            )
+
+
+            # weights
+            weights_array = None
+            if weights is not None:
+            
+                weights_array = np.ones(len(chan_list))
+            
+                for ichan, chan_split in enumerate(chan_list):
+                    if chan_split in weights:
+                        weights_array[ichan] = weights[chan_split]
+            
+            # let's do first overall all PSD
+            traces, traces_metadata = self._get_traces(
+                chan_list,
+                nrandoms=nrandoms,
+                trace_length_msec=trace_length_msec,
+                trace_length_samples=trace_length_samples,
+                pretrigger_length_msec=pretrigger_length_msec,
+                pretrigger_length_samples=pretrigger_length_samples,
+                stream=stream
+            )
+
+            self._fs = traces_metadata['sample_rate_hz']
+
+            if separator == '+':
+                if weights_array is not None:
+                    weights_array = weights_array[np.newaxis, :, np.newaxis]
+                    traces = traces * weights_array
+                traces = np.sum(traces, axis=1)
+                           
+            elif separator == '-':
+                if weights_array is not None:
+                    traces = (traces[:,0,:]*weights_array[0]
+                              - traces[:,1,:]*weights_array[1])
+                else:
+                    traces = traces[:,0,:] - traces[:,1,:]
+                    
+            elif separator is not None:
+                raise ValueError('ERROR: PSD can only be calculated '
+                                 'on single channels!')
+                
+            if traces.ndim==3:
+                if traces.shape[1] != 1:
+                    raise ValueError('ERROR: Multiple channels. Expecting '
+                                     'only one. Something went wrong!')
+                traces = traces[:,0,:]
+                
+            # autocut_noise
+            cut = qp.autocuts_noise(traces, fs=self._fs)
+        
+            if np.sum(cut)==0:
+                raise ValueError('ERROR: No events selected after noise autocut! '
+                                 + 'Unable to calculate PSD')
+                
+            cut_eff = np.sum(cut)/len(cut)*100
+            if self._verbose:
+                print('INFO: Number of events after cuts = '
+                      '{}, efficiency = '
+                      '{:0.2f}%'.format(np.sum(cut), cut_eff))
+                
+            # calc PSD
+            freqs, psd = qp.calc_psd(traces[cut],
+                                     fs=self._fs,
+                                     folded_over=False)
+        
+            # calc baseline offset and uncertainty
+            offset, offset_err = qp.utils.calc_offset(traces[cut], fs=self._fs)
+            self._offset[chan] = offset
+            self._offset_err[chan] = offset_err
+                       
+            # metadata
+            traces_metadata['cut_efficiency'] = cut_eff
+
+            if weights is not None:
+                for wchan,wval in weights.items():
+                    param = f'weights_{wchan}'
+                    traces_metadata[param] = wval
+
+            self.set_psd(
+                chan,
+                psd,
+                freqs,
+                sample_rate=self._fs,
+                pretrigger_length_samples=traces_metadata.get(
+                    'nb_pretrigger_samples'
+                ),
+                metadata=traces_metadata,
+                tag=tag,
+            )
+         
+ 
+            
+    def calc_csd(self, channels, 
+                 stream=None,
+                 trace_length_msec=None,
+                 trace_length_samples=None,
+                 pretrigger_length_msec=None,
+                 pretrigger_length_samples=None,
+                 nrandoms=None,
+                 tag='default',
+                 use_hann_window=False):
+        """
+        Calculate two-sided and folded CSD. 
+       
+        """
+
+        # --------------------------------
+        # Check arguments
+        # --------------------------------
+
+        if nrandoms is None:
+            raise ValueError('ERROR: Maximum number of randoms required!'
+                             ' Add "nrandoms" argument.')
+        
+        # check raw data has been loaded
+        if self._raw_data_files is None:
+            raise ValueError('ERROR: No raw data available. Use '
+                             + '"set_randoms()" function first!')
+
+        # check trace length
+        if (trace_length_msec is not None
+            and trace_length_samples is not None):
+            raise ValueError('ERROR: Trace length need to be '
+                             'in msec OR samples, nto both')
+
+        if (pretrigger_length_msec is not None
+            and pretrigger_length_samples is not None):
+            raise ValueError('ERROR: Pretrigger length need to be '
+                             'in msec OR samples, nto both')
+
+
+        # channels
+        if isinstance(channels, str):
+            
+            if '|' in channels:
+                channels = channels.replace(' ','')
+                channels = channels.split('|')
+            else:
+                raise ValueError(
+                    'ERROR: At least 2 channels required to calculate csd'
+                )
+
+        # -------------
+        # Get data and
+        # apply cuts
+        # -------------
+        
+        traces, traces_metadata = self._get_traces(
+            channels,
+            nrandoms=nrandoms,
+            trace_length_msec=trace_length_msec,
+            trace_length_samples=trace_length_samples,
+            pretrigger_length_msec=pretrigger_length_msec,
+            pretrigger_length_samples=pretrigger_length_samples,
+            stream=stream
+        )
+        
+        # check shape
+        if traces.shape[1] != len(channels):
+            raise ValueError('ERROR: No all channels found in raw data. ')
+        
+        self._fs = traces_metadata['sample_rate_hz']
+        
+        # apply pileup cut
+        cut = np.ones(traces.shape[0], dtype=bool)
+        for ichan in range(len(channels)):
+            
+            traces_chan = traces[:, ichan,:]
+                                
+            # autocut_noise
+            cut_chan = qp.autocuts_noise(traces_chan, fs=self._fs)
+                      
+            if np.sum(cut_chan) == 0:
+                raise ValueError(f'ERROR: No events selected after pileup autocut '
+                                 f'for channel {channels[ichan]} ')
+            cut &= cut_chan
+                      
+        # check efficiency total cut 
+        if np.sum(cut) == 0:
+            raise ValueError(f'ERROR: No events selected after pileup cut!')
+
+        cut_eff = np.sum(cut)/len(cut)*100
+        if self._verbose:
+            print('INFO: Number of events after cuts = '
+                  '{}, efficiency = '
+                  '{:0.2f}%'.format(np.sum(cut), cut_eff))
+            
+        # calc CSD two-sided
+        freqs, csd = qp.calc_csd(traces[cut],
+                                 fs=self._fs,
+                                 folded_over=False,
+                                 use_hann_window=use_hann_window)
+        # metadata
+        traces_metadata['cut_efficiency'] = cut_eff
+
+        self.set_csd(
+            channels,
+            csd,
+            freqs,
+            sample_rate=self._fs,
+            pretrigger_length_samples=traces_metadata.get(
+                'nb_pretrigger_samples'
+            ),
+            metadata=traces_metadata,
+            tag=tag,
+        )
+               
+            
+    def _get_traces(self, channels,
+                    stream=None,
+                    trace_length_msec=None,
+                    trace_length_samples=None,
+                    pretrigger_length_msec=None,
+                    pretrigger_length_samples=None,
+                    nrandoms=5000):
+        """
+        Get raw data traces
+        """
+        
+        # channels
+        if isinstance(channels, str):
+            channels = [channels]
+        nb_channels = len(channels)
+
+        # filter catalog with stream
+        catalog = self._catalog
+        
+        # filter based on stream
+        first_stream = self._stream_list[0]
+        stream_num = None
+        stream_id =  None
+        if stream is not None:
+            stream_num = extract_stream_num(stream)
+            stream_id =  extract_stream_id(stream)
+            
+            # check if stream in list
+            if stream_id not in self._raw_data_files.keys():
+                raise ValueError(
+                    f'ERROR: stream {stream} not found! '
+                    f'Check raw data path input.')
+                        
+            catalog = catalog.filter(streams=stream_id)
+ 
+        # sample rate 
+        fs  = catalog.sample_rate_hz
+                  
+        # trace length
+        nb_samples = None
+        if trace_length_samples is not None:
+            nb_samples = trace_length_samples
+        elif trace_length_msec is not None:
+            nb_samples = convert_length_msec_to_samples(
+                trace_length_msec, fs)
+        else:
+            if self._dataframe is not None:
+                raise ValueError('ERROR: number of samples required!')
+
+        # pretrigger length
+        nb_pretrigger_samples = None
+        if pretrigger_length_samples is not None:
+            nb_pretrigger_samples = pretrigger_length_samples
+        elif pretrigger_length_msec is not None:
+            nb_pretrigger_samples = convert_length_msec_to_samples(
+                pretrigger_length_msec, fs)
+
+        if (nb_samples is not None
+            and nb_pretrigger_samples  is None):
+            nb_pretrigger_samples =  nb_samples//2
+
+            
+        # get (filtered) event list
+        record_list = None
+        stream_ids = None
+        if (self._record_list is not None
+            or self._dataframe is not None):
+            
+            record_list = self._get_filtered_record_list(
+                stream_num=stream_num,
+                nrandoms=nrandoms,
+                nb_samples=nb_samples,
+                nb_pretrigger_samples=nb_pretrigger_samples)
+
+            if len(record_list) == 0:
+                raise ValueError('ERROR: No events selected! Something '
+                                 + 'went wrong!')
+
+            if len(record_list)>nrandoms:
+                record_list = record_list[:nrandoms]
+                
+            nrandoms = len(record_list)
+            
+        else:
+            # case raw data directly -> get stream list
+          
+            # list of stream
+            stream_ids = list()
+            if stream_id is not None:
+                stream_ids.append(stream_id)
+            else:
+                for stream_id in self._raw_data_files.keys():
+                    stream_ids.append(stream_id)
+                    
+        # instantiate StreamReader
+        raw_path = self._raw_base_path + '/' + self._acquisition_name
+        reader = StreamReader(
+            raw_path,
+            streams=stream_ids,
+        )
+
+
+        # detector config
+        self._detector_config = reader.get_detector_settings()
+         
+            
+        trace_array = reader.read_records(
+            record_list=record_list,
+            n_records=nrandoms,
+            channels=channels,
+            trace_length_samples=nb_samples,
+            pretrigger_length_samples=nb_pretrigger_samples,
+            include_metadata=False,
+            units='amps',
+            stack=True)
+        
+        
+        if self._verbose:
+            chan_string = str(channels)
+            if len(channels)==1:
+                chan_string = channels[0]
+                print('INFO: ' + str(trace_array.shape[0])
+                      + ' events found in raw data'
+                      + ' for channel(s) '
+                      + str(chan_string))
+
+        # metadata
+        if nb_samples is None: # full trace
+            nb_samples = trace_array.shape[-1]
+            nb_pretrigger_samples =  nb_samples//2
+            
+        trace_metadata = dict()
+        trace_metadata['sample_rate_hz'] = fs
+        trace_metadata['nb_samples'] = nb_samples
+        trace_metadata['nb_pretrigger_samples'] = nb_pretrigger_samples
+        trace_metadata['nb_randoms'] = trace_array.shape[0]
+        trace_metadata['fridge_run'] =  catalog.summary()['fridge_run']
+        
+        return trace_array, trace_metadata
+
+
+    def _get_filtered_record_list(self,
+                                 stream_num=None,
+                                 nrandoms=None,
+                                 nb_samples=None,
+                                 nb_pretrigger_samples=None):
+        """
+        Get event list from dataframe or unfiltered 
+        event list
+        """
+
+        # check if data available
+        if (self._record_list is None 
+            and self._dataframe is None):
+            return record_list
+
+
+        # convert dataframe into an event list
+        record_list = list()
+        if self._dataframe is not None:
+
+            # case dataframe input
+            dataframe = self._dataframe.copy()
+
+            # only take randoms
+            cut  = dataframe.trigger_type == 3
+
+            # check stream number
+            if stream_num is not None:
+                cut = cut & (dataframe.stream_number == stream_num)
+
+            # filter
+            dataframe = dataframe.filter(cut)
+
+            # loop dataframe ans build event list
+            for idx in range(len(dataframe)):
+
+                # even record from dataframe
+                record = dataframe.to_records(idx)
+
+                # extract event parameters and stored in
+                # dictionary
+                stored_record = dict()
+                stored_record['acquisition_num'] = str(
+                        record['acquisition_number'])
+                
+                stored_record['stream_num'] = int(
+                    record['stream_number'])
+                
+                
+                if self._storage_format == 'hdf5':
+                    stored_record['global_segment_num'] = int(
+                        record['global_segment_number']
+                    )
+
+                    if 'segment_trigger_index' in record:
+                        stored_record['segment_trigger_index'] = int(
+                            record['segment_trigger_index']
+                        )
+                        stored_record['segment_length_samples'] = int(
+                            record['segment_length_samples']
+                        )
+                else:
+                    if 'stream_trigger_index' in record:
+                        stored_record['stream_trigger_index'] = (
+                            record['stream_trigger_index']
+                        )
+                        stored_record['stream_length_samples'] = int(
+                            record['stream_length_samples']
+                        )
+                                     
+                    elif ('partition_trigger_index' in record
+                          and  'partition_start_index'  in record):
+                        stored_record['partition_trigger_index'] = (
+                            record['partition_trigger_index']
+                        )
+                        stored_record['partition_start_index'] = (
+                            record['partition_start_index']
+                        )
+
+                        stored_record['partition_length_samples'] = (
+                            record['partition_length_samples']
+                        )
+                    
+                # append to list
+                record_list.append(stored_record)
+
+        else:
+            # loop list of  dictionaries and check stream num
+            for irecord, record in enumerate(self._record_list):
+                
+                event_stream_num = record.get('stream_number')
+                if event_stream_num is None:
+                    raise ValueError('ERROR: Wrong input event list format! It should '
+                                     'contain dictionaries with '
+                                     '"stream_number"!')
+                
+                event_stream_num = int(event_stream_num)
+                
+                # check if stream if resquested
+                if (stream_num is not None
+                    and event_stream_num != stream_num):
+                    continue
+
+                record['stream_num'] = event_stream_num
+                record_list.append(copy.deepcopy(record))
+                
+
+
+        # cut edge ot be able to extract trace
+        output_record_list = list()
+        if nb_samples is not None:
+            
+            for irecord, record in enumerate(record_list):
+                
+                min_index = 0
+                max_index = 0
+                full_length_samples = None
+                
+                if self._storage_format == 'hdf5':
+                    if 'segment_trigger_index':
+                        segment_trigger_index = record['segment_trigger_index']
+                        full_length_samples = record['segment_length_samples']
+                        min_index = segment_trigger_index - nb_pretrigger_samples
+                        max_index = min_index + nb_samples -1
+                else:
+
+                    if 'stream_trigger_index' in record:
+                        stream_trigger_index = record['stream_trigger_index']
+                        full_length_samples = record['stream_length_samples']
+                        min_index = stream_trigger_index - nb_pretrigger_samples
+                        max_index = min_index + nb_samples -1
+                    elif 'partition_trigger_index' in record:
+                        partition_trigger_index = record['partition_trigger_index']
+                        full_length_samples = record['partition_length_samples']
+                        min_index = partition_trigger_index - nb_pretrigger_samples
+                        max_index = min_index + nb_samples -1
+
+                if full_length_samples is not None:
+                    if (min_index < 0
+                        or max_index >= full_length_samples):
+                        continue # skip
+
+                output_record_list.append(record)
+        else:
+            output_record_list = record_list
+
+            
+        # check number of randoms
+        if (nrandoms is not None
+            and len(output_record_list)>nrandoms):
+                output_record_list = output_record_list[:nrandoms]
+                
+        return output_record_list
+            
+    
+    
+    def _load_dataframe(self, dataframe_path):
+        """
+        Load vaex dataframe (hdf5), path can be a file 
+        or directory
+        """
+
+        file_list = []
+
+        if os.path.isdir(dataframe_path):
+            file_list = glob(f'{dataframe_path}/*.hdf5')
+            if len(file_list) == 0:
+                raise ValueError('ERROR: No vaex file found. Check path!')
+            
+        elif os.path.isfile(dataframe_path):
+            if dataframe_path.find('.hdf5') != -1:
+                raise ValueError(
+                    f'ERROR: dataframe file {dataframe_path} '
+                    f'not recognized!'
+                )
+            file_list = [dataframe_path]
+
+
+        # get dataframe
+        dataframe = vx.open_many(file_list)
+        return dataframe
+
+
+    
+