data.py

import mne
import matplotlib.pyplot as plt
import numpy as np
import scipy
from scipy.fft import fft, ifft, fftfreq, rfftfreq, rfft, irfft

class BVDataset:
    def __init__(self, path= '../../BRAINvISION_FILES/NZ0001/20220817/speech002.vhdr', start=100*2048, step_size=3*2048,scale=10e5,sampleRate=5):
        self.path=path        
        self.start = start
        self.step_size = step_size
        self.scale=scale
        self.sampleRate=sampleRate

        data = mne.io.read_raw_brainvision(self.path, preload=False)
        self.data_chunk, times = data.get_data(picks='all', return_times=True)
        self.info = data.info
    
    def getData(self,ch=range(0,100,10)):        
        #train and test data loading
        all_data = []
        for c in ch:  # Loop over c from 0 to 100    
            cur_data = self.scale*self.data_chunk[c][self.start : self.start + self.step_size:self.sampleRate]  # Extract data_chunk for current c
            all_data.append(cur_data)  # Append all_data to concatenated_data
        all_train_data = np.concatenate(all_data, axis=0)  # Concatenate all_data arrays into a single array
        return all_train_data
    def zeroMeanAndNotch(self,sequence):
        sequence = sequence - np.mean(sequence, axis=0)
        sequence=sequence.reshape(-1,1)
        sequence = self.notch_filter(sequence, self.info['sfreq'])

    def plot(self, data,title='train-data'):
        #plotting train and test data
        plt.figure()
        data_mean=np.mean(data)
        data_var=np.var(data)        
        fig, ax = plt.subplots(1, 1, figsize=(15, 7))
        ax.plot(data)
        ax.set_title(f"{title}: mean:{round(data_mean)},var:{round(data_var)}")        
       # plt.show()

    
    def notch_filter(self,X, f_sampling, f_notch=60, FILTER_HARMONICS=True, FFT=True):
        # Author: Alex Bujan
        # Cribbed and lightly modified from ecogVIS by JGM

        f_nyquist = f_sampling/2
        if FILTER_HARMONICS:
            notches = np.arange(f_notch, f_nyquist, f_notch)
        else:
            notches = np.array([f_notch])

        if FFT:
            fs = rfftfreq(X.shape[0], 1/f_sampling)
            delta = 1.0
            fd = rfft(X, axis=0)
        else:
            n_taps = 1001
            gain = [1, 1, 0, 0, 1, 1]
        for notch in notches:
            if FFT:
                window_mask = np.logical_and(fs > notch-delta, fs < notch+delta)
                window_size = window_mask.sum()
                window = np.hamming(window_size)
                fd[window_mask, :] *= (1 - window)[:, None]
            else:
                f_critical = np.array(
                    [0, notch-1, notch-.5, notch+.5, notch+1, f_nyquist]
                )/f_nyquist
                filt = scipy.signal.firwin2(n_taps, f_critical, gain)
                X = scipy.signal.filtfilt(filt, np.array([1]), X, axis=0)
        if FFT:
            X = irfft(fd, axis=0)
        return X