scitex 2.0.0__py2.py3-none-any.whl

scitex/dsp/_misc.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-05 01:03:32 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/dsp/_misc.py
+
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-04-05 12:14:08 (ywatanabe)"
+
+from ..decorators import torch_fn
+
+
+@torch_fn
+def ensure_3d(x):
+    if x.ndim == 1:  # assumes (seq_len,)
+        x = x.unsqueeze(0).unsqueeze(0)
+    elif x.ndim == 2:  # assumes (batch_siize, seq_len)
+        x = x.unsqueeze(1)
+    return x
+
+
+# @torch_fn
+# def unbias(x, dim=-1, fn="mean"):
+#     if fn == "mean":
+#         return x - x.mean(dim=dim, keepdims=True)
+#     if fn == "min":
+#         return x - x.min(dim=dim, keepdims=True)[0]
+
+
+# EOF

scitex/dsp/_mne.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-04 02:07:36 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/dsp/_mne.py
+
+import mne
+import pandas as pd
+from .params import EEG_MONTAGE_1020
+
+
+def get_eeg_pos(channel_names=EEG_MONTAGE_1020):
+    # Load the standard 10-20 montage
+    standard_montage = mne.channels.make_standard_montage("standard_1020")
+    standard_montage.ch_names = [
+        ch_name.upper() for ch_name in standard_montage.ch_names
+    ]
+
+    # Get the positions of the electrodes in the standard montage
+    positions = standard_montage.get_positions()
+
+    df = pd.DataFrame(positions["ch_pos"])[channel_names]
+
+    df.set_index(pd.Series(["x", "y", "z"]))
+
+    return df
+
+
+if __name__ == "__main__":
+    print(get_eeg_pos())
+
+
+# EOF

scitex/dsp/_modulation_index.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-04 02:09:55 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/dsp/_modulation_index.py
+
+import torch
+
+from ..decorators import signal_fn
+from ..nn._ModulationIndex import ModulationIndex
+
+
+@signal_fn
+def modulation_index(pha, amp, n_bins=18, amp_prob=False):
+    """
+    pha: (batch_size, n_chs, n_freqs_pha, n_segments, seq_len)
+    amp: (batch_size, n_chs, n_freqs_amp, n_segments, seq_len)
+    """
+    return ModulationIndex(n_bins=n_bins, amp_prob=amp_prob)(pha, amp)
+
+
+def _reshape(x, batch_size=2, n_chs=4):
+    return (
+        torch.tensor(x)
+        .float()
+        .unsqueeze(0)
+        .unsqueeze(0)
+        .repeat(batch_size, n_chs, 1, 1, 1)
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import matplotlib.pyplot as plt
+    import scitex
+
+    # Start
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt, fig_scale=3)
+
+    # Parameters
+    FS = 512
+    T_SEC = 5
+
+    # Demo signal
+    xx, tt, fs = scitex.dsp.demo_sig(fs=FS, t_sec=T_SEC, sig_type="tensorpac")
+    # xx.shape: (8, 19, 20, 512)
+
+    # Tensorpac
+    (
+        pha,
+        amp,
+        freqs_pha,
+        freqs_amp,
+        pac_tp,
+    ) = scitex.dsp.utils.pac.calc_pac_with_tensorpac(xx, fs, t_sec=T_SEC)
+
+    # GPU calculation with scitex.dsp.nn.ModulationIndex
+    pha, amp = _reshape(pha), _reshape(amp)
+    pac_scitex = scitex.dsp.modulation_index(pha, amp).cpu().numpy()
+    i_batch, i_ch = 0, 0
+    pac_scitex = pac_scitex[i_batch, i_ch]
+
+    # Plots
+    fig = scitex.dsp.utils.pac.plot_PAC_scitex_vs_tensorpac(
+        pac_scitex, pac_tp, freqs_pha, freqs_amp
+    )
+    fig.suptitle("MI (modulation index) calculation")
+    scitex.io.save(fig, "modulation_index.png")
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF
+
+"""
+/home/ywatanabe/proj/entrance/scitex/dsp/_modulation_index.py
+"""
+
+# EOF

scitex/dsp/_pac.py

@@ -0,0 +1,319 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-26 22:24:40 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/dsp/_pac.py
+
+THIS_FILE = "/home/ywatanabe/proj/scitex_repo/src/scitex/dsp/_pac.py"
+
+import sys
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from ..str import printc
+
+from ..decorators import signal_fn, batch_fn
+from ..nn._PAC import PAC
+
+"""
+scitex.dsp.pac function
+"""
+
+
+# @batch_fn
+@signal_fn
+def pac(
+    x,
+    fs,
+    pha_start_hz=2,
+    pha_end_hz=20,
+    pha_n_bands=100,
+    amp_start_hz=60,
+    amp_end_hz=160,
+    amp_n_bands=100,
+    device="cuda",
+    batch_size=1,
+    batch_size_ch=-1,
+    fp16=False,
+    trainable=False,
+    n_perm=None,
+    amp_prob=False,
+):
+    """
+    Compute the phase-amplitude coupling (PAC) for signals. This function automatically handles inputs as
+    PyTorch tensors, NumPy arrays, or pandas DataFrames.
+
+    Arguments:
+    - x (torch.Tensor | np.ndarray | pd.DataFrame): Input signal. Shape can be either (batch_size, n_chs, seq_len) or
+    - fs (float): Sampling frequency of the input signal.
+    - pha_start_hz (float, optional): Start frequency for phase bands. Default is 2 Hz.
+    - pha_end_hz (float, optional): End frequency for phase bands. Default is 20 Hz.
+    - pha_n_bands (int, optional): Number of phase bands. Default is 100.
+    - amp_start_hz (float, optional): Start frequency for amplitude bands. Default is 60 Hz.
+    - amp_end_hz (float, optional): End frequency for amplitude bands. Default is 160 Hz.
+    - amp_n_bands (int, optional): Number of amplitude bands. Default is 100.
+
+    Returns:
+    - torch.Tensor: PAC values. Shape: (batch_size, n_chs, pha_n_bands, amp_n_bands)
+    - numpy.ndarray: Phase bands used for the computation.
+    - numpy.ndarray: Amplitude bands used for the computation.
+
+    Example:
+        FS = 512
+        T_SEC = 4
+        xx, tt, fs = scitex.dsp.demo_sig(
+            batch_size=1, n_chs=1, fs=FS, t_sec=T_SEC, sig_type="tensorpac"
+        )
+        pac, pha_mids_hz, amp_mids_hz = scitex.dsp.pac(xx, fs)
+    """
+
+    def process_ch_batching(m, x, batch_size_ch, device):
+        n_chs = x.shape[1]
+        n_batches = (n_chs + batch_size_ch - 1) // batch_size_ch
+
+        agg = []
+        for ii in range(n_batches):
+            start, end = batch_size_ch * ii, min(batch_size_ch * (ii + 1), n_chs)
+            _pac = m(x[:, start:end, :].to(device)).detach().cpu()
+            agg.append(_pac)
+
+        # return np.concatenate(agg, axis=1)
+        return torch.cat(agg, dim=1)
+
+    m = PAC(
+        x.shape[-1],
+        fs,
+        pha_start_hz=pha_start_hz,
+        pha_end_hz=pha_end_hz,
+        pha_n_bands=pha_n_bands,
+        amp_start_hz=amp_start_hz,
+        amp_end_hz=amp_end_hz,
+        amp_n_bands=amp_n_bands,
+        fp16=fp16,
+        trainable=trainable,
+        n_perm=n_perm,
+        amp_prob=amp_prob,
+    ).to(device)
+
+    if batch_size_ch == -1:
+        return m(x.to(device)), m.PHA_MIDS_HZ, m.AMP_MIDS_HZ
+    else:
+        return (
+            process_ch_batching(m, x, batch_size_ch, device),
+            m.PHA_MIDS_HZ,
+            m.AMP_MIDS_HZ,
+        )
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+    import scitex
+
+    # Start
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt)
+
+    pac, freqs_pha, freqs_amp = scitex.dsp.pac(
+        np.random.rand(1, 16, 24000),
+        400,
+        batch_size=1,
+        batch_size_ch=8,
+        fp16=True,
+        n_perm=16,
+    )
+
+#     # Parameters
+#     FS = 512
+#     T_SEC = 4
+#     # IS_TRAINABLE = False
+#     # FP16 = True
+
+#     for IS_TRAINABLE in [True, False]:
+#         for FP16 in [True, False]:
+
+#             # Demo signal
+#             xx, tt, fs = scitex.dsp.demo_sig(
+#                 batch_size=1, n_chs=1, fs=FS, t_sec=T_SEC, sig_type="pac"
+#             )
+
+
+#             # scitex.str.print_debug()
+#             # xx = np.random.rand(1,16,24000)
+#             # fs = 400
+
+#             # scitex calculation
+#             pac_scitex, pha_mids_scitex, amp_mids_scitex = scitex.dsp.pac(
+#                 xx,
+#                 fs,
+#                 pha_n_bands=50,
+#                 amp_n_bands=30,
+#                 trainable=IS_TRAINABLE,
+#                 fp16=FP16,
+#             )
+#             i_batch, i_ch = 0, 0
+#             pac_scitex = pac_scitex[i_batch, i_ch]
+
+#             printc(type(pac_scitex))
+
+#             # Tensorpac calculation
+#             (
+#                 _,
+#                 _,
+#                 _pha_mids_tp,
+#                 _amp_mids_tp,
+#                 pac_tp,
+#             ) = scitex.dsp.utils.pac.calc_pac_with_tensorpac(xx, fs, T_SEC)
+
+#             # Validates the consitency in frequency definitions
+#             assert np.allclose(
+#                 pha_mids_scitex, _pha_mids_tp
+#             )
+#             assert np.allclose(
+#                 amp_mids_scitex, _amp_mids_tp
+#             )
+
+#             scitex.io.save(
+#                 (pac_scitex, pac_tp, pha_mids_scitex, amp_mids_scitex),
+#                 "./data/cache.npz",
+#             )
+
+#             # ################################################################################
+#             # # cache
+#             # pac_scitex, pac_tp, pha_mids_scitex, amp_mids_scitex = scitex.io.load(
+#             #     "./data/cache.npz"
+#             # )
+#             # ################################################################################
+
+#             # Plots
+#             fig = scitex.dsp.utils.pac.plot_PAC_scitex_vs_tensorpac(
+#                 pac_scitex, pac_tp, pha_mids_scitex, amp_mids_scitex
+#             )
+#             fig.suptitle(
+#                 "Phase-Amplitude Coupling calculation\n\n(Bandpass Filtering -> Hilbert Transformation-> Modulation Index)"
+#             )
+#             plt.show()
+
+#             scitex.gen.reload(scitex.dsp)
+
+#             # Saves the figure
+#             trainable_str = "trainable" if IS_TRAINABLE else "static"
+#             fp_str = "fp16" if FP16 else "fp32"
+#             scitex.io.save(
+#                 fig, f"pac_with_{trainable_str}_bandpass_{fp_str}.png"
+#             )
+
+
+# def run_method_tests():
+#     import scitex
+
+#     # Test parameters
+#     FS = 512
+#     T_SEC = 4
+
+#     class PACProcessor:
+#         @batch_torch_fn
+#         def process_pac(self, x, fs, **kwargs):
+#             return pac(x, fs, **kwargs)
+
+#         @signal_fn
+#         def process_signal(self, x):
+#             return x * 2
+
+#     def run_method_basic_tests():
+#         processor = PACProcessor()
+
+#         # Generate test signal
+#         xx, tt, fs = scitex.dsp.demo_sig(
+#             batch_size=1, n_chs=1, fs=FS, t_sec=T_SEC, sig_type="pac"
+#         )
+
+#         try:
+#             # Test method with batch processing
+#             result_batch, pha_mids, amp_mids = processor.process_pac(
+#                 xx, fs, pha_n_bands=50, amp_n_bands=30, batch_size=1
+#             )
+#             assert torch.is_tensor(result_batch)
+
+#             # Test basic torch method
+#             result_torch = processor.process_signal(xx)
+#             assert torch.is_tensor(result_torch)
+
+#             scitex.str.printc("Passed: Basic method tests", "yellow")
+#         except Exception as err:
+#             scitex.str.printc(f"Failed: Basic method tests - {str(err)}", "red")
+
+#     def run_method_cuda_tests():
+#         if not torch.cuda.is_available():
+#             scitex.str.printc(
+#                 "CUDA method tests skipped: No GPU available", "yellow"
+#             )
+#             return
+
+#         processor = PACProcessor()
+#         xx, tt, fs = scitex.dsp.demo_sig(
+#             batch_size=1, n_chs=1, fs=FS, t_sec=T_SEC, sig_type="pac"
+#         )
+
+#         try:
+#             # Test with CUDA
+#             result_cuda, _, _ = processor.process_pac(xx, fs, device="cuda")
+#             assert result_cuda.device.type == "cuda"
+
+#             result_torch = processor.process_signal(xx, device="cuda")
+#             assert result_torch.device.type == "cuda"
+
+#             scitex.str.printc("Passed: CUDA method tests", "yellow")
+#         except Exception as err:
+#             scitex.str.printc(f"Failed: CUDA method tests - {str(err)}", "red")
+
+#     def run_method_batch_size_tests():
+#         processor = PACProcessor()
+#         batch_sizes = [1, 2, 4]
+
+#         for batch_size in batch_sizes:
+#             try:
+#                 xx, tt, fs = scitex.dsp.demo_sig(
+#                     batch_size=batch_size,
+#                     n_chs=1,
+#                     fs=FS,
+#                     t_sec=T_SEC,
+#                     sig_type="pac",
+#                 )
+
+#                 result, _, _ = processor.process_pac(
+#                     xx, fs, batch_size=batch_size
+#                 )
+#                 assert result.shape[0] == batch_size
+
+#                 scitex.str.printc(
+#                     f"Passed: Method batch size test with size={batch_size}",
+#                     "yellow",
+#                 )
+#             except Exception as err:
+#                 scitex.str.printc(
+#                     f"Failed: Method batch size test with size={batch_size} - {str(err)}",
+#                     "red",
+#                 )
+
+#     # Execute method test suites
+#     test_suites = [
+#         ("Method Basic Tests", run_method_basic_tests),
+#         ("Method CUDA Tests", run_method_cuda_tests),
+#         ("Method Batch Size Tests", run_method_batch_size_tests),
+#     ]
+
+#     for test_name, test_func in test_suites:
+#         test_func()
+
+
+# if __name__ == "__main__":
+#     run_method_tests()
+
+# # EOF
+
+# """
+# python -m scitex.dsp._pac
+# """
+
+#
+
+# EOF

scitex/dsp/_psd.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-04 02:11:25 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/dsp/_psd.py
+
+"""This script does XYZ."""
+
+import torch
+
+from ..decorators import signal_fn
+from ..nn._PSD import PSD
+
+
+@signal_fn
+def psd(
+    x,
+    fs,
+    prob=False,
+    dim=-1,
+):
+    """
+    import matplotlib.pyplot as plt
+
+    x, t, fs = scitex.dsp.demo_sig()  # (batch_size, n_chs, seq_len)
+    pp, ff = psd(x, fs)
+
+    # Plots
+    plt, CC = scitex.plt.configure_mpl(plt)
+    fig, ax = scitex.plt.subplots()
+    ax.plot(fs, pp[0, 0])
+    ax.xlabel("Frequency [Hz]")
+    ax.ylabel("log(Power [uV^2 / Hz]) [a.u.]")
+    plt.show()
+    """
+    psd, freqs = PSD(fs, prob=prob, dim=dim)(x)
+    return psd, freqs
+
+
+def band_powers(self, psd):
+    """
+    Calculate the average power for specified frequency bands.
+    """
+    assert len(self.low_freqs) == len(self.high_freqs)
+
+    out = []
+    for ll, hh in zip(self.low_freqs, self.high_freqs):
+        band_indices = torch.where((freqs >= ll) & (freqs <= hh))[0].to(psd.device)
+        band_power = psd[..., band_indices].sum(dim=self.dim)
+        bandwidth = hh - ll
+        avg_band_power = band_power / bandwidth
+        out.append(avg_band_power)
+    out = torch.stack(out, dim=-1)
+    return out
+
+    # Average Power in Each Frequency Band
+    avg_band_powers = self.calc_band_avg_power(psd, freqs)
+    return (avg_band_powers,)
+
+
+if __name__ == "__main__":
+    import sys
+    import scitex
+
+    import matplotlib.pyplot as plt
+
+    # Start
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt)
+
+    # Parameters
+    SIG_TYPE = "chirp"
+
+    # Demo signal
+    xx, tt, fs = scitex.dsp.demo_sig(SIG_TYPE)  # (8, 19, 384)
+
+    # PSD calculation
+    pp, ff = psd(xx, fs, prob=True)
+
+    # Plots
+    fig, axes = scitex.plt.subplots(nrows=2)
+
+    axes[0].plot(tt, xx[0, 0], label=SIG_TYPE)
+    axes[1].set_title("Signal")
+    axes[0].set_xlabel("Time [s]")
+    axes[0].set_ylabel("Amplitude [?V]")
+
+    axes[1].plot(ff, pp[0, 0])
+    axes[1].set_title("PSD (power spectrum density)")
+    axes[1].set_xlabel("Frequency [Hz]")
+    axes[1].set_ylabel("Log(Power [?V^2 / Hz]) [a.u.]")
+
+    scitex.io.save(fig, "psd.png")
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF
+
+"""
+/home/ywatanabe/proj/entrance/scitex/dsp/_psd.py
+"""
+
+# EOF

scitex/dsp/_resample.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-04-13 02:35:11 (ywatanabe)"
+
+
+import torch
+import torchaudio.transforms as T
+from ..decorators import signal_fn
+
+
+@signal_fn
+def resample(x, src_fs, tgt_fs, t=None):
+    xr = T.Resample(src_fs, tgt_fs, dtype=x.dtype).to(x.device)(x)
+    if t is None:
+        return xr
+    if t is not None:
+        tr = torch.linspace(t[0], t[-1], xr.shape[-1])
+        return xr, tr
+
+
+if __name__ == "__main__":
+    import sys
+
+    import matplotlib.pyplot as plt
+
+    # Start
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt)
+
+    # Parameters
+    T_SEC = 1
+    SIG_TYPE = "chirp"
+    SRC_FS = 128
+    TGT_FS_UP = 256
+    TGT_FS_DOWN = 64
+    FREQS_HZ = [10, 30, 100, 300]
+
+    # Demo Signal
+    xx, tt, fs = scitex.dsp.demo_sig(
+        t_sec=T_SEC, fs=SRC_FS, freqs_hz=FREQS_HZ, sig_type=SIG_TYPE
+    )
+
+    # Resampling
+    xd, td = scitex.dsp.resample(xx, fs, TGT_FS_DOWN, t=tt)
+    xu, tu = scitex.dsp.resample(xx, fs, TGT_FS_UP, t=tt)
+
+    # Plots
+    i_batch, i_ch = 0, 0
+    fig, axes = plt.subplots(nrows=3, sharex=True, sharey=True)
+    axes[0].plot(tt, xx[i_batch, i_ch], label=f"Original ({SRC_FS} Hz)")
+    axes[1].plot(td, xd[i_batch, i_ch], label=f"Down-sampled ({TGT_FS_DOWN} Hz)")
+    axes[2].plot(tu, xu[i_batch, i_ch], label=f"Up-sampled ({TGT_FS_UP} Hz)")
+    for ax in axes:
+        ax.legend(loc="upper left")
+
+    axes[-1].set_xlabel("Time [s]")
+    fig.supylabel("Amplitude [?V]")
+    fig.suptitle("Resampling")
+    scitex.io.save(fig, "traces.png")
+    # plt.show()
+
+# EOF
+
+"""
+/home/ywatanabe/proj/entrance/scitex/dsp/_resample.py
+"""

scitex/dsp/_time.py

@@ -0,0 +1,36 @@
+#!./env/bin/python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-06-30 12:11:01 (ywatanabe)"
+# /mnt/ssd/ripple-wm-code/scripts/externals/scitex/src/scitex/dsp/_time.py
+
+
+import numpy as np
+
+
+def time(start_sec, end_sec, fs):
+    # return np.linspace(start_sec, end_sec, (end_sec - start_sec) * fs)
+    return scitex.gen.float_linspace(start_sec, end_sec, (end_sec - start_sec) * fs)
+
+
+def main():
+    out = time(10, 15, 256)
+    print(out)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import matplotlib.pyplot as plt
+
+    # # Argument Parser
+    # import argparse
+    # parser = argparse.ArgumentParser(description='')
+    # parser.add_argument('--var', '-v', type=int, default=1, help='')
+    # parser.add_argument('--flag', '-f', action='store_true', default=False, help='')
+    # args = parser.parse_args()
+    # Main
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt, verbose=False)
+    main()
+    scitex.gen.close(CONFIG, verbose=False, notify=False)
+
+# EOF