pylottone 0.2.2__py3-none-any.whl

pylottone/__init__.py

@@ -0,0 +1,58 @@
+from importlib import import_module
+
+_PT_EXPORTS = (
+    "est_dtft",
+    "dtft_sum",
+    "extract_raw_pt",
+    "sniffer_sub",
+    "plot_multich_comparison",
+    "pickcoilsbycorr",
+    "check_waveform_polarity",
+    "extract_pilottone_navs",
+    "calibrate_pt",
+    "apply_pt_calib",
+    "process_cplx_pt",
+    "pick_cardiac_source",
+    "pick_source_bypeak",
+    "pick_navigators_from_sources",
+    "pred_scan",
+)
+
+_SELFNAV_EXPORTS = (
+    "extract_selfnav_navs",
+)
+
+_TRIGGERING_EXPORTS = (
+    "repair_cardiac_triggers_rr",
+    "repair_ecg_triggers_with_pt",
+)
+
+__all__ = [*_PT_EXPORTS, *_SELFNAV_EXPORTS, *_TRIGGERING_EXPORTS, "main"]
+
+
+def __getattr__(name: str):
+    if name in _PT_EXPORTS:
+        module_name = ".model_selection" if name in {"pick_navigators_from_sources", "pred_scan"} else ".pt"
+        module = import_module(module_name, __name__)
+        value = getattr(module, name)
+        globals()[name] = value
+        return value
+    if name in _SELFNAV_EXPORTS:
+        module = import_module(".selfnav", __name__)
+        value = getattr(module, name)
+        globals()[name] = value
+        return value
+    if name in _TRIGGERING_EXPORTS:
+        module = import_module(".triggering", __name__)
+        value = getattr(module, name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__))
+
+
+def main() -> None:
+    print("Hello from pylottone!")

pylottone/constants.py

@@ -0,0 +1,8 @@
+PILOTTONE_WAVEFORM_ID = 1025
+PILOTTONE_CH = {'RESP': 0, 'CARDIAC': 1, 'CARDIAC_TRIGGERS': 2, 'CARDIAC_DERIVATIVE': 3, 'DERIVATIVE_TRIGGERS': 4}
+ECG_WAVEFORM_ID = 0
+PULSEOX_WAVEFORM_ID = 1
+RESP_WAVEFORM_ID = 2
+EXT1_WAVEFORM_ID = 3
+EXT2_WAVEFORM_ID = 4
+RESPPT_WAVEFORM_ID = 12

pylottone/editer.py

@@ -0,0 +1,252 @@
+from warnings import warn
+import numpy as np
+import numpy.typing as npt
+from numpy.fft import ifft, ifftshift
+import math
+
+try:
+    import cupy as cp
+except Exception as exc:
+    cp = None
+    _cupy_import_error = exc
+    warn(f'CuPy is unavailable; EDITER will use the CPU path. ({exc!r})')
+
+import scipy as sp
+
+
+def _cupy_ready() -> bool:
+    return cp is not None
+
+
+def _to_numpy(array):
+    if cp is not None and hasattr(array, 'get'):
+        return array.get()
+    return array
+
+def to_hybrid_kspace(indata):
+    '''Centered ifft on first dimension. Does not do fftshift before ifft, as it treats data as time signal.'''
+    return ifftshift(ifft(indata, None, axis=0), axes=0)
+
+def autopick_sensing_coils(data, f_emi, bw_emi, bw_sig, f_samp, ratio_th=None, n_sensing=None):
+
+    # Ensure CPU array for operations (Numba or SciPy paths expect numpy arrays)
+    if cp is not None and hasattr(data, '__cuda_array_interface__'):
+        data = _to_numpy(data)
+
+    n_samp = data.shape[0]
+    df = f_samp/n_samp
+    freq_axis = np.arange(0, f_samp, df) - (f_samp - (n_samp % 2)*df)/2 # Handles both even and odd length signals.
+    signal_mask = (freq_axis < bw_sig/2) & (freq_axis > -bw_sig/2)
+    signal_region = to_hybrid_kspace(data[:, 0,:].squeeze())*signal_mask[:,None]
+
+    emi_mask = (freq_axis < (f_emi+bw_emi/2)) & (freq_axis > (f_emi-bw_emi/2))
+    emi_region = to_hybrid_kspace(data[:, 0,:].squeeze())*emi_mask[:,None]
+
+    emi_energy = np.sum(np.abs(emi_region), axis=0)
+    signal_energy = np.sum(np.abs(signal_region), axis=0)
+    sratio = signal_energy/emi_energy
+
+    Isig = np.argsort(signal_energy)
+    Irat = np.argsort(sratio)
+
+    sratio /= sratio[Isig[-1]] # Normalize with the largest Signal coil to reduce dep on PT amplitude.
+
+    # plt.figure()
+    # plt.plot(np.abs(to_hybrid_kspace(data[:,0,:].squeeze())))
+
+    # fig, ax1 = plt.subplots()
+    # ax2 = ax1.twinx()
+    # ax1.plot(emi_energy, label='emi_energy')
+    # ax1.plot(signal_energy, label='signal_energy')
+    # ax1.legend()
+    # ax2.plot(sratio, 'green', label='ratio')
+    # ax2.legend()
+    # plt.show()
+    # plt.figure()
+    # plt.plot(np.abs(signal_region))
+    # plt.show()
+
+    # plt.figure()
+    # plt.plot(np.abs(emi_region))
+    # plt.show()
+
+    # print(f'Coil with highest signal {coil_name[Isig[-1]]}.')
+
+    # print(coil_name[Isig])
+
+    # print(coil_name[Irat])
+
+    # print(f'Coils with smallest ratio {coil_name[sratio < ratio_th]}')
+
+    n_ch = data.shape[2]
+    if n_sensing is not None and ratio_th is not None:
+        warn('Both n_sensing and ratio_th are set. Using n_sensing.')
+    if n_sensing is not None:
+        sensing_coils = Irat[:n_sensing]
+    elif ratio_th is not None:
+        sensing_coils = np.nonzero(sratio < ratio_th)[0]
+    else:
+        raise ValueError('Either n_sensing or ratio_th must be set.')
+
+    mri_coils = np.arange(n_ch)
+    mri_coils = mri_coils[~np.isin(mri_coils, sensing_coils)]
+    return mri_coils, sensing_coils
+
+
+
+def get_noise_mtx(
+    line_grp: npt.NDArray[np.complex64],
+    dk: list[int],
+    xp=None,
+):
+    """
+    Creates the shifted noise matrix for a given line group and kernel sizes.
+
+    Args:
+        line_grp (numpy.ndarray): Line group data of shape (Nsamples, Nlines, Nchannels).
+        dk (list or tuple): Kernel size [d_kx, d_ky].
+
+    Returns:
+        numpy.ndarray: Noise matrix of shape ((Nsamples * Nlines) x (Nchannels * (d_kx * 2 + 1) * (d_ky * 2 + 1))).
+    """
+    if xp is None:
+        if cp is not None and hasattr(line_grp, '__cuda_array_interface__'):
+            xp = cp.get_array_module(line_grp)
+        else:
+            xp = np
+    if cp is None and xp is not np:
+        xp = np
+    d_kx = dk[0]
+    d_ky = dk[1]
+
+    # noise_mat = np.zeros((line_grp.shape[0], line_grp.shape[1], n_ch*(d_kx*2+1)*(d_ky*2+1)), dtype=line_grp.dtype)
+    noise_mat = []
+
+    dfp = xp.pad(line_grp, ((d_kx, d_kx), (d_ky, d_ky), (0, 0)), mode='constant')
+    if d_ky == 0:
+        end_slc = None
+    else:
+        end_slc = d_ky
+    ii = 0
+    for col_shift in range(-d_kx, d_kx + 1):
+        for lin_shift in range(-d_ky, d_ky + 1):
+            dftmp = xp.roll(dfp, shift=(col_shift, lin_shift), axis=(0, 1))
+            cropped = dftmp[d_kx:-d_kx, d_ky:end_slc, :]
+            # noise_mat[:,:,(ii*n_ch):((ii+1)*n_ch)] = dftmp[d_kx:-d_kx, d_ky:end_slc, :]
+            ii += 1
+
+            noise_mat.append(cropped)
+
+    noise_mat = xp.concatenate(noise_mat, axis=2)
+    noise_mat = noise_mat.reshape(noise_mat.shape[0] * noise_mat.shape[1], -1)
+
+    return noise_mat
+
+
+def _denoise_sniffer_window(
+    snf,
+    denoise_rank: int,
+    xp,
+):
+    """Low-rank denoise of a sniffer window along coil dimension.
+
+    Args:
+        snf: Array of shape (Ncol, Nlin_window, Nsniffer_coils).
+        denoise_rank: Number of singular values/components to keep.
+        xp: Backend module (numpy or cupy).
+    """
+    if denoise_rank <= 0:
+        return snf
+
+    n_col, n_lin, n_ch = snf.shape
+    snf_mat = snf.reshape(n_col * n_lin, n_ch)
+    min_dim = min(snf_mat.shape)
+    keep_rank = min(denoise_rank, min_dim)
+
+    if keep_rank >= min_dim:
+        return snf
+
+    u, s, vh = xp.linalg.svd(snf_mat, full_matrices=False)
+    snf_denoised = (u[:, :keep_rank] * s[:keep_rank]) @ vh[:keep_rank, :]
+    return snf_denoised.reshape(n_col, n_lin, n_ch)
+
+
+def _est_emi_impl(signal_in, sniffer, line_grps, dk, w, denoise_sniffers, denoise_rank, xp, to_numpy):
+    Ncol, Nlin, Nc = sniffer.shape
+    emi_hat = xp.zeros((Ncol, Nlin), dtype=np.complex64)
+    kern = []
+
+    for cwin, pe_rng in enumerate(line_grps):
+        # Denoise sniffer channels per window before constructing the noise matrix.
+        snf = xp.asarray(sniffer[:, pe_rng, :])
+        if denoise_sniffers:
+            snf = _denoise_sniffer_window(snf, denoise_rank=denoise_rank, xp=xp)
+
+        # Build noise matrix and input vectors in the chosen backend
+        noise_mat = get_noise_mtx(snf, dk, xp=xp)
+
+        init_mat_sub = xp.reshape(xp.asarray(signal_in[:, pe_rng]), (Ncol * len(pe_rng), 1))
+        ww = xp.reshape(xp.asarray(w[:, pe_rng]), (Ncol * len(pe_rng), 1))
+
+        # Use SciPy's lstsq on CPU for consistency with previous implementation,
+        # and CuPy's lstsq on GPU when available.
+        if xp is np:
+            kern_ ,_,_,_ = sp.linalg.lstsq(ww * noise_mat, ww * init_mat_sub, cond=None, check_finite=False)
+        else:
+            kern_ ,_,_,_ = xp.linalg.lstsq(ww * noise_mat, ww * init_mat_sub, rcond=None)
+
+        kern.append((pe_rng, to_numpy(kern_)))
+        emi_hat[:, pe_rng] = xp.reshape(xp.dot(noise_mat, kern_), (Ncol, len(pe_rng)))
+
+    return to_numpy(emi_hat), kern
+
+def est_emi(
+    signal_in: npt.NDArray[np.complex64],
+    sniffer: npt.NDArray[np.complex64],
+    line_grps: list[npt.NDArray],
+    dk: list[int],
+    w: npt.NDArray[np.float32],
+    denoise_sniffers: bool = False,
+    denoise_rank: int = 1,
+):
+    return _est_emi_impl(signal_in, sniffer, line_grps, dk, w, denoise_sniffers, denoise_rank, np, lambda x: x)
+
+def est_emi_gpu(
+    signal_in: npt.NDArray[np.complex64],
+    sniffer: npt.NDArray[np.complex64],
+    line_grps: list[npt.NDArray],
+    dk: list[int],
+    w: npt.NDArray[np.float32],
+    denoise_sniffers: bool = False,
+    denoise_rank: int = 1,
+):
+    if cp is None:
+        warn('CuPy is unavailable; falling back to the CPU EDITER path.')
+        return est_emi(signal_in, sniffer, line_grps, dk, w, denoise_sniffers, denoise_rank)
+    return _est_emi_impl(signal_in, sniffer, line_grps, dk, w, denoise_sniffers, denoise_rank, cp, _to_numpy)
+
+def apply_editer(signal_in: npt.NDArray[np.complex64], sniffer: npt.NDArray[np.complex64], params, w) -> tuple[npt.NDArray[np.complex64], npt.NDArray[np.complex64]]:
+    max_lines = params['max_lines_per_group']
+    nlin = signal_in.shape[1]
+    if params['grouping_method'] == "uniform":
+        Ngrp = math.ceil(nlin/max_lines)
+        line_grps = []
+
+        for grp_i in range(Ngrp):
+            line_grps.append(np.arange(((grp_i)*max_lines), min(max_lines*(grp_i+1), nlin)))
+
+
+    denoise_sniffers = bool(params.get('denoise_sniffers', False))
+    denoise_rank = int(params.get('denoise_rank', 1))
+
+    if params['gpu'] == -1 or not _cupy_ready():
+        emi_hat, kernels = est_emi(signal_in, sniffer, line_grps, params['dk'], w, denoise_sniffers, denoise_rank)
+    else:
+        try:
+            with cp.cuda.Device(params['gpu']):
+                emi_hat, kernels = est_emi_gpu(signal_in, sniffer, line_grps, params['dk'], w, denoise_sniffers, denoise_rank)
+        except Exception as exc:
+            warn(f'CuPy/CUDA EDITER execution failed; falling back to CPU. ({exc!r})')
+            emi_hat, kernels = est_emi(signal_in, sniffer, line_grps, params['dk'], w, denoise_sniffers, denoise_rank)
+    return emi_hat, kernels
+    

pylottone/model_selection.py

@@ -0,0 +1,94 @@
+import joblib
+import numpy as np
+from scipy.signal import firwin, convolve
+from scipy.interpolate import pchip_interpolate
+import warnings
+
+def pred_scan(rocket_pipeline, scan, force_navpred=False):
+    ''' Predict navigators in the input scan using a pre-trained ROCKET classifier, given a set of sources.
+    A label of 0 indicates non-navigator, 1 indicates respiratory navigator, and 2 indicates cardiac navigator.
+    Parameters
+    ----------
+    rocket_pipeline : sklearn.pipeline.Pipeline
+        Pre-trained ROCKET classifier pipeline.
+    scan : np.ndarray
+        Input scan sources, shape (n_sources, n_samples)
+    force_navpred : bool
+        If True, forces the function to return a respiratory and cardiac navigator even if the classifier does not predict any.
+    Returns
+    -------
+    y_pred_ : np.ndarray
+        Predicted labels for each source, shape (n_sources,).
+    confs_ : np.ndarray
+        Confidence scores for each source, shape (n_sources, n_classes).
+    '''
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore") # Workaround, ignore sklearn tag warnings.
+        confs_ = rocket_pipeline.decision_function(scan[:,None,:])
+
+    y_pred_ = np.argmax(confs_, axis=1)
+
+    # Resolve multiple positive predictions
+    if np.sum(y_pred_ == 1) > 1:
+        resp_preds = (y_pred_ == 1).nonzero()[0]
+        conf_r_ = confs_[resp_preds, 1]
+        top_resp_idx = resp_preds[np.argmax(conf_r_)]
+        other_idxs = np.setdiff1d(resp_preds, np.array([top_resp_idx]))
+        y_pred_[other_idxs] = 0
+        warnings.warn(f"Multiple respiratory predictions found at indices {resp_preds}, keeping index {top_resp_idx} only.")
+    if np.sum(y_pred_ == 2) > 1:
+        card_preds = (y_pred_ == 2).nonzero()[0]
+        conf_r_ = confs_[card_preds, 2]
+        top_card_idx = card_preds[np.argmax(conf_r_)]
+        other_idxs = np.setdiff1d(card_preds, np.array([top_card_idx]))
+        y_pred_[other_idxs] = 0
+        warnings.warn(f"Multiple cardiac predictions found at indices {card_preds}, keeping index {top_card_idx} only.")
+    if force_navpred:
+        if np.sum(y_pred_ == 1) == 0:
+            y_pred_[np.argmax(confs_[:,1])] = 1
+            warnings.warn("No respiratory prediction was found, but force_navpred is True. Forcing the highest confidence prediction as respiratory.")
+        if np.sum(y_pred_ == 2) == 0:
+            y_pred_[np.argmax(confs_[:,2])] = 2
+            warnings.warn("No cardiac prediction was found, but force_navpred is True. Forcing the highest confidence prediction as cardiac.")
+    return y_pred_, confs_
+
+
+def pick_navigators_from_sources(sources, time_vec, classifier_path='rocket_pipeline.pkl', force_navpred=False):
+    ''' Pick respiratory and cardiac navigators from input sources using a pre-trained ROCKET classifier.
+    Parameters
+    ----------
+    sources : np.ndarray
+        Input sources, shape (n_sources, n_samples)
+    time_vec : np.ndarray
+        Time vector corresponding to the sources, unit is seconds, shape (n_samples,)
+    classifier_path : str
+        Path to the pre-trained ROCKET classifier. Must be compatible with joblib.load().
+    force_navpred : bool
+        If True, forces the function to return a respiratory and cardiac navigator even if the classifier does not predict any.
+    Returns
+    -------
+    resp_idx : int
+        Index of the respiratory navigator source in the input sources.
+    card_idx : int
+        Index of the cardiac navigator source in the input sources.
+    confs : np.ndarray
+        Confidence scores for each source, shape (n_sources, n_classes).
+    '''
+    rocket_pipeline = joblib.load(classifier_path)
+    n_samp = sources.shape[1]
+    dt_samp = time_vec[1] - time_vec[0]
+
+    h_denoise = firwin(2*(n_samp//8)-1, [0.1, 6], fs=1/dt_samp, window=('tukey', 1), pass_zero=False)
+    sources_filt = convolve(sources, h_denoise[None, :], mode='same')
+
+    dt_new = 10e-3  # 10 ms
+    n_samp_new = int(np.ceil(n_samp * dt_samp / dt_new))
+    time_new = np.arange(0, n_samp_new)*dt_new
+    sources_resampled = pchip_interpolate(time_vec, sources_filt, time_new, axis=1)
+    sources_resampled -= np.mean(sources_resampled, axis=1, keepdims=True)
+    sources_resampled /= np.std(sources_resampled, axis=1, keepdims=True)
+
+    y_pred, confs = pred_scan(rocket_pipeline, sources_resampled, force_navpred=force_navpred)
+    resp_idx = np.where(y_pred == 1)[0]
+    card_idx = np.where(y_pred == 2)[0]
+    return resp_idx, card_idx, confs