diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000..0f5169e --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,5 @@ +# Commits that only reformat code; git blame skips them. +# Enable locally with: git config blame.ignoreRevsFile .git-blame-ignore-revs + +# style: format the code base with ruff +99b337fe5f238f6d9244f4c6431f7070628de1d5 diff --git a/.github/workflows/Build.yml b/.github/workflows/Build.yml index 62aa039..8b4b725 100644 --- a/.github/workflows/Build.yml +++ b/.github/workflows/Build.yml @@ -3,6 +3,20 @@ name: CI on: [push, workflow_call] jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Check formatting with ruff + run: uv run --locked --only-group lint ruff format --check . + + - name: Check for obvious errors with ruff + run: uv run --locked --only-group lint ruff check . + build: runs-on: ubuntu-latest strategy: diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e3fea7..9ea3067 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Changed +- The whole code base is formatted with `ruff format` (line length 88). The change is purely cosmetic: the syntax tree of every file is unchanged, apart from whitespace inside docstrings, and the Python examples of `hypyp/sync/README.md` are formatted too. The vendored `hypyp/ext` and the tutorial notebooks are left untouched. The formatting commit is listed in `.git-blame-ignore-revs`, so `git blame` skips it (run `git config blame.ignoreRevsFile .git-blame-ignore-revs` once in your clone; GitHub applies it automatically) +- The CI now checks formatting (`ruff format --check`) and a small set of lint rules that only catch certain bugs: syntax errors, invalid comparisons and undefined names. Ruff comes from the new `lint` dependency group, which the `dev` group includes +- `black` is removed from the `dev` dependency group, since `ruff format` replaces it + ## [0.6.1] - 2026-10-03 ### Fixed diff --git a/hypyp/analyses.py b/hypyp/analyses.py index d13e058..73d94ce 100644 --- a/hypyp/analyses.py +++ b/hypyp/analyses.py @@ -35,36 +35,43 @@ from .sync import get_metric -def pow(epochs: mne.Epochs, fmin: float, fmax: float, n_fft: int, n_per_seg: int, epochs_average: bool) -> namedtuple: +def pow( + epochs: mne.Epochs, + fmin: float, + fmax: float, + n_fft: int, + n_per_seg: int, + epochs_average: bool, +) -> namedtuple: """ Computes the Power Spectral Density (PSD) on Epochs using Welch's method. - - This function calculates the power spectrum for each channel across the specified + + This function calculates the power spectrum for each channel across the specified frequency range. EOG channels are automatically dropped before computation. - + Parameters ---------- epochs : mne.Epochs A participant's Epochs object containing EEG data of shape (n_epochs, n_channels, n_times). - + fmin : float Minimum frequency in Hz to include in PSD calculation. - + fmax : float Maximum frequency in Hz to include in PSD calculation. - + n_fft : int - Length of FFT used. Must be >= n_per_seg. If larger, the segments will be + Length of FFT used. Must be >= n_per_seg. If larger, the segments will be zero-padded. If n_per_seg is None, n_fft must be <= number of time points. - + n_per_seg : int or None Length of each Welch segment (windowed with a Hamming window). If None, n_per_seg is set equal to n_fft. - + epochs_average : bool If True, PSD values are averaged over epochs. If False, PSD won't be averaged (the time course is maintained). - + Returns ------- psd_tuple : namedtuple @@ -73,15 +80,15 @@ def pow(epochs: mne.Epochs, fmin: float, fmax: float, n_fft: int, n_per_seg: int - psd: ndarray of PSD values with shape: - If epochs_average=True: (n_channels, n_frequencies) - If epochs_average=False: (n_epochs, n_channels, n_frequencies) - + Notes ----- This function can be iterated on groups and/or conditions: (for epochs in epochs['epochs_%s_%s_%s' % (subj, group, cond_name)]). - - The PSD values are computed using Welch's method with Hamming windows and + + The PSD values are computed using Welch's method with Hamming windows and expressed in µV²/Hz. - + Examples -------- >>> from mne import Epochs @@ -91,65 +98,82 @@ def pow(epochs: mne.Epochs, fmin: float, fmax: float, n_fft: int, n_per_seg: int """ # dropping EOG channels (incompatible with connectivity map model in stats) - for ch in epochs.info['chs']: - if ch['kind'] == 202: # FIFFV_EOG_CH - epochs.drop_channels([ch['ch_name']]) + for ch in epochs.info["chs"]: + if ch["kind"] == 202: # FIFFV_EOG_CH + epochs.drop_channels([ch["ch_name"]]) # computing power spectral density on epochs signal # average in the 1-second window around event (mean, but can choose 'median') - kwargs = dict(fmin=fmin, fmax=fmax, n_fft=n_fft, n_per_seg=n_per_seg, - tmin=None, tmax=None, method='welch', picks='all', exclude=[], - proj=False, remove_dc=True, n_jobs=1) + kwargs = dict( + fmin=fmin, + fmax=fmax, + n_fft=n_fft, + n_per_seg=n_per_seg, + tmin=None, + tmax=None, + method="welch", + picks="all", + exclude=[], + proj=False, + remove_dc=True, + n_jobs=1, + ) spectrum = EpochsSpectrum(epochs, **kwargs) psds = spectrum.get_data() freq_list = spectrum.freqs - + if epochs_average is True: # averaging power across epochs for each channel ch and each frequency f psd = np.mean(psds, axis=0) else: psd = psds - psd_tuple = namedtuple('PSD', ['freq_list', 'psd']) + psd_tuple = namedtuple("PSD", ["freq_list", "psd"]) - return psd_tuple(freq_list=freq_list, - psd=psd) + return psd_tuple(freq_list=freq_list, psd=psd) -def behav_corr(data: np.ndarray, behav: np.ndarray, data_name: str, behav_name: str, - p_thresh: float, multiple_corr: bool = True, verbose: bool = False) -> namedtuple: +def behav_corr( + data: np.ndarray, + behav: np.ndarray, + data_name: str, + behav_name: str, + p_thresh: float, + multiple_corr: bool = True, + verbose: bool = False, +) -> namedtuple: """ Correlates neural data with behavioral parameters using appropriate linear correlation methods. - - This function first checks for data normality, then applies either Pearson's or + + This function first checks for data normality, then applies either Pearson's or Spearman's correlation. For connectivity data, it can apply multiple comparison correction. - + Parameters ---------- data : np.ndarray Data to correlate with behavior. Can be either: - 1D array of values (same shape as behav) - 3D array of connectivity values with shape (n_dyads, n_channels, n_channels) - + behav : np.ndarray 1D array of behavioral values (e.g., timing, performance scores) - + data_name : str Description of the data (used for plot labeling if verbose=True) - + behav_name : str Description of the behavioral parameter (used for plot labeling if verbose=True) - + p_thresh : float Significance threshold for correlation tests (typically 0.05) - + multiple_corr : bool, optional Whether to apply multiple comparison correction (default=True). Uses the FDR-BH method when True. - + verbose : bool, optional Whether to generate visualization of the correlation (default=False) - + Returns ------- corr_tuple : namedtuple @@ -160,12 +184,12 @@ def behav_corr(data: np.ndarray, behav: np.ndarray, data_name: str, behav_name: - pvalue: p-value(s) (original or corrected if multiple_corr=True) - strat: string indicating which strategy was used ('normal', 'non_normal', or information about multiple comparison correction) - + Notes ----- - When correlating connectivity matrices with behavioral data, only the significant + When correlating connectivity matrices with behavioral data, only the significant correlations are returned in the r matrix, with non-significant values set to zero. - + Examples -------- >>> # Correlating average alpha power with reaction times @@ -175,7 +199,7 @@ def behav_corr(data: np.ndarray, behav: np.ndarray, data_name: str, behav_name: """ # storage for results - corr_tuple = namedtuple('corr_tuple', ['r', 'pvalue', 'strat']) + corr_tuple = namedtuple("corr_tuple", ["r", "pvalue", "strat"]) # simple correlation between vectors (data can be averaged PSD for example) if data.shape == behav.shape: @@ -185,18 +209,18 @@ def behav_corr(data: np.ndarray, behav: np.ndarray, data_name: str, behav_name: if min(pvalue1, pvalue2) < 0.05: # reject null hypothesis # (H0: data come from a normal distribution) - strat = 'non_normal' + strat = "non_normal" r, pvalue = scipy.stats.spearmanr(behav, data, axis=0) else: - strat = 'normal' + strat = "normal" r, pvalue = scipy.stats.pearsonr(behav, data) # can also use np.convolve, np.correlate, np.corrcoeff # vizualisation if verbose: plt.figure() plt.scatter(behav, data, label=str(r) + str(pvalue)) - plt.legend(loc='upper right') - plt.title('Linear correlation between ' + behav_name + ' and ' + data_name) + plt.legend(loc="upper right") + plt.title("Linear correlation between " + behav_name + " and " + data_name) plt.xlabel(behav_name) plt.ylabel(data_name) plt.show() @@ -217,11 +241,13 @@ def behav_corr(data: np.ndarray, behav: np.ndarray, data_name: str, behav_name: # correction for multiple comparisons if multiple_corr is True: # note: we reshape pvals to be able to use fdr_bh - pvals_corrected = statsmodels.stats.multitest.multipletests(pvals.flatten(), - alpha=0.05, - method='fdr_bh', - is_sorted=False, - returnsorted=False) + pvals_corrected = statsmodels.stats.multitest.multipletests( + pvals.flatten(), + alpha=0.05, + method="fdr_bh", + is_sorted=False, + returnsorted=False, + ) # put pval in original shape pvals_corrected = np.reshape(np.atleast_1d(pvals_corrected[1]), pvals.shape) # get r value for significant correlation only @@ -236,44 +262,44 @@ def behav_corr(data: np.ndarray, behav: np.ndarray, data_name: str, behav_name: if pvalue[i, j] < p_thresh: significant_corr[i, j] = rs[i, j] r = np.nan_to_num(significant_corr) - strat = 'correction for multiple comaprison ' + str(multiple_corr) + strat = "correction for multiple comaprison " + str(multiple_corr) return corr_tuple(r=r, pvalue=pvalue, strat=strat) def indices_connectivity_intrabrain(epochs: mne.Epochs) -> List[Tuple[int, int]]: """ Computes indices for intrabrain connectivity analysis between all EEG channels. - + This function generates all possible pairs of EEG channel indices for a single participant, excluding EOG channels. - + Parameters ---------- epochs : mne.Epochs A participant's Epochs object containing channel information - + Returns ------- channels : List[Tuple[int, int]] List of tuples, each containing a pair of channel indices (i, j) where i < j. These indices correspond to all possible pairs of EEG channels for the participant. - + Notes ----- This function automatically removes EOG channels before generating the pairs. The resulting indices can be used as input for connectivity analyses within a single brain. - + Examples -------- >>> channel_pairs = indices_connectivity_intrabrain(participant_epochs) >>> print(f"Number of channel pairs: {len(channel_pairs)}") """ - names = copy.deepcopy(epochs.info['ch_names']) - for ch in epochs.info['chs']: - if ch['kind'] == FIFF.FIFFV_EOG_CH: - names.remove(ch['ch_name']) + names = copy.deepcopy(epochs.info["ch_names"]) + for ch in epochs.info["chs"]: + if ch["kind"] == FIFF.FIFFV_EOG_CH: + names.remove(ch["ch_name"]) n = len(names) bin = 0 @@ -292,31 +318,31 @@ def indices_connectivity_intrabrain(epochs: mne.Epochs) -> List[Tuple[int, int]] def indices_connectivity_interbrain(epoch_hyper: mne.Epochs) -> List[Tuple[int, int]]: """ Computes indices for interbrain connectivity analysis between EEG channels of two participants. - + This function generates all possible pairs of channel indices where the first channel belongs to participant 1 and the second to participant 2, based on a merged epochs object. - + Parameters ---------- epoch_hyper : mne.Epochs A merged Epochs object containing data from both participants, with channels ordered such that the first half belongs to participant 1 and the second half to participant 2. - + Returns ------- channels : List[Tuple[int, int]] List of tuples, each containing a pair of channel indices (i, j) where: - i is an index of a channel from participant 1 (in range 0 to n_channels/2 - 1) - j is an index of a channel from participant 2 (in range n_channels/2 to n_channels - 1) - + Notes ----- This function assumes that the channels in epoch_hyper are organized as: [participant1_ch1, ..., participant1_chN, participant2_ch1, ..., participant2_chN] - + EOG channels are automatically excluded before generating the pairs. - + Examples -------- >>> interbrain_pairs = indices_connectivity_interbrain(merged_epochs) @@ -324,10 +350,10 @@ def indices_connectivity_interbrain(epoch_hyper: mne.Epochs) -> List[Tuple[int, """ channels = [] - names = copy.deepcopy(epoch_hyper.info['ch_names']) - for ch in epoch_hyper.info['chs']: - if ch['kind'] == FIFF.FIFFV_EOG_CH: - names.remove(ch['ch_name']) + names = copy.deepcopy(epoch_hyper.info["ch_names"]) + for ch in epoch_hyper.info["chs"]: + if ch["kind"] == FIFF.FIFFV_EOG_CH: + names.remove(ch["ch_name"]) l = list(range(0, int(len(names) / 2))) # l = list(range(0,62)) @@ -342,30 +368,34 @@ def indices_connectivity_interbrain(epoch_hyper: mne.Epochs) -> List[Tuple[int, return channels -def pair_connectivity(data: Union[list, np.ndarray], sampling_rate: int, - frequencies: Union[dict, list], mode: str, - epochs_average: bool = True) -> np.ndarray: +def pair_connectivity( + data: Union[list, np.ndarray], + sampling_rate: int, + frequencies: Union[dict, list], + mode: str, + epochs_average: bool = True, +) -> np.ndarray: """ Computes frequency-domain connectivity measures between two participants. - - This function is a high-level interface that processes EEG data, computes analytic + + This function is a high-level interface that processes EEG data, computes analytic signals, and calculates connectivity metrics between all possible channel pairs. - + Parameters ---------- data : Union[list, np.ndarray] EEG data from two participants with shape (2, n_epochs, n_channels, n_times) - + sampling_rate : int Sampling rate of the EEG data in Hz - + frequencies : Union[dict, list] Specification of frequency bands of interest: - If dict: {'band_name': [fmin, fmax], ...} defining frequency bands e.g., {'alpha': [8, 12], 'beta': [12, 20]} - If list: [fmin, fmax] for each integer frequency in the range e.g., [5, 30] for all integer frequencies from 5 to 30 Hz - + mode : str Connectivity measure to compute. Options: - 'envelope_corr': envelope correlation @@ -377,56 +407,60 @@ def pair_connectivity(data: Union[list, np.ndarray], sampling_rate: int, - 'imaginary_coh': imaginary coherence - 'pli': phase lag index - 'wpli': weighted phase lag index - + epochs_average : bool, optional If True, connectivity values are averaged across epochs (default) If False, epoch-by-epoch connectivity is preserved - + Returns ------- result : np.ndarray Connectivity matrix with shape: - If epochs_average=True: (n_freq, 2*n_channels, 2*n_channels) - If epochs_average=False: (n_freq, n_epochs, 2*n_channels, 2*n_channels) - + The channels are ordered as: [participant1_ch1, ..., participant1_chN, participant2_ch1, ..., participant2_chN] - + To extract only interbrain connectivity values, slice the last two dimensions: result[:, :, 0:n_channels, n_channels:2*n_channels] - + Notes ----- This function handles the complete process from raw EEG data to connectivity metrics: 1. Computes analytic signals for the specified frequencies 2. Applies the selected connectivity measure - + Each connectivity metric has different mathematical properties and interpretations: - PLV measures consistency in phase differences regardless of amplitude - Coherence measures linear relationship between signals in frequency domain - Imaginary coherence is less susceptible to volume conduction - PLI and wPLI are robust against common sources and amplitude effects - + Time complexity scales with O(n_channels^2 × n_epochs × n_frequencies × n_times) - + Examples -------- >>> # Computing alpha band PLV >>> alpha_plv = pair_connectivity( - ... [subj1_data, subj2_data], - ... sampling_rate=256, - ... frequencies={'alpha': [8, 12]}, + ... [subj1_data, subj2_data], + ... sampling_rate=256, + ... frequencies={'alpha': [8, 12]}, ... mode='plv' ... ) """ # Data consists of two lists of np.array (n_epochs, n_channels, epoch_size) - assert data[0].shape[0] == data[1].shape[0], "Two streams much have the same lengths." + assert data[0].shape[0] == data[1].shape[0], ( + "Two streams much have the same lengths." + ) # compute instantaneous analytic signal from EEG data if type(frequencies) == list: # average over tapers - values = np.mean(compute_single_freq(data, sampling_rate, frequencies), 3).squeeze() + values = np.mean( + compute_single_freq(data, sampling_rate, frequencies), 3 + ).squeeze() elif type(frequencies) == dict: values = compute_freq_bands(data, sampling_rate, frequencies) else: @@ -438,9 +472,13 @@ def pair_connectivity(data: Union[list, np.ndarray], sampling_rate: int, return result -def compute_sync(complex_signal: np.ndarray, mode: str, epochs_average: bool = True, - optimization: Optional[str] = None, - priority: Optional[list] = None) -> np.ndarray: +def compute_sync( + complex_signal: np.ndarray, + mode: str, + epochs_average: bool = True, + optimization: Optional[str] = None, + priority: Optional[list] = None, +) -> np.ndarray: """ Computes frequency-domain connectivity measures from analytic signals. @@ -530,25 +568,33 @@ def compute_sync(complex_signal: np.ndarray, mode: str, epochs_average: bool = T For full mathematical details and references, see ``hypyp/sync/README.md``. """ - n_epoch, n_ch, n_freq, n_samp = complex_signal.shape[1], complex_signal.shape[2], \ - complex_signal.shape[3], complex_signal.shape[4] + n_epoch, n_ch, n_freq, n_samp = ( + complex_signal.shape[1], + complex_signal.shape[2], + complex_signal.shape[3], + complex_signal.shape[4], + ) # calculate all epochs at once, the only downside is that the disk may not have enough space - complex_signal = complex_signal.transpose((1, 3, 0, 2, 4)).reshape(n_epoch, n_freq, 2 * n_ch, n_samp) + complex_signal = complex_signal.transpose((1, 3, 0, 2, 4)).reshape( + n_epoch, n_freq, 2 * n_ch, n_samp + ) transpose_axes = (0, 1, 3, 2) # Normalize mode names (handle aliases) mode_lower = mode.lower() mode_map = { - 'envelope_corr': 'envcorr', - 'pow_corr': 'powcorr', - 'imaginary_coh': 'imcoh', + "envelope_corr": "envcorr", + "pow_corr": "powcorr", + "imaginary_coh": "imcoh", } mode_normalized = mode_map.get(mode_lower, mode_lower) # Get the metric from the sync module try: - metric = get_metric(mode_normalized, optimization=optimization, priority=priority) + metric = get_metric( + mode_normalized, optimization=optimization, priority=priority + ) con = metric.compute(complex_signal, n_samp, transpose_axes) except ValueError: raise ValueError(f'Metric type "{mode}" not supported.') @@ -560,76 +606,80 @@ def compute_sync(complex_signal: np.ndarray, mode: str, epochs_average: bool = T return con -def compute_conn_mvar(complex_signal: np.ndarray, mvar_params: dict, - ica_params: dict, measure_params: dict, - check_stability: bool = True) -> np.ndarray: +def compute_conn_mvar( + complex_signal: np.ndarray, + mvar_params: dict, + ica_params: dict, + measure_params: dict, + check_stability: bool = True, +) -> np.ndarray: """ Computes connectivity measures based on multivariate autoregressive (MVAR) modeling. - - This function fits MVAR models to the data and computes directed connectivity + + This function fits MVAR models to the data and computes directed connectivity measures based on those models, with optional ICA for source separation. - + Parameters ---------- complex_signal : np.ndarray Complex analytic signals with shape (2, n_epochs, n_channels, n_freq_bins, n_times) - + mvar_params : dict Dictionary of MVAR model parameters with keys: - 'mvar_order': int, order of the MVAR model - 'fitting_method': str, method for fitting the MVAR model - 'delta': float, regularization parameter - + ica_params : dict Dictionary of ICA parameters with keys: - 'method': str, ICA algorithm to use - 'random_state': int, random seed for reproducibility - + measure_params : dict Dictionary defining the connectivity measure to compute with keys: - 'name': str, name of the connectivity measure (e.g., 'PDC', 'DTF') - 'n_fft': int, number of FFT points - + check_stability : bool, optional Whether to verify the stability of the MVAR model (default=True) - If True, the function will check if the model is stable and prompt for + If True, the function will check if the model is stable and prompt for continuation or epoch merging if necessary. - + Returns ------- connectivity : np.ndarray Connectivity measure matrix with shape: (n_epochs, n_freq, n_channels, n_channels, n_fft) or (1, n_freq, n_channels, n_channels, n_fft) if epochs are merged - + Notes ----- - MVAR-based connectivity measures provide information about directed (causal) - relationships between signals, unlike most of the other connectivity metrics + MVAR-based connectivity measures provide information about directed (causal) + relationships between signals, unlike most of the other connectivity metrics implemented in this module. - + The function relies on the MVARICA approach: 1. Apply ICA to separate independent sources 2. Fit an MVAR model to the sources 3. Compute connectivity measures in the source space 4. Back-project to the original signal space - + For stability, the number of time samples should be substantially larger than: mvar_order × n_channels^2 - + Warnings -------- This function is computationally intensive, especially for high model orders and large numbers of channels. - - If the MVAR model is unstable, the function will offer to merge epochs to + + If the MVAR model is unstable, the function will offer to merge epochs to increase the sample size for more stable estimation. - + References ---------- - Baccalá, L. A., & Sameshima, K. (2001). Partial directed coherence: a new + Baccalá, L. A., & Sameshima, K. (2001). Partial directed coherence: a new concept in neural structure determination. Biological cybernetics, 84(6), 463-474. - + Examples -------- >>> mvar_params = {'mvar_order': 5, 'fitting_method': 'least_squares', 'delta': 0.0} @@ -638,14 +688,22 @@ def compute_conn_mvar(complex_signal: np.ndarray, mvar_params: dict, >>> conn = compute_conn_mvar(complex_signal, mvar_params, ica_params, measure_params) """ - n_epoch, n_ch, n_freq, n_samp = complex_signal.shape[1], complex_signal.shape[2], \ - complex_signal.shape[3], complex_signal.shape[4] + n_epoch, n_ch, n_freq, n_samp = ( + complex_signal.shape[1], + complex_signal.shape[2], + complex_signal.shape[3], + complex_signal.shape[4], + ) - complex_signal = complex_signal.transpose((1, 3, 0, 2, 4)).reshape(n_epoch, n_freq, 2 * n_ch, n_samp) + complex_signal = complex_signal.transpose((1, 3, 0, 2, 4)).reshape( + n_epoch, n_freq, 2 * n_ch, n_samp + ) real_signal = np.real(complex_signal) aug_signal = real_signal[np.newaxis, ...] - mvar = MVAR(mvar_params["mvar_order"], mvar_params["fitting_method"], mvar_params["delta"]) + mvar = MVAR( + mvar_params["mvar_order"], mvar_params["fitting_method"], mvar_params["delta"] + ) if check_stability: c = True @@ -653,193 +711,247 @@ def compute_conn_mvar(complex_signal: np.ndarray, mvar_params: dict, is_stable = fit_mvar.stability() while c: - if is_stable: print("MVAR model is stable.") inp = input("Do you want to continue? ") if inp.lower() == "yes": - aux_3 = np.zeros((aug_signal.shape[1], aug_signal.shape[2], real_signal.shape[2], - real_signal.shape[2], measure_params["n_fft"]), dtype=np.complex128) + aux_3 = np.zeros( + ( + aug_signal.shape[1], + aug_signal.shape[2], + real_signal.shape[2], + real_signal.shape[2], + measure_params["n_fft"], + ), + dtype=np.complex128, + ) for e in range(aug_signal.shape[1]): for f in range(aug_signal.shape[2]): aux_sig = aug_signal[:, e, f, :, :] - conn_freq_epoch = connectivity_mvarica(real_signal=aux_sig, ica_params=ica_params, - measure_name=measure_params["name"], - n_fft=measure_params["n_fft"], var_model=mvar) + conn_freq_epoch = connectivity_mvarica( + real_signal=aux_sig, + ica_params=ica_params, + measure_name=measure_params["name"], + n_fft=measure_params["n_fft"], + var_model=mvar, + ) d_type = conn_freq_epoch.dtype aux_3[e, f, :, :, :] = conn_freq_epoch return np.asarray(np.real(aux_3), dtype=d_type) else: - return None else: - counter = 0 if counter == 0: - - print("MVAR model is not stable: number of time samples may be too small!") + print( + "MVAR model is not stable: number of time samples may be too small!" + ) print("\n") - nes_sample = mvar_params["mvar_order"] * real_signal.shape[2] * real_signal.shape[2] - print("At least " + str(nes_sample) + " samples are required for fitting MVAR model.") + nes_sample = ( + mvar_params["mvar_order"] + * real_signal.shape[2] + * real_signal.shape[2] + ) + print( + "At least " + + str(nes_sample) + + " samples are required for fitting MVAR model." + ) print("\n") inp = input("Do you want to merge the epochs?") if inp.lower() == "yes": - - merged_signal = aug_signal.reshape(1, real_signal.shape[1], real_signal.shape[2], - real_signal.shape[3] * real_signal.shape[0]) + merged_signal = aug_signal.reshape( + 1, + real_signal.shape[1], + real_signal.shape[2], + real_signal.shape[3] * real_signal.shape[0], + ) fit_mvar = mvar.fit(merged_signal[:, 0, 0, :][np.newaxis, ...]) is_stable = fit_mvar.stability() aug_signal = merged_signal[np.newaxis, ...] counter += 1 else: - return None else: - return "epochs are already merged" else: - - aux_3 = np.zeros((aug_signal.shape[1], aug_signal.shape[2], real_signal.shape[2], real_signal.shape[2], - measure_params["n_fft"]), dtype=np.complex128) + aux_3 = np.zeros( + ( + aug_signal.shape[1], + aug_signal.shape[2], + real_signal.shape[2], + real_signal.shape[2], + measure_params["n_fft"], + ), + dtype=np.complex128, + ) for e in range(aug_signal.shape[1]): for f in range(aug_signal.shape[2]): aux_sig = aug_signal[:, e, f, :, :] - conn_freq_epoch = connectivity_mvarica(real_signal=aux_sig, ica_params=ica_params, - measure_name=measure_params["name"], - n_fft=measure_params["n_fft"], var_model=mvar) + conn_freq_epoch = connectivity_mvarica( + real_signal=aux_sig, + ica_params=ica_params, + measure_name=measure_params["name"], + n_fft=measure_params["n_fft"], + var_model=mvar, + ) d_type = conn_freq_epoch.dtype aux_3[e, f, :, :, :] = conn_freq_epoch return np.asarray(aux_3, dtype=d_type) -def compute_single_freq(data: np.ndarray, sampling_rate: int, freq_range: List[float]) -> np.ndarray: +def compute_single_freq( + data: np.ndarray, sampling_rate: int, freq_range: List[float] +) -> np.ndarray: """ Computes analytic signals for each frequency in the specified range using multitaper method. - - This function calculates complex-valued time-frequency representations for each + + This function calculates complex-valued time-frequency representations for each integer frequency in the given range. - + Parameters ---------- data : np.ndarray EEG data from two participants with shape (2, n_epochs, n_channels, n_times) - + sampling_rate : int Sampling rate of the EEG data in Hz - + freq_range : List[float] A list [fmin, fmax] specifying the frequency range to analyze. Every integer frequency from fmin to fmax will be included. - + Returns ------- complex_signal : np.ndarray Complex-valued analytic signals with shape: (2, n_epochs, n_channels, n_tapers, n_frequencies, n_times) - + Notes ----- This function uses MNE's implementation of the multitaper method with 4 cycles per frequency, which provides a good balance between time and frequency resolution. - + Time-frequency decomposition is calculated for each participant separately. - + Examples -------- >>> complex_tf = compute_single_freq(data, sampling_rate=256, freq_range=[8, 30]) >>> print(f"Shape: {complex_tf.shape}") # (2, n_epochs, n_channels, n_tapers, 23, n_times) """ - complex_signal = np.array([mne.time_frequency.tfr_array_multitaper(data[participant], sfreq=sampling_rate, - freqs=np.arange( - freq_range[0], freq_range[1], 1), - n_cycles=4, zero_mean=False, use_fft=True, - decim=1, - output='complex') - for participant in range(2)]) + complex_signal = np.array( + [ + mne.time_frequency.tfr_array_multitaper( + data[participant], + sfreq=sampling_rate, + freqs=np.arange(freq_range[0], freq_range[1], 1), + n_cycles=4, + zero_mean=False, + use_fft=True, + decim=1, + output="complex", + ) + for participant in range(2) + ] + ) return complex_signal -def compute_freq_bands(data: np.ndarray, sampling_rate: int, freq_bands: dict, - filter_signal: bool = True, **filter_options) -> np.ndarray: +def compute_freq_bands( + data: np.ndarray, + sampling_rate: int, + freq_bands: dict, + filter_signal: bool = True, + **filter_options, +) -> np.ndarray: """ Computes analytic signals for specified frequency bands using FIR filtering and Hilbert transform. - + This function performs bandpass filtering for each defined frequency band, then applies the Hilbert transform to obtain complex-valued analytic signals. - + Parameters ---------- data : np.ndarray EEG data from two participants with shape (2, n_epochs, n_channels, n_times) - + sampling_rate : int Sampling rate of the EEG data in Hz - + freq_bands : dict Dictionary defining frequency bands: {'band_name': [fmin, fmax], ...} e.g., {'alpha': [8, 12], 'beta': [12, 20]} - + filter_signal : bool, optional Whether to apply bandpass filtering (default=True) If False, the Hilbert transform is applied directly to the input signals - + **filter_options Additional arguments for mne.filter.filter_data, such as: - filter_length: Length of the FIR filter in samples or 's' for seconds - l_trans_bandwidth: Width of the transition band at the low cut-off frequency - h_trans_bandwidth: Width of the transition band at the high cut-off frequency - + Returns ------- complex_signal : np.ndarray Complex-valued analytic signals with shape: (2, n_epochs, n_channels, n_freq_bands, n_times) where n_freq_bands corresponds to the number of frequency bands in freq_bands - + Notes ----- Unlike compute_single_freq which uses the multitaper method, this function: 1. Applies a bandpass filter to isolate the specified frequency band 2. Uses the Hilbert transform to compute the analytic signal - - This approach may be computationally more efficient when analyzing broader + + This approach may be computationally more efficient when analyzing broader frequency bands rather than individual frequencies. - - The order of frequency bands in the output corresponds to their order in the + + The order of frequency bands in the output corresponds to their order in the freq_bands dictionary. - + Examples -------- >>> bands = {'theta': [4, 8], 'alpha': [8, 13], 'beta': [13, 30]} - >>> complex_signals = compute_freq_bands(data, 256, bands, + >>> complex_signals = compute_freq_bands(data, 256, bands, ... filter_length='auto', l_trans_bandwidth=1) """ - assert data[0].shape[0] == data[1].shape[0], "Two data streams should have the same number of trials." + assert data[0].shape[0] == data[1].shape[0], ( + "Two data streams should have the same number of trials." + ) data = np.array(data) # filtering and Hilbert transform complex_signal = [] for freq_band in freq_bands.values(): if filter_signal: - filtered = np.array([mne.filter.filter_data(data[participant], - sampling_rate, l_freq=freq_band[0], h_freq=freq_band[1], - **filter_options, - verbose=False) - for participant in range(2) - # for each participant - ]) + filtered = np.array( + [ + mne.filter.filter_data( + data[participant], + sampling_rate, + l_freq=freq_band[0], + h_freq=freq_band[1], + **filter_options, + verbose=False, + ) + for participant in range(2) + # for each participant + ] + ) else: - filtered=np.array([data[participant] for participant in range(2)]) + filtered = np.array([data[participant] for participant in range(2)]) hilb = signal.hilbert(filtered) complex_signal.append(hilb) @@ -848,70 +960,82 @@ def compute_freq_bands(data: np.ndarray, sampling_rate: int, freq_bands: dict, return complex_signal -def compute_nmPLV(data: np.ndarray, sampling_rate: int, - freq_range1: List[float], freq_range2: List[float]) -> np.ndarray: +def compute_nmPLV( + data: np.ndarray, + sampling_rate: int, + freq_range1: List[float], + freq_range2: List[float], +) -> np.ndarray: """ Computes n:m Phase-Locking Value for cross-frequency coupling between participants. - + This function calculates the phase synchronization between two participants when their brain oscillations operate at different frequency ranges with an integer ratio relationship (n:m). - + Parameters ---------- data : np.ndarray EEG data from two participants with shape (2, n_epochs, n_channels, n_times) - + sampling_rate : int Sampling rate of the EEG data in Hz - + freq_range1 : List[float] Frequency range [fmin, fmax] for participant 1 - + freq_range2 : List[float] Frequency range [fmin, fmax] for participant 2 - + Returns ------- con : np.ndarray n:m PLV connectivity matrix with shape (n_freq, 2*n_channels, 2*n_channels) - + Notes ----- - n:m Phase Locking Value measures synchronization between oscillations at - different frequencies when there exists an n:m ratio between them. For example, + n:m Phase Locking Value measures synchronization between oscillations at + different frequencies when there exists an n:m ratio between them. For example, if participant 1 shows activity at 10 Hz and participant 2 at 20 Hz, this would be a 1:2 relationship. - + The function: 1. Computes the ratio between the mean frequencies of the two ranges 2. Applies appropriate phase multiplication to account for the frequency ratio 3. Calculates the phase locking value between the adjusted phases - + This measure is useful for studying cross-frequency coupling between brains, where different participants might exhibit coordination at different frequency bands. - + References ---------- Palva, J. M., Palva, S., & Kaila, K. (2005). Phase synchrony among neuronal oscillations in the human cortex. Journal of Neuroscience, 25(15), 3962-3972. - + Examples -------- >>> # Computing 1:2 PLV between alpha (participant 1) and beta (participant 2) - >>> nm_plv = compute_nmPLV(data, sampling_rate=256, + >>> nm_plv = compute_nmPLV(data, sampling_rate=256, ... freq_range1=[8, 12], freq_range2=[16, 24]) """ - r = np.mean(freq_range2)/np.mean(freq_range1) + r = np.mean(freq_range2) / np.mean(freq_range1) freq_range = [np.min(freq_range1), np.max(freq_range2)] - complex_signal = np.mean(compute_single_freq(data, sampling_rate, freq_range),3).squeeze() - - n_epoch, n_ch, n_freq, n_samp = complex_signal.shape[1], complex_signal.shape[2], \ - complex_signal.shape[3], complex_signal.shape[4] + complex_signal = np.mean( + compute_single_freq(data, sampling_rate, freq_range), 3 + ).squeeze() + + n_epoch, n_ch, n_freq, n_samp = ( + complex_signal.shape[1], + complex_signal.shape[2], + complex_signal.shape[3], + complex_signal.shape[4], + ) # calculate all epochs at once, the only downside is that the disk may not have enough space - complex_signal = complex_signal.transpose((1, 3, 0, 2, 4)).reshape(n_epoch, n_freq, 2 * n_ch, n_samp) + complex_signal = complex_signal.transpose((1, 3, 0, 2, 4)).reshape( + n_epoch, n_freq, 2 * n_ch, n_samp + ) transpose_axes = (0, 1, 3, 2) phase = complex_signal / np.abs(complex_signal) @@ -931,92 +1055,109 @@ def compute_nmPLV(data: np.ndarray, sampling_rate: int, return con -def xwt(sig1: mne.Epochs, sig2: mne.Epochs, freqs: Union[int, np.ndarray], - n_cycles=5.0, mode: str = "xwt") -> np.ndarray: +def xwt( + sig1: mne.Epochs, + sig2: mne.Epochs, + freqs: Union[int, np.ndarray], + n_cycles=5.0, + mode: str = "xwt", +) -> np.ndarray: """ Performs cross-wavelet transform or wavelet coherence analysis between two signals. - + This function computes time-frequency representations of two signals and their cross-spectrum using continuous wavelet transform. - + Parameters ---------- sig1 : mne.Epochs EEG data of first participant - + sig2 : mne.Epochs EEG data of second participant - + freqs : Union[int, np.ndarray] Frequencies of interest in Hz - + n_cycles : float, optional Number of cycles in the Morlet wavelet (default=5.0) Controls the time-frequency resolution trade-off - + mode : str, optional Type of analysis to perform (default="xwt"). Options: - 'power': absolute value of cross-wavelet transform - 'phase': phase angles of cross-wavelet transform - 'xwt': raw cross-wavelet transform (complex-valued) - 'wtc': wavelet coherence - + Returns ------- data : np.ndarray Result of the wavelet analysis with shape: (n_chans1, n_chans2, n_epochs, n_freqs, n_samples) - + Notes ----- This function provides different types of wavelet-based analyses: - + - Cross-wavelet transform (XWT): Reveals common power and relative phase between two time series in time-frequency space - + - Wavelet coherence (WTC): Measures the correlation between two signals in the time-frequency domain normalized between 0 and 1 - + Unlike Fourier-based methods, wavelet analysis maintains time resolution, allowing for the study of non-stationary signals and transient relationships. - + The function automatically checks that both signals have the same sampling rate, number of epochs, channels, and samples. - + References ---------- - Grinsted, A., Moore, J. C., & Jevrejeva, S. (2004). Application of the cross - wavelet transform and wavelet coherence to geophysical time series. + Grinsted, A., Moore, J. C., & Jevrejeva, S. (2004). Application of the cross + wavelet transform and wavelet coherence to geophysical time series. Nonlinear processes in geophysics, 11(5/6), 561-566. - - Maraun, D., & Kurths, J. (2004). Cross wavelet analysis: significance testing + + Maraun, D., & Kurths, J. (2004). Cross wavelet analysis: significance testing and pitfalls. Nonlinear Processes in Geophysics, 11(4), 505-514. - + Examples -------- >>> # Computing wavelet coherence between two EEG signals - >>> coherence = xwt(subj1_epochs, subj2_epochs, freqs=np.arange(4, 40, 1), + >>> coherence = xwt(subj1_epochs, subj2_epochs, freqs=np.arange(4, 40, 1), ... n_cycles=7, mode='wtc') """ - + # Set parameters for the output n_freqs = len(freqs) - sfreq = sig1.info['sfreq'] - assert sig1.info['sfreq'] == sig2.info['sfreq'], "Sig1 et sig2 should have the same sfreq value." + sfreq = sig1.info["sfreq"] + assert sig1.info["sfreq"] == sig2.info["sfreq"], ( + "Sig1 et sig2 should have the same sfreq value." + ) n_epochs1, n_chans1, n_samples1 = sig1.get_data(copy=False).shape n_epochs2, n_chans2, n_samples2 = sig2.get_data(copy=False).shape - assert n_epochs1 == n_epochs2, "n_epochs1 and n_epochs2 should have the same number of epochs." - assert n_chans1 == n_chans2, "n_chans1 and n_chans2 should have the same number of channels." - assert n_samples1 == n_samples2, "n_samples1 and n_samples2 should have the same number of samples." + assert n_epochs1 == n_epochs2, ( + "n_epochs1 and n_epochs2 should have the same number of epochs." + ) + assert n_chans1 == n_chans2, ( + "n_chans1 and n_chans2 should have the same number of channels." + ) + assert n_samples1 == n_samples2, ( + "n_samples1 and n_samples2 should have the same number of samples." + ) - cross_sigs = np.zeros((n_chans1, n_chans2, n_epochs1, n_freqs, n_samples1), dtype=complex) * np.nan + cross_sigs = ( + np.zeros((n_chans1, n_chans2, n_epochs1, n_freqs, n_samples1), dtype=complex) + * np.nan + ) wtcs = np.zeros((n_chans1, n_chans2, n_epochs1, n_freqs, n_samples1)) * np.nan # Set the mother wavelet - Ws = mne.time_frequency.tfr.morlet(sfreq, freqs, - n_cycles=n_cycles, sigma=None, zero_mean=True) + Ws = mne.time_frequency.tfr.morlet( + sfreq, freqs, n_cycles=n_cycles, sigma=None, zero_mean=True + ) # Wavelet scales in samples (Morlet scale-frequency relation) scales = n_cycles * sfreq / (2 * np.pi * freqs) @@ -1033,12 +1174,18 @@ def xwt(sig1: mne.Epochs, sig2: mne.Epochs, freqs: Union[int, np.ndarray], for ind1, ch_label1 in enumerate(sig1.ch_names): for ind2, ch_label2 in enumerate(sig2.ch_names): # Extract the channel's data for both participants and apply cwt - cur_sig1 = np.squeeze(sig1.get_data(mne.pick_channels(sig1.ch_names, [ch_label1]))) - cwt1 = mne.time_frequency.tfr.cwt(cur_sig1, Ws, use_fft=True, - mode='same', decim=1) - cur_sig2 = np.squeeze(sig2.get_data(mne.pick_channels(sig2.ch_names, [ch_label2]))) - cwt2 = mne.time_frequency.tfr.cwt(cur_sig2, Ws, use_fft=True, - mode='same', decim=1) + cur_sig1 = np.squeeze( + sig1.get_data(mne.pick_channels(sig1.ch_names, [ch_label1])) + ) + cwt1 = mne.time_frequency.tfr.cwt( + cur_sig1, Ws, use_fft=True, mode="same", decim=1 + ) + cur_sig2 = np.squeeze( + sig2.get_data(mne.pick_channels(sig2.ch_names, [ch_label2])) + ) + cwt2 = mne.time_frequency.tfr.cwt( + cur_sig2, Ws, use_fft=True, mode="same", decim=1 + ) # Compute cross-spectrum and auto-spectra cross_sig = cwt1 * cwt2.conj() @@ -1050,166 +1197,185 @@ def xwt(sig1: mne.Epochs, sig2: mne.Epochs, freqs: Union[int, np.ndarray], # following Grinsted et al. (2004) for fi, win in enumerate(smooth_wins): wps1[:, fi, :] = np.apply_along_axis( - lambda x: np.convolve(x, win, mode='same'), -1, wps1[:, fi, :]) + lambda x: np.convolve(x, win, mode="same"), -1, wps1[:, fi, :] + ) wps2[:, fi, :] = np.apply_along_axis( - lambda x: np.convolve(x, win, mode='same'), -1, wps2[:, fi, :]) + lambda x: np.convolve(x, win, mode="same"), -1, wps2[:, fi, :] + ) cross_sig[:, fi, :] = np.apply_along_axis( - lambda x: np.convolve(x, win, mode='same'), -1, cross_sig[:, fi, :]) + lambda x: np.convolve(x, win, mode="same"), -1, cross_sig[:, fi, :] + ) wtc = np.abs(cross_sig) ** 2 / (wps1 * wps2) wtcs[ind1, ind2, :, :, :] = wtc - if mode == 'power': + if mode == "power": data = np.abs(cross_sigs) - elif mode == 'phase': + elif mode == "phase": data = np.angle(cross_sigs) - elif mode == 'xwt': + elif mode == "xwt": data = cross_sigs - elif mode == 'wtc': - data = wtcs + elif mode == "wtc": + data = wtcs else: - data = 'Please specify a valid mode: power, phase, xwt, or wtc.' + data = "Please specify a valid mode: power, phase, xwt, or wtc." print(data) return data # helper function -def _multiply_conjugate(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) -> np.ndarray: +def _multiply_conjugate( + real: np.ndarray, imag: np.ndarray, transpose_axes: tuple +) -> np.ndarray: """ Computes the product of a complex array and its conjugate efficiently. - + .. deprecated:: 0.5.0 This function is deprecated and will be removed in version 1.0.0. Use :func:`hypyp.sync.multiply_conjugate` instead. - + This helper function performs matrix multiplication between complex arrays represented by their real and imaginary parts, collapsing the last dimension. - + Parameters ---------- real : np.ndarray Real part of the complex array - + imag : np.ndarray Imaginary part of the complex array - + transpose_axes : tuple Axes to transpose for matrix multiplication - + Returns ------- product : np.ndarray Product of the array and its complex conjugate - + Notes ----- This function implements the formula: product = (real × real.T + imag × imag.T) - i(real × imag.T - imag × real.T) - + Using einsum for efficient computation without explicitly creating complex arrays. """ warnings.warn( "_multiply_conjugate is deprecated and will be removed in version 1.0.0. " "Use hypyp.sync.multiply_conjugate instead.", DeprecationWarning, - stacklevel=2 + stacklevel=2, ) - formula = 'jilm,jimk->jilk' - product = np.einsum(formula, real, real.transpose(transpose_axes)) + \ - np.einsum(formula, imag, imag.transpose(transpose_axes)) - 1j * \ - (np.einsum(formula, real, imag.transpose(transpose_axes)) - \ - np.einsum(formula, imag, real.transpose(transpose_axes))) + formula = "jilm,jimk->jilk" + product = ( + np.einsum(formula, real, real.transpose(transpose_axes)) + + np.einsum(formula, imag, imag.transpose(transpose_axes)) + - 1j + * ( + np.einsum(formula, real, imag.transpose(transpose_axes)) + - np.einsum(formula, imag, real.transpose(transpose_axes)) + ) + ) return product # helper function -def _multiply_conjugate_time(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) -> np.ndarray: +def _multiply_conjugate_time( + real: np.ndarray, imag: np.ndarray, transpose_axes: tuple +) -> np.ndarray: """ Computes the product of a complex array and its conjugate without collapsing time dimension. - + .. deprecated:: 0.5.0 This function is deprecated and will be removed in version 1.0.0. Use :func:`hypyp.sync.multiply_conjugate_time` instead. - + Similar to _multiply_conjugate, but preserves the time dimension, which is needed for certain connectivity metrics like wPLI. - + Parameters ---------- real : np.ndarray Real part of the complex array - + imag : np.ndarray Imaginary part of the complex array - + transpose_axes : tuple Axes to transpose for matrix multiplication - + Returns ------- product : np.ndarray Product of the array and its complex conjugate with time dimension preserved - + Notes ----- This function uses a different einsum formula than _multiply_conjugate: 'jilm,jimk->jilkm' instead of 'jilm,jimk->jilk' - - This preserves the time dimension (m) in the output, which is necessary for - computing metrics that require individual time point values rather than + + This preserves the time dimension (m) in the output, which is necessary for + computing metrics that require individual time point values rather than time-averaged products. """ warnings.warn( "_multiply_conjugate_time is deprecated and will be removed in version 1.0.0. " "Use hypyp.sync.multiply_conjugate_time instead.", DeprecationWarning, - stacklevel=2 + stacklevel=2, ) - formula = 'jilm,jimk->jilkm' - product = np.einsum(formula, real, real.transpose(transpose_axes)) + \ - np.einsum(formula, imag, imag.transpose(transpose_axes)) - 1j * \ - (np.einsum(formula, real, imag.transpose(transpose_axes)) - \ - np.einsum(formula, imag, real.transpose(transpose_axes))) - + formula = "jilm,jimk->jilkm" + product = ( + np.einsum(formula, real, real.transpose(transpose_axes)) + + np.einsum(formula, imag, imag.transpose(transpose_axes)) + - 1j + * ( + np.einsum(formula, real, imag.transpose(transpose_axes)) + - np.einsum(formula, imag, real.transpose(transpose_axes)) + ) + ) + return product # helper function -def _multiply_product(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) -> np.ndarray: +def _multiply_product( + real: np.ndarray, imag: np.ndarray, transpose_axes: tuple +) -> np.ndarray: """ Computes the product of two complex arrays (not conjugate) efficiently. - + .. deprecated:: 0.5.0 This function is deprecated and will be removed in version 1.0.0. Use :func:`hypyp.sync.multiply_product` instead. - + This helper function performs matrix multiplication between complex arrays represented by their real and imaginary parts, collapsing the last dimension. Unlike _multiply_conjugate, this computes z1 * z2 instead of z1 * conj(z2). - + Parameters ---------- real : np.ndarray Real part of the complex array - + imag : np.ndarray Imaginary part of the complex array - + transpose_axes : tuple Axes to transpose for matrix multiplication - + Returns ------- product : np.ndarray Product of the array with itself (non-conjugate) - + Notes ----- This function implements the formula for z1 * z2: product = (real × real.T - imag × imag.T) + i(real × imag.T + imag × real.T) - + Using einsum for efficient computation without explicitly creating complex arrays. This is used in the adjusted circular correlation (accorr) metric. """ @@ -1217,67 +1383,73 @@ def _multiply_product(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) "_multiply_product is deprecated and will be removed in version 1.0.0. " "Use hypyp.sync.multiply_product instead.", DeprecationWarning, - stacklevel=2 + stacklevel=2, + ) + formula = "jilm,jimk->jilk" + product = ( + np.einsum(formula, real, real.transpose(transpose_axes)) + - np.einsum(formula, imag, imag.transpose(transpose_axes)) + + 1j + * ( + np.einsum(formula, real, imag.transpose(transpose_axes)) + + np.einsum(formula, imag, real.transpose(transpose_axes)) + ) ) - formula = 'jilm,jimk->jilk' - product = np.einsum(formula, real, real.transpose(transpose_axes)) - \ - np.einsum(formula, imag, imag.transpose(transpose_axes)) + 1j * \ - (np.einsum(formula, real, imag.transpose(transpose_axes)) + \ - np.einsum(formula, imag, real.transpose(transpose_axes))) return product # helper function -def _accorr_hybrid(complex_signal: np.ndarray, epochs_average: bool = True, - show_progress: bool = True) -> np.ndarray: +def _accorr_hybrid( + complex_signal: np.ndarray, epochs_average: bool = True, show_progress: bool = True +) -> np.ndarray: """ Computes Adjusted Circular Correlation using a hybrid approach. - + .. deprecated:: 0.5.0 This function is deprecated and will be removed in version 1.0.0. Use :class:`hypyp.sync.ACCorr` instead. - + This function calculates the adjusted circular correlation coefficient between all channel pairs. It uses a vectorized computation for the numerator and an exact loop-based computation for the denominator. - + Parameters ---------- complex_signal : np.ndarray Complex analytic signals with shape (n_epochs, n_freq, 2*n_channels, n_times) Note: This is the already reshaped signal from compute_sync. - + epochs_average : bool, optional If True, connectivity values are averaged across epochs (default) If False, epoch-by-epoch connectivity is preserved - + show_progress : bool, optional If True, display a progress bar during computation (default) If False, no progress bar is shown - + Returns ------- con : np.ndarray Adjusted circular correlation matrix with shape: - If epochs_average=True: (n_freq, 2*n_channels, 2*n_channels) - If epochs_average=False: (n_freq, n_epochs, 2*n_channels, 2*n_channels) - + Notes ----- The adjusted circular correlation is computed as: - + 1. Numerator (vectorized): Uses the difference between the absolute values of the conjugate product and the direct product of normalized complex signals. - + 2. Denominator (loop): For each channel pair, computes optimal phase centering parameters (m_adj, n_adj) that minimize the denominator, then calculates the normalization factor. - + This metric provides a more accurate measure of circular correlation by adjusting the phase centering for each channel pair individually, rather than using a global circular mean. - + References ---------- Zimmermann, M., Schultz-Nielsen, K., Dumas, G., & Konvalinka, I. (2024). @@ -1289,64 +1461,70 @@ def _accorr_hybrid(complex_signal: np.ndarray, epochs_average: bool = True, "_accorr_hybrid is deprecated and will be removed in version 1.0.0. " "Use hypyp.sync.ACCorr instead.", DeprecationWarning, - stacklevel=2 + stacklevel=2, ) n_epochs = complex_signal.shape[0] n_freq = complex_signal.shape[1] n_ch_total = complex_signal.shape[2] - + transpose_axes = (0, 1, 3, 2) - + # Numerator (vectorized) z = complex_signal / np.abs(complex_signal) c, s = np.real(z), np.imag(z) - + cross_conj = _multiply_conjugate(c, s, transpose_axes=transpose_axes) r_minus = np.abs(cross_conj) - + cross_prod = _multiply_product(c, s, transpose_axes=transpose_axes) r_plus = np.abs(cross_prod) - + num = r_minus - r_plus - + # Denominator (loop) angle = np.angle(complex_signal) den = np.zeros((n_epochs, n_freq, n_ch_total, n_ch_total)) - + total_pairs = (n_ch_total * (n_ch_total + 1)) // 2 - pbar = tqdm(total=total_pairs, desc=" accorr (denominator)", - disable=not show_progress, leave=False) - + pbar = tqdm( + total=total_pairs, + desc=" accorr (denominator)", + disable=not show_progress, + leave=False, + ) + for i in range(n_ch_total): for j in range(i, n_ch_total): alpha1 = angle[:, :, i, :] alpha2 = angle[:, :, j, :] - + phase_diff = alpha1 - alpha2 phase_sum = alpha1 + alpha2 - - mean_diff = np.angle(np.mean(np.exp(1j * phase_diff), axis=2, keepdims=True)) + + mean_diff = np.angle( + np.mean(np.exp(1j * phase_diff), axis=2, keepdims=True) + ) mean_sum = np.angle(np.mean(np.exp(1j * phase_sum), axis=2, keepdims=True)) - + n_adj = -1 * (mean_diff - mean_sum) / 2 m_adj = mean_diff + n_adj - + x_sin = np.sin(alpha1 - m_adj) y_sin = np.sin(alpha2 - n_adj) - + den_ij = 2 * np.sqrt(np.sum(x_sin**2, axis=2) * np.sum(y_sin**2, axis=2)) den[:, :, i, j] = den_ij den[:, :, j, i] = den_ij - + pbar.update(1) - + pbar.close() - + den = np.where(den == 0, 1, den) con = num / den con = con.swapaxes(0, 1) # n_freq x n_epoch x 2*n_ch x 2*n_ch - + if epochs_average: con = np.nanmean(con, axis=1) - - return con \ No newline at end of file + + return con diff --git a/hypyp/fnirs/__init__.py b/hypyp/fnirs/__init__.py index 64351ca..2b4092b 100644 --- a/hypyp/fnirs/__init__.py +++ b/hypyp/fnirs/__init__.py @@ -1,5 +1,8 @@ from .preprocessor.base_preprocessor import BasePreprocessor, BaseStep -from .preprocessor.implementations.mne_preprocessor_raw_to_haemo import MnePreprocessorRawToHaemo, MneStep +from .preprocessor.implementations.mne_preprocessor_raw_to_haemo import ( + MnePreprocessorRawToHaemo, + MneStep, +) from .preprocessor.implementations.mne_preprocessor_as_is import MnePreprocessorAsIs from .data_browser import DataBrowser from .recording import Recording @@ -8,14 +11,14 @@ from .channel_roi import ChannelROI __all__ = [ - 'Recording', - 'Dyad', - 'Study', - 'DataBrowser', - 'BasePreprocessor', - 'BaseStep', - 'MnePreprocessorAsIs', - 'MnePreprocessorRawToHaemo', - 'MneStep', - 'ChannelROI', + "Recording", + "Dyad", + "Study", + "DataBrowser", + "BasePreprocessor", + "BaseStep", + "MnePreprocessorAsIs", + "MnePreprocessorRawToHaemo", + "MneStep", + "ChannelROI", ] diff --git a/hypyp/fnirs/channel_roi.py b/hypyp/fnirs/channel_roi.py index 40e53b2..1db082e 100644 --- a/hypyp/fnirs/channel_roi.py +++ b/hypyp/fnirs/channel_roi.py @@ -3,10 +3,11 @@ import scipy.io + class ChannelROI: rois: OrderedDict - def __init__(self, rois:OrderedDict|None=None): + def __init__(self, rois: OrderedDict | None = None): """ Structure for grouping channels into regions of interest @@ -19,9 +20,9 @@ def __init__(self, rois:OrderedDict|None=None): self.rois = rois @staticmethod - def from_lionirs_file(roi_file_path:str): + def from_lionirs_file(roi_file_path: str): """ - Create a new ChannelROI object from a file exported + Create a new ChannelROI object from a file exported Args: roi_file_path (str): The path of the file exported by LIONirs toolbox for Region of Interests @@ -36,15 +37,15 @@ def from_lionirs_file(roi_file_path:str): row_channel_name_map = dict() - for i, channel_row in enumerate(mat['zone'][0][0].ml): - row_id = i+1 + for i, channel_row in enumerate(mat["zone"][0][0].ml): + row_id = i + 1 source_id = channel_row[0] detector_id = channel_row[1] # The last column is 1 or 2, which represents hbo or hbr, but we don't need that info here - row_channel_name_map[row_id] = f'S{source_id}_D{detector_id}' + row_channel_name_map[row_id] = f"S{source_id}_D{detector_id}" - mat_rois = mat['zone'][0][0].label[0] - mat_rows_per_rois = mat['zone'][0][0].plotLst[0] + mat_rois = mat["zone"][0][0].label[0] + mat_rows_per_rois = mat["zone"][0][0].plotLst[0] for mat_roi, mat_rows_per_roi in zip(mat_rois, mat_rows_per_rois): roi_name = mat_roi[0] @@ -54,14 +55,14 @@ def from_lionirs_file(roi_file_path:str): channels_in_roi.append(row_channel_name_map[row]) obj.rois[roi_name] = channels_in_roi return obj - + @property def ordered_ch_names(self): channel_names = [] for channel_names_in_roi in self.rois.values(): channel_names = channel_names + channel_names_in_roi return channel_names - + @property def group_boundaries_sizes(self): boundaries = [0] @@ -70,12 +71,12 @@ def group_boundaries_sizes(self): boundaries.append(len(self.ordered_ch_names)) return boundaries - def get_ch_names_in_order(self, ch_names:str) -> List[str]: + def get_ch_names_in_order(self, ch_names: str) -> List[str]: """ Given a list of channel names, return them in order of Region of Interest Args: - names (str): + names (str): Returns: List[str]: all the received ch_names, ordered according to defined regions of interest @@ -88,11 +89,11 @@ def get_ch_names_in_order(self, ch_names:str) -> List[str]: for name_found in names_found: if name_found not in all_names_found: all_names_found.append(name_found) - + for unordered_ch_name in ch_names: if not unordered_ch_name in all_names_found: all_names_found.append(unordered_ch_name) - + return all_names_found def get_roi_from_channel(self, ch_name): @@ -101,6 +102,4 @@ def get_roi_from_channel(self, ch_name): names_found = [name for name in v if ch_name.startswith(name)] if len(names_found) > 0: return k - return '' - - + return "" diff --git a/hypyp/fnirs/data_browser.py b/hypyp/fnirs/data_browser.py index 6d20596..9d7aec5 100644 --- a/hypyp/fnirs/data_browser.py +++ b/hypyp/fnirs/data_browser.py @@ -17,6 +17,7 @@ "downloads", ) + class DataBrowser: paths: List[str] @@ -28,7 +29,7 @@ def __init__(self): """ self.paths = [] - def add_source(self, path:str): + def add_source(self, path: str): """ Add a folder to look for NIRS files @@ -42,16 +43,16 @@ def add_source(self, path:str): return self @staticmethod - def is_path_nirx(path:str) -> bool: - return os.path.isfile(path) and path.endswith('.hdr') + def is_path_nirx(path: str) -> bool: + return os.path.isfile(path) and path.endswith(".hdr") @staticmethod - def is_path_fif(path:str) -> bool: - return os.path.isfile(path) and path.endswith('.fif') + def is_path_fif(path: str) -> bool: + return os.path.isfile(path) and path.endswith(".fif") @staticmethod - def is_path_snirf(path:str) -> bool: - return os.path.isfile(path) and path.endswith('.snirf') + def is_path_snirf(path: str) -> bool: + return os.path.isfile(path) and path.endswith(".snirf") def list_all_files(self) -> List[str]: """ @@ -64,7 +65,7 @@ def list_all_files(self) -> List[str]: for root_path in self.paths: if not Path(root_path).exists(): continue - for path in Path(root_path).rglob('*'): + for path in Path(root_path).rglob("*"): if DataBrowser.is_path_fif(str(path)): file_paths.append(str(path.absolute())) @@ -98,13 +99,12 @@ def download_demo_dataset(self) -> str: progressbar=True, ) - target_path = os.path.join(extract_path, 'fathers') + target_path = os.path.join(extract_path, "fathers") if not os.path.exists(target_path): - with ZipFile(zip_path, 'r') as zip: - print(f'Extracting to {extract_path}, (target: {target_path})') + with ZipFile(zip_path, "r") as zip: + print(f"Extracting to {extract_path}, (target: {target_path})") zip.extractall(path=extract_path) self.add_source(target_path) return target_path - diff --git a/hypyp/fnirs/dyad.py b/hypyp/fnirs/dyad.py index 2b27876..ac42a84 100644 --- a/hypyp/fnirs/dyad.py +++ b/hypyp/fnirs/dyad.py @@ -22,10 +22,14 @@ ) PairChannelMatchSingleType = str | List[str] | re.Pattern -PairChannelMatchType = PairChannelMatchSingleType | Tuple[PairChannelMatchSingleType, PairChannelMatchSingleType] +PairChannelMatchType = ( + PairChannelMatchSingleType + | Tuple[PairChannelMatchSingleType, PairChannelMatchSingleType] +) MIN_SECTION_LENGTH = 10 + class Dyad: """ The Dyad object is a pair of recordings (per subject) of an hyperscanning recording. @@ -37,15 +41,20 @@ class Dyad: label (str, optional): Custom label for the dyad. Defaults to `s1.label`-`s2.label`. is_pseudo (bool, optional): If the dyad is a permutated pair created for comparison. Used to track dyad "type" in results. Defaults to False. """ + s1: Recording s2: Recording label: str is_pseudo: bool - tasks: TaskList # intersection of tasks of subject 1 and subject 2 - wtcs: List[WTC] | None # the computed Wavelet Transform Coherence for each channel pairs in the dyad - df: CoherenceDataFrame | None # pandas dataframe from computed coherence - - def __init__(self, s1:Recording, s2:Recording, label:str='', is_pseudo:bool=False): + tasks: TaskList # intersection of tasks of subject 1 and subject 2 + wtcs: ( + List[WTC] | None + ) # the computed Wavelet Transform Coherence for each channel pairs in the dyad + df: CoherenceDataFrame | None # pandas dataframe from computed coherence + + def __init__( + self, s1: Recording, s2: Recording, label: str = "", is_pseudo: bool = False + ): self.s1 = s1 self.s2 = s2 self.wtcs = None @@ -54,7 +63,7 @@ def __init__(self, s1:Recording, s2:Recording, label:str='', is_pseudo:bool=Fals self.is_intra = s1 == s2 self.label = label - if self.label == '': + if self.label == "": self.label = Dyad._get_label_from_recordings(s1, s2) # Intersect the tasks @@ -69,25 +78,25 @@ def __init__(self, s1:Recording, s2:Recording, label:str='', is_pseudo:bool=Fals if task_name in s2_tasks_names: self.tasks.append(task) found_tasks_names.append(task_name) - - @property + + @property def recordings(self) -> Tuple[Recording, Recording]: return (self.s1, self.s2) - + @property def is_preprocessed(self) -> bool: for recording in self.recordings: if not recording.is_preprocessed: return False return True - + @property def is_wtc_computed(self): return self.wtcs is not None @staticmethod - def _get_label_from_recordings(s1:Recording, s2:Recording) -> str: - return f'{s1.subject_label}-{s2.subject_label}' + def _get_label_from_recordings(s1: Recording, s2: Recording) -> str: + return f"{s1.subject_label}-{s2.subject_label}" def preprocess(self, preprocessor: BasePreprocessor): """ @@ -103,29 +112,31 @@ def preprocess(self, preprocessor: BasePreprocessor): recording.preprocess(preprocessor) return self - def _append_pairs(self, - label_dyad:str, - s1_ch_names:List[str], - s2_ch_names:List[str], - s1_task_data:np.ndarray, - s2_task_data:np.ndarray, - s1:Recording, - s2:Recording, - task_name:str, - epoch_id:int, - is_intra_of:int|None, - is_pseudo:bool, - pairs:List[PairSignals]): + def _append_pairs( + self, + label_dyad: str, + s1_ch_names: List[str], + s2_ch_names: List[str], + s1_task_data: np.ndarray, + s2_task_data: np.ndarray, + s1: Recording, + s2: Recording, + task_name: str, + epoch_id: int, + is_intra_of: int | None, + is_pseudo: bool, + pairs: List[PairSignals], + ): n = s1_task_data.shape[1] - x = np.linspace(0, n/s1.preprocessed.info['sfreq'], n) + x = np.linspace(0, n / s1.preprocessed.info["sfreq"], n) for s1_i, s1_ch_name in enumerate(s1_ch_names): for s2_i, s2_ch_name in enumerate(s2_ch_names): - y1 = s1_task_data[s1_i,:] - y2 = s2_task_data[s2_i,:] + y1 = s1_task_data[s1_i, :] + y2 = s2_task_data[s2_i, :] # Crop signals in case they are not the same length stop = min(len(y1), len(y2)) - y1 = y1[:stop] - y2 = y2[:stop] + y1 = y1[:stop] + y2 = y2[:stop] # Look for NaN, and split in section section_id = 0 @@ -133,15 +144,29 @@ def _append_pairs(self, has_nan = np.any(nan_mask) if has_nan: - nan_idx = np.where(nan_mask)[0]+1 + nan_idx = np.where(nan_mask)[0] + 1 # We have to drop the first item of every split but the first split, because it contains the NaN - x_sections = [section for section in np.split(x, nan_idx) if len(section)>MIN_SECTION_LENGTH] - y1_sections = [section for section in np.split(y1, nan_idx) if len(section)>MIN_SECTION_LENGTH] - y2_sections = [section for section in np.split(y2, nan_idx) if len(section)>MIN_SECTION_LENGTH] - + x_sections = [ + section + for section in np.split(x, nan_idx) + if len(section) > MIN_SECTION_LENGTH + ] + y1_sections = [ + section + for section in np.split(y1, nan_idx) + if len(section) > MIN_SECTION_LENGTH + ] + y2_sections = [ + section + for section in np.split(y2, nan_idx) + if len(section) > MIN_SECTION_LENGTH + ] + # remove the lurking NaN in edges of our arrays for i in range(len(x_sections)): - delete_flags = np.isnan(y1_sections[i]) | np.isnan(y2_sections[i]) + delete_flags = np.isnan(y1_sections[i]) | np.isnan( + y2_sections[i] + ) x_sections[i] = np.delete(x_sections[i], delete_flags) y1_sections[i] = np.delete(y1_sections[i], delete_flags) y2_sections[i] = np.delete(y2_sections[i], delete_flags) @@ -151,27 +176,36 @@ def _append_pairs(self, y2_sections = [y2] for section_id in range(len(x_sections)): - pairs.append(PairSignals( - x_sections[section_id], - y1_sections[section_id], - y2_sections[section_id], - label_ch1=s1_ch_name, - label_ch2=s2_ch_name, - label_roi1=s1.get_roi_from_channel(s1_ch_name), - label_roi2=s2.get_roi_from_channel(s2_ch_name), - label_s1=s1.subject_label, - label_s2=s2.subject_label, - label_dyad=label_dyad, - label_task=task_name, - epoch_idx=epoch_id, - section_idx=section_id, - is_intra=(is_intra_of is not None), - is_intra_of=is_intra_of, - is_pseudo=is_pseudo, - )) - - - def get_pairs(self, s1:Recording, s2:Recording, label_dyad:str|None=None, ch_match:PairChannelMatchType|None=None, is_intra_of:int|None=None, is_pseudo:bool=False) -> List[PairSignals]: + pairs.append( + PairSignals( + x_sections[section_id], + y1_sections[section_id], + y2_sections[section_id], + label_ch1=s1_ch_name, + label_ch2=s2_ch_name, + label_roi1=s1.get_roi_from_channel(s1_ch_name), + label_roi2=s2.get_roi_from_channel(s2_ch_name), + label_s1=s1.subject_label, + label_s2=s2.subject_label, + label_dyad=label_dyad, + label_task=task_name, + epoch_idx=epoch_id, + section_idx=section_id, + is_intra=(is_intra_of is not None), + is_intra_of=is_intra_of, + is_pseudo=is_pseudo, + ) + ) + + def get_pairs( + self, + s1: Recording, + s2: Recording, + label_dyad: str | None = None, + ch_match: PairChannelMatchType | None = None, + is_intra_of: int | None = None, + is_pseudo: bool = False, + ) -> List[PairSignals]: """ Generate all the signal pairs between the 2 subjects and returns them in a format suitable for signal processing @@ -192,8 +226,8 @@ def get_pairs(self, s1:Recording, s2:Recording, label_dyad:str|None=None, ch_mat pairs: List[PairSignals] = [] - if s1.preprocessed.info['sfreq'] != s2.preprocessed.info['sfreq']: - raise RuntimeError('Recordings must have the same sampling frequency') + if s1.preprocessed.info["sfreq"] != s2.preprocessed.info["sfreq"]: + raise RuntimeError("Recordings must have the same sampling frequency") # Force match in tuple for leaner code below if not isinstance(ch_match, tuple): @@ -208,9 +242,16 @@ def check_match(ch_name, m): return ch_name in m return m in ch_name - - s1_ch_names = [ch_name for ch_name in s1.ordered_ch_names if check_match(ch_name, ch_match[0])] - s2_ch_names = [ch_name for ch_name in s2.ordered_ch_names if check_match(ch_name, ch_match[1])] + s1_ch_names = [ + ch_name + for ch_name in s1.ordered_ch_names + if check_match(ch_name, ch_match[0]) + ] + s2_ch_names = [ + ch_name + for ch_name in s2.ordered_ch_names + if check_match(ch_name, ch_match[1]) + ] seen_tasks = set() for task in self.tasks: @@ -239,12 +280,14 @@ def check_match(ch_name, m): epochs1 = s1.get_epochs_for_task(task.name).copy().pick(s1_ch_names) epochs2 = s2.get_epochs_for_task(task.name).copy().pick(s2_ch_names) if len(epochs1) != len(epochs2): - warnings.warn("The 2 recordings do not have the same epochs count. Some epochs will be skipped in pairs.") - + warnings.warn( + "The 2 recordings do not have the same epochs count. Some epochs will be skipped in pairs." + ) + # add one pair for each epoch for i in range(min(len(epochs1), len(epochs2))): - s1_task_data = epochs1.get_data(copy=False)[i,:,:] - s2_task_data = epochs2.get_data(copy=False)[i,:,:] + s1_task_data = epochs1.get_data(copy=False)[i, :, :] + s2_task_data = epochs2.get_data(copy=False)[i, :, :] epoch_id = i self._append_pairs( label_dyad, @@ -262,12 +305,12 @@ def check_match(ch_name, m): ) return pairs - + def compute_wtcs( self, wavelet: BaseWavelet | None = None, ch_match: PairChannelMatchType | None = None, - only_time_range: Tuple[float,float] | None = None, + only_time_range: Tuple[float, float] | None = None, bin_seconds: float | None = None, period_cuts: List[float] | None = None, frequency_cuts: List[float] | None = None, @@ -300,11 +343,13 @@ def compute_wtcs( wavelet = ComplexMorletWavelet() if period_cuts is not None and frequency_cuts is not None: - raise RuntimeError('Cannot specify both period_cuts and frequency_cuts') + raise RuntimeError("Cannot specify both period_cuts and frequency_cuts") self.wtcs = [] - pairs = self.get_pairs(self.s1, self.s2, ch_match=ch_match, is_pseudo=self.is_pseudo) + pairs = self.get_pairs( + self.s1, self.s2, ch_match=ch_match, is_pseudo=self.is_pseudo + ) if self.is_intra: # Force with_intra when subject 1 is the same as subject 2 @@ -312,10 +357,17 @@ def compute_wtcs( else: for pair in pairs: if verbose: - print(f'Running Wavelet Coherence for dyad "{self.label}" on pair "{pair.label}"') + print( + f'Running Wavelet Coherence for dyad "{self.label}" on pair "{pair.label}"' + ) if only_time_range is not None: pair = pair.sub(only_time_range) - wtc = wavelet.wtc(pair, bin_seconds=bin_seconds, period_cuts=period_cuts, frequency_cuts=frequency_cuts) + wtc = wavelet.wtc( + pair, + bin_seconds=bin_seconds, + period_cuts=period_cuts, + frequency_cuts=frequency_cuts, + ) if downsample is not None: wtc.downsample_in_time(downsample) @@ -334,15 +386,28 @@ def compute_wtcs( ch_match_intra = ch_match[i] else: ch_match_intra = ch_match - + # We have to keep track of the subject identifier when we are intra, since they would be displayed differently (parent VS child for example) - is_intra_of = i+1 - for pair in self.get_pairs(recording, recording, f'{recording.subject_label}(intra)', ch_match=ch_match_intra, is_intra_of=is_intra_of): + is_intra_of = i + 1 + for pair in self.get_pairs( + recording, + recording, + f"{recording.subject_label}(intra)", + ch_match=ch_match_intra, + is_intra_of=is_intra_of, + ): if verbose: - print(f'Running Wavelet Coherence intra-subject "{recording.subject_label}" on pair "{pair.label}"') + print( + f'Running Wavelet Coherence intra-subject "{recording.subject_label}" on pair "{pair.label}"' + ) if only_time_range is not None: pair = pair.sub(only_time_range) - wtc = wavelet.wtc(pair, bin_seconds=bin_seconds, period_cuts=period_cuts, frequency_cuts=frequency_cuts) + wtc = wavelet.wtc( + pair, + bin_seconds=bin_seconds, + period_cuts=period_cuts, + frequency_cuts=frequency_cuts, + ) if downsample is not None: wtc.downsample_in_time(downsample) recording.intra_wtcs.append(wtc) @@ -355,11 +420,13 @@ def compute_wtcs( self.s1.intra_wtcs = [] return self - + def _get_coherence_df(self, with_intra=False) -> pd.DataFrame: if with_intra: if not self.s1.is_wtc_computed or not self.s2.is_wtc_computed: - raise RuntimeError('Intra subject WTCs are not computed. Please check "compute_wtcs" arguments') + raise RuntimeError( + 'Intra subject WTCs are not computed. Please check "compute_wtcs" arguments' + ) wtcs = self.wtcs + self.s1.intra_wtcs + self.s2.intra_wtcs else: wtcs = self.wtcs @@ -379,7 +446,7 @@ def reset(self): # # Plots - # + # def plot_wtc(self, wtc: WTC, **kwargs): """ Plot the Wavelet Transform Coherence @@ -394,9 +461,16 @@ def plot_wtc(self, wtc: WTC, **kwargs): wtc.coif, wtc.sfreq, title=wtc.label_pair, - **kwargs) + **kwargs, + ) - def plot_coherence_matrix(self, field1:str='channel1', field2:str='channel2', query:str | None=None, **kwargs): + def plot_coherence_matrix( + self, + field1: str = "channel1", + field2: str = "channel2", + query: str | None = None, + **kwargs, + ): """ Plot the computed coherence metric for pair of fields (channel or roi) in a matrix format @@ -408,10 +482,12 @@ def plot_coherence_matrix(self, field1:str='channel1', field2:str='channel2', qu df = self.df if query is not None: df = df.query(query) - + ch_names1 = self.s1.ordered_ch_names ch_names2 = self.s2.ordered_ch_names - ordered_names = ch_names1 + [name for name in ch_names2 if name not in ch_names1] + ordered_names = ch_names1 + [ + name for name in ch_names2 if name not in ch_names1 + ] return plot_coherence_matrix( df, @@ -420,35 +496,28 @@ def plot_coherence_matrix(self, field1:str='channel1', field2:str='channel2', qu field1, field2, ordered_names, - **kwargs) - - def plot_coherence_matrix_per_channel(self, query:str|None=None, **kwargs): + **kwargs, + ) + + def plot_coherence_matrix_per_channel(self, query: str | None = None, **kwargs): """ Wraps plot_coherence_matrix to plot per channel Args: query (str | None, optional): pandas query to filter the dataframe. Defaults to None. """ - return self.plot_coherence_matrix( - 'channel1', - 'channel2', - query, - **kwargs) - - def plot_coherence_matrix_per_roi(self, query:str|None=None, **kwargs): + return self.plot_coherence_matrix("channel1", "channel2", query, **kwargs) + + def plot_coherence_matrix_per_roi(self, query: str | None = None, **kwargs): """ Wraps plot_coherence_matrix to plot per region of interest Args: query (str | None, optional): pandas query to filter the dataframe. Defaults to None. """ - return self.plot_coherence_matrix( - 'roi1', - 'roi2', - query, - **kwargs) - - def plot_coherence_matrix_per_channel_for_task(self, task:str, **kwargs): + return self.plot_coherence_matrix("roi1", "roi2", query, **kwargs) + + def plot_coherence_matrix_per_channel_for_task(self, task: str, **kwargs): """ Wraps plot_coherence_matrix_per_channel to plot for a specific task @@ -456,12 +525,10 @@ def plot_coherence_matrix_per_channel_for_task(self, task:str, **kwargs): task (str): task name """ return self.plot_coherence_matrix( - 'channel1', - 'channel2', - query=f'task=="{task}"', - **kwargs) - - def plot_coherence_matrix_per_roi_for_task(self, task:str, **kwargs): + "channel1", "channel2", query=f'task=="{task}"', **kwargs + ) + + def plot_coherence_matrix_per_roi_for_task(self, task: str, **kwargs): """ Wraps plot_coherence_matrix_per_roi to plot for a specific task @@ -469,69 +536,61 @@ def plot_coherence_matrix_per_roi_for_task(self, task:str, **kwargs): task (str): task name """ return self.plot_coherence_matrix( - 'roi1', - 'roi2', - query=f'task=="{task}"', - **kwargs) - + "roi1", "roi2", query=f'task=="{task}"', **kwargs + ) + def plot_coherence_bars_per_task(self, **kwargs): """ Plot coherence metric per task for comparison """ - return plot_coherence_bars_per_task( - self.df, - **kwargs) - + return plot_coherence_bars_per_task(self.df, **kwargs) + # # Plot connectogram (Proof of Concept) - # - def plot_coherence_connectogram_intra(self, recording:Recording, query:str|None=None, **kwargs): + # + def plot_coherence_connectogram_intra( + self, recording: Recording, query: str | None = None, **kwargs + ): df = self.df - selector = df['is_intra']==True + selector = df["is_intra"] == True df_filtered = df[selector] if query is not None: df_filtered = df_filtered.query(query) - pivot = df_filtered.pivot_table(index='roi1', columns='roi2', values='coherence', aggfunc='mean') - if 'title' not in kwargs: - kwargs['title'] = recording.subject_label + pivot = df_filtered.pivot_table( + index="roi1", columns="roi2", values="coherence", aggfunc="mean" + ) + if "title" not in kwargs: + kwargs["title"] = recording.subject_label - return plot_coherence_connectogram( - pivot, - **kwargs) + return plot_coherence_connectogram(pivot, **kwargs) - def plot_coherence_connectogram_s1(self, query:str|None=None, **kwargs): - return self.plot_coherence_connectogram_intra( - self.s1, - query, - **kwargs) + def plot_coherence_connectogram_s1(self, query: str | None = None, **kwargs): + return self.plot_coherence_connectogram_intra(self.s1, query, **kwargs) - def plot_coherence_connectogram_s2(self, query:str|None=None, **kwargs): - return self.plot_coherence_connectogram_intra( - self.s2, - query, - **kwargs) + def plot_coherence_connectogram_s2(self, query: str | None = None, **kwargs): + return self.plot_coherence_connectogram_intra(self.s2, query, **kwargs) - def plot_coherence_connectogram(self, query:str|None=None, title:str|None=None, **kwargs): + def plot_coherence_connectogram( + self, query: str | None = None, title: str | None = None, **kwargs + ): df = self.df.copy() - selector = df['is_intra']==False + selector = df["is_intra"] == False df_filtered = df[selector] if query is not None: df_filtered = df_filtered.query(query) # rename to have them separated in the plot - df_filtered.loc[:, 'roi1'] = 's1_' + df_filtered['roi1'].astype(str) - df_filtered.loc[:, 'roi2'] = 's2_' + df_filtered['roi2'].astype(str) + df_filtered.loc[:, "roi1"] = "s1_" + df_filtered["roi1"].astype(str) + df_filtered.loc[:, "roi2"] = "s2_" + df_filtered["roi2"].astype(str) - pivot = df_filtered.pivot_table(index='roi1', columns='roi2', values='coherence', aggfunc='mean') + pivot = df_filtered.pivot_table( + index="roi1", columns="roi2", values="coherence", aggfunc="mean" + ) if title is None: - title=f'{self.s1.subject_label} / {self.s2.subject_label}' - - return plot_coherence_connectogram_split( - pivot, - title=title, - **kwargs) + title = f"{self.s1.subject_label} / {self.s2.subject_label}" + return plot_coherence_connectogram_split(pivot, title=title, **kwargs) diff --git a/hypyp/fnirs/legacy/fnirs_tools.py b/hypyp/fnirs/legacy/fnirs_tools.py index 4802a9e..db1cefb 100644 --- a/hypyp/fnirs/legacy/fnirs_tools.py +++ b/hypyp/fnirs/legacy/fnirs_tools.py @@ -6,82 +6,108 @@ from scipy.io import loadmat -#Loading fnirs data - -def load_fnirs(path1: str, path2: str, attr: dict = None, preload: bool = False, verbose = None) -> mne.io.Raw: - """ - Loads fNIRS data files - - Arguments - ---------- - path1: str - participant #1 fNIRS data path (directory) - - part2: str - participant #2 fNIRS data path (directory) - - attr: dict, optional - dictionary containing optional attributes using to load different fNIRS file - (default is None, which returns MNE's default attributes) - - preload: bool, optional - Preload data into memory for data manipulation and faster indexing. - If True, the data will be preloaded into memory (fast, requires large amount of memory). - If preload is a string, preload is the file name of a memory-mapped file which is used to - store the data on the hard drive: slower, requires less memory. (default is False) - - verbose: bool, optional - Control verbosity of the logging output. If None, use the default verbosity level - - Returns - -------- - raw: - instance of RawSNIRF; a Raw object containing fNIRS data. - - """ - if ".snirf" in path1: - if attr is None: - data_1 = mne.io.read_raw_snirf(path1, optode_frame='unknown', preload=preload, verbose=verbose) - data_2 = mne.io.read_raw_snirf(path2, optode_frame='unknown', preload=preload, verbose=verbose) - else: - data_1 = mne.io.read_raw_snirf(path1, optode_frame = attr["optode_frame"], preload=preload, verbose=verbose) - data_2 = mne.io.read_raw_snirf(path2, optode_frame = attr["optode_frame"], preload=preload, verbose=verbose) +# Loading fnirs data - elif os.path.isdir(path1): - if attr is None: - data_1 = mne.io.read_raw_nirx(path1, saturated='annotate', preload=preload, verbose=verbose) - data_2 = mne.io.read_raw_nirx(path2, saturated='annotate', preload=preload, verbose=verbose) - else: - data_1 = mne.io.read_raw_nirx(path1, saturated = attr["saturated"], preload=preload, verbose=verbose) - data_2 = mne.io.read_raw_nirx(path2, saturated = attr["saturated"], preload=preload, verbose=verbose) - elif ".csv" in path1: - data_1 = mne.io.read_raw_hitachi(path1, preload=preload, verbose=verbose) - data_2 = mne.io.read_raw_hitachi(path2, preload=preload, verbose=verbose) +def load_fnirs( + path1: str, path2: str, attr: dict = None, preload: bool = False, verbose=None +) -> mne.io.Raw: + """ + Loads fNIRS data files + + Arguments + ---------- + path1: str + participant #1 fNIRS data path (directory) + + part2: str + participant #2 fNIRS data path (directory) + + attr: dict, optional + dictionary containing optional attributes using to load different fNIRS file + (default is None, which returns MNE's default attributes) - elif ".txt" in path1: - data_1 = mne.io.read_raw_boxy(path1, preload=preload, verbose=verbose) - data_2 = mne.io.read_raw_boxy(path2, preload=preload, verbose=verbose) + preload: bool, optional + Preload data into memory for data manipulation and faster indexing. + If True, the data will be preloaded into memory (fast, requires large amount of memory). + If preload is a string, preload is the file name of a memory-mapped file which is used to + store the data on the hard drive: slower, requires less memory. (default is False) - else: - print("data type is not supported") - - return (data_1, data_2) + verbose: bool, optional + Control verbosity of the logging output. If None, use the default verbosity level + Returns + -------- + raw: + instance of RawSNIRF; a Raw object containing fNIRS data. + + """ + if ".snirf" in path1: + if attr is None: + data_1 = mne.io.read_raw_snirf( + path1, optode_frame="unknown", preload=preload, verbose=verbose + ) + data_2 = mne.io.read_raw_snirf( + path2, optode_frame="unknown", preload=preload, verbose=verbose + ) + else: + data_1 = mne.io.read_raw_snirf( + path1, + optode_frame=attr["optode_frame"], + preload=preload, + verbose=verbose, + ) + data_2 = mne.io.read_raw_snirf( + path2, + optode_frame=attr["optode_frame"], + preload=preload, + verbose=verbose, + ) + + elif os.path.isdir(path1): + if attr is None: + data_1 = mne.io.read_raw_nirx( + path1, saturated="annotate", preload=preload, verbose=verbose + ) + data_2 = mne.io.read_raw_nirx( + path2, saturated="annotate", preload=preload, verbose=verbose + ) + else: + data_1 = mne.io.read_raw_nirx( + path1, saturated=attr["saturated"], preload=preload, verbose=verbose + ) + data_2 = mne.io.read_raw_nirx( + path2, saturated=attr["saturated"], preload=preload, verbose=verbose + ) + + elif ".csv" in path1: + data_1 = mne.io.read_raw_hitachi(path1, preload=preload, verbose=verbose) + data_2 = mne.io.read_raw_hitachi(path2, preload=preload, verbose=verbose) + + elif ".txt" in path1: + data_1 = mne.io.read_raw_boxy(path1, preload=preload, verbose=verbose) + data_2 = mne.io.read_raw_boxy(path2, preload=preload, verbose=verbose) + + else: + print("data type is not supported") + + return (data_1, data_2) + + +# Building the montage -#Building the montage def make_fnirs_montage( - source_labels:list, - detector_labels:list, - prob_directory: str, - Nz:list, - RPA:list, - - LPA:list, - head_size:float, - create_montage:bool = True, - mne_standard:str = None) -> mne.channels.DigMontage: + source_labels: list, + detector_labels: list, + prob_directory: str, + Nz: list, + RPA: list, + LPA: list, + head_size: float, + create_montage: bool = True, + mne_standard: str = None, +) -> mne.channels.DigMontage: """ Builds a compatible montage with MNE functions @@ -112,7 +138,7 @@ def make_fnirs_montage( head_size: float Head size in mm - creat_montage: bool, optional + creat_montage: bool, optional if the montage is already compatible, this argument should be set to False indicating there is no need to build the montage from the scratch. (default is True) @@ -121,11 +147,11 @@ def make_fnirs_montage( Returns ------- - montage: + montage: instance of DigMontage, a compatible montage with mne standards Note: In MNE-Python the naming of channels MUST follow the structure S#_D# type - where # is replaced by the appropriate source and detector numbers and type is either hbo, hbr or the wavelength. + where # is replaced by the appropriate source and detector numbers and type is either hbo, hbr or the wavelength. """ if create_montage: @@ -134,38 +160,66 @@ def make_fnirs_montage( f.write("# ASA optode file\n") f.write("ReferenceLabel avg\n") f.write("UnitPosition mm\n") - f.write("NumberPositions= " + str(prob_mat['probeInfo']['probes'].item()['nChannel0'].item()+3)+'\n') + f.write( + "NumberPositions= " + + str(prob_mat["probeInfo"]["probes"].item()["nChannel0"].item() + 3) + + "\n" + ) f.write("Positions\n") - f.write(str(Nz[0]) + ' ' + str(Nz[1]) + ' ' +str(Nz[2])+ '\n') - f.write(str(RPA[0]) + ' ' + str(RPA[1]) + ' ' + str(RPA[2]) + '\n') - f.write(str(LPA[0]) + ' ' + str(LPA[1]) + ' ' + str(LPA[2]) + '\n') - sensor_coord = prob_mat['probeInfo']['probes'].item()['coords_s3'].item() - detector_coord = prob_mat['probeInfo']['probes'].item()['coords_d3'].item() + f.write(str(Nz[0]) + " " + str(Nz[1]) + " " + str(Nz[2]) + "\n") + f.write(str(RPA[0]) + " " + str(RPA[1]) + " " + str(RPA[2]) + "\n") + f.write(str(LPA[0]) + " " + str(LPA[1]) + " " + str(LPA[2]) + "\n") + sensor_coord = prob_mat["probeInfo"]["probes"].item()["coords_s3"].item() + detector_coord = prob_mat["probeInfo"]["probes"].item()["coords_d3"].item() for j in range(len(sensor_coord)): - f.write(str(sensor_coord[j][0]) + ' ' + str(sensor_coord[j][1]) + ' ' +str(sensor_coord[j][2])+ '\n') + f.write( + str(sensor_coord[j][0]) + + " " + + str(sensor_coord[j][1]) + + " " + + str(sensor_coord[j][2]) + + "\n" + ) for i in range(len(detector_coord)): - f.write(str(detector_coord [i][0]) + ' ' + str(detector_coord [i][1]) + ' ' +str(detector_coord [i][2])+ '\n') - f.write('Labels\n' + 'Nz\n' + 'RPA\n' + 'LPA\n') + f.write( + str(detector_coord[i][0]) + + " " + + str(detector_coord[i][1]) + + " " + + str(detector_coord[i][2]) + + "\n" + ) + f.write("Labels\n" + "Nz\n" + "RPA\n" + "LPA\n") for k in range(len(detector_labels)): - f.write(str(detector_labels[k]) +'\n') + f.write(str(detector_labels[k]) + "\n") for z in range(len(source_labels)): - f.write(str(source_labels[z]) +'\n') + f.write(str(source_labels[z]) + "\n") f.close() - loc = mne.channels.read_custom_montage('fnirs_montage.elc', head_size=head_size) - os.remove('fnirs_montage.elc') + loc = mne.channels.read_custom_montage("fnirs_montage.elc", head_size=head_size) + os.remove("fnirs_montage.elc") else: if mne_standard is not None: loc = mne.channels.make_standard_montage(mne_standard) else: - loc = mne.channels.read_custom_montage(prob_directory, head_size=0.095, coord_frame='mri') + loc = mne.channels.read_custom_montage( + prob_directory, head_size=0.095, coord_frame="mri" + ) return loc -#Epochs -def fnirs_epoch(fnirs_participant_1: mne.io.Raw, fnirs_participant_2: mne.io.Raw, tmin: float = -0.1, tmax: float = 1, - baseline: tuple = (None, 0), preload: bool = True, event_repeated: str = 'merge'): +# Epochs + +def fnirs_epoch( + fnirs_participant_1: mne.io.Raw, + fnirs_participant_2: mne.io.Raw, + tmin: float = -0.1, + tmax: float = 1, + baseline: tuple = (None, 0), + preload: bool = True, + event_repeated: str = "merge", +): """ Extracts epochs from the raw instances @@ -194,63 +248,84 @@ def fnirs_epoch(fnirs_participant_1: mne.io.Raw, fnirs_participant_2: mne.io.Raw event_repeated: str, optional How to handle duplicates in events[:, 0]. - Can be 'error', to raise an error, + Can be 'error', to raise an error, 'drop' to only retain the row occurring first in the events, or 'merge' to combine the coinciding events (=duplicates) into a new event. (default is 'merge') Returns -------- - a tuple containing mne.Epoch data type of participant #1 and mne.Epoch data type of participant #2 + a tuple containing mne.Epoch data type of participant #1 and mne.Epoch data type of participant #2 """ fnirs_raw_1 = fnirs_participant_1 event1 = mne.events_from_annotations(fnirs_raw_1) events1 = event1[0] events1id = event1[1] - fnirs_epo1 = mne.Epochs(fnirs_raw_1, events1, events1id, tmin=tmin, tmax=tmax, - baseline=baseline, preload=preload, event_repeated=event_repeated) + fnirs_epo1 = mne.Epochs( + fnirs_raw_1, + events1, + events1id, + tmin=tmin, + tmax=tmax, + baseline=baseline, + preload=preload, + event_repeated=event_repeated, + ) fnirs_raw_2 = fnirs_participant_2 event2 = mne.events_from_annotations(fnirs_raw_2) events2 = event2[0] events2id = event2[1] - fnirs_epo2 = mne.Epochs(fnirs_raw_2, events2, events2id, tmin=tmin, tmax=tmax, - baseline=baseline, preload=preload, event_repeated=event_repeated) + fnirs_epo2 = mne.Epochs( + fnirs_raw_2, + events2, + events2id, + tmin=tmin, + tmax=tmax, + baseline=baseline, + preload=preload, + event_repeated=event_repeated, + ) return (fnirs_epo1, fnirs_epo2) -#Get the inputs for building the montage in a more user freindly way -def fnirs_montage_ui(): +# Get the inputs for building the montage in a more user freindly way - """ - Get the inputs for building the montage in a more user freindly way - Arguments - --------- - None +def fnirs_montage_ui(): + """ + Get the inputs for building the montage in a more user freindly way - Returns - -------- - source_labels, detector_labels, Nz, RPA, LPA, head_size: - see make_fnirs_montage corresponding inputs' description + Arguments + --------- + None - """ - source_labels = input("please enter sources names with the S# format: ").split() + Returns + -------- + source_labels, detector_labels, Nz, RPA, LPA, head_size: + see make_fnirs_montage corresponding inputs' description + + """ + source_labels = input("please enter sources names with the S# format: ").split() - detector_labels = input("please enter detectors names with the D# format: ").split() + detector_labels = input("please enter detectors names with the D# format: ").split() - Nz = input("please enter 3D Coordination of tip of the nose x y z in mm: ").split() - for i in range(len(Nz)): - Nz[i] = float(Nz[i]) + Nz = input("please enter 3D Coordination of tip of the nose x y z in mm: ").split() + for i in range(len(Nz)): + Nz[i] = float(Nz[i]) - RPA = input("please enter 3D Coordination of the right preauricular x y z in mm: ").split() - for i in range(len(RPA)): - RPA[i] = float(RPA[i]) + RPA = input( + "please enter 3D Coordination of the right preauricular x y z in mm: " + ).split() + for i in range(len(RPA)): + RPA[i] = float(RPA[i]) - LPA = input("please enter 3D Coordination of the left preauricular x y z in mm: ").split() - for i in range(len(LPA)): - LPA[i] = float(LPA[i]) + LPA = input( + "please enter 3D Coordination of the left preauricular x y z in mm: " + ).split() + for i in range(len(LPA)): + LPA[i] = float(LPA[i]) - head_size = float(input("please enter the head size in mm ")) + head_size = float(input("please enter the head size in mm ")) - return source_labels, detector_labels, Nz, RPA, LPA, head_size + return source_labels, detector_labels, Nz, RPA, LPA, head_size diff --git a/hypyp/fnirs/preprocessor/base_preprocessor.py b/hypyp/fnirs/preprocessor/base_preprocessor.py index ff37239..8416926 100644 --- a/hypyp/fnirs/preprocessor/base_preprocessor.py +++ b/hypyp/fnirs/preprocessor/base_preprocessor.py @@ -4,11 +4,12 @@ from .base_step import BaseStep # Generic type for underlying fnirs implementation (mne raw / cedalion recording) -T = TypeVar('T') +T = TypeVar("T") + class BasePreprocessor(ABC, Generic[T]): @abstractmethod - def read_file(self, path:str, verbose:bool=False) -> T: + def read_file(self, path: str, verbose: bool = False) -> T: """ Load the file to be preprocessed. Does not run the preprocess yet @@ -22,7 +23,7 @@ def read_file(self, path:str, verbose:bool=False) -> T: pass @abstractmethod - def run(self, raw:T, verbose:bool=False): + def run(self, raw: T, verbose: bool = False): """ run all the preprocessing steps on a raw object @@ -34,4 +35,3 @@ def run(self, raw:T, verbose:bool=False): List[BasePreprocessStep[T]]: list of all the intermediary steps of the preprocessing """ pass - diff --git a/hypyp/fnirs/preprocessor/base_step.py b/hypyp/fnirs/preprocessor/base_step.py index 8bee864..30acc44 100644 --- a/hypyp/fnirs/preprocessor/base_step.py +++ b/hypyp/fnirs/preprocessor/base_step.py @@ -1,23 +1,24 @@ from abc import ABC, abstractmethod from typing import List, Generic, TypeVar -PREPROCESS_STEP_BASE_KEY = 'base' -PREPROCESS_STEP_BASE_DESC = 'Loaded data' +PREPROCESS_STEP_BASE_KEY = "base" +PREPROCESS_STEP_BASE_DESC = "Loaded data" -PREPROCESS_STEP_OD_KEY = 'od' -PREPROCESS_STEP_OD_DESC = 'Optical density' +PREPROCESS_STEP_OD_KEY = "od" +PREPROCESS_STEP_OD_DESC = "Optical density" -PREPROCESS_STEP_OD_CLEAN_KEY = 'od_clean' -PREPROCESS_STEP_OD_CLEAN_DESC = 'Optical density cleaned' +PREPROCESS_STEP_OD_CLEAN_KEY = "od_clean" +PREPROCESS_STEP_OD_CLEAN_DESC = "Optical density cleaned" -PREPROCESS_STEP_HAEMO_KEY = 'haemo' -PREPROCESS_STEP_HAEMO_DESC = 'Hemoglobin' +PREPROCESS_STEP_HAEMO_KEY = "haemo" +PREPROCESS_STEP_HAEMO_DESC = "Hemoglobin" -PREPROCESS_STEP_HAEMO_FILTERED_KEY = 'haemo_filtered' -PREPROCESS_STEP_HAEMO_FILTERED_DESC = 'Hemoglobin Band-pass Filtered' +PREPROCESS_STEP_HAEMO_FILTERED_KEY = "haemo_filtered" +PREPROCESS_STEP_HAEMO_FILTERED_DESC = "Hemoglobin Band-pass Filtered" # Generic type for underlying fnirs implementation (mne raw / cedalion recording) -T = TypeVar('T') +T = TypeVar("T") + class BaseStep(ABC, Generic[T]): """ @@ -28,11 +29,12 @@ class BaseStep(ABC, Generic[T]): key (str): identifier for the step desc (str | None, optional): description of the setup. Defaults to "key" value. """ + obj: T key: str desc: str - def __init__(self, obj:T, key:str, desc:str|None=None): + def __init__(self, obj: T, key: str, desc: str | None = None): self.obj = obj self.key = key if desc is None: diff --git a/hypyp/fnirs/preprocessor/implementations/cedalion_preprocessor.py b/hypyp/fnirs/preprocessor/implementations/cedalion_preprocessor.py index aa19423..ed52430 100644 --- a/hypyp/fnirs/preprocessor/implementations/cedalion_preprocessor.py +++ b/hypyp/fnirs/preprocessor/implementations/cedalion_preprocessor.py @@ -25,19 +25,19 @@ class CedalionStep(BaseStep[cdc.Recording]): @property def n_times(self): - return len(self.obj['time']) - + return len(self.obj["time"]) + @property def sfreq(self): # TODO We probably have this information directly - return 1 / float(self.obj['time'][1] - self.obj['time'][0]) - + return 1 / float(self.obj["time"][1] - self.obj["time"][0]) + @property def ch_names(self): # TODO learn how to use xarray efficiently # TODO maybe we should send channel and wavelength return list(self.obj.channel.to_series()) - + def plot(self, **_kwargs): # select which time series we work with @@ -51,50 +51,60 @@ def plot(self, **_kwargs): ax[i].plot(ts.time, ts.sel(channel=ch, chromo="HbR"), "b-", label="HbR") ax[i].set_title(f"Ch. {ch}") # add stim markers using Cedalion's plot_stim_markers function - #cedalion.plots.plot_stim_markers(ax[i], rec.stim, y=1) + # cedalion.plots.plot_stim_markers(ax[i], rec.stim, y=1) ax[i].set_ylabel(r"$\Delta$ c / uM") ax[0].legend(ncol=6) ax[3].set_label("time / s") - ax[3].set_xlim(0,100) - #plt.tight_layout() + ax[3].set_xlim(0, 100) + # plt.tight_layout() return f class CedalionPreprocessor(BasePreprocessor[cdc.Recording]): - def read_file(self, path, verbose:bool=False) -> cdc.Recording: + def read_file(self, path, verbose: bool = False) -> cdc.Recording: # TODO honor verbose arg if not DataBrowser.is_path_snirf(path): - raise RuntimeError('Not implemented: only snirf file is supported for now') + raise RuntimeError( + "Not implemented: only snirf file is supported for now" + ) recordings = cedalion.io.read_snirf(path) - rec: cdc.Recording = recordings[0] + rec: cdc.Recording = recordings[0] # TODO we should force units only if they are not set - rec['amp'] = rec['amp'].pint.dequantify().pint.quantify("V") - rec['amp']['time'] = rec['amp']['time'].pint.dequantify().pint.quantify("second") + rec["amp"] = rec["amp"].pint.dequantify().pint.quantify("V") + rec["amp"]["time"] = ( + rec["amp"]["time"].pint.dequantify().pint.quantify("second") + ) return rec def run(self, rec, verbose: bool = False) -> list[CedalionStep]: # TODO honor verbose steps = [] - amp = rec['amp'] - steps.append(CedalionStep(amp, PREPROCESS_STEP_BASE_KEY, PREPROCESS_STEP_BASE_DESC)) + amp = rec["amp"] + steps.append( + CedalionStep(amp, PREPROCESS_STEP_BASE_KEY, PREPROCESS_STEP_BASE_DESC) + ) od = cedalion.nirs.int2od(amp) - steps.append(CedalionStep(od, PREPROCESS_STEP_OD_KEY, PREPROCESS_STEP_OD_DESC)) - + steps.append( + CedalionStep(od, PREPROCESS_STEP_OD_KEY, PREPROCESS_STEP_OD_DESC) + ) dpf = xr.DataArray( - [6., 6.], - dims="wavelength", - coords={"wavelength" : [760., 850.]}) # TODO unhardcode wavelengts + [6.0, 6.0], dims="wavelength", coords={"wavelength": [760.0, 850.0]} + ) # TODO unhardcode wavelengts - haemo = cedalion.nirs.beer_lambert(rec['amp'], rec.geo3d, dpf) - steps.append(CedalionStep(haemo, PREPROCESS_STEP_HAEMO_KEY, PREPROCESS_STEP_HAEMO_DESC)) + haemo = cedalion.nirs.beer_lambert(rec["amp"], rec.geo3d, dpf) + steps.append( + CedalionStep( + haemo, PREPROCESS_STEP_HAEMO_KEY, PREPROCESS_STEP_HAEMO_DESC + ) + ) return steps - -except: + +except: CedalionPreprocesStep = None - CedalionPreprocessor = None \ No newline at end of file + CedalionPreprocessor = None diff --git a/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_as_is.py b/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_as_is.py index 5cab9f9..7152e85 100644 --- a/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_as_is.py +++ b/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_as_is.py @@ -5,21 +5,23 @@ from ..base_preprocessor import BasePreprocessor from .mne_step import MneStep + # This is the same as MnePreprocessorBasic, but without the default preprocessor class MnePreprocessorAsIs(BasePreprocessor[mne.io.Raw]): """ The MnePreprocessorAsIs class uses the loaded data as-is, which means that one should provide haemoglobin concentration already prepared in a previous - preprocessing pipeline. + preprocessing pipeline. Use this preprocessor if you already cleaned your data and have haemoglobin concentration. This is the default preprocessor. """ + def __init__(self): super().__init__() - - def read_file(self, path:str, verbose:bool=False): + + def read_file(self, path: str, verbose: bool = False): if DataBrowser.is_path_fif(path): return mne.io.read_raw_fif(path, preload=True, verbose=verbose) @@ -30,12 +32,11 @@ def read_file(self, path:str, verbose:bool=False): return mne.io.read_raw_snirf(path, preload=True, verbose=verbose) raise ValueError(f"No reader for file {path}") - + def run(self, raw: mne.io.Raw, verbose: bool = False) -> list[MneStep]: if verbose: - print('Using MnePreprocessorAsIs, using raw data as already preprocessed') + print("Using MnePreprocessorAsIs, using raw data as already preprocessed") step = MneStep(raw, PREPROCESS_STEP_BASE_KEY, PREPROCESS_STEP_BASE_DESC) steps = [step] return steps - diff --git a/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_raw_to_haemo.py b/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_raw_to_haemo.py index 8a8223e..17a0d0d 100644 --- a/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_raw_to_haemo.py +++ b/hypyp/fnirs/preprocessor/implementations/mne_preprocessor_raw_to_haemo.py @@ -6,10 +6,11 @@ from .mne_step import MneStep from .mne_preprocessor_as_is import MnePreprocessorAsIs + class MnePreprocessorRawToHaemo(MnePreprocessorAsIs): """ The MnePreprocessorRawToHaemo class uses mne fNIRS features to run some basic preprocessing steps. - + It does these: 1. Store the raw data as a step @@ -26,20 +27,27 @@ class MnePreprocessorRawToHaemo(MnePreprocessorAsIs): already preprocessed and will be returned as-is. """ - def __init__(self, ppf=6.0, hp_filter=0.02, lp_filter=0.7, quality_sci_threshold=0.1): + + def __init__( + self, ppf=6.0, hp_filter=0.02, lp_filter=0.7, quality_sci_threshold=0.1 + ): super().__init__() self.ppf = ppf self.hp_filter = hp_filter self.lp_filter = lp_filter self.quality_sci_threshold = quality_sci_threshold - - def run(self, raw:mne.io.Raw, verbose:bool=False) -> list[MneStep]: + + def run(self, raw: mne.io.Raw, verbose: bool = False) -> list[MneStep]: # If have hbo or hbr, it means it is already in concentration - if len(mne.pick_types(raw.info, fnirs=['hbo', 'hbr'])) > 0: - raise ValueError('Loaded data seems to already be in concentrations. Use MnePreprocessorAsIs instead.') + if len(mne.pick_types(raw.info, fnirs=["hbo", "hbr"])) > 0: + raise ValueError( + "Loaded data seems to already be in concentrations. Use MnePreprocessorAsIs instead." + ) if verbose: - print('Using MnePreprocessorRawToHaemo, converting fNIRS raw data to haemoglobin concentrations') + print( + "Using MnePreprocessorRawToHaemo, converting fNIRS raw data to haemoglobin concentrations" + ) steps = [] steps.append(MneStep(raw, PREPROCESS_STEP_BASE_KEY, PREPROCESS_STEP_BASE_DESC)) @@ -48,17 +56,36 @@ def run(self, raw:mne.io.Raw, verbose:bool=False) -> list[MneStep]: steps.append(MneStep(raw_od, PREPROCESS_STEP_OD_KEY, PREPROCESS_STEP_OD_DESC)) quality_sci = mne.preprocessing.nirs.scalp_coupling_index(raw_od) - raw_od.info['bads'] = list(compress(raw_od.ch_names, quality_sci < self.quality_sci_threshold)) - picks = mne.pick_types(raw_od.info, fnirs=True, exclude='bads') + raw_od.info["bads"] = list( + compress(raw_od.ch_names, quality_sci < self.quality_sci_threshold) + ) + picks = mne.pick_types(raw_od.info, fnirs=True, exclude="bads") raw_od_clean = raw_od.copy().pick(picks) - steps.append(MneStep(raw_od_clean, PREPROCESS_STEP_OD_CLEAN_KEY, PREPROCESS_STEP_OD_CLEAN_DESC)) + steps.append( + MneStep( + raw_od_clean, + PREPROCESS_STEP_OD_CLEAN_KEY, + PREPROCESS_STEP_OD_CLEAN_DESC, + ) + ) # For partial pathlength, see https://doi.org/10.1117/1.JBO.18.10.105004, Table 1 - raw_haemo: mne.io.Raw = mne.preprocessing.nirs.beer_lambert_law(raw_od_clean.copy(), ppf=self.ppf) - steps.append(MneStep(raw_haemo, PREPROCESS_STEP_HAEMO_KEY, PREPROCESS_STEP_HAEMO_DESC)) + raw_haemo: mne.io.Raw = mne.preprocessing.nirs.beer_lambert_law( + raw_od_clean.copy(), ppf=self.ppf + ) + steps.append( + MneStep(raw_haemo, PREPROCESS_STEP_HAEMO_KEY, PREPROCESS_STEP_HAEMO_DESC) + ) - raw_haemo_filtered = raw_haemo.copy().filter(self.hp_filter, self.lp_filter, verbose=verbose, method='iir') - steps.append(MneStep(raw_haemo_filtered, PREPROCESS_STEP_HAEMO_FILTERED_KEY, PREPROCESS_STEP_HAEMO_FILTERED_DESC)) + raw_haemo_filtered = raw_haemo.copy().filter( + self.hp_filter, self.lp_filter, verbose=verbose, method="iir" + ) + steps.append( + MneStep( + raw_haemo_filtered, + PREPROCESS_STEP_HAEMO_FILTERED_KEY, + PREPROCESS_STEP_HAEMO_FILTERED_DESC, + ) + ) return steps - \ No newline at end of file diff --git a/hypyp/fnirs/preprocessor/implementations/mne_step.py b/hypyp/fnirs/preprocessor/implementations/mne_step.py index 201c983..d940f9b 100644 --- a/hypyp/fnirs/preprocessor/implementations/mne_step.py +++ b/hypyp/fnirs/preprocessor/implementations/mne_step.py @@ -2,19 +2,19 @@ from ..base_step import BaseStep + class MneStep(BaseStep[mne.io.Raw]): @property def n_times(self): return self.obj.n_times - + @property def sfreq(self): - return self.obj.info['sfreq'] - + return self.obj.info["sfreq"] + @property def ch_names(self): return self.obj.ch_names - + def plot(self, **kwargs): return self.obj.plot(**kwargs) - diff --git a/hypyp/fnirs/recording.py b/hypyp/fnirs/recording.py index 411fa6d..3686f1d 100644 --- a/hypyp/fnirs/recording.py +++ b/hypyp/fnirs/recording.py @@ -24,6 +24,7 @@ from ..wavelet.cwt import CWT from ..wavelet.wtc import WTC + class Recording: """ The Recording object encapsulates the logic around the recording for one participant. @@ -34,20 +35,21 @@ class Recording: tasks (TaskList, optional): list of tasks during the recording of the participant, that will be extracted from events in the raw files to build epochs. Defaults to []. channel_roi (ChannelROI | None, optional): region of interest object to group channels. Defaults to None. """ + filepath: str | None subject_label: str channel_roi: ChannelROI | None mne_raw: mne.io.Raw | None - intra_wtcs: List[WTC] | None # intra-subject wtc + intra_wtcs: List[WTC] | None # intra-subject wtc epochs_per_task: List[mne.Epochs] | None tasks: TaskList preprocess_steps: List[BaseStep] | None def __init__( self, - subject_label:str='', - tasks:TaskList=[], - channel_roi:ChannelROI|None=None + subject_label: str = "", + tasks: TaskList = [], + channel_roi: ChannelROI | None = None, ): self.filepath = None self.subject_label = subject_label @@ -61,24 +63,36 @@ def __init__( if len(tasks) == 0: # Use tasks with special values instead of time_range, # since we don't know yet the duration of the record - self.tasks = [Task(TASK_NAME_WHOLE_RECORD, onset_event_id=TASK_BEGINNING, offset_event_id=TASK_END)] + self.tasks = [ + Task( + TASK_NAME_WHOLE_RECORD, + onset_event_id=TASK_BEGINNING, + offset_event_id=TASK_END, + ) + ] def _assert_is_preprocessed(self): if not self.is_preprocessed: - raise RuntimeError('Recording is not preprocessed. Did you run preprocess() ?') + raise RuntimeError( + "Recording is not preprocessed. Did you run preprocess() ?" + ) def _assert_is_epochs_loaded(self): if not self.is_epochs_loaded: - raise RuntimeError('Recording does not have epochs loaded. Did you run populate_epochs_from_tasks() ?') + raise RuntimeError( + "Recording does not have epochs loaded. Did you run populate_epochs_from_tasks() ?" + ) def _assert_is_wtcs_computed(self): if not self.intra_wtcs: - raise RuntimeError('Recording has no intra_wtcs. Did you run compute_wtcs(with_intra=True) ?') + raise RuntimeError( + "Recording has no intra_wtcs. Did you run compute_wtcs(with_intra=True) ?" + ) @staticmethod def get_default_preprocessor(): return MnePreprocessorAsIs() - + @property def task_keys(self): return [task.name for task in self.tasks] @@ -86,11 +100,11 @@ def task_keys(self): @property def is_preprocessed(self) -> bool: return self.preprocess_steps is not None - + @property def is_epochs_loaded(self) -> bool: return self.epochs_per_task is not None - + @property def is_wtc_computed(self): return self.intra_wtcs is not None @@ -104,7 +118,7 @@ def preprocessed(self) -> mne.io.Raw: self._assert_is_preprocessed() # We want the last step of all the preprocessing return self.preprocess_steps[-1].obj - + @property def cwts(self) -> List[CWT]: self._assert_is_wtcs_computed() @@ -114,11 +128,10 @@ def cwts(self) -> List[CWT]: seen.append(wtc.cwt1) return seen - @property def pre(self) -> mne.io.Raw: return self.preprocessed - + @property def mne_preprocessed(self) -> mne.io.Raw: return self.preprocessed @@ -144,7 +157,7 @@ def preprocess_step_keys(self): self._assert_is_preprocessed() # get in the reverse order so that the last step is first in list keys = [] - for i in range(len(self.preprocess_steps)-1, -1, -1): + for i in range(len(self.preprocess_steps) - 1, -1, -1): keys.append(self.preprocess_steps[i].key) return keys @@ -152,24 +165,24 @@ def preprocess_step_keys(self): def preprocess_step_choices(self): self._assert_is_preprocessed() steps_dict = dict() - for i in range(len(self.preprocess_steps)-1, -1, -1): + for i in range(len(self.preprocess_steps) - 1, -1, -1): step = self.preprocess_steps[i] steps_dict[step.key] = step.desc return steps_dict - + def _fill_subject_label(self): - if self.subject_label == '': - self.subject_label = self.mne_raw.info['subject_info']['his_id'] - - if self.subject_label == '': + if self.subject_label == "": + self.subject_label = self.mne_raw.info["subject_info"]["his_id"] + + if self.subject_label == "": self.subject_label = generate_random_label(10) - + def load_file( self, - filepath:str, - preprocessor:BasePreprocessor|None=None, + filepath: str, + preprocessor: BasePreprocessor | None = None, preprocess=True, - verbose=False + verbose=False, ): """ Load a raw NIRS file as the recording of the Recording @@ -190,9 +203,9 @@ def load_file( preprocessor = Recording.get_default_preprocessor() if not Path(filepath).is_file(): - raise RuntimeError(f'Cannot find file {filepath}') + raise RuntimeError(f"Cannot find file {filepath}") - self.filepath = filepath + self.filepath = filepath self.mne_raw = preprocessor.read_file(filepath, verbose=verbose) self._fill_subject_label() @@ -201,13 +214,13 @@ def load_file( self.preprocess(preprocessor, verbose=verbose) return self - + def load_raw( self, - raw:mne.io.Raw, - preprocessor:BasePreprocessor|None=None, + raw: mne.io.Raw, + preprocessor: BasePreprocessor | None = None, preprocess=True, - verbose=False + verbose=False, ): """ Use an existing Raw object to load the data into a Recording @@ -231,8 +244,8 @@ def load_raw( self.preprocess(preprocessor, verbose=verbose) return self - - def preprocess(self, preprocessor:BasePreprocessor, verbose:bool=False): + + def preprocess(self, preprocessor: BasePreprocessor, verbose: bool = False): """ Run the preprocessing for the raw recording of the Recording @@ -247,7 +260,7 @@ def preprocess(self, preprocessor:BasePreprocessor, verbose:bool=False): self.populate_epochs_from_tasks(verbose=verbose) return self - def get_preprocess_step(self, key:str) -> BaseStep: + def get_preprocess_step(self, key: str) -> BaseStep: """ Get a specific step of the preprocessing pipeline @@ -263,8 +276,8 @@ def get_preprocess_step(self, key:str) -> BaseStep: return step raise RuntimeError(f'No preprocess step named "{key}"') - - def populate_epochs_from_tasks(self, verbose:bool=False): + + def populate_epochs_from_tasks(self, verbose: bool = False): """ Given the list of tasks (annotations and timed) of the Recording, find the given data in the preprocessed channels and load as epochs that will be compared between subjects @@ -276,10 +289,12 @@ def populate_epochs_from_tasks(self, verbose:bool=False): Returns: self: the Recording object itself. Useful for chaining operations """ - self.epochs_per_task = epochs_from_tasks(self.preprocessed, self.tasks, verbose=verbose) + self.epochs_per_task = epochs_from_tasks( + self.preprocessed, self.tasks, verbose=verbose + ) return self - - def get_epochs_for_task(self, task_name:str) -> List[mne.Epochs]: + + def get_epochs_for_task(self, task_name: str) -> List[mne.Epochs]: """ Given a task name, return all the Epochs object @@ -292,7 +307,9 @@ def get_epochs_for_task(self, task_name:str) -> List[mne.Epochs]: self._assert_is_epochs_loaded() id = None try: - id = [i for i in range(len(self.task_keys)) if self.task_keys[i] == task_name][0] + id = [ + i for i in range(len(self.task_keys)) if self.task_keys[i] == task_name + ][0] except: pass @@ -300,8 +317,8 @@ def get_epochs_for_task(self, task_name:str) -> List[mne.Epochs]: raise RuntimeError(f'Cannot find epochs for task "{task_name}"') return self.epochs_per_task[id] - - def get_roi_from_channel(self, ch_name:str) -> str: + + def get_roi_from_channel(self, ch_name: str) -> str: """ Given a channel name, return the region of interest it belongs to, if any @@ -312,17 +329,17 @@ def get_roi_from_channel(self, ch_name:str) -> str: str: name of the region of interest """ if self.channel_roi is None: - return '' + return "" return self.channel_roi.get_roi_from_channel(ch_name) - - def get_channel_to_standard_montage_map(self, as_data_frame:bool=True): + + def get_channel_to_standard_montage_map(self, as_data_frame: bool = True): my_montage = self.mne_raw.get_montage() - my_positions_dict = my_montage.get_positions()['ch_pos'] + my_positions_dict = my_montage.get_positions()["ch_pos"] # See all montages with mne.channels.get_builtin_montages() # print(mne.channels.get_builtin_montages()) - standard_montage = mne.channels.make_standard_montage('standard_1020') - standard_positions_dict = standard_montage.get_positions()['ch_pos'] + standard_montage = mne.channels.make_standard_montage("standard_1020") + standard_positions_dict = standard_montage.get_positions()["ch_pos"] standard_keys = list(standard_positions_dict.keys()) standard_positions = list(standard_positions_dict.values()) @@ -349,30 +366,34 @@ def get_channel_to_standard_montage_map(self, as_data_frame:bool=True): # Find all distances, and nearest in standard for k in my_positions_lookup.keys(): ch_pos_2d = my_positions_lookup[k].reshape(1, -1) # reshape to 2D - + # Get nearest standard electrode dists = cdist(ch_pos_2d, standard_positions) - + nearest_idx = np.argmin(dists) mapped_dict[k] = { - 'name': standard_keys[nearest_idx], - 'dist': dists[0, nearest_idx], + "name": standard_keys[nearest_idx], + "dist": dists[0, nearest_idx], } if as_data_frame: - return key_dict_as_data_frame(mapped_dict, col_names=['Source/Detector', 'Standard 1020', 'Distance']) + return key_dict_as_data_frame( + mapped_dict, col_names=["Source/Detector", "Standard 1020", "Distance"] + ) return mapped_dict def plot_steps_for_channel(self, ch_name: str, show_cwt=True): # get only the "S1_D1" part of "S1_D1 760" - ch_base_name = re.sub(' .*$', '', ch_name) + ch_base_name = re.sub(" .*$", "", ch_name) rows = len(self.preprocess_steps) if show_cwt: rows = rows + 1 height = rows * 3 - fig, axes = plt.subplots(rows, 1, figsize=(12, height), sharex=True, sharey=False) + fig, axes = plt.subplots( + rows, 1, figsize=(12, height), sharex=True, sharey=False + ) axes = np.atleast_1d(axes) for j, step in enumerate(self.preprocess_steps): @@ -381,31 +402,37 @@ def plot_steps_for_channel(self, ch_name: str, show_cwt=True): for k, local_ch_name in enumerate(raw.ch_names): if local_ch_name == f"{ch_base_name} 760": - axes[j].plot(raw.times, raw_data[k,:], color='g', label=local_ch_name) + axes[j].plot( + raw.times, raw_data[k, :], color="g", label=local_ch_name + ) if local_ch_name == f"{ch_base_name} 850": - axes[j].plot(raw.times, raw_data[k,:], color='C1', label=local_ch_name) + axes[j].plot( + raw.times, raw_data[k, :], color="C1", label=local_ch_name + ) if local_ch_name == f"{ch_base_name} hbo": - axes[j].plot(raw.times, raw_data[k,:], color='r', label=local_ch_name) + axes[j].plot( + raw.times, raw_data[k, :], color="r", label=local_ch_name + ) if local_ch_name == f"{ch_base_name} hbr": - axes[j].plot(raw.times, raw_data[k,:], color='b', label=local_ch_name) + axes[j].plot( + raw.times, raw_data[k, :], color="b", label=local_ch_name + ) axes[j].legend() axes[j].set_title(step.desc) found_cwt = None if show_cwt: for cwt in self.cwts: - cwt_base_name = re.sub(' .*$', '', cwt.label) + cwt_base_name = re.sub(" .*$", "", cwt.label) if cwt_base_name == ch_base_name: found_cwt = cwt break if found_cwt is None: raise ValueError(f"Cannot find CWT for ch_name '{ch_name}'") - cwt.plot(ax=axes[rows-1], show_colorbar=False) + cwt.plot(ax=axes[rows - 1], show_colorbar=False) fig.suptitle(f"Steps for {ch_name}") fig.tight_layout() return fig - - diff --git a/hypyp/fnirs/study.py b/hypyp/fnirs/study.py index fd4cc08..4610aa5 100644 --- a/hypyp/fnirs/study.py +++ b/hypyp/fnirs/study.py @@ -13,9 +13,10 @@ plot_coherence_connectogram_split, ) -class Study(): + +class Study: dyads: List[Dyad] - dyads_shuffled: List[Dyad]|None + dyads_shuffled: List[Dyad] | None def __init__(self, dyads: List[Dyad] = []): """ @@ -44,7 +45,7 @@ def is_wtc_shuffle_computed(self) -> bool: if not dyad.is_wtc_computed: return False return True - + @property def df(self) -> CoherenceDataFrame: """The pandas dataframe object from computed WTCs""" @@ -64,22 +65,22 @@ def preprocess(self, preprocessor: BasePreprocessor): if not dyad.is_preprocessed: dyad.preprocess(preprocessor) return self - + def compute_wtcs( self, # All the arguments are a copy from Dyad.compute_wtcs, but we need them explicitely for type hinting - wavelet: BaseWavelet|None = None, - ch_match: PairChannelMatchType|None = None, - only_time_range: Tuple[float,float]|None = None, - bin_seconds: float|None = None, - period_cuts: List[float]|None = None, - frequency_cuts: List[float]|None = None, + wavelet: BaseWavelet | None = None, + ch_match: PairChannelMatchType | None = None, + only_time_range: Tuple[float, float] | None = None, + bin_seconds: float | None = None, + period_cuts: List[float] | None = None, + frequency_cuts: List[float] | None = None, verbose: bool = False, with_intra: bool = True, - downsample: int|None = None, + downsample: int | None = None, keep_wtcs: bool = True, show_time_estimation: bool = True, - ): + ): """ Wraps the `compute_wtcs` of all the dyads. Arguments are directly passed to the dyads method @@ -122,13 +123,17 @@ def compute_wtcs( tracker.stop() if show_time_estimation: self._print_time_estimation(tracker.duration, len(self.dyads)) - + return self - + def _print_time_estimation(self, single_duration, count): - print(f'Time for computing one dyad: {TimeTracker.human_readable_duration(single_duration)}') - print(f'Expected time for {count} dyads: {TimeTracker.human_readable_duration(single_duration * count)}') - + print( + f"Time for computing one dyad: {TimeTracker.human_readable_duration(single_duration)}" + ) + print( + f"Expected time for {count} dyads: {TimeTracker.human_readable_duration(single_duration * count)}" + ) + def estimate_wtcs_run_time(self, *args, **kwargs): """ Computes the WTC for one dyad and print the expected run time for the whole study @@ -144,13 +149,13 @@ def estimate_wtcs_run_time(self, *args, **kwargs): self._print_time_estimation(tracker.duration, len(self.dyads)) return self - + def reset(self): self._clear_dyads_shuffle() for dyad in self.dyads: dyad.reset() return self - + def _clear_dyads_shuffle(self): """ Delete all the shuffle dyads that have been created @@ -160,7 +165,7 @@ def _clear_dyads_shuffle(self): """ self.dyads_shuffled = None return self - + def get_dyads_shuffle(self) -> List[Dyad]: """ Get a list of permutated recording pairs, useful for statistical analysis. @@ -173,7 +178,14 @@ def get_dyads_shuffle(self) -> List[Dyad]: for j, dyad2 in enumerate(self.dyads): if i == j: continue - dyads_shuffle.append(Dyad(dyad1.s1, dyad2.s2, label=f'shuffle s1:{dyad1.label}-s2:{dyad2.label}', is_pseudo=True)) + dyads_shuffle.append( + Dyad( + dyad1.s1, + dyad2.s2, + label=f"shuffle s1:{dyad1.label}-s2:{dyad2.label}", + is_pseudo=True, + ) + ) return dyads_shuffle # TODO add as argument the number of shuffle dyads @@ -188,7 +200,7 @@ def compute_wtcs_shuffle(self, *args, **kwargs): for dyad_shuffle in self.dyads_shuffled: dyad_shuffle.compute_wtcs(*args, **kwargs, with_intra=False) return self - + def get_coherence_df(self) -> CoherenceDataFrame: """ Loop over every dyad and concatenate all the pandas dataframes into one @@ -203,7 +215,7 @@ def get_coherence_df(self) -> CoherenceDataFrame: """ dfs = [] if not self.is_wtc_computed: - raise RuntimeError('wtc not computed') + raise RuntimeError("wtc not computed") for dyad in self.dyads: dfs.append(dyad.df) @@ -213,18 +225,19 @@ def get_coherence_df(self) -> CoherenceDataFrame: dfs.append(dyad_shuffle.df) return CoherenceDataFrame.concat(dfs) - + # # Plots # def plot_coherence_matrix( - self, - field1:str='channel1', - field2:str='channel2', - query:str | None=None, - s1_label:str='Subject1', - s2_label:str='Subject2', - **kwargs): + self, + field1: str = "channel1", + field2: str = "channel2", + query: str | None = None, + s1_label: str = "Subject1", + s2_label: str = "Subject2", + **kwargs, + ): """ Plot the computed coherence metric for pair of fields (channel or roi) in a matrix format @@ -236,7 +249,7 @@ def plot_coherence_matrix( df = self.df if query is not None: df = df.query(query) - + ch_names1 = [] ch_names2 = [] @@ -249,45 +262,34 @@ def plot_coherence_matrix( if ch_name not in ch_names2: ch_names2.append(ch_name) - ordered_names = ch_names1 + [name for name in ch_names2 if name not in ch_names1] + ordered_names = ch_names1 + [ + name for name in ch_names2 if name not in ch_names1 + ] # TODO: missing ordered roi return plot_coherence_matrix( - df, - s1_label, - s2_label, - field1, - field2, - ordered_names, - **kwargs) - - def plot_coherence_matrix_per_channel(self, query:str|None=None, **kwargs): + df, s1_label, s2_label, field1, field2, ordered_names, **kwargs + ) + + def plot_coherence_matrix_per_channel(self, query: str | None = None, **kwargs): """ Wraps plot_coherence_matrix to plot per channel Args: query (str | None, optional): pandas query to filter the dataframe. Defaults to None. """ - return self.plot_coherence_matrix( - 'channel1', - 'channel2', - query, - **kwargs) - - def plot_coherence_matrix_per_roi(self, query:str|None=None, **kwargs): + return self.plot_coherence_matrix("channel1", "channel2", query, **kwargs) + + def plot_coherence_matrix_per_roi(self, query: str | None = None, **kwargs): """ Wraps plot_coherence_matrix to plot per region of interest Args: query (str | None, optional): pandas query to filter the dataframe. Defaults to None. """ - return self.plot_coherence_matrix( - 'roi1', - 'roi2', - query, - **kwargs) - - def plot_coherence_matrix_per_channel_for_task(self, task:str, **kwargs): + return self.plot_coherence_matrix("roi1", "roi2", query, **kwargs) + + def plot_coherence_matrix_per_channel_for_task(self, task: str, **kwargs): """ Wraps plot_coherence_matrix_per_channel to plot for a specific task @@ -295,12 +297,10 @@ def plot_coherence_matrix_per_channel_for_task(self, task:str, **kwargs): task (str): task name """ return self.plot_coherence_matrix( - 'channel1', - 'channel2', - query=f'task=="{task}"', - **kwargs) - - def plot_coherence_matrix_per_roi_for_task(self, task:str, **kwargs): + "channel1", "channel2", query=f'task=="{task}"', **kwargs + ) + + def plot_coherence_matrix_per_roi_for_task(self, task: str, **kwargs): """ Wraps plot_coherence_matrix_per_roi to plot for a specific task @@ -308,62 +308,62 @@ def plot_coherence_matrix_per_roi_for_task(self, task:str, **kwargs): task (str): task name """ return self.plot_coherence_matrix( - 'roi1', - 'roi2', - query=f'task=="{task}"', - **kwargs) - + "roi1", "roi2", query=f'task=="{task}"', **kwargs + ) + def plot_coherence_bars_per_task(self, **kwargs): """ Plot coherence metric per task for comparison """ - return plot_coherence_bars_per_task( - self.df, - **kwargs) - + return plot_coherence_bars_per_task(self.df, **kwargs) + def plot_coherence_connectogram( - self, - query:str|None=None, - title:str|None=None, - s1_label:str='Subject1', - s2_label:str='Subject2', - **kwargs): + self, + query: str | None = None, + title: str | None = None, + s1_label: str = "Subject1", + s2_label: str = "Subject2", + **kwargs, + ): df = self.df.copy() - selector = df['is_intra']==False + selector = df["is_intra"] == False df_filtered = df[selector] if query is not None: df_filtered = df_filtered.query(query) # rename to have them separated in the plot - df_filtered.loc[:, 'roi1'] = s1_label + '_' + df_filtered['roi1'].astype(str) - df_filtered.loc[:, 'roi2'] = s2_label + '_' + df_filtered['roi2'].astype(str) + df_filtered.loc[:, "roi1"] = s1_label + "_" + df_filtered["roi1"].astype(str) + df_filtered.loc[:, "roi2"] = s2_label + "_" + df_filtered["roi2"].astype(str) - pivot = df_filtered.pivot_table(index='roi1', columns='roi2', values='coherence', aggfunc='mean') + pivot = df_filtered.pivot_table( + index="roi1", columns="roi2", values="coherence", aggfunc="mean" + ) if title is None: - title = f'{s1_label} / {s2_label}' + title = f"{s1_label} / {s2_label}" - return plot_coherence_connectogram_split( - pivot, - title=title, - **kwargs) + return plot_coherence_connectogram_split(pivot, title=title, **kwargs) - def plot_coherence_connectogram_intra(self, is_intra_of:int, query:str|None=None, **kwargs): + def plot_coherence_connectogram_intra( + self, is_intra_of: int, query: str | None = None, **kwargs + ): df = self.df - selector = (df['is_intra']==True) & (df['is_intra_of']==is_intra_of) + selector = (df["is_intra"] == True) & (df["is_intra_of"] == is_intra_of) df_filtered = df[selector] if query is not None: df_filtered = df_filtered.query(query) - pivot = df_filtered.pivot_table(index='roi1', columns='roi2', values='coherence', aggfunc='mean') + pivot = df_filtered.pivot_table( + index="roi1", columns="roi2", values="coherence", aggfunc="mean" + ) return plot_coherence_connectogram(pivot, **kwargs) - def plot_coherence_connectogram_s1(self, query:str|None=None, **kwargs): + def plot_coherence_connectogram_s1(self, query: str | None = None, **kwargs): return self.plot_coherence_connectogram_intra(1, query, **kwargs) - def plot_coherence_connectogram_s2(self, query:str|None=None, **kwargs): + def plot_coherence_connectogram_s2(self, query: str | None = None, **kwargs): return self.plot_coherence_connectogram_intra(2, query, **kwargs) # @@ -380,7 +380,7 @@ def from_pickle(file_path: str): Returns: Study: the Study object """ - with open(file_path, 'rb') as f: + with open(file_path, "rb") as f: return pickle.load(f) def save_pickle(self, file_path: str): @@ -390,7 +390,7 @@ def save_pickle(self, file_path: str): Args: file_path (str): disk path for the serialisation """ - with open(file_path, 'wb') as f: + with open(file_path, "wb") as f: pickle.dump(self, f) def save_feather(self, file_path: str): diff --git a/hypyp/mvarica.py b/hypyp/mvarica.py index f95dfaf..13ca002 100644 --- a/hypyp/mvarica.py +++ b/hypyp/mvarica.py @@ -1,4 +1,4 @@ -"""" +""" " Essential Functions for applying MVARICA process on input signals. """ @@ -10,61 +10,61 @@ class MVAR: """ Multivariate Vector Autoregressive Model implementation. - + This class implements methods for fitting, predicting with, and checking stability of MVAR models. MVAR models are useful for analyzing directed interactions between multiple time series by modeling how past values of all series affect current values of each series. - + Parameters ---------- model_order : int Order of the MVAR model, indicating how many past time points influence the current value. Higher orders can capture more complex temporal dependencies but require more data to fit. - + fitting_method : str or object, optional Method used for fitting the MVAR model (default='default'). Options: - 'default': Uses least squares method if delta=0, or regularized least squares if delta≠0 - custom object: Must implement a fit(x, y) method and a coef attribute - + delta : float, optional Ridge penalty parameter for regularization (default=0): - 0: No regularization (standard least squares) - >0: Adds L2 regularization to stabilize parameter estimation - + Attributes ---------- order : int The order of the MVAR model - + coeff : ndarray The estimated MVAR coefficients with shape (n_channels, n_channels * model_order) - + residuals : ndarray The residuals after fitting the model, with same shape as input signal - + Methods ------- fit(signal) Fits the MVAR model to the input signal - + predict(signal) Predicts values using the fitted MVAR model - + stability() Checks whether the MVAR model is stable - + copy() Creates a copy of the MVAR model instance - + Notes ----- Stability of an MVAR model is crucial for interpretability. An unstable model implies that the system would diverge over time, which is typically not physiologically plausible for brain activity. - + The fitting process constructs a system of equations based on the input signal and the specified model order, then solves for the coefficients. - + Examples -------- >>> import numpy as np @@ -81,18 +81,18 @@ class MVAR: >>> predicted = model.predict(data) """ - def __init__(self, model_order, fitting_method='default', delta=0): + def __init__(self, model_order, fitting_method="default", delta=0): """ Initialize an MVAR model with specified parameters. - + Parameters ---------- model_order : int Order of the MVAR model - + fitting_method : str or object, optional Method for fitting the MVAR model (default='default') - + delta : float, optional Ridge penalty parameter (default=0) """ @@ -107,7 +107,7 @@ def __init__(self, model_order, fitting_method='default', delta=0): def copy(self): """ Create a deep copy of the current MVAR model. - + Returns ------- mvar_copy : MVAR @@ -122,20 +122,20 @@ def copy(self): def predict(self, signal): """ Predict time series data using the fitted MVAR model. - + This method applies the fitted MVAR coefficients to predict values based on previous time points in the signal. - + Parameters ---------- signal : ndarray Input signal with shape (n_epochs, n_channels, n_samples) - + Returns ------- predicted : ndarray Predicted signal with the same shape as input - + Notes ----- Predictions start from the (model_order)th time point, as earlier @@ -148,29 +148,31 @@ def predict(self, signal): predicted = np.zeros(signal.shape) if epoch > sample - channel: for i in range(1, p + 1): - bp = self.coeff[:, (i - 1)::p] + bp = self.coeff[:, (i - 1) :: p] for j in range(p, sample): predicted[:, :, j] += np.dot(signal[:, :, j - i], bp.T) else: for i in range(1, p + 1): - bp = self.coeff[:, (i - 1)::p] + bp = self.coeff[:, (i - 1) :: p] for j in range(epoch): - predicted[j, :, p:] += np.dot(bp, signal[j, :, (p - i):(sample - i)]) + predicted[j, :, p:] += np.dot( + bp, signal[j, :, (p - i) : (sample - i)] + ) return predicted def stability(self): """ Check the stability of the fitted MVAR model. - + An MVAR model is stable if all eigenvalues of the coefficient matrix have modulus less than 1. Stability ensures that the model represents a stationary process. - + Returns ------- is_stable : bool True if the model is stable, False otherwise - + Notes ----- Stability is a necessary condition for valid connectivity analysis. @@ -179,7 +181,7 @@ def stability(self): co_0, co_1 = self.coeff.shape p = co_1 // co_0 - assert (co_1 == co_0 * p) + assert co_1 == co_0 * p top_block = [] for i in range(p): top_block.append(self.coeff[:, i::p]) @@ -196,26 +198,26 @@ def stability(self): def construct_equation(self, signal, delta_1=None): """ Construct the system of equations for MVAR model fitting. - + This method reorganizes the input signal into a form suitable for least squares estimation of MVAR coefficients. - + Parameters ---------- signal : ndarray Input signal with shape (n_epochs, n_channels, n_samples) - + delta_1 : float or None, optional Ridge penalty parameter for regularization - + Returns ------- x : ndarray Design matrix containing lagged versions of the signal - + y : ndarray Target matrix containing the values to be predicted - + Notes ----- If delta_1 is provided, regularization terms are added to the design matrix @@ -229,7 +231,9 @@ def construct_equation(self, signal, delta_1=None): x = np.zeros((rows, channel * mvar_order)) for i in range(channel): for j in range(1, mvar_order + 1): - x[:n, i * mvar_order + j - 1] = np.reshape(signal[:, i, mvar_order - j:-j].T, n) + x[:n, i * mvar_order + j - 1] = np.reshape( + signal[:, i, mvar_order - j : -j].T, n + ) if delta_1 is not None: np.fill_diagonal(x[n:, :], delta_1) y = np.zeros((rows, channel)) @@ -240,31 +244,31 @@ def construct_equation(self, signal, delta_1=None): def fit(self, signal): """ Fit the MVAR model to input signal data. - + This method estimates the MVAR coefficients that best predict the signal values based on their past values. - + Parameters ---------- signal : ndarray Input signal with shape (n_epochs, n_channels, n_samples) - + Returns ------- self : MVAR The fitted MVAR model instance - + Notes ----- The fitting method depends on the 'fitting_method' parameter: - If 'default' and delta=0: Standard least squares - If 'default' and delta≠0: Regularized least squares - If custom object: Uses the object's fit method - + After fitting, the model coefficients and residuals are stored as attributes. """ - if self.fit_method.lower() == 'default': + if self.fit_method.lower() == "default": if self.delta == 0 or self.delta is None: x, y = self.construct_equation(signal) else: @@ -285,47 +289,47 @@ def fit(self, signal): return self -def ica_wrapper(ica_input, ica_method='infomax_extended', random_state=None): +def ica_wrapper(ica_input, ica_method="infomax_extended", random_state=None): """ Performs Independent Component Analysis (ICA) on input data. - + This function serves as a unified interface to different ICA algorithms implemented in external packages like MNE and scikit-learn. - + Parameters ---------- ica_input : ndarray Input data matrix with shape (n_samples, n_features) - + ica_method : str, optional ICA algorithm to use (default='infomax_extended'). Options: - 'infomax_extended': Extended Infomax algorithm from MNE - 'infomax': Standard Infomax algorithm from MNE - 'fastica': FastICA algorithm from scikit-learn - + random_state : int or None, optional Seed for random number generator (default=None): - None: Use default random state - int: Set specific random seed for reproducibility - + Returns ------- unmixing_matrix : ndarray Unmixing matrix with shape (n_features, n_features) that transforms the input data into independent components - + Notes ----- The different ICA methods have varying properties: - Extended Infomax can separate both super- and sub-Gaussian sources - Standard Infomax works best for super-Gaussian sources - FastICA is generally faster but may be less stable for certain data types - + Raises ------ ValueError If an unsupported ICA method is specified - + Examples -------- >>> # Apply Extended Infomax ICA to random data @@ -335,61 +339,69 @@ def ica_wrapper(ica_input, ica_method='infomax_extended', random_state=None): >>> components = data @ unmixing.T """ - if ica_method.lower() == 'infomax_extended': + if ica_method.lower() == "infomax_extended": from mne.preprocessing.infomax_ import infomax + return infomax(ica_input, extended=True, random_state=random_state) - elif ica_method.lower() == 'infomax': + elif ica_method.lower() == "infomax": from mne.preprocessing.infomax_ import infomax + return infomax(ica_input, extended=False, random_state=random_state) - elif ica_method.lower() == 'fastica': + elif ica_method.lower() == "fastica": from sklearn.decomposition import FastICA + aux = FastICA(random_state=random_state) aux.fit(ica_input) return aux.components_ else: raise ValueError( - 'This method is not defined!' + '\n' + 'supported methods: infomax, fastica, picard, infomax_extended') + "This method is not defined!" + + "\n" + + "supported methods: infomax, fastica, picard, infomax_extended" + ) -def connectivity_mvarica(real_signal, ica_params, measure_name, n_fft=512, var_model=MVAR): +def connectivity_mvarica( + real_signal, ica_params, measure_name, n_fft=512, var_model=MVAR +): """ Applies MVARICA approach to estimate connectivity between brain sources. - + MVARICA (Multivariate Autoregressive Independent Component Analysis) combines MVAR modeling with ICA to jointly estimate source activities and their causal interactions. This function implements the full pipeline from signal to connectivity measures. - + Parameters ---------- real_signal : ndarray Input signal with shape (n_epochs, n_channels, n_samples) - + ica_params : dict Dictionary of ICA parameters with keys: - 'method': str, ICA algorithm to use (see ica_wrapper for options) - 'random_state': int or None, random seed for reproducibility - + measure_name : str Connectivity measure to compute. Options: - 'mvar_spectral': Spectral representation of VAR coefficients - 'mvar_tf': Transfer function - 'pdc': Partial directed coherence - 'dtf': Directed transfer function - + n_fft : int, optional Number of frequency bins for connectivity computation (default=512) - + var_model : MVAR, optional Pre-initialized MVAR model instance (default=MVAR) - + Returns ------- result : ndarray Connectivity measure matrix with shape dependent on the measure: - For all measures: (n_channels, n_channels, n_fft) Where each [i, j] entry represents connectivity from channel j to channel i - + Notes ----- Process steps: @@ -398,21 +410,21 @@ def connectivity_mvarica(real_signal, ica_params, measure_name, n_fft=512, var_m 3. Apply ICA to the residuals to estimate the mixing matrix 4. Transform the MVAR coefficients to the source space 5. Compute the specified connectivity measure in the frequency domain - + The different connectivity measures have different interpretations: - PDC (Partial Directed Coherence): Measures direct influence from j to i normalized by the total outflow from j - DTF (Directed Transfer Function): Measures total influence from j to i normalized by the total inflow to i - + References ---------- Baccalá, L. A., & Sameshima, K. (2001). Partial directed coherence: a new concept in neural structure determination. Biological cybernetics, 84(6), 463-474. - + Kaminski, M., & Blinowska, K. J. (1991). A new method of the description of the information flow in the brain structures. Biological cybernetics, 65(3), 203-210. - + Examples -------- >>> # Estimate PDC connectivity from simulated data @@ -421,40 +433,53 @@ def connectivity_mvarica(real_signal, ica_params, measure_name, n_fft=512, var_m >>> ica_params = {'method': 'infomax_extended', 'random_state': 42} >>> pdc = connectivity_mvarica(data, ica_params, 'pdc', n_fft=128, var_model=mvar_model) """ - + fit_var = var_model.fit(real_signal) res = real_signal - var_model.predict(real_signal) - unmix_matrix = ica_wrapper(np.concatenate(np.split(res, res.shape[0], 0), axis=2).squeeze(0).T, - ica_method=ica_params["method"], random_state=ica_params["random_state"]).T + unmix_matrix = ica_wrapper( + np.concatenate(np.split(res, res.shape[0], 0), axis=2).squeeze(0).T, + ica_method=ica_params["method"], + random_state=ica_params["random_state"], + ).T mix_matrix = sp.linalg.pinv(unmix_matrix) trns_unmix_matrix = unmix_matrix.T - e = np.concatenate([trns_unmix_matrix.dot(res[i, ...])[np.newaxis, ...] for i in range(res.shape[0])]) + e = np.concatenate( + [ + trns_unmix_matrix.dot(res[i, ...])[np.newaxis, ...] + for i in range(res.shape[0]) + ] + ) fit_var_b = fit_var.copy() for k in range(0, fit_var.order): - fit_var_b.coeff[:, k::fit_var.order] = mix_matrix.dot(fit_var.coeff[:, k::fit_var.order].transpose()).dot( - unmix_matrix).transpose() + fit_var_b.coeff[:, k :: fit_var.order] = ( + mix_matrix.dot(fit_var.coeff[:, k :: fit_var.order].transpose()) + .dot(unmix_matrix) + .transpose() + ) - noise_cov = np.cov(np.concatenate(np.split(e, e.shape[0], 0), axis=2).squeeze(0).T, rowvar=False) + noise_cov = np.cov( + np.concatenate(np.split(e, e.shape[0], 0), axis=2).squeeze(0).T, rowvar=False + ) coeffs = fit_var_b.coeff coeffs = np.asarray(coeffs) coshape_0, coshape_1 = coeffs.shape p = coshape_1 // coshape_0 - assert (coshape_1 == coshape_0 * p) + assert coshape_1 == coshape_0 * p - re_coeffs = np.reshape(coeffs, (coshape_0, coshape_0, p), 'c') + re_coeffs = np.reshape(coeffs, (coshape_0, coshape_0, p), "c") a = fft(np.dstack([np.eye(coshape_0), -re_coeffs]), n_fft * 2 - 1)[:, :, :n_fft] h = np.array([sp.linalg.solve(a, np.eye(a.shape[0])) for a in a.T]).T - if measure_name.lower() == 'mvar_spectral': + if measure_name.lower() == "mvar_spectral": result = a - elif measure_name.lower() == 'mvar_tf': + elif measure_name.lower() == "mvar_tf": result = h - elif measure_name.lower() == 'pdc': + elif measure_name.lower() == "pdc": result = np.abs(a / np.sqrt(np.sum(a.conj() * a, axis=0, keepdims=True))) - elif measure_name.lower() == 'dtf': + elif measure_name.lower() == "dtf": result = np.abs(h / np.sqrt(np.sum(h * h.conj(), axis=1, keepdims=True))) return result diff --git a/hypyp/plots.py b/hypyp/plots.py index 4989f90..5772ad3 100644 --- a/hypyp/plots.py +++ b/hypyp/plots.py @@ -7,19 +7,21 @@ from mne_connectivity.viz import plot_connectivity_circle from mne.viz import circular_layout + # Define a custom locator and formatter for periods def custom_locator_yticks(ymin, ymax): ticks = [] if ymin < 1: - values_01 = np.arange(math.ceil(ymin*10), 10) / 10 + values_01 = np.arange(math.ceil(ymin * 10), 10) / 10 ticks.extend(values_01) ticks.extend(range(math.ceil(ymin), 11)) ticks.extend(range(12, 21, 2)) ticks.extend(range(25, 40 + 1, 5)) ticks.extend(range(50, int(ymax) + 1, 10)) return ticks - + + def plot_wavelet_transform_weights( W, times, @@ -35,14 +37,14 @@ def plot_wavelet_transform_weights( show_cone_of_influence=True, show_nyquist=True, show_bins=True, - cmap='jet', + cmap="jet", ): # create the figure if needed if ax is None: fig, ax = plt.subplots() else: fig = ax.get_figure() - + nyquist = np.ones((len(times),)) * (sfreq / 2) if use_periods: @@ -54,7 +56,7 @@ def plot_wavelet_transform_weights( nyquist_y2 = np.min(periods) y_lims = [periods.max(), periods.min()] y_cuts = 1 / np.array(frequency_cuts) - y_label = 'Period (s)' + y_label = "Period (s)" else: y_items = freqs coi_y = coif @@ -63,41 +65,47 @@ def plot_wavelet_transform_weights( nyquist_y2 = np.max(freqs) y_lims = [freqs.min(), freqs.max()] y_cuts = frequency_cuts - y_label = 'Frequency (Hz)' + y_label = "Frequency (Hz)" xx, yy = np.meshgrid(times, y_items) - - #im = ax.pcolor(xx, yy, W, vmin=0, vmax=1) + + # im = ax.pcolor(xx, yy, W, vmin=0, vmax=1) im = ax.pcolor(xx, yy, np.abs(W), cmap=cmap) - ax.set_yscale('log') - ax.set_xlabel('Time (s)') + ax.set_yscale("log") + ax.set_xlabel("Time (s)") ax.set_ylabel(y_label) - color_invalid = 'C0' + color_invalid = "C0" # Cone of influence if show_cone_of_influence: ax.plot(times, coi_y, color=color_invalid) - ax.fill_between(times, coi_y, y2=coi_y2, step="mid", color=color_invalid, alpha=0.4) + ax.fill_between( + times, coi_y, y2=coi_y2, step="mid", color=color_invalid, alpha=0.4 + ) if show_nyquist: ax.plot(times, nyquist_y, color=color_invalid) - ax.fill_between(times, nyquist_y, y2=nyquist_y2, step="mid", color=color_invalid, alpha=0.4) - + ax.fill_between( + times, nyquist_y, y2=nyquist_y2, step="mid", color=color_invalid, alpha=0.4 + ) + if show_bins: if bin_seconds is not None: for time_cut in np.arange(0, max(times), bin_seconds): - plt.axvline(x=time_cut, color='red', lw=0.5) + plt.axvline(x=time_cut, color="red", lw=0.5) if y_cuts is not None: for y_cut in y_cuts: - plt.axhline(y=y_cut, color='red', lw=0.5) - + plt.axhline(y=y_cut, color="red", lw=0.5) + # Dynamically set ticks based on the current range ymin, ymax = ax.get_ylim() # Get the y-axis limits ax.set_yticks(custom_locator_yticks(ymin, ymax)) - ax.yaxis.set_major_formatter(FuncFormatter(lambda y,_: f"{int(y)}" if y >= 1 else f"{y:.1f}")) - #ax.yaxis.get_major_formatter().set_scientific(False) # Disable scientific notation + ax.yaxis.set_major_formatter( + FuncFormatter(lambda y, _: f"{int(y)}" if y >= 1 else f"{y:.1f}") + ) + # ax.yaxis.get_major_formatter().set_scientific(False) # Disable scientific notation ax.set_xlim(times.min(), times.max()) ax.set_ylim(*y_lims) @@ -108,11 +116,11 @@ def plot_wavelet_transform_weights( if title is not None: ax.set_title(title) else: - ax.set_title('CWT Weights') - + ax.set_title("CWT Weights") return fig + def subplot_heatmap_from_pivot(pivot, ordered_fields, ax): index_order = [i for i in ordered_fields if i in pivot.index] # Append indices not in ordered_fields @@ -125,86 +133,139 @@ def subplot_heatmap_from_pivot(pivot, ordered_fields, ax): column_order.append(c) pivot_reordered = pivot.reindex(index=index_order, columns=column_order) - heatmap = sns.heatmap(pivot_reordered, cmap='viridis', vmin=0, vmax=1, cbar=False, ax=ax) + heatmap = sns.heatmap( + pivot_reordered, cmap="viridis", vmin=0, vmax=1, cbar=False, ax=ax + ) ax.set_xticks(ticks=range(len(pivot_reordered.columns))) - ax.set_xticklabels(pivot_reordered.columns, rotation=90, ha='left', fontsize=6 if len(column_order)>15 else 10) + ax.set_xticklabels( + pivot_reordered.columns, + rotation=90, + ha="left", + fontsize=6 if len(column_order) > 15 else 10, + ) ax.set_yticks(ticks=range(len(pivot_reordered.index))) - ax.set_yticklabels(pivot_reordered.index, rotation=0, va='top', fontsize=6 if len(index_order)>15 else 10) - ax.tick_params(axis='both', which='both', length=0) + ax.set_yticklabels( + pivot_reordered.index, + rotation=0, + va="top", + fontsize=6 if len(index_order) > 15 else 10, + ) + ax.tick_params(axis="both", which="both", length=0) return heatmap + def plot_coherence_matrix( df, s1_label, s2_label, - field1, # roi1 or channel1 - field2, # roi2 or channel2 + field1, # roi1 or channel1 + field2, # roi2 or channel2 ordered_fields, ): # We don't sharex and sharey because the list of channels might be different in the 2 subjects - dyad_selector = (df['is_intra']==False) - s1_selector = (df['is_intra']==True) & (df['is_intra_of']==1) & (df['channel1']!=df['channel2']) - s2_selector = (df['is_intra']==True) & (df['is_intra_of']==2) & (df['channel1']!=df['channel2']) + dyad_selector = df["is_intra"] == False + s1_selector = ( + (df["is_intra"] == True) + & (df["is_intra_of"] == 1) + & (df["channel1"] != df["channel2"]) + ) + s2_selector = ( + (df["is_intra"] == True) + & (df["is_intra_of"] == 2) + & (df["channel1"] != df["channel2"]) + ) df_dyad = df[dyad_selector] df_s1 = df[s1_selector] df_s2 = df[s2_selector] - pivot_s1 = df_s1.pivot_table(index=field1, columns=field2, values='coherence', aggfunc='mean', observed=False) - pivot_s2 = df_s2.pivot_table(index=field2, columns=field1, values='coherence', aggfunc='mean', observed=False) - pivot_dyad = df_dyad.pivot_table(index=field1, columns=field2, values='coherence', aggfunc='mean', observed=False) - - if np.all(df['is_intra']): + pivot_s1 = df_s1.pivot_table( + index=field1, columns=field2, values="coherence", aggfunc="mean", observed=False + ) + pivot_s2 = df_s2.pivot_table( + index=field2, columns=field1, values="coherence", aggfunc="mean", observed=False + ) + pivot_dyad = df_dyad.pivot_table( + index=field1, columns=field2, values="coherence", aggfunc="mean", observed=False + ) + + if np.all(df["is_intra"]): fig, ax = plt.subplots(1, 1, figsize=(8, 8), sharex=False, sharey=False) - subplot_heatmap_from_pivot(pivot_s1.rename_axis(index=s1_label, columns=s1_label), ordered_fields=ordered_fields, ax=ax) - elif not np.any(df['is_intra']): + subplot_heatmap_from_pivot( + pivot_s1.rename_axis(index=s1_label, columns=s1_label), + ordered_fields=ordered_fields, + ax=ax, + ) + elif not np.any(df["is_intra"]): fig, ax = plt.subplots(1, 1, figsize=(8, 8), sharex=False, sharey=False) - subplot_heatmap_from_pivot(pivot_dyad.rename_axis(index=s1_label, columns=s2_label), ordered_fields=ordered_fields, ax=ax) + subplot_heatmap_from_pivot( + pivot_dyad.rename_axis(index=s1_label, columns=s2_label), + ordered_fields=ordered_fields, + ax=ax, + ) else: fig, axes = plt.subplots(2, 2, figsize=(8, 8), sharex=False, sharey=False) - subplot_heatmap_from_pivot(pivot_s1.rename_axis(index=s1_label, columns=s1_label), ordered_fields=ordered_fields, ax=axes[0,0]) # top left - subplot_heatmap_from_pivot(pivot_dyad.rename_axis(index=s1_label, columns=s2_label), ordered_fields=ordered_fields, ax=axes[0,1]) # top right - subplot_heatmap_from_pivot(pivot_dyad.T.rename_axis(index=s2_label, columns=s1_label), ordered_fields=ordered_fields, ax=axes[1,0]) # bottom left - subplot_heatmap_from_pivot(pivot_s2.rename_axis(index=s2_label, columns=s2_label), ordered_fields=ordered_fields, ax=axes[1,1]) # bottom right - - #fig.subplots_adjust(wspace=0.1, hspace=0.1) + subplot_heatmap_from_pivot( + pivot_s1.rename_axis(index=s1_label, columns=s1_label), + ordered_fields=ordered_fields, + ax=axes[0, 0], + ) # top left + subplot_heatmap_from_pivot( + pivot_dyad.rename_axis(index=s1_label, columns=s2_label), + ordered_fields=ordered_fields, + ax=axes[0, 1], + ) # top right + subplot_heatmap_from_pivot( + pivot_dyad.T.rename_axis(index=s2_label, columns=s1_label), + ordered_fields=ordered_fields, + ax=axes[1, 0], + ) # bottom left + subplot_heatmap_from_pivot( + pivot_s2.rename_axis(index=s2_label, columns=s2_label), + ordered_fields=ordered_fields, + ax=axes[1, 1], + ) # bottom right + + # fig.subplots_adjust(wspace=0.1, hspace=0.1) plt.tight_layout() return fig - -def plot_coherence_connectogram(df_pivot, title='', ax=None): + +def plot_coherence_connectogram(df_pivot, title="", ax=None): if ax is None: - fig, ax = plt.subplots(1, 1, subplot_kw={'projection': 'polar'}) + fig, ax = plt.subplots(1, 1, subplot_kw={"projection": "polar"}) else: fig = ax.get_figure() - - #node_angles = circular_layout( + + # node_angles = circular_layout( # df.columns, list(df.columns), start_pos=90, group_boundaries=[0, len(df.columns) // 2] - #) - plot_connectivity_circle(df_pivot.values, + # ) + plot_connectivity_circle( + df_pivot.values, df_pivot.columns, - #node_angles=node_angles, + # node_angles=node_angles, title=title, - #vmin=0, - #vmax=1, - colormap='Greys', + # vmin=0, + # vmax=1, + colormap="Greys", ax=ax, - facecolor='white', - textcolor='black', - node_edgecolor='black', - node_colors=['white'], + facecolor="white", + textcolor="black", + node_edgecolor="black", + node_colors=["white"], ) return fig - -def plot_coherence_connectogram_split(df_pivot, title='', ax=None): + + +def plot_coherence_connectogram_split(df_pivot, title="", ax=None): if ax is None: - fig, ax = plt.subplots(1, 1, subplot_kw={'projection': 'polar'}) + fig, ax = plt.subplots(1, 1, subplot_kw={"projection": "polar"}) else: fig = ax.get_figure() - + values = df_pivot.values.flatten() s1_roi_list = df_pivot.index.to_list() s2_roi_list = df_pivot.columns.to_list() @@ -216,7 +277,7 @@ def plot_coherence_connectogram_split(df_pivot, title='', ax=None): for i in range(len(s1_roi_list)): for j in range(len(s2_roi_list)): s1_roi_idx.append(i) - s2_roi_idx.append(j+len(s1_roi_list)) + s2_roi_idx.append(j + len(s1_roi_list)) indices = ( np.array(s1_roi_idx), @@ -227,44 +288,52 @@ def plot_coherence_connectogram_split(df_pivot, title='', ax=None): for i in range(len(s1_roi_list)): node_names_ordered.append(s1_roi_list[i]) for i in range(len(s2_roi_list)): - node_names_ordered.append(s2_roi_list[len(s2_roi_list)-i-1]) + node_names_ordered.append(s2_roi_list[len(s2_roi_list) - i - 1]) node_angles = circular_layout( - node_names, node_names_ordered, start_pos=90, group_boundaries=[0, len(s1_roi_list)] + node_names, + node_names_ordered, + start_pos=90, + group_boundaries=[0, len(s1_roi_list)], ) plot_connectivity_circle( values, node_names, - indices = indices, + indices=indices, node_angles=node_angles, title=title, - #vmin=0, - #vmax=1, - colormap='Greys', + # vmin=0, + # vmax=1, + colormap="Greys", ax=ax, - facecolor='white', - textcolor='black', - node_edgecolor='black', - node_colors=['white'], + facecolor="white", + textcolor="black", + node_edgecolor="black", + node_colors=["white"], ) return fig - + def plot_coherence_bars_per_task(df): - selector = (df['roi1'] == df['roi2']) + selector = df["roi1"] == df["roi2"] filtered_df = df[selector] # remove intra-subject "same channel" coherence to avoid counting them in means - intra_same_ch_selector = (filtered_df['is_intra']==True) & (filtered_df['channel1']==filtered_df['channel2']) + intra_same_ch_selector = (filtered_df["is_intra"] == True) & ( + filtered_df["channel1"] == filtered_df["channel2"] + ) filtered_df[intra_same_ch_selector] = np.nan p = sns.catplot( - data=filtered_df, kind="bar", - x="roi1", y="coherence", hue="task", + data=filtered_df, + kind="bar", + x="roi1", + y="coherence", + hue="task", col="is_intra", - #palette="dark", alpha=.6, height=6 + # palette="dark", alpha=.6, height=6 ) p.despine(left=True) p.set_axis_labels("", "Coherence") @@ -273,9 +342,7 @@ def plot_coherence_bars_per_task(df): p.legend.set_title("Task") - p.set_titles('Is intra: {col_name}') + p.set_titles("Is intra: {col_name}") plt.subplots_adjust(bottom=0.5) return p - - diff --git a/hypyp/prep.py b/hypyp/prep.py index 3e7f051..abd7d8f 100644 --- a/hypyp/prep.py +++ b/hypyp/prep.py @@ -11,7 +11,6 @@ | date | 2020-03-18 | """ - import numpy as np import copy import matplotlib.pyplot as plt @@ -20,28 +19,29 @@ from mne.preprocessing import ICA, corrmap from typing import List, Tuple, TypedDict, Union + class DicAR(TypedDict): """ Dictionary type for storing epoch rejection information. - + This type provides structured information about the rejection process, including the rejection strategy, threshold, and percentage of epochs rejected for each participant and the dyad. - + Attributes ---------- strategy : str The strategy used for epoch rejection ('union' or 'intersection') - + threshold : float The maximum allowed percentage of epochs to be rejected - + S1 : float Percentage of epochs rejected for participant 1 - + S2 : float Percentage of epochs rejected for participant 2 - + dyad : float Overall percentage of epochs rejected across both participants """ @@ -53,46 +53,48 @@ class DicAR(TypedDict): dyad: float -def filt(raw_S: List[mne.io.Raw], - freqs: Tuple[Union[float, None], Union[float, None]] = (2., None)) -> List[mne.io.Raw]: +def filt( + raw_S: List[mne.io.Raw], + freqs: Tuple[Union[float, None], Union[float, None]] = (2.0, None), +) -> List[mne.io.Raw]: """ Filter a list of raw EEG data to remove slow drifts or other unwanted frequency components. - + This function applies a high-pass or band-pass filter to each Raw object in the input list. Filtering helps to remove low-frequency drifts, power line noise, or other frequency-specific artifacts from the EEG data. - + Parameters ---------- raw_S : List[mne.io.Raw] List of Raw objects containing continuous EEG data - + freqs : Tuple[Union[float, None], Union[float, None]], optional Frequency range for filtering (default=(2., None)): - First element: Lower frequency bound (high-pass filter cutoff) - Second element: Upper frequency bound (low-pass filter cutoff) - None for either bound means no filtering in that direction - + Returns ------- raws : List[mne.io.Raw] List of filtered Raw objects with the same structure as the input - + Notes ----- By default, a 2 Hz high-pass filter is applied, which effectively removes slow drifts while preserving most of the EEG signal content. This is appropriate for most EEG analyses but may need adjustment for specific paradigms. - + Examples -------- >>> # Filter data with a high-pass at 1 Hz >>> filtered_raws = filt(raw_list, freqs=(1., None)) - + >>> # Apply a band-pass filter between 1 and 40 Hz >>> filtered_raws = filt(raw_list, freqs=(1., 40.)) """ - + raws = [mne.io.Raw.filter(raw, l_freq=freqs[0], h_freq=freqs[1]) for raw in raw_S] return raws @@ -100,24 +102,24 @@ def filt(raw_S: List[mne.io.Raw], def ICA_choice_comp(icas: List[ICA], epochs: List[mne.Epochs]) -> List[mne.Epochs]: """ Select ICA components for artifact rejection and apply ICA cleaning to epochs. - + This interactive function plots the Independent Components for each participant, lets the user choose relevant components for artifact rejection, and applies the cleaning to the epochs data. - + Parameters ---------- icas : List[ICA] List of fitted ICA objects (one for each participant) - + epochs : List[mne.Epochs] List of Epochs objects to clean (one for each participant) - + Returns ------- cleaned_epochs_ICA : List[mne.Epochs] List of ICA-cleaned Epochs objects (one for each participant) - + Notes ----- The function uses an interactive approach: @@ -125,10 +127,10 @@ def ICA_choice_comp(icas: List[ICA], epochs: List[mne.Epochs]) -> List[mne.Epoch 2. It prompts the user to select a participant and component to use as a template 3. It uses corrmap to find similar components across participants 4. It removes the identified components from all participants' data - + If you don't want to apply ICA cleaning, simply press Enter without typing anything when prompted for participant and component selection. - + Examples -------- >>> # This function is interactive and prompts the user for input @@ -141,62 +143,71 @@ def ICA_choice_comp(icas: List[ICA], epochs: List[mne.Epochs]) -> List[mne.Epoch # choosing participant and its component as a template for the other participant # if do not want to apply ICA on the data, do not fill the answer - subject_id = input("Which participant ICA do you want" - " to use as a template for artifact rejection?" - " Index begins at zero. (If you do not want to apply" - " ICA on your data, do not enter nothing and press enter.)") - - component_id = input("Which IC do you want to use as a template?" - " Index begins at zero. (If you did not choose" - " a participant number at first question," - " then do not enter nothing and press enter again" - " to not apply ICA on your data)") - - if (len(subject_id) == 0 or len(component_id) == 0): + subject_id = input( + "Which participant ICA do you want" + " to use as a template for artifact rejection?" + " Index begins at zero. (If you do not want to apply" + " ICA on your data, do not enter nothing and press enter.)" + ) + + component_id = input( + "Which IC do you want to use as a template?" + " Index begins at zero. (If you did not choose" + " a participant number at first question," + " then do not enter nothing and press enter again" + " to not apply ICA on your data)" + ) + + if len(subject_id) == 0 or len(component_id) == 0: return epochs return ICA_apply(icas, int(subject_id), int(component_id), epochs) -def ICA_apply(icas: List[ICA], subject_id: int, component_id: int, - epochs: List[mne.Epochs], plot: bool = True) -> List[mne.Epochs]: +def ICA_apply( + icas: List[ICA], + subject_id: int, + component_id: int, + epochs: List[mne.Epochs], + plot: bool = True, +) -> List[mne.Epochs]: """ Apply ICA artifact rejection using a template component. - + This function uses a component from one participant as a template to identify similar artifact components in all participants, then removes these components from the data. - + Parameters ---------- icas : List[ICA] List of fitted ICA objects (one for each participant) - + subject_id : int Index of the participant whose component will serve as the template - + component_id : int Index of the component to use as the template - + epochs : List[mne.Epochs] List of Epochs objects to clean (one for each participant) - + plot : bool, optional Whether to plot the identified components (default=True) - + Returns ------- cleaned_epochs_ICA : List[mne.Epochs] List of ICA-cleaned Epochs objects (one for each participant) - + Notes ----- This function uses MNE's corrmap function to identify components similar to the template component. It then labels these components as 'blink' artifacts and removes them from the data. - + The threshold for component similarity is set to 0.9 by default. - + Examples -------- >>> # Use the first component of the first participant as template @@ -209,11 +220,12 @@ def ICA_apply(icas: List[ICA], subject_id: int, component_id: int, template_eog_component = icas[subject_id].get_components()[:, component_id] # applying corrmap with at least 1 component detected for each subj - corrmap(icas, + corrmap( + icas, template=template_eog_component, threshold=0.9, - label='blink', - ch_type='eeg', + label="blink", + ch_type="eeg", plot=plot, ) @@ -222,77 +234,82 @@ def ICA_apply(icas: List[ICA], subject_id: int, component_id: int, # selecting ICA components after viz for ica in icas: - ica.exclude = ica.labels_['blink'] + ica.exclude = ica.labels_["blink"] # applying ica on clean_epochs # for each participant for subject_id, ica in zip(range(0, len(epochs)), icas): # taking all channels to apply ICA epochs_subj = mne.Epochs.copy(epochs[subject_id]) - bads_keep = epochs_subj.info['bads'] - epochs_subj.info['bads'] = [] + bads_keep = epochs_subj.info["bads"] + epochs_subj.info["bads"] = [] ica.apply(epochs_subj) - epochs_subj.info['bads'] = bads_keep + epochs_subj.info["bads"] = bads_keep cleaned_epochs_ICA.append(epochs_subj) return cleaned_epochs_ICA -def ICA_fit(epochs: List[mne.Epochs], n_components: int, method: str, - fit_params: dict, random_state: int) -> List[ICA]: +def ICA_fit( + epochs: List[mne.Epochs], + n_components: int, + method: str, + fit_params: dict, + random_state: int, +) -> List[ICA]: """ Compute Independent Component Analysis (ICA) on epochs for artifact rejection. - + This function applies global Autoreject to establish rejection thresholds, then fits ICA on the cleaned data. ICA is a commonly used technique to identify and remove artifacts such as eye blinks, muscle activity, and cardiac artifacts. - + Parameters ---------- epochs : List[mne.Epochs] List of Epochs objects (one for each participant) - + n_components : int Number of principal components to pass to the ICA algorithm For a first estimation, a value around 15 is often appropriate - + method : str ICA method to use. Options: - 'fastica': FastICA algorithm (most commonly used) - 'infomax': Infomax algorithm - 'picard': Picard algorithm - + fit_params : dict Additional parameters passed to the ICA estimator For Extended Infomax, use method='infomax' and fit_params=dict(extended=True) - + random_state : int Random seed for reproducible results For 15 components, random_state=97 works well For 20 components, random_state=0 works well - + Returns ------- icas : List[ICA] List of fitted ICA objects (one for each participant) - + Notes ----- Pre-requisites: - Install autoreject: https://api.github.com/repos/autoreject/autoreject/zipball/master - Filter the Epochs between 2 and 30 Hz before ICA fitting - + If Autoreject and ICA take too much time, try changing the 'decim' value in the ICA initialization to downsample the data during fitting. - + Examples -------- >>> # Fit ICA with 15 components using FastICA - >>> icas = ICA_fit(epochs_list, n_components=15, method='fastica', + >>> icas = ICA_fit(epochs_list, n_components=15, method='fastica', ... fit_params=None, random_state=97) - + >>> # Fit ICA with Extended Infomax - >>> icas = ICA_fit(epochs_list, n_components=15, method='infomax', + >>> icas = ICA_fit(epochs_list, n_components=15, method='infomax', ... fit_params=dict(extended=True), random_state=97) """ @@ -300,57 +317,63 @@ def ICA_fit(epochs: List[mne.Epochs], n_components: int, method: str, for epochs_subj in epochs: # per subj # applying AR to find global rejection threshold - reject = get_rejection_threshold(epochs_subj, ch_types='eeg') + reject = get_rejection_threshold(epochs_subj, ch_types="eeg") # if very long, can change decim value print(f"The rejection dictionary is {reject}") # fitting ICA on filt_raw after AR - ica = ICA(n_components=n_components, - method=method, - fit_params= fit_params, - random_state=random_state) + ica = ICA( + n_components=n_components, + method=method, + fit_params=fit_params, + random_state=random_state, + ) # take bad channels into account in ICA fit epochs_fit = mne.Epochs.copy(epochs_subj) - epochs_fit.info['bads'] = [] + epochs_fit.info["bads"] = [] epochs_fit.drop_bad(reject=reject, flat=None) icas.append(ica.fit(epochs_fit)) return icas -def AR_local(cleaned_epochs_ICA: List[mne.Epochs], strategy: str = 'union', - threshold: float = 50.0, verbose: bool = False) -> Tuple[mne.Epochs, DicAR]: +def AR_local( + cleaned_epochs_ICA: List[mne.Epochs], + strategy: str = "union", + threshold: float = 50.0, + verbose: bool = False, +) -> Tuple[mne.Epochs, DicAR]: """ Apply local Autoreject to repair or reject bad epochs. - + After ICA cleaning, this function identifies remaining problematic epochs and either repairs them (by interpolating bad channels) or rejects them entirely. It can use different strategies to handle epochs that are bad in only one participant. - + Parameters ---------- cleaned_epochs_ICA : List[mne.Epochs] List of Epochs objects after ICA cleaning (one for each participant) - + strategy : str, optional Strategy for handling bad epochs (default='union'): - 'union': Reject epochs that are bad in either participant - 'intersection': Reject only epochs that are bad in both participants, attempt to repair other bad epochs individually - + threshold : float, optional Maximum acceptable percentage of rejected epochs (default=50.0) If more epochs would be rejected, raises an error - + verbose : bool, optional Whether to plot before/after comparisons (default=False) - + Returns ------- cleaned_epochs_AR : List[mne.Epochs] List of Epochs objects after Autoreject cleaning - + dic_AR : DicAR Dictionary with information about the rejection process: - 'strategy': Strategy used ('union' or 'intersection') @@ -358,26 +381,26 @@ def AR_local(cleaned_epochs_ICA: List[mne.Epochs], strategy: str = 'union', - 'S1': Percentage of epochs rejected for participant 1 - 'S2': Percentage of epochs rejected for participant 2 - 'dyad': Overall percentage of epochs rejected - + Notes ----- This function uses the autoreject package to automatically identify and handle problematic epochs. The identification is based on statistical properties of the data and is more objective than manual rejection. - + With the 'intersection' strategy, the function also equalizes the number of epochs between participants after cleaning. - + Raises ------ RuntimeError If the percentage of rejected epochs exceeds the specified threshold - + Examples -------- >>> # Apply Autoreject with default settings >>> cleaned_epochs, rejection_info = AR_local(ica_cleaned_epochs) - + >>> # Use 'intersection' strategy with lower threshold >>> cleaned_epochs, rejection_info = AR_local( ... ica_cleaned_epochs, strategy='intersection', threshold=30.0, verbose=True @@ -388,8 +411,8 @@ def AR_local(cleaned_epochs_ICA: List[mne.Epochs], strategy: str = 'union', reject_logs: List[RejectLog] = [] AR: List[AutoReject] = [] dic_AR: DicAR = {} - dic_AR['strategy'] = strategy - dic_AR['threshold'] = threshold + dic_AR["strategy"] = strategy + dic_AR["threshold"] = threshold # defaults values for n_interpolates and consensus_percs n_interpolates = np.array([1, 4, 32]) @@ -406,14 +429,17 @@ def AR_local(cleaned_epochs_ICA: List[mne.Epochs], strategy: str = 'union', eeg=True, stim=False, eog=False, - exclude=[]) + exclude=[], + ) - ar = AutoReject(n_interpolates, + ar = AutoReject( + n_interpolates, consensus_percs, picks=picks, - thresh_method='random_search', + thresh_method="random_search", random_state=42, - verbose='tqdm_notebook') + verbose="tqdm_notebook", + ) # fitting AR to get bad epochs ar.fit(clean_epochs_subj) @@ -426,16 +452,16 @@ def AR_local(cleaned_epochs_ICA: List[mne.Epochs], strategy: str = 'union', bad1_idx = np.where(reject_logs[0].bad_epochs == True)[0].tolist() bad2_idx = np.where(reject_logs[1].bad_epochs == True)[0].tolist() - if strategy == 'union': + if strategy == "union": bad_idx = list(set(bad1_idx).union(set(bad2_idx))) - elif strategy == 'intersection': + elif strategy == "intersection": bad_idx = list(set(bad1_idx).intersection(set(bad2_idx))) else: - raise RuntimeError('not good strategy input!') + raise RuntimeError("not good strategy input!") # storing the percentage of epochs rejection - dic_AR['S1'] = float((len(bad1_idx) / len(cleaned_epochs_ICA[0])) * 100) - dic_AR['S2'] = float((len(bad2_idx) / len(cleaned_epochs_ICA[1])) * 100) + dic_AR["S1"] = float((len(bad1_idx) / len(cleaned_epochs_ICA[0])) * 100) + dic_AR["S2"] = float((len(bad2_idx) / len(cleaned_epochs_ICA[1])) * 100) # picking good epochs for the two subj cleaned_epochs_AR: List[mne.Epochs] = [] @@ -449,40 +475,48 @@ def AR_local(cleaned_epochs_ICA: List[mne.Epochs], strategy: str = 'union', epochs_AR = ar.transform(epochs_subj) cleaned_epochs_AR.append(epochs_AR) - if strategy == 'intersection': + if strategy == "intersection": # equalizing epochs length between two participants mne.epochs.equalize_epoch_counts(cleaned_epochs_AR) n_epochs_ICA = len(cleaned_epochs_ICA[0]) n_epochs_AR = len(cleaned_epochs_AR[0]) - dic_AR['dyad'] = float(((n_epochs_ICA - n_epochs_AR) / n_epochs_ICA) * 100) - if dic_AR['dyad'] >= threshold: - raise RuntimeError(f"percentage of rejected epochs ({dic_AR['dyad']}) above threshold ({threshold})! ") + dic_AR["dyad"] = float(((n_epochs_ICA - n_epochs_AR) / n_epochs_ICA) * 100) + if dic_AR["dyad"] >= threshold: + raise RuntimeError( + f"percentage of rejected epochs ({dic_AR['dyad']}) above threshold ({threshold})! " + ) if verbose: print(f"{dic_AR['dyad']} percent of bad epochs") # Vizualisation before after AR - evoked_before: List[mne.Evoked] = [epochs.average() for epochs in cleaned_epochs_ICA] - evoked_after_AR: List[mne.Evoked] = [epochs.average() for epochs in cleaned_epochs_AR] + evoked_before: List[mne.Evoked] = [ + epochs.average() for epochs in cleaned_epochs_ICA + ] + evoked_after_AR: List[mne.Evoked] = [ + epochs.average() for epochs in cleaned_epochs_AR + ] if verbose: - for evoked_before_subj, evoked_after_AR_subj in zip(evoked_before, evoked_after_AR): + for evoked_before_subj, evoked_after_AR_subj in zip( + evoked_before, evoked_after_AR + ): fig, axes = plt.subplots(2, 1, figsize=(6, 6)) for ax in axes: - ax.tick_params(axis='x', which='both', bottom='off', top='off') - ax.tick_params(axis='y', which='both', left='off', right='off') + ax.tick_params(axis="x", which="both", bottom="off", top="off") + ax.tick_params(axis="y", which="both", left="off", right="off") ylim = dict(grad=(-170, 200)) evoked_before_subj.pick_types(eeg=True, exclude=[]) evoked_before_subj.plot(exclude=[], axes=axes[0], ylim=ylim, show=False) - axes[0].set_title('Before autoreject') + axes[0].set_title("Before autoreject") evoked_after_AR_subj.pick_types(eeg=True, exclude=[]) evoked_after_AR_subj.plot(exclude=[], axes=axes[1], ylim=ylim) # Problème titre ne s'affiche pas pour le deuxieme axe !!! - axes[1].set_title('After autoreject') + axes[1].set_title("After autoreject") plt.tight_layout() diff --git a/hypyp/profiling.py b/hypyp/profiling.py index e91b027..8f4297f 100644 --- a/hypyp/profiling.py +++ b/hypyp/profiling.py @@ -1,6 +1,7 @@ import tracemalloc import time + class MemoryMonitor(object): def __init__(self): """ @@ -13,7 +14,7 @@ def __init__(self): """ pass - + def format_memory_size(self, size_bytes): if size_bytes >= 1 << 30: # Greater than or equal to 1 GiB return f"{size_bytes / (1 << 30):.2f} G" @@ -29,11 +30,14 @@ def __enter__(self): def __exit__(self, *args): res = tracemalloc.get_traced_memory() - print(f"[MemoryMonitor] allocated: {self.format_memory_size(res[0])}, peak: {self.format_memory_size(res[1])}", ) + print( + f"[MemoryMonitor] allocated: {self.format_memory_size(res[0])}, peak: {self.format_memory_size(res[1])}", + ) tracemalloc.stop() + class TimeTracker(object): - def __init__(self, label:str='time_tracker'): + def __init__(self, label: str = "time_tracker"): """ Track code execution time. Usage: @@ -77,4 +81,4 @@ def __enter__(self): def __exit__(self, *args): self.stop() - print(f"--- [{self.label}] {self.duration} seconds ---") \ No newline at end of file + print(f"--- [{self.label}] {self.duration} seconds ---") diff --git a/hypyp/shiny/study_viewer.py b/hypyp/shiny/study_viewer.py index f0d3088..4a4ac99 100644 --- a/hypyp/shiny/study_viewer.py +++ b/hypyp/shiny/study_viewer.py @@ -9,7 +9,7 @@ import matplotlib import pandas as pd -root = os.path.join(Path(__file__).parent, '..', '..') +root = os.path.join(Path(__file__).parent, "..", "..") sys.path.append(root) import hypyp.plots @@ -18,23 +18,25 @@ HARDCODED_RESULTS_PATH = "./data/results" -SIDEBAR_WIDTH = 400 # px +SIDEBAR_WIDTH = 400 # px DEFAULT_PLOT_COHERENCE_MATRIX_HEIGHT = 1000 DEFAULT_PLOT_COHERENCE_PER_TASK_HEIGHT = 600 DEFAULT_PLOT_WTC_HEIGHT = 600 DEFAULT_PLOT_CONNECTOGRAM_HEIGHT = 1000 -STR_ALL_DYADS = 'All dyads' -STR_ALL_TASKS = 'All tasks' +STR_ALL_DYADS = "All dyads" +STR_ALL_TASKS = "All tasks" # This is to avoid having external windows launched -matplotlib.use('Agg') +matplotlib.use("Agg") + def mne_figure_as_image(fig): temp_img_path = tempfile.NamedTemporaryFile(suffix=".png").name fig.savefig(temp_img_path) return {"src": temp_img_path, "alt": "MNE Plot"} - + + def ui_option_row(label, ui_element, sizes=(6, 6), center=False): label_style = "" if center: @@ -44,37 +46,37 @@ def ui_option_row(label, ui_element, sizes=(6, 6), center=False): ui.column(sizes[1], ui_element), ) + # UI app_ui = ui.page_fluid( ui.page_navbar( ui.nav_spacer(), ui.nav_panel( "Study Info", - ui.row( - ui.column( - 12, - ui.output_table('table_info_study') - ) - ), + ui.row(ui.column(12, ui.output_table("table_info_study"))), ), ui.nav_panel( "Coherence Matrix", ui.row( ui.column( 10, - ui.output_plot('plot_coherence_matrix', height=DEFAULT_PLOT_COHERENCE_MATRIX_HEIGHT) + ui.output_plot( + "plot_coherence_matrix", + height=DEFAULT_PLOT_COHERENCE_MATRIX_HEIGHT, + ), ), ui.column( 2, ui.input_select( - 'coherence_select_grouping', - 'Coherence grouping', + "coherence_select_grouping", + "Coherence grouping", choices={ - 'roi': 'Region of Interest', - 'channel': 'Individual Channels', - 'channel_roi': 'Channel-ROI', - 'roi_channel': 'ROI-Channel', - }) + "roi": "Region of Interest", + "channel": "Individual Channels", + "channel_roi": "Channel-ROI", + "roi_channel": "ROI-Channel", + }, + ), ), ), ), @@ -83,15 +85,21 @@ def ui_option_row(label, ui_element, sizes=(6, 6), center=False): ui.row( ui.column( 4, - ui.output_plot('plot_connectogram_s1', height=DEFAULT_PLOT_CONNECTOGRAM_HEIGHT) + ui.output_plot( + "plot_connectogram_s1", height=DEFAULT_PLOT_CONNECTOGRAM_HEIGHT + ), ), ui.column( 4, - ui.output_plot('plot_connectogram_s2', height=DEFAULT_PLOT_CONNECTOGRAM_HEIGHT) + ui.output_plot( + "plot_connectogram_s2", height=DEFAULT_PLOT_CONNECTOGRAM_HEIGHT + ), ), ui.column( 4, - ui.output_plot('plot_connectogram', height=DEFAULT_PLOT_CONNECTOGRAM_HEIGHT) + ui.output_plot( + "plot_connectogram", height=DEFAULT_PLOT_CONNECTOGRAM_HEIGHT + ), ), ), ), @@ -100,56 +108,56 @@ def ui_option_row(label, ui_element, sizes=(6, 6), center=False): ui.row( ui.column( 12, - ui.output_plot('plot_coherence_per_task', height=DEFAULT_PLOT_COHERENCE_PER_TASK_HEIGHT) + ui.output_plot( + "plot_coherence_per_task", + height=DEFAULT_PLOT_COHERENCE_PER_TASK_HEIGHT, + ), ), ), ), ui.nav_panel( "Data Frame", ui.row( - ui.column( - 12, - ui.output_data_frame('data_frame') - ), + ui.column(12, ui.output_data_frame("data_frame")), ), ), ui.nav_panel( "Wavelet Transform Coherence", ui.row( ui.column( - 10, - ui.output_plot('plot_wtc', height=DEFAULT_PLOT_WTC_HEIGHT) + 10, ui.output_plot("plot_wtc", height=DEFAULT_PLOT_WTC_HEIGHT) ), ui.column( 2, - ui.output_ui('ui_input_select_wtc'), + ui.output_ui("ui_input_select_wtc"), ), ), ), ui.nav_spacer(), - selected='Study Info', - #selected='Coherence Matrix', - #selected='Connectograms', - #selected='Coherence Per Task', - #selected='Wavelet Transform Coherence', - id='main_nav', + selected="Study Info", + # selected='Coherence Matrix', + # selected='Connectograms', + # selected='Coherence Per Task', + # selected='Wavelet Transform Coherence', + id="main_nav", sidebar=ui.sidebar( - ui.output_ui('ui_input_study_file'), - ui.output_ui('ui_input_select_dyad'), - ui.output_ui('ui_input_select_task'), + ui.output_ui("ui_input_study_file"), + ui.output_ui("ui_input_select_dyad"), + ui.output_ui("ui_input_select_task"), width=SIDEBAR_WIDTH, ), title="HyPyP fNIRS results viewer", fillable=True, - bg="transparent" + bg="transparent", ) ) + def server(input: Inputs, output: Outputs, session: Session): @reactive.calc def get_study() -> Study: study_file_name = input.select_study_file() - if study_file_name == '': + if study_file_name == "": return None study_file_path = os.path.join(HARDCODED_RESULTS_PATH, study_file_name) return Study.from_pickle(study_file_path) @@ -162,7 +170,7 @@ def get_dyad() -> Dyad: if input.select_dyad() is None or input.select_dyad() == STR_ALL_DYADS: return None - + dyads = [dyad for dyad in study.dyads if dyad.label == input.select_dyad()] if len(dyads) == 0: return None @@ -180,17 +188,19 @@ def get_wtc() -> WTC: @render.ui def ui_input_study_file(): - my_list = [d for d in os.listdir(HARDCODED_RESULTS_PATH) if d.endswith('.pickle')] + my_list = [ + d for d in os.listdir(HARDCODED_RESULTS_PATH) if d.endswith(".pickle") + ] my_list.sort() - my_list = [''] + my_list - #print(my_list) + my_list = [""] + my_list + # print(my_list) return ui.input_select( "select_study_file", f"Study file ({HARDCODED_RESULTS_PATH})", choices=my_list, selected="fnirs_study_lionlab.pickle", ) - + @render.ui def ui_input_select_dyad(): study = get_study() @@ -201,7 +211,7 @@ def ui_input_select_dyad(): f"Select Dyad", choices=[STR_ALL_DYADS] + [dyad.label for dyad in study.dyads], ) - + @render.ui def ui_input_select_task(): study = get_study() @@ -218,7 +228,7 @@ def ui_input_select_task(): f"Select Task", choices=[STR_ALL_TASKS] + [task.name for task in tasks], ) - + @render.ui def ui_input_select_wtc(): dyad = get_dyad() @@ -228,23 +238,32 @@ def ui_input_select_wtc(): "select_wtc", f"Select Channel Pair", # TODO the filtering here is sketchy - choices=[wtc.label_pair for wtc in dyad.wtcs if wtc.label_pair.startswith(input.select_task())] + choices=[ + wtc.label_pair + for wtc in dyad.wtcs + if wtc.label_pair.startswith(input.select_task()) + ], ) - + @render.table def table_info_study(): study = get_study() if study is None: return None - return pd.DataFrame({ - 'Dyad Label': [dyad.label for dyad in study.dyads], - 'Tasks': [', '.join([task.name for task in dyad.tasks]) for dyad in study.dyads], - 'Subject 1': [dyad.s1.subject_label for dyad in study.dyads], - 'Subject 2': [dyad.s2.subject_label for dyad in study.dyads], - }) - + return pd.DataFrame( + { + "Dyad Label": [dyad.label for dyad in study.dyads], + "Tasks": [ + ", ".join([task.name for task in dyad.tasks]) + for dyad in study.dyads + ], + "Subject 1": [dyad.s1.subject_label for dyad in study.dyads], + "Subject 2": [dyad.s2.subject_label for dyad in study.dyads], + } + ) + def get_query(): - task = input.select_task() + task = input.select_task() q = f'task == "{task}"' if task != STR_ALL_TASKS else None return q @@ -257,16 +276,16 @@ def plot_coherence_matrix(): return None grouping = input.coherence_select_grouping() - if grouping == 'roi': + if grouping == "roi": return obj.plot_coherence_matrix_per_roi(query=get_query()) - elif grouping == 'channel': + elif grouping == "channel": return obj.plot_coherence_matrix_per_channel(query=get_query()) - elif grouping == 'roi_channel': - return obj.plot_coherence_matrix('roi1', 'channel2', query=get_query()) - elif grouping == 'channel_roi': - return obj.plot_coherence_matrix('channel1', 'roi2', query=get_query()) + elif grouping == "roi_channel": + return obj.plot_coherence_matrix("roi1", "channel2", query=get_query()) + elif grouping == "channel_roi": + return obj.plot_coherence_matrix("channel1", "roi2", query=get_query()) else: - raise RuntimeError(f'Unknown grouping {grouping}') + raise RuntimeError(f"Unknown grouping {grouping}") @render.plot def plot_connectogram_s1(): @@ -275,8 +294,8 @@ def plot_connectogram_s1(): obj = get_study() if obj is None: return None - return obj.plot_coherence_connectogram_s1(query=get_query(), title='Subject1') - + return obj.plot_coherence_connectogram_s1(query=get_query(), title="Subject1") + @render.plot def plot_connectogram_s2(): obj = get_dyad() @@ -284,8 +303,8 @@ def plot_connectogram_s2(): obj = get_study() if obj is None: return None - return obj.plot_coherence_connectogram_s2(query=get_query(), title='Subject2') - + return obj.plot_coherence_connectogram_s2(query=get_query(), title="Subject2") + @render.plot def plot_connectogram(): obj = get_dyad() @@ -294,7 +313,7 @@ def plot_connectogram(): if obj is None: return None return obj.plot_coherence_connectogram(query=get_query()) - + @render.data_frame def data_frame(): obj = get_dyad() @@ -304,7 +323,7 @@ def data_frame(): return None # have a maximum so that it loads return render.DataGrid(obj.df[:1000], selection_mode="rows") - + @render.plot def plot_coherence_per_task(): obj = get_dyad() @@ -313,14 +332,14 @@ def plot_coherence_per_task(): if obj is None: return None return obj.plot_coherence_bars_per_task() - + @render.plot def plot_wtc(): fig, ax = plt.subplots() wtc = get_wtc() if wtc is None: return None - + wtc.plot( ax=ax, show_colorbar=True, @@ -328,6 +347,6 @@ def plot_wtc(): show_nyquist=True, ) return fig - -app = App(app_ui, server) \ No newline at end of file + +app = App(app_ui, server) diff --git a/hypyp/shiny/wavelet_explorer.py b/hypyp/shiny/wavelet_explorer.py index b51e164..bac104f 100644 --- a/hypyp/shiny/wavelet_explorer.py +++ b/hypyp/shiny/wavelet_explorer.py @@ -11,27 +11,37 @@ from hypyp.wavelet.pair_signals import PairSignals from hypyp.fnirs.data_browser import DataBrowser from hypyp.fnirs.recording import Recording -from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_raw_to_haemo import MnePreprocessorRawToHaemo -from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_as_is import MnePreprocessorAsIs -from hypyp.fnirs.preprocessor.implementations.cedalion_preprocessor import CedalionPreprocessor +from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_raw_to_haemo import ( + MnePreprocessorRawToHaemo, +) +from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_as_is import ( + MnePreprocessorAsIs, +) +from hypyp.fnirs.preprocessor.implementations.cedalion_preprocessor import ( + CedalionPreprocessor, +) from hypyp.signal import SyntheticSignal from hypyp.wavelet.base_wavelet import BaseWavelet from hypyp.wavelet.implementations.matlab_wavelet import MatlabWavelet -from hypyp.wavelet.implementations.pywavelets_wavelet import ComplexGaussianWavelet, ComplexMorletWavelet +from hypyp.wavelet.implementations.pywavelets_wavelet import ( + ComplexGaussianWavelet, + ComplexMorletWavelet, +) from hypyp.wavelet.implementations.pycwt_wavelet import PycwtWavelet -DEFAULT_PLOT_SIGNAL_HEIGHT = 150 # px -DEFAULT_PLOT_MNE_HEIGHT = 1200 # px -DEFAULT_SIGNAL_DURATION_SELECTION = 60 # seconds +DEFAULT_PLOT_SIGNAL_HEIGHT = 150 # px +DEFAULT_PLOT_MNE_HEIGHT = 1200 # px +DEFAULT_SIGNAL_DURATION_SELECTION = 60 # seconds # Use this environment variable to add a lookup path for NIRS data HYPYP_NIRS_DATA_PATH = os.getenv("HYPYP_NIRS_DATA_PATH", "") -STR_SAME_AS_SUBJECT_1 = 'Same as Subject 1' +STR_SAME_AS_SUBJECT_1 = "Same as Subject 1" # This is to avoid having external windows launched -matplotlib.use('Agg') +matplotlib.use("Agg") + def ui_option_row(label, ui_element, sizes=(6, 6), center=False): label_style = "" @@ -42,157 +52,222 @@ def ui_option_row(label, ui_element, sizes=(6, 6), center=False): ui.column(sizes[1], ui_element), ) + # UI app_ui = ui.page_fluid( ui.page_navbar( ui.nav_spacer(), ui.nav_panel( "Data Browser", - ui.tags.h4(ui.output_text('text_s1_file_path')), - ui.output_ui('ui_input_s1_mne_duration_slider'), - ui.output_ui('ui_preprocess_steps'), + ui.tags.h4(ui.output_text("text_s1_file_path")), + ui.output_ui("ui_input_s1_mne_duration_slider"), + ui.output_ui("ui_preprocess_steps"), ), ui.nav_panel( "Wavelet Coherence", ui.row( - ui.column(12, + ui.column( + 12, ui.card( - ui.tags.strong('Choice of signals'), - ui.output_ui('ui_input_signal_slider'), - ui.output_plot('plot_signals', height=DEFAULT_PLOT_SIGNAL_HEIGHT), - ui.output_plot('plot_signals_spectrum', height=DEFAULT_PLOT_SIGNAL_HEIGHT), + ui.tags.strong("Choice of signals"), + ui.output_ui("ui_input_signal_slider"), + ui.output_plot( + "plot_signals", height=DEFAULT_PLOT_SIGNAL_HEIGHT + ), + ui.output_plot( + "plot_signals_spectrum", height=DEFAULT_PLOT_SIGNAL_HEIGHT + ), ), ), ), ui.row( - ui.column(6, + ui.column( + 6, ui.card( - ui.tags.strong('Wavelet'), - ui.output_plot('plot_mother_wavelet'), + ui.tags.strong("Wavelet"), + ui.output_plot("plot_mother_wavelet"), ), ), - ui.column(6, + ui.column( + 6, ui.card( - ui.tags.strong('Wavelet Coherence'), - ui.output_plot('plot_wtc'), + ui.tags.strong("Wavelet Coherence"), + ui.output_plot("plot_wtc"), ), ), ), ), ui.nav_spacer(), - #selected='Data Browser', - selected='Wavelet Coherence', - id='main_nav', + # selected='Data Browser', + selected="Wavelet Coherence", + id="main_nav", sidebar=ui.sidebar( - ui.tags.strong('Signal parameters'), + ui.tags.strong("Signal parameters"), ui.input_select( "signal_type", "", choices={ - 'data_files': 'Local data files', - 'testing': 'Synthetic signals', + "data_files": "Local data files", + "testing": "Synthetic signals", }, ), - ui.output_ui('ui_input_signal_choice'), - ui.output_ui('ui_input_signal_choice_property'), - ui.output_ui('ui_input_signal_options'), - - ui.tags.strong('Wavelet parameters'), + ui.output_ui("ui_input_signal_choice"), + ui.output_ui("ui_input_signal_choice_property"), + ui.output_ui("ui_input_signal_options"), + ui.tags.strong("Wavelet parameters"), ui.input_select( "wavelet_library", "", choices={ - 'pywavelets': 'Use pywavelets', - 'pycwt': 'Use pycwt (based on matlab code)', - 'matlab': 'Use Matlab engine', + "pywavelets": "Use pywavelets", + "pycwt": "Use pycwt (based on matlab code)", + "matlab": "Use Matlab engine", }, ), - ui.output_ui('ui_input_wavelet_type'), - ui.output_ui('ui_input_wavelet_options'), - ui.output_ui('ui_input_coherence_options'), - - ui.input_action_button("button_action_compute_wtc", label="Compute Wavelet Transform Coherence"), - - ui.tags.strong('Display parameters'), - ui_option_row("Signal offset y", ui.input_checkbox("display_signal_offset_y", "", value=True), sizes=(8,4)), - ui_option_row("Downsample", ui.input_checkbox("display_downsample", "", value=True), sizes=(8,4)), - ui_option_row("Show COI", ui.input_checkbox("display_show_coi", "", value=True), sizes=(8,4)), - ui_option_row("Show Nyquist", ui.input_checkbox("display_show_nyquist", "", value=True), sizes=(8,4)), + ui.output_ui("ui_input_wavelet_type"), + ui.output_ui("ui_input_wavelet_options"), + ui.output_ui("ui_input_coherence_options"), + ui.input_action_button( + "button_action_compute_wtc", label="Compute Wavelet Transform Coherence" + ), + ui.tags.strong("Display parameters"), + ui_option_row( + "Signal offset y", + ui.input_checkbox("display_signal_offset_y", "", value=True), + sizes=(8, 4), + ), + ui_option_row( + "Downsample", + ui.input_checkbox("display_downsample", "", value=True), + sizes=(8, 4), + ), + ui_option_row( + "Show COI", + ui.input_checkbox("display_show_coi", "", value=True), + sizes=(8, 4), + ), + ui_option_row( + "Show Nyquist", + ui.input_checkbox("display_show_nyquist", "", value=True), + sizes=(8, 4), + ), open="always", width=400, ), title="HyPyP fNIRS explorer", fillable=True, - bg="transparent" + bg="transparent", ) ) + def server(input: Inputs, output: Outputs, session: Session): @reactive.calc() def get_signals() -> PairSignals: pair = None - if input.signal_type() == 'testing': + if input.signal_type() == "testing": fs = input.signal_sampling_frequency() N = input.signal_n() - x = np.linspace(0, N/fs, N) + x = np.linspace(0, N / fs, N) tmax = N / fs - if input.signal_testing_choice() == 'sinusoid': - y1 = SyntheticSignal(duration=tmax, n_points=N).add_sin(input.signal_sinusoid_freq1()).y - y2 = SyntheticSignal(duration=tmax, n_points=N).add_sin(input.signal_sinusoid_freq2()).y - elif input.signal_testing_choice() == 'sinusoid_almost_similar': - y1 = SyntheticSignal(duration=tmax, n_points=N).add_sin(input.signal_sinusoid_almost_similar_freq1()).y - y2 = SyntheticSignal(duration=tmax, n_points=N).add_sin(input.signal_sinusoid_almost_similar_freq1() - 0.001).y - elif input.signal_testing_choice() == 'sinusoid_dephased': - y1 = SyntheticSignal(duration=tmax, n_points=N).add_sin(input.signal_sinusoid_dephased_freq1()).y + if input.signal_testing_choice() == "sinusoid": + y1 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_sin(input.signal_sinusoid_freq1()) + .y + ) + y2 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_sin(input.signal_sinusoid_freq2()) + .y + ) + elif input.signal_testing_choice() == "sinusoid_almost_similar": + y1 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_sin(input.signal_sinusoid_almost_similar_freq1()) + .y + ) + y2 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_sin(input.signal_sinusoid_almost_similar_freq1() - 0.001) + .y + ) + elif input.signal_testing_choice() == "sinusoid_dephased": + y1 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_sin(input.signal_sinusoid_dephased_freq1()) + .y + ) y1_fft = np.fft.fft(y1) y1_magnitude = np.abs(y1_fft) y1_phase = np.angle(y1_fft) y1_random_phase = np.random.uniform(-np.pi, np.pi, len(y1_phase)) y1_randomized_fft = y1_magnitude * np.exp(1j * y1_random_phase) y2 = np.fft.ifft(y1_randomized_fft).real - y2 = y2 * 1 / (np.abs(np.max(y2)-np.min(y2)) / 2) - elif input.signal_testing_choice() == 'chirp': + y2 = y2 * 1 / (np.abs(np.max(y2) - np.min(y2)) / 2) + elif input.signal_testing_choice() == "chirp": signal1 = SyntheticSignal(duration=tmax, n_points=N) - signal1.add_chirp(input.signal_chirp_freq1(), input.signal_chirp_freq2()) + signal1.add_chirp( + input.signal_chirp_freq1(), input.signal_chirp_freq2() + ) y1 = signal1.y y2 = np.flip(y1) - elif input.signal_testing_choice() == 'chirp_sinusoid': - y1 = SyntheticSignal(duration=tmax, n_points=N).add_sin(input.signal_chirp_sinusoid_freq1()).y - y2 = SyntheticSignal(duration=tmax, n_points=N).add_chirp(input.signal_chirp_sinusoid_freq2(), input.signal_chirp_sinusoid_freq3()).y + elif input.signal_testing_choice() == "chirp_sinusoid": + y1 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_sin(input.signal_chirp_sinusoid_freq1()) + .y + ) + y2 = ( + SyntheticSignal(duration=tmax, n_points=N) + .add_chirp( + input.signal_chirp_sinusoid_freq2(), + input.signal_chirp_sinusoid_freq3(), + ) + .y + ) noise_level = input.signal_noise_level() - y1 += noise_level * np.random.normal(0, 1, len(x)) - y2 += noise_level * np.random.normal(0, 1, len(x)) + y1 += noise_level * np.random.normal(0, 1, len(x)) + y2 += noise_level * np.random.normal(0, 1, len(x)) - elif input.signal_type() == 'data_files': + elif input.signal_type() == "data_files": # TODO this only works with MNE s1_raw = get_recording_s1_step().obj.copy() s2_raw = get_recording_s2_step().obj.copy() - s1_selected_ch = s1_raw.pick(mne.pick_channels(s1_raw.ch_names, include = [input.signal_data_files_s1_channel()])) - s2_selected_ch = s2_raw.pick(mne.pick_channels(s2_raw.ch_names, include = [input.signal_data_files_s2_channel()])) + s1_selected_ch = s1_raw.pick( + mne.pick_channels( + s1_raw.ch_names, include=[input.signal_data_files_s1_channel()] + ) + ) + s2_selected_ch = s2_raw.pick( + mne.pick_channels( + s2_raw.ch_names, include=[input.signal_data_files_s2_channel()] + ) + ) - y1 = np.ma.masked_invalid(s1_selected_ch.get_data())[0,:] - y2 = np.ma.masked_invalid(s2_selected_ch.get_data())[0,:] + y1 = np.ma.masked_invalid(s1_selected_ch.get_data())[0, :] + y2 = np.ma.masked_invalid(s2_selected_ch.get_data())[0, :] # Crop signals stop = min(len(y1), len(y2)) - y1 = y1[:stop] - y2 = y2[:stop] - + y1 = y1[:stop] + y2 = y2[:stop] + # Force N for this signal N = len(y1) - fs = s1_selected_ch.info['sfreq'] - x = np.linspace(0, N/fs, N) + fs = s1_selected_ch.info["sfreq"] + x = np.linspace(0, N / fs, N) if pair is None: pair = PairSignals( x, y1, y2, - label_s1='s1', - label_s2='s2', - label_task='foo', + label_s1="s1", + label_s2="s2", + label_task="foo", ) try: @@ -200,46 +275,53 @@ def get_signals() -> PairSignals: return pair.sub(input.signal_range()) except: pass - + return pair def get_signal_duration(): - if input.signal_type() == 'testing': + if input.signal_type() == "testing": fs = input.signal_sampling_frequency() N = input.signal_n() - return N/fs - if input.signal_type() == 'data_files': - return min([get_recording_s1_step().duration, get_recording_s2_step().duration]) - + return N / fs + if input.signal_type() == "data_files": + return min( + [get_recording_s1_step().duration, get_recording_s2_step().duration] + ) @reactive.calc() def get_wavelet(): - if input.wavelet_library() == 'pywavelets': + if input.wavelet_library() == "pywavelets": wavelet_name = input.wavelet_name() - if wavelet_name == 'cmor': + if wavelet_name == "cmor": return ComplexMorletWavelet( bandwidth_frequency=input.wavelet_bandwidth(), center_frequency=input.wavelet_center_frequency(), - period_range=(input.wavelet_period_range_low(), input.wavelet_period_range_high()), + period_range=( + input.wavelet_period_range_low(), + input.wavelet_period_range_high(), + ), disable_caching=True, ) - elif wavelet_name == 'cgau': + elif wavelet_name == "cgau": return ComplexGaussianWavelet( degree=int(input.wavelet_degree()), - period_range=(input.wavelet_period_range_low(), input.wavelet_period_range_high()), + period_range=( + input.wavelet_period_range_low(), + input.wavelet_period_range_high(), + ), wtc_smoothing_win_size=input.smoothing_win_size(), disable_caching=True, ) else: - raise RuntimeError(f'Unknown wavelet_name: {wavelet_name}') + raise RuntimeError(f"Unknown wavelet_name: {wavelet_name}") - if input.wavelet_library() == 'pycwt': + if input.wavelet_library() == "pycwt": return PycwtWavelet() - if input.wavelet_library() == 'matlab': + if input.wavelet_library() == "matlab": return MatlabWavelet() - raise RuntimeError(f'Unknown wavelet library: {input.wavelet_library()}') + raise RuntimeError(f"Unknown wavelet library: {input.wavelet_library()}") @reactive.event(input.button_action_compute_wtc) def compute_coherence(): @@ -252,7 +334,7 @@ def compute_coherence(): def get_signal_data_files_s1_path(): return input.signal_data_files_s1_path() - + def get_signal_data_files_s2_path(): value = input.signal_data_files_s2_path() if value == STR_SAME_AS_SUBJECT_1: @@ -284,41 +366,48 @@ def get_data_browser(): def get_preprocessor(): value = input.recording_preprocessor() - if value == 'upstream': + if value == "upstream": return MnePreprocessorAsIs() - if value == 'mne': + if value == "mne": return MnePreprocessorRawToHaemo() - if value == 'cedalion': + if value == "cedalion": return CedalionPreprocessor() raise RuntimeError(f'Unknown preprocessor type "{value}"') - @reactive.calc() def get_recording_s1(): - return Recording().load_file(get_signal_data_files_s1_path(), get_preprocessor()) + return Recording().load_file( + get_signal_data_files_s1_path(), get_preprocessor() + ) @reactive.calc() def get_recording_s2(): - return Recording().load_file(get_signal_data_files_s2_path(), get_preprocessor()) + return Recording().load_file( + get_signal_data_files_s2_path(), get_preprocessor() + ) @reactive.calc() def get_recording_s1_step(): - return get_recording_s1().get_preprocess_step(input.signal_data_files_analysis_property()) + return get_recording_s1().get_preprocess_step( + input.signal_data_files_analysis_property() + ) @reactive.calc() def get_recording_s2_step(): - return get_recording_s2().get_preprocess_step(input.signal_data_files_analysis_property()) + return get_recording_s2().get_preprocess_step( + input.signal_data_files_analysis_property() + ) def get_mne_raw_plot_kwargs(step, duration_range): try: range_start, range_end = duration_range start = range_start - duration = (range_end - range_start) + duration = range_end - range_start return dict( n_channels=len(step.ch_names), start=start, duration=duration, - scalings='auto', + scalings="auto", show_scrollbars=False, block=True, show=False, @@ -338,12 +427,21 @@ def ui_preprocess_steps(): # Because of order of execution, this cannot be directly in the loop def bind_plot_mne_figure(step: int): def plot_mne_figure(): - return mne_figure_as_image(step.plot(**get_mne_raw_plot_kwargs(step, input.input_s1_mne_duration_slider()))) + return mne_figure_as_image( + step.plot( + **get_mne_raw_plot_kwargs( + step, input.input_s1_mne_duration_slider() + ) + ) + ) + # need to rename the function because every "output plot" must have a unique name - plot_mne_figure.__name__ = f'{plot_mne_figure.__name__}_{step.key}' + plot_mne_figure.__name__ = f"{plot_mne_figure.__name__}_{step.key}" renderer = render.image(plot_mne_figure) # This is needed to avoid having the scrollbar on the right - renderer._auto_output_ui_kwargs = dict(height=f'{DEFAULT_PLOT_MNE_HEIGHT}px') + renderer._auto_output_ui_kwargs = dict( + height=f"{DEFAULT_PLOT_MNE_HEIGHT}px" + ) return renderer nav_panels = [] @@ -351,99 +449,213 @@ def plot_mne_figure(): for step in recording.preprocess_steps: nav_panels.append(ui.nav_panel(step.desc, bind_plot_mne_figure(step))) - return ui.navset_card_tab(*nav_panels, selected=recording.preprocess_steps[-1].desc) + return ui.navset_card_tab( + *nav_panels, selected=recording.preprocess_steps[-1].desc + ) @render.plot def plot_s1_mne_sci(): recording = get_recording_s1() fig, ax = plt.subplots() ax.hist(recording.quality_sci) - ax.set(xlabel='Scalp Coupling Index', ylabel='Count', xlim=[0, 1]) - ax.title.set_text('Scalp Coupling Index') + ax.set(xlabel="Scalp Coupling Index", ylabel="Count", xlim=[0, 1]) + ax.title.set_text("Scalp Coupling Index") return fig @render.ui def ui_input_signal_choice(): choices = [] - if input.signal_type() == 'testing': - choices.append(ui.input_select( - "signal_testing_choice", - "", - choices={ - 'chirp': 'Crossing chirps', - 'chirp_sinusoid': 'One sinusoid, one chirp', - 'sinusoid': 'Sinusoid different', - 'sinusoid_almost_similar': 'Sinusoid almost similar', - 'sinusoid_dephased': 'Sinusoid dephased', - }, - )) - - elif input.signal_type() == 'data_files': + if input.signal_type() == "testing": + choices.append( + ui.input_select( + "signal_testing_choice", + "", + choices={ + "chirp": "Crossing chirps", + "chirp_sinusoid": "One sinusoid, one chirp", + "sinusoid": "Sinusoid different", + "sinusoid_almost_similar": "Sinusoid almost similar", + "sinusoid_dephased": "Sinusoid dephased", + }, + ) + ) + + elif input.signal_type() == "data_files": browser = get_data_browser() browser.download_demo_dataset() - choices.append(ui_option_row("Subject 1 file", ui.input_select( - "signal_data_files_s1_path", - "", - choices=browser.list_all_files(), - # for comparison with jupyter notebook - #selected="/home/patrice/work/ppsp/HyPyP-synchro/data/NIRS/downloads/fathers/FCS28/parent/NIRS-2019-11-10_003.hdr", - ))) - choices.append(ui_option_row("Subject 2 file", ui.input_select( - "signal_data_files_s2_path", - "", - choices=[STR_SAME_AS_SUBJECT_1] + browser.list_all_files(), - ))) - choices.append(ui_option_row("Preprocessor Class", ui.input_select( - "recording_preprocessor", - "", - choices={ - 'upstream': 'MNE (no preprocessing, load as-is)', - 'mne': 'MNE (basic fNIRS preprocessing)', - 'cedalion': 'Cedalion (proof of concept preprocessing)', - } - ))) - + choices.append( + ui_option_row( + "Subject 1 file", + ui.input_select( + "signal_data_files_s1_path", + "", + choices=browser.list_all_files(), + # for comparison with jupyter notebook + # selected="/home/patrice/work/ppsp/HyPyP-synchro/data/NIRS/downloads/fathers/FCS28/parent/NIRS-2019-11-10_003.hdr", + ), + ) + ) + choices.append( + ui_option_row( + "Subject 2 file", + ui.input_select( + "signal_data_files_s2_path", + "", + choices=[STR_SAME_AS_SUBJECT_1] + browser.list_all_files(), + ), + ) + ) + choices.append( + ui_option_row( + "Preprocessor Class", + ui.input_select( + "recording_preprocessor", + "", + choices={ + "upstream": "MNE (no preprocessing, load as-is)", + "mne": "MNE (basic fNIRS preprocessing)", + "cedalion": "Cedalion (proof of concept preprocessing)", + }, + ), + ) + ) + return choices @render.ui def ui_input_signal_choice_property(): options = [] - if input.signal_type() == 'data_files': - options.append(ui_option_row( - "Which step to use for analysis", - ui.input_select( - "signal_data_files_analysis_property", - "", - choices=get_recording_s1().preprocess_step_choices, + if input.signal_type() == "data_files": + options.append( + ui_option_row( + "Which step to use for analysis", + ui.input_select( + "signal_data_files_analysis_property", + "", + choices=get_recording_s1().preprocess_step_choices, + ), ) - )) + ) return options @render.ui def ui_input_signal_options(): options = [] - if input.signal_type() == 'testing': - if input.signal_testing_choice() == 'sinusoid': - options.append(ui_option_row("Freq 1", ui.input_numeric("signal_sinusoid_freq1", "", value=1))), - options.append(ui_option_row("Freq 2", ui.input_numeric("signal_sinusoid_freq2", "", value=0.8))), - elif input.signal_testing_choice() == 'sinusoid_almost_similar': - options.append(ui_option_row("Freq", ui.input_numeric("signal_sinusoid_almost_similar_freq1", "", value=0.2))), - elif input.signal_testing_choice() == 'sinusoid_dephased': - options.append(ui_option_row("Freq", ui.input_numeric("signal_sinusoid_dephased_freq1", "", value=0.02))), - elif input.signal_testing_choice() == 'chirp': - options.append(ui_option_row("Freq Chirp from", ui.input_numeric("signal_chirp_freq1", "", value=0.2))), - options.append(ui_option_row("Freq Chirp to", ui.input_numeric("signal_chirp_freq2", "", value=0.05))), - elif input.signal_testing_choice() == 'chirp_sinusoid': - options.append(ui_option_row("Freq Sinusoid", ui.input_numeric("signal_chirp_sinusoid_freq1", "", value=1))), - options.append(ui_option_row("Freq Chirp from", ui.input_numeric("signal_chirp_sinusoid_freq2", "", value=0.2))), - options.append(ui_option_row("Freq Chirp to", ui.input_numeric("signal_chirp_sinusoid_freq3", "", value=2))), - - options.append(ui_option_row("Sampling freq. (Hz)", ui.input_numeric("signal_sampling_frequency", "", value=5))), - options.append(ui_option_row("Nb. points", ui.input_numeric("signal_n", "", value=2000))), - options.append(ui_option_row("Noise level", ui.input_numeric("signal_noise_level", "", value=0.1))), - - if input.signal_type() == 'data_files': + if input.signal_type() == "testing": + if input.signal_testing_choice() == "sinusoid": + ( + options.append( + ui_option_row( + "Freq 1", + ui.input_numeric("signal_sinusoid_freq1", "", value=1), + ) + ), + ) + ( + options.append( + ui_option_row( + "Freq 2", + ui.input_numeric("signal_sinusoid_freq2", "", value=0.8), + ) + ), + ) + elif input.signal_testing_choice() == "sinusoid_almost_similar": + ( + options.append( + ui_option_row( + "Freq", + ui.input_numeric( + "signal_sinusoid_almost_similar_freq1", "", value=0.2 + ), + ) + ), + ) + elif input.signal_testing_choice() == "sinusoid_dephased": + ( + options.append( + ui_option_row( + "Freq", + ui.input_numeric( + "signal_sinusoid_dephased_freq1", "", value=0.02 + ), + ) + ), + ) + elif input.signal_testing_choice() == "chirp": + ( + options.append( + ui_option_row( + "Freq Chirp from", + ui.input_numeric("signal_chirp_freq1", "", value=0.2), + ) + ), + ) + ( + options.append( + ui_option_row( + "Freq Chirp to", + ui.input_numeric("signal_chirp_freq2", "", value=0.05), + ) + ), + ) + elif input.signal_testing_choice() == "chirp_sinusoid": + ( + options.append( + ui_option_row( + "Freq Sinusoid", + ui.input_numeric( + "signal_chirp_sinusoid_freq1", "", value=1 + ), + ) + ), + ) + ( + options.append( + ui_option_row( + "Freq Chirp from", + ui.input_numeric( + "signal_chirp_sinusoid_freq2", "", value=0.2 + ), + ) + ), + ) + ( + options.append( + ui_option_row( + "Freq Chirp to", + ui.input_numeric( + "signal_chirp_sinusoid_freq3", "", value=2 + ), + ) + ), + ) + + ( + options.append( + ui_option_row( + "Sampling freq. (Hz)", + ui.input_numeric("signal_sampling_frequency", "", value=5), + ) + ), + ) + ( + options.append( + ui_option_row( + "Nb. points", ui.input_numeric("signal_n", "", value=2000) + ) + ), + ) + ( + options.append( + ui_option_row( + "Noise level", + ui.input_numeric("signal_noise_level", "", value=0.1), + ) + ), + ) + + if input.signal_type() == "data_files": # TODO this try-except is here to have a PoC of Cedalion integration try: ch_names1 = get_recording_s1().preprocessed.ch_names @@ -452,23 +664,22 @@ def ui_input_signal_options(): ch_names1 = [] ch_names2 = [] - options.append(ui_option_row( - "Subject 1 channel", - ui.input_select( - "signal_data_files_s1_channel", - "", - choices=ch_names1 + options.append( + ui_option_row( + "Subject 1 channel", + ui.input_select( + "signal_data_files_s1_channel", "", choices=ch_names1 + ), ) - - )) - options.append(ui_option_row( - "Subject 2 channel", - ui.input_select( - "signal_data_files_s2_channel", - "", - choices=ch_names2 + ) + options.append( + ui_option_row( + "Subject 2 channel", + ui.input_select( + "signal_data_files_s2_channel", "", choices=ch_names2 + ), ) - )) + ) return options @@ -476,13 +687,27 @@ def ui_input_signal_options(): def ui_input_signal_slider(): duration = int(get_signal_duration()) default_duration = min([DEFAULT_SIGNAL_DURATION_SELECTION, duration]) - return ui.input_slider("signal_range", "", min=0, max=duration, value=[0, default_duration], width='100%') + return ui.input_slider( + "signal_range", + "", + min=0, + max=duration, + value=[0, default_duration], + width="100%", + ) @render.ui def ui_input_s1_mne_duration_slider(): duration = int(get_signal_duration()) default_duration = min([DEFAULT_SIGNAL_DURATION_SELECTION, duration]) - return ui.input_slider("input_s1_mne_duration_slider", "", min=0, max=duration, value=[0, default_duration], width='100%') + return ui.input_slider( + "input_s1_mne_duration_slider", + "", + min=0, + max=duration, + value=[0, default_duration], + width="100%", + ) @render.plot(height=DEFAULT_PLOT_SIGNAL_HEIGHT) def plot_signals(): @@ -494,7 +719,6 @@ def plot_signals(): ax.plot(pair.x, pair.y1 + offset) ax.plot(pair.x, pair.y2 - offset) return fig - @render.plot(height=DEFAULT_PLOT_SIGNAL_HEIGHT) def plot_signals_spectrum(): @@ -503,14 +727,14 @@ def plot_signals_spectrum(): N = pair.n yf = fft.fft(pair.y1) - xf = fft.fftfreq(N, pair.dt)[:N//2] + xf = fft.fftfreq(N, pair.dt)[: N // 2] yf2 = fft.fft(pair.y2) - ax.plot(xf, 2.0/N * np.abs(yf[0:N//2])) - ax.plot(xf, 2.0/N * np.abs(yf2[0:N//2])) - ax.set_xscale('log') - ax.set_yscale('log') - ax.title.set_text('Spectrum') + ax.plot(xf, 2.0 / N * np.abs(yf[0 : N // 2])) + ax.plot(xf, 2.0 / N * np.abs(yf2[0 : N // 2])) + ax.set_xscale("log") + ax.set_yscale("log") + ax.title.set_text("Spectrum") ax.grid() return fig @@ -522,7 +746,7 @@ def plot_mother_wavelet(): ax.plot(wav.psi_x, np.imag(wav.psi)) ax.plot(wav.psi_x, np.abs(wav.psi)) ax.title.set_text(f"mother wavelet ({wav.wavelet_name_with_args})") - ax.legend(['real', 'imag', 'abs']) + ax.legend(["real", "imag", "abs"]) return fig @render.plot() @@ -539,48 +763,94 @@ def plot_wtc(): @render.ui def ui_input_wavelet_type(): options = [] - if input.wavelet_library() == 'pywavelets': - options.append(ui.input_select( - "wavelet_name", - "", - choices={ - 'cmor': 'Complex Morlet', - 'cgau': 'Complex Gaussian', - }, - )) + if input.wavelet_library() == "pywavelets": + options.append( + ui.input_select( + "wavelet_name", + "", + choices={ + "cmor": "Complex Morlet", + "cgau": "Complex Gaussian", + }, + ) + ) return options - @render.ui def ui_input_wavelet_options(): options = [] - if input.wavelet_library() == 'pywavelets': - if input.wavelet_name() == 'cmor': - options.append(ui_option_row("Bandwidth/Center Freq.", ui.row( - ui.column(6, ui.input_numeric("wavelet_bandwidth", "", value=ComplexMorletWavelet.default_bandwidth_frequency)), - ui.column(6, ui.input_numeric("wavelet_center_frequency", "", value=ComplexMorletWavelet.default_center_frequency)), - ))) - if input.wavelet_name() == 'cgau': - options.append(ui_option_row("Gaussian degree", ui.row( - ui.input_select( - "wavelet_degree", - "", - choices=list(range(1,9)), - selected=ComplexGaussianWavelet.default_degree, - )))) + if input.wavelet_library() == "pywavelets": + if input.wavelet_name() == "cmor": + options.append( + ui_option_row( + "Bandwidth/Center Freq.", + ui.row( + ui.column( + 6, + ui.input_numeric( + "wavelet_bandwidth", + "", + value=ComplexMorletWavelet.default_bandwidth_frequency, + ), + ), + ui.column( + 6, + ui.input_numeric( + "wavelet_center_frequency", + "", + value=ComplexMorletWavelet.default_center_frequency, + ), + ), + ), + ) + ) + if input.wavelet_name() == "cgau": + options.append( + ui_option_row( + "Gaussian degree", + ui.row( + ui.input_select( + "wavelet_degree", + "", + choices=list(range(1, 9)), + selected=ComplexGaussianWavelet.default_degree, + ) + ), + ) + ) - options.append(ui_option_row("Period range", ui.row( - ui.column(6, ui.input_numeric("wavelet_period_range_low", "", value=BaseWavelet.default_period_range[0])), - ui.column(6, ui.input_numeric("wavelet_period_range_high", "", value=BaseWavelet.default_period_range[1])), - ))) + options.append( + ui_option_row( + "Period range", + ui.row( + ui.column( + 6, + ui.input_numeric( + "wavelet_period_range_low", + "", + value=BaseWavelet.default_period_range[0], + ), + ), + ui.column( + 6, + ui.input_numeric( + "wavelet_period_range_high", + "", + value=BaseWavelet.default_period_range[1], + ), + ), + ), + ) + ) return options - + @render.ui def ui_input_coherence_options(): options = [] - if input.wavelet_library() in ['pycwt']: + if input.wavelet_library() in ["pycwt"]: pass return options - -app = App(app_ui, server) \ No newline at end of file + + +app = App(app_ui, server) diff --git a/hypyp/signal.py b/hypyp/signal.py index b364445..ee119d5 100644 --- a/hypyp/signal.py +++ b/hypyp/signal.py @@ -6,6 +6,7 @@ from scipy.fft import fft, fftfreq from scipy.signal import chirp + class SyntheticSignal: x: np.ndarray y: np.ndarray @@ -27,7 +28,7 @@ def __init__(self, duration: float = 100, n_points: int = 1000): self.period = 1.0 / self.sampling_rate self.y = np.zeros_like(self.x) - def add_chirp(self, f0:float, f1:float): + def add_chirp(self, f0: float, f1: float): """ Add a chirp signal from frequencies f0 to f1 @@ -38,9 +39,9 @@ def add_chirp(self, f0:float, f1:float): Returns: SyntheticSignal: object itself, for chaining """ - self.y += chirp(self.x, f0=f0, f1=f1, t1=np.max(self.x), method='linear') + self.y += chirp(self.x, f0=f0, f1=f1, t1=np.max(self.x), method="linear") return self - + def add_sin(self, freq: float): """ Add a sinusoid to the signal at a specific frequency @@ -53,8 +54,8 @@ def add_sin(self, freq: float): """ self.y += np.sin(self.x * 2 * np.pi * freq) return self - - def add_noise(self, level:float=0.1): + + def add_noise(self, level: float = 0.1): """ Add noise to the signal @@ -66,8 +67,8 @@ def add_noise(self, level:float=0.1): """ self.y += level * np.random.normal(0, 1, self.n_points) return self - - def add_custom(self, y:np.ndarray): + + def add_custom(self, y: np.ndarray): """ Add all the received values to the signal @@ -79,11 +80,11 @@ def add_custom(self, y:np.ndarray): """ self.y += y return self - + # # Plots # - def plot(self, t:np.ndarray|None=None, ax:Axes|None=None) -> Figure: + def plot(self, t: np.ndarray | None = None, ax: Axes | None = None) -> Figure: if ax is None: fig, ax = plt.subplots() else: @@ -94,25 +95,22 @@ def plot(self, t:np.ndarray|None=None, ax:Axes|None=None) -> Figure: else: range = (int(t[0] * self.sampling_rate), int(t[1] * self.sampling_rate)) - ax.plot(self.x[range[0]:range[1]], self.y[range[0]:range[1]]) - ax.set_xlabel('Time (s)') + ax.plot(self.x[range[0] : range[1]], self.y[range[0] : range[1]]) + ax.set_xlabel("Time (s)") return fig - - def plot_fft(self, ax:Axes|None=None) -> Figure: + + def plot_fft(self, ax: Axes | None = None) -> Figure: if ax is None: fig, ax = plt.subplots() else: fig = ax.get_figure() - xf = fftfreq(self.n_points, self.period)[:self.n_points//2] + xf = fftfreq(self.n_points, self.period)[: self.n_points // 2] yf = fft(self.y) - #plt.plot(xf, 2.0/self.n_points * np.imag(yf[0:self.n_points//2])) - #plt.plot(xf, 2.0/self.n_points * np.imag(yf[0:self.n_points//2])) - ax.plot(xf, 2.0/self.n_points * np.abs(yf[0:self.n_points//2])) - ax.set_xscale('log') - ax.set_xlabel('Freq (Hz)') + # plt.plot(xf, 2.0/self.n_points * np.imag(yf[0:self.n_points//2])) + # plt.plot(xf, 2.0/self.n_points * np.imag(yf[0:self.n_points//2])) + ax.plot(xf, 2.0 / self.n_points * np.abs(yf[0 : self.n_points // 2])) + ax.set_xscale("log") + ax.set_xlabel("Freq (Hz)") return fig - - - diff --git a/hypyp/stats.py b/hypyp/stats.py index 4c659d3..293665e 100644 --- a/hypyp/stats.py +++ b/hypyp/stats.py @@ -11,7 +11,6 @@ | date | 2020-03-18 | """ - from collections import namedtuple from typing import List, Tuple import numpy as np @@ -23,32 +22,34 @@ from mne.stats import permutation_cluster_test -def statsCond(data: np.ndarray, epochs: mne.Epochs, n_permutations: int, alpha: float) -> namedtuple: +def statsCond( + data: np.ndarray, epochs: mne.Epochs, n_permutations: int, alpha: float +) -> namedtuple: """ Perform statistical t-test on participant measures (e.g., PSD) for a condition. - + This function tests whether the observed mean significantly deviates from 0 using a permutation-based t-test with False Discovery Rate (FDR) correction for multiple comparisons. - + Parameters ---------- data : np.ndarray Array of participant measures with shape (n_samples, n_tests, n_freq) where n_tests typically represents channels and n_freq the frequencies. Values will be averaged across frequencies for statistics. - + epochs : mne.Epochs Epochs object for a condition from a random participant, used only to access information like channel positions. - + n_permutations : int Number of permutations for the statistical test Should be at least 2*n_samples, e.g., 50000 - + alpha : float Significance threshold for the test, e.g., 0.05 - + Returns ------- statsCondTuple : namedtuple @@ -58,22 +59,22 @@ def statsCond(data: np.ndarray, epochs: mne.Epochs, n_permutations: int, alpha: - H0: T-statistics obtained by permutations, array of shape (n_permutations) - adj_p: tuple with boolean assessment of significance and FDR-corrected p-values - T_obs_plot: statistical values for significant sensors, array of shape (n_tests) - + Notes ----- This test calculates if the observed mean significantly deviates from 0; it doesn't compare two periods, but tests one period against the null hypothesis. - + Randomized data are generated with random sign flips, and the test is two-tailed by default (the alternative hypothesis is that the data mean is different from 0). - + To reduce false positives due to multiple comparisons, False Discovery Rate (FDR) correction is applied to the p-values. - + The frequency dimension is reduced to one for the test (average in the frequency band-of-interest). To take frequencies into account, use cluster statistics (see statscondCluster function). - + Examples -------- >>> # Independent t-test between two groups @@ -86,7 +87,7 @@ def statsCond(data: np.ndarray, epochs: mne.Epochs, n_permutations: int, alpha: ... n_permutations=10000, ... alpha=0.05 ... ) - >>> + >>> >>> # 2×2 repeated measures ANOVA (within-subjects design) >>> # Factor 1: Condition (2 levels), Factor 2: Time (2 levels) >>> anova_results = statscluster( @@ -101,53 +102,54 @@ def statsCond(data: np.ndarray, epochs: mne.Epochs, n_permutations: int, alpha: """ # checking whether data have the same size - assert(len(data.shape) == 3), "PSD does not have the appropriate shape!" + assert len(data.shape) == 3, "PSD does not have the appropriate shape!" # averaging across frequencies (compute stats only in ch space) power = np.mean(data, axis=2) - T_obs, p_values, H0 = mne.stats.permutation_t_test(power, n_permutations, - tail=0, n_jobs=1) - adj_p = mne.stats.fdr_correction(p_values, alpha=alpha, method='indep') + T_obs, p_values, H0 = mne.stats.permutation_t_test( + power, n_permutations, tail=0, n_jobs=1 + ) + adj_p = mne.stats.fdr_correction(p_values, alpha=alpha, method="indep") T_obs_plot = np.where(adj_p[0], T_obs, 0) # retrieving sensor position pos = np.array([[0, 0]]) - for i in range(0, len(epochs.info['ch_names'])): - cor = np.array([epochs.info['chs'][i]['loc'][0:2]]) + for i in range(0, len(epochs.info["ch_names"])): + cor = np.array([epochs.info["chs"][i]["loc"][0:2]]) pos = np.concatenate((pos, cor), axis=0) pos = pos[1:] statsCondTuple = namedtuple( - 'statsCond', ['T_obs', 'p_values', 'H0', 'adj_p', 'T_obs_plot']) + "statsCond", ["T_obs", "p_values", "H0", "adj_p", "T_obs_plot"] + ) return statsCondTuple( - T_obs=T_obs, - p_values=p_values, - H0=H0, - adj_p=adj_p, - T_obs_plot=T_obs_plot) + T_obs=T_obs, p_values=p_values, H0=H0, adj_p=adj_p, T_obs_plot=T_obs_plot + ) -def con_matrix(epochs: mne.Epochs, freqs_mean: List[float], draw: bool = False) -> namedtuple: +def con_matrix( + epochs: mne.Epochs, freqs_mean: List[float], draw: bool = False +) -> namedtuple: """ Compute a priori channel connectivity across space and frequencies. - + This function creates connectivity matrices that define which channels and frequencies should be considered neighbors for cluster-based statistics. - + Parameters ---------- epochs : mne.Epochs Epochs object containing channel information - + freqs_mean : List[float] List of frequencies in the frequency-band-of-interest used for power or coherence spectral density calculation - + draw : bool, optional Whether to plot the connectivity matrices (default=False) - + Returns ------- con_matrixTuple : namedtuple @@ -157,20 +159,20 @@ def con_matrix(epochs: mne.Epochs, freqs_mean: List[float], draw: bool = False) - ch_con_freq: Connectivity matrix between channels across space and frequencies, scipy.sparse.csr_matrix of shape (n_channels*len(freqs_mean), n_channels*len(freqs_mean)) - + Notes ----- The channel connectivity matrix (ch_con) is based on the spatial adjacency of EEG electrodes - channels that are physically adjacent are considered connected. - + The frequency-space connectivity matrix (ch_con_freq) extends this spatial adjacency to include frequency adjacency - neighboring frequencies for the same channel are also considered connected. - + These connectivity matrices are used as inputs to cluster-based statistical functions to define the neighborhood structure for clustering. - + Examples -------- >>> # Create connectivity matrices for alpha band frequencies @@ -181,8 +183,7 @@ def con_matrix(epochs: mne.Epochs, freqs_mean: List[float], draw: bool = False) """ # creating channel-to-channel connectivity matrix in space - ch_con, ch_names_con = find_ch_adjacency(epochs.info, - ch_type='eeg') + ch_con, ch_names_con = find_ch_adjacency(epochs.info, ch_type="eeg") ch_con_arr = ch_con.toarray() @@ -190,11 +191,14 @@ def con_matrix(epochs: mne.Epochs, freqs_mean: List[float], draw: bool = False) # to take channel connectivity across neighboring frequencies into # account l_freq = len(freqs_mean) - init = np.zeros((l_freq*len(ch_names_con), - l_freq*len(ch_names_con))) - for i in range(0, l_freq*len(ch_names_con)): - for p in range(0, l_freq*len(ch_names_con)): - if (p//len(ch_names_con) == i//len(ch_names_con)) or (p//len(ch_names_con) == i//len(ch_names_con) + 1) or (p//len(ch_names_con) == i//len(ch_names_con) - 1): + init = np.zeros((l_freq * len(ch_names_con), l_freq * len(ch_names_con))) + for i in range(0, l_freq * len(ch_names_con)): + for p in range(0, l_freq * len(ch_names_con)): + if ( + (p // len(ch_names_con) == i // len(ch_names_con)) + or (p // len(ch_names_con) == i // len(ch_names_con) + 1) + or (p // len(ch_names_con) == i // len(ch_names_con) - 1) + ): init[i][p] = 1 ch_con_mult = np.tile(ch_con_arr, (l_freq, l_freq)) @@ -210,37 +214,39 @@ def con_matrix(epochs: mne.Epochs, freqs_mean: List[float], draw: bool = False) plt.spy(ch_con_freq) plt.title("Meta-connectivity matrix") - con_matrixTuple = namedtuple('con_matrix', ['ch_con', 'ch_con_freq']) + con_matrixTuple = namedtuple("con_matrix", ["ch_con", "ch_con_freq"]) - return con_matrixTuple( - ch_con=ch_con, - ch_con_freq=ch_con_freq) + return con_matrixTuple(ch_con=ch_con, ch_con_freq=ch_con_freq) -def metaconn_matrix_2brains(electrodes: List[Tuple[int, int]], ch_con: scipy.sparse.csr_matrix, - freqs_mean: List[float], plot: bool = False) -> namedtuple: +def metaconn_matrix_2brains( + electrodes: List[Tuple[int, int]], + ch_con: scipy.sparse.csr_matrix, + freqs_mean: List[float], + plot: bool = False, +) -> namedtuple: """ Compute a priori connectivity matrices for hyperscanning analyses. - + This function creates connectivity matrices for pairs of channels across two brains (participants), taking into account spatial adjacency within each brain but assuming no direct connectivity between brains. - + Parameters ---------- electrodes : List[Tuple[int, int]] List of electrode pairs for which connectivity indices have been computed. Each tuple contains the indices of channels from participant 1 and participant 2. - + ch_con : scipy.sparse.csr_matrix Connectivity matrix between channels in space, typically from con_matrix() - + freqs_mean : List[float] List of frequencies in the frequency-band-of-interest - + plot : bool, optional Whether to plot the connectivity matrices (default=False) - + Returns ------- metaconn_matrix_2brainsTuple : namedtuple @@ -250,17 +256,17 @@ def metaconn_matrix_2brains(electrodes: List[Tuple[int, int]], ch_con: scipy.spa - metaconn_freq: Connectivity matrix between channel pairs across space and frequencies, array of shape (len(electrodes)*len(freqs_mean), len(electrodes)*len(freqs_mean)) - + Notes ----- This function assumes there is no a priori connectivity between channels from different participants. It considers two channel pairs connected if: 1. The respective channels within each participant are connected, or 2. Some channels are identical across the pairs - + The resulting connectivity matrices define the neighborhood structure for cluster-based statistics on hyperscanning data. - + Examples -------- >>> # Create metaconnectivity matrices for interbrain connectivity @@ -270,27 +276,32 @@ def metaconn_matrix_2brains(electrodes: List[Tuple[int, int]], ch_con: scipy.spa ... ) """ - n = np.max(electrodes, axis=0)[0]+1 + n = np.max(electrodes, axis=0)[0] + 1 # n = 62 metaconn = np.zeros((len(electrodes), len(electrodes))) for ne1, (e11, e12) in enumerate(electrodes): for ne2, (e21, e22) in enumerate(electrodes): # print(ne1,e11,e12,ne2,e21,e22) # considering no a priori connectivity between the 2 brains - metaconn[ne1, ne2] = (((ch_con[e11, e21]) and (ch_con[e12-n, e22-n])) or - ((ch_con[e11, e21]) and (e12 == e22)) or - ((ch_con[e12-n, e22-n]) and (e11 == e21)) or - ((e12 == e22) and (e11 == e21))) + metaconn[ne1, ne2] = ( + ((ch_con[e11, e21]) and (ch_con[e12 - n, e22 - n])) + or ((ch_con[e11, e21]) and (e12 == e22)) + or ((ch_con[e12 - n, e22 - n]) and (e11 == e21)) + or ((e12 == e22) and (e11 == e21)) + ) # duplicating the array 'freqs_mean' times to take channel connectivity # across neighboring frequencies into account l_freq = len(freqs_mean) - init = np.zeros((l_freq*len(electrodes), - l_freq*len(electrodes))) - for i in range(0, l_freq*len(electrodes)): - for p in range(0, l_freq*len(electrodes)): - if (p//len(electrodes) == i//len(electrodes)) or (p//len(electrodes) == i//len(electrodes) + 1) or (p//len(electrodes) == i//len(electrodes) - 1): + init = np.zeros((l_freq * len(electrodes), l_freq * len(electrodes))) + for i in range(0, l_freq * len(electrodes)): + for p in range(0, l_freq * len(electrodes)): + if ( + (p // len(electrodes) == i // len(electrodes)) + or (p // len(electrodes) == i // len(electrodes) + 1) + or (p // len(electrodes) == i // len(electrodes) - 1) + ): init[i][p] = 1 metaconn_mult = np.tile(metaconn, (l_freq, l_freq)) @@ -303,33 +314,35 @@ def metaconn_matrix_2brains(electrodes: List[Tuple[int, int]], ch_con: scipy.spa plt.title("Meta-connectivity matrix") metaconn_matrix_2brainsTuple = namedtuple( - 'metaconn_matrix_2brains', ['metaconn', 'metaconn_freq']) + "metaconn_matrix_2brains", ["metaconn", "metaconn_freq"] + ) - return metaconn_matrix_2brainsTuple( - metaconn=metaconn, - metaconn_freq=metaconn_freq) + return metaconn_matrix_2brainsTuple(metaconn=metaconn, metaconn_freq=metaconn_freq) -def metaconn_matrix(electrodes: List[Tuple[int, int]], ch_con: scipy.sparse.csr_matrix, - freqs_mean: List[float]) -> namedtuple: +def metaconn_matrix( + electrodes: List[Tuple[int, int]], + ch_con: scipy.sparse.csr_matrix, + freqs_mean: List[float], +) -> namedtuple: """ Compute a priori connectivity between pairs of sensors within one brain. - + This function creates connectivity matrices for pairs of channels within a single brain, taking into account spatial adjacency for cluster-based statistics. - + Parameters ---------- electrodes : List[Tuple[int, int]] List of electrode pairs for which connectivity indices have been computed. Each tuple contains the indices of two channels from the same participant. - + ch_con : scipy.sparse.csr_matrix Connectivity matrix between channels in space, typically from con_matrix() - + freqs_mean : List[float] List of frequencies in the frequency-band-of-interest - + Returns ------- metaconn_matrixTuple : namedtuple @@ -339,16 +352,16 @@ def metaconn_matrix(electrodes: List[Tuple[int, int]], ch_con: scipy.sparse.csr_ - metaconn_freq: Connectivity matrix between channel pairs across space and frequencies, array of shape (len(electrodes)*len(freqs_mean), len(electrodes)*len(freqs_mean)) - + Notes ----- This function determines whether two channel pairs are connected based on - the spatial adjacency of their constituent channels. It considers various + the spatial adjacency of their constituent channels. It considers various combinations of adjacency between the channels. - + The resulting connectivity matrices define the neighborhood structure for cluster-based statistics on connectivity data within a single brain. - + Examples -------- >>> # Create metaconnectivity matrices for intrabrain connectivity @@ -362,22 +375,27 @@ def metaconn_matrix(electrodes: List[Tuple[int, int]], ch_con: scipy.sparse.csr_ for ne1, (e11, e12) in enumerate(electrodes): for ne2, (e21, e22) in enumerate(electrodes): # print(ne1,e11,e12,ne2,e21,e22) - metaconn[ne1, ne2] = (((ch_con[e11, e21]) and (ch_con[e12, e22])) or - ((ch_con[e11, e22]) and (ch_con[e12, e21])) or - ((ch_con[e11, e21]) and (e12 == e22)) or - ((ch_con[e11, e22]) and (e12 == e21)) or - ((ch_con[e12, e21]) and (e11 == e22)) or - ((ch_con[e12, e22]) and (e11 == e21))) + metaconn[ne1, ne2] = ( + ((ch_con[e11, e21]) and (ch_con[e12, e22])) + or ((ch_con[e11, e22]) and (ch_con[e12, e21])) + or ((ch_con[e11, e21]) and (e12 == e22)) + or ((ch_con[e11, e22]) and (e12 == e21)) + or ((ch_con[e12, e21]) and (e11 == e22)) + or ((ch_con[e12, e22]) and (e11 == e21)) + ) # duplicating the array 'freqs_mean' times to take channels connectivity # across neighboring frequencies into account l_freq = len(freqs_mean) - init = np.zeros((l_freq*len(electrodes), - l_freq*len(electrodes))) - for i in range(0, l_freq*len(electrodes)): - for p in range(0, l_freq*len(electrodes)): - if (p//len(electrodes) == i//len(electrodes)) or (p//len(electrodes) == i//len(electrodes) + 1) or (p//len(electrodes) == i//len(electrodes) - 1): + init = np.zeros((l_freq * len(electrodes), l_freq * len(electrodes))) + for i in range(0, l_freq * len(electrodes)): + for p in range(0, l_freq * len(electrodes)): + if ( + (p // len(electrodes) == i // len(electrodes)) + or (p // len(electrodes) == i // len(electrodes) + 1) + or (p // len(electrodes) == i // len(electrodes) - 1) + ): init[i][p] = 1 metaconn_mult = np.tile(metaconn, (l_freq, l_freq)) @@ -387,48 +405,51 @@ def metaconn_matrix(electrodes: List[Tuple[int, int]], ch_con: scipy.sparse.csr_ # vizualising the array plt.spy(metaconn_freq) - metaconn_matrixTuple = namedtuple( - 'metaconn_matrix', ['metaconn', 'metaconn_freq']) + metaconn_matrixTuple = namedtuple("metaconn_matrix", ["metaconn", "metaconn_freq"]) - return metaconn_matrixTuple( - metaconn=metaconn, - metaconn_freq=metaconn_freq) + return metaconn_matrixTuple(metaconn=metaconn, metaconn_freq=metaconn_freq) -def statscondCluster(data: list, freqs_mean: list, ch_con_freq: scipy.sparse.csr_matrix, - tail: int, n_permutations: int, alpha: float = 0.05) -> namedtuple: +def statscondCluster( + data: list, + freqs_mean: list, + ch_con_freq: scipy.sparse.csr_matrix, + tail: int, + n_permutations: int, + alpha: float = 0.05, +) -> namedtuple: """ Perform cluster-level statistical permutation test on neurophysiological data. - + This function applies a cluster-based permutation test to identify significant differences between conditions, correcting for multiple comparisons by taking into account connectivity across space and frequencies. - + Parameters ---------- data : list List of arrays containing values from different conditions or groups to compare. Each array has shape (n_observations, n_features) - + freqs_mean : list Frequencies in the frequency-band-of-interest - + ch_con_freq : scipy.sparse.csr_matrix Connectivity or metaconnectivity matrix defining adjacency across space and frequencies, typically from con_matrix() or metaconn_matrix() - + tail : int Direction of the test: - 0: two-tailed test - 1: one-tailed test (greater) - -1: one-tailed test (less) - + n_permutations : int Number of permutations for the statistical test, e.g., 50000 - + alpha : float, optional Significance threshold for clusters (default=0.05) - + Returns ------- statscondClusterTuple : namedtuple @@ -438,17 +459,17 @@ def statscondCluster(data: list, freqs_mean: list, ch_con_freq: scipy.sparse.csr - cluster_p_values: p-value for each identified cluster - H0: Max cluster-level statistics under permutation, array of shape (n_permutations) - F_obs_plot: F-values for significant sensors, array of shape (n_features) - + Notes ----- This function uses MNE's permutation_cluster_test to perform the analysis. - + With t_power=1 (default), each location is weighted by its statistical score within a cluster, which gives more weight to stronger effects. - + The function automatically calculates appropriate thresholds for cluster formation based on the F-distribution for two-tailed tests. - + Examples -------- >>> # Compare alpha power between two conditions @@ -473,13 +494,17 @@ def statscondCluster(data: list, freqs_mean: list, ch_con_freq: scipy.sparse.csr threshold = f_dist.ppf(1 - alpha, dfn, dfd) else: threshold = None # One-tailed test uses MNE's default - + # computing the cluster permutation t test - F_obs, clusters, cluster_p_values, H0 = permutation_cluster_test(data, - threshold=threshold, - n_permutations=n_permutations, - tail=tail, adjacency=ch_con_freq, - t_power=1, out_type='mask') + F_obs, clusters, cluster_p_values, H0 = permutation_cluster_test( + data, + threshold=threshold, + n_permutations=n_permutations, + tail=tail, + adjacency=ch_con_freq, + t_power=1, + out_type="mask", + ) # t_power = 1 weighs each location by its statistical score, # when set to 0 it gives a count of locations in each cluster @@ -487,67 +512,78 @@ def statscondCluster(data: list, freqs_mean: list, ch_con_freq: scipy.sparse.csr F_obs_plot = np.zeros(F_obs.shape) for cluster_p in cluster_p_values: if cluster_p <= alpha: - sensors_plot = clusters[np.where(cluster_p_values == cluster_p)[ - 0][0]].astype('uint8') - F_values = sensors_plot*F_obs + sensors_plot = clusters[ + np.where(cluster_p_values == cluster_p)[0][0] + ].astype("uint8") + F_values = sensors_plot * F_obs # taking maximum F value if a sensor belongs to many clusters F_obs_plot = np.maximum(F_obs_plot, F_values) - statscondClusterTuple = namedtuple('statscondCluster', [ - 'F_obs', 'clusters', 'cluster_p_values', 'H0', 'F_obs_plot']) + statscondClusterTuple = namedtuple( + "statscondCluster", + ["F_obs", "clusters", "cluster_p_values", "H0", "F_obs_plot"], + ) return statscondClusterTuple( F_obs=F_obs, clusters=clusters, cluster_p_values=cluster_p_values, H0=H0, - F_obs_plot=F_obs_plot) - - -def statscluster(data: list, test: str, factor_levels: List[int], ch_con_freq: scipy.sparse.csr_matrix, - tail: int, n_permutations: int, alpha: float = 0.05) -> namedtuple: + F_obs_plot=F_obs_plot, + ) + + +def statscluster( + data: list, + test: str, + factor_levels: List[int], + ch_con_freq: scipy.sparse.csr_matrix, + tail: int, + n_permutations: int, + alpha: float = 0.05, +) -> namedtuple: """ Perform cluster-based statistical tests with various test statistics. - + This function provides a flexible interface to cluster-based permutation tests, supporting different statistical tests including t-tests, one-way ANOVA, and multiple-way ANOVA for complex experimental designs. - + Parameters ---------- data : list or np.ndarray For t-tests and one-way ANOVA: list of arrays containing values from different groups or conditions to compare For multiple-way ANOVA: numpy array organized by factors - + test : str Type of statistical test to use: - 'ind ttest': Independent samples t-test - 'rel ttest': Related (paired) samples t-test - 'f oneway': One-way ANOVA - 'f multipleway': Multiple-way ANOVA - + factor_levels : List[int] or None For multiple-way ANOVA, list specifying the number of levels for each factor (e.g., [2, 3] for a 2×3 design with 2 levels of factor 1 and 3 levels of factor 2) Set to None for other tests - + ch_con_freq : scipy.sparse.csr_matrix Connectivity or metaconnectivity matrix defining adjacency across space and frequencies - + tail : int Direction of the test: - 0: two-tailed test (must be used for f oneway) - 1: one-tailed test (greater) (must be used for f multipleway) - -1: one-tailed test (less) - + n_permutations : int Number of permutations for the statistical test, e.g., 50000 - + alpha : float, optional Significance threshold for clusters (default=0.05) - + Returns ------- statscondClusterTuple : namedtuple @@ -557,19 +593,19 @@ def statscluster(data: list, test: str, factor_levels: List[int], ch_con_freq: s - cluster_p_values: p-value for each identified cluster - H0: Max cluster-level statistics under permutation - Stat_obs_plot: Statistical values for significant sensors - + Notes ----- For multiple-way ANOVA with connectivity values, the last dimensions may need to be flattened from shape (n_sensors, n_sensors) to a vector using np.reshape. - + The function applies different thresholding approaches based on the test type: - For t-tests: Uses the alpha value directly - For one-way ANOVA: Uses the alpha value directly - For multiple-way ANOVA: Calculates appropriate F-thresholds based on factor levels - + With t_power=1, each location is weighted by its statistical score within a cluster. - + Examples -------- >>> # Independent t-test between two groups @@ -582,7 +618,7 @@ def statscluster(data: list, test: str, factor_levels: List[int], ch_con_freq: s ... n_permutations=10000, ... alpha=0.05 ... ) - >>> + >>> >>> # 2×2 repeated measures ANOVA (within-subjects design) >>> # Factor 1: Condition (2 levels), Factor 2: Time (2 levels) >>> anova_results = statscluster( @@ -597,65 +633,83 @@ def statscluster(data: list, test: str, factor_levels: List[int], ch_con_freq: s """ # type of test - if test == 'ind ttest': + if test == "ind ttest": + def stat_fun(*arg): - return(scipy.stats.ttest_ind(arg[0], arg[1], equal_var=False)[0]) + return scipy.stats.ttest_ind(arg[0], arg[1], equal_var=False)[0] + df = len(data[0]) + len(data[1]) - 2 p = alpha / 2 if tail == 0 else alpha threshold = scipy.stats.t.ppf(1 - p, df) - elif test == 'rel ttest': + elif test == "rel ttest": + def stat_fun(*arg): - return(scipy.stats.ttest_rel(arg[0], arg[1])[0]) + return scipy.stats.ttest_rel(arg[0], arg[1])[0] + df = len(data[0]) - 1 p = alpha / 2 if tail == 0 else alpha threshold = scipy.stats.t.ppf(1 - p, df) - elif test == 'f oneway': + elif test == "f oneway": + def stat_fun(*arg): - return(scipy.stats.f_oneway(arg[0], arg[1])[0]) + return scipy.stats.f_oneway(arg[0], arg[1])[0] + dfn = len(data) - 1 dfd = sum(len(d) for d in data) - len(data) threshold = f_dist.ppf(1 - alpha, dfn, dfd) - elif test == 'f multipleway': + elif test == "f multipleway": if max(factor_levels) > 2: correction = True else: correction = False + def stat_fun(*arg): - return(mne.stats.f_mway_rm(np.swapaxes(arg, 1, 0), - factor_levels, - effects='all', - correction=correction, - return_pvals=False)[0]) - threshold = mne.stats.f_threshold_mway_rm(n_subjects=data.shape[1], - factor_levels=factor_levels, - effects='all', - pvalue=alpha) + return mne.stats.f_mway_rm( + np.swapaxes(arg, 1, 0), + factor_levels, + effects="all", + correction=correction, + return_pvals=False, + )[0] + + threshold = mne.stats.f_threshold_mway_rm( + n_subjects=data.shape[1], + factor_levels=factor_levels, + effects="all", + pvalue=alpha, + ) # computing the cluster permutation t test - Stat_obs, clusters, cluster_p_values, H0 = permutation_cluster_test(data, - stat_fun=stat_fun, - threshold=threshold, - tail=tail, - n_permutations=n_permutations, - adjacency=ch_con_freq, - t_power=1, - out_type='mask') + Stat_obs, clusters, cluster_p_values, H0 = permutation_cluster_test( + data, + stat_fun=stat_fun, + threshold=threshold, + tail=tail, + n_permutations=n_permutations, + adjacency=ch_con_freq, + t_power=1, + out_type="mask", + ) # getting F values for sensors belonging to a significant cluster Stat_obs_plot = np.zeros(Stat_obs.shape) for cluster_p in cluster_p_values: if cluster_p <= alpha: - sensors_plot = clusters[np.where(cluster_p_values == cluster_p)[ - 0][0]].astype('uint8') - Stat_values = sensors_plot*Stat_obs + sensors_plot = clusters[ + np.where(cluster_p_values == cluster_p)[0][0] + ].astype("uint8") + Stat_values = sensors_plot * Stat_obs # taking maximum statistical value if a sensor is in many clusters Stat_obs_plot = np.maximum(Stat_obs_plot, Stat_values) - statscondClusterTuple = namedtuple('statscondCluster', [ - 'Stat_obs', 'clusters', 'cluster_p_values', 'H0', 'Stat_obs_plot']) + statscondClusterTuple = namedtuple( + "statscondCluster", + ["Stat_obs", "clusters", "cluster_p_values", "H0", "Stat_obs_plot"], + ) return statscondClusterTuple( Stat_obs=Stat_obs, clusters=clusters, cluster_p_values=cluster_p_values, H0=H0, - Stat_obs_plot=Stat_obs_plot) + Stat_obs_plot=Stat_obs_plot, + ) diff --git a/hypyp/sync/README.md b/hypyp/sync/README.md index 5dd6b4b..4bb35a3 100644 --- a/hypyp/sync/README.md +++ b/hypyp/sync/README.md @@ -205,7 +205,7 @@ benchmark data compiled from Mac M4 Max (131 runs) and Narval A100 (111 runs). The priority can be overridden per-call: ```python -get_metric('plv', optimization='auto', priority=['torch', 'cuda_kernel']) +get_metric("plv", optimization="auto", priority=["torch", "cuda_kernel"]) ``` If no GPU backend is available, `'auto'` falls back to numba, then numpy. @@ -275,20 +275,22 @@ pip install "hypyp[cupy]" from hypyp.analyses import compute_sync # Standard (numpy) -con = compute_sync(complex_signal, 'plv') +con = compute_sync(complex_signal, "plv") # Best available GPU backend -con = compute_sync(complex_signal, 'plv', optimization='auto') +con = compute_sync(complex_signal, "plv", optimization="auto") # Specific backend -con = compute_sync(complex_signal, 'pli', optimization='metal') +con = compute_sync(complex_signal, "pli", optimization="metal") # Custom priority -con = compute_sync(complex_signal, 'coh', optimization='auto', - priority=['torch', 'cuda_kernel']) +con = compute_sync( + complex_signal, "coh", optimization="auto", priority=["torch", "cuda_kernel"] +) # Direct class instantiation from hypyp.sync import get_metric -metric = get_metric('accorr', optimization='auto') + +metric = get_metric("accorr", optimization="auto") con = metric.compute(complex_signal_internal, n_samp, transpose_axes) ``` diff --git a/hypyp/sync/__init__.py b/hypyp/sync/__init__.py index aeca197..1b0b419 100644 --- a/hypyp/sync/__init__.py +++ b/hypyp/sync/__init__.py @@ -44,8 +44,12 @@ from typing import Optional from .base import ( - BaseMetric, multiply_conjugate, multiply_conjugate_time, multiply_product, - multiply_conjugate_torch, multiply_conjugate_time_torch, + BaseMetric, + multiply_conjugate, + multiply_conjugate_time, + multiply_product, + multiply_conjugate_torch, + multiply_conjugate_time_torch, ) from .plv import PLV from .ccorr import CCorr @@ -59,43 +63,44 @@ # Metric registry: maps mode names to metric classes METRICS = { - 'plv': PLV, - 'ccorr': CCorr, - 'accorr': ACCorr, - 'coh': Coh, - 'imcoh': ImCoh, - 'pli': PLI, - 'wpli': WPLI, - 'envcorr': EnvCorr, - 'powcorr': PowCorr, + "plv": PLV, + "ccorr": CCorr, + "accorr": ACCorr, + "coh": Coh, + "imcoh": ImCoh, + "pli": PLI, + "wpli": WPLI, + "envcorr": EnvCorr, + "powcorr": PowCorr, } __all__ = [ # Base classes and utilities - 'BaseMetric', - 'multiply_conjugate', - 'multiply_conjugate_time', - 'multiply_product', - 'multiply_conjugate_torch', - 'multiply_conjugate_time_torch', + "BaseMetric", + "multiply_conjugate", + "multiply_conjugate_time", + "multiply_product", + "multiply_conjugate_torch", + "multiply_conjugate_time_torch", # Metric classes - 'PLV', - 'CCorr', - 'ACCorr', - 'Coh', - 'ImCoh', - 'PLI', - 'WPLI', - 'EnvCorr', - 'PowCorr', + "PLV", + "CCorr", + "ACCorr", + "Coh", + "ImCoh", + "PLI", + "WPLI", + "EnvCorr", + "PowCorr", # Registry - 'METRICS', - 'get_metric', + "METRICS", + "get_metric", ] -def get_metric(mode: str, optimization: Optional[str] = None, - priority: Optional[list] = None) -> BaseMetric: +def get_metric( + mode: str, optimization: Optional[str] = None, priority: Optional[list] = None +) -> BaseMetric: """ Get a connectivity metric instance by name. @@ -130,7 +135,7 @@ def get_metric(mode: str, optimization: Optional[str] = None, """ mode_lower = mode.lower() if mode_lower not in METRICS: - available = ', '.join(METRICS.keys()) + available = ", ".join(METRICS.keys()) raise ValueError(f"Unknown metric mode '{mode}'. Available: {available}") return METRICS[mode_lower](optimization=optimization, priority=priority) diff --git a/hypyp/sync/accorr.py b/hypyp/sync/accorr.py index 835b9d1..dd83c2d 100644 --- a/hypyp/sync/accorr.py +++ b/hypyp/sync/accorr.py @@ -38,8 +38,12 @@ from tqdm import tqdm from .base import ( - BaseMetric, multiply_conjugate, multiply_product, - TORCH_AVAILABLE, MPS_AVAILABLE, NUMBA_AVAILABLE, + BaseMetric, + multiply_conjugate, + multiply_product, + TORCH_AVAILABLE, + MPS_AVAILABLE, + NUMBA_AVAILABLE, ) # Conditional imports for optional backends @@ -76,14 +80,18 @@ class ACCorr(BaseMetric): name = "accorr" - def __init__(self, optimization: Optional[str] = None, - priority: Optional[list] = None, - show_progress: bool = True): + def __init__( + self, + optimization: Optional[str] = None, + priority: Optional[list] = None, + show_progress: bool = True, + ): super().__init__(optimization, priority) self.show_progress = show_progress - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Adjusted Circular Correlation. @@ -101,31 +109,36 @@ def compute(self, complex_signal: np.ndarray, n_samp: int, con : np.ndarray ACCorr connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'metal': + if self._backend == "metal": return self._compute_metal(complex_signal, n_samp, transpose_axes) - elif self._backend == 'cuda_kernel': + elif self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) else: return self._compute_numpy(complex_signal, n_samp, transpose_axes) - def _compute_metal(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_metal( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """Metal compute shader for ACCorr on Apple Silicon GPU.""" from .kernels.metal_accorr import accorr_metal + return accorr_metal(complex_signal) - def _compute_cuda(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_cuda( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """CUDA kernel for ACCorr on NVIDIA GPU.""" from .kernels.cuda_accorr import accorr_cuda + return accorr_cuda(complex_signal) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ NumPy implementation of ACCorr with precompute optimization. @@ -157,8 +170,12 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, den = np.zeros((n_epochs, n_freq, n_ch_total, n_ch_total)) total_pairs = (n_ch_total * (n_ch_total + 1)) // 2 - pbar = tqdm(total=total_pairs, desc=" accorr (denominator)", - disable=not self.show_progress, leave=False) + pbar = tqdm( + total=total_pairs, + desc=" accorr (denominator)", + disable=not self.show_progress, + leave=False, + ) for i in range(n_ch_total): for j in range(i, n_ch_total): @@ -171,7 +188,9 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, x_sin = np.sin(alpha1 - m_adj) y_sin = np.sin(alpha2 - n_adj) - den_ij = 2 * np.sqrt(np.sum(x_sin**2, axis=2) * np.sum(y_sin**2, axis=2)) + den_ij = 2 * np.sqrt( + np.sum(x_sin**2, axis=2) * np.sum(y_sin**2, axis=2) + ) den[:, :, i, j] = den_ij den[:, :, j, i] = den_ij @@ -184,8 +203,9 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba-optimized implementation of ACCorr with precompute. @@ -214,8 +234,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, # Denominator via numba JIT angle = np.angle(complex_signal) - den = _accorr_den_numba(n_epochs, n_freq, n_ch_total, angle, - m_adj_all, n_adj_all) + den = _accorr_den_numba( + n_epochs, n_freq, n_ch_total, angle, m_adj_all, n_adj_all + ) den = np.where(den == 0, 1, den) con = num / den @@ -241,8 +262,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, value is empirical; re-derive if you change the upstream tensor layout. """ - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of ACCorr with precompute optimization. @@ -255,7 +277,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, """ device = self._device - if device == 'mps': + if device == "mps": float_type = torch.float32 complex_type = torch.complex64 bytes_per_elem = 4 @@ -264,7 +286,9 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, complex_type = torch.complex128 bytes_per_elem = 8 - complex_tensor = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) + complex_tensor = torch.from_numpy(complex_signal).to( + device=device, dtype=complex_type + ) n_epochs, n_freq, n_ch_total, n_times = complex_tensor.shape # Numerator (vectorized with torch einsum) @@ -272,7 +296,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, c, s = z.real, z.imag # Factorized: 4 einsum shared between cross_conj and cross_prod - formula = 'efit,efjt->efij' + formula = "efit,efjt->efij" cc = torch.einsum(formula, c, c) ss = torch.einsum(formula, s, s) cs = torch.einsum(formula, c, s) @@ -293,14 +317,24 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Denominator — choose vectorized or loop based on memory angle = torch.angle(complex_tensor) - tensor_5d_bytes = n_epochs * n_freq * n_ch_total * n_ch_total * n_times * bytes_per_elem + tensor_5d_bytes = ( + n_epochs * n_freq * n_ch_total * n_ch_total * n_times * bytes_per_elem + ) use_vectorized = tensor_5d_bytes < self._VRAM_THRESHOLD if use_vectorized: den = self._den_vectorized(angle, m_adj_all, n_adj_all, device, float_type) else: - den = self._den_loop(angle, m_adj_all, n_adj_all, device, float_type, - n_epochs, n_freq, n_ch_total) + den = self._den_loop( + angle, + m_adj_all, + n_adj_all, + device, + float_type, + n_epochs, + n_freq, + n_ch_total, + ) den = torch.where(den == 0, torch.ones_like(den), den) con = num / den @@ -323,25 +357,39 @@ def _den_vectorized(self, angle, m_adj_all, n_adj_all, device, float_type): x_sin = torch.sin(angle.unsqueeze(3) - m_adj_all.unsqueeze(-1)) # (E,F,C,C,T) y_sin = torch.sin(angle.unsqueeze(2) - n_adj_all.unsqueeze(-1)) # (E,F,C,C,T) - sum_x2 = torch.sum(x_sin ** 2, dim=-1) # (E, F, C, C) - sum_y2 = torch.sum(y_sin ** 2, dim=-1) # (E, F, C, C) + sum_x2 = torch.sum(x_sin**2, dim=-1) # (E, F, C, C) + sum_y2 = torch.sum(y_sin**2, dim=-1) # (E, F, C, C) return 2.0 * torch.sqrt(sum_x2 * sum_y2) - def _den_loop(self, angle, m_adj_all, n_adj_all, device, float_type, - n_epochs, n_freq, n_ch_total): + def _den_loop( + self, + angle, + m_adj_all, + n_adj_all, + device, + float_type, + n_epochs, + n_freq, + n_ch_total, + ): """ Loop-based denominator (fallback for large data). Iterates over channel pairs when the vectorized 5D tensor would exceed the VRAM threshold. """ - den = torch.zeros((n_epochs, n_freq, n_ch_total, n_ch_total), - device=device, dtype=float_type) + den = torch.zeros( + (n_epochs, n_freq, n_ch_total, n_ch_total), device=device, dtype=float_type + ) total_pairs = (n_ch_total * (n_ch_total + 1)) // 2 - pbar = tqdm(total=total_pairs, desc=" accorr_torch (denominator)", - disable=not self.show_progress, leave=False) + pbar = tqdm( + total=total_pairs, + desc=" accorr_torch (denominator)", + disable=not self.show_progress, + leave=False, + ) for i in range(n_ch_total): for j in range(i, n_ch_total): @@ -354,7 +402,9 @@ def _den_loop(self, angle, m_adj_all, n_adj_all, device, float_type, x_sin = torch.sin(alpha1 - m_adj) y_sin = torch.sin(alpha2 - n_adj) - den_ij = 2 * torch.sqrt(torch.sum(x_sin**2, dim=2) * torch.sum(y_sin**2, dim=2)) + den_ij = 2 * torch.sqrt( + torch.sum(x_sin**2, dim=2) * torch.sum(y_sin**2, dim=2) + ) den[:, :, i, j] = den_ij den[:, :, j, i] = den_ij @@ -366,6 +416,7 @@ def _den_loop(self, angle, m_adj_all, n_adj_all, device, float_type, # Numba JIT-compiled helper (defined at module level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _accorr_den_numba(n_epochs, n_freq, n_ch_total, angle, m_adj_all, n_adj_all): """ diff --git a/hypyp/sync/base.py b/hypyp/sync/base.py index 72cfaec..955d3f7 100644 --- a/hypyp/sync/base.py +++ b/hypyp/sync/base.py @@ -52,6 +52,7 @@ # Check optional dependency availability try: import torch + TORCH_AVAILABLE = True MPS_AVAILABLE = torch.backends.mps.is_available() CUDA_AVAILABLE = torch.cuda.is_available() @@ -62,6 +63,7 @@ try: import numba + NUMBA_AVAILABLE = True except ImportError: NUMBA_AVAILABLE = False @@ -95,21 +97,23 @@ AUTO_PRIORITY = { # einsum metrics — torch wins on MPS, cuda_kernel safe-first on CUDA # (torch OOMs at ≥512ch on CUDA; cuda_kernel computes pairwise) - 'plv': {'mps': ['torch'], 'cuda': ['cuda_kernel', 'torch']}, - 'ccorr': {'mps': ['torch'], 'cuda': ['cuda_kernel', 'torch']}, - 'coh': {'mps': ['torch'], 'cuda': ['cuda_kernel', 'torch']}, - 'imcoh': {'mps': ['torch'], 'cuda': ['cuda_kernel', 'torch']}, - 'envcorr': {'mps': ['torch'], 'cuda': ['cuda_kernel', 'torch']}, - 'powcorr': {'mps': ['torch'], 'cuda': ['cuda_kernel', 'torch']}, + "plv": {"mps": ["torch"], "cuda": ["cuda_kernel", "torch"]}, + "ccorr": {"mps": ["torch"], "cuda": ["cuda_kernel", "torch"]}, + "coh": {"mps": ["torch"], "cuda": ["cuda_kernel", "torch"]}, + "imcoh": {"mps": ["torch"], "cuda": ["cuda_kernel", "torch"]}, + "envcorr": {"mps": ["torch"], "cuda": ["cuda_kernel", "torch"]}, + "powcorr": {"mps": ["torch"], "cuda": ["cuda_kernel", "torch"]}, # sign-based — custom kernels beat torch on both platforms - 'pli': {'mps': ['metal', 'torch'], 'cuda': ['cuda_kernel', 'torch']}, - 'wpli': {'mps': ['metal', 'torch'], 'cuda': ['cuda_kernel', 'torch']}, + "pli": {"mps": ["metal", "torch"], "cuda": ["cuda_kernel", "torch"]}, + "wpli": {"mps": ["metal", "torch"], "cuda": ["cuda_kernel", "torch"]}, # accorr — Metal wins on MPS (circular correlation), cuda_kernel safe on CUDA - 'accorr': {'mps': ['metal', 'torch'], 'cuda': ['cuda_kernel', 'torch']}, + "accorr": {"mps": ["metal", "torch"], "cuda": ["cuda_kernel", "torch"]}, } -def multiply_conjugate(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) -> np.ndarray: +def multiply_conjugate( + real: np.ndarray, imag: np.ndarray, transpose_axes: tuple +) -> np.ndarray: """ Computes the product of a complex array and its conjugate efficiently. @@ -131,16 +135,23 @@ def multiply_conjugate(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple ----- Implements: product = (real x real.T + imag x imag.T) - i(real x imag.T - imag x real.T) """ - formula = 'jilm,jimk->jilk' - product = np.einsum(formula, real, real.transpose(transpose_axes)) + \ - np.einsum(formula, imag, imag.transpose(transpose_axes)) - 1j * \ - (np.einsum(formula, real, imag.transpose(transpose_axes)) - \ - np.einsum(formula, imag, real.transpose(transpose_axes))) + formula = "jilm,jimk->jilk" + product = ( + np.einsum(formula, real, real.transpose(transpose_axes)) + + np.einsum(formula, imag, imag.transpose(transpose_axes)) + - 1j + * ( + np.einsum(formula, real, imag.transpose(transpose_axes)) + - np.einsum(formula, imag, real.transpose(transpose_axes)) + ) + ) return product -def multiply_conjugate_time(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) -> np.ndarray: +def multiply_conjugate_time( + real: np.ndarray, imag: np.ndarray, transpose_axes: tuple +) -> np.ndarray: """ Computes the product of a complex array and its conjugate without collapsing time dimension. @@ -161,16 +172,23 @@ def multiply_conjugate_time(real: np.ndarray, imag: np.ndarray, transpose_axes: product : np.ndarray Product of the array and its complex conjugate with time dimension preserved """ - formula = 'jilm,jimk->jilkm' - product = np.einsum(formula, real, real.transpose(transpose_axes)) + \ - np.einsum(formula, imag, imag.transpose(transpose_axes)) - 1j * \ - (np.einsum(formula, real, imag.transpose(transpose_axes)) - \ - np.einsum(formula, imag, real.transpose(transpose_axes))) + formula = "jilm,jimk->jilkm" + product = ( + np.einsum(formula, real, real.transpose(transpose_axes)) + + np.einsum(formula, imag, imag.transpose(transpose_axes)) + - 1j + * ( + np.einsum(formula, real, imag.transpose(transpose_axes)) + - np.einsum(formula, imag, real.transpose(transpose_axes)) + ) + ) return product -def multiply_product(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) -> np.ndarray: +def multiply_product( + real: np.ndarray, imag: np.ndarray, transpose_axes: tuple +) -> np.ndarray: """ Computes the product of two complex arrays (not conjugate) efficiently. @@ -191,11 +209,16 @@ def multiply_product(real: np.ndarray, imag: np.ndarray, transpose_axes: tuple) product : np.ndarray Product of the array with itself (non-conjugate) """ - formula = 'jilm,jimk->jilk' - product = np.einsum(formula, real, real.transpose(transpose_axes)) - \ - np.einsum(formula, imag, imag.transpose(transpose_axes)) + 1j * \ - (np.einsum(formula, real, imag.transpose(transpose_axes)) + \ - np.einsum(formula, imag, real.transpose(transpose_axes))) + formula = "jilm,jimk->jilk" + product = ( + np.einsum(formula, real, real.transpose(transpose_axes)) + - np.einsum(formula, imag, imag.transpose(transpose_axes)) + + 1j + * ( + np.einsum(formula, real, imag.transpose(transpose_axes)) + + np.einsum(formula, imag, real.transpose(transpose_axes)) + ) + ) return product @@ -219,10 +242,12 @@ def multiply_conjugate_torch(c, s): torch.Tensor Complex product, shape (E, F, C, C). """ - formula = 'efit,efjt->efij' + formula = "efit,efjt->efij" import torch - return (torch.einsum(formula, c, c) + torch.einsum(formula, s, s)) - 1j * \ - (torch.einsum(formula, c, s) - torch.einsum(formula, s, c)) + + return (torch.einsum(formula, c, c) + torch.einsum(formula, s, s)) - 1j * ( + torch.einsum(formula, c, s) - torch.einsum(formula, s, c) + ) def multiply_conjugate_time_torch(c, s): @@ -244,10 +269,12 @@ def multiply_conjugate_time_torch(c, s): torch.Tensor Complex product, shape (E, F, C, C, T). """ - formula = 'efit,efjt->efijt' + formula = "efit,efjt->efijt" import torch - return (torch.einsum(formula, c, c) + torch.einsum(formula, s, s)) - 1j * \ - (torch.einsum(formula, c, s) - torch.einsum(formula, s, c)) + + return (torch.einsum(formula, c, c) + torch.einsum(formula, s, s)) - 1j * ( + torch.einsum(formula, c, s) - torch.einsum(formula, s, c) + ) class BaseMetric(ABC): @@ -276,17 +303,17 @@ class BaseMetric(ABC): name: str = "base" - def __init__(self, optimization: Optional[str] = None, - priority: Optional[list] = None): + def __init__( + self, optimization: Optional[str] = None, priority: Optional[list] = None + ): self.optimization = optimization self._priority = priority - self._backend, self._device = self._resolve_optimization( - optimization, priority - ) + self._backend, self._device = self._resolve_optimization(optimization, priority) @classmethod - def _resolve_optimization(cls, optimization: Optional[str] = None, - priority: Optional[list] = None) -> tuple: + def _resolve_optimization( + cls, optimization: Optional[str] = None, priority: Optional[list] = None + ) -> tuple: """ Resolves an optimization value to (backend, device). @@ -335,50 +362,54 @@ def _resolve_optimization(cls, optimization: Optional[str] = None, requested backend → numpy (with UserWarning) """ if optimization is None: - return 'numpy', 'cpu' + return "numpy", "cpu" - if optimization == 'auto': + if optimization == "auto": return cls._resolve_auto(priority) - if optimization == 'numba': + if optimization == "numba": if NUMBA_AVAILABLE: - return 'numba', 'cpu' + return "numba", "cpu" warnings.warn( "numba not installed, falling back to numpy. " "Install with: poetry install --with optim_numba", - UserWarning, stacklevel=3 + UserWarning, + stacklevel=3, ) - return 'numpy', 'cpu' + return "numpy", "cpu" - if optimization == 'torch': + if optimization == "torch": if TORCH_AVAILABLE: return cls._resolve_torch() warnings.warn( "torch not installed, falling back to numpy. " "Install with: poetry install --with optim_torch", - UserWarning, stacklevel=3 + UserWarning, + stacklevel=3, ) - return 'numpy', 'cpu' + return "numpy", "cpu" - if optimization == 'metal': + if optimization == "metal": if METAL_AVAILABLE: - return 'metal', 'mps' + return "metal", "mps" warnings.warn( "PyObjC Metal not available, falling back to numpy. " "Install with: pip install pyobjc-framework-Metal", - UserWarning, stacklevel=3 + UserWarning, + stacklevel=3, ) - return 'numpy', 'cpu' + return "numpy", "cpu" - if optimization == 'cuda_kernel': + if optimization == "cuda_kernel": if CUPY_AVAILABLE: - return 'cuda_kernel', 'cuda' + return "cuda_kernel", "cuda" warnings.warn( "CuPy not available, falling back to numpy. " "Install with: pip install cupy-cuda12x", - UserWarning, stacklevel=3 + UserWarning, + stacklevel=3, ) - return 'numpy', 'cpu' + return "numpy", "cpu" raise ValueError( f"Unknown optimization '{optimization}'. " @@ -412,41 +443,43 @@ def _resolve_auto(cls, priority: Optional[list] = None) -> tuple: On CPU-only machines, warns and falls back to numba → numpy. """ if MPS_AVAILABLE: - platform = 'mps' + platform = "mps" elif CUDA_AVAILABLE: - platform = 'cuda' + platform = "cuda" else: # No GPU — warn and fall back to CPU warnings.warn( "No GPU available. optimization='auto' selects the best GPU " "backend. Use optimization='numba' for CPU parallelism or " "optimization=None for numpy.", - UserWarning, stacklevel=4 + UserWarning, + stacklevel=4, ) if NUMBA_AVAILABLE: - return 'numba', 'cpu' - return 'numpy', 'cpu' + return "numba", "cpu" + return "numpy", "cpu" if priority is None: priority = AUTO_PRIORITY.get(cls.name, {}).get(platform, []) for backend in priority: - if backend == 'torch' and TORCH_AVAILABLE: + if backend == "torch" and TORCH_AVAILABLE: return cls._resolve_torch() - if backend == 'metal' and METAL_AVAILABLE: - return 'metal', 'mps' - if backend == 'cuda_kernel' and CUPY_AVAILABLE: - return 'cuda_kernel', 'cuda' + if backend == "metal" and METAL_AVAILABLE: + return "metal", "mps" + if backend == "cuda_kernel" and CUPY_AVAILABLE: + return "cuda_kernel", "cuda" # No GPU backend from priority list available — fall back warnings.warn( f"No GPU backend available for {cls.name!r} on platform " f"'{platform}'. Falling back to CPU.", - UserWarning, stacklevel=4 + UserWarning, + stacklevel=4, ) if NUMBA_AVAILABLE: - return 'numba', 'cpu' - return 'numpy', 'cpu' + return "numba", "cpu" + return "numpy", "cpu" @staticmethod def _resolve_torch() -> tuple: @@ -474,18 +507,16 @@ def _resolve_torch() -> tuple: should apply a looser tolerance when MPS is the active device. """ if MPS_AVAILABLE: - return 'torch', 'mps' + return "torch", "mps" if CUDA_AVAILABLE: - return 'torch', 'cuda' - warnings.warn( - "No GPU found, using torch on CPU", - UserWarning, stacklevel=4 - ) - return 'torch', 'cpu' + return "torch", "cuda" + warnings.warn("No GPU found, using torch on CPU", UserWarning, stacklevel=4) + return "torch", "cpu" @abstractmethod - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute the connectivity metric. diff --git a/hypyp/sync/ccorr.py b/hypyp/sync/ccorr.py index ceec2c7..b859776 100644 --- a/hypyp/sync/ccorr.py +++ b/hypyp/sync/ccorr.py @@ -28,53 +28,56 @@ class CCorr(BaseMetric): """ Circular Correlation connectivity metric. - + CCorr measures the circular correlation coefficient between the phases of two signals, using the circular mean for phase centering. - + References ---------- Fisher, N. I. (1995). Statistical analysis of circular data. Cambridge University Press. """ - + name = "ccorr" - - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Circular Correlation. - + Parameters ---------- complex_signal : np.ndarray Complex analytic signals with shape (n_epochs, n_freq, 2*n_channels, n_times). - + n_samp : int Number of time samples. - + transpose_axes : tuple Axes to transpose for matrix multiplication. - + Returns ------- con : np.ndarray CCorr connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'cuda_kernel': + if self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) def _compute_cuda(self, complex_signal, n_samp, transpose_axes): """CUDA kernel for CCorr.""" from .kernels.cuda_phase import ccorr_cuda + return ccorr_cuda(complex_signal) - - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ NumPy implementation of CCorr. @@ -95,14 +98,20 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, ).reshape(n_epoch, n_freq, n_ch_total, 1) angle = np.sin(angle - mu_angle) - formula = 'nilm,nimk->nilk' - con = np.abs(np.einsum(formula, angle, angle.transpose(transpose_axes)) / - np.sqrt(np.einsum('nil,nik->nilk', np.sum(angle ** 2, axis=3), - np.sum(angle ** 2, axis=3)))) + formula = "nilm,nimk->nilk" + con = np.abs( + np.einsum(formula, angle, angle.transpose(transpose_axes)) + / np.sqrt( + np.einsum( + "nil,nik->nilk", np.sum(angle**2, axis=3), np.sum(angle**2, axis=3) + ) + ) + ) return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of CCorr using angle-free reformulation. @@ -115,8 +124,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, s = np.imag(phase) return _ccorr_numba_kernel(c, s) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of CCorr using angle-free reformulation. @@ -132,7 +142,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, after the initial phase normalization, improving MPS float32 precision. """ device = self._device - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) @@ -149,16 +159,17 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, d = s * C_bar - c * S_bar # (E, F, C, T) # Correlation via einsum - formula = 'efit,efjt->efij' + formula = "efit,efjt->efij" num = torch.einsum(formula, d, d) - sum_sq = torch.sum(d ** 2, dim=3) - den = torch.sqrt(torch.einsum('efi,efj->efij', sum_sq, sum_sq)) + sum_sq = torch.sum(d**2, dim=3) + den = torch.sqrt(torch.einsum("efi,efj->efij", sum_sq, sum_sq)) con = torch.abs(num / den) return con.cpu().numpy() - def _compute_torch_cpu_circmean(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch_cpu_circmean( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Hybrid approach: circular mean in float64 on CPU, correlation on GPU. @@ -167,24 +178,25 @@ def _compute_torch_cpu_circmean(self, complex_signal: np.ndarray, n_samp: int, Kept for comparison with the angle-free reformulation. """ device = self._device - float_type = torch.float32 if device == 'mps' else torch.float64 + float_type = torch.float32 if device == "mps" else torch.float64 # Step 1: Circular mean in float64 on CPU (precision-critical) angle = np.angle(complex_signal) mu_angle = np.arctan2( np.mean(np.sin(angle), axis=3), np.mean(np.cos(angle), axis=3), - ).reshape(complex_signal.shape[0], complex_signal.shape[1], - complex_signal.shape[2], 1) + ).reshape( + complex_signal.shape[0], complex_signal.shape[1], complex_signal.shape[2], 1 + ) centered = np.sin(angle - mu_angle) # float64, precise # Step 2: Transfer centered signal to GPU for einsum d = torch.from_numpy(centered).to(device=device, dtype=float_type) - formula = 'efit,efjt->efij' + formula = "efit,efjt->efij" num = torch.einsum(formula, d, d) - sum_sq = torch.sum(d ** 2, dim=3) - den = torch.sqrt(torch.einsum('efi,efj->efij', sum_sq, sum_sq)) + sum_sq = torch.sum(d**2, dim=3) + den = torch.sqrt(torch.einsum("efi,efj->efij", sum_sq, sum_sq)) con = torch.abs(num / den) return con.cpu().numpy() @@ -192,6 +204,7 @@ def _compute_torch_cpu_circmean(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _ccorr_numba_kernel(c, s): """ diff --git a/hypyp/sync/coh.py b/hypyp/sync/coh.py index 0fe2698..3cd37bf 100644 --- a/hypyp/sync/coh.py +++ b/hypyp/sync/coh.py @@ -30,68 +30,73 @@ class Coh(BaseMetric): """ Coherence connectivity metric. - + Coherence measures the linear relationship between two signals in the frequency domain, normalized by their power. - + Mathematical formulation: Coh = |⟨XY*⟩|² / (⟨|X|²⟩⟨|Y|²⟩) - + References ---------- Nunez, P. L., & Srinivasan, R. (2006). Electric fields of the brain: the neurophysics of EEG. Oxford University Press. """ - + name = "coh" - - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Coherence. - + Parameters ---------- complex_signal : np.ndarray Complex analytic signals with shape (n_epochs, n_freq, 2*n_channels, n_times). - + n_samp : int Number of time samples. - + transpose_axes : tuple Axes to transpose for matrix multiplication. - + Returns ------- con : np.ndarray Coherence connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'cuda_kernel': + if self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) def _compute_cuda(self, complex_signal, n_samp, transpose_axes): """CUDA kernel for Coherence.""" from .kernels.cuda_amplitude import coh_cuda + return coh_cuda(complex_signal) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of Coherence.""" c = np.real(complex_signal) s = np.imag(complex_signal) amp = np.abs(complex_signal) ** 2 dphi = multiply_conjugate(c, s, transpose_axes=transpose_axes) - con = np.abs(dphi) / np.sqrt(np.einsum('nil,nik->nilk', np.nansum(amp, axis=3), - np.nansum(amp, axis=3))) + con = np.abs(dphi) / np.sqrt( + np.einsum("nil,nik->nilk", np.nansum(amp, axis=3), np.nansum(amp, axis=3)) + ) return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of Coherence with parallel epoch processing. @@ -103,8 +108,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, s = np.imag(complex_signal) return _coh_numba_kernel(c, s) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of Coherence. @@ -113,7 +119,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, MPS uses float32; CPU/CUDA uses float64. """ device = self._device - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) c, s = sig.real, sig.imag @@ -124,7 +130,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Power normalization: sqrt(sum|X_i|² * sum|X_j|²) amp = torch.abs(sig) ** 2 power = torch.nansum(amp, dim=3) - den = torch.sqrt(torch.einsum('efi,efj->efij', power, power)) + den = torch.sqrt(torch.einsum("efi,efj->efij", power, power)) con = torch.abs(dphi) / den return con.cpu().numpy() @@ -132,6 +138,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _coh_numba_kernel(c, s): """ @@ -162,11 +169,17 @@ def _coh_numba_kernel(c, s): re_sum = 0.0 im_sum = 0.0 for t in range(n_t): - re_sum += c[e, f, i, t] * c[e, f, j, t] + s[e, f, i, t] * s[e, f, j, t] - im_sum += s[e, f, i, t] * c[e, f, j, t] - c[e, f, i, t] * s[e, f, j, t] + re_sum += ( + c[e, f, i, t] * c[e, f, j, t] + + s[e, f, i, t] * s[e, f, j, t] + ) + im_sum += ( + s[e, f, i, t] * c[e, f, j, t] + - c[e, f, i, t] * s[e, f, j, t] + ) denom = np.sqrt(power[i] * power[j]) if denom > 0: - val = np.sqrt(re_sum ** 2 + im_sum ** 2) / denom + val = np.sqrt(re_sum**2 + im_sum**2) / denom else: val = 0.0 con[e, f, i, j] = val diff --git a/hypyp/sync/envelope_corr.py b/hypyp/sync/envelope_corr.py index 8144b9b..234c611 100644 --- a/hypyp/sync/envelope_corr.py +++ b/hypyp/sync/envelope_corr.py @@ -48,8 +48,9 @@ class EnvCorr(BaseMetric): name = "envcorr" - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Envelope Correlation. @@ -70,32 +71,38 @@ def compute(self, complex_signal: np.ndarray, n_samp: int, Envelope Correlation connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'cuda_kernel': + if self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) def _compute_cuda(self, complex_signal, n_samp, transpose_axes): """CUDA kernel for Envelope Correlation.""" from .kernels.cuda_amplitude import envcorr_cuda + return envcorr_cuda(complex_signal) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of Envelope Correlation.""" n_epoch, n_freq, n_ch_total = complex_signal.shape[:3] env = np.abs(complex_signal) mu_env = np.mean(env, axis=3).reshape(n_epoch, n_freq, n_ch_total, 1) env = env - mu_env - con = np.einsum('nilm,nimk->nilk', env, env.transpose(transpose_axes)) / \ - np.sqrt(np.einsum('nil,nik->nilk', np.sum(env ** 2, axis=3), np.sum(env ** 2, axis=3))) + con = np.einsum( + "nilm,nimk->nilk", env, env.transpose(transpose_axes) + ) / np.sqrt( + np.einsum("nil,nik->nilk", np.sum(env**2, axis=3), np.sum(env**2, axis=3)) + ) return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of Envelope Correlation. @@ -106,8 +113,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, env = np.abs(complex_signal) return _envcorr_numba_kernel(env) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of Envelope Correlation. @@ -116,8 +124,8 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, float64 (CPU/CUDA). No complex arithmetic on GPU. """ device = self._device - float_type = torch.float32 if device == 'mps' else torch.float64 - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + float_type = torch.float32 if device == "mps" else torch.float64 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) env = torch.abs(sig).to(dtype=float_type) # real envelope @@ -128,11 +136,11 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, env = env - mu # Pearson numerator: sum_t(env_i(t) * env_j(t)) - num = torch.einsum('efit,efjt->efij', env, env) + num = torch.einsum("efit,efjt->efij", env, env) # Denominator: sqrt(sum_t(env_i²) * sum_t(env_j²)) - sum_sq = torch.sum(env ** 2, dim=3) - den = torch.sqrt(torch.einsum('efi,efj->efij', sum_sq, sum_sq)) + sum_sq = torch.sum(env**2, dim=3) + den = torch.sqrt(torch.einsum("efi,efj->efij", sum_sq, sum_sq)) con = num / den return con.cpu().numpy() @@ -140,6 +148,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _envcorr_numba_kernel(env): """ diff --git a/hypyp/sync/imaginary_coh.py b/hypyp/sync/imaginary_coh.py index e3a0ae1..630d4c5 100644 --- a/hypyp/sync/imaginary_coh.py +++ b/hypyp/sync/imaginary_coh.py @@ -31,70 +31,75 @@ class ImCoh(BaseMetric): """ Imaginary Coherence connectivity metric. - + Imaginary Coherence uses only the imaginary part of the cross-spectrum, making it robust to volume conduction effects. - + Mathematical formulation: ImCoh = Im(⟨XY*⟩) / √(⟨|X|²⟩⟨|Y|²⟩) - + References ---------- Nolte, G., Bai, O., Wheaton, L., Mari, Z., Vorbach, S., & Hallett, M. (2004). Identifying true brain interaction from EEG data using the imaginary part of coherency. Clinical Neurophysiology, 115(10), 2292-2307. """ - + name = "imcoh" - - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Imaginary Coherence. - + Parameters ---------- complex_signal : np.ndarray Complex analytic signals with shape (n_epochs, n_freq, 2*n_channels, n_times). - + n_samp : int Number of time samples. - + transpose_axes : tuple Axes to transpose for matrix multiplication. - + Returns ------- con : np.ndarray Imaginary Coherence connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'cuda_kernel': + if self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) def _compute_cuda(self, complex_signal, n_samp, transpose_axes): """CUDA kernel for Imaginary Coherence.""" from .kernels.cuda_amplitude import imcoh_cuda + return imcoh_cuda(complex_signal) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of Imaginary Coherence.""" c = np.real(complex_signal) s = np.imag(complex_signal) amp = np.abs(complex_signal) ** 2 dphi = multiply_conjugate(c, s, transpose_axes=transpose_axes) - con = np.abs(np.imag(dphi)) / np.sqrt(np.einsum('nil,nik->nilk', np.nansum(amp, axis=3), - np.nansum(amp, axis=3))) + con = np.abs(np.imag(dphi)) / np.sqrt( + np.einsum("nil,nik->nilk", np.nansum(amp, axis=3), np.nansum(amp, axis=3)) + ) return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of Imaginary Coherence. @@ -106,8 +111,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, s = np.imag(complex_signal) return _imcoh_numba_kernel(c, s) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of Imaginary Coherence. @@ -119,21 +125,22 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, MPS uses float32; CPU/CUDA uses float64. """ device = self._device - float_type = torch.float32 if device == 'mps' else torch.float64 + float_type = torch.float32 if device == "mps" else torch.float64 - sig = torch.from_numpy(complex_signal).to(device=device, - dtype=torch.complex64 if device == 'mps' - else torch.complex128) + sig = torch.from_numpy(complex_signal).to( + device=device, + dtype=torch.complex64 if device == "mps" else torch.complex128, + ) c, s = sig.real, sig.imag # Im(X_i * conj(X_j)) = s_i*c_j - c_i*s_j — 2 einsum, no complex tensor - formula = 'efit,efjt->efij' + formula = "efit,efjt->efij" im_cross = torch.einsum(formula, s, c) - torch.einsum(formula, c, s) # Power normalization: sqrt(sum|X_i|² * sum|X_j|²) - amp = c ** 2 + s ** 2 # |X|² without creating complex abs + amp = c**2 + s**2 # |X|² without creating complex abs power = torch.sum(amp, dim=3) - den = torch.sqrt(torch.einsum('efi,efj->efij', power, power)) + den = torch.sqrt(torch.einsum("efi,efj->efij", power, power)) con = torch.abs(im_cross) / den return con.cpu().numpy() @@ -141,6 +148,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _imcoh_numba_kernel(c, s): """ @@ -168,7 +176,10 @@ def _imcoh_numba_kernel(c, s): im_sum = 0.0 for t in range(n_t): # Im(X_i * conj(X_j)) = s_i*c_j - c_i*s_j - im_sum += s[e, f, i, t] * c[e, f, j, t] - c[e, f, i, t] * s[e, f, j, t] + im_sum += ( + s[e, f, i, t] * c[e, f, j, t] + - c[e, f, i, t] * s[e, f, j, t] + ) denom = np.sqrt(power[i] * power[j]) if denom > 0: val = np.abs(im_sum) / denom diff --git a/hypyp/sync/kernels/__init__.py b/hypyp/sync/kernels/__init__.py index dc7003c..7c71bed 100644 --- a/hypyp/sync/kernels/__init__.py +++ b/hypyp/sync/kernels/__init__.py @@ -9,6 +9,7 @@ # Metal availability (Apple Silicon via PyObjC) try: import Metal as _Metal + METAL_AVAILABLE = True except ImportError: METAL_AVAILABLE = False @@ -16,6 +17,7 @@ # CUDA availability (NVIDIA via CuPy) try: import cupy as _cp + CUPY_AVAILABLE = True except ImportError: CUPY_AVAILABLE = False diff --git a/hypyp/sync/kernels/_cuda_dispatch.py b/hypyp/sync/kernels/_cuda_dispatch.py index b24ccd3..85ed605 100644 --- a/hypyp/sync/kernels/_cuda_dispatch.py +++ b/hypyp/sync/kernels/_cuda_dispatch.py @@ -52,10 +52,12 @@ def run_pairwise_kernel(complex_signal, get_kernel_fn): c_flat = cp.asarray( np.ascontiguousarray(np.real(complex_signal).reshape(n_ef, C, T)), - dtype=cp.float64) + dtype=cp.float64, + ) s_flat = cp.asarray( np.ascontiguousarray(np.imag(complex_signal).reshape(n_ef, C, T)), - dtype=cp.float64) + dtype=cp.float64, + ) # Upper-triangle pair indices idx_i, idx_j = [], [] @@ -74,9 +76,9 @@ def run_pairwise_kernel(complex_signal, get_kernel_fn): grid_size = (total_threads + block_size - 1) // block_size kernel( - (grid_size,), (block_size,), - (s_flat, c_flat, out, pairs_i, pairs_j, - n_ef, C, T, n_pairs) + (grid_size,), + (block_size,), + (s_flat, c_flat, out, pairs_i, pairs_j, n_ef, C, T, n_pairs), ) result = cp.asnumpy(out).reshape(E, F, C, C) diff --git a/hypyp/sync/kernels/_metal_dispatch.py b/hypyp/sync/kernels/_metal_dispatch.py index e791e34..323a9e7 100644 --- a/hypyp/sync/kernels/_metal_dispatch.py +++ b/hypyp/sync/kernels/_metal_dispatch.py @@ -55,7 +55,8 @@ def make_const_buffer(device, value): """Create a Metal buffer containing a single uint32 constant.""" return device.newBufferWithBytes_length_options_( - struct.pack('I', value), 4, Metal.MTLResourceStorageModeShared) + struct.pack("I", value), 4, Metal.MTLResourceStorageModeShared + ) @lru_cache(maxsize=8) @@ -93,10 +94,12 @@ def run_pairwise_kernel(complex_signal, compile_fn): E, F, C, T = complex_signal.shape n_ef = E * F - c_flat = np.ascontiguousarray(np.real(complex_signal).reshape(n_ef, C, T), - dtype=np.float32) - s_flat = np.ascontiguousarray(np.imag(complex_signal).reshape(n_ef, C, T), - dtype=np.float32) + c_flat = np.ascontiguousarray( + np.real(complex_signal).reshape(n_ef, C, T), dtype=np.float32 + ) + s_flat = np.ascontiguousarray( + np.imag(complex_signal).reshape(n_ef, C, T), dtype=np.float32 + ) # Upper-triangle pair indices idx_i, idx_j = [], [] @@ -110,16 +113,21 @@ def run_pairwise_kernel(complex_signal, compile_fn): # Metal buffers buf_s = device.newBufferWithBytes_length_options_( - s_flat.tobytes(), s_flat.nbytes, Metal.MTLResourceStorageModeShared) + s_flat.tobytes(), s_flat.nbytes, Metal.MTLResourceStorageModeShared + ) buf_c = device.newBufferWithBytes_length_options_( - c_flat.tobytes(), c_flat.nbytes, Metal.MTLResourceStorageModeShared) + c_flat.tobytes(), c_flat.nbytes, Metal.MTLResourceStorageModeShared + ) out_nbytes = n_ef * C * C * 4 buf_out = device.newBufferWithLength_options_( - out_nbytes, Metal.MTLResourceStorageModeShared) + out_nbytes, Metal.MTLResourceStorageModeShared + ) buf_pi = device.newBufferWithBytes_length_options_( - idx_i.tobytes(), idx_i.nbytes, Metal.MTLResourceStorageModeShared) + idx_i.tobytes(), idx_i.nbytes, Metal.MTLResourceStorageModeShared + ) buf_pj = device.newBufferWithBytes_length_options_( - idx_j.tobytes(), idx_j.nbytes, Metal.MTLResourceStorageModeShared) + idx_j.tobytes(), idx_j.nbytes, Metal.MTLResourceStorageModeShared + ) # Constant buffers held in named locals so they can be released in the # finally block (passing them inline to setBuffer would leak them). @@ -152,7 +160,8 @@ def run_pairwise_kernel(complex_signal, compile_fn): encoder.dispatchThreads_threadsPerThreadgroup_( Metal.MTLSize(total_threads, 1, 1), - Metal.MTLSize(threads_per_group, 1, 1)) + Metal.MTLSize(threads_per_group, 1, 1), + ) encoder.endEncoding() cmd_buffer.commit() diff --git a/hypyp/sync/kernels/cuda_accorr.py b/hypyp/sync/kernels/cuda_accorr.py index 95f750d..48d83e5 100644 --- a/hypyp/sync/kernels/cuda_accorr.py +++ b/hypyp/sync/kernels/cuda_accorr.py @@ -73,6 +73,8 @@ """ _accorr_kernel = None + + def _get_accorr(): global _accorr_kernel if _accorr_kernel is None: @@ -94,12 +96,15 @@ def accorr_cuda(complex_signal): z = complex_signal / np.abs(complex_signal) c_flat = cp.asarray( - np.ascontiguousarray(np.real(z).reshape(n_ef, C, T)), dtype=cp.float64) + np.ascontiguousarray(np.real(z).reshape(n_ef, C, T)), dtype=cp.float64 + ) s_flat = cp.asarray( - np.ascontiguousarray(np.imag(z).reshape(n_ef, C, T)), dtype=cp.float64) + np.ascontiguousarray(np.imag(z).reshape(n_ef, C, T)), dtype=cp.float64 + ) angle_flat = cp.asarray( np.ascontiguousarray(np.angle(complex_signal).reshape(n_ef, C, T)), - dtype=cp.float64) + dtype=cp.float64, + ) idx_i, idx_j = [], [] for i in range(C): @@ -117,9 +122,9 @@ def accorr_cuda(complex_signal): grid_size = (total_threads + block_size - 1) // block_size kernel( - (grid_size,), (block_size,), - (s_flat, c_flat, angle_flat, out, pairs_i, pairs_j, - n_ef, C, T, n_pairs) + (grid_size,), + (block_size,), + (s_flat, c_flat, angle_flat, out, pairs_i, pairs_j, n_ef, C, T, n_pairs), ) result = cp.asnumpy(out).reshape(E, F, C, C) diff --git a/hypyp/sync/kernels/cuda_amplitude.py b/hypyp/sync/kernels/cuda_amplitude.py index 1e75443..f4cd6ca 100644 --- a/hypyp/sync/kernels/cuda_amplitude.py +++ b/hypyp/sync/kernels/cuda_amplitude.py @@ -56,12 +56,15 @@ """ _coh_kernel = None + + def _get_coh(): global _coh_kernel if _coh_kernel is None: _coh_kernel = cp.RawKernel(_COH_SOURCE, "coh_kernel") return _coh_kernel + def coh_cuda(complex_signal): """Coh via CUDA. Float64.""" return run_pairwise_kernel(complex_signal, _get_coh) @@ -99,12 +102,15 @@ def coh_cuda(complex_signal): """ _imcoh_kernel = None + + def _get_imcoh(): global _imcoh_kernel if _imcoh_kernel is None: _imcoh_kernel = cp.RawKernel(_IMCOH_SOURCE, "imcoh_kernel") return _imcoh_kernel + def imcoh_cuda(complex_signal): """ImCoh via CUDA. Float64.""" return run_pairwise_kernel(complex_signal, _get_imcoh) @@ -152,12 +158,15 @@ def imcoh_cuda(complex_signal): """ _envcorr_kernel = None + + def _get_envcorr(): global _envcorr_kernel if _envcorr_kernel is None: _envcorr_kernel = cp.RawKernel(_ENVCORR_SOURCE, "envcorr_kernel") return _envcorr_kernel + def envcorr_cuda(complex_signal): """EnvCorr via CUDA. Pearson on envelopes. Float64.""" return run_pairwise_kernel(complex_signal, _get_envcorr) @@ -205,12 +214,15 @@ def envcorr_cuda(complex_signal): """ _powcorr_kernel = None + + def _get_powcorr(): global _powcorr_kernel if _powcorr_kernel is None: _powcorr_kernel = cp.RawKernel(_POWCORR_SOURCE, "powcorr_kernel") return _powcorr_kernel + def powcorr_cuda(complex_signal): """PowCorr via CUDA. Pearson on power. Float64.""" return run_pairwise_kernel(complex_signal, _get_powcorr) diff --git a/hypyp/sync/kernels/cuda_phase.py b/hypyp/sync/kernels/cuda_phase.py index 92f5d95..3b5df70 100644 --- a/hypyp/sync/kernels/cuda_phase.py +++ b/hypyp/sync/kernels/cuda_phase.py @@ -52,12 +52,15 @@ """ _pli_kernel = None + + def _get_pli(): global _pli_kernel if _pli_kernel is None: _pli_kernel = cp.RawKernel(_PLI_SOURCE, "pli_kernel") return _pli_kernel + def pli_cuda(complex_signal): """PLI via CUDA. Float64.""" return run_pairwise_kernel(complex_signal, _get_pli) @@ -92,12 +95,15 @@ def pli_cuda(complex_signal): """ _wpli_kernel = None + + def _get_wpli(): global _wpli_kernel if _wpli_kernel is None: _wpli_kernel = cp.RawKernel(_WPLI_SOURCE, "wpli_kernel") return _wpli_kernel + def wpli_cuda(complex_signal): """wPLI via CUDA. Float64.""" return run_pairwise_kernel(complex_signal, _get_wpli) @@ -133,12 +139,15 @@ def wpli_cuda(complex_signal): """ _plv_kernel = None + + def _get_plv(): global _plv_kernel if _plv_kernel is None: _plv_kernel = cp.RawKernel(_PLV_SOURCE, "plv_kernel") return _plv_kernel + def plv_cuda(complex_signal): """PLV via CUDA. Phase-normalizes then cross-spectrum. Float64.""" z = complex_signal / np.abs(complex_signal) @@ -183,12 +192,15 @@ def plv_cuda(complex_signal): """ _ccorr_kernel = None + + def _get_ccorr(): global _ccorr_kernel if _ccorr_kernel is None: _ccorr_kernel = cp.RawKernel(_CCORR_SOURCE, "ccorr_kernel") return _ccorr_kernel + def ccorr_cuda(complex_signal): """CCorr via CUDA. Phase-normalizes then angle-free Pearson. Float64.""" z = complex_signal / np.abs(complex_signal) diff --git a/hypyp/sync/kernels/metal_accorr.py b/hypyp/sync/kernels/metal_accorr.py index f24a55f..ab1fd22 100644 --- a/hypyp/sync/kernels/metal_accorr.py +++ b/hypyp/sync/kernels/metal_accorr.py @@ -99,7 +99,9 @@ def _compile_accorr(): device = Metal.MTLCreateSystemDefaultDevice() options = Metal.MTLCompileOptions.new() - library, error = device.newLibraryWithSource_options_error_(_ACCORR_SHADER, options, None) + library, error = device.newLibraryWithSource_options_error_( + _ACCORR_SHADER, options, None + ) if error: raise RuntimeError(f"Metal ACCorr shader failed: {error}") fn = library.newFunctionWithName_("accorr_kernel") @@ -124,7 +126,8 @@ def accorr_metal(complex_signal: np.ndarray) -> np.ndarray: c_flat = np.ascontiguousarray(np.real(z).reshape(n_ef, C, T), dtype=np.float32) s_flat = np.ascontiguousarray(np.imag(z).reshape(n_ef, C, T), dtype=np.float32) angle_flat = np.ascontiguousarray( - np.angle(complex_signal).reshape(n_ef, C, T), dtype=np.float32) + np.angle(complex_signal).reshape(n_ef, C, T), dtype=np.float32 + ) idx_i, idx_j = [], [] for i in range(C): @@ -137,18 +140,24 @@ def accorr_metal(complex_signal: np.ndarray) -> np.ndarray: # Metal buffers — extended layout for ACCorr buf_s = device.newBufferWithBytes_length_options_( - s_flat.tobytes(), s_flat.nbytes, Metal.MTLResourceStorageModeShared) + s_flat.tobytes(), s_flat.nbytes, Metal.MTLResourceStorageModeShared + ) buf_c = device.newBufferWithBytes_length_options_( - c_flat.tobytes(), c_flat.nbytes, Metal.MTLResourceStorageModeShared) + c_flat.tobytes(), c_flat.nbytes, Metal.MTLResourceStorageModeShared + ) buf_angle = device.newBufferWithBytes_length_options_( - angle_flat.tobytes(), angle_flat.nbytes, Metal.MTLResourceStorageModeShared) + angle_flat.tobytes(), angle_flat.nbytes, Metal.MTLResourceStorageModeShared + ) out_nbytes = n_ef * C * C * 4 buf_out = device.newBufferWithLength_options_( - out_nbytes, Metal.MTLResourceStorageModeShared) + out_nbytes, Metal.MTLResourceStorageModeShared + ) buf_pi = device.newBufferWithBytes_length_options_( - idx_i.tobytes(), idx_i.nbytes, Metal.MTLResourceStorageModeShared) + idx_i.tobytes(), idx_i.nbytes, Metal.MTLResourceStorageModeShared + ) buf_pj = device.newBufferWithBytes_length_options_( - idx_j.tobytes(), idx_j.nbytes, Metal.MTLResourceStorageModeShared) + idx_j.tobytes(), idx_j.nbytes, Metal.MTLResourceStorageModeShared + ) # Constant buffers held in named locals so they can be released in the # finally block (passing them inline to setBuffer would leak them). @@ -182,7 +191,8 @@ def accorr_metal(complex_signal: np.ndarray) -> np.ndarray: encoder.dispatchThreads_threadsPerThreadgroup_( Metal.MTLSize(total_threads, 1, 1), - Metal.MTLSize(threads_per_group, 1, 1)) + Metal.MTLSize(threads_per_group, 1, 1), + ) encoder.endEncoding() cmd_buffer.commit() diff --git a/hypyp/sync/kernels/metal_phase.py b/hypyp/sync/kernels/metal_phase.py index 73f5345..26f2d5b 100644 --- a/hypyp/sync/kernels/metal_phase.py +++ b/hypyp/sync/kernels/metal_phase.py @@ -81,7 +81,9 @@ def _compile_pli(): device = Metal.MTLCreateSystemDefaultDevice() options = Metal.MTLCompileOptions.new() - library, error = device.newLibraryWithSource_options_error_(_PLI_SHADER, options, None) + library, error = device.newLibraryWithSource_options_error_( + _PLI_SHADER, options, None + ) if error: raise RuntimeError(f"Metal PLI shader failed: {error}") fn = library.newFunctionWithName_("pli_kernel") @@ -151,7 +153,9 @@ def pli_metal(complex_signal): def _compile_wpli(): device = Metal.MTLCreateSystemDefaultDevice() options = Metal.MTLCompileOptions.new() - library, error = device.newLibraryWithSource_options_error_(_WPLI_SHADER, options, None) + library, error = device.newLibraryWithSource_options_error_( + _WPLI_SHADER, options, None + ) if error: raise RuntimeError(f"Metal wPLI shader failed: {error}") fn = library.newFunctionWithName_("wpli_kernel") diff --git a/hypyp/sync/pli.py b/hypyp/sync/pli.py index c832cf9..669b88c 100644 --- a/hypyp/sync/pli.py +++ b/hypyp/sync/pli.py @@ -49,8 +49,9 @@ class PLI(BaseMetric): name = "pli" - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Phase Lag Index. @@ -70,18 +71,19 @@ def compute(self, complex_signal: np.ndarray, n_samp: int, con : np.ndarray PLI connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'metal': + if self._backend == "metal": return self._compute_metal(complex_signal, n_samp, transpose_axes) - elif self._backend == 'cuda_kernel': + elif self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of Phase Lag Index.""" c = np.real(complex_signal) s = np.imag(complex_signal) @@ -89,8 +91,9 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, con = np.abs(np.mean(np.sign(np.imag(dphi)), axis=4)) return con - def _compute_metal(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_metal( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Metal compute shader implementation of PLI on Apple Silicon GPU. @@ -101,20 +104,24 @@ def _compute_metal(self, complex_signal: np.ndarray, n_samp: int, Requires: pip install pyobjc-framework-Metal """ from .kernels.metal_phase import pli_metal + return pli_metal(complex_signal) - def _compute_cuda(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_cuda( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ CUDA kernel implementation of PLI on NVIDIA GPU. Requires: pip install cupy-cuda12x """ from .kernels.cuda_phase import pli_cuda + return pli_cuda(complex_signal) - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of PLI with fused kernel. @@ -129,8 +136,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, s = np.imag(complex_signal) return _pli_numba_kernel(c, s) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of Phase Lag Index using per-channel broadcast. @@ -144,20 +152,21 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, MPS uses float32 precision; CPU/CUDA uses float64. """ device = self._device - float_type = torch.float32 if device == 'mps' else torch.float64 - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + float_type = torch.float32 if device == "mps" else torch.float64 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) n_epochs, n_freq, n_ch, n_times = sig.shape c, s = sig.real, sig.imag # (E, F, C, T) - con = torch.zeros((n_epochs, n_freq, n_ch, n_ch), - device=device, dtype=float_type) + con = torch.zeros( + (n_epochs, n_freq, n_ch, n_ch), device=device, dtype=float_type + ) for i in range(n_ch): # s[:,:,i:i+1,:] is a VIEW (contiguous slice), no copy # Broadcasting against (E, F, C, T) produces (E, F, C, T) - im = s[:, :, i:i+1, :] * c - c[:, :, i:i+1, :] * s # (E, F, C, T) + im = s[:, :, i : i + 1, :] * c - c[:, :, i : i + 1, :] * s # (E, F, C, T) con[:, :, i, :] = torch.abs(torch.mean(torch.sign(im), dim=-1)) return con.cpu().numpy() @@ -165,6 +174,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _pli_numba_kernel(c, s): """ @@ -184,8 +194,10 @@ def _pli_numba_kernel(c, s): for j in range(i, n_ch): sign_sum = 0.0 for t in range(n_t): - im = s[e, f, i, t] * c[e, f, j, t] \ - - c[e, f, i, t] * s[e, f, j, t] + im = ( + s[e, f, i, t] * c[e, f, j, t] + - c[e, f, i, t] * s[e, f, j, t] + ) if im > 0: sign_sum += 1.0 elif im < 0: diff --git a/hypyp/sync/plv.py b/hypyp/sync/plv.py index befd540..7fc6194 100644 --- a/hypyp/sync/plv.py +++ b/hypyp/sync/plv.py @@ -30,57 +30,60 @@ class for the public API and references. class PLV(BaseMetric): """ Phase Locking Value connectivity metric. - + PLV measures the consistency of phase differences between two signals across time, regardless of amplitude. - + Mathematical formulation: PLV = |⟨e^(i(φₓ-φᵧ))⟩| - + References ---------- Lachaux, J. P., Rodriguez, E., Martinerie, J., & Varela, F. J. (1999). Measuring phase synchrony in brain signals. Human brain mapping, 8(4), 194-208. """ - + name = "plv" - - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Phase Locking Value. - + Parameters ---------- complex_signal : np.ndarray Complex analytic signals with shape (n_epochs, n_freq, 2*n_channels, n_times). - + n_samp : int Number of time samples. - + transpose_axes : tuple Axes to transpose for matrix multiplication. - + Returns ------- con : np.ndarray PLV connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'cuda_kernel': + if self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) def _compute_cuda(self, complex_signal, n_samp, transpose_axes): """CUDA kernel for PLV.""" from .kernels.cuda_phase import plv_cuda + return plv_cuda(complex_signal) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of PLV.""" phase = complex_signal / np.abs(complex_signal) c = np.real(phase) @@ -89,8 +92,9 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, con = abs(dphi) / n_samp return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of PLV with parallel epoch processing. @@ -108,8 +112,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, s = np.imag(phase) return _plv_numba_kernel(c, s, n_samp) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of PLV. @@ -120,7 +125,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, MPS uses float32; CPU/CUDA uses float64. """ device = self._device - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) @@ -137,6 +142,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _plv_numba_kernel(c, s, n_samp): """ @@ -158,8 +164,14 @@ def _plv_numba_kernel(c, s, n_samp): im_sum = 0.0 for t in range(n_t): # z_i * conj(z_j) = (c_i*c_j + s_i*s_j) + i*(s_i*c_j - c_i*s_j) - re_sum += c[e, f, i, t] * c[e, f, j, t] + s[e, f, i, t] * s[e, f, j, t] - im_sum += s[e, f, i, t] * c[e, f, j, t] - c[e, f, i, t] * s[e, f, j, t] + re_sum += ( + c[e, f, i, t] * c[e, f, j, t] + + s[e, f, i, t] * s[e, f, j, t] + ) + im_sum += ( + s[e, f, i, t] * c[e, f, j, t] + - c[e, f, i, t] * s[e, f, j, t] + ) con[e, f, i, j] = np.sqrt(re_sum**2 + im_sum**2) / n_samp return con diff --git a/hypyp/sync/pow_corr.py b/hypyp/sync/pow_corr.py index 9804269..21d403c 100644 --- a/hypyp/sync/pow_corr.py +++ b/hypyp/sync/pow_corr.py @@ -50,8 +50,9 @@ class PowCorr(BaseMetric): name = "powcorr" - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Power Correlation. @@ -72,32 +73,38 @@ def compute(self, complex_signal: np.ndarray, n_samp: int, Power Correlation connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'cuda_kernel': + if self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) def _compute_cuda(self, complex_signal, n_samp, transpose_axes): """CUDA kernel for Power Correlation.""" from .kernels.cuda_amplitude import powcorr_cuda + return powcorr_cuda(complex_signal) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of Power Correlation.""" n_epoch, n_freq, n_ch_total = complex_signal.shape[:3] env = np.abs(complex_signal) ** 2 mu_env = np.mean(env, axis=3).reshape(n_epoch, n_freq, n_ch_total, 1) env = env - mu_env - con = np.einsum('nilm,nimk->nilk', env, env.transpose(transpose_axes)) / \ - np.sqrt(np.einsum('nil,nik->nilk', np.sum(env ** 2, axis=3), np.sum(env ** 2, axis=3))) + con = np.einsum( + "nilm,nimk->nilk", env, env.transpose(transpose_axes) + ) / np.sqrt( + np.einsum("nil,nik->nilk", np.sum(env**2, axis=3), np.sum(env**2, axis=3)) + ) return con - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of Power Correlation. @@ -108,8 +115,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, env = np.abs(complex_signal) ** 2 return _powcorr_numba_kernel(env) - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of Power Correlation. @@ -123,8 +131,8 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, powers that underflow in float32. """ device = self._device - float_type = torch.float32 if device == 'mps' else torch.float64 - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + float_type = torch.float32 if device == "mps" else torch.float64 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) # Power: |z|² via abs then square @@ -137,16 +145,17 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # L2-normalize per channel: env_hat = env / ||env|| # Then dot(env_hat_i, env_hat_j) = Pearson correlation directly - norm = torch.sqrt(torch.sum(env ** 2, dim=3, keepdim=True)) + norm = torch.sqrt(torch.sum(env**2, dim=3, keepdim=True)) norm = torch.where(norm == 0, torch.ones_like(norm), norm) env = env / norm - con = torch.einsum('efit,efjt->efij', env, env) + con = torch.einsum("efit,efjt->efij", env, env) return con.cpu().numpy() # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _powcorr_numba_kernel(env): """ diff --git a/hypyp/sync/wpli.py b/hypyp/sync/wpli.py index dd72619..a42ca2d 100644 --- a/hypyp/sync/wpli.py +++ b/hypyp/sync/wpli.py @@ -48,8 +48,9 @@ class WPLI(BaseMetric): name = "wpli" - def compute(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def compute( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Compute Weighted Phase Lag Index. @@ -69,30 +70,35 @@ def compute(self, complex_signal: np.ndarray, n_samp: int, con : np.ndarray wPLI connectivity matrix with shape (n_epoch, n_freq, 2*n_ch, 2*n_ch). """ - if self._backend == 'metal': + if self._backend == "metal": return self._compute_metal(complex_signal, n_samp, transpose_axes) - elif self._backend == 'cuda_kernel': + elif self._backend == "cuda_kernel": return self._compute_cuda(complex_signal, n_samp, transpose_axes) - elif self._backend == 'torch': + elif self._backend == "torch": return self._compute_torch(complex_signal, n_samp, transpose_axes) - elif self._backend == 'numba': + elif self._backend == "numba": return self._compute_numba(complex_signal, n_samp, transpose_axes) return self._compute_numpy(complex_signal, n_samp, transpose_axes) - def _compute_metal(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_metal( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """Metal compute shader for wPLI on Apple Silicon GPU.""" from .kernels.metal_phase import wpli_metal + return wpli_metal(complex_signal) - def _compute_cuda(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_cuda( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """CUDA kernel for wPLI on NVIDIA GPU.""" from .kernels.cuda_phase import wpli_cuda + return wpli_cuda(complex_signal) - def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numba( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ Numba JIT implementation of wPLI with fused kernel. @@ -106,8 +112,9 @@ def _compute_numba(self, complex_signal: np.ndarray, n_samp: int, s = np.imag(complex_signal) return _wpli_numba_kernel(c, s) - def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_numpy( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """NumPy implementation of Weighted Phase Lag Index.""" c = np.real(complex_signal) s = np.imag(complex_signal) @@ -119,8 +126,9 @@ def _compute_numpy(self, complex_signal: np.ndarray, n_samp: int, con = con_num / con_den return con - def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, - transpose_axes: tuple) -> np.ndarray: + def _compute_torch( + self, complex_signal: np.ndarray, n_samp: int, transpose_axes: tuple + ) -> np.ndarray: """ PyTorch implementation of Weighted Phase Lag Index. @@ -135,41 +143,46 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, MPS uses float32 precision; CPU/CUDA uses float64. """ device = self._device - float_type = torch.float32 if device == 'mps' else torch.float64 - complex_type = torch.complex64 if device == 'mps' else torch.complex128 + float_type = torch.float32 if device == "mps" else torch.float64 + complex_type = torch.complex64 if device == "mps" else torch.complex128 sig = torch.from_numpy(complex_signal).to(device=device, dtype=complex_type) n_epochs, n_freq, n_ch, n_times = sig.shape c, s = sig.real, sig.imag - con = torch.zeros((n_epochs, n_freq, n_ch, n_ch), - device=device, dtype=float_type) + con = torch.zeros( + (n_epochs, n_freq, n_ch, n_ch), device=device, dtype=float_type + ) # Chunk by epoch — each chunk is (F, C, C, T), 5x fewer iterations # than (epoch, freq) chunking. Falls back to double loop if chunk # would exceed MPS INT_MAX. chunk_elements = n_freq * n_ch * n_ch * n_times - if device == 'mps' and chunk_elements > 2_000_000_000: + if device == "mps" and chunk_elements > 2_000_000_000: # Fallback: (epoch, freq) chunking for very large configs - formula = 'it,jt->ijt' + formula = "it,jt->ijt" for e in range(n_epochs): for f in range(n_freq): c_ef = c[e, f] s_ef = s[e, f] - im_dphi = torch.einsum(formula, s_ef, c_ef) - \ - torch.einsum(formula, c_ef, s_ef) + im_dphi = torch.einsum(formula, s_ef, c_ef) - torch.einsum( + formula, c_ef, s_ef + ) con_num = torch.abs(torch.mean(im_dphi, dim=-1)) con_den = torch.mean(torch.abs(im_dphi), dim=-1) - con_den = torch.where(con_den == 0, torch.ones_like(con_den), con_den) + con_den = torch.where( + con_den == 0, torch.ones_like(con_den), con_den + ) con[e, f] = con_num / con_den else: # Fast path: epoch-only chunking - formula = 'fit,fjt->fijt' + formula = "fit,fjt->fijt" for e in range(n_epochs): c_e = c[e] # (F, C, T) s_e = s[e] - im_dphi = torch.einsum(formula, s_e, c_e) - \ - torch.einsum(formula, c_e, s_e) + im_dphi = torch.einsum(formula, s_e, c_e) - torch.einsum( + formula, c_e, s_e + ) # |Im| * sign(Im) = Im con_num = torch.abs(torch.mean(im_dphi, dim=-1)) con_den = torch.mean(torch.abs(im_dphi), dim=-1) @@ -181,6 +194,7 @@ def _compute_torch(self, complex_signal: np.ndarray, n_samp: int, # Numba JIT kernel (module-level for caching) if NUMBA_AVAILABLE: + @njit(parallel=True, cache=True) def _wpli_numba_kernel(c, s): """ @@ -201,8 +215,10 @@ def _wpli_numba_kernel(c, s): im_sum = 0.0 abs_sum = 0.0 for t in range(n_t): - im = s[e, f, i, t] * c[e, f, j, t] \ - - c[e, f, i, t] * s[e, f, j, t] + im = ( + s[e, f, i, t] * c[e, f, j, t] + - c[e, f, i, t] * s[e, f, j, t] + ) im_sum += im abs_sum += abs(im) if abs_sum > 0: diff --git a/hypyp/utils.py b/hypyp/utils.py index 45aeae0..d742176 100644 --- a/hypyp/utils.py +++ b/hypyp/utils.py @@ -11,7 +11,6 @@ | date | 2020-03-18 | """ - from typing import Tuple, List import math import random @@ -26,30 +25,32 @@ from mne import create_info, EpochsArray -def create_epochs(raw_S1: mne.io.Raw, raw_S2: mne.io.Raw, duration: float) -> Tuple[mne.Epochs, mne.Epochs]: +def create_epochs( + raw_S1: mne.io.Raw, raw_S2: mne.io.Raw, duration: float +) -> Tuple[mne.Epochs, mne.Epochs]: """ Create epochs from continuous raw EEG data for two participants. - + This function segments continuous EEG data into fixed-length epochs for both participants, interpolates bad channels, and prepares the data for further analysis. - + Parameters ---------- raw_S1 : mne.io.Raw Raw EEG data for participant 1 - + raw_S2 : mne.io.Raw Raw EEG data for participant 2 - + duration : float Duration of each epoch in seconds - + Returns ------- tuple : (List[mne.Epochs], List[mne.Epochs]) A tuple containing two lists of Epochs objects, one for each participant - + Notes ----- The function performs the following steps: @@ -57,10 +58,10 @@ def create_epochs(raw_S1: mne.io.Raw, raw_S2: mne.io.Raw, duration: float) -> Tu 2. Segments the continuous data into epochs based on these events 3. Interpolates bad channels if present 4. Removes bad channel labels after interpolation - + The returned epochs have no baseline correction applied (baseline=None), and no automatic rejection criteria (reject=None). - + Examples -------- >>> # Create 2-second epochs from raw data @@ -74,42 +75,60 @@ def create_epochs(raw_S1: mne.io.Raw, raw_S2: mne.io.Raw, duration: float) -> Tu for raw1, raw2 in zip(raw_S1, raw_S2): # creating fixed events - fixed_events1 = mne.make_fixed_length_events(raw1, - id=1, - start=0, - stop=None, - duration=duration, - first_samp=True, - overlap=0.0) - fixed_events2 = mne.make_fixed_length_events(raw2, - id=1, - start=0, - stop=None, - duration=duration, - first_samp=True, - overlap=0.0) + fixed_events1 = mne.make_fixed_length_events( + raw1, + id=1, + start=0, + stop=None, + duration=duration, + first_samp=True, + overlap=0.0, + ) + fixed_events2 = mne.make_fixed_length_events( + raw2, + id=1, + start=0, + stop=None, + duration=duration, + first_samp=True, + overlap=0.0, + ) # epoching the events per time window - epoch1 = mne.Epochs(raw1, fixed_events1, event_id=1, tmin=0, tmax=duration, - baseline=None, preload=True, reject=None, proj=True) + epoch1 = mne.Epochs( + raw1, + fixed_events1, + event_id=1, + tmin=0, + tmax=duration, + baseline=None, + preload=True, + reject=None, + proj=True, + ) # reject=reject_criteria, no baseline correction # preload needed after - epoch2 = mne.Epochs(raw2, fixed_events2, event_id=1, tmin=0, tmax=duration, - baseline=None, preload=True, reject=None, proj=True) + epoch2 = mne.Epochs( + raw2, + fixed_events2, + event_id=1, + tmin=0, + tmax=duration, + baseline=None, + preload=True, + reject=None, + proj=True, + ) # interpolating bad channels and removing the label - if len(epoch1.info['bads']) > 0: - epoch1 = mne.Epochs.interpolate_bads(epoch1, - reset_bads=True, - mode='accurate', - origin='auto', - verbose=None) - if len(epoch2.info['bads']) > 0: - epoch2 = mne.Epochs.interpolate_bads(epoch2, - reset_bads=True, - mode='accurate', - origin='auto', - verbose=None) + if len(epoch1.info["bads"]) > 0: + epoch1 = mne.Epochs.interpolate_bads( + epoch1, reset_bads=True, mode="accurate", origin="auto", verbose=None + ) + if len(epoch2.info["bads"]) > 0: + epoch2 = mne.Epochs.interpolate_bads( + epoch2, reset_bads=True, mode="accurate", origin="auto", verbose=None + ) epoch_S1.append(epoch1) epoch_S2.append(epoch2) @@ -120,38 +139,38 @@ def create_epochs(raw_S1: mne.io.Raw, raw_S2: mne.io.Raw, duration: float) -> Tu def merge(epoch_S1: mne.Epochs, epoch_S2: mne.Epochs) -> mne.Epochs: """ Merge epochs from two participants into a single hyperscanning dataset. - + This function combines the EEG data from two participants into a single Epochs object, where channels from each participant are labeled with suffixes "_S1" and "_S2" respectively. - + Parameters ---------- epoch_S1 : mne.Epochs Epochs object for participant 1 - + epoch_S2 : mne.Epochs Epochs object for participant 2 - + Returns ------- ep_hyper : mne.Epochs Merged Epochs object containing data from both participants - + Notes ----- Prior to merging, any bad channels are interpolated and their labels are removed from the 'bads' list. - + The function assumes that the time alignment between participants has already been performed at the raw data creation stage. - + After merging, average referencing cannot be applied to the data, and standard topographic plotting is not possible. Use specialized hyperscanning visualization tools instead. - + The function preserves channel types (EEG/EOG) from the original data. - + Examples -------- >>> # Merge epochs from two participants @@ -162,42 +181,42 @@ def merge(epoch_S1: mne.Epochs, epoch_S2: mne.Epochs) -> mne.Epochs: # checking bad ch for epochs, interpolating # and removing them from 'bads' if needed - if len(epoch_S1.info['bads']) > 0: - epoch_S1 = mne.Epochs.interpolate_bads(epoch_S1, - reset_bads=True, - mode='accurate', - origin='auto', - verbose=None) + if len(epoch_S1.info["bads"]) > 0: + epoch_S1 = mne.Epochs.interpolate_bads( + epoch_S1, reset_bads=True, mode="accurate", origin="auto", verbose=None + ) # head-digitization-based origin fit - if len(epoch_S2.info['bads']) > 0: - epoch_S2 = mne.Epochs.interpolate_bads(epoch_S2, - reset_bads=True, - mode='accurate', - origin='auto', - verbose=None) + if len(epoch_S2.info["bads"]) > 0: + epoch_S2 = mne.Epochs.interpolate_bads( + epoch_S2, reset_bads=True, mode="accurate", origin="auto", verbose=None + ) - sfreq = epoch_S1[0].info['sfreq'] - ch_names = epoch_S1[0].info['ch_names'] + sfreq = epoch_S1[0].info["sfreq"] + ch_names = epoch_S1[0].info["ch_names"] # creating channels label for each participant ch_names1 = [] for i in ch_names: - ch_names1.append(i+'_S1') + ch_names1.append(i + "_S1") ch_names2 = [] for i in ch_names: - ch_names2.append(i+'_S2') + ch_names2.append(i + "_S2") merges = [] # Verify both epochs have same filter settings - if (epoch_S1.info['highpass'] != epoch_S2.info['highpass'] or - epoch_S1.info['lowpass'] != epoch_S2.info['lowpass']): + if ( + epoch_S1.info["highpass"] != epoch_S2.info["highpass"] + or epoch_S1.info["lowpass"] != epoch_S2.info["lowpass"] + ): import warnings + warnings.warn("Filter settings differ between participants. Using S1 settings.") - + # checking wether data have the same size - assert(len(epoch_S1) == len(epoch_S2) - ), "Epochs from S1 and S2 should have the same size!" + assert len(epoch_S1) == len(epoch_S2), ( + "Epochs from S1 and S2 should have the same size!" + ) # picking data per epoch for l in range(0, len(epoch_S1)): @@ -220,35 +239,35 @@ def merge(epoch_S1: mne.Epochs, epoch_S2: mne.Epochs) -> mne.Epochs: merges.append(merge_arr) merged = np.array(merges) - ch_names_merged = ch_names1+ch_names2 + ch_names_merged = ch_names1 + ch_names2 # Create info object with filter information preserved - info = mne.create_info(ch_names_merged, sfreq, ch_types='eeg', verbose=None) + info = mne.create_info(ch_names_merged, sfreq, ch_types="eeg", verbose=None) # Also preserve other relevant metadata - if 'description' in epoch_S1.info: - info['description'] = epoch_S1.info['description'] - elif 'description' in epoch_S2.info: - info['description'] = epoch_S2.info['description'] + if "description" in epoch_S1.info: + info["description"] = epoch_S1.info["description"] + elif "description" in epoch_S2.info: + info["description"] = epoch_S2.info["description"] ep_hyper = mne.EpochsArray(merged, info) # Preserve filter information from source epochs with ep_hyper.info._unlock(): - ep_hyper.info['highpass'] = epoch_S1.info['highpass'] - ep_hyper.info['lowpass'] = epoch_S1.info['lowpass'] - + ep_hyper.info["highpass"] = epoch_S1.info["highpass"] + ep_hyper.info["lowpass"] = epoch_S1.info["lowpass"] + # setting channels type EOG_ch = [] - for ch in epoch_S1.info['chs']: - if ch['kind'] == FIFF.FIFFV_EOG_CH: - EOG_ch.append(ch['ch_name']) + for ch in epoch_S1.info["chs"]: + if ch["kind"] == FIFF.FIFFV_EOG_CH: + EOG_ch.append(ch["ch_name"]) - for ch in ep_hyper.info['chs']: - if ch['ch_name'].split('_')[0] in EOG_ch: - ch['kind'] = FIFF.FIFFV_EOG_CH + for ch in ep_hyper.info["chs"]: + if ch["ch_name"].split("_")[0] in EOG_ch: + ch["kind"] = FIFF.FIFFV_EOG_CH else: - ch['kind'] = FIFF.FIFFV_EEG_CH + ch["kind"] = FIFF.FIFFV_EEG_CH return ep_hyper @@ -256,22 +275,22 @@ def merge(epoch_S1: mne.Epochs, epoch_S2: mne.Epochs) -> mne.Epochs: def split(raw_merge: mne.io.Raw) -> Tuple[mne.io.Raw, mne.io.Raw]: """ Split merged raw data back into separate datasets for each participant. - + This function reverses the merging process, extracting individual participants' data from a merged hyperscanning dataset based on channel name suffixes. - + Parameters ---------- raw_merge : mne.io.Raw Merged Raw data for both participants, with channels having suffixes "_S1" and "_S2" (or "_1" and "_2") - + Returns ------- tuple : (mne.io.Raw, mne.io.Raw) A tuple containing two Raw objects, one for each participant, with standard 10-20 montage applied - + Notes ----- The function performs the following steps: @@ -279,10 +298,10 @@ def split(raw_merge: mne.io.Raw) -> Tuple[mne.io.Raw, mne.io.Raw]: 2. Creates new Raw objects for each participant 3. Applies the standard 10-20 montage to both datasets 4. Sets the EEG reference to the average of all channels - + Channel types (EEG/EOG) are preserved from the original data. Bad channel labels are transferred to the appropriate participant. - + Examples -------- >>> # Split previously merged raw data @@ -294,11 +313,11 @@ def split(raw_merge: mne.io.Raw) -> Tuple[mne.io.Raw, mne.io.Raw]: ch_S1 = [] ch_S2 = [] ch = [] - for name in raw_merge.info['ch_names']: - if name.endswith('S1') or name.endswith('_1'): + for name in raw_merge.info["ch_names"]: + if name.endswith("S1") or name.endswith("_1"): ch_S1.append(name) - ch.append(name.split('_')[0]) - elif name.endswith('S2') or name.endswith('_2'): + ch.append(name.split("_")[0]) + elif name.endswith("S2") or name.endswith("_2"): ch_S2.append(name) # picking individual participant data @@ -306,67 +325,75 @@ def split(raw_merge: mne.io.Raw) -> Tuple[mne.io.Raw, mne.io.Raw]: data_S2 = raw_merge.get_data(picks=ch_S2) # creating info for raws - info = mne.create_info(ch, raw_merge.info['sfreq'], ch_types='eeg', verbose=None) + info = mne.create_info(ch, raw_merge.info["sfreq"], ch_types="eeg", verbose=None) raw_S1 = mne.io.RawArray(data_S1, info) raw_S2 = mne.io.RawArray(data_S2, info) # setting info about channels and task - raw_S1.info['bads'] = [ - ch.split('_')[0] for ch in ch_S1 if ch in raw_merge.info['bads']] - raw_S2.info['bads'] = [ - ch.split('_')[0] for ch in ch_S2 if ch in raw_merge.info['bads']] + raw_S1.info["bads"] = [ + ch.split("_")[0] for ch in ch_S1 if ch in raw_merge.info["bads"] + ] + raw_S2.info["bads"] = [ + ch.split("_")[0] for ch in ch_S2 if ch in raw_merge.info["bads"] + ] for raws in (raw_S1, raw_S2): - raws.info['description'] = raw_merge.info['description'] - - # setting montage 94 electrodes (ignore somes to correspond to our data) - for ch in raws.info['chs']: - if ch['ch_name'].startswith('MOh') or ch['ch_name'].startswith('MOb') or ('EOG' in ch['ch_name']): + raws.info["description"] = raw_merge.info["description"] + + # setting montage 94 electrodes (ignore somes to correspond to our data) + for ch in raws.info["chs"]: + if ( + ch["ch_name"].startswith("MOh") + or ch["ch_name"].startswith("MOb") + or ("EOG" in ch["ch_name"]) + ): # print('emg') - ch['kind'] = FIFF.FIFFV_EOG_CH + ch["kind"] = FIFF.FIFFV_EOG_CH else: - ch['kind'] = FIFF.FIFFV_EEG_CH - montage = mne.channels.make_standard_montage('standard_1020') - raw_1020_S1 = raw_S1.copy().set_montage(montage, on_missing='ignore') - raw_1020_S2 = raw_S2.copy().set_montage(montage, on_missing='ignore') + ch["kind"] = FIFF.FIFFV_EEG_CH + montage = mne.channels.make_standard_montage("standard_1020") + raw_1020_S1 = raw_S1.copy().set_montage(montage, on_missing="ignore") + raw_1020_S2 = raw_S2.copy().set_montage(montage, on_missing="ignore") # set reference to electrodes average # (instate of initial ref to avoid ref biais) # and storing it in raw.info['projs']: applied when Epochs - raw_1020_S1, _ = mne.set_eeg_reference(raw_1020_S1, 'average', projection=True) - raw_1020_S2, _ = mne.set_eeg_reference(raw_1020_S2, 'average', projection=True) + raw_1020_S1, _ = mne.set_eeg_reference(raw_1020_S1, "average", projection=True) + raw_1020_S2, _ = mne.set_eeg_reference(raw_1020_S2, "average", projection=True) return raw_1020_S1, raw_1020_S2 -def concatenate_epochs(epoch_S1: mne.Epochs, epoch_S2: mne.Epochs) -> Tuple[mne.Epochs, mne.Epochs]: +def concatenate_epochs( + epoch_S1: mne.Epochs, epoch_S2: mne.Epochs +) -> Tuple[mne.Epochs, mne.Epochs]: """ Concatenate multiple epochs objects into single epochs objects for each participant. - + This function combines multiple epoch instances (e.g., from different experimental blocks or sessions) into a single epochs object for each participant. - + Parameters ---------- epoch_S1 : list List of Epochs objects for participant 1 - + epoch_S2 : list List of Epochs objects for participant 2 - + Returns ------- tuple : (mne.Epochs, mne.Epochs) A tuple containing two concatenated Epochs objects, one for each participant - + Notes ----- This function is useful when you have recorded multiple experimental blocks or conditions and want to combine them for analysis. - + The epochs to be concatenated must be compatible in terms of sampling rate, channel names, and other attributes. - + Examples -------- >>> # Concatenate epochs from two experimental blocks @@ -386,38 +413,38 @@ def concatenate_epochs(epoch_S1: mne.Epochs, epoch_S2: mne.Epochs) -> Tuple[mne. def normalizing(baseline: np.ndarray, task: np.ndarray, type: str) -> np.ndarray: """ Normalize data (PSD or CSD values) relative to a baseline condition. - + This function computes Z-scores or log-ratios between task and baseline conditions, which is useful for comparing spectral power or connectivity changes relative to a reference state. - + Parameters ---------- baseline : np.ndarray Baseline condition data with shape (n_epochs, n_channels, n_frequencies) - + task : np.ndarray Task condition data with shape (n_epochs, n_channels, n_frequencies) - + type : str Normalization method to use: - 'Zscore': (task - baseline_mean) / baseline_std - 'Logratio': log10(task / baseline_mean) - + Returns ------- Normed_task : np.ndarray Normalized data with shape (n_channels, n_frequencies) - + Notes ----- For 'Logratio' normalization, only positive values can be used as input. If your data contains negative values, consider taking the absolute value before normalization. - + Both normalization methods average across epochs before computing the normalization, resulting in a single normalized value per channel and frequency. - + Examples -------- >>> # Normalize alpha power using Z-scores @@ -433,48 +460,51 @@ def normalizing(baseline: np.ndarray, task: np.ndarray, type: str) -> np.ndarray m_task = np.mean(task, axis=0) std_baseline = np.std(baseline, axis=0) # normalizing power during task by baseline average power across events - if type == 'Zscore': + if type == "Zscore": s = np.subtract(m_task, m_baseline) Normed_task = np.divide(s, std_baseline) - if type == 'Logratio': + if type == "Logratio": d = np.divide(m_task, m_baseline) Normed_task = np.log10(d) return Normed_task -def generate_random_epoch(epoch: mne.Epochs, mu: float=0, sigma: float=2.0) -> mne.Epochs: + +def generate_random_epoch( + epoch: mne.Epochs, mu: float = 0, sigma: float = 2.0 +) -> mne.Epochs: """ Generate epochs with random data following a normal distribution. - + This function creates a new Epochs object with the same structure as the input epochs, but with random data values drawn from a normal distribution. - + Parameters ---------- epoch : mne.Epochs Template Epochs object to copy structure from - + mu : float, optional Mean of the normal distribution (default=0) - + sigma : float, optional Standard deviation of the normal distribution (default=2.0) - + Returns ------- random_epochs : mne.Epochs New Epochs object with random data values - + Notes ----- This function is useful for: - Creating null data for testing analysis pipelines - Generating surrogate data for statistical tests - Simulating baseline noise with known properties - + The random data has the same dimensions as the input epoch data: (n_epochs, n_channels, n_times) - + Examples -------- >>> # Generate random epochs with the same structure as real data @@ -485,8 +515,8 @@ def generate_random_epoch(epoch: mne.Epochs, mu: float=0, sigma: float=2.0) -> m >>> print(f"Real data mean: {real_mean}, Random data mean: {random_mean}") """ - # Get epoch information - info = epoch.info #create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types) + # Get epoch information + info = epoch.info # create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types) # Get epochs as a 3D NumPy array of shape (n_epochs, n_channels, n_times) # Get the arrays’ shape @@ -499,59 +529,64 @@ def generate_random_epoch(epoch: mne.Epochs, mu: float=0, sigma: float=2.0) -> m return EpochsArray(data=r_epoch, info=info) -def generate_virtual_epoch(epoch: mne.Epochs, W: np.ndarray, frequency_mean: float=10, - frequency_std: float=0.2, noise_phase_level: float=0.005, - noise_amplitude_level: float=0.1) -> mne.Epochs: +def generate_virtual_epoch( + epoch: mne.Epochs, + W: np.ndarray, + frequency_mean: float = 10, + frequency_std: float = 0.2, + noise_phase_level: float = 0.005, + noise_amplitude_level: float = 0.1, +) -> mne.Epochs: """ Generate epochs with simulated data using Kuramoto oscillators. - + This function creates a new Epochs object with simulated EEG data based on a network of coupled Kuramoto oscillators, which can be used to model brain oscillations and synchronization. - + Parameters ---------- epoch : mne.Epochs Template Epochs object to copy structure from - + W : np.ndarray Coupling matrix between oscillators, with shape (n_channels, n_channels) - + frequency_mean : float, optional Mean frequency of oscillators in Hz (default=10) - + frequency_std : float, optional Standard deviation of oscillator frequencies in Hz (default=0.2) - + noise_phase_level : float, optional Amount of noise added to the phase (default=0.005) - + noise_amplitude_level : float, optional Amount of noise added to the amplitude (default=0.1) - + Returns ------- simulated_epochs : mne.Epochs New Epochs object with simulated data values - + Notes ----- The Kuramoto model is a mathematical model used to describe synchronization in a network of coupled oscillators. Each oscillator has its own natural frequency drawn from a normal distribution with mean `frequency_mean` and standard deviation `frequency_std`. - + The coupling between oscillators is defined by the matrix W, where W[i,j] represents the strength of the connection from oscillator j to oscillator i. - + The simulation uses the scipy.integrate.solve_ivp function to solve the differential equations of the Kuramoto model. - + This function is useful for: - Testing connectivity measures with known ground truth - Simulating synchronization phenomena - Generating data with controlled properties for method validation - + Examples -------- >>> # Create a simple coupling matrix (3 oscillators) @@ -571,32 +606,47 @@ def generate_virtual_epoch(epoch: mne.Epochs, W: np.ndarray, frequency_mean: flo """ n_epo, n_chan, n_samp = epoch.get_data().shape - sfreq = epoch.info['sfreq'] + sfreq = epoch.info["sfreq"] Nt = n_samp * n_epo tmax = n_samp / sfreq * n_epo # s - tv = np.linspace(0., tmax, Nt) + tv = np.linspace(0.0, tmax, Nt) freq = frequency_mean + frequency_std * np.random.randn(n_chan) - omega = 2. * np.pi * freq + omega = 2.0 * np.pi * freq def fp(t, p): p = np.atleast_2d(p) - coupling = np.squeeze((np.sin(p) * np.matmul(W, np.cos(p).T).T) - (np.cos(p) * np.matmul(W, np.sin(p).T).T)) + coupling = np.squeeze( + (np.sin(p) * np.matmul(W, np.cos(p).T).T) + - (np.cos(p) * np.matmul(W, np.sin(p).T).T) + ) dotp = omega - coupling + noise_phase_level * np.random.randn(n_chan) / n_samp return dotp - p0 = 2 * np.pi * np.block([np.zeros(n_chan//2), np.zeros(n_chan//2) + np.random.rand(n_chan//2) + 0.5]) + p0 = ( + 2 + * np.pi + * np.block( + [ + np.zeros(n_chan // 2), + np.zeros(n_chan // 2) + np.random.rand(n_chan // 2) + 0.5, + ] + ) + ) ans = solve_ivp(fun=fp, t_span=(tv[0], tv[-1]), y0=p0, t_eval=tv) - phi = ans['y'].T % (2*np.pi) + phi = ans["y"].T % (2 * np.pi) eeg = np.sin(phi) + noise_amplitude_level * np.random.randn(*phi.shape) - + simulation = epoch.copy() - simulation._data = np.transpose(np.reshape(eeg.T, [n_chan, n_epo, n_samp]), (1, 0, 2)) - + simulation._data = np.transpose( + np.reshape(eeg.T, [n_chan, n_epo, n_samp]), (1, 0, 2) + ) + return simulation + class Task: name: str onset_event_id: int | None @@ -604,14 +654,14 @@ class Task: offset_event_id: int | None duration: float | None - + def __init__( self, - name:str, - onset_event_id:int|None=None, - onset_time:float|None=None, - offset_event_id:int|None=None, - duration:float|None=None, + name: str, + onset_event_id: int | None = None, + onset_time: float | None = None, + offset_event_id: int | None = None, + duration: float | None = None, ): self.name = name self.onset_event_id = None @@ -630,26 +680,26 @@ def __init__( if duration is not None: self.duration = duration - + # Need one start if self.onset_event_id is None and self.onset_time is None: - raise RuntimeError('Must set either onset_event_id or onset_time') + raise RuntimeError("Must set either onset_event_id or onset_time") if self.onset_event_id is not None and self.onset_time is not None: - raise RuntimeError('Cannot set both onset_event_id and onset_time') + raise RuntimeError("Cannot set both onset_event_id and onset_time") # Need one end if self.offset_event_id is None and self.duration is None: - raise RuntimeError('Must set either offset_event_id or duration') + raise RuntimeError("Must set either offset_event_id or duration") if self.offset_event_id is not None and self.duration is not None: - raise RuntimeError('Cannot set both offset_event_id and duration') + raise RuntimeError("Cannot set both offset_event_id and duration") # invalid if self.onset_time is not None and self.offset_event_id is not None: - raise RuntimeError('Cannot use offset_event_id with onset_time') + raise RuntimeError("Cannot use offset_event_id with onset_time") - #if self.onset_event_id is not None and self.onset + # if self.onset_event_id is not None and self.onset @property def is_event_based(self): @@ -659,22 +709,26 @@ def is_event_based(self): def is_time_based(self): return self.onset_time is not None + # typing -TaskTuple = tuple[str, int, int|None] +TaskTuple = tuple[str, int, int | None] TaskList = list[Task] # Constants for task description TASK_NEXT_EVENT = -2 TASK_BEGINNING = -1 TASK_END = -1 -TASK_NAME_WHOLE_RECORD = 'whole_record' +TASK_NAME_WHOLE_RECORD = "whole_record" -def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) -> List[mne.Epochs]: + +def epochs_from_tasks( + raw: mne.io.Raw, tasks: TaskList, verbose: bool = False +) -> List[mne.Epochs]: events, events_map = mne.events_from_annotations(raw) - #print(events) + # print(events) all_epochs = [] - sfreq = raw.info['sfreq'] + sfreq = raw.info["sfreq"] for task in [t for t in tasks if t.is_event_based]: task_key = task.name @@ -685,10 +739,14 @@ def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) - task_events = [] if onset_event_id is not None and not isinstance(onset_event_id, int): - raise ValueError(f'onset_event_id must be an integer. Received {type(onset_event_id)}') + raise ValueError( + f"onset_event_id must be an integer. Received {type(onset_event_id)}" + ) if offset_event_id is not None and not isinstance(offset_event_id, int): - raise ValueError(f'offset_event_id must be an integer. Received {type(offset_event_id)}') + raise ValueError( + f"offset_event_id must be an integer. Received {type(offset_event_id)}" + ) # To handle start of raw as "event" for task if onset_event_id == TASK_BEGINNING: @@ -698,7 +756,7 @@ def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) - for event_start in events_loop: t_start = event_start[0] - if task.duration is not None: # will not use offset_event_id + if task.duration is not None: # will not use offset_event_id t_duration = task.duration else: # Find end of task @@ -715,8 +773,10 @@ def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) - event_end = events[where] if len(event_end) == 0: - raise RuntimeError(f'Cannot find end of task "{task_key}" with trigger_id "{offset_event_id}" (event_id "{offset_event_id}")') - t_end = event_end[0, 0] # use the first + raise RuntimeError( + f'Cannot find end of task "{task_key}" with trigger_id "{offset_event_id}" (event_id "{offset_event_id}")' + ) + t_end = event_end[0, 0] # use the first t_min = (t_start - raw.first_samp) / sfreq t_max = (t_end - raw.first_samp) / sfreq @@ -726,23 +786,25 @@ def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) - t_durations.append(t_duration) task_events.append([t_start, 0, onset_event_id]) - all_epochs.append(mne.Epochs( - raw, - task_events, - event_id={ task_key: onset_event_id}, - tmin=0, - tmax=min(t_durations), - baseline=None, - preload=True, - event_repeated='merge', - verbose=verbose, - )) + all_epochs.append( + mne.Epochs( + raw, + task_events, + event_id={task_key: onset_event_id}, + tmin=0, + tmax=min(t_durations), + baseline=None, + preload=True, + event_repeated="merge", + verbose=verbose, + ) + ) # time based events_per_task = dict() duration_per_task = dict() event_id_per_task = dict() - next_event_id = 1000 # start our event_id at 1000 for time based + next_event_id = 1000 # start our event_id at 1000 for time based for i, task in enumerate([t for t in tasks if t.is_time_based]): task_name = task.name @@ -755,13 +817,15 @@ def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) - next_event_id += 1 event_id = event_id_per_task[task_name] - events = mne.make_fixed_length_events(raw, - id=event_id, - start=task_start, - stop=task_end, - duration=duration, - first_samp=True, - overlap=0.0) + events = mne.make_fixed_length_events( + raw, + id=event_id, + start=task_start, + stop=task_end, + duration=duration, + first_samp=True, + overlap=0.0, + ) if not task_name in events_per_task.keys(): events_per_task[task_name] = events @@ -770,23 +834,25 @@ def epochs_from_tasks(raw: mne.io.Raw, tasks: TaskList, verbose: bool = False) - events_per_task[task_name] = np.vstack([events_per_task[task_name], events]) duration_per_task[task_name] = min(duration_per_task[task_name], duration) - for task_name in events_per_task.keys(): event_id_map = {task_name: event_id_per_task[task_name]} - epochs = mne.Epochs(raw, - events_per_task[task_name], - event_id=event_id_map, - tmin=0, - tmax=duration_per_task[task_name], - baseline=None, - preload=True, - reject=None, - proj=True, - verbose=verbose) + epochs = mne.Epochs( + raw, + events_per_task[task_name], + event_id=event_id_map, + tmin=0, + tmax=duration_per_task[task_name], + baseline=None, + preload=True, + reject=None, + proj=True, + verbose=verbose, + ) all_epochs.append(epochs) return all_epochs + def downsample_in_time(times, *args, bins=500): ret = [] # We assume time is always the last column @@ -794,23 +860,32 @@ def downsample_in_time(times, *args, bins=500): if factor == 1: return [times, *args, 1] - + # First deal with times. Need to pad (cval) with max value, we don't want to "go back in time" for the last values ret.append(block_reduce(times, block_size=factor, func=np.min, cval=np.max(times))) - + for item in args: if len(item.shape) == 1: - ret.append(block_reduce(item, block_size=factor, func=np.mean, cval=np.mean(item))) + ret.append( + block_reduce(item, block_size=factor, func=np.mean, cval=np.mean(item)) + ) elif len(item.shape) == 2: - ret.append(block_reduce(item, block_size=(1,factor), func=np.mean, cval=np.mean(item))) + ret.append( + block_reduce( + item, block_size=(1, factor), func=np.mean, cval=np.mean(item) + ) + ) else: - raise RuntimeError(f'Unsupported number of column for downsampling: {len(item)}') - + raise RuntimeError( + f"Unsupported number of column for downsampling: {len(item)}" + ) + ret.append(factor) return ret - -def generate_random_label(length:int) -> str: + + +def generate_random_label(length: int) -> str: """ Generate a random label of a specific length @@ -821,11 +896,12 @@ def generate_random_label(length:int) -> str: str: a unique label """ letters = string.ascii_lowercase - return ''.join(random.choice(letters) for i in range(length)) + return "".join(random.choice(letters) for i in range(length)) + def key_dict_as_data_frame(my_dict, col_names=[]) -> pd.DataFrame: d = dict() - col = col_names[0] if len(col_names) > 0 else 'key' + col = col_names[0] if len(col_names) > 0 else "key" d[col] = list(my_dict.keys()) values = list(my_dict.values()) @@ -834,7 +910,7 @@ def key_dict_as_data_frame(my_dict, col_names=[]) -> pd.DataFrame: return pd.DataFrame() for i, k in enumerate(values[0].keys()): - col = col_names[i+1] if len(col_names) > 1 else k + col = col_names[i + 1] if len(col_names) > 1 else k d[col] = [v[k] for v in values] - + return pd.DataFrame(d) diff --git a/hypyp/viz.py b/hypyp/viz.py index 52aebdc..24e515e 100644 --- a/hypyp/viz.py +++ b/hypyp/viz.py @@ -30,49 +30,56 @@ from hypyp.analyses import xwt -def transform(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: float=0.5, - rotY: float=(np.pi)/2, rotZ: float=(np.pi)/2, children: bool=False, - child_head: bool=False) -> np.ndarray: +def transform( + locs: np.ndarray, + traX: float = 0.15, + traY: float = 0, + traZ: float = 0.5, + rotY: float = (np.pi) / 2, + rotZ: float = (np.pi) / 2, + children: bool = False, + child_head: bool = False, +) -> np.ndarray: """ Apply a series of transformations to 3D coordinates for visualization. - - This function transforms electrode coordinates to fit properly into the head - models for hyperscanning visualizations. It can adjust for adult or child head + + This function transforms electrode coordinates to fit properly into the head + models for hyperscanning visualizations. It can adjust for adult or child head sizes and positions. - + Parameters ---------- locs : np.ndarray Array of shape (n_sensors, 3) containing the 3D coordinates to transform - + traX : float, optional Translation along X-axis in meters (default=0.15) - + traY : float, optional Translation along Y-axis in meters (default=0) - + traZ : float, optional Translation along Z-axis in meters (default=0.5) - + rotY : float, optional Rotation around Y-axis in radians (default=π/2) - + rotZ : float, optional Rotation around Z-axis in radians (default=π/2) - + children : bool, optional Whether to apply transformations specific to children's head models (default=False) - + child_head : bool, optional If True and children=True, apply transformations for a child's head model If False and children=True, apply transformations for an adult's head model with a child (default=False) - + Returns ------- np.ndarray The transformed 3D coordinates with the same shape as input - + Notes ----- The transformation sequence is: @@ -80,22 +87,22 @@ def transform(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: float=0.5 2. Scaling (with different factors for children) 3. X/Y/Z translations 4. Y-axis rotation - + When using with hyperscanning visualization functions, this transformation - is typically applied with different parameters to position each participant's + is typically applied with different parameters to position each participant's sensors appropriately in the visualization. - + Examples -------- >>> # Transform adult electrode positions >>> transformed_locs = transform(electrode_positions, traX=-0.17, rotZ=-np.pi/2) - >>> + >>> >>> # Transform child electrode positions >>> child_locs = transform( - ... electrode_positions, - ... traX=0.17, - ... rotZ=np.pi/2, - ... children=True, + ... electrode_positions, + ... traX=0.17, + ... rotZ=np.pi/2, + ... children=True, ... child_head=True ... ) """ @@ -109,11 +116,11 @@ def transform(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: float=0.5 scale = 1.03 trans = -0.01 yfactor = 1.2 - else: + else: scale = 1.03 trans = -0.01 yfactor = 1.2 - + # Z rotation newX = locs[:, 0] * np.cos(rotZ) - locs[:, 1] * np.sin(rotZ) newY = locs[:, 0] * np.sin(rotZ) + locs[:, 1] * np.cos(rotZ) @@ -123,12 +130,16 @@ def transform(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: float=0.5 locs[:, 1] = locs[:, 1] + traY # Reduce the size of the eeg headsets - newZ = locs[:, 0] * np.cos(rotZ) + locs[:, 1] * np.cos(rotZ) + locs[:, 2] * np.cos(rotZ/2) + newZ = ( + locs[:, 0] * np.cos(rotZ) + + locs[:, 1] * np.cos(rotZ) + + locs[:, 2] * np.cos(rotZ / 2) + ) locs[:, 2] = newZ * scale # Y rotation newX = locs[:, 0] * np.cos(rotY) + locs[:, 2] * np.sin(rotY) - newZ = - locs[:, 0] * np.sin(rotY) + locs[:, 2] * np.cos(rotY) + newZ = -locs[:, 0] * np.sin(rotY) + locs[:, 2] * np.cos(rotY) locs[:, 0] = newX * scale - trans locs[:, 2] = newZ * scale * yfactor locs[:, 2] = locs[:, 2] + traZ @@ -136,70 +147,76 @@ def transform(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: float=0.5 return locs -def transform_2d_intra(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: float=0, - rotZ: float=(np.pi)/2, children: bool=False, - child_head: bool=False) -> np.ndarray: +def transform_2d_intra( + locs: np.ndarray, + traX: float = 0.15, + traY: float = 0, + traZ: float = 0, + rotZ: float = (np.pi) / 2, + children: bool = False, + child_head: bool = False, +) -> np.ndarray: """ Transform electrode coordinates for 2D intra-brain visualizations. - + This function applies transformations to electrode coordinates for 2D visualizations of intra-brain connectivity. It's a simplified version of the transform() function with fewer rotations. - + Parameters ---------- locs : np.ndarray Array of shape (n_sensors, 3) containing the 3D coordinates to transform - + traX : float, optional Translation along X-axis in meters (default=0.15) - + traY : float, optional Translation along Y-axis in meters (default=0) - + traZ : float, optional Translation along Z-axis in meters (default=0) - + rotZ : float, optional Rotation around Z-axis in radians (default=π/2) - + children : bool, optional Whether to apply transformations specific to children's head models (default=False) - + child_head : bool, optional If True and children=True, apply transformations for a child's head model If False and children=True, apply transformations for an adult's head model with a child (default=False) - + Returns ------- np.ndarray The transformed 3D coordinates with the same shape as input - + Notes ----- This function is primarily used for 2D intra-brain connectivity visualizations where the full 3D rotation from transform() is not needed. - + The main differences from transform() are: 1. No Y-axis rotation 2. Simpler scaling for the Z dimension - + Examples -------- >>> # Transform adult electrode positions for 2D intra-brain visualization >>> transformed_locs = transform_2d_intra( - ... electrode_positions, - ... traX=-0.178, + ... electrode_positions, + ... traX=-0.178, ... traY=0.012 ... ) - >>> + >>> >>> # Transform child electrode positions >>> child_locs = transform_2d_intra( - ... electrode_positions, - ... traX=0.178, - ... traY=0.012, - ... children=True, + ... electrode_positions, + ... traX=0.178, + ... traY=0.012, + ... children=True, ... child_head=True ... ) """ @@ -214,65 +231,70 @@ def transform_2d_intra(locs: np.ndarray, traX: float=0.15, traY: float=0, traZ: else: scale = 1.0 - # Apply scaling to all coordinates + # Apply scaling to all coordinates locs = locs * scale - + # translation locs[:, 0] = locs[:, 0] + traX locs[:, 1] = locs[:, 1] + traY locs[:, 2] = locs[:, 2] + traZ # Reduce the size of the eeg headsets - newZ = locs[:, 0] * np.cos(rotZ) + locs[:, 1] * np.cos(rotZ) + locs[:, 2] * np.cos(rotZ/2) + newZ = ( + locs[:, 0] * np.cos(rotZ) + + locs[:, 1] * np.cos(rotZ) + + locs[:, 2] * np.cos(rotZ / 2) + ) locs[:, 2] = newZ return locs + def bezier_interpolation(t, p0, p1, c0, c1=None): """ Calculate points on a cubic Bezier curve for smooth connection visualization. - + This function implements cubic Bezier curve interpolation to create smooth curves for visualizing connectivity between electrodes. It's especially useful for creating aesthetically pleasing inter-brain connectivity visualizations. - + Parameters ---------- t : float Interpolation parameter between 0 and 1 0 represents the start point, 1 represents the end point - + p0 : array-like Starting point coordinates - + p1 : array-like Ending point coordinates - + c0 : array-like Control point for the starting point - + c1 : array-like, optional Control point for the ending point If None, c0 is used for both control points (for intra-brain curves) (default=None) - + Returns ------- array-like Point coordinates on the Bezier curve at parameter t - + Notes ----- The function implements the cubic Bezier formula: B(t) = (1-t)³P₀ + 3(1-t)²t(2P₀-C₀) + 3(1-t)t²(2P₁-C₁) + t³P₁ - + When visualizing intra-brain connectivity (within one brain), typically c1 is set to None, which makes the curve symmetric. - + For inter-brain connectivity, different control points are used to create a curve that properly arcs between the two head models. - + Examples -------- >>> # Calculate 10 points along a Bezier curve for smooth connection @@ -280,7 +302,7 @@ def bezier_interpolation(t, p0, p1, c0, c1=None): >>> for i in range(10): ... t = i / 9 # 0 to 1 ... point = bezier_interpolation( - ... t, + ... t, ... [0, 0, 0], # start point ... [1, 1, 0], # end point ... [0.5, 0, 0] # control point @@ -290,46 +312,53 @@ def bezier_interpolation(t, p0, p1, c0, c1=None): if c1 is None: c1 = c0 - - return ((1 - t) ** 3 * p0 + - 3 * (1 - t) ** 2 * t * (2 * p0 - c0) + - 3 * (1 - t) * t ** 2 * (2 * p1 - c1) + - t ** 3 * p1) - -def plot_sensors_2d_inter(epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = True, - children: bool=False, child_head: bool=False): + return ( + (1 - t) ** 3 * p0 + + 3 * (1 - t) ** 2 * t * (2 * p0 - c0) + + 3 * (1 - t) * t**2 * (2 * p1 - c1) + + t**3 * p1 + ) + + +def plot_sensors_2d_inter( + epo1: mne.Epochs, + epo2: mne.Epochs, + lab: bool = True, + children: bool = False, + child_head: bool = False, +): """ Plot sensors from two participants in 2D for inter-brain visualizations. - + This function plots the sensor positions from two participants in a 2D view, with special markers for bad channels. It's typically used as part of inter-brain connectivity visualizations. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + lab : bool, optional Whether to plot channel labels (default=True) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the sensors in the current matplotlib Axes - + Notes ----- The function: @@ -339,10 +368,10 @@ def plot_sensors_2d_inter(epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = True, - Circle markers for good channels - X markers for bad channels (those in info['bads']) 4. Optionally adds channel labels - + This function is typically used as part of a larger visualization function like viz_2D_topomap_inter() rather than called directly. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -354,77 +383,116 @@ def plot_sensors_2d_inter(epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = True, """ # extract sensor info and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) if children: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2), children=children, child_head=not child_head) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2, children=children, child_head=child_head) + loc1 = transform( + loc1, + traX=-0.17, + traY=0, + traZ=0.08, + rotY=(-np.pi / 12), + rotZ=(-np.pi / 2), + children=children, + child_head=not child_head, + ) + loc2 = transform( + loc2, + traX=0.17, + traY=0, + traZ=0.08, + rotY=(np.pi / 12), + rotZ=np.pi / 2, + children=children, + child_head=child_head, + ) else: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2)) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2) + loc1 = transform( + loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi / 12), rotZ=(-np.pi / 2) + ) + loc2 = transform( + loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi / 12), rotZ=np.pi / 2 + ) lab1 = [ch for ch in epo1.ch_names] lab2 = [ch for ch in epo2.ch_names] - bads_epo1 = epo1.info['bads'] - bads_epo2 = epo2.info['bads'] + bads_epo1 = epo1.info["bads"] + bads_epo2 = epo2.info["bads"] # plot sensors ('x' for bads) for ch in epo1.ch_names: index_ch = epo1.ch_names.index(ch) x1, y1, z1 = loc1[index_ch, :] if ch in bads_epo1: - plt.plot(x1, y1, marker='x', color='dimgrey') + plt.plot(x1, y1, marker="x", color="dimgrey") else: - plt.plot(x1, y1, marker='o', color='dimgrey') + plt.plot(x1, y1, marker="o", color="dimgrey") if lab: - plt.text(x1+0.012, y1+0.012, lab1[index_ch], horizontalalignment='center', verticalalignment='center') + plt.text( + x1 + 0.012, + y1 + 0.012, + lab1[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) for ch in epo2.ch_names: index_ch = epo2.ch_names.index(ch) x2, y2, z2 = loc2[index_ch, :] if ch in bads_epo2: - plt.plot(x2, y2, marker='x', color='dimgrey') + plt.plot(x2, y2, marker="x", color="dimgrey") else: - plt.plot(x2, y2, marker='o', color='dimgrey') + plt.plot(x2, y2, marker="o", color="dimgrey") if lab: - plt.text(x2+0.012, y2+0.012, lab2[index_ch], horizontalalignment='center', verticalalignment='center') - - -def plot_sensors_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = False, - children: bool=False, child_head: bool=False): + plt.text( + x2 + 0.012, + y2 + 0.012, + lab2[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) + + +def plot_sensors_2d_intra( + epo1: mne.Epochs, + epo2: mne.Epochs, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Plot sensors from two participants in 2D for intra-brain visualizations. - + This function plots the sensor positions from two participants for visualizing - intra-brain connectivity (connectivity within each brain separately). The sensors + intra-brain connectivity (connectivity within each brain separately). The sensors are positioned side by side rather than facing each other as in inter-brain plots. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the sensors in the current matplotlib Axes - + Notes ----- The function: @@ -434,11 +502,11 @@ def plot_sensors_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = False, - Circle markers for good channels - X markers for bad channels (those in info['bads']) 4. Optionally adds channel labels - + Unlike plot_sensors_2d_inter(), this function positions the heads side by side rather than facing each other, which is appropriate for visualizing intra-brain connectivity where comparisons are made between participants. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -450,80 +518,118 @@ def plot_sensors_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = False, """ # extract sensor info and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) if children: - loc1 = transform_2d_intra(loc1, traX=-0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2), children=children, child_head=not child_head) - loc2 = transform_2d_intra(loc2, traX=0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2), children=children, child_head=child_head) + loc1 = transform_2d_intra( + loc1, + traX=-0.178, + traY=0.012, + traZ=0, + rotZ=(-np.pi / 2), + children=children, + child_head=not child_head, + ) + loc2 = transform_2d_intra( + loc2, + traX=0.178, + traY=0.012, + traZ=0, + rotZ=(-np.pi / 2), + children=children, + child_head=child_head, + ) else: - loc1 = transform_2d_intra(loc1, traX=-0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2)) - loc2 = transform_2d_intra(loc2, traX=0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2)) + loc1 = transform_2d_intra( + loc1, traX=-0.178, traY=0.012, traZ=0, rotZ=(-np.pi / 2) + ) + loc2 = transform_2d_intra( + loc2, traX=0.178, traY=0.012, traZ=0, rotZ=(-np.pi / 2) + ) lab1 = [ch for ch in epo1.ch_names] lab2 = [ch for ch in epo2.ch_names] - bads_epo1 = epo1.info['bads'] - bads_epo2 = epo2.info['bads'] + bads_epo1 = epo1.info["bads"] + bads_epo2 = epo2.info["bads"] # plot sensors ('x' for bads) for ch in epo1.ch_names: index_ch = epo1.ch_names.index(ch) x1, y1, z1 = loc1[index_ch, :] if ch in bads_epo1: - plt.plot(x1, y1, marker='x', color='dimgrey') + plt.plot(x1, y1, marker="x", color="dimgrey") else: - plt.plot(x1, y1, marker='o', color='dimgrey') + plt.plot(x1, y1, marker="o", color="dimgrey") if lab: - plt.text(x1+0.012, y1+0.012, lab1[index_ch], horizontalalignment='center', verticalalignment='center') + plt.text( + x1 + 0.012, + y1 + 0.012, + lab1[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) for ch in epo2.ch_names: index_ch = epo2.ch_names.index(ch) x2, y2, z2 = loc2[index_ch, :] if ch in bads_epo2: - plt.plot(x2, y2, marker='x', color='dimgrey') + plt.plot(x2, y2, marker="x", color="dimgrey") else: - plt.plot(x2, y2, marker='o', color='dimgrey') + plt.plot(x2, y2, marker="o", color="dimgrey") if lab: - plt.text(x2+0.012, y2+0.012, lab2[index_ch], horizontalalignment='center', verticalalignment='center') - - -def plot_sensors_3d_inter(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = False, - children: bool=False, child_head: bool=False): + plt.text( + x2 + 0.012, + y2 + 0.012, + lab2[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) + + +def plot_sensors_3d_inter( + ax: str, + epo1: mne.Epochs, + epo2: mne.Epochs, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Plot sensors from two participants in 3D for inter-brain visualizations. - + This function plots the sensor positions from two participants in a 3D view, with special markers for bad channels. It's typically used as part of inter-brain connectivity 3D visualizations. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib 3D axis (created with projection='3d') - + epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the sensors in the provided matplotlib 3D Axes - + Notes ----- The function: @@ -533,13 +639,13 @@ def plot_sensors_3d_inter(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, lab: bool - Sphere markers for good channels - X markers for bad channels (those in info['bads']) 4. Optionally adds channel labels in 3D space - + This function is typically used as part of a larger visualization function like viz_3D_inter() rather than called directly. - + The main difference from the 2D equivalent is that this function requires a 3D axis and uses 3D scatter plots instead of 2D plots. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -552,81 +658,122 @@ def plot_sensors_3d_inter(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, lab: bool >>> plt.show() """ - # extract sensor infos and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) + # extract sensor infos and transform loc to fit with headmodel + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) if children: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2), children=children, child_head=not child_head) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2, children=children, child_head=child_head) + loc1 = transform( + loc1, + traX=-0.17, + traY=0, + traZ=0.08, + rotY=(-np.pi / 12), + rotZ=(-np.pi / 2), + children=children, + child_head=not child_head, + ) + loc2 = transform( + loc2, + traX=0.17, + traY=0, + traZ=0.08, + rotY=(np.pi / 12), + rotZ=np.pi / 2, + children=children, + child_head=child_head, + ) else: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2)) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2) + loc1 = transform( + loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi / 12), rotZ=(-np.pi / 2) + ) + loc2 = transform( + loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi / 12), rotZ=np.pi / 2 + ) lab1 = [ch for ch in epo1.ch_names] lab2 = [ch for ch in epo2.ch_names] - bads_epo1 = epo1.info['bads'] - bads_epo2 = epo2.info['bads'] + bads_epo1 = epo1.info["bads"] + bads_epo2 = epo2.info["bads"] # plot sensors ('x' for bads) for ch in epo1.ch_names: index_ch = epo1.ch_names.index(ch) x1, y1, z1 = loc1[index_ch, :] if ch in bads_epo1: - ax.scatter(x1, y1, z1, marker='x', color='dimgrey') + ax.scatter(x1, y1, z1, marker="x", color="dimgrey") else: - ax.scatter(x1, y1, z1, marker='o', color='dimgrey') + ax.scatter(x1, y1, z1, marker="o", color="dimgrey") if lab: - ax.text(x1+0.012, y1+0.012, z1, lab1[index_ch], horizontalalignment='center', verticalalignment='center') + ax.text( + x1 + 0.012, + y1 + 0.012, + z1, + lab1[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) for ch in epo2.ch_names: index_ch = epo2.ch_names.index(ch) x2, y2, z2 = loc2[index_ch, :] if ch in bads_epo2: - ax.scatter(x2, y2, z2, marker='x', color='dimgrey') + ax.scatter(x2, y2, z2, marker="x", color="dimgrey") else: - ax.scatter(x2, y2, z2, marker='o', color='dimgrey') + ax.scatter(x2, y2, z2, marker="o", color="dimgrey") if lab: - ax.text(x2+0.012, y2+0.012, z2, lab2[index_ch], horizontalalignment='center', verticalalignment='center') - - - -def plot_sensors_3d_intra(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, lab: bool = False, - children: bool = False, child_head: bool = False): + ax.text( + x2 + 0.012, + y2 + 0.012, + z2, + lab2[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) + + +def plot_sensors_3d_intra( + ax: str, + epo1: mne.Epochs, + epo2: mne.Epochs, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Plot sensors from two participants in 3D for intra-brain visualizations. - + This function plots the sensor positions from two participants in a 3D view - for visualizing intra-brain connectivity. The sensors are positioned side by side + for visualizing intra-brain connectivity. The sensors are positioned side by side rather than facing each other as in inter-brain plots. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib 3D axis (created with projection='3d') - + epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the sensors in the provided matplotlib 3D Axes - + Notes ----- The function: @@ -636,11 +783,11 @@ def plot_sensors_3d_intra(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, lab: bool - Sphere markers for good channels - X markers for bad channels (those in info['bads']) 4. Optionally adds channel labels in 3D space - + Unlike plot_sensors_3d_inter(), this function positions the heads side by side rather than facing each other, which is appropriate for visualizing intra-brain connectivity where comparisons are made between participants. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -653,90 +800,128 @@ def plot_sensors_3d_intra(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, lab: bool >>> plt.show() """ - # extract sensor infos and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) + # extract sensor infos and transform loc to fit with headmodel + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) if children: - loc1 = transform(loc1, traX=0, traY=0, traZ=0.04, rotY=0, rotZ=(-np.pi/2), children=children, child_head=not child_head) - loc2 = transform(loc2, traX=0, traY=0.5, traZ=0.04, rotY=0, rotZ=(-np.pi/2), children=children, child_head=child_head) + loc1 = transform( + loc1, + traX=0, + traY=0, + traZ=0.04, + rotY=0, + rotZ=(-np.pi / 2), + children=children, + child_head=not child_head, + ) + loc2 = transform( + loc2, + traX=0, + traY=0.5, + traZ=0.04, + rotY=0, + rotZ=(-np.pi / 2), + children=children, + child_head=child_head, + ) else: - loc1 = transform(loc1, traX=0, traY=0, traZ=0.04, rotY=0, rotZ=(-np.pi/2)) - loc2 = transform(loc2, traX=0, traY=0.5, traZ=0.04, rotY=0, rotZ=(-np.pi/2)) + loc1 = transform(loc1, traX=0, traY=0, traZ=0.04, rotY=0, rotZ=(-np.pi / 2)) + loc2 = transform(loc2, traX=0, traY=0.5, traZ=0.04, rotY=0, rotZ=(-np.pi / 2)) lab1 = [ch for ch in epo1.ch_names] lab2 = [ch for ch in epo2.ch_names] - bads_epo1 = epo1.info['bads'] - bads_epo2 = epo2.info['bads'] + bads_epo1 = epo1.info["bads"] + bads_epo2 = epo2.info["bads"] # plot sensors('x' for bads) for ch in epo1.ch_names: index_ch = epo1.ch_names.index(ch) x1, y1, z1 = loc1[index_ch, :] if ch in bads_epo1: - ax.scatter(x1, y1, z1, marker='x', color='dimgrey') + ax.scatter(x1, y1, z1, marker="x", color="dimgrey") else: - ax.scatter(x1, y1, z1, marker='o', color='dimgrey') + ax.scatter(x1, y1, z1, marker="o", color="dimgrey") if lab: - ax.text(x1+0.012, y1+0.012, z1, lab1[index_ch], horizontalalignment='center', verticalalignment='center') + ax.text( + x1 + 0.012, + y1 + 0.012, + z1, + lab1[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) for ch in epo2.ch_names: index_ch = epo2.ch_names.index(ch) x2, y2, z2 = loc2[index_ch, :] if ch in bads_epo2: - ax.scatter(x2, y2, z2, marker='x', color='dimgrey') + ax.scatter(x2, y2, z2, marker="x", color="dimgrey") else: - ax.scatter(x2, y2, z2, marker='o', color='dimgrey') + ax.scatter(x2, y2, z2, marker="o", color="dimgrey") if lab: - ax.text(x2+0.012, y2+0.012, z2, lab2[index_ch], horizontalalignment='center', verticalalignment='center') - - -def plot_links_2d_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, - threshold: str='auto', steps: int=10, children: bool=False, - child_head: bool=False): + ax.text( + x2 + 0.012, + y2 + 0.012, + z2, + lab2[index_ch], + horizontalalignment="center", + verticalalignment="center", + ) + + +def plot_links_2d_inter( + epo1: mne.Epochs, + epo2: mne.Epochs, + C: np.ndarray, + threshold: str = "auto", + steps: int = 10, + children: bool = False, + child_head: bool = False, +): """ Plot inter-brain connectivity links between two participants in 2D. - + This function visualizes the connectivity between electrodes of two participants using bezier curves. Only connections above a specified threshold are displayed, with line color and thickness indicating the strength and sign of the connection. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_2) containing connectivity values between all pairs of electrodes across participants - + threshold : float or str, optional Threshold for displaying connections (default='auto') - If float: Only connections with absolute value above this threshold are displayed - If 'auto': Threshold is set to the maximum median by column plus the maximum standard error by column - + steps : int, optional Number of steps for bezier curve interpolation (default=10) - If steps < 3: Straight lines are drawn instead of curves - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the connectivity links in the current matplotlib Axes - + Notes ----- The function: @@ -747,10 +932,10 @@ def plot_links_2d_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, - Red color scale for positive connectivity values - Blue color scale for negative connectivity values - Line thickness proportional to connection strength - + This function is typically used as part of a larger visualization function like viz_2D_topomap_inter() rather than called directly. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -762,18 +947,18 @@ def plot_links_2d_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, >>> fig, ax = plt.subplots(figsize=(10, 8)) >>> # Plot 2D connectivity links >>> plot_links_2d_inter( - ... epochs_subj1, - ... epochs_subj2, - ... connectivity, - ... threshold=0.5, + ... epochs_subj1, + ... epochs_subj2, + ... connectivity, + ... threshold=0.5, ... steps=15 ... ) >>> plt.show() """ - + def plot_bezier_curve(x1, y1, x2, y2, ctr1, ctr2, color, weight, steps): if steps <= 2: - plt.plot([x1, x2], [y1, y2], '-', color=color, linewidth=weight) + plt.plot([x1, x2], [y1, y2], "-", color=color, linewidth=weight) else: alphas = np.linspace(0, 1, steps) for idx in range(len(alphas) - 1): @@ -782,33 +967,55 @@ def plot_bezier_curve(x1, y1, x2, y2, ctr1, ctr2, color, weight, steps): xnn = bezier_interpolation(b, x1, x2, ctr1[0], ctr2[0]) yn = bezier_interpolation(a, y1, y2, ctr1[1], ctr2[1]) ynn = bezier_interpolation(b, y1, y2, ctr1[1], ctr2[1]) - plt.plot([xn, xnn], [yn, ynn], '-', color=color, linewidth=weight) - + plt.plot([xn, xnn], [yn, ynn], "-", color=color, linewidth=weight) + # extract sensor infos and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) - + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) + if children: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2), children=True, child_head=not child_head) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2, children=True, child_head=child_head) + loc1 = transform( + loc1, + traX=-0.17, + traY=0, + traZ=0.08, + rotY=(-np.pi / 12), + rotZ=(-np.pi / 2), + children=True, + child_head=not child_head, + ) + loc2 = transform( + loc2, + traX=0.17, + traY=0, + traZ=0.08, + rotY=(np.pi / 12), + rotZ=np.pi / 2, + children=True, + child_head=child_head, + ) else: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2)) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2) + loc1 = transform( + loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi / 12), rotZ=(-np.pi / 2) + ) + loc2 = transform( + loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi / 12), rotZ=np.pi / 2 + ) ctr1 = np.nanmean(loc1, 0) ctr2 = np.nanmean(loc2, 0) # Calculate automatic threshold - if threshold == 'auto': + if threshold == "auto": threshold = np.max(np.median(C, 0)) + np.max(np.std(C, 0)) else: threshold = threshold # define colormap - cmap_p = matplotlib.colormaps['Reds'] + cmap_p = matplotlib.colormaps["Reds"] norm_p = matplotlib.colors.Normalize(vmin=threshold, vmax=np.nanmax(C[:])) - cmap_n = matplotlib.colormaps['Blues_r'] + cmap_n = matplotlib.colormaps["Blues_r"] norm_n = matplotlib.colors.Normalize(vmin=np.min(C[:]), vmax=-threshold) @@ -819,64 +1026,75 @@ def plot_bezier_curve(x1, y1, x2, y2, ctr1, ctr2, color, weight, steps): x2, y2 = loc2[e2, 0], loc2[e2, 1] if C[e1, e2] >= threshold: color_p = cmap_p(norm_p(C[e1, e2])) - weight = 0.2 + 1.6 * ((C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold)) + weight = 0.2 + 1.6 * ( + (C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold) + ) plot_bezier_curve(x1, y1, x2, y2, ctr1, ctr2, color_p, weight, steps) elif C[e1, e2] <= -threshold: color_n = cmap_n(norm_n(C[e1, e2])) - weight = 0.2 + 1.6 * ((-C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold)) + weight = 0.2 + 1.6 * ( + (-C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold) + ) plot_bezier_curve(x1, y1, x2, y2, ctr1, ctr2, color_n, weight, steps) -def plot_links_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndarray, - threshold: str='auto', steps: int=10, children: bool=False, - child_head: bool=False): +def plot_links_2d_intra( + epo1: mne.Epochs, + epo2: mne.Epochs, + C1: np.ndarray, + C2: np.ndarray, + threshold: str = "auto", + steps: int = 10, + children: bool = False, + child_head: bool = False, +): """ Plot intra-brain connectivity links for two participants in 2D. - + This function visualizes the connectivity within each participant's brain - separately using bezier curves. Only connections above a specified threshold - are displayed, with line color and thickness indicating the strength and sign + separately using bezier curves. Only connections above a specified threshold + are displayed, with line color and thickness indicating the strength and sign of the connection. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C1 : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_1) containing connectivity values between all pairs of electrodes for participant 1 - + C2 : np.ndarray Connectivity matrix with shape (n_channels_2, n_channels_2) containing connectivity values between all pairs of electrodes for participant 2 - + threshold : float or str, optional Threshold for displaying connections (default='auto') - If float: Only connections with absolute value above this threshold are displayed - If 'auto': Threshold is set to the maximum of the median values plus the maximum of standard errors across both participants - + steps : int, optional Number of steps for bezier curve interpolation (default=10) - If steps < 3: Straight lines are drawn instead of curves - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the connectivity links in the current matplotlib Axes - + Notes ----- The function: @@ -887,15 +1105,15 @@ def plot_links_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: - Red color scale for positive connectivity values - Blue color scale for negative connectivity values - Line thickness proportional to connection strength - + Unlike plot_links_2d_inter(), this function: - Handles two separate connectivity matrices (one per participant) - Uses the same color scale for both participants for consistent visualization - Uses separate control points for each participant's connectivity curves - + This function is typically used as part of a larger visualization function like viz_2D_topomap_intra() rather than called directly. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -909,11 +1127,11 @@ def plot_links_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: >>> fig, ax = plt.subplots(figsize=(10, 8)) >>> # Plot 2D intra-brain connectivity links >>> plot_links_2d_intra( - ... epochs_subj1, - ... epochs_subj2, + ... epochs_subj1, + ... epochs_subj2, ... connectivity_1, ... connectivity_2, - ... threshold=0.5, + ... threshold=0.5, ... steps=15 ... ) >>> plt.show() @@ -921,7 +1139,7 @@ def plot_links_2d_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: def plot_bezier_curve(x1, y1, x2, y2, ctr, color, weight, steps): if steps <= 2: - plt.plot([x1, x2], [y1, y2], '-', color=color, linewidth=weight) + plt.plot([x1, x2], [y1, y2], "-", color=color, linewidth=weight) else: alphas = np.linspace(0, 1, steps) for idx in range(len(alphas) - 1): @@ -930,37 +1148,59 @@ def plot_bezier_curve(x1, y1, x2, y2, ctr, color, weight, steps): xnn = bezier_interpolation(b, x1, x2, ctr[0]) yn = bezier_interpolation(a, y1, y2, ctr[1]) ynn = bezier_interpolation(b, y1, y2, ctr[1]) # Fixed: ctr1 -> ctr - plt.plot([xn, xnn], [yn, ynn], '-', color=color, linewidth=weight) + plt.plot([xn, xnn], [yn, ynn], "-", color=color, linewidth=weight) # extract sensor infos and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) - if children: - loc1 = transform_2d_intra(loc1, traX=-0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2), children=children, child_head=not child_head) - loc2 = transform_2d_intra(loc2, traX=0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2), children=children, child_head=child_head) - else: - loc1 = transform_2d_intra(loc1, traX=-0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2)) - loc2 = transform_2d_intra(loc2, traX=0.178, traY=0.012, traZ=0, rotZ=(-np.pi/2)) + if children: + loc1 = transform_2d_intra( + loc1, + traX=-0.178, + traY=0.012, + traZ=0, + rotZ=(-np.pi / 2), + children=children, + child_head=not child_head, + ) + loc2 = transform_2d_intra( + loc2, + traX=0.178, + traY=0.012, + traZ=0, + rotZ=(-np.pi / 2), + children=children, + child_head=child_head, + ) + else: + loc1 = transform_2d_intra( + loc1, traX=-0.178, traY=0.012, traZ=0, rotZ=(-np.pi / 2) + ) + loc2 = transform_2d_intra( + loc2, traX=0.178, traY=0.012, traZ=0, rotZ=(-np.pi / 2) + ) ctr1 = np.nanmean(loc1, 0) ctr2 = np.nanmean(loc2, 0) - + # Calculate vmin and vmax for colormap as min and max [C1, C2] vmax = np.nanmax([np.nanmax(C1), np.nanmax(C2)]) vmin = np.nanmin([np.nanmin(C1), np.nanmin(C2)]) # Calculate automatic threshold - if threshold == 'auto': - threshold = np.max([np.median(C1, 0), np.median(C2, 0)]) + np.max([np.std(C1, 0), np.std(C2, 0)]) + if threshold == "auto": + threshold = np.max([np.median(C1, 0), np.median(C2, 0)]) + np.max( + [np.std(C1, 0), np.std(C2, 0)] + ) else: threshold = threshold # Define colormap for both participants - cmap_p = matplotlib.colormaps['Reds'] + cmap_p = matplotlib.colormaps["Reds"] norm_p = matplotlib.colors.Normalize(vmin=threshold, vmax=vmax) - cmap_n = matplotlib.colormaps['Blues_r'] + cmap_n = matplotlib.colormaps["Blues_r"] norm_n = matplotlib.colors.Normalize(vmin=vmin, vmax=-threshold) # plot links for participant 1 @@ -992,54 +1232,61 @@ def plot_bezier_curve(x1, y1, x2, y2, ctr, color, weight, steps): plot_bezier_curve(x1, y1, x2, y2, ctr2, color_n, weight, steps) -def plot_links_3d_inter(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, - threshold: str='auto', steps: int=10, children: bool=False, - child_head: bool=False): +def plot_links_3d_inter( + ax: str, + epo1: mne.Epochs, + epo2: mne.Epochs, + C: np.ndarray, + threshold: str = "auto", + steps: int = 10, + children: bool = False, + child_head: bool = False, +): """ Plot inter-brain connectivity links between two participants in 3D. - + This function visualizes the connectivity between electrodes of two participants using 3D bezier curves. Only connections above a specified threshold are displayed, with line color and thickness indicating the strength and sign of the connection. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib 3D axis (created with projection='3d') - + epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_2) containing connectivity values between all pairs of electrodes across participants - + threshold : float or str, optional Threshold for displaying connections (default='auto') - If float: Only connections with absolute value above this threshold are displayed - If 'auto': Threshold is set to the maximum median by column plus the maximum standard error by column - + steps : int, optional Number of steps for bezier curve interpolation (default=10) - If steps < 3: Straight lines are drawn instead of curves - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the connectivity links in the provided matplotlib 3D Axes - + Notes ----- The function: @@ -1050,11 +1297,11 @@ def plot_links_3d_inter(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarr - Red color scale for positive connectivity values - Blue color scale for negative connectivity values - Line thickness proportional to connection strength - + This function is similar to plot_links_2d_inter() but works in 3D space. It is typically used as part of a larger visualization function like viz_3D_inter() rather than called directly. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -1069,18 +1316,18 @@ def plot_links_3d_inter(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarr >>> # Plot 3D connectivity links >>> plot_links_3d_inter( ... ax, - ... epochs_subj1, - ... epochs_subj2, - ... connectivity, - ... threshold=0.5, + ... epochs_subj1, + ... epochs_subj2, + ... connectivity, + ... threshold=0.5, ... steps=15 ... ) >>> plt.show() """ - + def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, steps): if steps <= 2: - ax.plot([x1, x2], [y1, y2], [z1, z2], '-', color=color, linewidth=weight) + ax.plot([x1, x2], [y1, y2], [z1, z2], "-", color=color, linewidth=weight) else: alphas = np.linspace(0, 1, steps) for idx in range(len(alphas) - 1): @@ -1091,18 +1338,42 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step ynn = bezier_interpolation(b, y1, y2, ctr1[1], ctr2[1]) zn = bezier_interpolation(a, z1, z2, ctr1[2], ctr2[2]) znn = bezier_interpolation(b, z1, z2, ctr1[2], ctr2[2]) - ax.plot([xn, xnn], [yn, ynn], [zn, znn], '-', color=color, linewidth=weight) + ax.plot( + [xn, xnn], [yn, ynn], [zn, znn], "-", color=color, linewidth=weight + ) + + # extract sensor infos and transform loc to fit with headmodel + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) - # extract sensor infos and transform loc to fit with headmodel - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) - if children: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2), children=True, child_head=not child_head) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2, children=True, child_head=child_head) + loc1 = transform( + loc1, + traX=-0.17, + traY=0, + traZ=0.08, + rotY=(-np.pi / 12), + rotZ=(-np.pi / 2), + children=True, + child_head=not child_head, + ) + loc2 = transform( + loc2, + traX=0.17, + traY=0, + traZ=0.08, + rotY=(np.pi / 12), + rotZ=np.pi / 2, + children=True, + child_head=child_head, + ) else: - loc1 = transform(loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi/12), rotZ=(-np.pi/2)) - loc2 = transform(loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi/12), rotZ=np.pi/2) + loc1 = transform( + loc1, traX=-0.17, traY=0, traZ=0.08, rotY=(-np.pi / 12), rotZ=(-np.pi / 2) + ) + loc2 = transform( + loc2, traX=0.17, traY=0, traZ=0.08, rotY=(np.pi / 12), rotZ=np.pi / 2 + ) ctr1 = np.nanmean(loc1, 0) ctr1[2] -= 0.2 @@ -1110,16 +1381,16 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step ctr2[2] -= 0.2 # Calculate automatic threshold - if threshold == 'auto': + if threshold == "auto": threshold = np.max(np.median(C, 0)) + np.max(np.std(C, 0)) else: threshold = threshold # define colormap - cmap_p = matplotlib.colormaps['Reds'] + cmap_p = matplotlib.colormaps["Reds"] norm_p = matplotlib.colors.Normalize(vmin=threshold, vmax=np.nanmax(C[:])) - cmap_n = matplotlib.colormaps['Blues_r'] + cmap_n = matplotlib.colormaps["Blues_r"] norm_n = matplotlib.colors.Normalize(vmin=np.min(C[:]), vmax=-threshold) # plot links @@ -1129,67 +1400,83 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step x2, y2, z2 = loc2[e2, :] if C[e1, e2] >= threshold: color_p = cmap_p(norm_p(C[e1, e2])) - weight = 0.2 + 1.6 * ((C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold)) - plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color_p, weight, steps) + weight = 0.2 + 1.6 * ( + (C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold) + ) + plot_bezier_curve_3d( + x1, y1, z1, x2, y2, z2, ctr1, ctr2, color_p, weight, steps + ) elif C[e1, e2] <= -threshold: color_n = cmap_n(norm_n(C[e1, e2])) - weight = 0.2 + 1.6 * ((-C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold)) - plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color_n, weight, steps) - - -def plot_links_3d_intra(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndarray, - threshold: str='auto', steps: int=10, children: bool=False, - child_head: bool=False): + weight = 0.2 + 1.6 * ( + (-C[e1, e2] - threshold) / (np.nanmax(C[:]) - threshold) + ) + plot_bezier_curve_3d( + x1, y1, z1, x2, y2, z2, ctr1, ctr2, color_n, weight, steps + ) + + +def plot_links_3d_intra( + ax: str, + epo1: mne.Epochs, + epo2: mne.Epochs, + C1: np.ndarray, + C2: np.ndarray, + threshold: str = "auto", + steps: int = 10, + children: bool = False, + child_head: bool = False, +): """ Plot intra-brain connectivity links for two participants in 3D. - + This function visualizes the connectivity within each participant's brain - separately using 3D bezier curves. Only connections above a specified threshold - are displayed, with line color and thickness indicating the strength and sign + separately using 3D bezier curves. Only connections above a specified threshold + are displayed, with line color and thickness indicating the strength and sign of the connection. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib 3D axis (created with projection='3d') - + epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C1 : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_1) containing connectivity values between all pairs of electrodes for participant 1 - + C2 : np.ndarray Connectivity matrix with shape (n_channels_2, n_channels_2) containing connectivity values between all pairs of electrodes for participant 2 - + threshold : float or str, optional Threshold for displaying connections (default='auto') - If float: Only connections with absolute value above this threshold are displayed - If 'auto': Threshold is set to the maximum of the median values plus the maximum of standard errors across both participants - + steps : int, optional Number of steps for bezier curve interpolation (default=10) - If steps < 3: Straight lines are drawn instead of curves - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- None Plots the connectivity links in the provided matplotlib 3D Axes - + Notes ----- The function: @@ -1200,15 +1487,15 @@ def plot_links_3d_intra(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndar - Red color scale for positive connectivity values - Blue color scale for negative connectivity values - Line thickness proportional to connection strength - + Unlike plot_links_3d_inter(), this function: - Handles two separate connectivity matrices (one per participant) - Uses the same color scale for both participants for consistent visualization - Uses separate control points for each participant's connectivity curves - + This function is typically used as part of a larger visualization function like viz_3D_intra() rather than called directly. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -1225,19 +1512,19 @@ def plot_links_3d_intra(ax: str, epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndar >>> # Plot 3D intra-brain connectivity links >>> plot_links_3d_intra( ... ax, - ... epochs_subj1, - ... epochs_subj2, + ... epochs_subj1, + ... epochs_subj2, ... connectivity_1, ... connectivity_2, - ... threshold=0.5, + ... threshold=0.5, ... steps=15 ... ) >>> plt.show() """ - + def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, steps): if steps <= 2: - ax.plot([x1, x2], [y1, y2], [z1, z2], '-', color=color, linewidth=weight) + ax.plot([x1, x2], [y1, y2], [z1, z2], "-", color=color, linewidth=weight) else: alphas = np.linspace(0, 1, steps) for idx in range(len(alphas) - 1): @@ -1248,17 +1535,37 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step ynn = bezier_interpolation(b, y1, y2, ctr1[1], ctr2[1]) zn = bezier_interpolation(a, z1, z2, ctr1[2], ctr2[2]) znn = bezier_interpolation(b, z1, z2, ctr1[2], ctr2[2]) - ax.plot([xn, xnn], [yn, ynn], [zn, znn], '-', color=color, linewidth=weight) - - loc1 = copy(np.array([ch['loc'][:3] for ch in epo1.info['chs']])) - loc2 = copy(np.array([ch['loc'][:3] for ch in epo2.info['chs']])) - - if children: - loc1 = transform(loc1, traX=0, traY=0, traZ=0.04, rotY=0, rotZ=(-np.pi/2), children=children, child_head=not child_head) - loc2 = transform(loc2, traX=0, traY=0.5, traZ=0.04, rotY=0, rotZ=(-np.pi/2), children=children, child_head=child_head) - else: - loc1 = transform(loc1, traX=0, traY=0, traZ=0.04, rotY=0, rotZ=(-np.pi/2)) - loc2 = transform(loc2, traX=0, traY=0.5, traZ=0.04, rotY=0, rotZ=(-np.pi/2)) + ax.plot( + [xn, xnn], [yn, ynn], [zn, znn], "-", color=color, linewidth=weight + ) + + loc1 = copy(np.array([ch["loc"][:3] for ch in epo1.info["chs"]])) + loc2 = copy(np.array([ch["loc"][:3] for ch in epo2.info["chs"]])) + + if children: + loc1 = transform( + loc1, + traX=0, + traY=0, + traZ=0.04, + rotY=0, + rotZ=(-np.pi / 2), + children=children, + child_head=not child_head, + ) + loc2 = transform( + loc2, + traX=0, + traY=0.5, + traZ=0.04, + rotY=0, + rotZ=(-np.pi / 2), + children=children, + child_head=child_head, + ) + else: + loc1 = transform(loc1, traX=0, traY=0, traZ=0.04, rotY=0, rotZ=(-np.pi / 2)) + loc2 = transform(loc2, traX=0, traY=0.5, traZ=0.04, rotY=0, rotZ=(-np.pi / 2)) ctr1 = np.nanmean(loc1, 0) ctr1[2] -= 0.2 @@ -1272,16 +1579,18 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step vmin = np.min(Cmin) # Calculate automatic threshold - if threshold == 'auto': - threshold = np.max([np.median(C1, 0), np.median(C2, 0)]) + np.max([np.std(C1, 0), np.std(C2, 0)]) + if threshold == "auto": + threshold = np.max([np.median(C1, 0), np.median(C2, 0)]) + np.max( + [np.std(C1, 0), np.std(C2, 0)] + ) else: threshold = threshold # Define colormap for both participant - cmap_p = matplotlib.colormaps['Reds'] + cmap_p = matplotlib.colormaps["Reds"] norm_p = matplotlib.colors.Normalize(vmin=threshold, vmax=vmax) - cmap_n = matplotlib.colormaps['Blues_r'] + cmap_n = matplotlib.colormaps["Blues_r"] norm_n = matplotlib.colors.Normalize(vmin=vmin, vmax=-threshold) for e1 in range(len(loc1)): @@ -1291,11 +1600,15 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step if C1[e1, e2] >= threshold: color_p = cmap_p(norm_p(C1[e1, e2])) weight = 0.2 + 1.6 * ((C1[e1, e2] - threshold) / (vmax - threshold)) - plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr1, color_p, weight, steps) + plot_bezier_curve_3d( + x1, y1, z1, x2, y2, z2, ctr1, ctr1, color_p, weight, steps + ) elif C1[e1, e2] <= -threshold: color_n = cmap_n(norm_n(C1[e1, e2])) weight = 0.2 + 1.6 * ((-C1[e1, e2] - threshold) / (vmax - threshold)) - plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr1, color_n, weight, steps) + plot_bezier_curve_3d( + x1, y1, z1, x2, y2, z2, ctr1, ctr1, color_n, weight, steps + ) for e1 in range(len(loc2)): x1, y1, z1 = loc2[e1, :] @@ -1304,35 +1617,39 @@ def plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr1, ctr2, color, weight, step if C2[e1, e2] >= threshold: color_p = cmap_p(norm_p(C2[e1, e2])) weight = 0.2 + 1.6 * ((C2[e1, e2] - threshold) / (vmax - threshold)) - plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr2, ctr2, color_p, weight, steps) + plot_bezier_curve_3d( + x1, y1, z1, x2, y2, z2, ctr2, ctr2, color_p, weight, steps + ) elif C2[e1, e2] <= -threshold: color_n = cmap_n(norm_n(C2[e1, e2])) weight = 0.2 + 1.6 * ((-C2[e1, e2] - threshold) / (vmax - threshold)) - plot_bezier_curve_3d(x1, y1, z1, x2, y2, z2, ctr2, ctr2, color_n, weight, steps) + plot_bezier_curve_3d( + x1, y1, z1, x2, y2, z2, ctr2, ctr2, color_n, weight, steps + ) -def plot_2d_topomap_inter(ax, children: bool=False): +def plot_2d_topomap_inter(ax, children: bool = False): """ Plot 2D head topomap outlines for inter-brain visualizations. - + This function draws the head outlines for two participants facing each other - in a 2D representation. It creates the basic head shapes including circles + in a 2D representation. It creates the basic head shapes including circles for the heads, ellipses for the ears, and polygons for the noses. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib axis where the topomap will be plotted - + children : bool, optional Whether to use child head model proportions (default=False) If True, one head will be scaled to child size - + Returns ------- None Plots the head outlines in the provided matplotlib Axes - + Notes ----- The function: @@ -1340,16 +1657,16 @@ def plot_2d_topomap_inter(ax, children: bool=False): 2. Adds ears and nose features to each head 3. If children=True, scales one head (right side) to child proportions 4. Removes axis ticks and spines for cleaner visualization - + This function handles only the head outlines - it does not plot sensors or connectivity. It is typically used as part of a larger visualization function like viz_2D_topomap_inter() which then adds sensors and links. - + Key coordinates: - Head centers: [-0.178, 0] and [0.178, 0] - Adult head radius: 0.1 - Child head radius (when children=True): 0.065 - + Examples -------- >>> import matplotlib.pyplot as plt @@ -1358,27 +1675,89 @@ def plot_2d_topomap_inter(ax, children: bool=False): >>> # Plot head outlines >>> plot_2d_topomap_inter(ax, children=False) >>> plt.show() - >>> + >>> >>> # Plot adult-child head outlines >>> fig, ax = plt.subplots(figsize=(8, 6)) >>> plot_2d_topomap_inter(ax, children=True) >>> plt.show() """ - def add_head(ax, center, radius, ear_width, ear_height, nose_coords, ear_offset, zorder=0): - circle = matplotlib.patches.Circle(xy=center, radius=radius, edgecolor="k", facecolor="w") + def add_head( + ax, center, radius, ear_width, ear_height, nose_coords, ear_offset, zorder=0 + ): + circle = matplotlib.patches.Circle( + xy=center, radius=radius, edgecolor="k", facecolor="w" + ) ax.add_patch(circle) - ax.add_patch(matplotlib.patches.Ellipse(xy=[center[0], center[1] + radius + ear_offset], width=ear_height, height=ear_width, edgecolor="k", facecolor="w", zorder=zorder)) - ax.add_patch(matplotlib.patches.Ellipse(xy=[center[0], center[1] - radius - ear_offset], width=ear_height, height=ear_width, edgecolor="k", facecolor="w", zorder=zorder)) - ax.add_patch(matplotlib.patches.Polygon(xy=nose_coords, edgecolor="k", facecolor="w", zorder=zorder)) + ax.add_patch( + matplotlib.patches.Ellipse( + xy=[center[0], center[1] + radius + ear_offset], + width=ear_height, + height=ear_width, + edgecolor="k", + facecolor="w", + zorder=zorder, + ) + ) + ax.add_patch( + matplotlib.patches.Ellipse( + xy=[center[0], center[1] - radius - ear_offset], + width=ear_height, + height=ear_width, + edgecolor="k", + facecolor="w", + zorder=zorder, + ) + ) + ax.add_patch( + matplotlib.patches.Polygon( + xy=nose_coords, edgecolor="k", facecolor="w", zorder=zorder + ) + ) if children: scaling_factor = 0.95 - add_head(ax, center=[-0.178, 0], radius=0.065, ear_width=0.01625, ear_height=0.0325, nose_coords=[[-0.130 * scaling_factor, -0.02 * scaling_factor], [-0.130 * scaling_factor, 0.02 * scaling_factor], [-0.107 * scaling_factor, 0]], ear_offset=0.005) - add_head(ax, center=[0.178, 0], radius=0.1, ear_width=0.025, ear_height=0.05, nose_coords=[[0.087, -0.027], [0.087, 0.027], [0.068, 0]], ear_offset=0.005) - else: - add_head(ax, center=[-0.178, 0], radius=0.1, ear_width=0.025, ear_height=0.05, nose_coords=[[-0.087, -0.027], [-0.087, 0.027], [-0.068, 0]], ear_offset=0.005) - add_head(ax, center=[0.178, 0], radius=0.1, ear_width=0.025, ear_height=0.05, nose_coords=[[0.087, -0.027], [0.087, 0.027], [0.068, 0]], ear_offset=0.005) + add_head( + ax, + center=[-0.178, 0], + radius=0.065, + ear_width=0.01625, + ear_height=0.0325, + nose_coords=[ + [-0.130 * scaling_factor, -0.02 * scaling_factor], + [-0.130 * scaling_factor, 0.02 * scaling_factor], + [-0.107 * scaling_factor, 0], + ], + ear_offset=0.005, + ) + add_head( + ax, + center=[0.178, 0], + radius=0.1, + ear_width=0.025, + ear_height=0.05, + nose_coords=[[0.087, -0.027], [0.087, 0.027], [0.068, 0]], + ear_offset=0.005, + ) + else: + add_head( + ax, + center=[-0.178, 0], + radius=0.1, + ear_width=0.025, + ear_height=0.05, + nose_coords=[[-0.087, -0.027], [-0.087, 0.027], [-0.068, 0]], + ear_offset=0.005, + ) + add_head( + ax, + center=[0.178, 0], + radius=0.1, + ear_width=0.025, + ear_height=0.05, + nose_coords=[[0.087, -0.027], [0.087, 0.027], [0.068, 0]], + ear_offset=0.005, + ) for loc, spine in ax.spines.items(): spine.set_linewidth(0) @@ -1386,28 +1765,28 @@ def add_head(ax, center, radius, ear_width, ear_height, nose_coords, ear_offset, ax.set_yticks([]) -def plot_2d_topomap_intra(ax, children: bool=False): +def plot_2d_topomap_intra(ax, children: bool = False): """ Plot 2D head topomap outlines for intra-brain visualizations. - + This function draws the head outlines for two participants positioned side by side - in a 2D representation. It creates the basic head shapes including circles + in a 2D representation. It creates the basic head shapes including circles for the heads, ellipses for the ears, and polygons for the noses. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib axis where the topomap will be plotted - + children : bool, optional Whether to use child head model proportions (default=False) If True, one head will be scaled to child size - + Returns ------- None Plots the head outlines in the provided matplotlib Axes - + Notes ----- The function: @@ -1415,17 +1794,17 @@ def plot_2d_topomap_intra(ax, children: bool=False): 2. Adds ears and nose features to each head 3. If children=True, scales one head (right side) to child proportions 4. Removes axis ticks and spines for cleaner visualization - + Unlike plot_2d_topomap_inter() which positions heads facing each other, this function positions heads side by side with noses pointing upward. This orientation is appropriate for intra-brain connectivity visualization where each participant's brain connectivity is analyzed separately. - + Key coordinates: - Head centers: [-0.178, 0] and [0.178, 0] - Adult head radius: 0.1 - Child head radius (when children=True): 0.065 - + Examples -------- >>> import matplotlib.pyplot as plt @@ -1434,27 +1813,89 @@ def plot_2d_topomap_intra(ax, children: bool=False): >>> # Plot head outlines for intra-brain visualization >>> plot_2d_topomap_intra(ax, children=False) >>> plt.show() - >>> + >>> >>> # Plot adult-child head outlines >>> fig, ax = plt.subplots(figsize=(8, 6)) >>> plot_2d_topomap_intra(ax, children=True) >>> plt.show() """ - def add_head(ax, center, radius, ear_width, ear_height, nose_coords, ear_offset, zorder=0): - circle = matplotlib.patches.Circle(xy=center, radius=radius, edgecolor="k", facecolor="w") + def add_head( + ax, center, radius, ear_width, ear_height, nose_coords, ear_offset, zorder=0 + ): + circle = matplotlib.patches.Circle( + xy=center, radius=radius, edgecolor="k", facecolor="w" + ) ax.add_patch(circle) - ax.add_patch(matplotlib.patches.Ellipse(xy=[center[0] + radius + ear_offset, center[1]], width=ear_height, height=ear_width, edgecolor="k", facecolor="w", zorder=zorder)) - ax.add_patch(matplotlib.patches.Ellipse(xy=[center[0] - radius - ear_offset, center[1]], width=ear_height, height=ear_width, edgecolor="k", facecolor="w", zorder=zorder)) - ax.add_patch(matplotlib.patches.Polygon(xy=nose_coords, edgecolor="k", facecolor="w", zorder=zorder)) + ax.add_patch( + matplotlib.patches.Ellipse( + xy=[center[0] + radius + ear_offset, center[1]], + width=ear_height, + height=ear_width, + edgecolor="k", + facecolor="w", + zorder=zorder, + ) + ) + ax.add_patch( + matplotlib.patches.Ellipse( + xy=[center[0] - radius - ear_offset, center[1]], + width=ear_height, + height=ear_width, + edgecolor="k", + facecolor="w", + zorder=zorder, + ) + ) + ax.add_patch( + matplotlib.patches.Polygon( + xy=nose_coords, edgecolor="k", facecolor="w", zorder=zorder + ) + ) if children: scaling_factor = 0.95 - add_head(ax, center=[-0.178, 0], radius=0.065, ear_width=0.01625, ear_height=0.0325, nose_coords=[[-0.130 * scaling_factor, -0.02 * scaling_factor], [-0.130 * scaling_factor, 0.02 * scaling_factor], [-0.107 * scaling_factor, 0]], ear_offset=0.005) - add_head(ax, center=[0.178, 0], radius=0.1, ear_width=0.025, ear_height=0.05, nose_coords=[[0.151, 0.091], [0.205, 0.091], [0.178, 0.11]], ear_offset=0.005) + add_head( + ax, + center=[-0.178, 0], + radius=0.065, + ear_width=0.01625, + ear_height=0.0325, + nose_coords=[ + [-0.130 * scaling_factor, -0.02 * scaling_factor], + [-0.130 * scaling_factor, 0.02 * scaling_factor], + [-0.107 * scaling_factor, 0], + ], + ear_offset=0.005, + ) + add_head( + ax, + center=[0.178, 0], + radius=0.1, + ear_width=0.025, + ear_height=0.05, + nose_coords=[[0.151, 0.091], [0.205, 0.091], [0.178, 0.11]], + ear_offset=0.005, + ) else: - add_head(ax, center=[-0.178, 0], radius=0.1, ear_width=0.025, ear_height=0.05, nose_coords=[[-0.151, 0.091], [-0.205, 0.091], [-0.178, 0.11]], ear_offset=0.005) - add_head(ax, center=[0.178, 0], radius=0.1, ear_width=0.025, ear_height=0.05, nose_coords=[[0.151, 0.091], [0.205, 0.091], [0.178, 0.11]], ear_offset=0.005) + add_head( + ax, + center=[-0.178, 0], + radius=0.1, + ear_width=0.025, + ear_height=0.05, + nose_coords=[[-0.151, 0.091], [-0.205, 0.091], [-0.178, 0.11]], + ear_offset=0.005, + ) + add_head( + ax, + center=[0.178, 0], + radius=0.1, + ear_width=0.025, + ear_height=0.05, + nose_coords=[[0.151, 0.091], [0.205, 0.091], [0.178, 0.11]], + ear_offset=0.005, + ) for loc, spine in ax.spines.items(): spine.set_linewidth(0) @@ -1465,42 +1906,42 @@ def add_head(ax, center, radius, ear_width, ear_height, nose_coords, ear_offset, def plot_3d_heads(ax, vertices, faces): """ Plot 3D head models using vertices and faces. - - This function draws 3D head models defined by vertices and faces in a + + This function draws 3D head models defined by vertices and faces in a wireframe representation. It plots the edges between vertices that define the faces of the 3D model. - + Parameters ---------- ax : matplotlib.axes.Axes Matplotlib 3D axis (created with projection='3d') - + vertices : np.ndarray Array of shape (n_vertices, 3) containing the 3D coordinates of vertices - + faces : np.ndarray Array of shape (n_faces, 4) containing the vertex indices for each face Each face is defined by 4 vertex indices - + Returns ------- None Plots the 3D head model wireframe in the provided matplotlib 3D Axes - + Notes ----- The function: 1. Extracts vertex coordinates (with some reordering of x, y, z for correct orientation) 2. For each face, plots lines connecting its vertices 3. Uses a thin gray line style for the wireframe - + This function creates a lightweight wireframe representation rather than a fully shaded 3D model, which is sufficient for visualization purposes while maintaining good performance. - + The wireframe approach allows electrode and connectivity visualization to be clearly visible in front of or behind the head model. - + Examples -------- >>> import matplotlib.pyplot as plt @@ -1526,51 +1967,67 @@ def plot_3d_heads(ax, vertices, faces): V1 = faces[F, 1] V2 = faces[F, 2] V3 = faces[F, 3] - ax.plot([x_V[V0], x_V[V1]], - [y_V[V0], y_V[V1]], - [z_V[V0], z_V[V1]], - '-', color= 'grey', linewidth=0.3) - ax.plot([x_V[V1], x_V[V2]], - [y_V[V1], y_V[V2]], - [z_V[V1], z_V[V2]], - '-', color= 'grey', linewidth=0.3) - ax.plot([x_V[V2], x_V[V3]], - [y_V[V2], y_V[V3]], - [z_V[V2], z_V[V3]], - '-', color= 'grey', linewidth=0.3) - ax.plot([x_V[V3], x_V[V1]], - [y_V[V3], y_V[V1]], - [z_V[V3], z_V[V1]], - '-', color= 'grey', linewidth=0.3) - - -def get_3d_heads_inter(children: bool=False, child_head: bool=False): + ax.plot( + [x_V[V0], x_V[V1]], + [y_V[V0], y_V[V1]], + [z_V[V0], z_V[V1]], + "-", + color="grey", + linewidth=0.3, + ) + ax.plot( + [x_V[V1], x_V[V2]], + [y_V[V1], y_V[V2]], + [z_V[V1], z_V[V2]], + "-", + color="grey", + linewidth=0.3, + ) + ax.plot( + [x_V[V2], x_V[V3]], + [y_V[V2], y_V[V3]], + [z_V[V2], z_V[V3]], + "-", + color="grey", + linewidth=0.3, + ) + ax.plot( + [x_V[V3], x_V[V1]], + [y_V[V3], y_V[V1]], + [z_V[V3], z_V[V1]], + "-", + color="grey", + linewidth=0.3, + ) + + +def get_3d_heads_inter(children: bool = False, child_head: bool = False): """ Get vertices and faces for 3D head models in inter-brain configuration. - + This function returns the vertices and faces defining two 3D head models positioned facing each other, as needed for inter-brain connectivity visualization. - + Parameters ---------- children : bool, optional Whether to use child head model scaling (default=False) - + child_head : bool, optional If True and children=True, scales the second head to child size If False and children=True, scales the first head to child size (default=False) - + Returns ------- vertices : np.ndarray Array of shape (n_vertices, 3) containing the 3D coordinates of all vertices from both head models - + faces : np.ndarray Array of shape (n_faces, 4) containing the vertex indices for each face from both head models - + Notes ----- The function: @@ -1579,29 +2036,29 @@ def get_3d_heads_inter(children: bool=False, child_head: bool=False): 3. Rotates and translates the models to position them facing each other 4. If children=True and child_head=True, applies scaling to make the second head smaller 5. Combines vertices and faces into unified arrays - + For the faces array in the second head, vertex indices are offset by the number of vertices in the first head to maintain correct indexing in the combined array. - - This function is typically used to provide head models for the + + This function is typically used to provide head models for the viz_3D_inter() visualization function. - + Examples -------- >>> # Get vertices and faces for two adult heads >>> vertices, faces = get_3d_heads_inter() >>> print(f"Number of vertices: {len(vertices)}") >>> print(f"Number of faces: {len(faces)}") - >>> + >>> >>> # Get vertices and faces for adult and child heads >>> vertices, faces = get_3d_heads_inter(children=True, child_head=True) """ - + # Extract vertices and faces for the first head file_manager = ExitStack() atexit.register(file_manager.close) - ref = importlib_resources.files('hypyp') / 'data/Basehead.obj' + ref = importlib_resources.files("hypyp") / "data/Basehead.obj" filename = file_manager.enter_context(importlib_resources.as_file(ref)) mesh = meshio.read(Path(filename).resolve()) @@ -1640,33 +2097,33 @@ def get_3d_heads_inter(children: bool=False, child_head: bool=False): return vertices, faces -def get_3d_heads_intra(children: bool=False, child_head: bool=False): +def get_3d_heads_intra(children: bool = False, child_head: bool = False): """ Get vertices and faces for 3D head models in intra-brain configuration. - + This function returns the vertices and faces defining two 3D head models positioned side by side, as needed for intra-brain connectivity visualization. - + Parameters ---------- children : bool, optional Whether to use child head model scaling (default=False) - + child_head : bool, optional If True and children=True, scales the second head to child size If False and children=True, scales the first head to child size (default=False) - + Returns ------- vertices : np.ndarray Array of shape (n_vertices, 3) containing the 3D coordinates of all vertices from both head models - + faces : np.ndarray Array of shape (n_faces, 4) containing the vertex indices for each face from both head models - + Notes ----- The function: @@ -1675,33 +2132,33 @@ def get_3d_heads_intra(children: bool=False, child_head: bool=False): 3. Translates the models to position them side by side along the X-axis 4. If children=True and child_head=True, applies scaling to make the second head smaller 5. Combines vertices and faces into unified arrays - + Unlike get_3d_heads_inter() which rotates one head to face the other, this function simply positions the heads side by side without rotation. This configuration is appropriate for comparing intra-brain connectivity between participants. - + For the faces array in the second head, vertex indices are offset by the number of vertices in the first head to maintain correct indexing in the combined array. - - This function is typically used to provide head models for the + + This function is typically used to provide head models for the viz_3D_intra() visualization function. - + Examples -------- >>> # Get vertices and faces for two adult heads side by side >>> vertices, faces = get_3d_heads_intra() >>> print(f"Number of vertices: {len(vertices)}") >>> print(f"Number of faces: {len(faces)}") - >>> + >>> >>> # Get vertices and faces for adult and child heads side by side >>> vertices, faces = get_3d_heads_intra(children=True, child_head=True) """ - + # Extract vertices and faces for the first head file_manager = ExitStack() atexit.register(file_manager.close) - ref = importlib_resources.files('hypyp') / 'data/Basehead.obj' + ref = importlib_resources.files("hypyp") / "data/Basehead.obj" filename = file_manager.enter_context(importlib_resources.as_file(ref)) mesh = meshio.read(Path(filename).resolve()) @@ -1732,33 +2189,35 @@ def get_3d_heads_intra(children: bool=False, child_head: bool=False): return vertices, faces -def plot_significant_sensors(T_obs_plot: np.ndarray, epochs: mne.Epochs, significant: np.ndarray = None): +def plot_significant_sensors( + T_obs_plot: np.ndarray, epochs: mne.Epochs, significant: np.ndarray = None +): """ Plot statistical values on a topographic map for significant sensors. - - This function visualizes the results of statistical tests (such as t-tests or + + This function visualizes the results of statistical tests (such as t-tests or cluster-corrected t-tests) performed on sensor data across participants or conditions. It creates a topographic plot with statistical values displayed for significant sensors. - + Parameters ---------- T_obs_plot : np.ndarray Array of shape (n_sensors,) containing the statistical values (T or F statistics) to plot for significant sensors. Typically, this array has zeros for non-significant sensors and the actual statistical values for significant ones. - + epochs : mne.Epochs Epochs object to get channel positions from - + significant : np.ndarray, optional Optional mask array of shape (n_sensors,) where non-zero values indicate significant sensors to highlight (default=None) - + Returns ------- None Creates a topographic plot of the statistical values - + Notes ----- The function: @@ -1766,10 +2225,10 @@ def plot_significant_sensors(T_obs_plot: np.ndarray, epochs: mne.Epochs, signifi 2. Determines appropriate color scale limits (symmetric around zero) 3. Creates a topographic plot using MNE's plot_topomap function 4. If significant mask is provided, adds white circles around the significant sensors - + This function is designed for single-participant statistical visualization. For inter-brain connectivity statistics, use the plot_links_* functions instead. - + Examples -------- >>> import numpy as np @@ -1783,7 +2242,7 @@ def plot_significant_sensors(T_obs_plot: np.ndarray, epochs: mne.Epochs, signifi >>> t_values[significant_indices] = np.random.randn(len(significant_indices)) * 3 >>> # Plot the significant sensors >>> plot_significant_sensors(t_values, epochs) - >>> + >>> >>> # Alternatively, use a significance mask >>> significance_mask = np.zeros(n_channels) >>> significance_mask[significant_indices] = 1 @@ -1792,8 +2251,8 @@ def plot_significant_sensors(T_obs_plot: np.ndarray, epochs: mne.Epochs, signifi # getting sensors position pos = np.array([[0, 0]]) - for i in range(0, len(epochs.info['ch_names'])): - cor = np.array([epochs.info['chs'][i]['loc'][0:2]]) + for i in range(0, len(epochs.info["ch_names"])): + cor = np.array([epochs.info["chs"][i]["loc"][0:2]]) pos = np.concatenate((pos, cor), axis=0) pos = pos[1:] # topoplot of significant sensors @@ -1803,68 +2262,79 @@ def plot_significant_sensors(T_obs_plot: np.ndarray, epochs: mne.Epochs, signifi else: vmax = None vmin = None - + fig = plt.figure() ax = fig.add_subplot() if significant is not None: - mne.viz.plot_topomap(T_obs_plot, pos, vlim=(vmin, vmax), sensors=False, axes=ax, show=False) - loc1 = copy(np.array([ch['loc'][:3] for ch in epochs.info['chs']])) + mne.viz.plot_topomap( + T_obs_plot, pos, vlim=(vmin, vmax), sensors=False, axes=ax, show=False + ) + loc1 = copy(np.array([ch["loc"][:3] for ch in epochs.info["chs"]])) for ich, ch in enumerate(epochs.ch_names): index_ch = epochs.ch_names.index(ch) x1, y1, z1 = loc1[index_ch, :] if significant[ich] != 0: - ax.plot(x1, y1, marker='o', color='white') - else: - mne.viz.plot_topomap(T_obs_plot, pos, vlim=(vmin, vmax), sensors=True, axes=ax, show=True) + ax.plot(x1, y1, marker="o", color="white") + else: + mne.viz.plot_topomap( + T_obs_plot, pos, vlim=(vmin, vmax), sensors=True, axes=ax, show=True + ) return None -def viz_2D_headmodel_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, - threshold: float=0.95, steps: int=10, lab: bool = True, - children: bool=False, child_head: bool=False): +def viz_2D_headmodel_inter( + epo1: mne.Epochs, + epo2: mne.Epochs, + C: np.ndarray, + threshold: float = 0.95, + steps: int = 10, + lab: bool = True, + children: bool = False, + child_head: bool = False, +): """ Create a complete 2D head model visualization with inter-brain connectivity. - + This function provides a high-level interface for creating a 2D visualization of two head models with sensors and inter-brain connectivity links. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_2) containing connectivity values between all pairs of electrodes across participants - + threshold : float, optional Threshold for displaying connections (default=0.95) Only connections with absolute value above this threshold are displayed - + steps : int, optional Number of steps for bezier curve interpolation (default=10) If steps < 3, straight lines are drawn instead of curves - + lab : bool, optional Whether to plot channel labels (default=True) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- ax : matplotlib.axes.Axes The matplotlib Axes object containing the visualization - + Notes ----- This function combines several lower-level visualization functions: @@ -1872,11 +2342,11 @@ def viz_2D_headmodel_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, 2. Generates 3D head models and applies a camera transformation for 2D view 3. Plots sensors using plot_sensors_2d_inter() 4. Plots connectivity links using plot_links_2d_inter() - + Unlike viz_2D_topomap_inter() which uses simpler 2D head outlines, this function uses a 3D head model rendered in a 2D view, which can provide a more realistic representation. - + Examples -------- >>> import numpy as np @@ -1899,78 +2369,100 @@ def viz_2D_headmodel_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, """ fig, ax = plt.subplots(1, 1) - ax.axis("off") + ax.axis("off") if children: vertices, faces = get_3d_heads_inter(children=True, child_head=child_head) else: vertices, faces = get_3d_heads_inter() - + camera = Camera("ortho", theta=90, phi=180, scale=1) - mesh = Mesh(ax, camera.transform @ glm.yrotate(90), vertices, faces, - facecolors='white', edgecolors='black', linewidths=.25) - camera.connect(ax, mesh.update) - plt.gca().set_aspect('equal', 'box') - plt.axis('off') + mesh = Mesh( + ax, + camera.transform @ glm.yrotate(90), + vertices, + faces, + facecolors="white", + edgecolors="black", + linewidths=0.25, + ) + camera.connect(ax, mesh.update) + plt.gca().set_aspect("equal", "box") + plt.axis("off") if children: plot_sensors_2d_inter(epo1, epo2, lab=lab, children=True, child_head=child_head) - plot_links_2d_inter(epo1, epo2, C=C, threshold=threshold, steps=steps, children=True, child_head=child_head) + plot_links_2d_inter( + epo1, + epo2, + C=C, + threshold=threshold, + steps=steps, + children=True, + child_head=child_head, + ) else: plot_sensors_2d_inter(epo1, epo2, lab=lab) plot_links_2d_inter(epo1, epo2, C=C, threshold=threshold, steps=steps) - plt.tight_layout() - plt.show() - + plt.tight_layout() + plt.show() + return ax -def viz_2D_topomap_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, - threshold: float=0.95, steps: int=10, lab: bool = False, - children: bool=False, child_head: bool=False): +def viz_2D_topomap_inter( + epo1: mne.Epochs, + epo2: mne.Epochs, + C: np.ndarray, + threshold: float = 0.95, + steps: int = 10, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Create a 2D topographic map visualization with inter-brain connectivity. - + This function provides a high-level interface for creating a 2D topographic visualization of two head outlines with sensors and inter-brain connectivity links. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_2) containing connectivity values between all pairs of electrodes across participants - + threshold : float, optional Threshold for displaying connections (default=0.95) Only connections with absolute value above this threshold are displayed - + steps : int, optional Number of steps for bezier curve interpolation (default=10) If steps < 3, straight lines are drawn instead of curves - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- ax : matplotlib.axes.Axes The matplotlib Axes object containing the visualization - + Notes ----- This function combines several lower-level visualization functions: @@ -1978,11 +2470,11 @@ def viz_2D_topomap_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, 2. Draws 2D head outlines using plot_2d_topomap_inter() 3. Plots sensors using plot_sensors_2d_inter() 4. Plots connectivity links using plot_links_2d_inter() - + This function uses simpler 2D head outlines compared to viz_2D_headmodel_inter() which uses a projected 3D model. The 2D outlines are computationally lighter and may be preferred for quick visualizations or publications. - + Examples -------- >>> import numpy as np @@ -2011,7 +2503,15 @@ def viz_2D_topomap_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, if children: plot_2d_topomap_inter(ax, children=True) plot_sensors_2d_inter(epo1, epo2, lab=lab, children=True, child_head=child_head) - plot_links_2d_inter(epo1, epo2, C=C, threshold=threshold, steps=steps, children=True, child_head=child_head) + plot_links_2d_inter( + epo1, + epo2, + C=C, + threshold=threshold, + steps=steps, + children=True, + child_head=child_head, + ) else: plot_2d_topomap_inter(ax) plot_sensors_2d_inter(epo1, epo2, lab=lab) @@ -2022,56 +2522,64 @@ def viz_2D_topomap_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, return ax -def viz_2D_topomap_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndarray, - threshold: float=0.95, steps: int=2, lab: bool = False, - children: bool=False, child_head: bool=False): +def viz_2D_topomap_intra( + epo1: mne.Epochs, + epo2: mne.Epochs, + C1: np.ndarray, + C2: np.ndarray, + threshold: float = 0.95, + steps: int = 2, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Create a 2D topographic map visualization with intra-brain connectivity. - + This function provides a high-level interface for creating a 2D topographic visualization of two head outlines side by side, displaying intra-brain connectivity within each participant separately. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C1 : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_1) containing connectivity values between pairs of electrodes for participant 1 - + C2 : np.ndarray Connectivity matrix with shape (n_channels_2, n_channels_2) containing connectivity values between pairs of electrodes for participant 2 - + threshold : float, optional Threshold for displaying connections (default=0.95) Only connections with absolute value above this threshold are displayed - + steps : int, optional Number of steps for bezier curve interpolation (default=2) If steps < 3, straight lines are drawn instead of curves - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- ax : matplotlib.axes.Axes The matplotlib Axes object containing the visualization - + Notes ----- This function combines several lower-level visualization functions: @@ -2079,11 +2587,11 @@ def viz_2D_topomap_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: 2. Draws 2D head outlines using plot_2d_topomap_intra() 3. Plots sensors using plot_sensors_2d_intra() 4. Plots connectivity links using plot_links_2d_intra() - + Unlike viz_2D_topomap_inter() which visualizes connections between participants, this function displays connections within each participant's brain separately, allowing for comparison of intra-brain connectivity patterns. - + Examples -------- >>> import numpy as np @@ -2110,7 +2618,7 @@ def viz_2D_topomap_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: fig = plt.figure() ax = fig.add_subplot(111, aspect=1) ax.axis("off") - + if children: plot_2d_topomap_intra(ax, children=True) # bads are represented as squares @@ -2118,63 +2626,79 @@ def viz_2D_topomap_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: # plotting links according to sign (red for positive values, # blue for negative) and value (line thickness increases # with the strength of connectivity) - plot_links_2d_intra(epo1, epo2, C1=C1, C2=C2, threshold=threshold, steps=steps, children=True, child_head=child_head) + plot_links_2d_intra( + epo1, + epo2, + C1=C1, + C2=C2, + threshold=threshold, + steps=steps, + children=True, + child_head=child_head, + ) else: plot_2d_topomap_intra(ax) plot_sensors_2d_intra(epo1, epo2, lab=lab) plot_links_2d_intra(epo1, epo2, C1=C1, C2=C2, threshold=threshold, steps=steps) - + plt.tight_layout() plt.show() return ax -def viz_3D_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, - threshold: float=0.95, steps: int=10, lab: bool = False, - children: bool=False, child_head: bool=False): +def viz_3D_inter( + epo1: mne.Epochs, + epo2: mne.Epochs, + C: np.ndarray, + threshold: float = 0.95, + steps: int = 10, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Create a complete 3D visualization with inter-brain connectivity. - + This function provides a high-level interface for creating a 3D visualization of two head models with sensors and inter-brain connectivity links. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_2) containing connectivity values between all pairs of electrodes across participants - + threshold : float, optional Threshold for displaying connections (default=0.95) Only connections with absolute value above this threshold are displayed - + steps : int, optional Number of steps for bezier curve interpolation (default=10) If steps < 3, straight lines are drawn instead of curves - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- ax : matplotlib.axes.Axes The matplotlib 3D Axes object containing the visualization - + Notes ----- This function combines several lower-level visualization functions: @@ -2183,11 +2707,11 @@ def viz_3D_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, 3. Plots 3D head models using plot_3d_heads() 4. Plots sensors using plot_sensors_3d_inter() 5. Plots connectivity links using plot_links_3d_inter() - + The 3D visualization allows for interactive rotation and zoom, providing a more comprehensive view of the connectivity patterns compared to 2D visualizations. - + Examples -------- >>> import numpy as np @@ -2214,13 +2738,24 @@ def viz_3D_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, vertices, faces = get_3d_heads_inter() fig = plt.figure() - ax = fig.add_subplot(projection='3d') + ax = fig.add_subplot(projection="3d") ax.axis("off") plot_3d_heads(ax, vertices, faces) if children: - plot_sensors_3d_inter(ax, epo1, epo2, lab=lab, children=True, child_head=child_head) - plot_links_3d_inter(ax, epo1, epo2, C=C, threshold=threshold, steps=steps, children=True, child_head=child_head) + plot_sensors_3d_inter( + ax, epo1, epo2, lab=lab, children=True, child_head=child_head + ) + plot_links_3d_inter( + ax, + epo1, + epo2, + C=C, + threshold=threshold, + steps=steps, + children=True, + child_head=child_head, + ) else: plot_sensors_3d_inter(ax, epo1, epo2, lab=lab) plot_links_3d_inter(ax, epo1, epo2, C=C, threshold=threshold, steps=steps) @@ -2231,56 +2766,64 @@ def viz_3D_inter(epo1: mne.Epochs, epo2: mne.Epochs, C: np.ndarray, return ax -def viz_3D_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndarray, - threshold: float=0.95, steps: int=10, lab: bool = False, - children: bool=False, child_head: bool=False): +def viz_3D_intra( + epo1: mne.Epochs, + epo2: mne.Epochs, + C1: np.ndarray, + C2: np.ndarray, + threshold: float = 0.95, + steps: int = 10, + lab: bool = False, + children: bool = False, + child_head: bool = False, +): """ Create a complete 3D visualization with intra-brain connectivity. - + This function provides a high-level interface for creating a 3D visualization of two head models side by side, displaying intra-brain connectivity within each participant separately. - + Parameters ---------- epo1 : mne.Epochs Epochs object for participant 1, containing channel information - + epo2 : mne.Epochs Epochs object for participant 2, containing channel information - + C1 : np.ndarray Connectivity matrix with shape (n_channels_1, n_channels_1) containing connectivity values between pairs of electrodes for participant 1 - + C2 : np.ndarray Connectivity matrix with shape (n_channels_2, n_channels_2) containing connectivity values between pairs of electrodes for participant 2 - + threshold : float, optional Threshold for displaying connections (default=0.95) Only connections with absolute value above this threshold are displayed - + steps : int, optional Number of steps for bezier curve interpolation (default=10) If steps < 3, straight lines are drawn instead of curves - + lab : bool, optional Whether to plot channel labels (default=False) - + children : bool, optional Whether to apply transformations for child head models (default=False) - + child_head : bool, optional If True and children=True, treats participant 2 as a child If False and children=True, treats participant 1 as a child (default=False) - + Returns ------- ax : matplotlib.axes.Axes The matplotlib 3D Axes object containing the visualization - + Notes ----- This function combines several lower-level visualization functions: @@ -2289,11 +2832,11 @@ def viz_3D_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndar 3. Plots 3D head models using plot_3d_heads() 4. Plots sensors using plot_sensors_3d_intra() 5. Plots connectivity links using plot_links_3d_intra() - + Unlike viz_3D_inter() which visualizes connections between participants, this function displays connections within each participant's brain separately, allowing for comparison of intra-brain connectivity patterns in 3D. - + Examples -------- >>> import numpy as np @@ -2322,20 +2865,34 @@ def viz_3D_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndar vertices, faces = get_3d_heads_intra() fig = plt.figure() - ax = fig.add_subplot(projection='3d') + ax = fig.add_subplot(projection="3d") ax.axis("off") plot_3d_heads(ax, vertices, faces) if children: # bads are represented as squares - plot_sensors_3d_intra(ax, epo1, epo2, lab=lab, children=True, child_head=child_head) + plot_sensors_3d_intra( + ax, epo1, epo2, lab=lab, children=True, child_head=child_head + ) # plotting links according to sign (red for positive values, # blue for negative) and value (line thickness increases # with the strength of connectivity) - plot_links_3d_intra(ax, epo1, epo2, C1=C1, C2=C2, threshold=threshold, steps=steps, children=True, child_head=child_head) + plot_links_3d_intra( + ax, + epo1, + epo2, + C1=C1, + C2=C2, + threshold=threshold, + steps=steps, + children=True, + child_head=child_head, + ) else: plot_sensors_3d_intra(ax, epo1, epo2, lab=lab) - plot_links_3d_intra(ax, epo1, epo2, C1=C1, C2=C2, threshold=threshold, steps=steps) + plot_links_3d_intra( + ax, epo1, epo2, C1=C1, C2=C2, threshold=threshold, steps=steps + ) plt.tight_layout() plt.show() @@ -2343,56 +2900,61 @@ def viz_3D_intra(epo1: mne.Epochs, epo2: mne.Epochs, C1: np.ndarray, C2: np.ndar return ax -def plot_xwt(sig1: mne.Epochs, sig2: mne.Epochs, - sfreq: Union[int, float], - freqs: Union[int, float, np.ndarray], - time: int, analysis: str, - figsize: tuple = (30, 8), tmin: int = 0, - x_units: Union[int, float] = 100): +def plot_xwt( + sig1: mne.Epochs, + sig2: mne.Epochs, + sfreq: Union[int, float], + freqs: Union[int, float, np.ndarray], + time: int, + analysis: str, + figsize: tuple = (30, 8), + tmin: int = 0, + x_units: Union[int, float] = 100, +): """ Plot the results of Cross-Wavelet Transform (XWT) analysis. - + This function visualizes time-frequency representations of cross-wavelet transform results between two participants' EEG signals. It can display phase angles, power, or wavelet coherence. - + Parameters ---------- sig1 : mne.Epochs EEG data for the first participant - + sig2 : mne.Epochs EEG data for the second participant - + sfreq : int or float Sampling frequency of the EEG data in Hz - + freqs : int, float, or np.ndarray Frequencies of interest in Hz - + time : int Time duration of the sample in seconds - + analysis : str Type of analysis to visualize, options: - 'phase': Phase angle differences - 'power': Cross-wavelet power - 'wtc': Wavelet coherence - + figsize : tuple, optional Figure size in inches (default=(30, 8)) - + tmin : int, optional Start time for the plot in seconds (default=0) - + x_units : int or float, optional Distance between x-axis ticks in data points (default=100) - + Returns ------- fig : matplotlib.figure.Figure The matplotlib Figure object containing the visualization - + Notes ----- The function: @@ -2400,15 +2962,15 @@ def plot_xwt(sig1: mne.Epochs, sig2: mne.Epochs, 2. Computes the cone of influence (COI) to indicate regions with edge effects 3. Creates a colormap-based visualization of the XWT results 4. Masks regions outside the COI with hatching to indicate less reliable values - + Different colormaps are used based on the analysis type: - 'phase': 'hsv' colormap for phase angles - 'power': 'viridis' colormap for power values - 'wtc': 'plasma' colormap for coherence values - - This function is not meant to be called independently but is usually called + + This function is not meant to be called independently but is usually called from higher-level functions like plot_xwt_crosspower or plot_xwt_phase_angle. - + Examples -------- >>> # Plot cross-wavelet power @@ -2422,7 +2984,7 @@ def plot_xwt(sig1: mne.Epochs, sig2: mne.Epochs, ... figsize=(12, 8) ... ) >>> plt.show() - >>> + >>> >>> # Plot wavelet coherence >>> fig = plot_xwt( ... epochs_subj1, @@ -2436,16 +2998,16 @@ def plot_xwt(sig1: mne.Epochs, sig2: mne.Epochs, >>> plt.show() """ - dt = 1/sfreq - xmax = time/dt + dt = 1 / sfreq + xmax = time / dt xmin = tmin * dt - tick_units = xmax/x_units - unit_conv = time/x_units + tick_units = xmax / x_units + unit_conv = time / x_units xticks = np.arange(xmin, xmax, tick_units) x_labels = np.arange(tmin, time, unit_conv) xmark = [] for xlabel in x_labels: - if xlabel != '': + if xlabel != "": mark = str(round(xlabel, 3)) xmark.append(mark) else: @@ -2453,71 +3015,70 @@ def plot_xwt(sig1: mne.Epochs, sig2: mne.Epochs, coi = [] for f in freqs: - dt_f = 1/f - f_coi_init = math.sqrt(2*2*dt_f) - f_coi = f_coi_init/dt + dt_f = 1 / f + f_coi_init = math.sqrt(2 * 2 * dt_f) + f_coi = f_coi_init / dt coi.append(f_coi) coi_check = [] for item in coi: - if item >= (time/dt): + if item >= (time / dt): coi_check.append(False) else: coi_check.append(True) if False in coi_check: - print('Warning: your epoch length seems too short for the wavelet transform!') + print("Warning: your epoch length seems too short for the wavelet transform!") else: - print('Epoch length is appropriate for wavelet transform') + print("Epoch length is appropriate for wavelet transform") coi_index = np.arange(0, len(freqs)) rev_coi = [] for f in freqs: - dt_f = 1/f - f_coi = math.sqrt(2*2*dt_f) - sub_max = (time-f_coi)/dt + dt_f = 1 / f + f_coi = math.sqrt(2 * 2 * dt_f) + sub_max = (time - f_coi) / dt rev_coi.append(sub_max) fig = plt.figure(figsize=figsize) plt.subplot(122) - if analysis == 'phase': - data = xwt(sig1, sig2, freqs, mode='phase') - analysis_title = 'Cross Wavelet Transform (Phase Angle)' - cbar_title = 'Phase Difference' - my_cm = matplotlib.colormaps['hsv'] + if analysis == "phase": + data = xwt(sig1, sig2, freqs, mode="phase") + analysis_title = "Cross Wavelet Transform (Phase Angle)" + cbar_title = "Phase Difference" + my_cm = matplotlib.colormaps["hsv"] - plt.imshow(data, aspect='auto', cmap=my_cm, interpolation='nearest') + plt.imshow(data, aspect="auto", cmap=my_cm, interpolation="nearest") - elif analysis == 'power': - data = xwt(sig1, sig2, freqs, mode='power') + elif analysis == "power": + data = xwt(sig1, sig2, freqs, mode="power") normed_data = (data - np.min(data)) / (np.max(data) - np.min(data)) - analysis_title = 'Cross Wavelet Transform (Power)' - cbar_title = 'Cross Power' - my_cm = matplotlib.colormaps['viridis'] + analysis_title = "Cross Wavelet Transform (Power)" + cbar_title = "Cross Power" + my_cm = matplotlib.colormaps["viridis"] - plt.imshow(normed_data, aspect='auto', cmap=my_cm, - interpolation='lanczos') + plt.imshow(normed_data, aspect="auto", cmap=my_cm, interpolation="lanczos") - elif analysis == 'wtc': - data = xwt(sig1, sig2, freqs, mode='wtc') - analysis_title = 'Wavelet Coherence' - cbar_title = 'Coherence' - my_cm = matplotlib.colormaps['plasma'] + elif analysis == "wtc": + data = xwt(sig1, sig2, freqs, mode="wtc") + analysis_title = "Wavelet Coherence" + cbar_title = "Coherence" + my_cm = matplotlib.colormaps["plasma"] - plt.imshow(data, aspect='auto', cmap=my_cm, interpolation='lanczos') + plt.imshow(data, aspect="auto", cmap=my_cm, interpolation="lanczos") else: - raise ValueError('Analysis must be set as phase, power, or wtc.') + raise ValueError("Analysis must be set as phase, power, or wtc.") plt.gca().invert_yaxis() - plt.ylabel('Frequencies (Hz)') - plt.xlabel('Time (s)') + plt.ylabel("Frequencies (Hz)") + plt.xlabel("Time (s)") ylabels = np.linspace((freqs[0]), (freqs[-1]), len(freqs[0:])) ymark = [] for ylabel in ylabels: - if ylabel != '': + if ylabel != "": mark = str(round(ylabel, 3)) ymark.append(mark) else: @@ -2527,16 +3088,16 @@ def plot_xwt(sig1: mne.Epochs, sig2: mne.Epochs, plt.gca().set_xticks(ticks=xticks, labels=xmark) plt.xlim(tmin, xmax) plt.ylim(0, int(len(ylabels[0:-1]))) - plt.plot(coi, coi_index, 'w') - plt.plot(rev_coi, coi_index, 'w') + plt.plot(coi, coi_index, "w") + plt.plot(rev_coi, coi_index, "w") - plt.fill_between(coi, coi_index, hatch='X', fc='w', alpha=0.5) - plt.fill_between(rev_coi, coi_index, hatch='X', fc='w', alpha=0.5) + plt.fill_between(coi, coi_index, hatch="X", fc="w", alpha=0.5) + plt.fill_between(rev_coi, coi_index, hatch="X", fc="w", alpha=0.5) - plt.axvspan(xmin, min(coi), hatch='X', fc='w', alpha=0.5) - plt.axvspan(xmax, max(rev_coi), hatch='X', fc='w', alpha=0.5) + plt.axvspan(xmin, min(coi), hatch="X", fc="w", alpha=0.5) + plt.axvspan(xmax, max(rev_coi), hatch="X", fc="w", alpha=0.5) plt.title(analysis_title) plt.colorbar(label=cbar_title) - return fig \ No newline at end of file + return fig diff --git a/hypyp/wavelet/__init__.py b/hypyp/wavelet/__init__.py index 4087097..7863835 100644 --- a/hypyp/wavelet/__init__.py +++ b/hypyp/wavelet/__init__.py @@ -1,17 +1,20 @@ from .base_wavelet import BaseWavelet -from .implementations.pywavelets_wavelet import ComplexMorletWavelet, ComplexGaussianWavelet +from .implementations.pywavelets_wavelet import ( + ComplexMorletWavelet, + ComplexGaussianWavelet, +) from .coherence_data_frame import CoherenceDataFrame from .cwt import CWT from .pair_signals import PairSignals from .wtc import WTC __all__ = [ - 'BaseWavelet', + "BaseWavelet", # Only expose wavelets from pywavelets_wavelet to avoid confusion. This is the default implementation. - 'ComplexMorletWavelet', - 'ComplexGaussianWavelet', - 'CoherenceDataFrame', - 'CWT', - 'PairSignals', - 'WTC', + "ComplexMorletWavelet", + "ComplexGaussianWavelet", + "CoherenceDataFrame", + "CWT", + "PairSignals", + "WTC", ] diff --git a/hypyp/wavelet/base_wavelet.py b/hypyp/wavelet/base_wavelet.py index 0fadbee..9bc8987 100644 --- a/hypyp/wavelet/base_wavelet.py +++ b/hypyp/wavelet/base_wavelet.py @@ -19,7 +19,8 @@ DEFAULT_SMOOTH_WIN_SIZE = 0.6 DEFAULT_PERIOD_RANGE = (2, 20) -DEFAULT_PERIOD_DJ = 1/12 +DEFAULT_PERIOD_DJ = 1 / 12 + class BaseWavelet(ABC): """ @@ -40,7 +41,7 @@ class BaseWavelet(ABC): dj: float wtc_smoothing_win_size: float period_range: Tuple[float, float] - cache: dict|None + cache: dict | None cache_is_disabled: bool default_period_range: Tuple[float, float] = DEFAULT_PERIOD_RANGE @@ -49,14 +50,14 @@ class BaseWavelet(ABC): def __init__( self, - period_range:Tuple[float, float]|None=None, - frequency_range:Tuple[float, float]|None=None, - evaluate:bool=True, - cache:dict|None=None, - disable_caching:bool=False, - verbose:bool=False, - dj:float=DEFAULT_PERIOD_DJ, - wtc_smoothing_win_size:float=DEFAULT_SMOOTH_WIN_SIZE, + period_range: Tuple[float, float] | None = None, + frequency_range: Tuple[float, float] | None = None, + evaluate: bool = True, + cache: dict | None = None, + disable_caching: bool = False, + verbose: bool = False, + dj: float = DEFAULT_PERIOD_DJ, + wtc_smoothing_win_size: float = DEFAULT_SMOOTH_WIN_SIZE, ): self._wavelet = None self._psi_x = None @@ -67,11 +68,11 @@ def __init__( self.wtc_smoothing_win_size = wtc_smoothing_win_size if dj > 1: - raise RuntimeError(f'dj must be smaller than 1. Received {dj}') + raise RuntimeError(f"dj must be smaller than 1. Received {dj}") self.dj = dj if period_range is not None and frequency_range is not None: - raise RuntimeError('Cannot specify both period_range and frequency_range') + raise RuntimeError("Cannot specify both period_range and frequency_range") if period_range is not None: self.period_range = period_range @@ -79,11 +80,11 @@ def __init__( self.period_range = (1 / frequency_range[0], 1 / frequency_range[1]) else: self.period_range = DEFAULT_PERIOD_RANGE - + # swap to have them in order if self.period_range[0] > self.period_range[1]: self.period_range = (self.period_range[1], self.period_range[0]) - + self.cache = cache self.cache_is_disabled = disable_caching @@ -94,39 +95,39 @@ def __init__( if evaluate: self.evaluate_psi() - + @property def psi(self): return self._psi - + @property def psi_x(self): return self._psi_x - + @property def psi_dx(self): if self._psi_x is None: - raise RuntimeError('Wavelet not evaluated yet') + raise RuntimeError("Wavelet not evaluated yet") return self._psi_x[1] - self._psi_x[0] @property def domain(self): if self._psi_x is None: - raise RuntimeError('Wavelet not evaluated yet') + raise RuntimeError("Wavelet not evaluated yet") return self._psi_x[0], self._psi_x[-1] @property def use_caching(self): return not self.cache_is_disabled and self.cache is not None - + @abstractmethod def evaluate_psi(self): pass - + @abstractmethod - def cwt(self, y, dt, label='') -> CWT: + def cwt(self, y, dt, label="") -> CWT: pass - + @property @abstractmethod def wavelet_library(self): @@ -136,12 +137,12 @@ def wavelet_library(self): @abstractmethod def wavelet_name(self): pass - + @property @abstractmethod def wavelet_name_with_args(self): pass - + @property @abstractmethod def flambda(self): @@ -154,26 +155,28 @@ def get_periods(self) -> np.ndarray: Returns: np.array: list of periods that should be used with this wavelet """ - low, high = self.period_range - n_scales = np.log2(high/low) + low, high = self.period_range + n_scales = np.log2(high / low) n_steps = int(np.round(n_scales / self.dj)) periods = np.logspace(np.log2(low), np.log2(high), n_steps, base=2) return periods - - def wtc(self, - pair: PairSignals, - bin_seconds:float|None=None, - period_cuts:List[float]|None=None, - frequency_cuts:List[float]|None=None, - cache_suffix:str='') -> WTC: + + def wtc( + self, + pair: PairSignals, + bin_seconds: float | None = None, + period_cuts: List[float] | None = None, + frequency_cuts: List[float] | None = None, + cache_suffix: str = "", + ) -> WTC: """ - Compute the Wavalet Transform Coherence for a pair of signals. - + Compute the Wavalet Transform Coherence for a pair of signals. + First computes the Continuous Wavelet Transform of both signals, then filters them and finally the coherence between the 2 signals. It saves to a cache in memory the intermediary results for reuse on other pairs. - For example, if there is a pair made of signal A and signal B and then a pair with + For example, if there is a pair made of signal A and signal B and then a pair with signal A and signal C, we don't have to re-compute the CWT of signal A twice. Args: @@ -191,7 +194,7 @@ def wtc(self, dt = pair.dt if len(y1) != len(y2): - raise RuntimeError(f'Arrays not same size. y1:{len(y1)}, y2:{len(y2)}') + raise RuntimeError(f"Arrays not same size. y1:{len(y1)}, y2:{len(y2)}") N = len(y1) @@ -203,18 +206,30 @@ def wtc(self, # TODO: maybe this should be in preprocessing instead y1 = (y1 - y1.mean()) / y1.std() y2 = (y2 - y2.mean()) / y2.std() - - cwt1_cached = self._get_cache_item(self._get_cache_key_pair(pair, 0, 'cwt', cache_suffix)) - cwt1: CWT = self.cwt(y1, dt, label=pair.label_ch1) if cwt1_cached is None else cwt1_cached - cwt2_cached = self._get_cache_item(self._get_cache_key_pair(pair, 1, 'cwt', cache_suffix)) - cwt2: CWT = self.cwt(y2, dt, label=pair.label_ch2) if cwt2_cached is None else cwt2_cached + cwt1_cached = self._get_cache_item( + self._get_cache_key_pair(pair, 0, "cwt", cache_suffix) + ) + cwt1: CWT = ( + self.cwt(y1, dt, label=pair.label_ch1) + if cwt1_cached is None + else cwt1_cached + ) + + cwt2_cached = self._get_cache_item( + self._get_cache_key_pair(pair, 1, "cwt", cache_suffix) + ) + cwt2: CWT = ( + self.cwt(y2, dt, label=pair.label_ch2) + if cwt2_cached is None + else cwt2_cached + ) if (cwt1.scales != cwt2.scales).any(): - raise RuntimeError('The two CWT have different scales') + raise RuntimeError("The two CWT have different scales") if (cwt1.frequencies != cwt2.frequencies).any(): - raise RuntimeError('The two CWT have different frequencies') + raise RuntimeError("The two CWT have different frequencies") W1 = cwt1.W W2 = cwt2.W @@ -223,7 +238,9 @@ def wtc(self, try: W12 = W1 * W2.conj() except ValueError as e: - warnings.warn("Wrong operand shapes could mean that a wrong cached value is used. Please check that cache keys contain all the relevant arguments") + warnings.warn( + "Wrong operand shapes could mean that a wrong cached value is used. Please check that cache keys contain all the relevant arguments" + ) raise e periods = cwt1.periods @@ -237,21 +254,41 @@ def wtc(self, cache_suffix=cache_suffix, ) - S1_cached = self._get_cache_item(self._get_cache_key_pair(pair, 0, 'smooth', cache_suffix)) - S1 = self.smoothing(np.abs(W1) ** 2 / scaleMatrix, **smoothing_kwargs) if S1_cached is None else S1_cached + S1_cached = self._get_cache_item( + self._get_cache_key_pair(pair, 0, "smooth", cache_suffix) + ) + S1 = ( + self.smoothing(np.abs(W1) ** 2 / scaleMatrix, **smoothing_kwargs) + if S1_cached is None + else S1_cached + ) - S2_cached = self._get_cache_item(self._get_cache_key_pair(pair, 1, 'smooth', cache_suffix)) - S2 = self.smoothing(np.abs(W2) ** 2 / scaleMatrix, **smoothing_kwargs) if S2_cached is None else S2_cached + S2_cached = self._get_cache_item( + self._get_cache_key_pair(pair, 1, "smooth", cache_suffix) + ) + S2 = ( + self.smoothing(np.abs(W2) ** 2 / scaleMatrix, **smoothing_kwargs) + if S2_cached is None + else S2_cached + ) S12 = np.abs(self.smoothing(W12 / scaleMatrix, **smoothing_kwargs)) - wtc = S12 ** 2 / (S1 * S2) + wtc = S12**2 / (S1 * S2) coi = self._get_and_cache_cone_of_influence(N, dt, cache_suffix=cache_suffix) - self._update_cache_if_none(self._get_cache_key_pair(pair, 0, 'cwt', cache_suffix), cwt1) - self._update_cache_if_none(self._get_cache_key_pair(pair, 1, 'cwt', cache_suffix), cwt2) - self._update_cache_if_none(self._get_cache_key_pair(pair, 0, 'smooth', cache_suffix), S1) - self._update_cache_if_none(self._get_cache_key_pair(pair, 1, 'smooth', cache_suffix), S2) + self._update_cache_if_none( + self._get_cache_key_pair(pair, 0, "cwt", cache_suffix), cwt1 + ) + self._update_cache_if_none( + self._get_cache_key_pair(pair, 1, "cwt", cache_suffix), cwt2 + ) + self._update_cache_if_none( + self._get_cache_key_pair(pair, 0, "smooth", cache_suffix), S1 + ) + self._update_cache_if_none( + self._get_cache_key_pair(pair, 1, "smooth", cache_suffix), S2 + ) return WTC( wtc, @@ -275,8 +312,8 @@ def _update_cache_if_none(self, key, value): if self._get_cache_item(key) is None: self._add_cache_item(key, value) - def _get_and_cache_cone_of_influence(self, N, dt, cache_suffix=''): - cache_key = self._get_cache_key('coi', N, dt, cache_suffix) + def _get_and_cache_cone_of_influence(self, N, dt, cache_suffix=""): + cache_key = self._get_cache_key("coi", N, dt, cache_suffix) coi = self._get_cache_item(cache_key) if coi is not None: return coi @@ -284,10 +321,9 @@ def _get_and_cache_cone_of_influence(self, N, dt, cache_suffix=''): self._add_cache_item(cache_key, coi) return coi - - def _get_cone_of_influence(self, N:int, dt:float) -> np.ndarray: + def _get_cone_of_influence(self, N: int, dt: float) -> np.ndarray: """ - Get the cone of influence of a Continuous Wavelet Transform, + Get the cone of influence of a Continuous Wavelet Transform, that is the region where edge effects become significant due to the finite length of the signal Args: @@ -303,15 +339,17 @@ def _get_cone_of_influence(self, N:int, dt:float) -> np.ndarray: e_folding_time = 1.0 / np.sqrt(2) flambda = self.flambda - coi = (N / 2 - np.abs(np.arange(0, N) - (N - 1) / 2)) + coi = N / 2 - np.abs(np.arange(0, N) - (N - 1) / 2) coi = flambda * e_folding_time * dt * coi return coi - + # # Smoothing # - def smoothing(self, W:np.ndarray, dt:float, scales:np.ndarray, cache_suffix:str='') -> np.ndarray: + def smoothing( + self, W: np.ndarray, dt: float, scales: np.ndarray, cache_suffix: str = "" + ) -> np.ndarray: """ Smooth the weights of a continuous wavelet transform in time and frequency domains. @@ -328,48 +366,54 @@ def smoothing(self, W:np.ndarray, dt:float, scales:np.ndarray, cache_suffix:str= # Filter in time. # fft_kwargs = self._get_fft_kwargs(W[0, :]) - scales_norm = scales # scales are already normalized - - k = 2 * np.pi * fft.fftfreq(fft_kwargs['n']) - k2 = k ** 2 + scales_norm = scales # scales are already normalized + + k = 2 * np.pi * fft.fftfreq(fft_kwargs["n"]) + k2 = k**2 # Smoothing by Gaussian window (absolute value of wavelet function) # using the convolution theorem: multiplication by Gaussian curve in # Fourier domain for each scale, outer product of scale and frequency - gaus_fft = np.exp(-0.5 * (scales_norm[:, np.newaxis] ** 2) * k2) # Outer product + gaus_fft = np.exp( + -0.5 * (scales_norm[:, np.newaxis] ** 2) * k2 + ) # Outer product W_fft = fft.fft(W, axis=1, **fft_kwargs) smooth_fft = gaus_fft * W_fft smooth = fft.ifft(smooth_fft, axis=1, **fft_kwargs, overwrite_x=True) - T = smooth[:, :W.shape[1]] # Remove possibly padded region due to FFT + T = smooth[:, : W.shape[1]] # Remove possibly padded region due to FFT if np.isreal(W).all(): T = T.real - + + # + # Filter in scale. # - # Filter in scale. - # - win_cache_key = self._get_cache_key('boxcar_window', self.wtc_smoothing_win_size, self.dj, cache_suffix) + win_cache_key = self._get_cache_key( + "boxcar_window", self.wtc_smoothing_win_size, self.dj, cache_suffix + ) win = self._get_cache_item(win_cache_key) if win is None: win = self._get_smoothing_window(self.wtc_smoothing_win_size, self.dj) self._add_cache_item(win_cache_key, win) - T = convolve1d(T, win, axis=0, mode='nearest') - #T = convolve2d(T, win[:, np.newaxis], 'same') # Scales are "vertical" + T = convolve1d(T, win, axis=0, mode="nearest") + # T = convolve2d(T, win[:, np.newaxis], 'same') # Scales are "vertical" return T @staticmethod def _get_fft_kwargs(signal, **kwargs): - return dict(**kwargs, n = int(2 ** np.ceil(np.log2(len(signal))))) + return dict(**kwargs, n=int(2 ** np.ceil(np.log2(len(signal))))) @staticmethod def _get_smoothing_window(boxcar_size, dj): # Copied from matlab # boxcar_size is "in scale" - #size_in_scales = boxcar_size - size_in_scales = boxcar_size * 2 # the pycwt code has a *2. TODO: find out why it is different from matlab code + # size_in_scales = boxcar_size + size_in_scales = ( + boxcar_size * 2 + ) # the pycwt code has a *2. TODO: find out why it is different from matlab code size_in_steps = size_in_scales / dj fraction = size_in_steps % 1 fraction_half = fraction / 2 @@ -389,8 +433,6 @@ def _get_smoothing_window(boxcar_size, dj): # normalize win /= win.sum() return win - - # # Caching @@ -405,15 +447,15 @@ def _get_cache_item(self, key): try: found = self.cache[key] if self.verbose: - print(f'Found cache key {key}') + print(f"Found cache key {key}") pass return found except: if self.verbose: - #print(f'Not found cache key {key}') + # print(f'Not found cache key {key}') pass return None - + def _add_cache_item(self, key, value): if self.use_caching: self.cache[key] = value @@ -421,45 +463,46 @@ def _add_cache_item(self, key, value): def clear_cache(self): self.cache = dict() - def _get_cache_key_pair(self, pair: PairSignals, subject_idx: int, obj_id: str, cache_suffix: str = ''): + def _get_cache_key_pair( + self, pair: PairSignals, subject_idx: int, obj_id: str, cache_suffix: str = "" + ): if not self.use_caching: return None if subject_idx == 0: subject_label = pair.label_s1 ch_name = pair.label_ch1 - # include signal length in key because it might have been croped to align y1 and y2 + # include signal length in key because it might have been croped to align y1 and y2 n = len(pair.y1) elif subject_idx == 1: subject_label = pair.label_s2 ch_name = pair.label_ch2 n = len(pair.y2) else: - raise RuntimeError(f'subject_idx must be 0 or 1') - + raise RuntimeError(f"subject_idx must be 0 or 1") + if len(subject_label) < 1: - raise RuntimeError(f'subjects must have labels to use caching') + raise RuntimeError(f"subjects must have labels to use caching") - if pair.label_task == '': - raise RuntimeError(f'must have task to have unique identifiers in caching') + if pair.label_task == "": + raise RuntimeError(f"must have task to have unique identifiers in caching") + + time_range_str = f"{pair.x[0]}-{pair.x[-1]}" + key = f"[{subject_label}][{ch_name}][{pair.label_task}][{pair.epoch_idx}][{n}][{pair.x[0]}-{pair.x[-1]}][{time_range_str}][{obj_id}]" + if cache_suffix != "": + key += f"_{cache_suffix}" - time_range_str = f'{pair.x[0]}-{pair.x[-1]}' - key = f'[{subject_label}][{ch_name}][{pair.label_task}][{pair.epoch_idx}][{n}][{pair.x[0]}-{pair.x[-1]}][{time_range_str}][{obj_id}]' - if cache_suffix != '': - key += f'_{cache_suffix}' - return key - + def _get_cache_key(self, *args): if not self.use_caching: return None - return f'key_{"_".join([str(arg) for arg in args])}' - + return f"key_{'_'.join([str(arg) for arg in args])}" # # Plots # - def plot_mother_wavelet(self, show_legend=True, ax:Axes|None=None) -> Figure: + def plot_mother_wavelet(self, show_legend=True, ax: Axes | None = None) -> Figure: if ax is None: fig, ax = plt.subplots() else: @@ -470,5 +513,5 @@ def plot_mother_wavelet(self, show_legend=True, ax:Axes|None=None) -> Figure: ax.plot(self.psi_x, np.abs(self.psi)) ax.title.set_text(f"mother wavelet ({self.wavelet_name_with_args})") if show_legend: - ax.legend(['real', 'imag', 'abs']) + ax.legend(["real", "imag", "abs"]) return fig diff --git a/hypyp/wavelet/coherence_data_frame.py b/hypyp/wavelet/coherence_data_frame.py index b20a364..7cd8fc9 100644 --- a/hypyp/wavelet/coherence_data_frame.py +++ b/hypyp/wavelet/coherence_data_frame.py @@ -4,28 +4,29 @@ import pyarrow.feather as feather COHERENCE_FRAME_COLUMNS = [ - 'dyad', - 'is_intra', - 'is_intra_of', - 'is_pseudo', - 'subject1', - 'subject2', - 'roi1', - 'roi2', - 'channel1', - 'channel2', - 'task', - 'epoch', - 'section', - 'bin', - 'coherence', - 'coherence_masked', - 'bin_time_range', - 'bin_period_range', - 'wavelet_library', - 'wavelet_name', + "dyad", + "is_intra", + "is_intra_of", + "is_pseudo", + "subject1", + "subject2", + "roi1", + "roi2", + "channel1", + "channel2", + "task", + "epoch", + "section", + "bin", + "coherence", + "coherence_masked", + "bin_time_range", + "bin_period_range", + "wavelet_library", + "wavelet_name", ] + class CoherenceDataFrame(TypedDict, total=False): # row properties dyad: pd.Categorical @@ -55,7 +56,7 @@ class CoherenceDataFrame(TypedDict, total=False): wavelet_name: pd.Categorical @staticmethod - def from_wtc_frame_rows(data:list): + def from_wtc_frame_rows(data: list): """ Get a typed pandas DataFrame from wavelet transform coherence data @@ -69,20 +70,19 @@ def from_wtc_frame_rows(data:list): data, columns=COHERENCE_FRAME_COLUMNS, ) - #CoherenceDataFrame._set_dtype_categories(df) + # CoherenceDataFrame._set_dtype_categories(df) return cast(CoherenceDataFrame, df) - + @staticmethod def _set_dtype_categories(df): - df['dyad'] = df['dyad'].astype('category') - df['task'] = df['task'].astype('category') - df['bin_time_range'] = df['bin_time_range'].astype('category') - df['bin_period_range'] = df['bin_period_range'].astype('category') - df['wavelet_library'] = df['wavelet_library'].astype('category') - df['wavelet_name'] = df['wavelet_name'].astype('category') - - + df["dyad"] = df["dyad"].astype("category") + df["task"] = df["task"].astype("category") + df["bin_time_range"] = df["bin_time_range"].astype("category") + df["bin_period_range"] = df["bin_period_range"].astype("category") + df["wavelet_library"] = df["wavelet_library"].astype("category") + df["wavelet_name"] = df["wavelet_name"].astype("category") + @staticmethod def concat(dfs: list[pd.DataFrame]): """ @@ -96,9 +96,9 @@ def concat(dfs: list[pd.DataFrame]): """ df = pd.concat(dfs, ignore_index=True) return cast(CoherenceDataFrame, df) - + @staticmethod - def from_feather(feather_path:str): + def from_feather(feather_path: str): """ Read pandas feather file and return as CoherenceDataFrame @@ -108,13 +108,13 @@ def from_feather(feather_path:str): Returns: CoherenceDataFrame: a typed pandas DataFrame """ - with open(feather_path, 'rb') as f: + with open(feather_path, "rb") as f: df = feather.read_feather(f) - + return cast(CoherenceDataFrame, df) - + @staticmethod - def save_feather(df:pd.DataFrame, feather_path:str): + def save_feather(df: pd.DataFrame, feather_path: str): """ Save to disk @@ -122,6 +122,5 @@ def save_feather(df:pd.DataFrame, feather_path:str): df (CoherenceDataFrame): the pandas dataframe to save feather_path (str): path on disk """ - with open(feather_path, 'wb') as f: + with open(feather_path, "wb") as f: feather.write_feather(df, f) - diff --git a/hypyp/wavelet/cwt.py b/hypyp/wavelet/cwt.py index abf5ae1..73d448d 100644 --- a/hypyp/wavelet/cwt.py +++ b/hypyp/wavelet/cwt.py @@ -3,6 +3,7 @@ from ..plots import plot_wavelet_transform_weights + class CWT: W: np.ndarray times: np.ndarray @@ -10,13 +11,21 @@ class CWT: scales: np.ndarray periods: np.ndarray frequencies: np.ndarray - coi: np.ndarray # Cone of influence, in periods - coif: np.ndarray # Cone of influence, in frequencies + coi: np.ndarray # Cone of influence, in periods + coif: np.ndarray # Cone of influence, in frequencies label: str - def __init__(self, weights:np.ndarray, times:np.ndarray, scales:np.ndarray, periods:np.ndarray, coi:np.ndarray, label:str=''): + def __init__( + self, + weights: np.ndarray, + times: np.ndarray, + scales: np.ndarray, + periods: np.ndarray, + coi: np.ndarray, + label: str = "", + ): """ - The CWT object holds the results of a Continuous Wavelet Transform + The CWT object holds the results of a Continuous Wavelet Transform Args: weights (np.ndarray): weights of the transforms @@ -46,6 +55,12 @@ def plot(self, **kwargs) -> Figure: Returns: Figure: matplotlib.Figure """ - return plot_wavelet_transform_weights(self.W, self.times, self.frequencies, self.coif, self.sfreq, title=f"CWT of {self.label}", **kwargs) - - + return plot_wavelet_transform_weights( + self.W, + self.times, + self.frequencies, + self.coif, + self.sfreq, + title=f"CWT of {self.label}", + **kwargs, + ) diff --git a/hypyp/wavelet/implementations/matlab_wavelet.py b/hypyp/wavelet/implementations/matlab_wavelet.py index e2d63de..1435406 100644 --- a/hypyp/wavelet/implementations/matlab_wavelet.py +++ b/hypyp/wavelet/implementations/matlab_wavelet.py @@ -33,58 +33,69 @@ def __init__( print("Matlab engine started") current_file_path = Path(__file__).resolve() - root_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))) - wavelet_coherence_code_path = os.path.join(root_path, 'wavelet-coherence') + root_path = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(current_file_path))) + ) + wavelet_coherence_code_path = os.path.join(root_path, "wavelet-coherence") self.eng.cd(wavelet_coherence_code_path) - print(f'Changed matlab working directory to {self.eng.pwd()}') + print(f"Changed matlab working directory to {self.eng.pwd()}") super().__init__(**kwargs) @property def wavelet_library(self): - return 'matlab' + return "matlab" @property def wavelet_name(self): - return 'cmor' + return "cmor" @property def wavelet_name_with_args(self): - return 'cmor' + return "cmor" @property def flambda(self): - raise RuntimeError('Not implemented') + raise RuntimeError("Not implemented") def evaluate_psi(self): self._psi_x = np.arange(10) - self._psi = np.zeros((10, )) + self._psi = np.zeros((10,)) return self._psi, self.psi_x - + def cwt(self, y, dt): - raise RuntimeError('Not implemented') + raise RuntimeError("Not implemented") - def wtc(self, pair: PairSignals, bin_seconds:float|None=None, period_cuts:List[float]|None=None, cache_suffix=''): + def wtc( + self, + pair: PairSignals, + bin_seconds: float | None = None, + period_cuts: List[float] | None = None, + cache_suffix="", + ): y1 = pair.y1 y2 = pair.y2 dt = pair.dt - self.eng.workspace['y1'] = np.ascontiguousarray(y1) - self.eng.workspace['y2'] = np.ascontiguousarray(y2) - self.eng.eval("[Rsq, period, scale, coi, sig] = wtc(y1, y2, 'mcc', 0, 'MakeFigure', 1, 'ArrowSize', 0.2, 'S0', 2, 'MaxScale', 1000); ", nargout=0) - - wtc = np.array(self.eng.workspace['Rsq']) + self.eng.workspace["y1"] = np.ascontiguousarray(y1) + self.eng.workspace["y2"] = np.ascontiguousarray(y2) + self.eng.eval( + "[Rsq, period, scale, coi, sig] = wtc(y1, y2, 'mcc', 0, 'MakeFigure', 1, 'ArrowSize', 0.2, 'S0', 2, 'MaxScale', 1000); ", + nargout=0, + ) + + wtc = np.array(self.eng.workspace["Rsq"]) N = len(y1) - times = np.linspace(0, N*dt, N) - periods = np.array(self.eng.workspace['period']).flatten() * dt - scales = np.array(self.eng.workspace['scale']).flatten() - coi = np.array(self.eng.workspace['coi']).flatten() + times = np.linspace(0, N * dt, N) + periods = np.array(self.eng.workspace["period"]).flatten() * dt + scales = np.array(self.eng.workspace["scale"]).flatten() + coi = np.array(self.eng.workspace["coi"]).flatten() # Remove the periods that are out of our period_range start = np.searchsorted(periods, self.period_range[0]) stop = np.searchsorted(periods, self.period_range[1]) periods = periods[start:stop] scales = scales[start:stop] - wtc = wtc[start:stop,:] + wtc = wtc[start:stop, :] return WTC( wtc, @@ -96,7 +107,8 @@ def wtc(self, pair: PairSignals, bin_seconds:float|None=None, period_cuts:List[f wavelet_library=self.wavelet_library, wavelet_name_with_args=self.wavelet_name_with_args, bin_seconds=bin_seconds, - period_cuts=period_cuts) + period_cuts=period_cuts, + ) except: MatlabWavelet = None diff --git a/hypyp/wavelet/implementations/pycwt_wavelet.py b/hypyp/wavelet/implementations/pycwt_wavelet.py index 42f6266..ae5e64d 100644 --- a/hypyp/wavelet/implementations/pycwt_wavelet.py +++ b/hypyp/wavelet/implementations/pycwt_wavelet.py @@ -33,49 +33,61 @@ def __init__( @property def wavelet_library(self): - return 'pycwt' + return "pycwt" @property def wavelet_name(self): - return f'cmor' + return f"cmor" @property def wavelet_name_with_args(self): - return f'cmor{self.center_frequency}' + return f"cmor{self.center_frequency}" @property def flambda(self): return self._wavelet.flambda() def evaluate_psi(self): - self._psi_x = np.linspace(self.lower_bound, self.upper_bound, 2**self.precision) + self._psi_x = np.linspace( + self.lower_bound, self.upper_bound, 2**self.precision + ) wavelet = pycwt.wavelet.Morlet() # TODO: with this parameter, the deltaj0 should change in pycwt code wavelet.f0 = self.center_frequency self._wavelet = wavelet self._psi = wavelet.psi(self._psi_x) return self._psi, self._psi_x - + def cwt(self, y, dt) -> CWT: periods = self.get_periods() - W, scales, freqs, coi, _, _ = pycwt.cwt(y, dt=dt, freqs=1/periods, wavelet=self._wavelet) + W, scales, freqs, coi, _, _ = pycwt.cwt( + y, dt=dt, freqs=1 / periods, wavelet=self._wavelet + ) periods = 1 / freqs times = np.arange(len(y)) * dt return CWT(weights=W, times=times, scales=scales, periods=periods, coi=coi) - def smoothing(self, W, dt, scales, _win_size=0, _cache_suffix=''): + def smoothing(self, W, dt, scales, _win_size=0, _cache_suffix=""): return self._wavelet.smooth(W, dt, self.dj, scales) - def wtc(self, pair: PairSignals, bin_seconds:float|None=None, period_cuts:List[float]|None=None, cache_suffix=''): + def wtc( + self, + pair: PairSignals, + bin_seconds: float | None = None, + period_cuts: List[float] | None = None, + cache_suffix="", + ): y1 = pair.y1 y2 = pair.y2 dt = pair.dt - assert (len(y1) == len(y2)), "error: arrays not same size" + assert len(y1) == len(y2), "error: arrays not same size" N = len(y1) times = np.arange(N) * dt - wtc, _awtc, coi, freqs, _sig = pycwt.wct(y1, y2, dt=dt, wavelet=self._wavelet, sig=False) + wtc, _awtc, coi, freqs, _sig = pycwt.wct( + y1, y2, dt=dt, wavelet=self._wavelet, sig=False + ) periods = 1 / freqs # Remove the periods that are out of our period_range @@ -84,7 +96,7 @@ def wtc(self, pair: PairSignals, bin_seconds:float|None=None, period_cuts:List[f stop = np.searchsorted(periods, self.period_range[1]) periods = periods[start:stop] freqs = freqs[start:stop] - wtc = wtc[start:stop,:] + wtc = wtc[start:stop, :] # TODO get scales to send to WTC return WTC( @@ -97,7 +109,8 @@ def wtc(self, pair: PairSignals, bin_seconds:float|None=None, period_cuts:List[f wavelet_library=self.wavelet_library, wavelet_name_with_args=self.wavelet_name_with_args, bin_seconds=bin_seconds, - period_cuts=period_cuts) + period_cuts=period_cuts, + ) except: - PycwtWavelet = None \ No newline at end of file + PycwtWavelet = None diff --git a/hypyp/wavelet/implementations/pywavelets_wavelet.py b/hypyp/wavelet/implementations/pywavelets_wavelet.py index 557cc6b..07b7b7b 100644 --- a/hypyp/wavelet/implementations/pywavelets_wavelet.py +++ b/hypyp/wavelet/implementations/pywavelets_wavelet.py @@ -10,6 +10,7 @@ DEFAULT_MORLET_CENTER_FREQUENCY = 1 DEFAULT_GAUSSIAN_DEGREE = 2 + class PywaveletsWavelet(BaseWavelet): """ Parent class for the default Wavelet implementation, using Pywavelets library. @@ -22,6 +23,7 @@ class PywaveletsWavelet(BaseWavelet): upper_bound (float, optional): upper bound for mother wavelet evaluation. Defaults to 8. cwt_params (dict | None, optional): params to be sent to cwt(), the Continuous Wavelet Transform. Defaults to None. """ + cwt_params: dict _wavelet_name: str lower_bound: float @@ -30,12 +32,12 @@ class PywaveletsWavelet(BaseWavelet): def __init__( self, - wavelet_name:str=f'cmor{DEFAULT_MORLET_BANDWIDTH_FREQUENCY},{DEFAULT_MORLET_CENTER_FREQUENCY}', - lower_bound:float=-8, - upper_bound:float=8, - cwt_params:dict|None=None, - period_range:Tuple[float, float]|None=None, - frequency_range:Tuple[float, float]|None=None, + wavelet_name: str = f"cmor{DEFAULT_MORLET_BANDWIDTH_FREQUENCY},{DEFAULT_MORLET_CENTER_FREQUENCY}", + lower_bound: float = -8, + upper_bound: float = 8, + cwt_params: dict | None = None, + period_range: Tuple[float, float] | None = None, + frequency_range: Tuple[float, float] | None = None, **kwargs, ): self._wavelet_name = wavelet_name @@ -47,17 +49,19 @@ def __init__( self.cwt_params = cwt_params # explicitly have period_range and frequency_range arguments, even if they belong to parent class - super().__init__(period_range=period_range, frequency_range=frequency_range, **kwargs) + super().__init__( + period_range=period_range, frequency_range=frequency_range, **kwargs + ) @property def wavelet_library(self): - return 'pywavelets' + return "pywavelets" @property def wavelet_name_with_args(self): name = self._wavelet_name if self.wtc_smoothing_win_size is not None: - name += f'[win:{self.wtc_smoothing_win_size}]' + name += f"[win:{self.wtc_smoothing_win_size}]" return name @property @@ -75,19 +79,29 @@ def evaluate_psi(self): def get_scales(self, dt): frequencies = 1 / self.get_periods() - scales = pywt.frequency2scale(self._wavelet, frequencies*dt) + scales = pywt.frequency2scale(self._wavelet, frequencies * dt) return scales - def cwt(self, y, dt, label='', cache_suffix:str='') -> CWT: + def cwt(self, y, dt, label="", cache_suffix: str = "") -> CWT: N = len(y) times = np.arange(N) * dt scales = self.get_scales(dt) - W, freqs = pywt.cwt(y, scales, self._wavelet, sampling_period=dt, method='fft', **self.cwt_params) + W, freqs = pywt.cwt( + y, + scales, + self._wavelet, + sampling_period=dt, + method="fft", + **self.cwt_params, + ) periods = 1 / freqs coi = self._get_and_cache_cone_of_influence(N, dt, cache_suffix=cache_suffix) - - return CWT(weights=W, times=times, scales=scales, periods=periods, coi=coi, label=label) + + return CWT( + weights=W, times=times, scales=scales, periods=periods, coi=coi, label=label + ) + class ComplexMorletWavelet(PywaveletsWavelet): """ @@ -99,6 +113,7 @@ class ComplexMorletWavelet(PywaveletsWavelet): bandwidth_frequency (float, optional): Defaults to 2. center_frequency (float, optional): Defaults to 1. """ + bandwidth_frequency: float center_frequency: float @@ -107,30 +122,30 @@ class ComplexMorletWavelet(PywaveletsWavelet): def __init__( self, - bandwidth_frequency:float=DEFAULT_MORLET_BANDWIDTH_FREQUENCY, - center_frequency:float=DEFAULT_MORLET_CENTER_FREQUENCY, - period_range:Tuple[float, float]|None=None, - frequency_range:Tuple[float, float]|None=None, + bandwidth_frequency: float = DEFAULT_MORLET_BANDWIDTH_FREQUENCY, + center_frequency: float = DEFAULT_MORLET_CENTER_FREQUENCY, + period_range: Tuple[float, float] | None = None, + frequency_range: Tuple[float, float] | None = None, **kwargs, ): self.bandwidth_frequency = bandwidth_frequency self.center_frequency = center_frequency # explicitly have period_range and frequency_range arguments, even if they belong to parent class return super().__init__( - wavelet_name=f'cmor{bandwidth_frequency},{center_frequency}', - period_range=period_range, - frequency_range=frequency_range, + wavelet_name=f"cmor{bandwidth_frequency},{center_frequency}", + period_range=period_range, + frequency_range=frequency_range, **kwargs, ) - + @property def flambda(self): # Equations come from "A Practical Guide to Wavelet Analysis" from Torrence and Compo (1998), Table 1 f0 = 2 * np.pi * self.center_frequency flambda = 2 * np.pi / f0 - #flambda = 4 * np.pi / (f0 + np.sqrt(2 + f0**2)) + # flambda = 4 * np.pi / (f0 + np.sqrt(2 + f0**2)) return flambda - + class ComplexGaussianWavelet(PywaveletsWavelet): """ @@ -141,21 +156,22 @@ class ComplexGaussianWavelet(PywaveletsWavelet): Args: degree (int, optional): the "degree" of "Degree of Gaussian" (DOG). Defaults to 2. """ + degree: int default_degree: int = DEFAULT_GAUSSIAN_DEGREE def __init__( self, - degree:int=DEFAULT_GAUSSIAN_DEGREE, - period_range:Tuple[float, float]|None=None, - frequency_range:Tuple[float, float]|None=None, - **kwargs + degree: int = DEFAULT_GAUSSIAN_DEGREE, + period_range: Tuple[float, float] | None = None, + frequency_range: Tuple[float, float] | None = None, + **kwargs, ): self.degree = degree # explicitly have period_range and frequency_range arguments, even if they belong to parent class return super().__init__( - wavelet_name=f'cgau{degree}', + wavelet_name=f"cgau{degree}", period_range=period_range, frequency_range=frequency_range, **kwargs, diff --git a/hypyp/wavelet/pair_signals.py b/hypyp/wavelet/pair_signals.py index 446caf6..1157ac0 100644 --- a/hypyp/wavelet/pair_signals.py +++ b/hypyp/wavelet/pair_signals.py @@ -2,6 +2,7 @@ import numpy as np + class PairSignals: x: np.ndarray y1: np.ndarray @@ -27,24 +28,25 @@ class PairSignals: label_roi1: str label_roi2: str - def __init__(self, - x, - y1, - y2, - label_ch1='', - label_ch2='', - label_roi1='', - label_roi2='', - label_s1='', - label_s2='', - label_dyad='', - label_task='', - epoch_idx=0, - section_idx=0, - is_intra:bool=False, - is_intra_of:int=0, - is_pseudo:bool=False, - ): + def __init__( + self, + x, + y1, + y2, + label_ch1="", + label_ch2="", + label_roi1="", + label_roi2="", + label_s1="", + label_s2="", + label_dyad="", + label_task="", + epoch_idx=0, + section_idx=0, + is_intra: bool = False, + is_intra_of: int = 0, + is_pseudo: bool = False, + ): """ A pair of signal that are already aligned and can be compared @@ -90,24 +92,24 @@ def __init__(self, self.label_s2 = label_s2 self.label_roi1 = label_roi1 self.label_roi2 = label_roi2 - + @property def label(self): - ret = f'{self.label_ch1} - {self.label_ch2}' + ret = f"{self.label_ch1} - {self.label_ch2}" prefix = self.label_task if self.epoch_idx > 0: - prefix = f'{prefix}[{self.epoch_idx}]' + prefix = f"{prefix}[{self.epoch_idx}]" if self.section_idx > 0: - prefix = f'{prefix}(section:{self.section_idx})' + prefix = f"{prefix}(section:{self.section_idx})" + + if prefix != "": + ret = f"{prefix} - {ret}" - if prefix != '': - ret = f'{prefix} - {ret}' - return ret - def sub(self, time_range:Tuple[float, float], section_idx:int|None=None): + def sub(self, time_range: Tuple[float, float], section_idx: int | None = None): """ Get a new PairSignals from a portion of the initial PairSignals @@ -118,7 +120,7 @@ def sub(self, time_range:Tuple[float, float], section_idx:int|None=None): Returns: PairSignals: a new PairSignals """ - if time_range[0] == 0 and time_range[1] == self.n/self.sfreq: + if time_range[0] == 0 and time_range[1] == self.n / self.sfreq: return self signal_from = int(self.sfreq * time_range[0]) @@ -126,7 +128,7 @@ def sub(self, time_range:Tuple[float, float], section_idx:int|None=None): if section_idx is None: section_idx = self.section_idx - + return PairSignals( self.x[signal_from:signal_to], self.y1[signal_from:signal_to], @@ -145,8 +147,6 @@ def sub(self, time_range:Tuple[float, float], section_idx:int|None=None): epoch_idx=self.epoch_idx, section_idx=section_idx, ) - - def __repr__(self): - return f'Pair({self.label})' - + def __repr__(self): + return f"Pair({self.label})" diff --git a/hypyp/wavelet/wtc.py b/hypyp/wavelet/wtc.py index 9e5e86e..0112443 100644 --- a/hypyp/wavelet/wtc.py +++ b/hypyp/wavelet/wtc.py @@ -13,6 +13,7 @@ MASK_THRESHOLD = 0.5 + class WTC: W: np.ndarray cwts: Tuple[CWT] @@ -50,8 +51,8 @@ class WTC: coherence_metric: float coherence_masked: float - bin_seconds: float|None - period_cuts: List[float]|None + bin_seconds: float | None + period_cuts: List[float] | None coherence_bins: List[Tuple[float, float, str, str]] wavelet_library: str @@ -59,18 +60,18 @@ class WTC: def __init__( self, - W:np.ndarray, - times:np.ndarray, - scales:np.ndarray, - periods:np.ndarray, - coi:np.ndarray, + W: np.ndarray, + times: np.ndarray, + scales: np.ndarray, + periods: np.ndarray, + coi: np.ndarray, pair: PairSignals, - bin_seconds:float|None=None, - period_cuts:List[float]|None=None, - frequency_cuts:List[float]|None=None, - wavelet_library:str='', - wavelet_name_with_args:str='', - cwts:Tuple[CWT, CWT]=None, + bin_seconds: float | None = None, + period_cuts: List[float] | None = None, + frequency_cuts: List[float] | None = None, + wavelet_library: str = "", + wavelet_name_with_args: str = "", + cwts: Tuple[CWT, CWT] = None, ): """ The WTC object holds the results of a Wavelet Transform Coherence @@ -117,15 +118,15 @@ def __init__( self.label_roi2 = pair.label_roi2 # These will not change when we downsample - dt = (times[1] - times[0]) + dt = times[1] - times[0] self.dt = dt self.sfreq = 1 / dt self.nyquist = self.sfreq / 2 self.bin_seconds = bin_seconds - + if period_cuts is not None and frequency_cuts is not None: - raise RuntimeError('Cannot specify both period_cuts and frequency_cuts') + raise RuntimeError("Cannot specify both period_cuts and frequency_cuts") if frequency_cuts is not None: self.period_cuts = list(1 / np.array(frequency_cuts)) @@ -133,7 +134,7 @@ def __init__( self.period_cuts = period_cuts else: self.period_cuts = None - + self.coherence_bins = [] self.wavelet_library = wavelet_library @@ -146,19 +147,19 @@ def frequency_cuts(self): if self.period_cuts is None: return None return 1 / np.array(self.period_cuts) - + @property def cwt1(self): if self.cwts is None: return None return self.cwts[0] - + @property def cwt2(self): if self.cwts is None: return None return self.cwts[1] - + @property def p_ranges(self): if self.period_cuts is None: @@ -180,7 +181,7 @@ def p_ranges(self): start_look_at = cursor break cursor += 1 - + # add one last for the remaining if len(p_ranges) == 0: p_ranges.append((0, len(self.periods))) @@ -188,18 +189,18 @@ def p_ranges(self): last_i = p_ranges[-1][1] if last_i < len(self.periods): p_ranges.append((last_i, len(self.periods))) - + return p_ranges - - + @property def p_ranges_str(self): ret = [] for p_start, p_end in self.p_ranges: - ret.append(f"{self.frequencies[p_start]:.2f}Hz-{self.frequencies[p_end-1]:.2f}Hz") + ret.append( + f"{self.frequencies[p_start]:.2f}Hz-{self.frequencies[p_end - 1]:.2f}Hz" + ) return ret - - + @property def t_ranges(self): # don't use self.dt because we want to deal with downsampled data as well @@ -236,34 +237,41 @@ def _compute_coherence_in_coi(self): # Time and period bins split self.coherence_bins = [] - + # loop over time bins and period bins for t_start, t_stop in self.t_ranges: for p_start, p_stop in self.p_ranges: wtc_bin = self.wtc_masked[p_start:p_stop, t_start:t_stop] coherence_bin = np.mean(wtc_bin) coherence_masked = np.mean(wtc_bin.mask) - if np.ma.is_masked(coherence_bin) or not np.isfinite(coherence_bin) or coherence_masked > MASK_THRESHOLD: + if ( + np.ma.is_masked(coherence_bin) + or not np.isfinite(coherence_bin) + or coherence_masked > MASK_THRESHOLD + ): coherence_bin = np.nan - self.coherence_bins.append(( - coherence_bin, - coherence_masked, - f'{np.round(t_start*dt):.0f}-{np.round(t_stop*dt):.0f}', - f'{self.periods[p_start]:.1f}-{self.periods[p_stop-1]:.1f}', - )) - - - def downsample_in_time(self, bins:int): + self.coherence_bins.append( + ( + coherence_bin, + coherence_masked, + f"{np.round(t_start * dt):.0f}-{np.round(t_stop * dt):.0f}", + f"{self.periods[p_start]:.1f}-{self.periods[p_stop - 1]:.1f}", + ) + ) + + def downsample_in_time(self, bins: int): """ Merge weights together over time to save memory and allow for faster displaying Args: bins (int): number of bins to keep """ - self.times, self.W, self.coi, self.coif, _factor = downsample_in_time(self.times, self.W, self.coi, self.coif, bins=bins) + self.times, self.W, self.coi, self.coif, _factor = downsample_in_time( + self.times, self.W, self.coi, self.coif, bins=bins + ) # must recompute coherence in cone of interest self._compute_coherence_in_coi() - + def _moving_average_1d(self, arr, window_size): kernel = np.ones(window_size) / window_size @@ -271,9 +279,9 @@ def _moving_average_1d(self, arr, window_size): mask = np.array(arr.mask) valid_indices = np.where(~arr.mask)[0] if len(valid_indices) > 0: - arr[:valid_indices[0]] = arr[valid_indices[0]] - arr[valid_indices[-1]:] = arr[valid_indices[-1]] - result = np.convolve(arr, kernel, mode='same') # keep same size array + arr[: valid_indices[0]] = arr[valid_indices[0]] + arr[valid_indices[-1] :] = arr[valid_indices[-1]] + result = np.convolve(arr, kernel, mode="same") # keep same size array return np.ma.masked_array(result, mask) def get_as_time_series(self, window_size=10) -> np.ndarray: @@ -281,8 +289,8 @@ def get_as_time_series(self, window_size=10) -> np.ndarray: data = np.ma.zeros((len(p_ranges), self.wtc_masked.shape[1])) for i in range(len(p_ranges)): p_start, p_stop = p_ranges[i] - data[i,:] = self.wtc_masked[p_start:p_stop].mean(axis=0) - #print(data.shape) + data[i, :] = self.wtc_masked[p_start:p_stop].mean(axis=0) + # print(data.shape) result = np.apply_along_axis(self._moving_average_1d, 1, data, window_size) return result @@ -291,30 +299,32 @@ def as_frame_rows(self) -> List[List]: # IMPORTANT: must match the ordering of COHERENCE_FRAME_COLUMNS frames = [] for bin_id in range(len(self.coherence_bins)): - frames.append([ - self.label_dyad, - self.is_intra, - self.is_intra_of, - self.is_pseudo, - self.label_s1, - self.label_s2, - self.label_roi1, - self.label_roi2, - self.label_ch1, - self.label_ch2, - self.task, - self.epoch_id, - self.section_id, - bin_id, - self.coherence_bins[bin_id][0], # metric - self.coherence_bins[bin_id][1], # masked - self.coherence_bins[bin_id][2], # time range - self.coherence_bins[bin_id][3], # period range - self.wavelet_library, - self.wavelet_name_with_args, - ]) + frames.append( + [ + self.label_dyad, + self.is_intra, + self.is_intra_of, + self.is_pseudo, + self.label_s1, + self.label_s2, + self.label_roi1, + self.label_roi2, + self.label_ch1, + self.label_ch2, + self.task, + self.epoch_id, + self.section_id, + bin_id, + self.coherence_bins[bin_id][0], # metric + self.coherence_bins[bin_id][1], # masked + self.coherence_bins[bin_id][2], # time range + self.coherence_bins[bin_id][3], # period range + self.wavelet_library, + self.wavelet_name_with_args, + ] + ) return frames - + def to_frame(self) -> CoherenceDataFrame: """ Get a typed pandas DataFrame from the WTC @@ -335,8 +345,10 @@ def plot(self, **kwargs) -> Figure: Returns: Figure: matplotlib.Figure """ - if 'title' not in kwargs: - kwargs['title'] = f'{self.label_s1}[{self.label_ch1}] - {self.label_s2}[{self.label_ch2}]' + if "title" not in kwargs: + kwargs["title"] = ( + f"{self.label_s1}[{self.label_ch1}] - {self.label_s2}[{self.label_ch2}]" + ) return plot_wavelet_transform_weights( self.W, @@ -346,5 +358,5 @@ def plot(self, **kwargs) -> Figure: self.sfreq, bin_seconds=self.bin_seconds, frequency_cuts=self.frequency_cuts, - **kwargs) - + **kwargs, + ) diff --git a/hypyp/xdf/__init__.py b/hypyp/xdf/__init__.py index 724985a..822dfc9 100644 --- a/hypyp/xdf/__init__.py +++ b/hypyp/xdf/__init__.py @@ -2,4 +2,3 @@ from .xdf_stream import XDFStream __all__ = ["XDFImport", "XDFStream"] - diff --git a/hypyp/xdf/xdf_import.py b/hypyp/xdf/xdf_import.py index ccd3ea3..f87bf35 100644 --- a/hypyp/xdf/xdf_import.py +++ b/hypyp/xdf/xdf_import.py @@ -6,7 +6,8 @@ from .xdf_stream import XDFStream, Markers -class XDFImport(): + +class XDFImport: """ Read an XDF file and (optionally) create mne Raws and Annotations @@ -19,14 +20,15 @@ class XDFImport(): convert_to_mne: Flag to disable the automatic conversion to mne Raws and Annotations verbose: Verbose flag """ - + file_path: str mne_type_map: dict | None selected_stream_indices: List[int] - map_id_to_idx: dict # Stream identifier to the index in our list + map_id_to_idx: dict # Stream identifier to the index in our list verbose: bool - def __init__(self, + def __init__( + self, file_path: str, select_type: str = None, select_matches: list = None, @@ -35,7 +37,7 @@ def __init__(self, verbose: bool = False, convert_to_mne: bool = True, ): - self.file_path = file_path + self.file_path = file_path self.scale = scale self.selected_stream_indices = [] @@ -44,9 +46,13 @@ def __init__(self, self.verbose = verbose # Load file - streams, self.header = pyxdf.load_xdf(file_path, verbose=(None if verbose == False else verbose)) + streams, self.header = pyxdf.load_xdf( + file_path, verbose=(None if verbose == False else verbose) + ) - self.available_streams = [XDFStream(stream, mne_type_map=mne_type_map) for stream in streams] + self.available_streams = [ + XDFStream(stream, mne_type_map=mne_type_map) for stream in streams + ] self.init_map_streams() @@ -60,30 +66,40 @@ def __init__(self, # Assert that we have found at least one real stream of selected type if len(self.selected_stream_indices) == 0: - raise ValueError(f"No stream selected. select_type: {select_type}, select_matches: {select_matches}") + raise ValueError( + f"No stream selected. select_type: {select_type}, select_matches: {select_matches}" + ) - if verbose: + if verbose: print(self) # Create MNE objects if convert_to_mne: self.convert_streams_to_mne() - + @property def selected_streams(self): """List of selected streams, given the select_match or select_type""" return [self.available_streams[idx] for idx in self.selected_stream_indices] - + @property def selected_signal_streams(self): """List of selected streams that are signal streams (srate>0), given the select_match or select_type""" - return [self.available_streams[idx] for idx in self.selected_stream_indices if self.available_streams[idx].is_mne_raw_compatible] - + return [ + self.available_streams[idx] + for idx in self.selected_stream_indices + if self.available_streams[idx].is_mne_raw_compatible + ] + @property def selected_markers_streams(self): """List of selected streams that are markers streams (srate==0), given the select_match or select_type""" - return [self.available_streams[idx] for idx in self.selected_stream_indices if self.available_streams[idx].is_mne_annotations_compatible] - + return [ + self.available_streams[idx] + for idx in self.selected_stream_indices + if self.available_streams[idx].is_mne_annotations_compatible + ] + @property def mne_raws_dict(self): """Dictionary of mne.io.RawArray object for selected streams, by stream name""" @@ -91,12 +107,12 @@ def mne_raws_dict(self): for stream in self.selected_signal_streams: ret[stream.name] = stream.mne_raw return ret - + @property def mne_raws(self): """List of mne.io.RawArray object for selected streams""" return [stream.mne_raw for stream in self.selected_signal_streams] - + @property def markers_dict(self) -> Dict[str, Markers]: """ @@ -108,7 +124,7 @@ def markers_dict(self) -> Dict[str, Markers]: for stream in self.selected_markers_streams: ret[stream.name] = stream.markers return ret - + @property def markers(self) -> Markers: """ @@ -123,15 +139,21 @@ def markers(self) -> Markers: else: markers += stream.markers return markers - + def get_stream_indices_for_type(self, stream_type: str): """For a given type of stream, get the list of index in the available streams which are of this type""" - return [idx for idx, stream in enumerate(self.available_streams) if stream.type == stream_type] - + return [ + idx + for idx, stream in enumerate(self.available_streams) + if stream.type == stream_type + ] + def get_stream_ids_for_type(self, stream_type: str): """For a given type of stream, get the XDF identifiers in the available streams which are of this type""" - return [stream.id for stream in self.available_streams if stream.type == stream_type] - + return [ + stream.id for stream in self.available_streams if stream.type == stream_type + ] + def init_map_streams(self): """ Create mapping between stream indentifiers and indices in our available_streams list @@ -143,20 +165,22 @@ def init_map_streams(self): for idx, stream in enumerate(self.available_streams): self.map_id_to_idx[stream.id] = idx - def select_all_streams(self) -> list: """ Find in the available streams loaded from the XDF file all the streams that can be converted to MNE - + Subsequent calls to class methods will only apply to the selected streams """ - self.selected_stream_indices = [idx for idx, stream in enumerate(self.available_streams) if stream.is_mne_compatible] - + self.selected_stream_indices = [ + idx + for idx, stream in enumerate(self.available_streams) + if stream.is_mne_compatible + ] def select_streams_by_type(self, stream_type: str) -> list: """ Find in the available streams loaded from the XDF file the streams that are of a specific type. - + Subsequent calls to class methods will only apply to the selected streams Parameters: @@ -165,12 +189,12 @@ def select_streams_by_type(self, stream_type: str) -> list: if self.verbose: print(f"Looking for streams of type '{stream_type}'") - + self.selected_stream_indices = self.get_stream_indices_for_type(stream_type) def select_streams_by_matches(self, keyword_matches: list): """ - Interpret the query made by the user (a list of indexes, or `str` that matches + Interpret the query made by the user (a list of indexes, or `str` that matches streams' name) into a list containing the indexes within the XDF file. Parameters: @@ -188,11 +212,11 @@ def select_streams_by_matches(self, keyword_matches: list): for idx, stream in enumerate(self.available_streams): if stream.name == keyword_match: found_idx = idx - + if found_idx is None: raise ValueError(f"No stream matching keyword '{keyword_match}'") - self.selected_stream_indices.append(found_idx) + self.selected_stream_indices.append(found_idx) def convert_streams_to_mne(self): """ @@ -200,29 +224,33 @@ def convert_streams_to_mne(self): mne.io.RawArray objects are then available using obj.mne_raws or obj.mne_raws_dict """ - + # Find if all the stream have unique names (true if any stream name is duplicated) names = [stream.name for stream in self.selected_streams] for stream in self.selected_streams: if self.verbose: - print(f'Converting {stream.name} to MNE') + print(f"Converting {stream.name} to MNE") has_duplicate_names = names.count(stream.name) > 1 if has_duplicate_names: - warnings.warn("Multiple streams have the same name. Adding original stream_id as suffixes to the generated raws") - + warnings.warn( + "Multiple streams have the same name. Adding original stream_id as suffixes to the generated raws" + ) + stream.convert_to_mne_raw(self.scale, append_stream_id=has_duplicate_names) - + # add annotations from markers streams to data streams for data_stream in self.selected_signal_streams: if data_stream.markers is not None: - mne_annotations = data_stream.markers.as_mne_annotations(data_stream.reference_time) + mne_annotations = data_stream.markers.as_mne_annotations( + data_stream.reference_time + ) data_stream.mne_raw.set_annotations(mne_annotations) if self.verbose: print("All convertion to MNE done.") - + def rename_channels(self, new_names): """ Set the name of all the channels for every signal streams. Useful when they were not correctly loaded (or not present) from the XDF file @@ -232,7 +260,7 @@ def rename_channels(self, new_names): """ for stream in self.selected_signal_streams: stream.rename_channels(new_names) - + def set_montage(self, montage): """ Set the montage of the raw(s) using a custom mne montage label, or the path to a dig.montage file. @@ -245,9 +273,8 @@ def set_montage(self, montage): names = [stream.name for stream in self.selected_signal_streams] print(f"Setting '{montage}' as the montage for streams: {','.join(names)}") - for stream in self.selected_signal_streams: - try: + try: stream.set_montage(montage) except ValueError as e: warnings.warn(f"Invalid montage given to mne.set_montage(): {montage}") @@ -263,14 +290,19 @@ def save_to_fif_files(self, dir_path): Returns: List[str]: The list of file names that have been created """ - return [stream.save_to_fif_file(dir_path) for stream in self.selected_signal_streams] - + return [ + stream.save_to_fif_file(dir_path) for stream in self.selected_signal_streams + ] + def __str__(self): - available_streams_str = "\n ".join([str(stream) for stream in self.available_streams]) - selected_streams_str = ",".join([stream.name for stream in self.selected_streams]) + available_streams_str = "\n ".join( + [str(stream) for stream in self.available_streams] + ) + selected_streams_str = ",".join( + [stream.name for stream in self.selected_streams] + ) return f"""XDFImport with {len(self.available_streams)} available streams and {len(self.selected_streams)} selected streams (loaded from {self.file_path}) Available streams: {available_streams_str} Selected streams: [{selected_streams_str}] """ - \ No newline at end of file diff --git a/hypyp/xdf/xdf_stream.py b/hypyp/xdf/xdf_stream.py index 117adcf..d5ce5fa 100644 --- a/hypyp/xdf/xdf_stream.py +++ b/hypyp/xdf/xdf_stream.py @@ -7,8 +7,9 @@ import numpy as np from pyxdf import pyxdf + @dataclass -class Markers(): +class Markers: """ Data class to keep the markers data from non-signal XDF streams @@ -16,6 +17,7 @@ class Markers(): timestamps (1D np.array): The time of markers descriptions (List[str]): The string value for every marker. Must have the same length as timestamps """ + timestamps: np.array descriptions: List[str] @@ -43,8 +45,12 @@ def as_mne_annotations(self, reference_time: float = 0.0): return empty durations = np.zeros_like(onset_times) - return mne.Annotations(onset=onset_times[i:], duration=durations[i:], description=self.descriptions[i:]) - + return mne.Annotations( + onset=onset_times[i:], + duration=durations[i:], + description=self.descriptions[i:], + ) + def __add__(self, other): if not isinstance(other, Markers): return NotImplemented @@ -52,7 +58,8 @@ def __add__(self, other): descriptions = self.descriptions + other.descriptions return Markers(timestamps, descriptions) -class XDFStream(): + +class XDFStream: """ Class representation of a single stream from a XDF import @@ -61,13 +68,16 @@ class XDFStream(): metadata_desc (dict): Channel info from xdf loaded stream mne_raw (mne.io.RawArray) """ + pyxdf_stream: pyxdf.StreamData metadata_desc: dict mne_raw: mne.io.RawArray | None markers: mne.Annotations | None @staticmethod - def get_default_mne_ch_type_for_stream_type(stream_type: str, type_map: dict = None): + def get_default_mne_ch_type_for_stream_type( + stream_type: str, type_map: dict = None + ): """ Get a mne channel type, given a stream type. @@ -82,10 +92,37 @@ def get_default_mne_ch_type_for_stream_type(stream_type: str, type_map: dict = N return type_map[stream_type] mne_types = [ - 'grad', 'mag', 'ref_meg', 'eeg', 'seeg', 'dbs', 'ecog', 'eog', 'emg', 'ecg', - 'resp', 'bio', 'misc', 'stim', 'exci', 'syst', 'ias', 'gof', 'dipole', 'chpi', - 'fnirs_cw_amplitude', 'fnirs_fd_ac_amplitude', 'fnirs_fd_phase', 'fnirs_od', 'hbo', 'hbr', - 'csd', 'temperature', 'gsr', 'eyegaze', 'pupil', + "grad", + "mag", + "ref_meg", + "eeg", + "seeg", + "dbs", + "ecog", + "eog", + "emg", + "ecg", + "resp", + "bio", + "misc", + "stim", + "exci", + "syst", + "ias", + "gof", + "dipole", + "chpi", + "fnirs_cw_amplitude", + "fnirs_fd_ac_amplitude", + "fnirs_fd_phase", + "fnirs_od", + "hbo", + "hbr", + "csd", + "temperature", + "gsr", + "eyegaze", + "pupil", ] type_lower = stream_type.lower() @@ -93,11 +130,11 @@ def get_default_mne_ch_type_for_stream_type(stream_type: str, type_map: dict = N return type_lower # A few extra cases for convenience - if type_lower == 'fnirs' or type_lower == 'nirs': - return 'fnirs_cw_amplitude' + if type_lower == "fnirs" or type_lower == "nirs": + return "fnirs_cw_amplitude" - if type_lower == 'markers' or type_lower == 'stim': - return 'stim' + if type_lower == "markers" or type_lower == "stim": + return "stim" # return the stream type received, maybe there are new mne types not listed above return type_lower @@ -118,19 +155,21 @@ def get_mne_ch_types(stream_type: str, ch_names: list, type_map: dict = None): List[str]: List of type for every channel """ ch_types = [] - default_ch_type = XDFStream.get_default_mne_ch_type_for_stream_type(stream_type, type_map=type_map) + default_ch_type = XDFStream.get_default_mne_ch_type_for_stream_type( + stream_type, type_map=type_map + ) for ch_name in ch_names: ch_type = default_ch_type - if stream_type.startswith('Acc'): - ch_type = 'misc' - if ch_name.startswith('Markers'): - ch_type = 'stim' - if ch_name.startswith('Acc'): - ch_type = 'misc' - if ch_name.startswith('Gyro'): - ch_type = 'misc' - if ch_name.startswith('Quat'): - ch_type = 'misc' + if stream_type.startswith("Acc"): + ch_type = "misc" + if ch_name.startswith("Markers"): + ch_type = "stim" + if ch_name.startswith("Acc"): + ch_type = "misc" + if ch_name.startswith("Gyro"): + ch_type = "misc" + if ch_name.startswith("Quat"): + ch_type = "misc" ch_types.append(ch_type) return ch_types @@ -144,10 +183,10 @@ def __init__(self, stream, mne_type_map: dict = None): if self.pyxdf_stream["info"]["desc"] and self.pyxdf_stream["info"]["desc"][0]: self.metadata_desc = self.pyxdf_stream["info"]["desc"][0] - + if self.is_markers_compatible: self.init_markers() - + @property def name(self): return self.pyxdf_stream["info"]["name"][0] @@ -155,7 +194,7 @@ def name(self): @property def type(self): return self.pyxdf_stream["info"]["type"][0] - + @property def id(self): return self.pyxdf_stream["info"]["stream_id"] @@ -163,22 +202,22 @@ def id(self): @property def srate(self): return float(self.pyxdf_stream["info"]["nominal_srate"][0]) - + @property def reference_time(self): - return self.pyxdf_stream['time_stamps'][0] + return self.pyxdf_stream["time_stamps"][0] @property def is_markers_compatible(self): return self.srate == 0.0 - + @property def is_mne_compatible(self): return self.is_mne_raw_compatible or self.is_mne_annotations_compatible @property def is_mne_raw_compatible(self): - if self.type.lower() == 'quality': + if self.type.lower() == "quality": return False return self.srate > 0 @@ -188,17 +227,19 @@ def is_mne_annotations_compatible(self): @property def time_series(self): - if self.pyxdf_stream["time_series"].shape[0] > self.pyxdf_stream["time_series"].shape[1]: + if ( + self.pyxdf_stream["time_series"].shape[0] + > self.pyxdf_stream["time_series"].shape[1] + ): return self.pyxdf_stream["time_series"].T return self.pyxdf_stream["time_series"] - @property def ch_names(self): # The specification for metadata is to have xml nodes as desc->channels->channel->label... # Some manufacturers have desc->channel->name... - # See https://github.com/sccn/xdf/wiki/EEG-Meta-Data + # See https://github.com/sccn/xdf/wiki/EEG-Meta-Data # Let's handle both lookup_fields = ["label", "name"] channels_children = None @@ -206,20 +247,20 @@ def ch_names(self): channels_children = self.metadata_desc["channels"][0]["channel"] elif "channel" in self.metadata_desc: channels_children = self.metadata_desc["channel"] - + if channels_children is not None: for field in lookup_fields: if field in channels_children[0]: return [channel[field][0] for channel in channels_children] - + # fallback to generic naming - n_channels = int(self.pyxdf_stream['info']['channel_count'][0]) - return [f'{self.type}_{n:03}' for n in range(1, 1+n_channels)] - + n_channels = int(self.pyxdf_stream["info"]["channel_count"][0]) + return [f"{self.type}_{n:03}" for n in range(1, 1 + n_channels)] + @property def ch_types(self): return XDFStream.get_mne_ch_types(self.type, self.ch_names, self.mne_type_map) - + def init_markers(self): """ Save markers from a stream of type "Markers" (which contains strings and have a sampling rate of 0) @@ -228,9 +269,9 @@ def init_markers(self): """ assert self.is_markers_compatible - timestamps = np.array(self.pyxdf_stream['time_stamps']) + timestamps = np.array(self.pyxdf_stream["time_stamps"]) - descriptions = self.pyxdf_stream['time_series'] + descriptions = self.pyxdf_stream["time_series"] descriptions = [str(desc[0]) for desc in descriptions] # Flatten description self.markers = Markers(timestamps, descriptions) @@ -238,18 +279,17 @@ def init_markers(self): def get_unique_name(self, append_stream_id: bool = False): unique_stream_name = self.name if append_stream_id: - unique_stream_name += f'-{self.id}' + unique_stream_name += f"-{self.id}" return unique_stream_name - - + def convert_to_mne_raw( self, scale: float | str | None, append_stream_id: bool = False, - verbose: bool = False + verbose: bool = False, ): """ - Create and store in self a mne.io.RawArray object from + Create and store in self a mne.io.RawArray object from Parameters: scale (float | str | None): Scaling factor or 'auto' for automatic scaling, None for no scaling. @@ -260,7 +300,7 @@ def convert_to_mne_raw( if self.is_mne_raw_compatible: mne_info = self.create_mne_info(verbose=verbose) scaled_data = self.scale_data(self.time_series, scale, verbose=verbose) - self.mne_raw = mne.io.RawArray(scaled_data, mne_info) + self.mne_raw = mne.io.RawArray(scaled_data, mne_info) def create_mne_info(self, verbose: bool = False): """Create a mne.info object from the XDF's stream metadata.""" @@ -268,18 +308,29 @@ def create_mne_info(self, verbose: bool = False): ch_types = self.ch_types if self.srate <= 0: - raise RuntimeError(f"invalid sampling frequency {self.srate} for stream '{self.id}' with name '{self.name}' and type '{self.type}'") + raise RuntimeError( + f"invalid sampling frequency {self.srate} for stream '{self.id}' with name '{self.name}' and type '{self.type}'" + ) + + mne_info = mne.create_info( + ch_names, ch_types=ch_types, sfreq=self.srate, verbose=verbose + ) - mne_info = mne.create_info(ch_names, ch_types=ch_types, sfreq=self.srate, verbose=verbose) - mne_info["subject_info"] = { - "id": self.id, # integer identifier of the subject - "his_id": self.unique_name, # string identifier of the subject + "id": self.id, # integer identifier of the subject + "his_id": self.unique_name, # string identifier of the subject } return mne_info - def scale_data(self, data, scale: float | str | None, std_threshold=1e-5, amplitude_threshold=1.0, verbose=False): + def scale_data( + self, + data, + scale: float | str | None, + std_threshold=1e-5, + amplitude_threshold=1.0, + verbose=False, + ): """ Automatically scale signal data to V if necessary based on both standard deviation and amplitude range. @@ -293,7 +344,7 @@ def scale_data(self, data, scale: float | str | None, std_threshold=1e-5, amplit Returns: Scaled signal data array. """ - + # Check for NaN and Inf values in the data data = data.astype(np.float64) @@ -309,21 +360,25 @@ def scale_data(self, data, scale: float | str | None, std_threshold=1e-5, amplit print(f"- Returning the data scaled with {scale}") return data * scale - if scale=='auto': + if scale == "auto": std = data.std() amplitude_range = data.max() - data.min() if std > std_threshold and amplitude_range > amplitude_threshold: # Data has both a small standard deviation and small amplitude range, indicating it's in mV - print(f"- Std ({std}) and amplitude range ({amplitude_range}) are larges >> converting to V (e.g., * 10e-6)") + print( + f"- Std ({std}) and amplitude range ({amplitude_range}) are larges >> converting to V (e.g., * 10e-6)" + ) return data * 10e-6 # Scale to V # Data has a large standard deviation or amplitude range, indicating it's already in V - print(f"- Std ({std}) or amplitude range ({amplitude_range}) are smalls >> NOT converting to V") + print( + f"- Std ({std}) or amplitude range ({amplitude_range}) are smalls >> NOT converting to V" + ) return data - + raise ValueError(f"Invalid value for scale: {scale}") - + def set_montage(self, montage): """ Add the mne montage @@ -332,7 +387,7 @@ def set_montage(self, montage): montage (str): See MNE documentation """ self.mne_raw.set_montage(montage) - + def rename_channels(self, new_names): """ Set the name of all the channels for signal stream. Useful when they were not correctly loaded (or not present) from the XDF file @@ -346,9 +401,9 @@ def rename_channels(self, new_names): for i, new_name in enumerate(new_names): mapping[old_names[i]] = new_name - + self.mne_raw.rename_channels(mapping) - + def save_to_fif_file(self, dir_path: str): """ Save the mne.io.RawArray object as .fif file, using the unique name for file naming @@ -363,11 +418,11 @@ def save_to_fif_file(self, dir_path: str): save_file_path = os.path.join(dir_path, f"{self.unique_name}_raw.fif") self.mne_raw.save(save_file_path, overwrite=True) return save_file_path - + def __str__(self): - ch_names = ','.join(self.ch_names) - ch_types = ','.join(self.ch_types) + ch_names = ",".join(self.ch_names) + ch_types = ",".join(self.ch_types) return f"""Stream id {self.id} of type '{self.type}' with name '{self.name}' Sampling Rate: {self.srate} Channel names: [{ch_names}] - Channel types: [{ch_types}]""" \ No newline at end of file + Channel types: [{ch_types}]""" diff --git a/pyproject.toml b/pyproject.toml index 35dcd09..c8fefe4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ Documentation = "https://hypyp.readthedocs.io" [dependency-groups] dev = [ {include-group = "docs"}, - "black>=25.1.0", + {include-group = "lint"}, "ipykernel>=6.19.4", "ipywidgets>=7.5.1", "jupyterlab>=4.3.5", @@ -73,6 +73,10 @@ docs = [ "mkdocstrings>=0.20.0", "mkdocstrings-python>=1.1.0", ] +# Only the formatter and linter, so that the CI lint job installs nothing else. +lint = [ + "ruff>=0.16.10", +] benchmark = [ "doit>=0.36.0", ] @@ -86,3 +90,15 @@ include = ["hypyp*"] [tool.setuptools.package-data] hypyp = ["data/*.obj"] + +[tool.ruff] +line-length = 88 +target-version = "py311" +# hypyp/ext is vendored third-party code and is kept as upstream wrote it. +# The tutorial notebooks are not formatted yet. +extend-exclude = ["hypyp/ext", "*.ipynb"] + +[tool.ruff.lint] +# Only what is certainly a bug: syntax errors, invalid comparisons and +# statements, undefined names. Style rules are deliberately left out for now. +select = ["E9", "F63", "F7", "F82"] diff --git a/tests/accorr_reference.py b/tests/accorr_reference.py index f978b21..a782791 100644 --- a/tests/accorr_reference.py +++ b/tests/accorr_reference.py @@ -74,7 +74,9 @@ def accorr_reference( phase_diff = alpha1 - alpha2 phase_sum = alpha1 + alpha2 - mean_diff = np.angle(np.mean(np.exp(1j * phase_diff), axis=2, keepdims=True)) + mean_diff = np.angle( + np.mean(np.exp(1j * phase_diff), axis=2, keepdims=True) + ) mean_sum = np.angle(np.mean(np.exp(1j * phase_sum), axis=2, keepdims=True)) n_adj = -1 * (mean_diff - mean_sum) / 2 @@ -83,9 +85,7 @@ def accorr_reference( x_sin = np.sin(alpha1 - m_adj) y_sin = np.sin(alpha2 - n_adj) - den_ij = 2 * np.sqrt( - np.sum(x_sin**2, axis=2) * np.sum(y_sin**2, axis=2) - ) + den_ij = 2 * np.sqrt(np.sum(x_sin**2, axis=2) * np.sum(y_sin**2, axis=2)) den[:, :, i, j] = den_ij den[:, :, j, i] = den_ij diff --git a/tests/conftest.py b/tests/conftest.py index c2ba472..2d31e61 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,7 @@ def epochs(): mne.epochs.equalize_epoch_counts([epo1, epo2]) epoch_merge = utils.merge(epo1, epo2) - epochsTuple = namedtuple('epochs', ['epo1', 'epo2', 'epoch_merge']) + epochsTuple = namedtuple("epochs", ["epo1", "epo2", "epoch_merge"]) return epochsTuple(epo1=epo1, epo2=epo2, epoch_merge=epoch_merge) @@ -35,8 +35,9 @@ def preprocessed_epochs(): mne.epochs.equalize_epoch_counts([epo1, epo2]) epoch_merge = utils.merge(epo1, epo2) - preprocessedTuple = namedtuple('preprocessed_epochs', - ['pepo1', 'pepo2', 'pepochs_merge']) + preprocessedTuple = namedtuple( + "preprocessed_epochs", ["pepo1", "pepo2", "pepochs_merge"] + ) return preprocessedTuple(pepo1=epo1, pepo2=epo2, pepochs_merge=epoch_merge) @@ -48,14 +49,12 @@ def complex_signal_raw(preprocessed_epochs): Returns complex_signal_raw with shape (2, n_epochs, n_channels, n_freq, n_times), ready for passing directly to compute_sync(). """ - freq_bands = OrderedDict({ - 'Alpha-Low': [7.5, 11], - 'Alpha-High': [11.5, 13] - }) + freq_bands = OrderedDict({"Alpha-Low": [7.5, 11], "Alpha-High": [11.5, 13]}) - data_inter = np.array([preprocessed_epochs.pepo1.get_data(), - preprocessed_epochs.pepo2.get_data()]) - sampling_rate = preprocessed_epochs.pepo1.info['sfreq'] + data_inter = np.array( + [preprocessed_epochs.pepo1.get_data(), preprocessed_epochs.pepo2.get_data()] + ) + sampling_rate = preprocessed_epochs.pepo1.info["sfreq"] return analyses.compute_freq_bands( data_inter, @@ -63,7 +62,7 @@ def complex_signal_raw(preprocessed_epochs): freq_bands, filter_length=int(sampling_rate), l_trans_bandwidth=5.0, - h_trans_bandwidth=5.0 + h_trans_bandwidth=5.0, ) @@ -75,10 +74,13 @@ def complex_signal(complex_signal_raw): Returns complex_signal with shape (n_epochs, n_freq, 2*n_channels, n_times), ready for metric computation via ACCorr.compute() or accorr_reference(). """ - n_epoch, n_ch, n_freq, n_samp = (complex_signal_raw.shape[1], - complex_signal_raw.shape[2], - complex_signal_raw.shape[3], - complex_signal_raw.shape[4]) + n_epoch, n_ch, n_freq, n_samp = ( + complex_signal_raw.shape[1], + complex_signal_raw.shape[2], + complex_signal_raw.shape[3], + complex_signal_raw.shape[4], + ) - return complex_signal_raw.transpose( - (1, 3, 0, 2, 4)).reshape(n_epoch, n_freq, 2 * n_ch, n_samp) + return complex_signal_raw.transpose((1, 3, 0, 2, 4)).reshape( + n_epoch, n_freq, 2 * n_ch, n_samp + ) diff --git a/tests/test_coherence_data_frame.py b/tests/test_coherence_data_frame.py index f36f361..0c5f615 100644 --- a/tests/test_coherence_data_frame.py +++ b/tests/test_coherence_data_frame.py @@ -11,39 +11,44 @@ from hypyp.wavelet.implementations.pywavelets_wavelet import ComplexMorletWavelet from hypyp.wavelet.wtc import WTC + def test_instanciate(): wavelet = ComplexMorletWavelet(disable_caching=True) signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() - res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y, label_ch1='foo', label_ch2='bar')) + res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y, label_ch1="foo", label_ch2="bar") + ) df = CoherenceDataFrame.from_wtc_frame_rows(res.as_frame_rows) - assert np.all(df['channel1'] == 'foo') - assert np.all(df['channel2'] == 'bar') + assert np.all(df["channel1"] == "foo") + assert np.all(df["channel2"] == "bar") # Make sure categorical is correct when merging data frames - res2 = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y, label_ch1='baz', label_ch2='bar')) + res2 = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y, label_ch1="baz", label_ch2="bar") + ) df2 = CoherenceDataFrame.from_wtc_frame_rows(res2.as_frame_rows) df_concat = CoherenceDataFrame.concat([df, df2]) - assert len(df_concat['channel1'].unique()) == 2 - assert len(df_concat['channel2'].unique()) == 1 + assert len(df_concat["channel1"].unique()) == 2 + assert len(df_concat["channel2"].unique()) == 1 # All the 'bar' should be set to the same category - #assert np.all(df_concat['channel2'].cat.codes == 0) + # assert np.all(df_concat['channel2'].cat.codes == 0) def test_save_load_feather(): wavelet = ComplexMorletWavelet(disable_caching=True) signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() - res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y, label_ch1='foo', label_ch2='bar')) + res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y, label_ch1="foo", label_ch2="bar") + ) df_before = CoherenceDataFrame.from_wtc_frame_rows(res.as_frame_rows) - with tempfile.TemporaryDirectory(prefix='hypyp-') as tmp_path: - feather_path = os.path.join(tmp_path, 'test.feather') + with tempfile.TemporaryDirectory(prefix="hypyp-") as tmp_path: + feather_path = os.path.join(tmp_path, "test.feather") CoherenceDataFrame.save_feather(df_before, feather_path) df_after = CoherenceDataFrame.from_feather(feather_path) - assert np.all(df_after['channel1'] == 'foo') - - + assert np.all(df_after["channel1"] == "foo") diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 140e2ca..13897af 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -17,6 +17,7 @@ # Cache / registry helpers # --------------------------------------------------------------------------- + def test_cache_dir_is_path(): assert isinstance(datasets.cache_dir(), Path) @@ -40,6 +41,7 @@ def test_fetch_unknown_key_raises(): # EEG helpers # --------------------------------------------------------------------------- + def test_eeg_epochs_participant1(): path = datasets.eeg_epochs(participant=1) assert Path(path).exists() @@ -62,6 +64,7 @@ def test_eeg_raw(): # EEGLAB helpers — .fdt companion must be co-located # --------------------------------------------------------------------------- + def test_eeglab_epochs_set_exists(): path = datasets.eeglab_epochs() assert Path(path).exists() @@ -90,6 +93,7 @@ def test_eeglab_raw_fdt_colocated(): # fNIRS helpers # --------------------------------------------------------------------------- + def test_fnirs_samples_default_returns_four(): paths = datasets.fnirs_samples() assert len(paths) == 4 @@ -145,6 +149,7 @@ def test_fnirs_fcs01_parent_companions_colocated(): # XDF helpers # --------------------------------------------------------------------------- + def test_xdf_dyad_noise(): path = datasets.xdf_dyad_noise() assert Path(path).exists() diff --git a/tests/test_fnirs.py b/tests/test_fnirs.py index a3c02a8..559dcbe 100644 --- a/tests/test_fnirs.py +++ b/tests/test_fnirs.py @@ -17,10 +17,17 @@ from hypyp.fnirs.channel_roi import ChannelROI from hypyp.fnirs.dyad import Dyad from hypyp.fnirs.data_browser import DataBrowser -from hypyp.fnirs.preprocessor.base_step import PREPROCESS_STEP_BASE_KEY, PREPROCESS_STEP_HAEMO_FILTERED_KEY +from hypyp.fnirs.preprocessor.base_step import ( + PREPROCESS_STEP_BASE_KEY, + PREPROCESS_STEP_HAEMO_FILTERED_KEY, +) from hypyp.fnirs.preprocessor.implementations.mne_step import MneStep -from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_raw_to_haemo import MnePreprocessorRawToHaemo -from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_as_is import MnePreprocessorAsIs +from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_raw_to_haemo import ( + MnePreprocessorRawToHaemo, +) +from hypyp.fnirs.preprocessor.implementations.mne_preprocessor_as_is import ( + MnePreprocessorAsIs, +) from hypyp.utils import TASK_NEXT_EVENT, Task from hypyp import datasets @@ -32,21 +39,29 @@ # avoid all the output from mne logging.disable() + # Test helpers def get_test_recording(): - tasks = [Task('task1', onset_time=0, duration=20)] + tasks = [Task("task1", onset_time=0, duration=20)] return Recording(tasks=tasks).load_file(snirf_file1) -def get_test_recordings(count:int=2): - tasks = [Task('task1', onset_time=0, duration=20)] - recordings = [Recording(tasks=tasks, subject_label=f's{i+1}').load_file(snirf_file1) for i in range(count)] + +def get_test_recordings(count: int = 2): + tasks = [Task("task1", onset_time=0, duration=20)] + recordings = [ + Recording(tasks=tasks, subject_label=f"s{i + 1}").load_file(snirf_file1) + for i in range(count) + ] return recordings + def get_test_ch_match_one(): - return 'S1_D1 760' + return "S1_D1 760" + def get_test_ch_match_few(): - return re.compile(r'^S1_.*760') + return re.compile(r"^S1_.*760") + # # Data Browser @@ -56,39 +71,44 @@ def test_list_paths(): paths = browser.paths assert len(paths) == 0 + def test_add_source(): browser = DataBrowser() previous_count = len(browser.paths) - browser.add_source('/foo') + browser.add_source("/foo") new_paths = browser.paths assert len(new_paths) == previous_count + 1 + def test_list_files(): import os + browser = DataBrowser() browser.add_source(os.path.dirname(snirf_file1)) assert len(browser.paths) > 0 assert len(browser.list_all_files()) > 0 # path should be absolute - assert browser.list_all_files()[0].startswith('/') + assert browser.list_all_files()[0].startswith("/") # # Preprocessing pipeline # + # Try to load every file types we have @pytest.mark.parametrize("file_path", fnirs_files) def test_data_loader_all_types(file_path): warnings.filterwarnings("ignore", category=RuntimeWarning) raw = MnePreprocessorRawToHaemo().read_file(file_path) - assert raw.info['sfreq'] > 0 + assert raw.info["sfreq"] > 0 assert len(raw.ch_names) > 0 - + + def test_preprocess_step(): - key = 'foo_key' - desc = 'foo_description' - raw = mne.io.RawArray(np.array([[1., 2.]]), mne.create_info(['foo'], 1)) + key = "foo_key" + desc = "foo_description" + raw = mne.io.RawArray(np.array([[1.0, 2.0]]), mne.create_info(["foo"], 1)) step = MneStep(raw, key, desc) assert step.obj.get_data().shape[0] == 1 assert step.obj.get_data().shape[1] == 2 @@ -96,88 +116,98 @@ def test_preprocess_step(): assert step.desc == desc assert step.duration == 2 + # # Subject Recordings # + def test_load_from_raw(): raw = mne.io.read_raw_snirf(snirf_file1, preload=True) recording = Recording().load_raw(raw) - assert recording.subject_label == 'default' + assert recording.subject_label == "default" assert len(recording.mne_raw.times) == len(raw.times) def test_recording(): filepath = snirf_file1 - recording = Recording(subject_label='my_subject') + recording = Recording(subject_label="my_subject") recording.load_file(filepath, MnePreprocessorRawToHaemo(), preprocess=False) - assert recording.subject_label == 'my_subject' + assert recording.subject_label == "my_subject" assert recording.filepath == filepath assert recording.mne_raw is not None - assert len(recording.tasks) == 1 # default task, which is the complete record + assert len(recording.tasks) == 1 # default task, which is the complete record assert recording.preprocess_steps is None assert recording.is_preprocessed == False - assert recording.epochs_per_task is None # need preprocessing to extract epochs + assert recording.epochs_per_task is None # need preprocessing to extract epochs + def test_recording_load_participant_name(): filepath = snirf_file1 recording = Recording() recording.load_file(filepath, preprocess=False) - assert recording.subject_label == 'default' + assert recording.subject_label == "default" + def test_recording_tasks(): - recording = Recording(tasks=[Task('my_task_in_time', onset_time=1, duration=2)]) + recording = Recording(tasks=[Task("my_task_in_time", onset_time=1, duration=2)]) assert len(recording.tasks) == 1 - assert recording.task_keys[0] == 'my_task_in_time' + assert recording.task_keys[0] == "my_task_in_time" + def test_recording_epochs(): tasks = [ - Task('task1', onset_event_id=1, offset_event_id=TASK_NEXT_EVENT), - Task('task2', onset_event_id=3, offset_event_id=TASK_NEXT_EVENT), + Task("task1", onset_event_id=1, offset_event_id=TASK_NEXT_EVENT), + Task("task2", onset_event_id=3, offset_event_id=TASK_NEXT_EVENT), ] recording = Recording(tasks=tasks) recording.load_file(snirf_file1, MnePreprocessorRawToHaemo()) - assert len(recording.get_epochs_for_task('task1')) == 2 - n_events = recording.get_epochs_for_task('task1').events.shape[0] + assert len(recording.get_epochs_for_task("task1")) == 2 + n_events = recording.get_epochs_for_task("task1").events.shape[0] assert n_events == 2 - assert len(recording.get_epochs_for_task('task2')) == 3 + assert len(recording.get_epochs_for_task("task2")) == 3 + def test_recording_time_range_task(): tasks = [ - Task('task1', onset_time=1, duration=1), - Task('task2', onset_time=4, duration=1), + Task("task1", onset_time=1, duration=1), + Task("task2", onset_time=4, duration=1), ] recording = Recording(tasks=tasks) recording.load_file(snirf_file1, MnePreprocessorRawToHaemo()) - epochs_task1 = recording.get_epochs_for_task('task1') - epochs_task2 = recording.get_epochs_for_task('task2') + epochs_task1 = recording.get_epochs_for_task("task1") + epochs_task2 = recording.get_epochs_for_task("task2") assert len(epochs_task1) == 1 n_events = epochs_task1.events.shape[0] assert n_events == 1 assert len(epochs_task2) == 1 - + + def test_recording_time_range_task_recurring_event(): tasks = [ - Task('task1', onset_time=1, duration=1), - Task('task1', onset_time=4, duration=1), - Task('task1', onset_time=8, duration=1), + Task("task1", onset_time=1, duration=1), + Task("task1", onset_time=4, duration=1), + Task("task1", onset_time=8, duration=1), ] recording = Recording(tasks=tasks) recording.load_file(snirf_file1, MnePreprocessorRawToHaemo()) - epochs = recording.get_epochs_for_task('task1') + epochs = recording.get_epochs_for_task("task1") assert len(epochs) == 3 n_events = epochs.events.shape[0] assert n_events == 3 - + def test_upstream_preprocessor(): - recording = Recording(tasks=[Task('task1', onset_event_id=1, offset_event_id=TASK_NEXT_EVENT)]).load_file(snirf_file1, MnePreprocessorAsIs()) + recording = Recording( + tasks=[Task("task1", onset_event_id=1, offset_event_id=TASK_NEXT_EVENT)] + ).load_file(snirf_file1, MnePreprocessorAsIs()) assert len(recording.preprocess_steps) == 1 assert recording.is_preprocessed == True assert recording.preprocess_steps[0].key == PREPROCESS_STEP_BASE_KEY assert recording.epochs_per_task is not None assert len(recording.epochs_per_task) > 0 + def test_mne_preprocessor(): preprocessor = MnePreprocessorRawToHaemo() recording = Recording().load_file(snirf_file1, preprocessor, preprocess=False) @@ -190,13 +220,21 @@ def test_mne_preprocessor(): assert recording.preprocess_step_keys[0] == recording.preprocess_steps[-1].key # can get step by key - assert recording.get_preprocess_step(PREPROCESS_STEP_BASE_KEY).key == PREPROCESS_STEP_BASE_KEY - assert recording.get_preprocess_step(PREPROCESS_STEP_HAEMO_FILTERED_KEY).key == PREPROCESS_STEP_HAEMO_FILTERED_KEY + assert ( + recording.get_preprocess_step(PREPROCESS_STEP_BASE_KEY).key + == PREPROCESS_STEP_BASE_KEY + ) + assert ( + recording.get_preprocess_step(PREPROCESS_STEP_HAEMO_FILTERED_KEY).key + == PREPROCESS_STEP_HAEMO_FILTERED_KEY + ) + # # Dyad # + def test_recording_dyad(): s1 = Recording().load_file(snirf_file1, preprocess=False) s2 = Recording().load_file(snirf_file2, preprocess=False) @@ -207,18 +245,19 @@ def test_recording_dyad(): assert dyad.is_preprocessed == True pairs = dyad.get_pairs(dyad.s1, dyad.s2) - #assert len(pairs) == n_channels * n_channels * n_epochs # TODO this is the test we with, with epochs + # assert len(pairs) == n_channels * n_channels * n_epochs # TODO this is the test we with, with epochs assert len(pairs) == len(s1.preprocessed.ch_names) * len(s2.preprocessed.ch_names) assert pairs[0].label is not None assert pairs[0].label_ch1 == s1.preprocessed.ch_names[0] assert pairs[0].label_ch2 == s2.preprocessed.ch_names[0] assert s1.subject_label in dyad.label + def test_dyad_pairs_recurring_event(): tasks = [ - Task('task1', onset_time=1, duration=1), - Task('task1', onset_time=3, duration=1), - Task('task1', onset_time=5, duration=1), + Task("task1", onset_time=1, duration=1), + Task("task1", onset_time=3, duration=1), + Task("task1", onset_time=5, duration=1), ] # Use the same file for the 2 subjects s1 = Recording(tasks=tasks).load_file(snirf_file1) @@ -232,41 +271,44 @@ def test_dyad_pairs_recurring_event(): assert pairs[0].epoch_idx == 0 assert pairs[1].epoch_idx == 1 assert pairs[2].epoch_idx == 2 - + + def test_dyad_tasks_intersection(): tasks = [ - Task('my_task1', onset_time=1, duration=1), - Task('my_task2', onset_time=3, duration=1), - Task('my_task3', onset_time=5, duration=1), + Task("my_task1", onset_time=1, duration=1), + Task("my_task2", onset_time=3, duration=1), + Task("my_task3", onset_time=5, duration=1), ] s1 = Recording(tasks=[tasks[0], tasks[1]]) s2 = Recording(tasks=[tasks[1], tasks[2]]) assert len(Dyad(s1, s1).tasks) == 2 assert len(Dyad(s2, s2).tasks) == 2 assert len(Dyad(s1, s2).tasks) == 1 - + + def test_dyad_tasks_with_same_name_different_definition(): # tasks with the same name should be considered to be the same task, even if they have a different "definition" (e.g. start at different times) tasks1 = [ - Task('my_task1', onset_event_id=1, duration=1), - Task('my_task2', onset_event_id=3, duration=1), - Task('my_task3', onset_event_id=5, duration=1), - Task('different_name', onset_event_id=99, duration=1), + Task("my_task1", onset_event_id=1, duration=1), + Task("my_task2", onset_event_id=3, duration=1), + Task("my_task3", onset_event_id=5, duration=1), + Task("different_name", onset_event_id=99, duration=1), ] tasks2 = [ - Task('my_task1', onset_event_id=11, duration=1), - Task('my_task2', onset_event_id=13, duration=1), - Task('my_task3', onset_event_id=15, duration=1), - Task('unknown_task', onset_event_id=99, duration=1), + Task("my_task1", onset_event_id=11, duration=1), + Task("my_task2", onset_event_id=13, duration=1), + Task("my_task3", onset_event_id=15, duration=1), + Task("unknown_task", onset_event_id=99, duration=1), ] s1 = Recording(tasks=tasks1) s2 = Recording(tasks=tasks2) merged_task_names = [task.name for task in Dyad(s1, s2).tasks] - assert 'my_task1' in merged_task_names - assert 'my_task2' in merged_task_names - assert 'my_task3' in merged_task_names - assert 'different_name' not in merged_task_names - assert 'unknown_task' not in merged_task_names + assert "my_task1" in merged_task_names + assert "my_task2" in merged_task_names + assert "my_task3" in merged_task_names + assert "different_name" not in merged_task_names + assert "unknown_task" not in merged_task_names + def test_dyad_check_sfreq_same(): s1 = Recording().load_file(snirf_file1) @@ -275,26 +317,30 @@ def test_dyad_check_sfreq_same(): dyad = Dyad(s1, s2) with pytest.raises(Exception): dyad.get_pairs(dyad.s1, dyad.s2) - + + def test_dyad_compute_pair_wtc(): # test with the same subject, so we can check we have a high coherence recordings = get_test_recordings() dyad = Dyad(*recordings) - pair = dyad.get_pairs(dyad.s1, dyad.s2)[0].sub((0, 10)) # Take 10% of the file + pair = dyad.get_pairs(dyad.s1, dyad.s2)[0].sub((0, 10)) # Take 10% of the file wtc = ComplexMorletWavelet().wtc(pair) # Should have a mean of 1 since the first pair is the same signal assert np.mean(wtc.W) == pytest.approx(1) assert wtc.label_dyad == dyad.label + def test_dyad_cwt_cache_during_wtc(): s1, s2 = get_test_recordings() dyad = Dyad(s1, s2) - pair = dyad.get_pairs(dyad.s1, dyad.s2)[0].sub((0, 10)) # Take 10% of the file + pair = dyad.get_pairs(dyad.s1, dyad.s2)[0].sub((0, 10)) # Take 10% of the file wavelet = ComplexMorletWavelet(cache=dict()) - with patch.object(wavelet, 'cwt', wraps=wavelet.cwt) as spy_method: + with patch.object(wavelet, "cwt", wraps=wavelet.cwt) as spy_method: wtc_no_cache = wavelet.wtc(pair) assert spy_method.call_count == 2 - wtc_with_cache = wavelet.wtc(pair) # this should not call cwt, but use the cache + wtc_with_cache = wavelet.wtc( + pair + ) # this should not call cwt, but use the cache assert spy_method.call_count == 2 assert np.all(wtc_no_cache.W == wtc_with_cache.W) # make sure we can clear the cache @@ -302,12 +348,15 @@ def test_dyad_cwt_cache_during_wtc(): wavelet.wtc(pair) assert spy_method.call_count == 4 + def test_dyad_coi_cache_during_wtc(): s1, s2 = get_test_recordings() dyad = Dyad(s1, s2) - pair = dyad.get_pairs(dyad.s1, dyad.s2)[0].sub((0, 10)) # Take 10% of the file + pair = dyad.get_pairs(dyad.s1, dyad.s2)[0].sub((0, 10)) # Take 10% of the file wavelet = ComplexMorletWavelet(cache=dict()) - with patch.object(wavelet, '_get_cone_of_influence', wraps=wavelet._get_cone_of_influence) as spy_method: + with patch.object( + wavelet, "_get_cone_of_influence", wraps=wavelet._get_cone_of_influence + ) as spy_method: wtc_no_cache = wavelet.wtc(pair) assert spy_method.call_count == 1 wtc_with_cache = wavelet.wtc(pair) @@ -319,102 +368,120 @@ def test_dyad_coi_cache_during_wtc(): _ = wavelet.wtc(pair) assert spy_method.call_count == 2 + def test_dyad_cwt_cache_with_different_times(): # When computing intra-subject, we cannot re-use the cache since the cwt might have been computed cropped to match lenght of both signal # We would have this error if the cache is not invalidated and we have different length # "ValueError: operands could not be broadcast together with shapes (40,79) (40,157)" # This can happen for annotation based tasks. Here we force a different task by changing the 2nd subject - s1 = Recording(subject_label='subject1', tasks=[Task('my_task', onset_time=0, duration=10)]).load_file(snirf_file1) - s2 = Recording(subject_label='subject2', tasks=[Task('my_task', onset_time=0, duration=20)]).load_file(snirf_file2) + s1 = Recording( + subject_label="subject1", tasks=[Task("my_task", onset_time=0, duration=10)] + ).load_file(snirf_file1) + s2 = Recording( + subject_label="subject2", tasks=[Task("my_task", onset_time=0, duration=20)] + ).load_file(snirf_file2) # Force a different task length for subject2 to have a different lenght dyad = Dyad(s1, s1) - dyad.s2 = s2 # hack for the test - dyad.is_intra = False # hack for the test + dyad.s2 = s2 # hack for the test + dyad.is_intra = False # hack for the test dyad.compute_wtcs(ch_match=get_test_ch_match_one(), with_intra=True) + def test_dyad_compute_all_wtc(): recordings = get_test_recordings() dyad = Dyad(*recordings) assert dyad.is_wtc_computed == False - dyad.compute_wtcs(only_time_range=(0,10)) + dyad.compute_wtcs(only_time_range=(0, 10)) assert dyad.is_wtc_computed == True - assert len(dyad.wtcs) == len(recordings[0].preprocessed.ch_names)**2 + assert len(dyad.wtcs) == len(recordings[0].preprocessed.ch_names) ** 2 # Should have a mean of 1 since the first pair is the same signal assert np.mean(dyad.wtcs[0].W) == pytest.approx(1) - assert len(dyad.df['channel1'].unique()) == 32 - + assert len(dyad.df["channel1"].unique()) == 32 + + def test_dyad_computes_intra_subject(): s1 = Recording().load_file(snirf_file1) s2 = Recording().load_file(snirf_file2) dyad = Dyad(s1, s2) - dyad.compute_wtcs(only_time_range=(0,10), with_intra=True) + dyad.compute_wtcs(only_time_range=(0, 10), with_intra=True) assert s1.is_wtc_computed == True assert s2.is_wtc_computed == True assert len(dyad.wtcs) == len(s1.intra_wtcs) assert len(dyad.wtcs) == len(s2.intra_wtcs) - + + def test_dyad_computes_intra_subject_channel_match(): s1 = Recording().load_file(snirf_file1) s2 = Recording().load_file(snirf_file2) dyad = Dyad(s1, s2) # have a different channel for each subject - ch_match = ('S1_D1 760', 'S1_D2 760') - dyad.compute_wtcs(ch_match=ch_match, only_time_range=(0,10), with_intra=True) - df_intra = dyad.df[dyad.df['is_intra'] == True] - df_inter = dyad.df[dyad.df['is_intra'] == False] + ch_match = ("S1_D1 760", "S1_D2 760") + dyad.compute_wtcs(ch_match=ch_match, only_time_range=(0, 10), with_intra=True) + df_intra = dyad.df[dyad.df["is_intra"] == True] + df_inter = dyad.df[dyad.df["is_intra"] == False] assert len(df_intra) > 0 - assert np.all(df_intra['channel1'] == df_intra['channel2']) - assert np.all(df_inter['channel1'] != df_inter['channel2']) + assert np.all(df_intra["channel1"] == df_intra["channel2"]) + assert np.all(df_inter["channel1"] != df_inter["channel2"]) # we keep track of the subject id for intra-subject coherence - assert np.sum(df_intra['is_intra_of'] == 1) == 1 - assert np.sum(df_intra['is_intra_of'] == 2) == 1 + assert np.sum(df_intra["is_intra_of"] == 1) == 1 + assert np.sum(df_intra["is_intra_of"] == 2) == 1 + def test_dyad_compute_str_match_wtc(): recordings = get_test_recordings() dyad = Dyad(*recordings) - dyad.compute_wtcs(ch_match='760', only_time_range=(0,10)) + dyad.compute_wtcs(ch_match="760", only_time_range=(0, 10)) assert dyad.is_wtc_computed == True - assert len(dyad.wtcs) == (len(recordings[0].preprocessed.pick('all').ch_names)/2)**2 + assert ( + len(dyad.wtcs) + == (len(recordings[0].preprocessed.pick("all").ch_names) / 2) ** 2 + ) + def test_dyad_compute_regex_match_wtc(): recordings = get_test_recordings() dyad = Dyad(*recordings) - dyad.compute_wtcs(ch_match=get_test_ch_match_few(), only_time_range=(0,10)) + dyad.compute_wtcs(ch_match=get_test_ch_match_few(), only_time_range=(0, 10)) assert len(dyad.wtcs) == 4 assert dyad.wtcs[0].label_pair == dyad.get_pairs(dyad.s1, dyad.s2)[0].label + def test_dyad_compute_tuple_match_wtc(): recordings = get_test_recordings() dyad = Dyad(*recordings) - regex1 = re.compile(r'^S1_D1.*760') - regex2 = re.compile(r'.*760') - dyad.compute_wtcs(ch_match=(regex1, regex2), only_time_range=(0,10)) + regex1 = re.compile(r"^S1_D1.*760") + regex2 = re.compile(r".*760") + dyad.compute_wtcs(ch_match=(regex1, regex2), only_time_range=(0, 10)) assert len(dyad.wtcs) == 16 - #[print(wtc.label) for wtc in dyad.wtcs] + # [print(wtc.label) for wtc in dyad.wtcs] + def test_dyad_compute_list_match_wtc(): recordings = get_test_recordings() dyad = Dyad(*recordings) - ch_list = ['S1_D1 760', 'S1_D2 760', 'S2_D1 760'] - dyad.compute_wtcs(ch_match=ch_list, only_time_range=(0,10)) + ch_list = ["S1_D1 760", "S1_D2 760", "S2_D1 760"] + dyad.compute_wtcs(ch_match=ch_list, only_time_range=(0, 10)) assert len(dyad.wtcs) == len(ch_list) * len(ch_list) - + + def test_dyad_compute_list_per_subject_match_wtc(): recordings = get_test_recordings() dyad = Dyad(*recordings) - ch_list1 = ['S1_D1 760', 'S1_D2 760'] - ch_list2 = ['S2_D1 760', 'S2_D2 760'] - dyad.compute_wtcs(ch_match=(ch_list1, ch_list2), only_time_range=(0,10), with_intra=False) + ch_list1 = ["S1_D1 760", "S1_D2 760"] + ch_list2 = ["S2_D1 760", "S2_D2 760"] + dyad.compute_wtcs( + ch_match=(ch_list1, ch_list2), only_time_range=(0, 10), with_intra=False + ) assert len(dyad.wtcs) == len(ch_list1) * len(ch_list2) - assert len(dyad.df['channel1'].unique()) == 2 + assert len(dyad.df["channel1"].unique()) == 2 + + assert ch_list1[0] in dyad.df["channel1"].unique() + assert ch_list2[0] not in dyad.df["channel1"].unique() - assert ch_list1[0] in dyad.df['channel1'].unique() - assert ch_list2[0] not in dyad.df['channel1'].unique() + assert ch_list1[0] not in dyad.df["channel2"].unique() + assert ch_list2[0] in dyad.df["channel2"].unique() - assert ch_list1[0] not in dyad.df['channel2'].unique() - assert ch_list2[0] in dyad.df['channel2'].unique() - def test_dyad_is_intra_when_same_subject(): recording = Recording().load_file(snirf_file1) @@ -423,7 +490,8 @@ def test_dyad_is_intra_when_same_subject(): # make sure wtcs are marked as "is_intra" dyad.compute_wtcs() - assert dyad.df['is_intra'][0] == True + assert dyad.df["is_intra"][0] == True + def test_access_cwts_after_computation(): recording = Recording().load_file(snirf_file1) @@ -431,10 +499,13 @@ def test_access_cwts_after_computation(): dyad.compute_wtcs() assert len(recording.cwts) == len(recording.mne_raw.ch_names) + def test_dyad_wtc_per_task(): tasks = [ - Task('task1', onset_event_id=1, offset_event_id=TASK_NEXT_EVENT), # these 2 events have different duration - Task('task3', onset_event_id=3, offset_event_id=TASK_NEXT_EVENT), + Task( + "task1", onset_event_id=1, offset_event_id=TASK_NEXT_EVENT + ), # these 2 events have different duration + Task("task3", onset_event_id=3, offset_event_id=TASK_NEXT_EVENT), ] recording1 = Recording(tasks=tasks).load_file(snirf_file1) recording2 = Recording(tasks=tasks).load_file(snirf_file2) @@ -446,13 +517,16 @@ def test_dyad_wtc_per_task(): dyad.compute_wtcs(ch_match=ch_name) assert len(dyad.wtcs) == len(pairs) # must compare the first and last wtcs to make sure we are on different tasks (otherwise we might compare 2 epochs of the same task) - assert dyad.wtcs[0].W.shape[1] != dyad.wtcs[-1].W.shape[1] # not the same duration - assert 'task1' in [wtc.task for wtc in dyad.wtcs] # order may have changed because of task intersection + assert dyad.wtcs[0].W.shape[1] != dyad.wtcs[-1].W.shape[1] # not the same duration + assert "task1" in [ + wtc.task for wtc in dyad.wtcs + ] # order may have changed because of task intersection + def test_dyad_task_annotations_and_time_range_combined(): tasks = [ - Task('task_annotation1', onset_event_id=1, offset_event_id=TASK_NEXT_EVENT), - Task('task_time_range', onset_time=10, duration=20), + Task("task_annotation1", onset_event_id=1, offset_event_id=TASK_NEXT_EVENT), + Task("task_time_range", onset_time=10, duration=20), ] recording1 = Recording(tasks=tasks) recording1.load_file(snirf_file1) @@ -467,33 +541,35 @@ def test_dyad_task_annotations_and_time_range_combined(): dyad.compute_wtcs(ch_match=ch_name) assert len(dyad.wtcs) == len(pairs) found_tasks = [wtc.task for wtc in dyad.wtcs] - assert 'task_annotation1' in found_tasks - assert 'task_time_range' in found_tasks + assert "task_annotation1" in found_tasks + assert "task_time_range" in found_tasks + def test_dyad_wtc_nan_channel_section(): tasks = [ - Task('task1', onset_time=0, duration=500), + Task("task1", onset_time=0, duration=500), ] # Use the same file for the 2 subjects s1 = Recording(tasks=tasks).load_file(snirf_file1) s2 = Recording(tasks=tasks).load_file(snirf_file2) dyad = Dyad(s1, s2) - epochs = s1.get_epochs_for_task('task1') + epochs = s1.get_epochs_for_task("task1") data = epochs.get_data() # Set some values to NaN to split in 3 sections data[0, 0, 20:40] = np.nan - t = int(100*epochs.info['sfreq']) - data[0, 0, t:t+1] = np.nan + t = int(100 * epochs.info["sfreq"]) + data[0, 0, t : t + 1] = np.nan epochs._data = data dyad.compute_wtcs(ch_match=epochs.ch_names[0], downsample=None) df = dyad.df assert len(dyad.wtcs) == 3 # the first section is too small, coherence should be NaN - assert np.all(np.isnan(df[df['section']==0]['coherence'].head(1))) + assert np.all(np.isnan(df[df["section"] == 0]["coherence"].head(1))) # the next 2 sections have enough data for wtc - assert np.all(np.isfinite(df[df['section']==1]['coherence'].head(1))) - assert np.all(np.isfinite(df[df['section']==2]['coherence'].head(1))) + assert np.all(np.isfinite(df[df["section"] == 1]["coherence"].head(1))) + assert np.all(np.isfinite(df[df["section"] == 2]["coherence"].head(1))) + def test_study_wtc(): s1, s2 = get_test_recordings() @@ -514,24 +590,27 @@ def test_study_wtc(): df = study.df assert study.is_wtc_computed == True assert len(dyad1.wtcs) == 1 - + # dyads shuffle are computed only when we want significance assert study.is_wtc_shuffle_computed == False assert study.dyads_shuffled is None - assert np.all(df['is_pseudo'] == False) + assert np.all(df["is_pseudo"] == False) study.compute_wtcs_shuffle(**wtcs_kwargs) df_with_shuffle = study.df - assert len(study.dyads_shuffled) == len(dyads)*(len(dyads)-1) + assert len(study.dyads_shuffled) == len(dyads) * (len(dyads) - 1) assert study.is_wtc_shuffle_computed == True assert len(study.dyads_shuffled[0].wtcs) == 1 - assert len(df_with_shuffle['is_pseudo'].unique()) == 2 + assert len(df_with_shuffle["is_pseudo"].unique()) == 2 + def test_study_is_pseudo_no_duplicate(): recordings = get_test_recordings(4) - study = Study([Dyad(recordings[0], recordings[1]), Dyad(recordings[2], recordings[3])]) + study = Study( + [Dyad(recordings[0], recordings[1]), Dyad(recordings[2], recordings[3])] + ) wtcs_kwargs = dict(ch_match=get_test_ch_match_one()) study.compute_wtcs(**wtcs_kwargs, show_time_estimation=False, with_intra=False) assert len(study.df) == 2 @@ -542,11 +621,11 @@ def test_study_is_pseudo_no_duplicate(): study.compute_wtcs(**wtcs_kwargs, show_time_estimation=False, with_intra=True) assert len(study.df) == 6 - assert len(study.df[study.df['is_intra']==True]) == 4 + assert len(study.df[study.df["is_intra"] == True]) == 4 study.compute_wtcs_shuffle(**wtcs_kwargs) # should not have more "is_intra" assert len(study.df) == 8 - assert len(study.df[study.df['is_intra']==True]) == 4 + assert len(study.df[study.df["is_intra"] == True]) == 4 def test_dyad_computes_whole_record_by_default(): @@ -555,25 +634,30 @@ def test_dyad_computes_whole_record_by_default(): dyad.compute_wtcs(ch_match=get_test_ch_match_one()) assert len(dyad.wtcs) == 1 + def test_dyad_does_not_compute_tasks_when_epochs_not_loaded(): - recording = Recording(tasks=[Task('task1', onset_event_id=1, offset_event_id=TASK_NEXT_EVENT)]) + recording = Recording( + tasks=[Task("task1", onset_event_id=1, offset_event_id=TASK_NEXT_EVENT)] + ) recording.load_file(snirf_file1, preprocess=False) dyad = Dyad(recording, recording) with pytest.raises(Exception): # This should raise an exception, since the epochs have not been loaded from annotations - dyad.compute_wtcs(ch_match=get_test_ch_match_one(), only_time_range=(0,10)) + dyad.compute_wtcs(ch_match=get_test_ch_match_one(), only_time_range=(0, 10)) + def test_dyad_coherence_pandas(): s1, s2 = get_test_recordings() dyad = Dyad(s1, s2) dyad.compute_wtcs(ch_match=get_test_ch_match_few()) df = dyad._get_coherence_df() - assert len(df['task'].unique()) == len(dyad.s1.task_keys) - assert len(df['channel1'].unique()) == 2 - assert len(df['channel2'].unique()) == 2 - assert len(df['dyad'].unique()) == 1 - assert df['subject1'].unique()[0] == s1.subject_label - assert df['subject2'].unique()[0] == s2.subject_label + assert len(df["task"].unique()) == len(dyad.s1.task_keys) + assert len(df["channel1"].unique()) == 2 + assert len(df["channel2"].unique()) == 2 + assert len(df["dyad"].unique()) == 1 + assert df["subject1"].unique()[0] == s1.subject_label + assert df["subject2"].unique()[0] == s2.subject_label + def test_dyad_coherence_pandas_with_intra(): s1, s2 = get_test_recordings() @@ -586,30 +670,36 @@ def test_dyad_coherence_pandas_with_intra(): dyad.compute_wtcs(ch_match=get_test_ch_match_few(), with_intra=True) df = dyad._get_coherence_df(with_intra=True) - assert len(df['task'].unique()) == len(dyad.s1.task_keys) - assert len(df['dyad'].unique()) == 3 - assert len(df['is_intra'].unique()) == 2 + assert len(df["task"].unique()) == len(dyad.s1.task_keys) + assert len(df["dyad"].unique()) == 3 + assert len(df["is_intra"].unique()) == 2 + def test_study_coherence_pandas(): s1, s2 = get_test_recordings() s3, _ = get_test_recordings() - dyad1 = Dyad(s1, s2, label='dyad1') - dyad2 = Dyad(s1, s3, label='dyad2') - dyad3 = Dyad(s2, s3, label='dyad3') + dyad1 = Dyad(s1, s2, label="dyad1") + dyad2 = Dyad(s1, s3, label="dyad2") + dyad3 = Dyad(s2, s3, label="dyad3") study = Study([dyad1, dyad2, dyad3]) - study.compute_wtcs(ch_match=get_test_ch_match_few(), with_intra=False, show_time_estimation=False) + study.compute_wtcs( + ch_match=get_test_ch_match_few(), with_intra=False, show_time_estimation=False + ) df = study.df - assert len(df['task'].unique()) == len(dyad1.s1.task_keys) - assert len(df['channel1'].unique()) == 2 - assert len(df['channel2'].unique()) == 2 - assert len(df['dyad'].unique()) == 3 - assert np.all(df['is_intra'] == False) + assert len(df["task"].unique()) == len(dyad1.s1.task_keys) + assert len(df["channel1"].unique()) == 2 + assert len(df["channel2"].unique()) == 2 + assert len(df["dyad"].unique()) == 3 + assert np.all(df["is_intra"] == False) + def test_study_coherence_pandas_with_intra(): - study = Study([Dyad(*get_test_recordings(), label='dyad1')]) - study.compute_wtcs(ch_match=get_test_ch_match_few(), with_intra=True, show_time_estimation=False) + study = Study([Dyad(*get_test_recordings(), label="dyad1")]) + study.compute_wtcs( + ch_match=get_test_ch_match_few(), with_intra=True, show_time_estimation=False + ) df = study.df - assert len(df['is_intra'].unique()) == 2 + assert len(df["is_intra"].unique()) == 2 def test_dyad_coherence_pandas_on_roi(): @@ -617,8 +707,9 @@ def test_dyad_coherence_pandas_on_roi(): dyad = Dyad(s1, s2) dyad.compute_wtcs(ch_match=get_test_ch_match_few()) df = dyad._get_coherence_df() - assert len(df['roi1']) == 4 - assert len(df['roi2']) == 4 + assert len(df["roi1"]) == 4 + assert len(df["roi2"]) == 4 + def test_wtc_downsampling(): recordings = get_test_recordings() @@ -626,7 +717,7 @@ def test_wtc_downsampling(): n = 100 Study([dyad]).compute_wtcs(ch_match=get_test_ch_match_one(), downsample=n) assert len(dyad.wtcs[0].times) <= n - + def test_save_study_to_disk(): subject = get_test_recording() @@ -637,22 +728,24 @@ def test_save_study_to_disk(): file_path = temp_file.name study.save_pickle(file_path) study_reloaded = Study.from_pickle(file_path) - + assert len(study_reloaded.dyads) == len(study.dyads) + def test_save_study_df_to_disk(): recordings = get_test_recordings() dyad = Dyad(*recordings) study = Study([dyad]) study.compute_wtcs(show_time_estimation=False, with_intra=False) - with tempfile.NamedTemporaryFile(suffix='.feather') as temp_file: + with tempfile.NamedTemporaryFile(suffix=".feather") as temp_file: file_path = temp_file.name study.save_feather(file_path) df = CoherenceDataFrame.from_feather(file_path) - assert np.all(df['subject1'] == recordings[0].subject_label) - + assert np.all(df["subject1"] == recordings[0].subject_label) + + def test_study_run_estimation(capsys): recordings = get_test_recordings() dyads = [] @@ -661,57 +754,60 @@ def test_study_run_estimation(capsys): study = Study(dyads) study.estimate_wtcs_run_time() out = capsys.readouterr() - assert 'time' in str(out) + assert "time" in str(out) + def test_lionirs_channel_grouping(): - roi_file_path = datasets.fetch('fnirs/lionirs/channel_grouping_7ROI.mat') + roi_file_path = datasets.fetch("fnirs/lionirs/channel_grouping_7ROI.mat") croi = ChannelROI.from_lionirs_file(roi_file_path) assert len(croi.rois.keys()) == 14 key1 = list(croi.rois.keys())[0] roi1 = croi.rois[key1] ch1 = roi1[0] - assert key1 == 'PreFr_L' + assert key1 == "PreFr_L" assert len(roi1) == 5 - assert ch1 == 'S1_D9' + assert ch1 == "S1_D9" assert len(croi.ordered_ch_names) == 50 # Actual ordering: # ['S1_D9', 'S2_D9', 'S2_D2', 'S1_D2', 'S2_D1', 'S3_D1', 'S3_D2', ...] - unordered_names = ['S1_D2 hbo', 'S1_D2 hbr', 'S1_D9 hbo', 'S1_D9 hbr', 'whatever'] + unordered_names = ["S1_D2 hbo", "S1_D2 hbr", "S1_D9 hbo", "S1_D9 hbr", "whatever"] ordered_names = croi.get_ch_names_in_order(unordered_names) - assert ordered_names[0] == 'S1_D9 hbo' - assert ordered_names[1] == 'S1_D9 hbr' - assert ordered_names[2] == 'S1_D2 hbo' - assert ordered_names[3] == 'S1_D2 hbr' - assert ordered_names[4] == 'whatever' + assert ordered_names[0] == "S1_D9 hbo" + assert ordered_names[1] == "S1_D9 hbr" + assert ordered_names[2] == "S1_D2 hbo" + assert ordered_names[3] == "S1_D2 hbr" + assert ordered_names[4] == "whatever" assert croi.group_boundaries_sizes[:2] == [0, 5] assert croi.group_boundaries_sizes[-1] == len(croi.ordered_ch_names) - assert croi.get_roi_from_channel(ordered_names[0]) == 'PreFr_L' + assert croi.get_roi_from_channel(ordered_names[0]) == "PreFr_L" + def test_ordered_recording_ch_names(): - roi_file_path = datasets.fetch('fnirs/lionirs/channel_grouping_7ROI.mat') + roi_file_path = datasets.fetch("fnirs/lionirs/channel_grouping_7ROI.mat") croi = ChannelROI.from_lionirs_file(roi_file_path) recording = Recording(channel_roi=croi).load_file(snirf_file1) ch_names = recording.ordered_ch_names - assert ch_names[0] == 'S2_D2 760' + assert ch_names[0] == "S2_D2 760" def test_positions_to_standard_montage(): - #recording = get_test_recording() + # recording = get_test_recording() recording = Recording().load_file(datasets.fnirs_slow_breathing()) map = recording.get_channel_to_standard_montage_map(as_data_frame=False) - assert map['S1']['name'] == 'AF7' - assert map['S1_D2']['name'] == 'AF7' + assert map["S1"]["name"] == "AF7" + assert map["S1_D2"]["name"] == "AF7" + def test_positions_to_standard_montage_as_data_frame(): - #recording = get_test_recording() + # recording = get_test_recording() recording = Recording().load_file(snirf_file1) df = recording.get_channel_to_standard_montage_map() - assert df[df.columns[0]][0] == 'S1' + assert df[df.columns[0]][0] == "S1" # Skip this test because it downloads data. We don't want this on the CI @@ -723,5 +819,5 @@ def test_download_demos(): new_paths = browser.paths assert len(new_paths) == previous_count + 1 -#def test_load_lionirs(): +# def test_load_lionirs(): diff --git a/tests/test_fnirs_cedalion.py b/tests/test_fnirs_cedalion.py index af4dbcc..1e6ee13 100644 --- a/tests/test_fnirs_cedalion.py +++ b/tests/test_fnirs_cedalion.py @@ -1,14 +1,17 @@ import pytest from hypyp.fnirs.preprocessor.base_step import PREPROCESS_STEP_BASE_KEY -from hypyp.fnirs.preprocessor.implementations.cedalion_preprocessor import CedalionPreprocessor +from hypyp.fnirs.preprocessor.implementations.cedalion_preprocessor import ( + CedalionPreprocessor, +) + +snirf_file = "./data/NIRS/DCARE_02_sub1.snirf" -snirf_file = './data/NIRS/DCARE_02_sub1.snirf' def test_cedalion_preprocessor(): if CedalionPreprocessor is None: pytest.skip("Optional dependency Cedalion is not installed") - + import xarray as xr preprocessor = CedalionPreprocessor() @@ -17,7 +20,7 @@ def test_cedalion_preprocessor(): step = steps[0] assert step.key == PREPROCESS_STEP_BASE_KEY assert isinstance(step.obj, xr.DataArray) - assert step.n_times == len(step.obj['time']) + assert step.n_times == len(step.obj["time"]) assert step.sfreq == 7.8125 assert len(step.ch_names) > 0 - assert step.ch_names[0] == 'S1D1' + assert step.ch_names[0] == "S1D1" diff --git a/tests/test_pdc.py b/tests/test_pdc.py index 60df6aa..27ac1e8 100644 --- a/tests/test_pdc.py +++ b/tests/test_pdc.py @@ -7,7 +7,6 @@ class pdc_test(unittest.TestCase): - def test_function(self): # three sources a <- b <- c # check connectivity is zero where expected @@ -15,17 +14,34 @@ def test_function(self): m, mp = var_coeffs_init.shape n_fft = 10 p = mp // m - var_coeffs = np.reshape(var_coeffs_init, (m, m, p), 'c') + var_coeffs = np.reshape(var_coeffs_init, (m, m, p), "c") cov_mat = np.eye(3) a = fft(np.dstack([np.eye(m), -var_coeffs]), n_fft * 2 - 1)[:, :, :n_fft] pdc_result = np.abs(a / np.sqrt(np.sum(a.conj() * a, axis=0, keepdims=True))) - assert_array_equal(np.isclose(np.sum(np.abs(a), 2), 0), np.isclose(var_coeffs_init + cov_mat, 0)) - self.assertFalse(np.all(np.sum(np.abs(a), 2) == np.sum(np.abs(np.transpose(a, axes=[1, 0, 2])), 2))) - assert_array_equal(np.isclose(np.sum(np.abs(pdc_result), 2), 0), np.isclose(var_coeffs_init + cov_mat, 0)) - self.assertFalse(np.all(np.sum(pdc_result, 2) == np.sum(np.transpose(pdc_result, axes=[1, 0, 2]), 2))) + assert_array_equal( + np.isclose(np.sum(np.abs(a), 2), 0), + np.isclose(var_coeffs_init + cov_mat, 0), + ) + self.assertFalse( + np.all( + np.sum(np.abs(a), 2) + == np.sum(np.abs(np.transpose(a, axes=[1, 0, 2])), 2) + ) + ) + assert_array_equal( + np.isclose(np.sum(np.abs(pdc_result), 2), 0), + np.isclose(var_coeffs_init + cov_mat, 0), + ) + self.assertFalse( + np.all( + np.sum(pdc_result, 2) + == np.sum(np.transpose(pdc_result, axes=[1, 0, 2]), 2) + ) + ) self.assertEqual(np.sum(pdc_result, 2)[0, 0], n_fft) self.assertEqual(np.sum(pdc_result, 2)[1, 1], np.sum(pdc_result, 2)[2, 2]) self.assertEqual(np.sum(pdc_result, 2)[0, 1], np.sum(pdc_result, 2)[1, 2]) -if __name__ == '__main__': - unittest.main() \ No newline at end of file + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prep.py b/tests/test_prep.py index a0da896..3340c5a 100644 --- a/tests/test_prep.py +++ b/tests/test_prep.py @@ -5,6 +5,7 @@ from hypyp import prep + def test_filt(): n_channels = 10 n_samples = 100000 @@ -13,81 +14,88 @@ def test_filt(): random_state = np.random.RandomState(seed=42) # create noise signal - info = mne.create_info(ch_names=n_channels, sfreq=sfreq, ch_types='eeg') + info = mne.create_info(ch_names=n_channels, sfreq=sfreq, ch_types="eeg") data = random_state.normal(size=(n_channels, n_samples)) raw = mne.io.RawArray(data, info) raw_psd = raw.compute_psd().get_data() # compare power on lowest freq # Take the first few frequencies and not just index 0 - raw_filt_default, = prep.filt([raw]) + (raw_filt_default,) = prep.filt([raw]) raw_filt_default_psd = raw_filt_default.compute_psd().get_data() - assert np.sum(raw_filt_default_psd[:,:3]) < np.sum(raw_psd[:,:3]) + assert np.sum(raw_filt_default_psd[:, :3]) < np.sum(raw_psd[:, :3]) # compare power on highest freq - raw_filt, = prep.filt([raw], (2., 10)) + (raw_filt,) = prep.filt([raw], (2.0, 10)) raw_filt_psd = raw_filt.compute_psd().get_data() - assert np.sum(raw_filt_psd[:,-1]) < np.sum(raw_filt_default_psd[:,-1]) - - -@pytest.mark.parametrize("fit_kwargs", [ - dict(method='fastica', fit_params=dict(tol=0.01)), # increase tolerance to converge - dict(method='infomax', fit_params=dict(extended=True)) -]) + assert np.sum(raw_filt_psd[:, -1]) < np.sum(raw_filt_default_psd[:, -1]) + + +@pytest.mark.parametrize( + "fit_kwargs", + [ + dict( + method="fastica", fit_params=dict(tol=0.01) + ), # increase tolerance to converge + dict(method="infomax", fit_params=dict(extended=True)), + ], +) def test_ICA(epochs, fit_kwargs): ep = [epochs.epo1, epochs.epo2] icas = prep.ICA_fit(ep, n_components=15, **fit_kwargs, random_state=97) - + assert len(icas) == len(ep) # check that the number of componenents is similar between the two participants - for i in range(0, len(icas)-1): - assert mne.preprocessing.ICA.get_components( - icas[i]).shape == mne.preprocessing.ICA.get_components(icas[i+1]).shape + for i in range(0, len(icas) - 1): + assert ( + mne.preprocessing.ICA.get_components(icas[i]).shape + == mne.preprocessing.ICA.get_components(icas[i + 1]).shape + ) cleaned_epochs_ICA = prep.ICA_apply(icas, 0, 0, ep, plot=False) # check bad channels are not deleted - assert epochs.epo1.info['ch_names'] == cleaned_epochs_ICA[0].info['ch_names'] - assert epochs.epo2.info['ch_names'] == cleaned_epochs_ICA[1].info['ch_names'] + assert epochs.epo1.info["ch_names"] == cleaned_epochs_ICA[0].info["ch_names"] + assert epochs.epo2.info["ch_names"] == cleaned_epochs_ICA[1].info["ch_names"] # check signal change by comparing total amplitude raw_amplitudes = np.mean(np.abs(epochs.epo1.get_data(copy=True)), axis=1) - processed_amplitudes = np.mean(np.abs(cleaned_epochs_ICA[0].get_data(copy=True)), axis=1) + processed_amplitudes = np.mean( + np.abs(cleaned_epochs_ICA[0].get_data(copy=True)), axis=1 + ) assert np.sum(processed_amplitudes) < np.sum(raw_amplitudes) -@pytest.mark.parametrize("AR_local_kwargs", [ - dict(strategy='union'), - dict(strategy='intersection'), -]) +@pytest.mark.parametrize( + "AR_local_kwargs", + [ + dict(strategy="union"), + dict(strategy="intersection"), + ], +) def test_AR_local(epochs, AR_local_kwargs): # test on epochs, but usually applied on cleaned epochs with ICA cleaned_epochs_AR, dic_AR = prep.AR_local( - [epochs.epo1, epochs.epo2], - **AR_local_kwargs, - threshold=50.0, - verbose=False + [epochs.epo1, epochs.epo2], **AR_local_kwargs, threshold=50.0, verbose=False ) assert len(epochs.epo1) >= len(cleaned_epochs_AR[0]) assert len(epochs.epo2) >= len(cleaned_epochs_AR[1]) assert len(cleaned_epochs_AR[0]) == len(cleaned_epochs_AR[1]) - assert dic_AR['S2'] + dic_AR['S1'] == dic_AR['dyad'] - assert dic_AR['S2'] <= dic_AR['dyad'] + assert dic_AR["S2"] + dic_AR["S1"] == dic_AR["dyad"] + assert dic_AR["S2"] <= dic_AR["dyad"] -@pytest.mark.parametrize("AR_local_kwargs", [ - dict(strategy='union'), - dict(strategy='intersection'), -]) +@pytest.mark.parametrize( + "AR_local_kwargs", + [ + dict(strategy="union"), + dict(strategy="intersection"), + ], +) def test_AR_local_exception(epochs, AR_local_kwargs): # test the threshold with pytest.raises(Exception): prep.AR_local( - [epochs.epo1, epochs.epo2], - **AR_local_kwargs, - threshold=0.0, - verbose=False + [epochs.epo1, epochs.epo2], **AR_local_kwargs, threshold=0.0, verbose=False ) - - diff --git a/tests/test_signal.py b/tests/test_signal.py index 56e1606..bd25b58 100644 --- a/tests/test_signal.py +++ b/tests/test_signal.py @@ -4,6 +4,7 @@ from hypyp.signal import SyntheticSignal + def test_instanciate(): signal = SyntheticSignal(duration=300) assert signal is not None @@ -11,28 +12,29 @@ def test_instanciate(): assert len(signal.y) == len(signal.x) assert np.sum(signal.y) == 0 + def add_noise(): signal = SyntheticSignal(duration=300) signal.add_noise() assert np.sum(np.abs(signal.y)) > 0 + def test_sin(): signal = SyntheticSignal(duration=300) signal.add_sin(freq=1) assert len(signal.y) == len(signal.x) assert signal.y[0] == 0 + def test_chirp(): signal = SyntheticSignal(duration=300) signal.add_chirp(1, 5) assert len(signal.y) == len(signal.x) # TODO should test the frequencies + def test_custom(): n = 300 signal = SyntheticSignal(duration=n, n_points=n) - signal.add_custom(np.ones((n, ))) + signal.add_custom(np.ones((n,))) assert np.sum(signal.y) == n - - - diff --git a/tests/test_stats.py b/tests/test_stats.py index 19cc4bd..5cb567b 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -20,43 +20,48 @@ def test_metaconn_matrix_2brains(epochs): # computing ch_con and sensors pairs for metaconn calculation con_matrixTuple = stats.con_matrix(epochs.epo1, freq, draw=False) ch_con_freq = con_matrixTuple.ch_con_freq - sensor_pairs = analyses.indices_connectivity_interbrain( - epochs.epoch_merge) + sensor_pairs = analyses.indices_connectivity_interbrain(epochs.epoch_merge) # computing metaconn_freq and test it - metaconn_matrix_2brainsTuple = stats.metaconn_matrix_2brains(sensor_pairs, - con_matrixTuple.ch_con, freq) + metaconn_matrix_2brainsTuple = stats.metaconn_matrix_2brains( + sensor_pairs, con_matrixTuple.ch_con, freq + ) metaconn_freq = metaconn_matrix_2brainsTuple.metaconn_freq # take a random ch_name: random.seed(20) # Init the random number generator for reproducibility # n = random.randrange(0, 63) # for our data taske into account EOG ch!!! - n = random.randrange(0, len(epochs.epo1.info['ch_names'])) - tot = len(epochs.epo1.info['ch_names']) - p = random.randrange(len(epochs.epo1.info['ch_names']), len( - epochs.epoch_merge.info['ch_names'])+1) + n = random.randrange(0, len(epochs.epo1.info["ch_names"])) + tot = len(epochs.epo1.info["ch_names"]) + p = random.randrange( + len(epochs.epo1.info["ch_names"]), len(epochs.epoch_merge.info["ch_names"]) + 1 + ) # checking for each pair in which ch_name is, # whether ch_name linked himself # (in neighbouring frequencies also) - assert metaconn_freq[n+tot, p] == metaconn_freq[n, p] - assert metaconn_freq[n-tot, p] == metaconn_freq[n, p] - assert metaconn_freq[n+tot, p+tot] == metaconn_freq[n, p] - assert metaconn_freq[n-tot, p-tot] == metaconn_freq[n, p] - assert metaconn_freq[n, p+tot] == metaconn_freq[n, p] - assert metaconn_freq[n, p-tot] == metaconn_freq[n, p] + assert metaconn_freq[n + tot, p] == metaconn_freq[n, p] + assert metaconn_freq[n - tot, p] == metaconn_freq[n, p] + assert metaconn_freq[n + tot, p + tot] == metaconn_freq[n, p] + assert metaconn_freq[n - tot, p - tot] == metaconn_freq[n, p] + assert metaconn_freq[n, p + tot] == metaconn_freq[n, p] + assert metaconn_freq[n, p - tot] == metaconn_freq[n, p] # and not in the other frequencies if metaconn_freq[n, p] == 1: for i in range(1, len(freq)): - assert metaconn_freq[n+tot*(i+1), p] != metaconn_freq[n, p] - assert metaconn_freq[n-tot*(i+1), p] != metaconn_freq[n, p] - assert metaconn_freq[n+tot * - (i+1), p+tot*(i+1)] != metaconn_freq[n, p] - assert metaconn_freq[n-tot * - (i+1), p-tot*(i+1)] != metaconn_freq[n, p] - assert metaconn_freq[n, p+tot*(i+1)] != metaconn_freq[n, p] - assert metaconn_freq[n, p-tot*(i+1)] != metaconn_freq[n, p] + assert metaconn_freq[n + tot * (i + 1), p] != metaconn_freq[n, p] + assert metaconn_freq[n - tot * (i + 1), p] != metaconn_freq[n, p] + assert ( + metaconn_freq[n + tot * (i + 1), p + tot * (i + 1)] + != metaconn_freq[n, p] + ) + assert ( + metaconn_freq[n - tot * (i + 1), p - tot * (i + 1)] + != metaconn_freq[n, p] + ) + assert metaconn_freq[n, p + tot * (i + 1)] != metaconn_freq[n, p] + assert metaconn_freq[n, p - tot * (i + 1)] != metaconn_freq[n, p] # check for each f if connects to the good other ch and not to more - assert metaconn_freq[n+tot*i, p+tot*i] == ch_con_freq[n, p-tot] + assert metaconn_freq[n + tot * i, p + tot * i] == ch_con_freq[n, p - tot] def test_PSD(epochs): @@ -65,24 +70,22 @@ def test_PSD(epochs): """ fmin = 10 fmax = 13 - psd_tuple = analyses.pow(epochs.epo1, - fmin, fmax, - n_fft=256, - n_per_seg=None, - epochs_average=True) + psd_tuple = analyses.pow( + epochs.epo1, fmin, fmax, n_fft=256, n_per_seg=None, epochs_average=True + ) psd = psd_tuple.psd freq_list = psd_tuple.freq_list assert type(psd) == np.ndarray - assert psd.shape == ( - len(epochs.epo1.info['ch_names']), len(freq_list)) - psd_tuple = analyses.pow(epochs.epo1, - fmin, fmax, - n_fft=256, - n_per_seg=None, - epochs_average=False) + assert psd.shape == (len(epochs.epo1.info["ch_names"]), len(freq_list)) + psd_tuple = analyses.pow( + epochs.epo1, fmin, fmax, n_fft=256, n_per_seg=None, epochs_average=False + ) psd = psd_tuple.psd - assert psd.shape == (len(epochs.epo1), len( - epochs.epo1.info['ch_names']), len(freq_list)) + assert psd.shape == ( + len(epochs.epo1), + len(epochs.epo1.info["ch_names"]), + len(freq_list), + ) def test_behav_corr(epochs): @@ -98,42 +101,54 @@ def test_behav_corr(epochs): p_thresh = 0.05 - corr_tuple = analyses.behav_corr(data, behav, - data_name='epochs', - behav_name='time', - p_thresh=p_thresh, - multiple_corr=False, - verbose=False) + corr_tuple = analyses.behav_corr( + data, + behav, + data_name="epochs", + behav_name="time", + p_thresh=p_thresh, + multiple_corr=False, + verbose=False, + ) assert pytest.approx(corr_tuple.r) in [-1, 1] # test for connectivity values data # generate artificial group of 2 subjects repeated mne.epochs.equalize_epoch_counts([epochs.epo1, epochs.epo2]) assert len(epochs.epo1) == len(epochs.epo2) - con_ind = analyses.pair_connectivity(np.array([epochs.epo1, epochs.epo1]), - sampling_rate=epochs.epo1.info['sfreq'], - frequencies=[8, 10], - mode='ccorr', - epochs_average=True) - con_subj = analyses.pair_connectivity(np.array([epochs.epo1, epochs.epo2]), - sampling_rate=epochs.epo1.info['sfreq'], - frequencies=[8, 10], - mode='ccorr', - epochs_average=True) + con_ind = analyses.pair_connectivity( + np.array([epochs.epo1, epochs.epo1]), + sampling_rate=epochs.epo1.info["sfreq"], + frequencies=[8, 10], + mode="ccorr", + epochs_average=True, + ) + con_subj = analyses.pair_connectivity( + np.array([epochs.epo1, epochs.epo2]), + sampling_rate=epochs.epo1.info["sfreq"], + frequencies=[8, 10], + mode="ccorr", + epochs_average=True, + ) # remove frequency dimension con_ind = np.mean(con_ind, axis=0) con_subj = np.mean(con_subj, axis=0) assert con_ind.shape == (62, 62) assert con_subj.shape == (62, 62) - data = np.stack((con_ind, con_subj, con_subj, con_subj, con_subj, con_subj, con_subj)) + data = np.stack( + (con_ind, con_subj, con_subj, con_subj, con_subj, con_subj, con_subj) + ) behav = np.array([0, 1, 1, 1, 1, 1, 1]) # correlate connectivity and behaviour across pairs without multiple comparison correction - corr_tuple = analyses.behav_corr(data, behav, - data_name='ccorr', - behav_name='imitation score', - p_thresh=p_thresh, - multiple_corr=False, - verbose=True) + corr_tuple = analyses.behav_corr( + data, + behav, + data_name="ccorr", + behav_name="imitation score", + p_thresh=p_thresh, + multiple_corr=False, + verbose=True, + ) # test that there is a correlation (repeated measures) significant_r = [] for i in range(0, corr_tuple.r.shape[0]): @@ -143,12 +158,15 @@ def test_behav_corr(epochs): assert len(significant_r) != 0 # correlate connectivity and behaviour across pairs with multiple comparison correction - corr_tuple = analyses.behav_corr(data, behav, - data_name='ccorr', - behav_name='imitation score', - p_thresh=p_thresh, - multiple_corr=True, - verbose=True) + corr_tuple = analyses.behav_corr( + data, + behav, + data_name="ccorr", + behav_name="imitation score", + p_thresh=p_thresh, + multiple_corr=True, + verbose=True, + ) # test that there is a correlation (repeated measures) significant_r = [] for i in range(0, corr_tuple.r.shape[0]): @@ -162,22 +180,27 @@ def test_behav_corr(epochs): for k in range(0, 5): random_r1 = utils.generate_random_epoch(epochs.epo1, mu=0, sigma=0.01) random_r2 = utils.generate_random_epoch(epochs.epo2, mu=4, sigma=0.01) - con = analyses.pair_connectivity(np.array([random_r1, random_r2]), - sampling_rate=epochs.epo1.info['sfreq'], - frequencies=[8, 10], - mode='ccorr', - epochs_average=True) + con = analyses.pair_connectivity( + np.array([random_r1, random_r2]), + sampling_rate=epochs.epo1.info["sfreq"], + frequencies=[8, 10], + mode="ccorr", + epochs_average=True, + ) data.append(np.mean(con, axis=0)) data = np.mean(np.array([data]), axis=0) # correlate connectivity and behaviour across pairs dyads = data.shape[0] behav = np.arange(0, dyads) - corr_tuple = analyses.behav_corr(data, behav, - data_name='ccorr', - behav_name='imitation score', - p_thresh=p_thresh, - multiple_corr=True, - verbose=False) + corr_tuple = analyses.behav_corr( + data, + behav, + data_name="ccorr", + behav_name="imitation score", + p_thresh=p_thresh, + multiple_corr=True, + verbose=False, + ) # test that there is no correlation (random measures) for i in range(0, corr_tuple.r.shape[0]): for j in range(0, corr_tuple.r.shape[1]): @@ -192,15 +215,14 @@ def test_indexes_connectivity(epochs): Test index intra- and inter-brains """ electrodes = analyses.indices_connectivity_intrabrain(epochs.epo1) - length = len(epochs.epo1.info['ch_names']) + length = len(epochs.epo1.info["ch_names"]) L = [] for i in range(1, length): - L.append(length-i) + L.append(length - i) tot = sum(L) assert len(electrodes) == tot - electrodes_hyper = analyses.indices_connectivity_interbrain( - epochs.epoch_merge) - assert len(electrodes_hyper) == length*length + electrodes_hyper = analyses.indices_connectivity_interbrain(epochs.epoch_merge) + assert len(electrodes_hyper) == length * length # format that do not work for mne.spectral_connectivity @@ -214,40 +236,35 @@ def test_stats(epochs): fmin = 10 fmax = 13 - psd_tuple = analyses.pow(random_r1, - fmin, fmax, - n_fft=256, - n_per_seg=None, - epochs_average=False) + psd_tuple = analyses.pow( + random_r1, fmin, fmax, n_fft=256, n_per_seg=None, epochs_average=False + ) psd = psd_tuple.psd statsCondTuple = stats.statsCond(psd, random_r1, 3000, 0.05) - assert statsCondTuple.T_obs.shape[0] == len(epochs.epo1.info['ch_names']) + assert statsCondTuple.T_obs.shape[0] == len(epochs.epo1.info["ch_names"]) for i in range(0, len(statsCondTuple.p_values)): assert statsCondTuple.p_values[i] <= statsCondTuple.adj_p[1][i] - assert statsCondTuple.T_obs_plot.shape[0] == len( - epochs.epo1.info['ch_names']) - - psd_tuple2 = analyses.pow(random_r2, - fmin, fmax, - n_fft=256, - n_per_seg=None, - epochs_average=False) + assert statsCondTuple.T_obs_plot.shape[0] == len(epochs.epo1.info["ch_names"]) + + psd_tuple2 = analyses.pow( + random_r2, fmin, fmax, n_fft=256, n_per_seg=None, epochs_average=False + ) psd2 = psd_tuple2.psd freq_list = psd_tuple2.freq_list data = [psd, psd2] con_matrixTuple = stats.con_matrix(random_r1, freq_list, draw=False) - statscondClusterTuple = stats.statscondCluster(data, - freq_list, - scipy.sparse.bsr_matrix( - con_matrixTuple.ch_con_freq), - tail=0, - n_permutations=3000, - alpha=0.05) - assert statscondClusterTuple.F_obs.shape[0] == len( - epochs.epo1.info['ch_names']) + statscondClusterTuple = stats.statscondCluster( + data, + freq_list, + scipy.sparse.bsr_matrix(con_matrixTuple.ch_con_freq), + tail=0, + n_permutations=3000, + alpha=0.05, + ) + assert statscondClusterTuple.F_obs.shape[0] == len(epochs.epo1.info["ch_names"]) # for i in range(0, len(statscondClusterTuple.clusters)): # assert len(np.where(statscondClusterTuple.clusters[i])=='True') < len( # epochs.epo1.info['ch_names']) @@ -262,19 +279,18 @@ def test_utils(epochs): ep_hyper = utils.merge(epochs.epo1, epochs.epo2) assert type(ep_hyper) == mne.epochs.EpochsArray # check channels number - assert len(ep_hyper.info['ch_names']) == 2 * \ - len(epochs.epo1.info['ch_names']) + assert len(ep_hyper.info["ch_names"]) == 2 * len(epochs.epo1.info["ch_names"]) # check EOG channels number # check data for S2 or 1 correspond in the ep_hyper, on channel n and # epoch n, randomnly assigned random.seed(10) - nch = random.randrange(0, len(epochs.epo1.info['ch_names'])) + nch = random.randrange(0, len(epochs.epo1.info["ch_names"])) ne = random.randrange(0, len(epochs.epo1)) - ch_name = epochs.epo1.info['ch_names'][nch] - liste = ep_hyper.info['ch_names'] - ch_index1 = liste.index(ch_name + '_S1') - ch_index2 = liste.index(ch_name + '_S2') + ch_name = epochs.epo1.info["ch_names"][nch] + liste = ep_hyper.info["ch_names"] + ch_index1 = liste.index(ch_name + "_S1") + ch_index2 = liste.index(ch_name + "_S2") ep_hyper_data = ep_hyper.get_data(copy=True) epo1_data = epochs.epo1.get_data(copy=True) epo2_data = epochs.epo2.get_data(copy=True) @@ -284,7 +300,9 @@ def test_utils(epochs): def test_compute_nmPLV(epochs): - result = analyses.compute_nmPLV(data=epochs, sampling_rate=500, freq_range1=[4, 8], freq_range2=[10, 14]) + result = analyses.compute_nmPLV( + data=epochs, sampling_rate=500, freq_range1=[4, 8], freq_range2=[10, 14] + ) hPLV = result[:, :31, 31:].mean() PLV1 = result[:, :31, :31].mean() PLV2 = result[:, 31:, 31:].mean() @@ -298,38 +316,38 @@ def test_pair_connectivity_accorr(epochs): Test adjusted circular correlation (accorr) connectivity metric. """ mne.epochs.equalize_epoch_counts([epochs.epo1, epochs.epo2]) - + # Test with epochs_average=True con = analyses.pair_connectivity( np.array([epochs.epo1, epochs.epo2]), - sampling_rate=epochs.epo1.info['sfreq'], + sampling_rate=epochs.epo1.info["sfreq"], frequencies=[8, 12], - mode='accorr', - epochs_average=True + mode="accorr", + epochs_average=True, ) - - n_channels = len(epochs.epo1.info['ch_names']) + + n_channels = len(epochs.epo1.info["ch_names"]) # Shape should be (n_freq, 2*n_channels, 2*n_channels) assert con.shape[1] == 2 * n_channels assert con.shape[2] == 2 * n_channels - + # Values should be bounded (typically between -1 and 1 for correlation) assert np.all(np.isfinite(con)) - + # Test with epochs_average=False con_no_avg = analyses.pair_connectivity( np.array([epochs.epo1, epochs.epo2]), - sampling_rate=epochs.epo1.info['sfreq'], + sampling_rate=epochs.epo1.info["sfreq"], frequencies=[8, 12], - mode='accorr', - epochs_average=False + mode="accorr", + epochs_average=False, ) - + # Shape should be (n_freq, n_epochs, 2*n_channels, 2*n_channels) assert con_no_avg.shape[1] == len(epochs.epo1) assert con_no_avg.shape[2] == 2 * n_channels assert con_no_avg.shape[3] == 2 * n_channels - + # Diagonal should have high values (self-correlation) diag_values = np.array([con[0, i, i] for i in range(con.shape[1])]) assert np.mean(diag_values) > 0.5 diff --git a/tests/test_sync.py b/tests/test_sync.py index 54d86f8..2e5ca5a 100644 --- a/tests/test_sync.py +++ b/tests/test_sync.py @@ -14,8 +14,12 @@ from hypyp.sync import get_metric from hypyp.sync.accorr import ACCorr from hypyp.sync.base import ( - BaseMetric, AUTO_PRIORITY, - NUMBA_AVAILABLE, TORCH_AVAILABLE, MPS_AVAILABLE, METAL_AVAILABLE, + BaseMetric, + AUTO_PRIORITY, + NUMBA_AVAILABLE, + TORCH_AVAILABLE, + MPS_AVAILABLE, + METAL_AVAILABLE, ) from hypyp.sync.kernels import CUPY_AVAILABLE from tests.accorr_reference import accorr_reference @@ -25,22 +29,30 @@ class TestAccorrReference: """Basic properties of the reference implementation.""" def test_reference_shape_no_average(self, complex_signal): - result = accorr_reference(complex_signal, epochs_average=False, show_progress=False) + result = accorr_reference( + complex_signal, epochs_average=False, show_progress=False + ) n_epochs, n_freq, n_ch, _ = complex_signal.shape assert result.shape == (n_freq, n_epochs, n_ch, n_ch) def test_reference_shape_with_average(self, complex_signal): - result = accorr_reference(complex_signal, epochs_average=True, show_progress=False) + result = accorr_reference( + complex_signal, epochs_average=True, show_progress=False + ) n_epochs, n_freq, n_ch, _ = complex_signal.shape assert result.shape == (n_freq, n_ch, n_ch) def test_reference_value_range(self, complex_signal): - result = accorr_reference(complex_signal, epochs_average=True, show_progress=False) + result = accorr_reference( + complex_signal, epochs_average=True, show_progress=False + ) assert np.all(result >= -1 - 1e-10) and np.all(result <= 1 + 1e-10) assert not np.any(np.isnan(result)) def test_reference_symmetry(self, complex_signal): - result = accorr_reference(complex_signal, epochs_average=True, show_progress=False) + result = accorr_reference( + complex_signal, epochs_average=True, show_progress=False + ) for freq_idx in range(result.shape[0]): matrix = result[freq_idx] np.testing.assert_allclose(matrix, matrix.T, rtol=1e-10, atol=1e-12) @@ -62,31 +74,43 @@ def _compute_with_class(self, complex_signal, optimization=None): return con.swapaxes(0, 1) def test_numpy_vs_reference(self, complex_signal): - result_reference = accorr_reference(complex_signal, epochs_average=False, show_progress=False) + result_reference = accorr_reference( + complex_signal, epochs_average=False, show_progress=False + ) result_numpy = self._compute_with_class(complex_signal, optimization=None) - np.testing.assert_allclose(result_numpy, result_reference, rtol=1e-9, atol=1e-10) + np.testing.assert_allclose( + result_numpy, result_reference, rtol=1e-9, atol=1e-10 + ) @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="Numba not available") def test_numba_vs_reference(self, complex_signal): - result_reference = accorr_reference(complex_signal, epochs_average=False, show_progress=False) - result_numba = self._compute_with_class(complex_signal, optimization='numba') - np.testing.assert_allclose(result_numba, result_reference, rtol=1e-9, atol=1e-10) + result_reference = accorr_reference( + complex_signal, epochs_average=False, show_progress=False + ) + result_numba = self._compute_with_class(complex_signal, optimization="numba") + np.testing.assert_allclose( + result_numba, result_reference, rtol=1e-9, atol=1e-10 + ) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_torch_vs_reference(self, complex_signal): - result_reference = accorr_reference(complex_signal, epochs_average=False, show_progress=False) - metric = ACCorr(optimization='torch', show_progress=False) + result_reference = accorr_reference( + complex_signal, epochs_average=False, show_progress=False + ) + metric = ACCorr(optimization="torch", show_progress=False) n_samp = complex_signal.shape[3] con = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) result_torch = con.swapaxes(0, 1) # MPS uses float32, so tolerance is lower - if metric._device == 'mps': - np.testing.assert_allclose(result_torch, result_reference, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric._device == "mps": + np.testing.assert_allclose( + result_torch, result_reference, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: - np.testing.assert_allclose(result_torch, result_reference, - rtol=1e-9, atol=1e-10) + np.testing.assert_allclose( + result_torch, result_reference, rtol=1e-9, atol=1e-10 + ) class TestAccorrViaComputeSync: @@ -96,26 +120,39 @@ class TestAccorrViaComputeSync: def test_compute_sync_default(self, complex_signal, complex_signal_raw): """compute_sync with optimization=None should match reference.""" - result_reference = accorr_reference(complex_signal, epochs_average=True, show_progress=False) - result = compute_sync(complex_signal_raw, 'accorr', optimization=None, epochs_average=True) + result_reference = accorr_reference( + complex_signal, epochs_average=True, show_progress=False + ) + result = compute_sync( + complex_signal_raw, "accorr", optimization=None, epochs_average=True + ) np.testing.assert_allclose(result, result_reference, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="Numba not available") def test_compute_sync_numba(self, complex_signal, complex_signal_raw): """compute_sync with optimization='numba' should match reference.""" - result_reference = accorr_reference(complex_signal, epochs_average=True, show_progress=False) - result = compute_sync(complex_signal_raw, 'accorr', optimization='numba', epochs_average=True) + result_reference = accorr_reference( + complex_signal, epochs_average=True, show_progress=False + ) + result = compute_sync( + complex_signal_raw, "accorr", optimization="numba", epochs_average=True + ) np.testing.assert_allclose(result, result_reference, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_compute_sync_torch(self, complex_signal, complex_signal_raw): """compute_sync with optimization='torch' should match reference.""" - result_reference = accorr_reference(complex_signal, epochs_average=True, show_progress=False) - result = compute_sync(complex_signal_raw, 'accorr', optimization='torch', epochs_average=True) + result_reference = accorr_reference( + complex_signal, epochs_average=True, show_progress=False + ) + result = compute_sync( + complex_signal_raw, "accorr", optimization="torch", epochs_average=True + ) # MPS uses float32, so a looser tolerance is required if MPS_AVAILABLE: - np.testing.assert_allclose(result, result_reference, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + np.testing.assert_allclose( + result, result_reference, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result, result_reference, rtol=1e-9, atol=1e-10) @@ -129,6 +166,7 @@ class TestPLV: def test_plv_shape(self, complex_signal): """PLV output shape should match input dimensions.""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] result = PLV().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) n_epochs, n_freq, n_ch, _ = complex_signal.shape @@ -137,6 +175,7 @@ def test_plv_shape(self, complex_signal): def test_plv_value_range(self, complex_signal): """PLV values should be in [0, 1].""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] result = PLV().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) assert np.all(result >= -1e-10) and np.all(result <= 1 + 1e-10) @@ -146,28 +185,38 @@ def test_plv_value_range(self, complex_signal): def test_plv_numba_vs_numpy(self, complex_signal): """Numba PLV should match numpy PLV exactly (both float64).""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] - result_np = PLV(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = PLV(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PLV(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = PLV(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_plv_torch_vs_numpy(self, complex_signal): """Torch PLV should match numpy PLV within MPS tolerance.""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] - result_np = PLV(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = PLV(optimization='torch') + result_np = PLV(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = PLV(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) def test_plv_symmetry(self, complex_signal): """PLV matrix should be symmetric (PLV(i,j) == PLV(j,i)).""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] result = PLV().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) for e in range(result.shape[0]): @@ -180,18 +229,28 @@ def test_plv_symmetry(self, complex_signal): def test_plv_metal_vs_numpy(self, complex_signal): """Metal PLV should match numpy PLV within float32 tolerance.""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] - result_np = PLV(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = PLV(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PLV(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = PLV(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_plv_cuda_vs_numpy(self, complex_signal): """CUDA PLV should match numpy PLV exactly (both float64).""" from hypyp.sync.plv import PLV + n_samp = complex_signal.shape[3] - result_np = PLV(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = PLV(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PLV(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = PLV(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -203,6 +262,7 @@ class TestCCorr: def test_ccorr_shape(self, complex_signal): """CCorr output shape should match input dimensions.""" from hypyp.sync.ccorr import CCorr + metric = CCorr() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -212,6 +272,7 @@ def test_ccorr_shape(self, complex_signal): def test_ccorr_value_range(self, complex_signal): """CCorr values should be non-negative (abs of correlation).""" from hypyp.sync.ccorr import CCorr + metric = CCorr() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -235,13 +296,17 @@ def test_ccorr_vs_scipy_reference(self, complex_signal): n_epoch, n_freq, n_ch_total, 1 ) angle_centered = np.sin(angle - mu_angle_scipy) - formula = 'nilm,nimk->nilk' + formula = "nilm,nimk->nilk" transpose_axes = self.TRANSPOSE_AXES result_scipy = np.abs( - np.einsum(formula, angle_centered, angle_centered.transpose(transpose_axes)) / - np.sqrt(np.einsum('nil,nik->nilk', - np.sum(angle_centered ** 2, axis=3), - np.sum(angle_centered ** 2, axis=3))) + np.einsum(formula, angle_centered, angle_centered.transpose(transpose_axes)) + / np.sqrt( + np.einsum( + "nil,nik->nilk", + np.sum(angle_centered**2, axis=3), + np.sum(angle_centered**2, axis=3), + ) + ) ) np.testing.assert_allclose(result_new, result_scipy, rtol=1e-12, atol=1e-14) @@ -249,6 +314,7 @@ def test_ccorr_vs_scipy_reference(self, complex_signal): def test_ccorr_symmetry(self, complex_signal): """CCorr matrix should be symmetric for each epoch/freq.""" from hypyp.sync.ccorr import CCorr + metric = CCorr() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -258,27 +324,34 @@ def test_ccorr_symmetry(self, complex_signal): result[e, f], result[e, f].T, rtol=1e-10, atol=1e-12 ) - @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="Numba not available") def test_ccorr_numba_vs_numpy(self, complex_signal): """Numba CCorr should match numpy CCorr exactly (both float64).""" from hypyp.sync.ccorr import CCorr + n_samp = complex_signal.shape[3] - result_np = CCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = CCorr(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = CCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = CCorr(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_ccorr_torch_vs_numpy(self, complex_signal): """Torch CCorr should match numpy CCorr within MPS tolerance.""" from hypyp.sync.ccorr import CCorr + n_samp = complex_signal.shape[3] - result_np = CCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = CCorr(optimization='torch') + result_np = CCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = CCorr(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) # Angle-free reformulation eliminates transcendental function chain, # bringing MPS precision in line with PLV. - if metric_torch._device == 'mps': + if metric_torch._device == "mps": np.testing.assert_allclose(result_torch, result_np, rtol=1e-5, atol=1e-5) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) @@ -290,18 +363,28 @@ def test_ccorr_metal_vs_numpy(self, complex_signal): Uses Kahan summation with fastMath=OFF to preserve IEEE-754 compliance. """ from hypyp.sync.ccorr import CCorr + n_samp = complex_signal.shape[3] - result_np = CCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = CCorr(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = CCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = CCorr(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_ccorr_cuda_vs_numpy(self, complex_signal): """CUDA CCorr should match numpy CCorr exactly (both float64).""" from hypyp.sync.ccorr import CCorr + n_samp = complex_signal.shape[3] - result_np = CCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = CCorr(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = CCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = CCorr(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -314,6 +397,7 @@ class TestCoh: def test_coh_shape(self, complex_signal): """Coh output shape should match input dimensions.""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] result = Coh().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) n_epochs, n_freq, n_ch, _ = complex_signal.shape @@ -322,6 +406,7 @@ def test_coh_shape(self, complex_signal): def test_coh_value_range(self, complex_signal): """Coh values should be in [0, 1].""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] result = Coh().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) assert np.all(result >= -1e-10) and np.all(result <= 1 + 1e-10) @@ -331,28 +416,38 @@ def test_coh_value_range(self, complex_signal): def test_coh_numba_vs_numpy(self, complex_signal): """Numba Coh should match numpy Coh exactly (both float64).""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] - result_np = Coh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = Coh(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = Coh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = Coh(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_coh_torch_vs_numpy(self, complex_signal): """Torch Coh should match numpy Coh within MPS tolerance.""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] - result_np = Coh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = Coh(optimization='torch') + result_np = Coh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = Coh(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) def test_coh_symmetry(self, complex_signal): """Coh matrix should be symmetric.""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] result = Coh().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) for e in range(result.shape[0]): @@ -365,18 +460,28 @@ def test_coh_symmetry(self, complex_signal): def test_coh_metal_vs_numpy(self, complex_signal): """Metal Coh should match numpy Coh within float32 tolerance.""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] - result_np = Coh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = Coh(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = Coh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = Coh(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_coh_cuda_vs_numpy(self, complex_signal): """CUDA Coh should match numpy Coh exactly (both float64).""" from hypyp.sync.coh import Coh + n_samp = complex_signal.shape[3] - result_np = Coh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = Coh(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = Coh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = Coh(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -389,6 +494,7 @@ class TestImCoh: def test_imcoh_shape(self, complex_signal): """ImCoh output shape should match input dimensions.""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] result = ImCoh().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) n_epochs, n_freq, n_ch, _ = complex_signal.shape @@ -397,6 +503,7 @@ def test_imcoh_shape(self, complex_signal): def test_imcoh_value_range(self, complex_signal): """ImCoh values should be in [0, 1].""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] result = ImCoh().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) assert np.all(result >= -1e-10) and np.all(result <= 1 + 1e-10) @@ -406,28 +513,38 @@ def test_imcoh_value_range(self, complex_signal): def test_imcoh_numba_vs_numpy(self, complex_signal): """Numba ImCoh should match numpy ImCoh exactly (both float64).""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] - result_np = ImCoh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = ImCoh(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = ImCoh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = ImCoh(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_imcoh_torch_vs_numpy(self, complex_signal): """Torch ImCoh should match numpy ImCoh within MPS tolerance.""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] - result_np = ImCoh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = ImCoh(optimization='torch') + result_np = ImCoh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = ImCoh(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) def test_imcoh_symmetry(self, complex_signal): """ImCoh matrix should be symmetric.""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] result = ImCoh().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) for e in range(result.shape[0]): @@ -440,18 +557,28 @@ def test_imcoh_symmetry(self, complex_signal): def test_imcoh_metal_vs_numpy(self, complex_signal): """Metal ImCoh should match numpy ImCoh within float32 tolerance.""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] - result_np = ImCoh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = ImCoh(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = ImCoh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = ImCoh(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_imcoh_cuda_vs_numpy(self, complex_signal): """CUDA ImCoh should match numpy ImCoh exactly (both float64).""" from hypyp.sync.imaginary_coh import ImCoh + n_samp = complex_signal.shape[3] - result_np = ImCoh(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = ImCoh(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = ImCoh(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = ImCoh(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -464,6 +591,7 @@ class TestEnvCorr: def test_envcorr_shape(self, complex_signal): """EnvCorr output shape should match input dimensions.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] result = EnvCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) n_epochs, n_freq, n_ch, _ = complex_signal.shape @@ -472,6 +600,7 @@ def test_envcorr_shape(self, complex_signal): def test_envcorr_value_range(self, complex_signal): """EnvCorr values should be in [-1, 1].""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] result = EnvCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) assert np.all(result >= -1 - 1e-10) and np.all(result <= 1 + 1e-10) @@ -481,28 +610,38 @@ def test_envcorr_value_range(self, complex_signal): def test_envcorr_numba_vs_numpy(self, complex_signal): """Numba EnvCorr should match numpy exactly (both float64).""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] - result_np = EnvCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = EnvCorr(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = EnvCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = EnvCorr(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_envcorr_torch_vs_numpy(self, complex_signal): """Torch EnvCorr should match numpy within MPS tolerance.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] - result_np = EnvCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = EnvCorr(optimization='torch') + result_np = EnvCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = EnvCorr(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) def test_envcorr_symmetry(self, complex_signal): """EnvCorr matrix should be symmetric.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] result = EnvCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) for e in range(result.shape[0]): @@ -524,6 +663,7 @@ class TestPLI: def test_pli_shape(self, complex_signal): """PLI output shape should match input dimensions.""" from hypyp.sync.pli import PLI + metric = PLI() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -533,6 +673,7 @@ def test_pli_shape(self, complex_signal): def test_pli_value_range(self, complex_signal): """PLI values should be in [0, 1].""" from hypyp.sync.pli import PLI + metric = PLI() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -543,24 +684,33 @@ def test_pli_value_range(self, complex_signal): def test_pli_numba_vs_numpy(self, complex_signal): """Numba PLI should match numpy PLI exactly (both float64).""" from hypyp.sync.pli import PLI + n_samp = complex_signal.shape[3] - result_np = PLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = PLI(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = PLI(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_pli_torch_vs_numpy(self, complex_signal): """Torch PLI should match numpy PLI.""" from hypyp.sync.pli import PLI + n_samp = complex_signal.shape[3] - result_np = PLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = PLI(optimization='torch') + result_np = PLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = PLI(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) @@ -568,12 +718,14 @@ def test_pli_torch_vs_numpy(self, complex_signal): def test_pli_torch_shape(self, complex_signal): """Torch PLI output shape should match numpy.""" from hypyp.sync.pli import PLI + n_samp = complex_signal.shape[3] - result = PLI(optimization='torch').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result = PLI(optimization="torch").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) n_epochs, n_freq, n_ch, _ = complex_signal.shape assert result.shape == (n_epochs, n_freq, n_ch, n_ch) - @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_pli_torch_large_channels(self): """PLI torch should handle 128ch/subject (256 total) without MPS INT_MAX crash.""" @@ -581,12 +733,15 @@ def test_pli_torch_large_channels(self): rng = np.random.default_rng(42) n_ch_per_subject = 128 - sig = rng.standard_normal((2, 1, 2 * n_ch_per_subject, 256)) + \ - 1j * rng.standard_normal((2, 1, 2 * n_ch_per_subject, 256)) + sig = rng.standard_normal( + (2, 1, 2 * n_ch_per_subject, 256) + ) + 1j * rng.standard_normal((2, 1, 2 * n_ch_per_subject, 256)) n_samp = sig.shape[3] result_np = PLI().compute(sig, n_samp, self.TRANSPOSE_AXES) - result_torch = PLI(optimization='torch').compute(sig, n_samp, self.TRANSPOSE_AXES) + result_torch = PLI(optimization="torch").compute( + sig, n_samp, self.TRANSPOSE_AXES + ) assert result_torch.shape == result_np.shape assert not np.any(np.isnan(result_torch)) @@ -595,9 +750,14 @@ def test_pli_torch_large_channels(self): def test_pli_metal_vs_numpy(self, complex_signal): """Metal PLI should match numpy PLI within float32 tolerance.""" from hypyp.sync.pli import PLI + n_samp = complex_signal.shape[3] - result_np = PLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = PLI(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = PLI(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) # Float32 precision — sign() near zero can flip np.testing.assert_allclose(result_metal, result_np, rtol=1e-2, atol=1e-2) @@ -605,10 +765,13 @@ def test_pli_metal_vs_numpy(self, complex_signal): def test_pli_metal_large_channels(self): """Metal PLI should handle 128ch/subject (256 total).""" from hypyp.sync.pli import PLI + rng = np.random.default_rng(42) - sig = rng.standard_normal((2, 1, 256, 256)) + 1j * rng.standard_normal((2, 1, 256, 256)) + sig = rng.standard_normal((2, 1, 256, 256)) + 1j * rng.standard_normal( + (2, 1, 256, 256) + ) n_samp = sig.shape[3] - result = PLI(optimization='metal').compute(sig, n_samp, self.TRANSPOSE_AXES) + result = PLI(optimization="metal").compute(sig, n_samp, self.TRANSPOSE_AXES) assert result.shape == (2, 1, 256, 256) assert not np.any(np.isnan(result)) assert np.allclose(np.diagonal(result[0, 0]), 0) # diagonal = 0 @@ -617,9 +780,14 @@ def test_pli_metal_large_channels(self): def test_pli_cuda_vs_numpy(self, complex_signal): """CUDA PLI should match numpy PLI exactly (both float64).""" from hypyp.sync.pli import PLI + n_samp = complex_signal.shape[3] - result_np = PLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = PLI(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = PLI(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) # Float64: should match to machine precision np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -634,6 +802,7 @@ class TestWPLI: def test_wpli_shape(self, complex_signal): """wPLI output shape should match input dimensions.""" from hypyp.sync.wpli import WPLI + metric = WPLI() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -643,6 +812,7 @@ def test_wpli_shape(self, complex_signal): def test_wpli_value_range(self, complex_signal): """wPLI values should be in [0, 1].""" from hypyp.sync.wpli import WPLI + metric = WPLI() n_samp = complex_signal.shape[3] result = metric.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) @@ -654,24 +824,33 @@ def test_wpli_value_range(self, complex_signal): def test_wpli_numba_vs_numpy(self, complex_signal): """Numba wPLI should match numpy wPLI exactly (both float64).""" from hypyp.sync.wpli import WPLI + n_samp = complex_signal.shape[3] - result_np = WPLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = WPLI(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = WPLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = WPLI(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_wpli_torch_vs_numpy(self, complex_signal): """Torch wPLI should match numpy wPLI.""" from hypyp.sync.wpli import WPLI + n_samp = complex_signal.shape[3] - result_np = WPLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = WPLI(optimization='torch') + result_np = WPLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = WPLI(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) @@ -679,12 +858,14 @@ def test_wpli_torch_vs_numpy(self, complex_signal): def test_wpli_torch_shape(self, complex_signal): """Torch wPLI output shape should match numpy.""" from hypyp.sync.wpli import WPLI + n_samp = complex_signal.shape[3] - result = WPLI(optimization='torch').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result = WPLI(optimization="torch").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) n_epochs, n_freq, n_ch, _ = complex_signal.shape assert result.shape == (n_epochs, n_freq, n_ch, n_ch) - @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_wpli_torch_large_channels(self): """wPLI torch should handle 128ch/subject (256 total) without MPS INT_MAX crash.""" @@ -692,12 +873,15 @@ def test_wpli_torch_large_channels(self): rng = np.random.default_rng(42) n_ch_per_subject = 128 - sig = rng.standard_normal((2, 1, 2 * n_ch_per_subject, 256)) + \ - 1j * rng.standard_normal((2, 1, 2 * n_ch_per_subject, 256)) + sig = rng.standard_normal( + (2, 1, 2 * n_ch_per_subject, 256) + ) + 1j * rng.standard_normal((2, 1, 2 * n_ch_per_subject, 256)) n_samp = sig.shape[3] result_np = WPLI().compute(sig, n_samp, self.TRANSPOSE_AXES) - result_torch = WPLI(optimization='torch').compute(sig, n_samp, self.TRANSPOSE_AXES) + result_torch = WPLI(optimization="torch").compute( + sig, n_samp, self.TRANSPOSE_AXES + ) assert result_torch.shape == result_np.shape assert not np.any(np.isnan(result_torch)) @@ -712,6 +896,7 @@ class TestEnvCorr: def test_envcorr_shape(self, complex_signal): """EnvCorr output shape should match input dimensions.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] result = EnvCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) n_epochs, n_freq, n_ch, _ = complex_signal.shape @@ -720,6 +905,7 @@ def test_envcorr_shape(self, complex_signal): def test_envcorr_value_range(self, complex_signal): """EnvCorr values should be in [-1, 1] (Pearson correlation).""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] result = EnvCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) assert np.all(result >= -1 - 1e-10) and np.all(result <= 1 + 1e-10) @@ -728,6 +914,7 @@ def test_envcorr_value_range(self, complex_signal): def test_envcorr_symmetry(self, complex_signal): """EnvCorr matrix should be symmetric.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] result = EnvCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) for e in range(result.shape[0]): @@ -740,22 +927,31 @@ def test_envcorr_symmetry(self, complex_signal): def test_envcorr_numba_vs_numpy(self, complex_signal): """Numba EnvCorr should match numpy EnvCorr exactly (both float64).""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] - result_np = EnvCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = EnvCorr(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = EnvCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = EnvCorr(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_envcorr_torch_vs_numpy(self, complex_signal): """Torch EnvCorr should match numpy EnvCorr within MPS tolerance.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] - result_np = EnvCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = EnvCorr(optimization='torch') + result_np = EnvCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = EnvCorr(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) @@ -763,18 +959,28 @@ def test_envcorr_torch_vs_numpy(self, complex_signal): def test_envcorr_metal_vs_numpy(self, complex_signal): """Metal EnvCorr should match numpy EnvCorr within float32 tolerance.""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] - result_np = EnvCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = EnvCorr(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = EnvCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = EnvCorr(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_envcorr_cuda_vs_numpy(self, complex_signal): """CUDA EnvCorr should match numpy EnvCorr exactly (both float64).""" from hypyp.sync.envelope_corr import EnvCorr + n_samp = complex_signal.shape[3] - result_np = EnvCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = EnvCorr(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = EnvCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = EnvCorr(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -787,6 +993,7 @@ class TestPowCorr: def test_powcorr_shape(self, complex_signal): """PowCorr output shape should match input dimensions.""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] result = PowCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) n_epochs, n_freq, n_ch, _ = complex_signal.shape @@ -795,6 +1002,7 @@ def test_powcorr_shape(self, complex_signal): def test_powcorr_value_range(self, complex_signal): """PowCorr values should be in [-1, 1] (Pearson correlation).""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] result = PowCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) assert np.all(result >= -1 - 1e-10) and np.all(result <= 1 + 1e-10) @@ -803,6 +1011,7 @@ def test_powcorr_value_range(self, complex_signal): def test_powcorr_symmetry(self, complex_signal): """PowCorr matrix should be symmetric.""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] result = PowCorr().compute(complex_signal, n_samp, self.TRANSPOSE_AXES) for e in range(result.shape[0]): @@ -815,60 +1024,88 @@ def test_powcorr_symmetry(self, complex_signal): def test_powcorr_numba_vs_numpy(self, complex_signal): """Numba PowCorr should match numpy PowCorr exactly (both float64).""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] - result_np = PowCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_numba = PowCorr(optimization='numba').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PowCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_numba = PowCorr(optimization="numba").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_numba, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not TORCH_AVAILABLE, reason="Torch not available") def test_powcorr_torch_vs_numpy(self, complex_signal): """Torch PowCorr should match numpy PowCorr within MPS tolerance.""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] - result_np = PowCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - metric_torch = PowCorr(optimization='torch') + result_np = PowCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + metric_torch = PowCorr(optimization="torch") result_torch = metric_torch.compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - if metric_torch._device == 'mps': - np.testing.assert_allclose(result_torch, result_np, - rtol=self.MPS_TOL, atol=self.MPS_TOL) + if metric_torch._device == "mps": + np.testing.assert_allclose( + result_torch, result_np, rtol=self.MPS_TOL, atol=self.MPS_TOL + ) else: np.testing.assert_allclose(result_torch, result_np, rtol=1e-9, atol=1e-10) - @pytest.mark.skipif(not METAL_AVAILABLE, reason="Metal not available") def test_powcorr_metal_vs_numpy(self, complex_signal): """Metal PowCorr should match numpy PowCorr within float32 tolerance.""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] - result_np = PowCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = PowCorr(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PowCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = PowCorr(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_powcorr_cuda_vs_numpy(self, complex_signal): """CUDA PowCorr should match numpy PowCorr exactly (both float64).""" from hypyp.sync.pow_corr import PowCorr + n_samp = complex_signal.shape[3] - result_np = PowCorr(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = PowCorr(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = PowCorr(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = PowCorr(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @pytest.mark.skipif(not METAL_AVAILABLE, reason="Metal not available") def test_wpli_metal_vs_numpy(self, complex_signal): """Metal wPLI should match numpy wPLI within float32 tolerance.""" from hypyp.sync.wpli import WPLI + n_samp = complex_signal.shape[3] - result_np = WPLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = WPLI(optimization='metal').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = WPLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = WPLI(optimization="metal").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-2, atol=1e-2) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_wpli_cuda_vs_numpy(self, complex_signal): """CUDA wPLI should match numpy wPLI exactly (both float64).""" from hypyp.sync.wpli import WPLI + n_samp = complex_signal.shape[3] - result_np = WPLI(optimization=None).compute(complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = WPLI(optimization='cuda_kernel').compute(complex_signal, n_samp, self.TRANSPOSE_AXES) + result_np = WPLI(optimization=None).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = WPLI(optimization="cuda_kernel").compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -881,22 +1118,28 @@ class TestAccorrKernels: def test_accorr_metal_vs_numpy(self, complex_signal): """Metal ACCorr should match numpy within float32 tolerance.""" from hypyp.sync.accorr import ACCorr + n_samp = complex_signal.shape[3] result_np = ACCorr(optimization=None, show_progress=False).compute( - complex_signal, n_samp, self.TRANSPOSE_AXES) - result_metal = ACCorr(optimization='metal', show_progress=False).compute( - complex_signal, n_samp, self.TRANSPOSE_AXES) + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_metal = ACCorr(optimization="metal", show_progress=False).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_metal, result_np, rtol=1e-5, atol=1e-5) @pytest.mark.skipif(not CUPY_AVAILABLE, reason="CuPy not available") def test_accorr_cuda_vs_numpy(self, complex_signal): """CUDA ACCorr should match numpy exactly (both float64).""" from hypyp.sync.accorr import ACCorr + n_samp = complex_signal.shape[3] result_np = ACCorr(optimization=None, show_progress=False).compute( - complex_signal, n_samp, self.TRANSPOSE_AXES) - result_cuda = ACCorr(optimization='cuda_kernel', show_progress=False).compute( - complex_signal, n_samp, self.TRANSPOSE_AXES) + complex_signal, n_samp, self.TRANSPOSE_AXES + ) + result_cuda = ACCorr(optimization="cuda_kernel", show_progress=False).compute( + complex_signal, n_samp, self.TRANSPOSE_AXES + ) np.testing.assert_allclose(result_cuda, result_np, rtol=1e-9, atol=1e-10) @@ -909,22 +1152,22 @@ def test_invalid_optimization(self): def test_numba_fallback_warning(self): """When numba is unavailable, optimization='numba' warns and falls back to numpy.""" - with patch('hypyp.sync.base.NUMBA_AVAILABLE', False): + with patch("hypyp.sync.base.NUMBA_AVAILABLE", False): with pytest.warns(UserWarning, match="numba not installed"): - metric = ACCorr(optimization='numba') - assert metric._backend == 'numpy' + metric = ACCorr(optimization="numba") + assert metric._backend == "numpy" def test_torch_fallback_warning(self): """When torch is unavailable, optimization='torch' warns and falls back to numpy.""" - with patch('hypyp.sync.base.TORCH_AVAILABLE', False): + with patch("hypyp.sync.base.TORCH_AVAILABLE", False): with pytest.warns(UserWarning, match="torch not installed"): - metric = ACCorr(optimization='torch') - assert metric._backend == 'numpy' + metric = ACCorr(optimization="torch") + assert metric._backend == "numpy" def test_auto_resolves(self): """optimization='auto' should resolve without error.""" - metric = ACCorr(optimization='auto') - assert metric._backend in ('numpy', 'numba', 'torch', 'metal', 'cuda_kernel') + metric = ACCorr(optimization="auto") + assert metric._backend in ("numpy", "numba", "torch", "metal", "cuda_kernel") class TestAutoDispatch: @@ -933,15 +1176,15 @@ class TestAutoDispatch: def test_auto_all_metrics_resolve(self): """optimization='auto' resolves for every metric without error.""" for metric_name in AUTO_PRIORITY: - m = get_metric(metric_name, optimization='auto') - assert m._backend in ('numpy', 'numba', 'torch', 'metal', 'cuda_kernel') + m = get_metric(metric_name, optimization="auto") + assert m._backend in ("numpy", "numba", "torch", "metal", "cuda_kernel") @pytest.mark.skipif(not MPS_AVAILABLE, reason="MPS not available") def test_auto_einsum_prefers_torch_on_mps(self): """Einsum metrics should prefer torch on Apple Silicon.""" - for metric_name in ['plv', 'ccorr', 'coh', 'imcoh', 'envcorr', 'powcorr']: - m = get_metric(metric_name, optimization='auto') - assert m._backend == 'torch' and m._device == 'mps', ( + for metric_name in ["plv", "ccorr", "coh", "imcoh", "envcorr", "powcorr"]: + m = get_metric(metric_name, optimization="auto") + assert m._backend == "torch" and m._device == "mps", ( f"{metric_name} auto: expected torch/mps, got {m._backend}/{m._device}" ) @@ -949,42 +1192,48 @@ def test_auto_einsum_prefers_torch_on_mps(self): @pytest.mark.skipif(not METAL_AVAILABLE, reason="Metal not available") def test_auto_sign_prefers_metal_on_mps(self): """PLI/wPLI/ACCorr should prefer Metal on Apple Silicon.""" - for metric_name in ['pli', 'wpli', 'accorr']: - m = get_metric(metric_name, optimization='auto') - assert m._backend == 'metal', ( + for metric_name in ["pli", "wpli", "accorr"]: + m = get_metric(metric_name, optimization="auto") + assert m._backend == "metal", ( f"{metric_name} auto: expected metal, got {m._backend}" ) def test_auto_priority_override(self): """Custom priority overrides the AUTO_PRIORITY table.""" - m = get_metric('plv', optimization='auto', priority=['numba']) + m = get_metric("plv", optimization="auto", priority=["numba"]) if NUMBA_AVAILABLE: - assert m._backend == 'numba' + assert m._backend == "numba" else: - assert m._backend == 'numpy' + assert m._backend == "numpy" def test_auto_priority_skips_unavailable(self): """Priority list gracefully skips unavailable backends.""" - with patch('hypyp.sync.base.METAL_AVAILABLE', False), \ - patch('hypyp.sync.base.CUPY_AVAILABLE', False): - m = get_metric('pli', optimization='auto', priority=['metal', 'cuda_kernel', 'numba']) + with ( + patch("hypyp.sync.base.METAL_AVAILABLE", False), + patch("hypyp.sync.base.CUPY_AVAILABLE", False), + ): + m = get_metric( + "pli", optimization="auto", priority=["metal", "cuda_kernel", "numba"] + ) if NUMBA_AVAILABLE: - assert m._backend == 'numba' + assert m._backend == "numba" else: - assert m._backend == 'numpy' + assert m._backend == "numpy" def test_auto_fallback_cpu_only(self): """On CPU-only machines, auto warns and falls back to numba or numpy.""" - with patch('hypyp.sync.base.MPS_AVAILABLE', False), \ - patch('hypyp.sync.base.CUDA_AVAILABLE', False): + with ( + patch("hypyp.sync.base.MPS_AVAILABLE", False), + patch("hypyp.sync.base.CUDA_AVAILABLE", False), + ): with pytest.warns(UserWarning, match="No GPU available"): - m = get_metric('plv', optimization='auto') + m = get_metric("plv", optimization="auto") if NUMBA_AVAILABLE: - assert m._backend == 'numba' + assert m._backend == "numba" else: - assert m._backend == 'numpy' + assert m._backend == "numpy" def test_priority_parameter_propagated_via_get_metric(self): """get_metric passes priority through to the metric class.""" - m = get_metric('accorr', optimization='auto', priority=['numba']) - assert m._priority == ['numba'] + m = get_metric("accorr", optimization="auto", priority=["numba"]) + assert m._priority == ["numba"] diff --git a/tests/test_utils.py b/tests/test_utils.py index 9b46aab..9b40b7a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -8,42 +8,48 @@ from hypyp.signal import SyntheticSignal from hypyp.wavelet.implementations.pywavelets_wavelet import ComplexMorletWavelet + def get_fake_raw(): sfreq = 100 duration = 10 n_channels = 5 n_samples = int(sfreq * duration) data = np.random.randn(n_channels, n_samples) - ch_names = [f'Foo {i}' for i in range(n_channels)] - ch_types = ['eeg'] * n_channels + ch_names = [f"Foo {i}" for i in range(n_channels)] + ch_types = ["eeg"] * n_channels info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types) raw = mne.io.RawArray(data, info) return raw + def test_task_with_onset_event_id(): - task = utils.Task('my_task', onset_event_id=1, offset_event_id=1) - assert task.name == 'my_task' + task = utils.Task("my_task", onset_event_id=1, offset_event_id=1) + assert task.name == "my_task" assert task.is_event_based == True assert task.is_time_based == False + def test_task_with_onset_time(): - task = utils.Task('my_task', onset_time=1, duration=10) - assert task.name == 'my_task' + task = utils.Task("my_task", onset_time=1, duration=10) + assert task.name == "my_task" assert task.is_event_based == False assert task.is_time_based == True + def test_epochs_per_task(): raw = get_fake_raw() # Define events using annotations (e.g., at 2s, 5s, and 8s with a duration of 0s) onsets = [2, 5, 8, 9] durations = [0, 0, 0, 0] - descriptions = np.arange(1, len(onsets)+1) - raw.set_annotations(mne.Annotations(onset=onsets, duration=durations, description=descriptions)) + descriptions = np.arange(1, len(onsets) + 1) + raw.set_annotations( + mne.Annotations(onset=onsets, duration=durations, description=descriptions) + ) tasks = [ - utils.Task('task1', onset_event_id=1, offset_event_id=2), - utils.Task('rest', onset_event_id=2, offset_event_id=3), - utils.Task('task2', onset_event_id=3, offset_event_id=4), + utils.Task("task1", onset_event_id=1, offset_event_id=2), + utils.Task("rest", onset_event_id=2, offset_event_id=3), + utils.Task("task2", onset_event_id=3, offset_event_id=4), ] all_epochs = utils.epochs_from_tasks(raw, tasks) @@ -51,30 +57,36 @@ def test_epochs_per_task(): assert all_epochs[0].tmax == pytest.approx(3) assert all_epochs[2].tmax == pytest.approx(1) + def test_epochs_recurring_task(): raw = get_fake_raw() onsets = [2, 5, 6, 9] durations = [0, 0, 0, 0] descriptions = [1, 2, 1, 2] - raw.set_annotations(mne.Annotations(onset=onsets, duration=durations, description=descriptions)) + raw.set_annotations( + mne.Annotations(onset=onsets, duration=durations, description=descriptions) + ) tasks = [ - utils.Task('task1', onset_event_id=1, offset_event_id=2), + utils.Task("task1", onset_event_id=1, offset_event_id=2), ] all_epochs = utils.epochs_from_tasks(raw, tasks) assert len(all_epochs) == 1 assert len(all_epochs[0]) == 2 + def test_epochs_recurring_task_crop_time(): raw = get_fake_raw() - onsets = [2, 5, 6, 7] # Have the 2nd task occurence shorter than first one + onsets = [2, 5, 6, 7] # Have the 2nd task occurence shorter than first one durations = [0, 0, 0, 0] descriptions = [1, 2, 1, 2] - raw.set_annotations(mne.Annotations(onset=onsets, duration=durations, description=descriptions)) + raw.set_annotations( + mne.Annotations(onset=onsets, duration=durations, description=descriptions) + ) tasks = [ - utils.Task('task1', onset_event_id=1, offset_event_id=2), + utils.Task("task1", onset_event_id=1, offset_event_id=2), ] all_epochs = utils.epochs_from_tasks(raw, tasks) @@ -83,25 +95,53 @@ def test_epochs_recurring_task_crop_time(): assert len(epochs.times) >= 100 assert len(epochs.times) < 200 + def test_task_start_end_combinaison(): raw = get_fake_raw() # Define events using annotations (e.g., at 2s, 5s, and 8s with a duration of 0s) onsets = [2, 5, 8, 9] durations = [0, 0, 0, 0] - descriptions = np.arange(1, len(onsets)+1) - raw.set_annotations(mne.Annotations(onset=onsets, duration=durations, description=descriptions)) - - assert utils.epochs_from_tasks(raw, [utils.Task('task1', onset_event_id=1, offset_event_id=utils.TASK_NEXT_EVENT)])[0].tmax == pytest.approx(3) - assert utils.epochs_from_tasks(raw, [utils.Task('task1', onset_event_id=1, offset_event_id=utils.TASK_END)])[0].tmax == pytest.approx(8, abs=0.01) - assert utils.epochs_from_tasks(raw, [utils.Task('task1', onset_event_id=utils.TASK_BEGINNING, offset_event_id=utils.TASK_NEXT_EVENT)])[0].tmax == pytest.approx(2, abs=0.01) - assert utils.epochs_from_tasks(raw, [utils.Task('task1', onset_event_id=utils.TASK_BEGINNING, offset_event_id=utils.TASK_END)])[0].tmax == pytest.approx(10, abs=0.01) + descriptions = np.arange(1, len(onsets) + 1) + raw.set_annotations( + mne.Annotations(onset=onsets, duration=durations, description=descriptions) + ) + + assert utils.epochs_from_tasks( + raw, + [utils.Task("task1", onset_event_id=1, offset_event_id=utils.TASK_NEXT_EVENT)], + )[0].tmax == pytest.approx(3) + assert utils.epochs_from_tasks( + raw, [utils.Task("task1", onset_event_id=1, offset_event_id=utils.TASK_END)] + )[0].tmax == pytest.approx(8, abs=0.01) + assert utils.epochs_from_tasks( + raw, + [ + utils.Task( + "task1", + onset_event_id=utils.TASK_BEGINNING, + offset_event_id=utils.TASK_NEXT_EVENT, + ) + ], + )[0].tmax == pytest.approx(2, abs=0.01) + assert utils.epochs_from_tasks( + raw, + [ + utils.Task( + "task1", + onset_event_id=utils.TASK_BEGINNING, + offset_event_id=utils.TASK_END, + ) + ], + )[0].tmax == pytest.approx(10, abs=0.01) def test_task_from_time_range(): raw = get_fake_raw() tasks = [ - utils.Task('task1', onset_time=0, duration=1), # task from start to 1 second - utils.Task('task2', onset_time=3, duration=2), # task from 3 seconds to 5 seconds + utils.Task("task1", onset_time=0, duration=1), # task from start to 1 second + utils.Task( + "task2", onset_time=3, duration=2 + ), # task from 3 seconds to 5 seconds ] all_epochs = utils.epochs_from_tasks(raw, tasks) @@ -109,11 +149,14 @@ def test_task_from_time_range(): duration = all_epochs[1].times[-1] - all_epochs[1].times[0] assert duration == tasks[1].duration + def test_task_from_time_range_recurring(): raw = get_fake_raw() tasks = [ - utils.Task('task1', onset_time=0, duration=1), # task from start to 1 second - utils.Task('task1', onset_time=3, duration=2), # task from 3 seconds to 5 seconds + utils.Task("task1", onset_time=0, duration=1), # task from start to 1 second + utils.Task( + "task1", onset_time=3, duration=2 + ), # task from 3 seconds to 5 seconds ] all_epochs = utils.epochs_from_tasks(raw, tasks) @@ -123,17 +166,19 @@ def test_task_from_time_range_recurring(): duration = all_epochs[0].times[-1] - all_epochs[0].times[0] assert duration == tasks[0].duration -#def test_task_onset_id_with_offset_duration(): + # def test_task_onset_id_with_offset_duration(): raw = get_fake_raw() # Define events using annotations (e.g., at 2s, 5s, and 8s with a duration of 0s) onsets = [2, 5, 8, 9] durations = [0, 0, 0, 0] - descriptions = np.arange(1, len(onsets)+1) - raw.set_annotations(mne.Annotations(onset=onsets, duration=durations, description=descriptions)) + descriptions = np.arange(1, len(onsets) + 1) + raw.set_annotations( + mne.Annotations(onset=onsets, duration=durations, description=descriptions) + ) tasks = [ - utils.Task('task1', onset_event_id=1, duration=2), - utils.Task('rest', onset_event_id=2, duration=1), + utils.Task("task1", onset_event_id=1, duration=2), + utils.Task("rest", onset_event_id=2, duration=1), ] all_epochs = utils.epochs_from_tasks(raw, tasks) @@ -141,39 +186,43 @@ def test_task_from_time_range_recurring(): assert all_epochs[0].tmax == pytest.approx(2) assert all_epochs[1].tmax == pytest.approx(1) + def test_mutually_exclusive_task_arguments(): # no start with pytest.raises(Exception): - utils.Task('task1') + utils.Task("task1") # both start with pytest.raises(Exception): - utils.Task('task1', onset_event_id=1, onset_time=1, duration=1) + utils.Task("task1", onset_event_id=1, onset_time=1, duration=1) # no end with pytest.raises(Exception): - utils.Task('task1', onset_event_id=1) + utils.Task("task1", onset_event_id=1) # both end with pytest.raises(Exception): - utils.Task('task1', onset_time=1, offset_event_id=2, duration=2) + utils.Task("task1", onset_time=1, offset_event_id=2, duration=2) # invalid combo with pytest.raises(Exception): - utils.Task('task1', onset_time=1, offset_event_id=1) + utils.Task("task1", onset_time=1, offset_event_id=1) + def test_task_with_events_as_float_str(): raw = get_fake_raw() onsets = [2, 5, 8, 9] durations = [0, 0, 0, 0] - descriptions = [f'{x}.0' for x in np.arange(1, len(onsets)+1)] - annotations = mne.Annotations(onset=onsets, duration=durations, description=descriptions) + descriptions = [f"{x}.0" for x in np.arange(1, len(onsets) + 1)] + annotations = mne.Annotations( + onset=onsets, duration=durations, description=descriptions + ) raw.set_annotations(annotations) tasks = [ - utils.Task('task1', onset_event_id=1, offset_event_id=2), - utils.Task('rest', onset_event_id=2, offset_event_id=3), - utils.Task('task2', onset_event_id=3, offset_event_id=4), + utils.Task("task1", onset_event_id=1, offset_event_id=2), + utils.Task("rest", onset_event_id=2, offset_event_id=3), + utils.Task("task2", onset_event_id=3, offset_event_id=4), ] all_epochs = utils.epochs_from_tasks(raw, tasks) @@ -181,19 +230,27 @@ def test_task_with_events_as_float_str(): assert all_epochs[0].tmax == pytest.approx(3) assert all_epochs[2].tmax == pytest.approx(1) + def test_task_enforce_id_as_int(): raw = get_fake_raw() onsets = [2, 5, 8, 9] durations = [0, 0, 0, 0] - descriptions = [f'{x}' for x in np.arange(1, len(onsets)+1)] - annotations = mne.Annotations(onset=onsets, duration=durations, description=descriptions) + descriptions = [f"{x}" for x in np.arange(1, len(onsets) + 1)] + annotations = mne.Annotations( + onset=onsets, duration=durations, description=descriptions + ) raw.set_annotations(annotations) with pytest.raises(ValueError): - utils.epochs_from_tasks(raw, [utils.Task('task1', onset_event_id='1', offset_event_id=2)]) + utils.epochs_from_tasks( + raw, [utils.Task("task1", onset_event_id="1", offset_event_id=2)] + ) with pytest.raises(ValueError): - utils.epochs_from_tasks(raw, [utils.Task('task1', onset_event_id=1, offset_event_id='2')]) + utils.epochs_from_tasks( + raw, [utils.Task("task1", onset_event_id=1, offset_event_id="2")] + ) + def test_downsampling(): wavelet = ComplexMorletWavelet() @@ -204,13 +261,16 @@ def test_downsampling(): assert res.W.shape[1] == len(signal.y) bins = 1000 - times, coif, W, factor = utils.downsample_in_time(res.times, res.coif, res.W, bins=bins) + times, coif, W, factor = utils.downsample_in_time( + res.times, res.coif, res.W, bins=bins + ) assert factor == len(signal.y) // bins assert len(times) == bins assert len(coif) == bins - assert W.shape[0] == res.W.shape[0] # no change + assert W.shape[0] == res.W.shape[0] # no change assert W.shape[1] == bins + def test_downsampling_low_values(): wavelet = ComplexMorletWavelet() signal = SyntheticSignal(duration=10, n_points=50).add_sin(1) @@ -222,48 +282,51 @@ def test_downsampling_low_values(): assert len(coif) == len(res.coif) assert W.shape[0] == res.W.shape[0] assert W.shape[1] == res.W.shape[1] - + + def test_downsampling_threshold(): bins = 1000 wavelet = ComplexMorletWavelet() - signal = SyntheticSignal(duration=10, n_points=bins+1).add_sin(1) + signal = SyntheticSignal(duration=10, n_points=bins + 1).add_sin(1) res = wavelet.cwt(signal.y, signal.period) - times, coif, W, factor = utils.downsample_in_time(res.times, res.coif, res.W, bins=bins) + times, coif, W, factor = utils.downsample_in_time( + res.times, res.coif, res.W, bins=bins + ) assert factor == 2 expected_len = bins / 2 + 1 assert len(times) == expected_len assert len(coif) == expected_len assert W.shape[1] == expected_len - + + def test_random_label(): assert len(utils.generate_random_label(10)) == 10 + def test_key_dict_as_data_frame(): my_dict = { - 'foo': { - 'col1': '1', - 'col2': '2', + "foo": { + "col1": "1", + "col2": "2", }, - 'bar': { - 'col1': '1', - 'col2': '2', + "bar": { + "col1": "1", + "col2": "2", }, } df = utils.key_dict_as_data_frame(my_dict) - assert df.columns[0] == 'key' - assert df.columns[1] == 'col1' - assert df.columns[2] == 'col2' - assert df['key'][0] == 'foo' - assert df['col1'][0] == '1' + assert df.columns[0] == "key" + assert df.columns[1] == "col1" + assert df.columns[2] == "col2" + assert df["key"][0] == "foo" + assert df["col1"][0] == "1" # Using column names - df = utils.key_dict_as_data_frame(my_dict, col_names=['a', 'b', 'c']) - assert df.columns[0] == 'a' - assert df.columns[1] == 'b' - assert df.columns[2] == 'c' - assert df['a'][0] == 'foo' - assert df['b'][0] == '1' - - \ No newline at end of file + df = utils.key_dict_as_data_frame(my_dict, col_names=["a", "b", "c"]) + assert df.columns[0] == "a" + assert df.columns[1] == "b" + assert df.columns[2] == "c" + assert df["a"][0] == "foo" + assert df["b"][0] == "1" diff --git a/tests/test_wavelet.py b/tests/test_wavelet.py index 0a72083..95537c6 100644 --- a/tests/test_wavelet.py +++ b/tests/test_wavelet.py @@ -5,36 +5,44 @@ from hypyp.signal import SyntheticSignal from hypyp.wavelet.pair_signals import PairSignals from hypyp.wavelet.base_wavelet import BaseWavelet -from hypyp.wavelet.implementations.pywavelets_wavelet import ComplexGaussianWavelet, ComplexMorletWavelet +from hypyp.wavelet.implementations.pywavelets_wavelet import ( + ComplexGaussianWavelet, + ComplexMorletWavelet, +) + def test_instanciate_complex_gaussian_wavelet(): wavelet = ComplexGaussianWavelet(degree=2) - assert wavelet.wavelet_name == 'cgau2' + assert wavelet.wavelet_name == "cgau2" assert wavelet.degree == 2 - assert len(wavelet.psi) == 2 ** 10 + assert len(wavelet.psi) == 2**10 assert isinstance(wavelet.cwt_params, dict) + def test_instanciate_complex_morlet_wavelet(): wavelet = ComplexMorletWavelet(2, 1) - assert wavelet.wavelet_name == 'cmor2,1' + assert wavelet.wavelet_name == "cmor2,1" assert wavelet.bandwidth_frequency == 2 assert wavelet.center_frequency == 1 + def test_resolution(): wavelet = ComplexMorletWavelet(evaluate=False) # need to be evaluated before calling .dt with pytest.raises(Exception): wavelet.psi_dx - + wavelet.evaluate_psi() assert wavelet.psi_dx > 0 assert wavelet.psi_dx < 1 + def test_default_domain(): wavelet = ComplexMorletWavelet() assert min(wavelet.psi_x) == wavelet.domain[0] assert max(wavelet.psi_x) == wavelet.domain[1] - + + def test_domain(): lower_bound = -1 upper_bound = 1 @@ -44,30 +52,38 @@ def test_domain(): assert max(wavelet.psi_x) == upper_bound assert wavelet.psi_dx > 0 assert wavelet.psi_dx < 1 - + + def test_psi(): wavelet = ComplexMorletWavelet(bandwidth_frequency=2, center_frequency=1) assert len(wavelet.psi_x) == len(wavelet.psi) - assert len(wavelet.psi_x) == 2 ** 10 + assert len(wavelet.psi_x) == 2**10 assert np.sum(np.real(wavelet.psi) * wavelet.psi_dx) != 0 assert np.sum(np.imag(wavelet.psi) * wavelet.psi_dx) != 0 assert np.sum(np.abs(wavelet.psi) * wavelet.psi_dx) == pytest.approx(1) + def test_number_of_scales(): assert len(ComplexMorletWavelet(period_range=(2, 4)).get_periods()) == 12 assert len(ComplexMorletWavelet(period_range=(4, 8)).get_periods()) == 12 assert len(ComplexMorletWavelet(period_range=(2, 8)).get_periods()) == 24 - assert len(ComplexMorletWavelet(dj=1/100, period_range=(2, 4)).get_periods()) == 100 + assert ( + len(ComplexMorletWavelet(dj=1 / 100, period_range=(2, 4)).get_periods()) == 100 + ) with pytest.raises(Exception): # must be < 1 ComplexMorletWavelet(dj=2) - -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) + + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_cwt(wavelet_class): wavelet = wavelet_class() signal = SyntheticSignal(duration=100).add_sin(0.05) @@ -76,6 +92,7 @@ def test_cwt(wavelet_class): assert len(res.scales) == len(res.frequencies) # TODO test something on res.W + def test_pair_signals(): signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() @@ -85,19 +102,23 @@ def test_pair_signals(): assert sub.x[-1] < 1.1 assert sub.x[-1] > 0.9 assert sub.section_idx == pair.section_idx - + + def test_pair_signals_epoch(): signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() pair = PairSignals(signal1.x, signal1.y, signal2.y) - sub = pair.sub((0, 1), section_idx=pair.section_idx+1) - assert sub.section_idx == pair.section_idx+1 - - -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) + sub = pair.sub((0, 1), section_idx=pair.section_idx + 1) + assert sub.section_idx == pair.section_idx + 1 + + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_wtc(wavelet_class): wavelet = wavelet_class(disable_caching=True) signal1 = SyntheticSignal().add_noise() @@ -106,10 +127,14 @@ def test_wtc(wavelet_class): assert res.coherence_metric > 0 assert res.coherence_metric < 1 -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_cwts(wavelet_class): wavelet = wavelet_class(disable_caching=True) signal1 = SyntheticSignal().add_noise() @@ -117,81 +142,140 @@ def test_cwts(wavelet_class): res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y)) assert res.cwt1 is not None assert res.cwt2 is not None - + + def test_cache_key(): - assert 'key_foo_bar' in ComplexMorletWavelet()._get_cache_key('foo', 'bar') - assert ComplexMorletWavelet(disable_caching=True)._get_cache_key('foo', 'bar') is None - + assert "key_foo_bar" in ComplexMorletWavelet()._get_cache_key("foo", "bar") + assert ( + ComplexMorletWavelet(disable_caching=True)._get_cache_key("foo", "bar") is None + ) + + def test_cache(): wavelet = ComplexMorletWavelet(cache=dict()) assert len(list(wavelet.cache.keys())) == 0 - wavelet._add_cache_item('foo', 'bar') + wavelet._add_cache_item("foo", "bar") assert len(list(wavelet.cache.keys())) == 1 - assert wavelet._get_cache_item('foo') == 'bar' + assert wavelet._get_cache_item("foo") == "bar" wavelet.clear_cache() assert len(list(wavelet.cache.keys())) == 0 - assert wavelet._get_cache_item('foo') == None + assert wavelet._get_cache_item("foo") == None zeros = np.zeros((10,)) x = np.arange(len(zeros)) # all these should yield a unique key pairs = [ - PairSignals(x, zeros, zeros, label_ch1='ch1', label_ch2='ch2', label_s1='subject1', label_s2='subject2', label_task='my_task'), - PairSignals(x[2:], zeros, zeros, label_ch1='ch1', label_ch2='ch2', label_s1='subject1', label_s2='subject2', label_task='my_task'), # other time range - PairSignals(x, zeros, zeros, label_ch1='ch1', label_ch2='ch2', label_s1='subject1', label_s2='subject2', label_task='my_other_task'), - PairSignals(x, zeros, zeros, label_ch1='ch3', label_ch2='ch4', label_s1='subject1', label_s2='subject2', label_task='my_task'), - PairSignals(x, zeros, zeros, label_ch1='ch1', label_ch2='ch2', label_s1='subject5.1', label_s2='subject5.2', label_task='my_task'), + PairSignals( + x, + zeros, + zeros, + label_ch1="ch1", + label_ch2="ch2", + label_s1="subject1", + label_s2="subject2", + label_task="my_task", + ), + PairSignals( + x[2:], + zeros, + zeros, + label_ch1="ch1", + label_ch2="ch2", + label_s1="subject1", + label_s2="subject2", + label_task="my_task", + ), # other time range + PairSignals( + x, + zeros, + zeros, + label_ch1="ch1", + label_ch2="ch2", + label_s1="subject1", + label_s2="subject2", + label_task="my_other_task", + ), + PairSignals( + x, + zeros, + zeros, + label_ch1="ch3", + label_ch2="ch4", + label_s1="subject1", + label_s2="subject2", + label_task="my_task", + ), + PairSignals( + x, + zeros, + zeros, + label_ch1="ch1", + label_ch2="ch2", + label_s1="subject5.1", + label_s2="subject5.2", + label_task="my_task", + ), ] # add all the keys to a list, then use a set to remove duplicates and make sure we still have the same count keys = [] for pair in pairs: - keys.append(wavelet._get_cache_key_pair(pair, 0, 'cwt')) - keys.append(wavelet._get_cache_key_pair(pair, 1, 'cwt')) + keys.append(wavelet._get_cache_key_pair(pair, 0, "cwt")) + keys.append(wavelet._get_cache_key_pair(pair, 1, "cwt")) - keys.append(wavelet._get_cache_key_pair(pairs[0], 0, 'another_suffix')) + keys.append(wavelet._get_cache_key_pair(pairs[0], 0, "another_suffix")) - #print(keys) + # print(keys) assert len(keys) == len(set(keys)) -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_wtc_coi_masked(wavelet_class): wavelet = wavelet_class(disable_caching=True) signal = SyntheticSignal().add_noise() res = wavelet.wtc(PairSignals(signal.x, signal.y, signal.y)) assert res.wtc_masked is not None - assert res.wtc_masked.mask[0,0] == True - assert res.wtc_masked.mask[0,len(signal.x)//2] == False - -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) -def test_period_frequency_range(wavelet_class): - frequency_range = np.array([5., 1.]) + assert res.wtc_masked.mask[0, 0] == True + assert res.wtc_masked.mask[0, len(signal.x) // 2] == False + + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) +def test_period_frequency_range(wavelet_class): + frequency_range = np.array([5.0, 1.0]) period_range = 1 / frequency_range signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() wavelet1 = wavelet_class(period_range=tuple(period_range), disable_caching=True) res1 = wavelet1.wtc(PairSignals(signal1.x, signal1.y, signal2.y)) - assert np.all(res1.frequencies[[0,-1]] == pytest.approx(frequency_range)) + assert np.all(res1.frequencies[[0, -1]] == pytest.approx(frequency_range)) - wavelet2 = wavelet_class(frequency_range=tuple(frequency_range), disable_caching=True) + wavelet2 = wavelet_class( + frequency_range=tuple(frequency_range), disable_caching=True + ) res2 = wavelet2.wtc(PairSignals(signal1.x, signal1.y, signal2.y)) - assert np.all(res2.frequencies[[0,-1]] == pytest.approx(frequency_range)) - + assert np.all(res2.frequencies[[0, -1]] == pytest.approx(frequency_range)) + + def test_smooth_in_scale_window(): assert np.sum(BaseWavelet._get_smoothing_window(0.6, 10)) == 1 - assert len(BaseWavelet._get_smoothing_window(1, 1/2)) == 4 - assert len(BaseWavelet._get_smoothing_window(1.1, 1/2)) == 5 + assert len(BaseWavelet._get_smoothing_window(1, 1 / 2)) == 4 + assert len(BaseWavelet._get_smoothing_window(1.1, 1 / 2)) == 5 # dirac - assert len(BaseWavelet._get_smoothing_window(1/12, 1/12)) == 2 + assert len(BaseWavelet._get_smoothing_window(1 / 12, 1 / 12)) == 2 win10 = BaseWavelet._get_smoothing_window(5, 1) assert len(win10) == 10 @@ -201,17 +285,22 @@ def test_smooth_in_scale_window(): assert len(win10_plus) == 11 assert np.mean(win10_plus) > win10_plus[0] + def test_fft_kwargs(): d = BaseWavelet._get_fft_kwargs(np.zeros((10,))) - assert d['n'] == 16 - - d = BaseWavelet._get_fft_kwargs(np.zeros((10,)), extra='foo') - assert d['extra'] == 'foo' - -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) + assert d["n"] == 16 + + d = BaseWavelet._get_fft_kwargs(np.zeros((10,)), extra="foo") + assert d["extra"] == "foo" + + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_smoothing(wavelet_class): dt = 0.1 dj = 0.1 @@ -223,11 +312,15 @@ def test_smoothing(wavelet_class): assert np.all(smoothed > 0) assert np.all(smoothed < 1) # TODO Should test more - -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) + + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_to_pandas_df(wavelet_class): signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() @@ -235,12 +328,16 @@ def test_to_pandas_df(wavelet_class): res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y)) df = res.to_frame() - assert df['coherence'][0] > 0 + assert df["coherence"][0] > 0 + -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_downsampling(wavelet_class): wavelet = wavelet_class(disable_caching=True) signal = SyntheticSignal(n_points=2000).add_sin(1) @@ -258,6 +355,7 @@ def test_downsampling(wavelet_class): assert wtc.sfreq == 1 / dt assert wtc.nyquist == nyquist + def test_wtc_time_slicing(): tmax = 100 n = 1000 @@ -270,11 +368,12 @@ def test_wtc_time_slicing(): assert df.shape[0] == 10 # first and last bins should be excluded (nan) since they do not have enough unmasked values - masked = df['coherence_masked'] + masked = df["coherence_masked"] assert masked[0] > masked[1] - assert np.isnan(df.at[0, 'coherence']) - assert np.isnan(df.at[df.shape[0]-1, 'coherence']) - + assert np.isnan(df.at[0, "coherence"]) + assert np.isnan(df.at[df.shape[0] - 1, "coherence"]) + + def test_wtc_period_slicing(): tmax = 100 n = 1000 @@ -282,15 +381,18 @@ def test_wtc_period_slicing(): signal2 = SyntheticSignal(duration=tmax, n_points=n).add_noise() wavelet = ComplexMorletWavelet(disable_caching=True) period_cuts = [3, 5, 10] - res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y), period_cuts=period_cuts) + res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y), period_cuts=period_cuts + ) df = res.to_frame() assert df.shape[0] == 4 # first and last bins should be excluded (nan) since they do not have enough unmasked values - masked = df['coherence_masked'] + masked = df["coherence_masked"] assert masked[0] < masked[1] + def test_wtc_period_vs_frequency_slicing(): tmax = 100 n = 1000 @@ -299,21 +401,28 @@ def test_wtc_period_vs_frequency_slicing(): wavelet = ComplexMorletWavelet(disable_caching=True) period_cuts = np.array([3, 5, 10]) frequency_cuts = 1 / period_cuts - period_res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y), period_cuts=period_cuts) - frequency_res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y), frequency_cuts=frequency_cuts) + period_res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y), period_cuts=period_cuts + ) + frequency_res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y), frequency_cuts=frequency_cuts + ) assert np.all(period_res.W == frequency_res.W) + def test_wtc_time_series(): signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() wavelet = ComplexMorletWavelet(disable_caching=True) period_cuts = [3, 5, 10] - res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y), period_cuts=period_cuts) + res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y), period_cuts=period_cuts + ) time_series = res.get_as_time_series() assert time_series.shape[0] == 4 assert time_series.shape[1] == len(signal1.x) - + def test_wtc_period_slicing_edge_cases(): tmax = 100 @@ -327,7 +436,8 @@ def test_wtc_period_slicing_edge_cases(): assert wavelet.wtc(pair, period_cuts=[0, 1, 1.5]).to_frame().shape[0] == 1 assert wavelet.wtc(pair, period_cuts=[3, 10]).to_frame().shape[0] == 3 assert wavelet.wtc(pair, period_cuts=[3, 3.001, 3.002, 10]).to_frame().shape[0] == 3 - + + def test_wtc_period_time_combined_slicing(): tmax = 100 n = 1000 @@ -336,16 +446,24 @@ def test_wtc_period_time_combined_slicing(): wavelet = ComplexMorletWavelet(disable_caching=True) bin_seconds = 20 period_cuts = [3, 5, 10] - res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y), bin_seconds=bin_seconds, period_cuts=period_cuts) + res = wavelet.wtc( + PairSignals(signal1.x, signal1.y, signal2.y), + bin_seconds=bin_seconds, + period_cuts=period_cuts, + ) df = res.to_frame() assert df.shape[0] == 5 * 4 - #print(df) - -@pytest.mark.parametrize("wavelet_class", [ - ComplexMorletWavelet, - ComplexGaussianWavelet, -]) + # print(df) + + +@pytest.mark.parametrize( + "wavelet_class", + [ + ComplexMorletWavelet, + ComplexGaussianWavelet, + ], +) def test_wtc_wavelet_info(wavelet_class): signal1 = SyntheticSignal().add_noise() signal2 = SyntheticSignal().add_noise() @@ -354,16 +472,20 @@ def test_wtc_wavelet_info(wavelet_class): res = wavelet.wtc(pair) df = res.to_frame() - assert df.at[0, 'wavelet_library'] == 'pywavelets' - assert df.at[0, 'wavelet_name'] == wavelet.wavelet_name_with_args - -@pytest.mark.parametrize("wavelet", [ - ComplexMorletWavelet(), - ComplexMorletWavelet(bandwidth_frequency=10, center_frequency=1), - ComplexGaussianWavelet(degree=1), - ComplexGaussianWavelet(degree=2), - ComplexGaussianWavelet(degree=3), -]) + assert df.at[0, "wavelet_library"] == "pywavelets" + assert df.at[0, "wavelet_name"] == wavelet.wavelet_name_with_args + + +@pytest.mark.parametrize( + "wavelet", + [ + ComplexMorletWavelet(), + ComplexMorletWavelet(bandwidth_frequency=10, center_frequency=1), + ComplexGaussianWavelet(degree=1), + ComplexGaussianWavelet(degree=2), + ComplexGaussianWavelet(degree=3), + ], +) def test_cone_of_influence(wavelet): n = 11 dt = 1 @@ -372,6 +494,3 @@ def test_cone_of_influence(wavelet): assert coi[0] < coi[1] assert np.argmax(coi) == n // 2 assert len(coi) == n - - - \ No newline at end of file diff --git a/tests/test_wavelet_implementations.py b/tests/test_wavelet_implementations.py index 972ddd3..d668c2d 100644 --- a/tests/test_wavelet_implementations.py +++ b/tests/test_wavelet_implementations.py @@ -11,10 +11,11 @@ # TODO: test values with sinusoid signal + def test_pywavelets(): wavelet = ComplexMorletWavelet(disable_caching=True) psi, x = wavelet.evaluate_psi() - assert psi.dtype.kind == 'c' + assert psi.dtype.kind == "c" assert min(x) == wavelet.lower_bound assert max(x) == wavelet.upper_bound assert len(x) == len(psi) @@ -22,13 +23,14 @@ def test_pywavelets(): signal2 = SyntheticSignal().add_noise() res = wavelet.wtc(PairSignals(signal1.x, signal1.y, signal2.y)) + def test_pycwt(): if PycwtWavelet is None: pytest.skip("Optional dependency Pycwt is not installed") - + wavelet = PycwtWavelet() psi, x = wavelet.evaluate_psi() - assert psi.dtype.kind == 'c' + assert psi.dtype.kind == "c" assert min(x) == wavelet.lower_bound assert max(x) == wavelet.upper_bound assert len(x) == len(psi) @@ -39,4 +41,3 @@ def test_pycwt(): # make sure we do not exceed the period range assert res.periods[0] >= wavelet.period_range[0] assert res.periods[-1] <= wavelet.period_range[1] - diff --git a/tests/test_xdf.py b/tests/test_xdf.py index 39762c4..c16a017 100644 --- a/tests/test_xdf.py +++ b/tests/test_xdf.py @@ -26,28 +26,29 @@ STREAM_ID_EEG_2 = 7 STREAM_ID_QUALITY_1 = 8 + def test_instanciate(): xdf = XDFImport(file_path) assert xdf.file_path == file_path assert len(xdf.available_streams) == 8 - assert xdf.available_streams[0].name == 'LSLOutletHS1-Markers' - assert xdf.available_streams[1].name == 'LSLOutletHS1-Accelerometer' - assert xdf.available_streams[2].name == 'LSLOutletHS1-Quality' - assert xdf.available_streams[3].name == 'LSLOutletHS1-EEG' - assert xdf.available_streams[4].name == 'LSLOutletHS2-Markers' - assert xdf.available_streams[5].name == 'LSLOutletHS2-Accelerometer' - assert xdf.available_streams[6].name == 'LSLOutletHS2-Quality' - assert xdf.available_streams[7].name == 'LSLOutletHS2-EEG' - - assert xdf.available_streams[0].type == 'Markers' - assert xdf.available_streams[1].type == 'Accelerometer' - assert xdf.available_streams[2].type == 'Quality' - assert xdf.available_streams[3].type == 'EEG' - assert xdf.available_streams[4].type == 'Markers' - assert xdf.available_streams[5].type == 'Accelerometer' - assert xdf.available_streams[6].type == 'Quality' - assert xdf.available_streams[7].type == 'EEG' + assert xdf.available_streams[0].name == "LSLOutletHS1-Markers" + assert xdf.available_streams[1].name == "LSLOutletHS1-Accelerometer" + assert xdf.available_streams[2].name == "LSLOutletHS1-Quality" + assert xdf.available_streams[3].name == "LSLOutletHS1-EEG" + assert xdf.available_streams[4].name == "LSLOutletHS2-Markers" + assert xdf.available_streams[5].name == "LSLOutletHS2-Accelerometer" + assert xdf.available_streams[6].name == "LSLOutletHS2-Quality" + assert xdf.available_streams[7].name == "LSLOutletHS2-EEG" + + assert xdf.available_streams[0].type == "Markers" + assert xdf.available_streams[1].type == "Accelerometer" + assert xdf.available_streams[2].type == "Quality" + assert xdf.available_streams[3].type == "EEG" + assert xdf.available_streams[4].type == "Markers" + assert xdf.available_streams[5].type == "Accelerometer" + assert xdf.available_streams[6].type == "Quality" + assert xdf.available_streams[7].type == "EEG" assert xdf.available_streams[0].id == STREAM_ID_MARKERS_1 assert xdf.available_streams[1].id == STREAM_ID_ACC_1 @@ -67,34 +68,40 @@ def test_instanciate(): assert xdf.map_id_to_idx[STREAM_ID_QUALITY_2] == 6 assert xdf.map_id_to_idx[STREAM_ID_EEG_2] == 7 + def test_convert_by_default(): - xdf = XDFImport(file_path, select_type='EEG') + xdf = XDFImport(file_path, select_type="EEG") assert xdf.file_path == file_path assert len(xdf.mne_raws) == 2 + def test_instanciate_no_matching_stream(): with pytest.raises(Exception): # use a non-existing stream_id non_existing_stream_id = 0 XDFImport(file_path, select_matches=[non_existing_stream_id]) + def test_xdf_stream_indices_per_type(): # should not be able to create mne object when using stream of type Markers xdf = XDFImport(file_path, convert_to_mne=False) - assert STREAM_ID_MARKERS_1 in xdf.get_stream_ids_for_type('Markers') - assert STREAM_ID_MARKERS_2 in xdf.get_stream_ids_for_type('Markers') - assert STREAM_ID_EEG_1 in xdf.get_stream_ids_for_type('EEG') - assert STREAM_ID_EEG_2 in xdf.get_stream_ids_for_type('EEG') + assert STREAM_ID_MARKERS_1 in xdf.get_stream_ids_for_type("Markers") + assert STREAM_ID_MARKERS_2 in xdf.get_stream_ids_for_type("Markers") + assert STREAM_ID_EEG_1 in xdf.get_stream_ids_for_type("EEG") + assert STREAM_ID_EEG_2 in xdf.get_stream_ids_for_type("EEG") + def test_all_eeg_streams(): xdf = XDFImport(file_path) assert len(xdf.mne_raws[0].ch_names) == len(xdf.selected_signal_streams[0].ch_names) + def test_match_stream_by_idx(): xdf = XDFImport(file_path, select_matches=[STREAM_ID_EEG_1]) assert len(xdf.mne_raws) == 1 assert len(xdf.mne_raws[0].ch_names) == len(xdf.selected_streams[0].ch_names) - + + def test_match_stream_by_name(): # first grab the name of the first channel xdf_tmp = XDFImport(file_path) @@ -104,6 +111,7 @@ def test_match_stream_by_name(): xdf = XDFImport(file_path, select_matches=[key]) assert list(xdf.mne_raws_dict.keys())[0] == key + def test_match_stream_by_both(): # first grab the name of the first channel xdf_tmp = XDFImport(file_path) @@ -113,57 +121,73 @@ def test_match_stream_by_both(): xdf = XDFImport(file_path, select_matches=[key, STREAM_ID_EEG_2]) assert len(list(xdf.mne_raws_dict.keys())) == 2 + def test_match_unexistent_type(): with pytest.raises(Exception): - xdf = XDFImport(file_path, select_type='foo') + xdf = XDFImport(file_path, select_type="foo") + def test_stream_type_to_mne_type(): - # Available mne types are : - #[ 'grad', 'mag', 'ref_meg', 'eeg', 'seeg', 'dbs', 'ecog', 'eog', 'emg', 'ecg', + # Available mne types are : + # [ 'grad', 'mag', 'ref_meg', 'eeg', 'seeg', 'dbs', 'ecog', 'eog', 'emg', 'ecg', # 'resp', 'bio', 'misc', 'stim', 'exci', 'syst', 'ias', 'gof', 'dipole', 'chpi', # 'fnirs_cw_amplitude', 'fnirs_fd_ac_amplitude', 'fnirs_fd_phase', 'fnirs_od', 'hbo', 'hbr', # 'csd', 'temperature', 'gsr', 'eyegaze', 'pupil' ] - assert XDFStream.get_default_mne_ch_type_for_stream_type('EEG') == 'eeg' - assert XDFStream.get_default_mne_ch_type_for_stream_type('fNIRS') == 'fnirs_cw_amplitude' - assert XDFStream.get_default_mne_ch_type_for_stream_type('markers') == 'stim' - assert XDFStream.get_default_mne_ch_type_for_stream_type('stim') == 'stim' + assert XDFStream.get_default_mne_ch_type_for_stream_type("EEG") == "eeg" + assert ( + XDFStream.get_default_mne_ch_type_for_stream_type("fNIRS") + == "fnirs_cw_amplitude" + ) + assert XDFStream.get_default_mne_ch_type_for_stream_type("markers") == "stim" + assert XDFStream.get_default_mne_ch_type_for_stream_type("stim") == "stim" + def test_stream_type_map_explicit(): - mne_force_type = 'misc' - my_map = {'EEG': mne_force_type} - assert XDFStream.get_default_mne_ch_type_for_stream_type('EEG', my_map) == mne_force_type + mne_force_type = "misc" + my_map = {"EEG": mne_force_type} + assert ( + XDFStream.get_default_mne_ch_type_for_stream_type("EEG", my_map) + == mne_force_type + ) # Force EEG data as 'misc' in mne xdf = XDFImport(file_path, mne_type_map=my_map) - assert xdf.available_streams[xdf.map_id_to_idx[STREAM_ID_EEG_1]].type == 'EEG' - assert xdf.mne_raws[0].get_channel_types()[0] == 'misc' + assert xdf.available_streams[xdf.map_id_to_idx[STREAM_ID_EEG_1]].type == "EEG" + assert xdf.mne_raws[0].get_channel_types()[0] == "misc" + def test_channel_names(): xdf = XDFImport(file_path, convert_to_mne=False) ch_names = xdf.available_streams[xdf.map_id_to_idx[STREAM_ID_EEG_1]].ch_names - assert ch_names[0] == 'Ch1' + assert ch_names[0] == "Ch1" + def test_channel_names_standard_metadata(): # the description of channels in the "file_path" xdf does not follow the spec here https://github.com/sccn/xdf/wiki/EEG-Meta-Data # let's try with another xdf file which follows the spec xdf = XDFImport(file_path_with_markers, convert_to_mne=False) ch_names = xdf.selected_signal_streams[0].ch_names - assert ch_names[0] == 'Ch1' - assert ch_names[1] == 'Ch2' - assert ch_names[2] == 'Ch3' + assert ch_names[0] == "Ch1" + assert ch_names[1] == "Ch2" + assert ch_names[2] == "Ch3" + def test_ch_names_to_ch_types(): # look at automatic detection of mne types - ch_types = XDFStream.get_mne_ch_types('EEG', ['Fp1', 'Fp2', 'AccX', 'GyroX', 'QuatX']) - assert ch_types == ['eeg', 'eeg', 'misc', 'misc', 'misc'] + ch_types = XDFStream.get_mne_ch_types( + "EEG", ["Fp1", "Fp2", "AccX", "GyroX", "QuatX"] + ) + assert ch_types == ["eeg", "eeg", "misc", "misc", "misc"] + def test_ch_types_in_raw(): - rawEEG = XDFImport(file_path, select_type='EEG').mne_raws[0] - assert rawEEG.get_channel_types()[0] == 'eeg' + rawEEG = XDFImport(file_path, select_type="EEG").mne_raws[0] + assert rawEEG.get_channel_types()[0] == "eeg" + + rawAcc = XDFImport(file_path, select_type="Accelerometer").mne_raws[0] + assert rawAcc.get_channel_types()[0] == "misc" - rawAcc = XDFImport(file_path, select_type='Accelerometer').mne_raws[0] - assert rawAcc.get_channel_types()[0] == 'misc' def test_fif_file_for_stream(): xdf = XDFImport(file_path) @@ -171,45 +195,66 @@ def test_fif_file_for_stream(): with tempfile.TemporaryDirectory() as tmp_dir: saved_file = xdf_stream.save_to_fif_file(tmp_dir) raw_reloaded = mne.io.read_raw_fif(saved_file) - assert raw_reloaded.get_data()[0][0] == xdf.mne_raws_dict[xdf_stream.name].get_data()[0][0] - + assert ( + raw_reloaded.get_data()[0][0] + == xdf.mne_raws_dict[xdf_stream.name].get_data()[0][0] + ) + + def test_mne_compatible_stream_types(): xdf = XDFImport(file_path, convert_to_mne=False) - marker_streams = [stream for stream in xdf.available_streams if stream.type == 'Markers'] + marker_streams = [ + stream for stream in xdf.available_streams if stream.type == "Markers" + ] assert marker_streams[0].is_mne_raw_compatible == False - + + def test_subject_unique_name(): xdf = XDFImport(file_path, select_matches=[1], convert_to_mne=False) - assert xdf.selected_streams[0].get_unique_name() == 'LSLOutletHS1-Accelerometer' - assert xdf.selected_streams[0].get_unique_name(append_stream_id=True) == 'LSLOutletHS1-Accelerometer-1' - + assert xdf.selected_streams[0].get_unique_name() == "LSLOutletHS1-Accelerometer" + assert ( + xdf.selected_streams[0].get_unique_name(append_stream_id=True) + == "LSLOutletHS1-Accelerometer-1" + ) + + def test_markers(): xdf = XDFImport(file_path_with_markers) keys = list(xdf.markers_dict) assert len(keys) == 1 - assert keys[0] == XDFImport(file_path_with_markers, select_type='Markers').selected_streams[0].name + assert ( + keys[0] + == XDFImport(file_path_with_markers, select_type="Markers") + .selected_streams[0] + .name + ) markers = list(xdf.markers_dict.values())[0] annotations = markers.as_mne_annotations() assert len(annotations) > 0 + def test_merge_markers(): xdf = XDFImport(file_path_with_markers) annotations = xdf.markers.as_mne_annotations() markers_dup = xdf.markers + xdf.markers assert len(markers_dup.as_mne_annotations()) == 2 * len(annotations) + def test_markers_timing(): xdf = XDFImport(file_path_with_markers) markers = xdf.markers offset = markers.timestamps[0] annotations = markers.as_mne_annotations(offset) - assert annotations[0]['onset'] == 0.0 + assert annotations[0]["onset"] == 0.0 - annotations_for_stream = markers.as_mne_annotations(xdf.selected_signal_streams[0].reference_time) - assert annotations_for_stream[0]['onset'] > 0 - assert annotations_for_stream[0]['onset'] < 10 + annotations_for_stream = markers.as_mne_annotations( + xdf.selected_signal_streams[0].reference_time + ) + assert annotations_for_stream[0]["onset"] > 0 + assert annotations_for_stream[0]["onset"] < 10 xdf.mne_raws[0].set_annotations(annotations_for_stream) + def test_markers_negative(): xdf = XDFImport(file_path_with_markers) annotations_all = xdf.markers.as_mne_annotations() @@ -221,16 +266,53 @@ def test_markers_negative(): # have a big time offset to make sure we receive no annotations assert len(xdf.markers.as_mne_annotations(first_time + 2000)) == 0 + def test_fif_files(): - xdf = XDFImport(file_path, select_type='EEG') + xdf = XDFImport(file_path, select_type="EEG") with tempfile.TemporaryDirectory() as tmp_dir: saved_files = xdf.save_to_fif_files(tmp_dir) assert len(saved_files) == 2 # make sure we can load without an error _ = mne.io.read_raw_fif(saved_files[0]) _ = mne.io.read_raw_fif(saved_files[1]) - + + def test_montage(): - xdf = XDFImport(file_path, select_type='EEG') - xdf.rename_channels(['Fp1', 'Fp2', 'F3', 'F4', 'C3', 'C4', 'P3', 'P4', 'O1', 'O2', 'F7', 'F8', 'T7', 'T8', 'P7', 'P8', 'Fz', 'Cz', 'Pz', 'POz', 'FC1', 'FC2', 'CP1', 'CP2', 'FC5', 'FC6', 'CP5', 'CP6', 'FT9', 'FT10', 'TP9', 'TP10']) - xdf.set_montage('standard_1020') + xdf = XDFImport(file_path, select_type="EEG") + xdf.rename_channels( + [ + "Fp1", + "Fp2", + "F3", + "F4", + "C3", + "C4", + "P3", + "P4", + "O1", + "O2", + "F7", + "F8", + "T7", + "T8", + "P7", + "P8", + "Fz", + "Cz", + "Pz", + "POz", + "FC1", + "FC2", + "CP1", + "CP2", + "FC5", + "FC6", + "CP5", + "CP6", + "FT9", + "FT10", + "TP9", + "TP10", + ] + ) + xdf.set_montage("standard_1020") diff --git a/uv.lock b/uv.lock index 4d4dc27..651bc51 100644 --- a/uv.lock +++ b/uv.lock @@ -192,43 +192,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1a/39/47f9197bdd44df24d67ac8893641e16f386c984a0619ef2ee4c51fbbc019/beautifulsoup4-4.14.3-py3-none-any.whl", hash = "sha256:0918bfe44902e6ad8d57732ba310582e98da931428d231a5ecb9e7c703a735bb", size = 107721, upload-time = "2025-11-30T15:08:24.087Z" }, ] -[[package]] -name = "black" -version = "26.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "pathspec" }, - { name = "platformdirs" }, - { name = "pytokens" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e1/c5/61175d618685d42b005847464b8fb4743a67b1b8fdb75e50e5a96c31a27a/black-26.3.1.tar.gz", hash = "sha256:2c50f5063a9641c7eed7795014ba37b0f5fa227f3d408b968936e24bc0566b07", size = 666155, upload-time = "2026-03-12T03:36:03.593Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/17/57/5f11c92861f9c92eb9dddf515530bc2d06db843e44bdcf1c83c1427824bc/black-26.3.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:28ef38aee69e4b12fda8dba75e21f9b4f979b490c8ac0baa7cb505369ac9e1ff", size = 1851987, upload-time = "2026-03-12T03:40:06.248Z" }, - { url = "https://files.pythonhosted.org/packages/54/aa/340a1463660bf6831f9e39646bf774086dbd8ca7fc3cded9d59bbdf4ad0a/black-26.3.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bf9bf162ed91a26f1adba8efda0b573bc6924ec1408a52cc6f82cb73ec2b142c", size = 1689499, upload-time = "2026-03-12T03:40:07.642Z" }, - { url = "https://files.pythonhosted.org/packages/f3/01/b726c93d717d72733da031d2de10b92c9fa4c8d0c67e8a8a372076579279/black-26.3.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:474c27574d6d7037c1bc875a81d9be0a9a4f9ee95e62800dab3cfaadbf75acd5", size = 1754369, upload-time = "2026-03-12T03:40:09.279Z" }, - { url = "https://files.pythonhosted.org/packages/e3/09/61e91881ca291f150cfc9eb7ba19473c2e59df28859a11a88248b5cbbc4d/black-26.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:5e9d0d86df21f2e1677cc4bd090cd0e446278bcbbe49bf3659c308c3e402843e", size = 1413613, upload-time = "2026-03-12T03:40:10.943Z" }, - { url = "https://files.pythonhosted.org/packages/16/73/544f23891b22e7efe4d8f812371ab85b57f6a01b2fc45e3ba2e52ba985b8/black-26.3.1-cp311-cp311-win_arm64.whl", hash = "sha256:9a5e9f45e5d5e1c5b5c29b3bd4265dcc90e8b92cf4534520896ed77f791f4da5", size = 1219719, upload-time = "2026-03-12T03:40:12.597Z" }, - { url = "https://files.pythonhosted.org/packages/dc/f8/da5eae4fc75e78e6dceb60624e1b9662ab00d6b452996046dfa9b8a6025b/black-26.3.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e6f89631eb88a7302d416594a32faeee9fb8fb848290da9d0a5f2903519fc1", size = 1895920, upload-time = "2026-03-12T03:40:13.921Z" }, - { url = "https://files.pythonhosted.org/packages/2c/9f/04e6f26534da2e1629b2b48255c264cabf5eedc5141d04516d9d68a24111/black-26.3.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:41cd2012d35b47d589cb8a16faf8a32ef7a336f56356babd9fcf70939ad1897f", size = 1718499, upload-time = "2026-03-12T03:40:15.239Z" }, - { url = "https://files.pythonhosted.org/packages/04/91/a5935b2a63e31b331060c4a9fdb5a6c725840858c599032a6f3aac94055f/black-26.3.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f76ff19ec5297dd8e66eb64deda23631e642c9393ab592826fd4bdc97a4bce7", size = 1794994, upload-time = "2026-03-12T03:40:17.124Z" }, - { url = "https://files.pythonhosted.org/packages/e7/0a/86e462cdd311a3c2a8ece708d22aba17d0b2a0d5348ca34b40cdcbea512e/black-26.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:ddb113db38838eb9f043623ba274cfaf7d51d5b0c22ecb30afe58b1bb8322983", size = 1420867, upload-time = "2026-03-12T03:40:18.83Z" }, - { url = "https://files.pythonhosted.org/packages/5b/e5/22515a19cb7eaee3440325a6b0d95d2c0e88dd180cb011b12ae488e031d1/black-26.3.1-cp312-cp312-win_arm64.whl", hash = "sha256:dfdd51fc3e64ea4f35873d1b3fb25326773d55d2329ff8449139ebaad7357efb", size = 1230124, upload-time = "2026-03-12T03:40:20.425Z" }, - { url = "https://files.pythonhosted.org/packages/f5/77/5728052a3c0450c53d9bb3945c4c46b91baa62b2cafab6801411b6271e45/black-26.3.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:855822d90f884905362f602880ed8b5df1b7e3ee7d0db2502d4388a954cc8c54", size = 1895034, upload-time = "2026-03-12T03:40:21.813Z" }, - { url = "https://files.pythonhosted.org/packages/52/73/7cae55fdfdfbe9d19e9a8d25d145018965fe2079fa908101c3733b0c55a0/black-26.3.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8a33d657f3276328ce00e4d37fe70361e1ec7614da5d7b6e78de5426cb56332f", size = 1718503, upload-time = "2026-03-12T03:40:23.666Z" }, - { url = "https://files.pythonhosted.org/packages/e1/87/af89ad449e8254fdbc74654e6467e3c9381b61472cc532ee350d28cfdafb/black-26.3.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f1cd08e99d2f9317292a311dfe578fd2a24b15dbce97792f9c4d752275c1fa56", size = 1793557, upload-time = "2026-03-12T03:40:25.497Z" }, - { url = "https://files.pythonhosted.org/packages/43/10/d6c06a791d8124b843bf325ab4ac7d2f5b98731dff84d6064eafd687ded1/black-26.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:c7e72339f841b5a237ff14f7d3880ddd0fc7f98a1199e8c4327f9a4f478c1839", size = 1422766, upload-time = "2026-03-12T03:40:27.14Z" }, - { url = "https://files.pythonhosted.org/packages/59/4f/40a582c015f2d841ac24fed6390bd68f0fc896069ff3a886317959c9daf8/black-26.3.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc622538b430aa4c8c853f7f63bc582b3b8030fd8c80b70fb5fa5b834e575c2", size = 1232140, upload-time = "2026-03-12T03:40:28.882Z" }, - { url = "https://files.pythonhosted.org/packages/d5/da/e36e27c9cebc1311b7579210df6f1c86e50f2d7143ae4fcf8a5017dc8809/black-26.3.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:2d6bfaf7fd0993b420bed691f20f9492d53ce9a2bcccea4b797d34e947318a78", size = 1889234, upload-time = "2026-03-12T03:40:30.964Z" }, - { url = "https://files.pythonhosted.org/packages/0e/7b/9871acf393f64a5fa33668c19350ca87177b181f44bb3d0c33b2d534f22c/black-26.3.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f89f2ab047c76a9c03f78d0d66ca519e389519902fa27e7a91117ef7611c0568", size = 1720522, upload-time = "2026-03-12T03:40:32.346Z" }, - { url = "https://files.pythonhosted.org/packages/03/87/e766c7f2e90c07fb7586cc787c9ae6462b1eedab390191f2b7fc7f6170a9/black-26.3.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b07fc0dab849d24a80a29cfab8d8a19187d1c4685d8a5e6385a5ce323c1f015f", size = 1787824, upload-time = "2026-03-12T03:40:33.636Z" }, - { url = "https://files.pythonhosted.org/packages/ac/94/2424338fb2d1875e9e83eed4c8e9c67f6905ec25afd826a911aea2b02535/black-26.3.1-cp314-cp314-win_amd64.whl", hash = "sha256:0126ae5b7c09957da2bdbd91a9ba1207453feada9e9fe51992848658c6c8e01c", size = 1445855, upload-time = "2026-03-12T03:40:35.442Z" }, - { url = "https://files.pythonhosted.org/packages/86/43/0c3338bd928afb8ee7471f1a4eec3bdbe2245ccb4a646092a222e8669840/black-26.3.1-cp314-cp314-win_arm64.whl", hash = "sha256:92c0ec1f2cc149551a2b7b47efc32c866406b6891b0ee4625e95967c8f4acfb1", size = 1258109, upload-time = "2026-03-12T03:40:36.832Z" }, - { url = "https://files.pythonhosted.org/packages/8e/0d/52d98722666d6fc6c3dd4c76df339501d6efd40e0ff95e6186a7b7f0befd/black-26.3.1-py3-none-any.whl", hash = "sha256:2bd5aa94fc267d38bb21a70d7410a89f1a1d318841855f698746f8e7f51acd1b", size = 207542, upload-time = "2026-03-12T03:36:01.668Z" }, -] - [[package]] name = "bleach" version = "6.4.0" @@ -978,7 +941,6 @@ benchmark = [ { name = "doit" }, ] dev = [ - { name = "black" }, { name = "ipykernel" }, { name = "ipywidgets" }, { name = "jupyterlab" }, @@ -995,6 +957,7 @@ dev = [ { name = "pytest" }, { name = "pytest-mock" }, { name = "pytest-watch" }, + { name = "ruff" }, { name = "tabulate" }, { name = "widgetsnbextension" }, ] @@ -1005,6 +968,9 @@ docs = [ { name = "mkdocstrings" }, { name = "mkdocstrings-python" }, ] +lint = [ + { name = "ruff" }, +] [package.metadata] requires-dist = [ @@ -1039,7 +1005,6 @@ provides-extras = ["shiny", "numba", "torch", "metal", "cupy"] [package.metadata.requires-dev] benchmark = [{ name = "doit", specifier = ">=0.36.0" }] dev = [ - { name = "black", specifier = ">=25.1.0" }, { name = "ipykernel", specifier = ">=6.19.4" }, { name = "ipywidgets", specifier = ">=7.5.1" }, { name = "jupyterlab", specifier = ">=4.3.5" }, @@ -1056,6 +1021,7 @@ dev = [ { name = "pytest", specifier = ">=7.2.0" }, { name = "pytest-mock", specifier = ">=3.14.0" }, { name = "pytest-watch", specifier = ">=4.2.0" }, + { name = "ruff", specifier = ">=0.16.10" }, { name = "tabulate", specifier = ">=0.9.0" }, { name = "widgetsnbextension", specifier = ">=3.5.1" }, ] @@ -1066,6 +1032,7 @@ docs = [ { name = "mkdocstrings", specifier = ">=0.20.0" }, { name = "mkdocstrings-python", specifier = ">=1.1.0" }, ] +lint = [{ name = "ruff", specifier = ">=0.16.10" }] [[package]] name = "idna" @@ -2122,15 +2089,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] -[[package]] -name = "mypy-extensions" -version = "1.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, -] - [[package]] name = "narwhals" version = "2.18.0" @@ -3233,40 +3191,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2c/29/b2be72ba559df54647e6ad2ff6cf1a94d488857ee67eeae28b252a6e734d/pytkdocs-0.16.5-py3-none-any.whl", hash = "sha256:7feb2535a8582d891fa5353cf2f4ecc7840915e631d322dbcbc1c8cc605da23a", size = 38823, upload-time = "2025-03-09T17:32:40.9Z" }, ] -[[package]] -name = "pytokens" -version = "0.4.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b6/34/b4e015b99031667a7b960f888889c5bd34ef585c85e1cb56a594b92836ac/pytokens-0.4.1.tar.gz", hash = "sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a", size = 23015, upload-time = "2026-01-30T01:03:45.924Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/92/790ebe03f07b57e53b10884c329b9a1a308648fc083a6d4a39a10a28c8fc/pytokens-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d70e77c55ae8380c91c0c18dea05951482e263982911fc7410b1ffd1dadd3440", size = 160864, upload-time = "2026-01-30T01:02:57.882Z" }, - { url = "https://files.pythonhosted.org/packages/13/25/a4f555281d975bfdd1eba731450e2fe3a95870274da73fb12c40aeae7625/pytokens-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a58d057208cb9075c144950d789511220b07636dd2e4708d5645d24de666bdc", size = 248565, upload-time = "2026-01-30T01:02:59.912Z" }, - { url = "https://files.pythonhosted.org/packages/17/50/bc0394b4ad5b1601be22fa43652173d47e4c9efbf0044c62e9a59b747c56/pytokens-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b49750419d300e2b5a3813cf229d4e5a4c728dae470bcc89867a9ad6f25a722d", size = 260824, upload-time = "2026-01-30T01:03:01.471Z" }, - { url = "https://files.pythonhosted.org/packages/4e/54/3e04f9d92a4be4fc6c80016bc396b923d2a6933ae94b5f557c939c460ee0/pytokens-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d9907d61f15bf7261d7e775bd5d7ee4d2930e04424bab1972591918497623a16", size = 264075, upload-time = "2026-01-30T01:03:04.143Z" }, - { url = "https://files.pythonhosted.org/packages/d1/1b/44b0326cb5470a4375f37988aea5d61b5cc52407143303015ebee94abfd6/pytokens-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:ee44d0f85b803321710f9239f335aafe16553b39106384cef8e6de40cb4ef2f6", size = 103323, upload-time = "2026-01-30T01:03:05.412Z" }, - { url = "https://files.pythonhosted.org/packages/41/5d/e44573011401fb82e9d51e97f1290ceb377800fb4eed650b96f4753b499c/pytokens-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:140709331e846b728475786df8aeb27d24f48cbcf7bcd449f8de75cae7a45083", size = 160663, upload-time = "2026-01-30T01:03:06.473Z" }, - { url = "https://files.pythonhosted.org/packages/f0/e6/5bbc3019f8e6f21d09c41f8b8654536117e5e211a85d89212d59cbdab381/pytokens-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d6c4268598f762bc8e91f5dbf2ab2f61f7b95bdc07953b602db879b3c8c18e1", size = 255626, upload-time = "2026-01-30T01:03:08.177Z" }, - { url = "https://files.pythonhosted.org/packages/bf/3c/2d5297d82286f6f3d92770289fd439956b201c0a4fc7e72efb9b2293758e/pytokens-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24afde1f53d95348b5a0eb19488661147285ca4dd7ed752bbc3e1c6242a304d1", size = 269779, upload-time = "2026-01-30T01:03:09.756Z" }, - { url = "https://files.pythonhosted.org/packages/20/01/7436e9ad693cebda0551203e0bf28f7669976c60ad07d6402098208476de/pytokens-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5ad948d085ed6c16413eb5fec6b3e02fa00dc29a2534f088d3302c47eb59adf9", size = 268076, upload-time = "2026-01-30T01:03:10.957Z" }, - { url = "https://files.pythonhosted.org/packages/2e/df/533c82a3c752ba13ae7ef238b7f8cdd272cf1475f03c63ac6cf3fcfb00b6/pytokens-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:3f901fe783e06e48e8cbdc82d631fca8f118333798193e026a50ce1b3757ea68", size = 103552, upload-time = "2026-01-30T01:03:12.066Z" }, - { url = "https://files.pythonhosted.org/packages/cb/dc/08b1a080372afda3cceb4f3c0a7ba2bde9d6a5241f1edb02a22a019ee147/pytokens-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b", size = 160720, upload-time = "2026-01-30T01:03:13.843Z" }, - { url = "https://files.pythonhosted.org/packages/64/0c/41ea22205da480837a700e395507e6a24425151dfb7ead73343d6e2d7ffe/pytokens-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f", size = 254204, upload-time = "2026-01-30T01:03:14.886Z" }, - { url = "https://files.pythonhosted.org/packages/e0/d2/afe5c7f8607018beb99971489dbb846508f1b8f351fcefc225fcf4b2adc0/pytokens-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1", size = 268423, upload-time = "2026-01-30T01:03:15.936Z" }, - { url = "https://files.pythonhosted.org/packages/68/d4/00ffdbd370410c04e9591da9220a68dc1693ef7499173eb3e30d06e05ed1/pytokens-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4", size = 266859, upload-time = "2026-01-30T01:03:17.458Z" }, - { url = "https://files.pythonhosted.org/packages/a7/c9/c3161313b4ca0c601eeefabd3d3b576edaa9afdefd32da97210700e47652/pytokens-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78", size = 103520, upload-time = "2026-01-30T01:03:18.652Z" }, - { url = "https://files.pythonhosted.org/packages/8f/a7/b470f672e6fc5fee0a01d9e75005a0e617e162381974213a945fcd274843/pytokens-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4a14d5f5fc78ce85e426aa159489e2d5961acf0e47575e08f35584009178e321", size = 160821, upload-time = "2026-01-30T01:03:19.684Z" }, - { url = "https://files.pythonhosted.org/packages/80/98/e83a36fe8d170c911f864bfded690d2542bfcfacb9c649d11a9e6eb9dc41/pytokens-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f50fd18543be72da51dd505e2ed20d2228c74e0464e4262e4899797803d7fa", size = 254263, upload-time = "2026-01-30T01:03:20.834Z" }, - { url = "https://files.pythonhosted.org/packages/0f/95/70d7041273890f9f97a24234c00b746e8da86df462620194cef1d411ddeb/pytokens-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc74c035f9bfca0255c1af77ddd2d6ae8419012805453e4b0e7513e17904545d", size = 268071, upload-time = "2026-01-30T01:03:21.888Z" }, - { url = "https://files.pythonhosted.org/packages/da/79/76e6d09ae19c99404656d7db9c35dfd20f2086f3eb6ecb496b5b31163bad/pytokens-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f66a6bbe741bd431f6d741e617e0f39ec7257ca1f89089593479347cc4d13324", size = 271716, upload-time = "2026-01-30T01:03:23.633Z" }, - { url = "https://files.pythonhosted.org/packages/79/37/482e55fa1602e0a7ff012661d8c946bafdc05e480ea5a32f4f7e336d4aa9/pytokens-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:b35d7e5ad269804f6697727702da3c517bb8a5228afa450ab0fa787732055fc9", size = 104539, upload-time = "2026-01-30T01:03:24.788Z" }, - { url = "https://files.pythonhosted.org/packages/30/e8/20e7db907c23f3d63b0be3b8a4fd1927f6da2395f5bcc7f72242bb963dfe/pytokens-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8fcb9ba3709ff77e77f1c7022ff11d13553f3c30299a9fe246a166903e9091eb", size = 168474, upload-time = "2026-01-30T01:03:26.428Z" }, - { url = "https://files.pythonhosted.org/packages/d6/81/88a95ee9fafdd8f5f3452107748fd04c24930d500b9aba9738f3ade642cc/pytokens-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79fc6b8699564e1f9b521582c35435f1bd32dd06822322ec44afdeba666d8cb3", size = 290473, upload-time = "2026-01-30T01:03:27.415Z" }, - { url = "https://files.pythonhosted.org/packages/cf/35/3aa899645e29b6375b4aed9f8d21df219e7c958c4c186b465e42ee0a06bf/pytokens-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d31b97b3de0f61571a124a00ffe9a81fb9939146c122c11060725bd5aea79975", size = 303485, upload-time = "2026-01-30T01:03:28.558Z" }, - { url = "https://files.pythonhosted.org/packages/52/a0/07907b6ff512674d9b201859f7d212298c44933633c946703a20c25e9d81/pytokens-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:967cf6e3fd4adf7de8fc73cd3043754ae79c36475c1c11d514fc72cf5490094a", size = 306698, upload-time = "2026-01-30T01:03:29.653Z" }, - { url = "https://files.pythonhosted.org/packages/39/2a/cbbf9250020a4a8dd53ba83a46c097b69e5eb49dd14e708f496f548c6612/pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918", size = 116287, upload-time = "2026-01-30T01:03:30.912Z" }, - { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" }, -] - [[package]] name = "pywavelets" version = "1.9.0" @@ -3678,6 +3602,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/b7/b95708304cd49b7b6f82fdd039f1748b66ec2b21d6a45180910802f1abf1/rpds_py-0.30.0-pp311-pypy311_pp73-musllinux_1_2_x86_64.whl", hash = "sha256:ac37f9f516c51e5753f27dfdef11a88330f04de2d564be3991384b2f3535d02e", size = 562191, upload-time = "2025-11-30T20:24:36.853Z" }, ] +[[package]] +name = "ruff" +version = "0.16.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/49/23802c45f093eb14bde54b141d2b2f058edfa63a06db7beed047308cc08f/ruff-0.16.10.tar.gz", hash = "sha256:eff4728c4eaae93f0955cd264d24b2ab348e74bf59986ccf282ba6dc16b3b017", size = 4958724, upload-time = "2026-10-01T18:03:21.697Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/21/ebce22e1d90cdb2cd691026b9c6e9bec6499f0b396089481755b5d49efff/ruff-0.16.10-py3-none-linux_armv6l.whl", hash = "sha256:488b0fe3f3574210e5cf80d9f59b9e3ab17a127a8155de3f307b392589cfb511", size = 10095557, upload-time = "2026-10-01T18:02:36.072Z" }, + { url = "https://files.pythonhosted.org/packages/cb/98/a54de85876a8b2612bfa0d84c7b9abfb39c6a3354aee7800b09c1649b9e3/ruff-0.16.10-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e748ff95c934c4e978783b8e687bc174e7bd84e8ad24e3243e1ecfcda5e0282d", size = 10340577, upload-time = "2026-10-01T18:02:39.258Z" }, + { url = "https://files.pythonhosted.org/packages/9f/16/1a5a4a2657effe4806110f2b907313802f1367fcdbb29e8122766e407fab/ruff-0.16.10-py3-none-macosx_11_0_arm64.whl", hash = "sha256:3031a4a2e8e7b8a46f70be45f198c35a11ece509a94b80334d8d397a33c67550", size = 9774282, upload-time = "2026-10-01T18:02:41.79Z" }, + { url = "https://files.pythonhosted.org/packages/6e/fb/470085af734da396e68cd80fb0a3e7c109a459716ae59588c0f6fab8a17d/ruff-0.16.10-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:494401c86df4c4c25f69b9419605d944467ee98c42fb6ad405ef4fa40b8fb67d", size = 9920895, upload-time = "2026-10-01T18:02:44.458Z" }, + { url = "https://files.pythonhosted.org/packages/57/de/f10cffe4f88a37ea6615ec460f01bf76f0bc1477737a35d9ee61a630a78b/ruff-0.16.10-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d203abc0ff2b773ee33d00ab8df0bb67046fbc7c07b119332f08b9b345cf8221", size = 9902707, upload-time = "2026-10-01T18:02:47.041Z" }, + { url = "https://files.pythonhosted.org/packages/ff/44/3fdcedf83ae60ef837dd239170e480dce606a9142cfa2db911afdf855fd9/ruff-0.16.10-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bd83d1235a5258d318477bc5b576303974cbdef5df0c01a1bff14efcc23bd12a", size = 10619287, upload-time = "2026-10-01T18:02:49.485Z" }, + { url = "https://files.pythonhosted.org/packages/c1/62/02e76a5574002153618eafbb70e159468addc72a5d2c488e65cbb4639d2d/ruff-0.16.10-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bc2610fb269fa56dd8a68669ae470fa6272902668c0fc2ebc3aa112b2633d5b8", size = 11339942, upload-time = "2026-10-01T18:02:52.008Z" }, + { url = "https://files.pythonhosted.org/packages/6b/c4/cde27d47ad8d4126c587e608c67c46e7feac11763f606d261a1bc8a489e6/ruff-0.16.10-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3e70175e29cc94c26ea296c80e470180b744b7419026898e58f520c6ab32578e", size = 10934316, upload-time = "2026-10-01T18:02:54.811Z" }, + { url = "https://files.pythonhosted.org/packages/e4/03/17234145f302a645a123e8c3bb2411ecf4669fbc350de3b0d803230a1729/ruff-0.16.10-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f33f43a864a8483eebd160e713336c8bab02c934feaff0a33cf5ccb41546d09a", size = 10387968, upload-time = "2026-10-01T18:02:57.497Z" }, + { url = "https://files.pythonhosted.org/packages/71/29/2493af60240ee7644b38a4b821f5f9c3b5a4fa3178770fe1d0c685217395/ruff-0.16.10-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:1dfc6f0088149fb6a362c1c446bcbb3fd2157b3852fe2fa68409276eab9ad9b3", size = 10537906, upload-time = "2026-10-01T18:03:00.006Z" }, + { url = "https://files.pythonhosted.org/packages/2e/9a/f56b28f3b143bb9e518c34fb89ab191b626bca8b70dbf8af0d2e2d473572/ruff-0.16.10-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6553498afc35f580f036030795810b9e6bcea31604b0fd9e8d352795473042e3", size = 10012938, upload-time = "2026-10-01T18:03:03.006Z" }, + { url = "https://files.pythonhosted.org/packages/c2/c4/fca37362848ea4d4d80712df13632e645e7c7cfb1bdedf140699ac7a090b/ruff-0.16.10-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:a3b8471dea115d37f123882be852bed13403746d3a76c11de4a19ec5f9ff5a03", size = 9897945, upload-time = "2026-10-01T18:03:05.818Z" }, + { url = "https://files.pythonhosted.org/packages/7d/c7/e0bc57664d6e0c61fd22f620af260fe9662d7e1ddb320ea7f160e76ef165/ruff-0.16.10-py3-none-musllinux_1_2_i686.whl", hash = "sha256:92e59a70bcbd9d3a5483656da906ec28edfdacfce00afd99edb8b4e9d15644be", size = 10333134, upload-time = "2026-10-01T18:03:08.25Z" }, + { url = "https://files.pythonhosted.org/packages/21/aa/5c9f3b68737c0e4a33a1db5dfab44f7b786d96ca91a7233112ddab1e6dc2/ruff-0.16.10-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7ae7375f803b5520dc9f546bed7e3a0acb70b91812e9bb4b19927de22f25b77d", size = 10741702, upload-time = "2026-10-01T18:03:10.775Z" }, + { url = "https://files.pythonhosted.org/packages/78/fa/0f9c2020dc316c53d983be011720b8157cc052b97e3f4dfa7db6d0f880a6/ruff-0.16.10-py3-none-win32.whl", hash = "sha256:2a12e01cb9156c10c466f63b46eaae5ecea28dfbd21b5836353ae498e7d1349a", size = 10139176, upload-time = "2026-10-01T18:03:13.267Z" }, + { url = "https://files.pythonhosted.org/packages/99/29/cfb0df9448d4d4ad48c2de029ada9ebd71baa6da983a6c77ee6c6cd0fe82/ruff-0.16.10-py3-none-win_amd64.whl", hash = "sha256:97f2015c92aa97105b0eab19eb5d224884399281cfc5da86a92db4ab5e7fb2ca", size = 10584734, upload-time = "2026-10-01T18:03:16.006Z" }, + { url = "https://files.pythonhosted.org/packages/fc/05/c16957eb287c3fc062e032619a25868d93d408a844b725bcf514f7a378ff/ruff-0.16.10-py3-none-win_arm64.whl", hash = "sha256:25a65fe998c4e6861ec079ada5826a2fc605e6cbccbe9dcd7fac1f54e791621b", size = 10366440, upload-time = "2026-10-01T18:03:19.04Z" }, +] + [[package]] name = "scikit-image" version = "0.26.0"