eegprep 0.0.2__py3-none-any.whl

eegprep/eeg_mne2eeglab_epochs.py

@@ -0,0 +1,187 @@
+# Example to export MNE epochs to EEGLAB dataset
+# Events are not handled correctly in this example but it works
+
+import mne
+from mne.datasets import sample
+from mne.preprocessing import ICA
+import math
+
+import numpy as np
+from scipy.io import savemat
+
+# Load example data
+def eeg_mne2eeglab_epochs(epochs, ica):
+    
+    # export to EEGLAB dataset
+    data = epochs.get_data()  # Get the data from the epochs
+    n_epochs, n_channels, n_times = data.shape
+    ica_weights = ica.get_components()  # ICA weights (n_components x n_channels)
+    
+    # create identity matrix of size n_channels x n_channels
+    ica_sphere = np.eye(n_channels)  # ICA sphere (n_channels x n_channels)
+
+    # Compute the mixing matrix (inverse weights)
+    ica_inverse_weights = np.linalg.pinv(ica_weights)  # Shape: (n_channels, n_components)
+    
+    ica_channels = ica.info['ch_names']
+    raw_channels = epochs.info['ch_names']  # Assuming you have the raw object
+    ica_channel_indices = [raw_channels.index(ch) for ch in ica_channels]
+    ica_channel_indices = np.array(ica_channel_indices)
+    
+    ica_act = ica.get_sources(epochs).get_data(copy=True).transpose(1, 2, 0)  # Get the ICA activations
+   
+    print('Reference conversion may not be accurate...')
+    if 'custom_ref_applied' in epochs.info and epochs.info['custom_ref_applied']:
+        ref = 'common'  # Custom reference was applied
+    else:
+        ref = 'average'  # Default to average reference
+    
+    eeglab_dict = {
+        'setname'         : '',
+        'filename'        : '',
+        'filepath'        : '',
+        'subject'         : '',
+        'group'           : '',
+        'condition'       : '',
+        'session'         : np.array([]),
+        'comments'        : '',
+        'nbchan'          : n_channels,
+        'trials'          : n_epochs,
+        'pnts'            : n_times,
+        'srate'           : epochs.info['sfreq'],
+        'xmin'            : epochs.times[0],
+        'xmax'            : epochs.times[-1],
+        'times'           : epochs.times,
+        'data'            : data,
+        'icaact'          : ica_act,
+        'icawinv'         : ica_inverse_weights,
+        'icasphere'       : ica_weights,
+        'icaweights'      : ica_sphere,
+        'icachansind'     : ica_channel_indices,
+        'chanlocs'        : np.array([]),
+        'urchanlocs'      : np.array([]),
+        'chaninfo'        : np.array([]),
+        'ref'             : ref,
+        'event'           : np.array([]),
+        'urevent'         : np.array([]),
+        'eventdescription': np.array([]),
+        'epoch'           : np.array([]),
+        'epochdescription': np.array([]),
+        'reject'          : np.array([]),
+        'stats'           : np.array([]),
+        'specdata'        : np.array([]),
+        'specicaact'      : np.array([]),
+        'splinefile'      : np.array([]),
+        'icasplinefile'   : np.array([]),
+        'dipfit'          : np.array([]),
+        'history'         : np.array([]),
+        'saved'           : np.array([]),
+        'etc'             : np.array([]),
+        'datfile'         : np.array([]),
+        'run'             : np.array([]),
+        'roi'             : np.array([]),
+    }
+
+    # create channel locations
+    ch_names = epochs.ch_names
+    
+    ch_locs = epochs.info['chs']
+
+    theta_all = []
+    radius_all = []
+    sph_theta_all  = []
+    sph_phi_all    = []
+    sph_radius_all = []
+    X_all = []
+    Y_all = []
+    Z_all = []
+    for ch in ch_locs:
+        if 'loc' in ch and ch['loc'] is not None:
+            X_all.append(ch['loc'][1]*1000)
+            Y_all.append(-ch['loc'][0]*1000)
+            Z_all.append(ch['loc'][2]*1000)
+            hypotxy = math.hypot(X_all[-1],Y_all[-1])
+            sph_radius_all.append(math.hypot(hypotxy,Z_all[-1]))
+            
+            az = math.atan2(Y_all[-1],X_all[-1])/math.pi*180
+            horiz = math.atan2(Z_all[-1],hypotxy)/math.pi*180
+            
+            sph_theta_all.append(az)
+            sph_phi_all.append(horiz)
+
+            theta_all.append(-az) # warning inverse notation compared to MATLAB to match
+            radius_all.append(0.5 - horiz/180) # warning inverse notation compared to MATLAB to match
+        
+    d_list = [{
+        'labels': ch_name,
+        'theta': theta,
+        'radius': radius,
+        'X': X,
+        'Y': Y,
+        'Z': Z,
+        'sph_theta': sph_theta,
+        'sph_phi': sph_phi,
+        'sph_radius': sph_radius,
+        'type': 'EEG',
+        'urchan': 0,
+        'ref': ''
+    } for ch_name, theta, radius, X, Y, Z, sph_theta, sph_phi, sph_radius in zip(ch_names, theta_all, radius_all, X_all, Y_all, Z_all, sph_theta_all, sph_phi_all, sph_radius_all)]
+    # Create the list of dictionaries with a string field
+    # d_list = [{
+    #     'labels': ch_name,
+    #     'theta': math.atan2(ch_loc[0], ch_loc[1]),
+    #     'radius': math.hypot(ch_loc[1], ch_loc[0]),
+    #     'X': ch_loc[1]*1000,
+    #     'Y': ch_loc[0]*1000,
+    #     'Z': ch_loc[2]*1000,
+    #     'sph_theta': 0,
+    #     'sph_phi': 0,
+    #     'sph_radius': 0,
+    #     'type': 'EEG',
+    #     'urchan': 0,
+    #     'ref': ''
+    # } for ch_name, ch_loc in zip(ch_names, ch_locs_xyz)]
+
+    # convert d_list to a numpy array
+    d_list = np.array(d_list)
+    eeglab_dict['chanlocs'] = d_list
+    
+    # # Step 4: Save the EEGLAB dataset as a .mat file
+    return eeglab_dict
+
+    #print("EEGLAB dataset saved successfully!")
+    
+def test_eeg_mne2eeeglab_epochs():
+    sample_data_folder = mne.datasets.sample.data_path()
+    sample_data_raw_file = (
+        sample_data_folder / "MEG" / "sample" / "sample_audvis_filt-0-40_raw.fif"
+    )
+
+    raw = mne.io.read_raw_fif(sample_data_raw_file)
+
+    # extract data epochs    
+    events = mne.find_events(raw, stim_channel="STI 014")
+    event_dict = {
+        "auditory/left": 1,
+        "auditory/right": 2,
+        "visual/left": 3,
+        "visual/right": 4,
+        "smiley": 5,
+        "buttonpress": 32,
+    }
+    epochs = mne.Epochs(
+        raw,
+        events,
+        event_id=event_dict,
+        tmin=-0.2,
+        tmax=0.5,
+        preload=True,
+    )    
+
+    ica = ICA(n_components=15, random_state=97, max_iter=800)
+    ica.fit(raw)
+    
+    EEG = eeg_mne2eeglab_epochs(epochs, ica)
+    savemat('output_file.mat', EEG) # use pop_saveset
+
+# test_eeg_mne2eeeglab_epochs()

eegprep/eeg_rpsd.py

@@ -0,0 +1,65 @@
+import numpy as np
+from numpy.fft import fft
+from scipy.signal.windows import hamming
+
+def eeg_rpsd(EEG, nfreqs=None, pct_data=100):
+    # clean input cutoff freq
+    nyquist = EEG['srate'] // 2
+    if nfreqs is None or nfreqs > nyquist:
+        nfreqs = nyquist
+    nfreqs = int(nfreqs)
+    
+    # setup constants
+    ncomp = EEG['icaweights'].shape[0]
+    
+    # Hamming window
+    n_points = min(EEG['pnts'], EEG['srate'])
+    m = n_points
+    isOddLength = m % 2
+    if isOddLength:
+        x = np.arange(0, (m - 1) / 2 + 1) / (m - 1)
+    else:
+        x = np.arange(0, m / 2) / (m - 1)
+    
+    a = 0.54
+    window = a - (1 - a) * np.cos(2 * np.pi * x)
+    if isOddLength:
+        window = np.concatenate([window, window[-2::-1]])
+    else:
+        window = np.concatenate([window, window[::-1]])
+    
+    cutoff = (EEG['pnts'] // n_points) * n_points
+    index = np.add.outer(np.ceil(np.arange(0, cutoff - n_points + 1, n_points/2)).astype(int), np.arange(0, n_points)).astype(int).transpose()
+
+    np.random.seed(0)  # rng('default') in MATLAB
+    n_seg = index.shape[1] * EEG['trials']
+    subset = np.random.permutation(n_seg)[:int(n_seg * pct_data / 100)]
+    
+    # calculate windowed spectrums
+    psdmed = np.zeros((ncomp, nfreqs))
+    for it in range(ncomp):
+        temp = np.reshape(EEG['icaact'][it, index, :], (1, index.shape[0], index.shape[1] * EEG['trials']))
+        temp = temp[:, :, subset] * window[:, np.newaxis]
+        temp = fft(temp, n_points, axis=1)
+        temp = np.abs(temp) ** 2
+        temp = temp[:, 1:nfreqs + 1, :] * 2 / (EEG['srate'] * np.sum(window ** 2))
+        if nfreqs == nyquist:
+            temp[:, -1, :] /= 2
+        psdmed[it, :] = 20 * np.log10(np.median(temp, axis=2))
+    
+    return psdmed
+
+def test_eeg_rpsd():
+    EEG = {
+        'srate': 256,
+        'icaweights': np.random.randn(10, 256),
+        'pnts': 1000,
+        'trials': 5,
+        'icaact': np.random.randn(10, 1000, 5)
+    }
+    
+    psdmed = eeg_rpsd(EEG, 100)
+    assert psdmed.shape == (10, 100)
+    assert np.all(np.isfinite(psdmed))
+
+# test_eeg_rpsd()

eegprep/eeg_rpsd_compare_helper.py

@@ -0,0 +1,17 @@
+from eeg_rpsd import eeg_rpsd
+from pop_loadset import pop_loadset
+import sys
+
+if len(sys.argv) > 1:
+    eeglab_file_path = sys.argv[1]
+else:
+    eeglab_file_path = '/System/Volumes/Data/data/matlab/eeglab/sample_data/eeglab_data_epochs_ica.set'
+EEG = pop_loadset(eeglab_file_path)
+
+# Print the loaded data
+res = eeg_rpsd(EEG, 100)
+
+# save in a matlab file
+import scipy.io
+scipy.io.savemat('eeg_rpsd_data.mat', {'grid': res})
+print('Saved eeg_rpsd_data.mat')

eegprep/eeglabcompat.py

@@ -0,0 +1,117 @@
+from oct2py import Oct2Py, get_log
+from pop_saveset import pop_saveset
+from pop_loadset import pop_loadset
+import logging
+import os
+import sys
+
+eeglab = Oct2Py(logger=get_log())
+oc = Oct2Py(logger=get_log())
+eeglab.logger = get_log("new_log")
+    
+def pop_resample( EEG, freq): # 2 additional parameters in MATLAB (never used)
+    
+    pop_saveset(EEG, 'tmp.set') # 0.8 seconds
+    EEG2 = eeglab.pop_loadset('tmp.set') # 2 seconds
+    EEG2 = eeglab.pop_resample(EEG2, freq) # 2.4 seconds
+    eeglab.pop_saveset(EEG2, 'tmp2.set') # 2.4 seconds
+    EEG3 = pop_loadset('tmp2.set') # 0.2 seconds
+    
+    # delete temporary files
+    os.remove('tmp.set')
+    os.remove('tmp2.set')
+    return EEG3
+
+
+def pop_eegfiltnew(locutoff=None,hicutoff=None,revfilt=False,plotfreqz=False):
+    # error if locutoff and hicutoff are none
+    if locutoff==None and hicutoff==None:
+        raise('Cannot have low cutoff and high cutoff not defined')
+    
+    pop_saveset(EEG, 'tmp.set') # 0.8 seconds
+    EEG2 = eeglab.pop_loadset('tmp.set') # 2 seconds
+    EEG3 = pop_eegfiltnew(EEG2, 'locutoff',locutoff,'hicutoff',hicutoff,'revfilt',revfilt,'plotfreqz',plotfreqz)
+    eeglab.pop_saveset(EEG2, 'tmp2.set') # 2.4 seconds
+    EEG4 = pop_loadset('tmp2.set') # 0.2 seconds
+    
+    # delete temporary files
+    os.remove('tmp.set')
+    os.remove('tmp2.set')
+    return EEG4
+
+def clean_artifacts( EEG, ChannelCriterion=False, LineNoiseCriterion=False, FlatlineCriterion=False, BurstCriterion=False, WindowCriterion=0, Highpass=[0.25, 0.75], WindowCriterionTolerances=[float('-inf'), 8]):
+
+    if ChannelCriterion == False or ChannelCriterion == 'off':
+        ChannelCriterion='off'
+    else:
+        ChannelCriterion='on'
+        
+    if LineNoiseCriterion == False or LineNoiseCriterion == 'off':
+        LineNoiseCriterion='off'
+    else:
+        LineNoiseCriterion='on'
+    
+    if FlatlineCriterion == False or FlatlineCriterion == 'off':
+        FlatlineCriterion='off'
+    else:
+        FlatlineCriterion='on'
+
+    if BurstCriterion == False or BurstCriterion == 'off':
+        BurstCriterion='off'
+    else:
+        BurstCriterion='on'
+
+    pop_saveset(EEG, 'tmp.set') # 0.8 seconds
+    EEG2 = eeglab.pop_loadset('tmp.set') # 2 seconds
+    EEG3 = eeglab.clean_artifacts(EEG2, 'ChannelCriterion', ChannelCriterion, \
+        'LineNoiseCriterion', LineNoiseCriterion, \
+        'FlatlineCriterion', FlatlineCriterion, \
+        'BurstCriterion', BurstCriterion,  \
+        'WindowCriterion', WindowCriterion, \
+        'Highpass',Highpass, \
+        'WindowCriterionTolerances', WindowCriterionTolerances)
+    eeglab.pop_saveset(EEG3, 'tmp2.set') # 2.4 seconds
+    EEG4 = pop_loadset('tmp2.set') # 0.2 seconds
+    
+    # delete temporary files
+    os.remove('tmp.set')
+    os.remove('tmp2.set')
+    return EEG3
+
+# sys.exit()
+
+eeglab.logger.setLevel(logging.WARNING)
+eeglab.warning('off', 'backtrace')
+
+# On the command line, type "octave-8.4.0" OCTAVE_EXECUTABLE or OCTAVE var
+path2eeglab = '/System/Volumes/Data/data/matlab/eeglab/' # init >10 seconds
+eeglab.addpath(path2eeglab + '/functions/guifunc')
+eeglab.addpath(path2eeglab + '/functions/popfunc')
+eeglab.addpath(path2eeglab + '/functions/adminfunc')
+eeglab.addpath(path2eeglab + '/plugins/firfilt')
+eeglab.addpath(path2eeglab + '/functions/sigprocfunc')
+eeglab.addpath(path2eeglab + '/functions/miscfunc')
+eeglab.addpath(path2eeglab + '/plugins/dipfit')
+eeglab.addpath(path2eeglab + '/plugins/clean_rawdata')
+res = eeglab.version()
+print('Running EEGLAB commands in compatibility mode with Octave ' + res)
+eeglab.logger.setLevel(logging.INFO)
+
+eeglab_file_path = '/System/Volumes/Data/data/matlab/eeglab/sample_data/eeglab_data_epochs_ica.set'
+
+EEG = pop_loadset(eeglab_file_path)
+EEG = pop_eegfiltnew(EEG, locutoff=5,hicutoff=25,revfilt=True,plotfreqz=False)
+
+# EEG = eeglab.pop_loadset(eeglab_file_path)
+# TMPEEG = eeglab.pop_eegfiltnew(EEG, 'locutoff',5,'hicutoff',25,'revfilt',1,'plotfreqz',0)
+# CLEANEDEEG = eeglab.clean_artifacts(TMPEEG, 'ChannelCriterion', 'off', 
+#     'LineNoiseCriterion', 'off',
+#     'FlatlineCriterion', 'off',
+#     'BurstCriterion', 'off',
+#     'WindowCriterion', 0,
+#     'Highpass',[0.25, 0.75],
+#     'WindowCriterionTolerances', [-10000000, 8])
+
+# clean_artifacts( EEG, ChannelCriterion='on' )
+
+

eegprep/iclabel.py

@@ -0,0 +1,51 @@
+from .ICL_feature_extractor import ICL_feature_extractor
+from .iclabel_net import ICLabelNet
+import torch
+import numpy as np
+
+def iclabel(EEG, algorithm='default'):
+    # issue error if algorithm is not 'default'
+    if algorithm != 'default':
+        raise ValueError('Only the default algorithm is supported')
+    
+    #ICLABEL Extract ICLabel features from an EEG dataset.
+    features = ICL_feature_extractor(EEG, True)
+    
+    # Equivalent of MATLAB code reshaping
+    features[0] = np.single(np.concatenate([features[0],-features[0],features[0][:, ::-1, :, :],-features[0][:, ::-1, :, :]], axis=3))
+    features[1] = np.single(np.tile(features[1], (1, 1, 1, 4)))
+    features[2] = np.single(np.tile(features[2], (1, 1, 1, 4)))
+    print('Feature 0 shape:', features[0].shape)
+    print('Feature 1 shape:', features[1].shape)
+    print('Feature 2 shape:', features[2].shape)
+
+    # Load the ICLabelNet model
+    model = ICLabelNet('netICL.mat')
+    
+    # Convert the features to torch tensors
+    image = torch.tensor(features[0]).permute(-1, 2, 0, 1)
+    psdmed = torch.tensor(features[1]).permute(-1, 2, 0, 1)
+    autocorr = torch.tensor(features[2]).permute(-1, 2, 0, 1)
+    
+    # Get the output from the model
+    output = model(image, psdmed, autocorr)
+    output_np = output.detach().numpy()
+    output_np = output_np.T  # Transpose the array
+    output_np = np.reshape(output_np, (-1, 4), order='F')  # Reshape to have 4 columns
+    output_np = np.mean(output_np, axis=1)  # Compute the mean along the second axis (columns)
+    output_np = np.reshape(output_np, (7, -1), order='F')  # Reshape to have 7 rows
+    output_np = output_np.T  # Transpose back
+    
+    if 'ic_classification' not in EEG['etc']:
+        EEG['etc']['ic_classification'] = {}
+    if 'ICLabel' not in EEG['etc']['ic_classification']:
+        EEG['etc']['ic_classification']['ICLabel'] = {}
+        
+    EEG['etc']['ic_classification']['ICLabel']['classes'] = np.array(
+        ['Brain', 'Muscle', 'Eye', 'Heart', 'Line Noise', 'Channel Noise', 'Other'], 
+        dtype=np.object
+    )
+    EEG['etc']['ic_classification']['ICLabel']['classifications'] = output_np
+    EEG['etc']['ic_classification']['ICLabel']['version'] = 'default'
+    
+    return EEG

eegprep/iclabel_helper.py

@@ -0,0 +1,26 @@
+from iclabel import iclabel
+from pop_loadset import pop_loadset
+from pop_saveset import pop_saveset
+import sys
+
+# check if there are 2 arguments otherwise issue an error
+if len(sys.argv) != 3:
+    raise ValueError('Please provide the input and output file paths as command line arguments')
+
+# get first command line argument
+eeglab_file_path_in  = sys.argv[1]
+eeglab_file_path_out = sys.argv[2]
+
+EEG = pop_loadset(eeglab_file_path_in)
+# EEG['data']   = EEG['data'].reshape(EEG['data'].shape[0], -1)
+# EEG['icaact'] = EEG['icaact'].reshape(EEG['icaact'].shape[0], -1)
+# EEG['nbchan'] = EEG['data'].shape[0]
+# EEG['pnts']   = EEG['data'].shape[1]
+# EEG['trials'] = 1
+
+# Print the loaded data
+print('Data shape:', EEG['data'].shape)
+EEG = iclabel(EEG, 'default')
+
+# save dataset
+pop_saveset(EEG, eeglab_file_path_out)

eegprep/iclabel_net.py

@@ -0,0 +1,136 @@
+import scipy.io
+import torch
+import scipy
+import numpy as np
+
+class Reshape(torch.nn.Module):
+        def __init__(self, shape):
+            super().__init__()
+            self.shape = shape
+
+        def forward(self, x):
+            return x.view(x.shape[0], *self.shape)
+
+class Concatenate(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+    
+    def forward(self, x: list):
+        return torch.cat(x, dim=self.dim)
+
+
+class ICLabelNet(torch.nn.Module):
+    def __init__(self, mat_path):
+        super().__init__()
+        iclabel_matlab = scipy.io.loadmat(mat_path)
+        params = iclabel_matlab['params'][0]
+        i = 11
+        print('shape of param', i, torch.tensor(params[i][1]).shape)
+        self.discriminator_image_layer1_conv = torch.nn.Conv2d(in_channels=1, out_channels=128, kernel_size=4, stride=2, padding=1, dilation=1)
+        print(self.discriminator_image_layer1_conv.weight.shape)
+        self.discriminator_image_layer1_conv.weight = torch.nn.Parameter(torch.tensor(params[0][1]).permute(3, 2, 0, 1))
+        self.discriminator_image_layer1_conv.bias = torch.nn.Parameter(torch.tensor(params[1][1]).squeeze())
+        self.discriminator_image_layer1_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_image_layer2_conv = torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, dilation=1)
+        self.discriminator_image_layer2_conv.weight = torch.nn.Parameter(torch.tensor(params[2][1]).permute(3, 2, 0, 1))
+        self.discriminator_image_layer2_conv.bias = torch.nn.Parameter(torch.tensor(params[3][1]).squeeze())
+        self.discriminator_image_layer2_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_image_layer3_conv = torch.nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=2, padding=1, dilation=1)
+        self.discriminator_image_layer3_conv.weight = torch.nn.Parameter(torch.tensor(params[4][1]).permute(3, 2, 0, 1))
+        self.discriminator_image_layer3_conv.bias = torch.nn.Parameter(torch.tensor(params[5][1]).squeeze())
+        self.discriminator_image_layer3_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_psdmed_layer1_conv_conv = torch.nn.Conv2d(in_channels=1, out_channels=128, kernel_size=(1,3), stride=1, padding=(0,1), dilation=1)
+        self.discriminator_psdmed_layer1_conv_conv.weight = torch.nn.Parameter(torch.tensor(params[6][1]).permute(3, 2, 0, 1))
+        self.discriminator_psdmed_layer1_conv_conv.bias = torch.nn.Parameter(torch.tensor(params[7][1]).squeeze())
+        self.discriminator_psdmed_layer1_conv_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_psdmed_layer2_conv_conv = torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(1,3), stride=1, padding=(0,1), dilation=1)
+        self.discriminator_psdmed_layer2_conv_conv.weight = torch.nn.Parameter(torch.tensor(params[8][1]).permute(3, 2, 0, 1))
+        self.discriminator_psdmed_layer2_conv_conv.bias = torch.nn.Parameter(torch.tensor(params[9][1]).squeeze())
+        self.discriminator_psdmed_layer2_conv_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_psdmed_layer3_conv_conv = torch.nn.Conv2d(in_channels=256, out_channels=1, kernel_size=(1,3), stride=1, padding=(0,1), dilation=1)
+        self.discriminator_psdmed_layer3_conv_conv.weight = torch.nn.Parameter(torch.tensor(params[10][1]).unsqueeze(3).permute(3, 2, 0, 1))
+        self.discriminator_psdmed_layer3_conv_conv.bias = torch.nn.Parameter(torch.tensor(params[11][1]).squeeze(1))
+        self.discriminator_psdmed_layer3_conv_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_autocorr_layer1_conv_conv = torch.nn.Conv2d(in_channels=1, out_channels=128, kernel_size=(1,3), stride=1, padding=(0,1), dilation=1)
+        self.discriminator_autocorr_layer1_conv_conv.weight = torch.nn.Parameter(torch.tensor(params[12][1]).permute(3, 2, 0, 1))
+        self.discriminator_autocorr_layer1_conv_conv.bias = torch.nn.Parameter(torch.tensor(params[13][1]).squeeze())
+        self.discriminator_autocorr_layer1_conv_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_autocorr_layer2_conv_conv = torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(1,3), stride=1, padding=(0,1), dilation=1)
+        self.discriminator_autocorr_layer2_conv_conv.weight = torch.nn.Parameter(torch.tensor(params[14][1]).permute(3, 2, 0, 1))
+        self.discriminator_autocorr_layer2_conv_conv.bias = torch.nn.Parameter(torch.tensor(params[15][1]).squeeze())
+        self.discriminator_autocorr_layer2_conv_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_autocorr_layer3_conv_conv = torch.nn.Conv2d(in_channels=256, out_channels=1, kernel_size=(1,3), stride=1, padding=(0,1), dilation=1)
+        self.discriminator_autocorr_layer3_conv_conv.weight = torch.nn.Parameter(torch.tensor(params[16][1]).unsqueeze(3).permute(3, 2, 0, 1))
+        self.discriminator_autocorr_layer3_conv_conv.bias = torch.nn.Parameter(torch.tensor(params[17][1]).squeeze(1))
+        self.discriminator_autocorr_layer3_conv_relu = torch.nn.LeakyReLU(0.2)
+        self.discriminator_psdmed_reshape = Reshape((100, 1, 1))
+        self.discriminator_psdmed_concat1 = Concatenate(dim=2)
+        self.discriminator_psdmed_concat2 = Concatenate(dim=3)
+        self.discriminator_autocorr_reshape = Reshape((100, 1, 1))
+        self.discriminator_autocorr_concat1 = Concatenate(dim=2)
+        self.discriminator_autocorr_concat2 = Concatenate(dim=3)
+        self.discriminator_concat = Concatenate(dim=1)
+        self.discriminator_conv = torch.nn.Conv2d(in_channels=712, out_channels=7, kernel_size=4, stride=1, padding=0, dilation=1)
+        self.discriminator_conv.weight = torch.nn.Parameter(torch.tensor(params[18][1]).permute(3, 2, 0, 1))
+        self.discriminator_conv.bias = torch.nn.Parameter(torch.tensor(params[19][1]).squeeze())
+        self.discriminator_softmax = torch.nn.Softmax(dim=1)
+
+    def forward(self, image, psdmed, autocorr):
+        x_image = self.discriminator_image_layer1_conv(image)
+        x_image = self.discriminator_image_layer1_relu(x_image)
+        x_image = self.discriminator_image_layer2_conv(x_image)
+        x_image = self.discriminator_image_layer2_relu(x_image)
+        x_image = self.discriminator_image_layer3_conv(x_image)
+        x_image = self.discriminator_image_layer3_relu(x_image)
+        print('x_image', x_image.shape)
+
+        x_psdmed = self.discriminator_psdmed_layer1_conv_conv(psdmed)
+        x_psdmed = self.discriminator_psdmed_layer1_conv_relu(x_psdmed)
+        x_psdmed = self.discriminator_psdmed_layer2_conv_conv(x_psdmed)
+        x_psdmed = self.discriminator_psdmed_layer2_conv_relu(x_psdmed)
+        x_psdmed = self.discriminator_psdmed_layer3_conv_conv(x_psdmed)
+        x_psdmed = self.discriminator_psdmed_layer3_conv_relu(x_psdmed)
+        x_psdmed = self.discriminator_psdmed_reshape(x_psdmed)
+        x_psdmed = self.discriminator_psdmed_concat1([x_psdmed]*4)
+        x_psdmed = self.discriminator_psdmed_concat2([x_psdmed]*4)
+        print('x_psdmed', x_psdmed.shape)
+
+        x_autocorr = self.discriminator_autocorr_layer1_conv_conv(autocorr)
+        x_autocorr = self.discriminator_autocorr_layer1_conv_relu(x_autocorr)
+        x_autocorr = self.discriminator_autocorr_layer2_conv_conv(x_autocorr)
+        x_autocorr = self.discriminator_autocorr_layer2_conv_relu(x_autocorr)
+        x_autocorr = self.discriminator_autocorr_layer3_conv_conv(x_autocorr)
+        x_autocorr = self.discriminator_autocorr_layer3_conv_relu(x_autocorr)
+        x_autocorr = self.discriminator_autocorr_reshape(x_autocorr)
+        x_autocorr = self.discriminator_autocorr_concat1([x_autocorr]*4)
+        x_autocorr = self.discriminator_autocorr_concat2([x_autocorr]*4)
+        print('x_autocorr', x_autocorr.shape)
+
+        x = self.discriminator_concat([x_image, x_psdmed, x_autocorr])
+        x = self.discriminator_conv(x)
+        print('x', x.shape)
+        # subtract max value to avoid overflow
+        x = x - torch.max(x, dim=1, keepdim=True).values
+        x = self.discriminator_softmax(x)
+        
+        return x
+    
+# if __name__ == "__main__":
+#     model = ICLabelNet('netICL.mat')
+#     image_mat = scipy.io.loadmat('net_vars.mat')['in_image']
+#     psdmed_mat = scipy.io.loadmat('net_vars.mat')['in_psdmed']
+#     autocorr_mat = scipy.io.loadmat('net_vars.mat')['in_autocorr']
+#     # assuming third dimension is trivial and last dimension is channel. First two dimensions (32 x 32) are size of topoplot
+#     image = torch.tensor(image_mat).permute(-1, 2, 0, 1)
+#     print('image shape', image.shape)
+#     psdmed = torch.tensor(psdmed_mat).permute(-1, 2, 0, 1)
+#     print('psd shape', psdmed.shape)
+#     autocorr = torch.tensor(autocorr_mat).permute(-1, 2, 0, 1)
+#     print('autocorr shape', autocorr.shape)
+#     output = model(image, psdmed, autocorr)
+#     print(output.shape)
+    
+#     # save the output to a mat file
+#     scipy.io.savemat('output4.mat', {'output': output.detach().numpy()})
+