Depression_TMS/algorithm_V0/datacollect/eeg_quality_check-mat.py

# -*- coding: utf-8 -*-
"""
EEG Data Quality Check - eeg_data0.mat
===================================
1. Time Domain Signal (Full Duration)
2. Amplitude Spectrum (FFT)
3. Power Spectral Density (Linear Scale)
4. Power Spectral Density (dB Scale)
"""

import numpy as np
import matplotlib.pyplot as plt
import mne
from scipy import signal
from scipy.io import loadmat


def load_and_preprocess(filepath):
    """Load .mat file (custom format) and basic preprocessing."""
    mat_data = loadmat(filepath, simplify_cells=True)
    eeg = mat_data['eeg']

    # Extract data (shape: samples x channels)
    data = eeg['data'].T  # Transpose to (channels x samples)
    sfreq = eeg['sample_rate']

    # Get channel names (try multiple possible keys)
    if 'chn' in eeg:
        ch_names = list(eeg['chn'])
    elif 'electrode_name' in eeg:
        ch_names = list(eeg['electrode_name'])
    else:
        n_channels = data.shape[0]
        ch_names = [f'Ch{i+1}' for i in range(n_channels)]

    # Create MNE Info object
    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types='eeg')
    raw = mne.io.RawArray(data, info)

    raw.filter(l_freq=0.5, h_freq=10, fir_design='firwin', verbose=False)
    return raw


def main():
    filepath = r"D:\Ivey\Code_New_Proj\brainplot\plot64\eeg_data0511.mat"
    output_path = r"D:\Ivey\Code_New_Proj\brainplot\plot64\eeg_quality_check_depression.png"
    raw = load_and_preprocess(filepath)

    # Print all channel names first
    print(f"\nAvailable channels ({len(raw.ch_names)}):")
    for i, ch in enumerate(raw.ch_names):
        print(f"  {i:3d}: {ch}")

    select_channel = ['AIN5']
    raw.pick(select_channel)

    # Use all channels, full duration
    ch_names = raw.ch_names
    n_channels = len(ch_names)
    data = raw.get_data()
    sfreq = raw.info['sfreq']
    n_samples = data.shape[1]
    duration = n_samples / sfreq

    print(f"Info: {n_channels} channels, {duration:.1f}s, {sfreq:.0f} Hz")

    # Compute frequency domain data 
    n_fft = 2**int(np.ceil(np.log2(n_samples)))
    freqs_fft = np.fft.rfftfreq(n_fft, 1 / sfreq)
    fft_vals = np.fft.rfft(data, n=n_fft)
    amplitude = np.abs(fft_vals) / n_fft * 2

    freqs_psd, psd = signal.welch(data, fs=sfreq, nperseg=4096,
                                  noverlap=2048, scaling='density')

    # Frequency mask: 0.5-80 Hz
    mask_fft = (freqs_fft >= 0.5) & (freqs_fft <= 80)
    mask_psd = (freqs_psd >= 0.5) & (freqs_psd <= 80)

    freq_fft = freqs_fft[mask_fft]
    freq_psd = freqs_psd[mask_psd]

    # Plot: 4 rows x 1 column
    fig, axes = plt.subplots(4, 1, figsize=(16, 20))
    fig.suptitle(f'EEG Data Quality Check — {", ".join(ch_names)}, '
                 f'Full Duration: {duration:.1f}s',
                 fontsize=16, fontweight='bold', y=0.995)

    # Colormap for distinct channel  
    cmap = plt.cm.tab10 if n_channels <= 10 else plt.cm.tab20
    colors = [cmap(i) for i in np.linspace(0, 1, n_channels)]

    # ---- Row 1: Time Domain Signal ----
    ax = axes[0]
    offset = 0
    step = max(100, np.std(data, axis=1).mean() * 1e6 * 4)

    # Downsample for display
    ds = max(1, n_samples // (int(duration) * 500))
    t = np.arange(0, n_samples, ds) / sfreq

    for i in range(n_channels):
        sig = data[i, ::ds] * 1e6 + offset
        ax.plot(t, sig, linewidth=0.5, alpha=0.9, color=colors[i], label=ch_names[i])
        ax.text(t[0] - 0.5, offset, ch_names[i], fontsize=7, va='center', ha='right', color=colors[i])
        offset += step

    ax.set_xlim(0, duration)
    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Amplitude (μV)')
    ax.set_title('1. Time Domain Signal (Full Duration)', fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', fontsize=7, ncol=max(1, n_channels // 3), framealpha=0.8)

    # ---- Row 2: Amplitude Spectrum (FFT) ----
    ax = axes[1]
    amp_data = amplitude[:, mask_fft] * 1e6  # (n_channels, n_freqs)
    for i in range(n_channels):
        ax.plot(freq_fft, amp_data[i], color=colors[i], linewidth=1.0, alpha=0.85, label=ch_names[i])
    ax.axvline(50, color='red', linestyle='--', alpha=0.6, label='50 Hz Mains')
    ax.set_xlim(0.5, 30)
    ax.set_xlabel('Frequency (Hz)')
    ax.set_ylabel('Amplitude (μV)')
    ax.set_title('2. Amplitude Spectrum (FFT)', fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', fontsize=7, ncol=max(1, n_channels // 3), framealpha=0.8)

    # ---- Row 3: PSD (Linear Scale) ----
    ax = axes[2]
    psd_data = psd[:, mask_psd] * 1e12  # (n_channels, n_freqs)
    for i in range(n_channels):
        ax.plot(freq_psd, psd_data[i], color=colors[i], linewidth=1.0, alpha=0.85, label=ch_names[i])
    ax.axvline(50, color='red', linestyle='--', alpha=0.6, label='50 Hz Mains')
    ax.set_xlim(0.5, 80)
    ax.set_xlabel('Frequency (Hz)')
    ax.set_ylabel('Power (μV²/Hz)')
    ax.set_title('3. Power Spectral Density (Linear Scale)', fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', fontsize=7, ncol=max(1, n_channels // 3), framealpha=0.8)

    # ---- Row 4: PSD (dB Scale) ----
    ax = axes[3]
    for i in range(n_channels):
        psd_dbi = 10 * np.log10(psd_data[i] + 1e-20)
        ax.plot(freq_psd, psd_dbi, color=colors[i], linewidth=1.0, alpha=0.85, label=ch_names[i])
    ax.axvline(50, color='red', linestyle='--', alpha=0.6, label='50 Hz Mains')
    ax.set_xlim(0.5, 80)
    ax.set_xlabel('Frequency (Hz)')
    ax.set_ylabel('Power (dB)')
    ax.set_title('4. Power Spectral Density (dB Scale)', fontweight='bold')
    ax.grid(True, alpha=0.3)
    ax.legend(loc='upper right', fontsize=7, ncol=max(1, n_channels // 3), framealpha=0.8)

    plt.tight_layout()
    plt.subplots_adjust(top=0.97)

    plt.savefig(output_path, dpi=150, bbox_inches='tight',
                facecolor='white', edgecolor='none')
    print(f"Figure saved to: {output_path}")
    plt.show()


if __name__ == "__main__":
    main()



from scipy.io import loadmat
import numpy as np

mat = loadmat(r'D:\Ivey\Code_New_Proj\brainplot\plot64\eeg_data0511.mat', simplify_cells=True)
data = mat['eeg']['data']  # (samples, channels)
sfreq = 250
seg1 = data[0:int(10*sfreq), :]  # 0-10s
seg2 = data[int(10*sfreq):int(20*sfreq), :]  # 10-20s

print('Segment 1 (0-10s) shape:', seg1.shape)
print('Segment 2 (10-20s) shape:', seg2.shape)
print('Are they equal?', np.allclose(seg1, seg2))
print('Max difference:', np.max(np.abs(seg1 - seg2)))
print('Mean difference:', np.mean(np.abs(seg1 - seg2)))

# Check correlation
corr = np.corrcoef(seg1.flatten(), seg2.flatten())[0, 1]
print(f'Correlation: {corr:.4f}')