scitex 2.0.0__py2.py3-none-any.whl

scitex/nn/_Spectrogram.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-04-02 09:21:12 (ywatanabe)"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..decorators import numpy_fn, torch_fn
+
+
+class Spectrogram(nn.Module):
+    def __init__(
+        self,
+        sampling_rate,
+        n_fft=256,
+        hop_length=None,
+        win_length=None,
+        window="hann",
+    ):
+        super().__init__()
+        self.sampling_rate = sampling_rate
+        self.n_fft = n_fft
+        self.hop_length = hop_length if hop_length is not None else n_fft // 4
+        self.win_length = win_length if win_length is not None else n_fft
+        if window == "hann":
+            self.window = torch.hann_window(window_length=self.win_length)
+        else:
+            raise ValueError(
+                "Unsupported window type. Extend this to support more window types."
+            )
+
+    def forward(self, x):
+        """
+        Computes the spectrogram for each channel in the input signal.
+
+        Parameters:
+        - signal (torch.Tensor): Input signal of shape (batch_size, n_chs, seq_len).
+
+        Returns:
+        - spectrograms (torch.Tensor): The computed spectrograms for each channel.
+        """
+
+        x = scitex.dsp.ensure_3d(x)
+
+        batch_size, n_chs, seq_len = x.shape
+        spectrograms = []
+
+        for ch in range(n_chs):
+            x_ch = x[:, ch, :].unsqueeze(1)  # Maintain expected input shape for stft
+            spec = torch.stft(
+                x_ch.squeeze(1),
+                n_fft=self.n_fft,
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                window=self.window.to(x.device),
+                center=True,
+                pad_mode="reflect",
+                normalized=False,
+                return_complex=True,
+            )
+            magnitude = torch.abs(spec).unsqueeze(1)  # Keep channel dimension
+            spectrograms.append(magnitude)
+
+        # Concatenate spectrograms along channel dimension
+        spectrograms = torch.cat(spectrograms, dim=1)
+
+        # Calculate frequencies (y-axis)
+        freqs = torch.linspace(0, self.sampling_rate / 2, steps=self.n_fft // 2 + 1)
+
+        # Calculate times (x-axis)
+        # The number of frames can be computed from the size of the last dimension of the spectrogram
+        n_frames = spectrograms.shape[-1]
+        # Time of each frame in seconds, considering the hop length and sampling rate
+        times_sec = torch.arange(0, n_frames) * (self.hop_length / self.sampling_rate)
+
+        return spectrograms, freqs, times_sec
+
+
+@torch_fn
+def spectrograms(x, fs, cuda=False):
+    return Spectrogram(fs)(x)
+
+
+@torch_fn
+def my_softmax(x, dim=-1):
+    return F.softmax(x, dim=dim)
+
+
+@torch_fn
+def unbias(x, func="min", dim=-1, cuda=False):
+    if func == "min":
+        return x - x.min(dim=dim, keepdims=True)[0]
+    if func == "mean":
+        return x - x.mean(dim=dim, keepdims=True)[0]
+
+
+@torch_fn
+def normalize(x, axis=-1, amp=1.0, cuda=False):
+    high = torch.abs(x.max(axis=axis, keepdims=True)[0])
+    low = torch.abs(x.min(axis=axis, keepdims=True)[0])
+    return amp * x / torch.maximum(high, low)
+
+
+@torch_fn
+def spectrograms(x, fs, dj=0.125, cuda=False):
+    from wavelets_pytorch.transform import (
+        WaveletTransformTorch,
+    )  # PyTorch version
+
+    dt = 1 / fs
+    # dj = 0.125
+    batch_size, n_chs, seq_len = x.shape
+
+    x = x.cpu().numpy()
+
+    # # Batch of signals to process
+    # batch = np.array([batch_size * seq_len])
+
+    # Initialize wavelet filter banks (scipy and torch implementation)
+    # wa_scipy = WaveletTransform(dt, dj)
+    wa_torch = WaveletTransformTorch(dt, dj, cuda=True)
+
+    # Performing wavelet transform (and compute scalogram)
+    # cwt_scipy = wa_scipy.cwt(batch)
+    x = x[:, 0][:, np.newaxis]
+    cwt_torch = wa_torch.cwt(x)
+
+    return cwt_torch
+
+
+if __name__ == "__main__":
+    import scitex
+    import seaborn as sns
+    import torchaudio
+
+    fs = 1024  # 128
+    t_sec = 10
+    x = scitex.dsp.np.demo_sig(t_sec=t_sec, fs=fs, type="ripple")
+
+    normalize(unbias(x, cuda=True), cuda=True)
+
+    # My implementtion
+    ss = spectrograms(x, fs, cuda=True)
+    fig, axes = plt.subplots(nrows=2)
+    axes[0].plot(np.arange(x[0, 0]) / fs, x[0, 0])
+    sns.heatmap(ss[0], ax=axes[1])
+    plt.show()
+
+    ss, ff, tt = spectrograms(x, fs, cuda=True)
+    fig, axes = plt.subplots(nrows=2)
+    axes[0].plot(np.arange(x[0, 0]) / fs, x[0, 0])
+    sns.heatmap(ss[0], ax=axes[1])
+    plt.show()
+
+    # Torch Audio
+    transform = torchaudio.transforms.Spectrogram(n_fft=16, normalized=True).cuda()
+    xx = torch.tensor(x).float().cuda()[0, 0]
+    ss = transform(xx)
+    sns.heatmap(ss.detach().cpu().numpy())
+
+    plt.show()

scitex/nn/_SwapChannels.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2023-05-04 21:21:19 (ywatanabe)"
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+import scitex
+import numpy as np
+import random
+
+
+class SwapChannels(nn.Module):
+    def __init__(self, dropout=0.5):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        """x: [batch_size, n_chs, seq_len]"""
+        if self.training:
+            orig_chs = torch.arange(x.shape[1])
+
+            indi_orig = self.dropout(torch.ones(x.shape[1])).bool()
+            chs_to_shuffle = orig_chs[~indi_orig]
+
+            rand_chs = random.sample(
+                list(np.array(chs_to_shuffle)), len(chs_to_shuffle)
+            )
+
+            swapped_chs = orig_chs.clone()
+            swapped_chs[~indi_orig] = torch.LongTensor(rand_chs)
+
+            x = x[:, swapped_chs.long(), :]
+
+        return x
+
+
+if __name__ == "__main__":
+    ## Demo data
+    bs, n_chs, seq_len = 16, 360, 1000
+    x = torch.rand(bs, n_chs, seq_len)
+
+    sc = SwapChannels()
+    print(sc(x).shape)  # [16, 19, 1000]
+
+    # sb = SubjectBlock(n_chs=n_chs)
+    # print(sb(x, s).shape) # [16, 270, 1000]
+
+    # summary(sb, x, s)

scitex/nn/_TransposeLayer.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-03-30 07:26:35 (ywatanabe)"
+
+import torch.nn as nn
+
+
+class TransposeLayer(nn.Module):
+    def __init__(
+        self,
+        axis1,
+        axis2,
+    ):
+        super().__init__()
+        self.axis1 = axis1
+        self.axis2 = axis2
+
+    def forward(self, x):
+        return x.transpose(self.axis1, self.axis2)

scitex/nn/_Wavelet.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-03 07:17:26 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/nn/_Wavelet.py
+
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-05-30 11:04:45 (ywatanabe)"
+
+
+import scitex
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..gen._to_even import to_even
+from ..gen._to_odd import to_odd
+
+
+class Wavelet(nn.Module):
+    def __init__(
+        self, samp_rate, kernel_size=None, freq_scale="linear", out_scale="log"
+    ):
+        super().__init__()
+        self.register_buffer("dummy", torch.tensor(0))
+        self.kernel = None
+        self.init_kernel(samp_rate, kernel_size=kernel_size, freq_scale=freq_scale)
+        self.out_scale = out_scale
+
+    def forward(self, x):
+        """Apply the 2D filter (n_filts, kernel_size) to input signal x with shape: (batch_size, n_chs, seq_len)"""
+        x = scitex.dsp.ensure_3d(x).to(self.dummy.device)
+        seq_len = x.shape[-1]
+
+        # Ensure the kernel is initialized
+        if self.kernel is None:
+            self.init_kernel()
+            if self.kernel is None:
+                raise ValueError("Filter kernel has not been initialized.")
+        assert self.kernel.ndim == 2
+        self.kernel = self.kernel.to(x.device)  # cuda, torch.complex128
+
+        # Edge handling and convolution
+        extension_length = self.radius
+        first_segment = x[:, :, :extension_length].flip(dims=[-1])
+        last_segment = x[:, :, -extension_length:].flip(dims=[-1])
+        extended_x = torch.cat([first_segment, x, last_segment], dim=-1)
+
+        # working??
+        kernel_batched = self.kernel.unsqueeze(1)
+        extended_x_reshaped = extended_x.view(-1, 1, extended_x.shape[-1])
+
+        filtered_x_real = F.conv1d(
+            extended_x_reshaped, kernel_batched.real.float(), groups=1
+        )
+        filtered_x_imag = F.conv1d(
+            extended_x_reshaped, kernel_batched.imag.float(), groups=1
+        )
+
+        filtered_x = torch.view_as_complex(
+            torch.stack([filtered_x_real, filtered_x_imag], dim=-1)
+        )
+
+        filtered_x = filtered_x.view(
+            x.shape[0], x.shape[1], kernel_batched.shape[0], -1
+        )
+        filtered_x = filtered_x.view(
+            x.shape[0], x.shape[1], kernel_batched.shape[0], -1
+        )
+        filtered_x = filtered_x[..., :seq_len]
+        assert filtered_x.shape[-1] == seq_len
+
+        pha = filtered_x.angle()
+        amp = filtered_x.abs()
+
+        # Repeats freqs
+        freqs = (
+            self.freqs.unsqueeze(0).unsqueeze(0).repeat(pha.shape[0], pha.shape[1], 1)
+        )
+
+        if self.out_scale == "log":
+            return pha, torch.log(amp + 1e-5), freqs
+        else:
+            return pha, amp, freqs
+
+    def init_kernel(self, samp_rate, kernel_size=None, freq_scale="log"):
+        device = self.dummy.device
+        morlets, freqs = self.gen_morlet_to_nyquist(
+            samp_rate, kernel_size=kernel_size, freq_scale=freq_scale
+        )
+        self.kernel = torch.tensor(morlets).to(device)
+        self.freqs = torch.tensor(freqs).float().to(device)
+
+    @staticmethod
+    def gen_morlet_to_nyquist(samp_rate, kernel_size=None, freq_scale="linear"):
+        """
+        Generates Morlet wavelets for exponentially increasing frequency bands up to the Nyquist frequency.
+
+        Parameters:
+        - samp_rate (int): The sampling rate of the signal, in Hertz.
+        - kernel_size (int): The size of the kernel, in number of samples.
+
+        Returns:
+        - np.ndarray: A 2D array of complex values representing the Morlet wavelets for each frequency band.
+        """
+        if kernel_size is None:
+            kernel_size = int(samp_rate)  # * 2.5)
+
+        nyquist_freq = samp_rate / 2
+
+        # Log freq_scale
+        def calc_freq_boundaries_log(nyquist_freq):
+            n_kernels = int(np.floor(np.log2(nyquist_freq)))
+            mid_hz = np.array([2 ** (n + 1) for n in range(n_kernels)])
+            width_hz = np.hstack([np.array([1]), np.diff(mid_hz) / 2]) + 1
+            low_hz = mid_hz - width_hz
+            high_hz = mid_hz + width_hz
+            low_hz[0] = 0.1
+            return low_hz, high_hz
+
+        def calc_freq_boundaries_linear(nyquist_freq):
+            n_kernels = int(nyquist_freq)
+            high_hz = np.linspace(1, nyquist_freq, n_kernels)
+            low_hz = high_hz - np.hstack([np.array(1), np.diff(high_hz)])
+            low_hz[0] = 0.1
+            return low_hz, high_hz
+
+        if freq_scale == "linear":
+            fn = calc_freq_boundaries_linear
+        if freq_scale == "log":
+            fn = calc_freq_boundaries_log
+        low_hz, high_hz = fn(nyquist_freq)
+
+        morlets = []
+        freqs = []
+
+        for _, (ll, hh) in enumerate(zip(low_hz, high_hz)):
+            if ll > nyquist_freq:
+                break
+
+            center_frequency = (ll + hh) / 2
+
+            t = np.arange(-kernel_size // 2, kernel_size // 2) / samp_rate
+            # Calculate standard deviation of the gaussian window for a given center frequency
+            sigma = 7 / (2 * np.pi * center_frequency)
+            sine_wave = np.exp(2j * np.pi * center_frequency * t)
+            gaussian_window = np.exp(-(t**2) / (2 * sigma**2))
+            morlet_wavelet = sine_wave * gaussian_window
+
+            freqs.append(center_frequency)
+            morlets.append(morlet_wavelet)
+
+        return np.array(morlets), np.array(freqs)
+
+    @property
+    def kernel_size(
+        self,
+    ):
+        return to_even(self.kernel.shape[-1])
+
+    @property
+    def radius(
+        self,
+    ):
+        return to_even(self.kernel_size // 2)
+
+
+if __name__ == "__main__":
+    import matplotlib.pyplot as plt
+    import scitex
+
+    xx, tt, fs = scitex.dsp.demo_sig(sig_type="chirp")
+
+    pha, amp, ff = scitex.dsp.wavelet(xx, fs)
+
+    fig, ax = scitex.plt.subplots()
+    ax.imshow2d(amp[0, 0].T)
+    ax = scitex.plt.ax.set_ticks(ax, xticks=tt, yticks=ff)
+    ax = scitex.plt.ax.set_n_ticks(ax)
+    plt.show()
+
+
+# EOF

scitex/nn/__init__.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+"""Scitex nn module."""
+
+from ._AxiswiseDropout import AxiswiseDropout
+from ._BNet import BHead, BNet, BNet_config
+from ._BNet_Res import BHead, BNet, BNet_config
+from ._ChannelGainChanger import ChannelGainChanger
+from ._DropoutChannels import DropoutChannels
+from ._Filters import BandPassFilter, BandStopFilter, BaseFilter1D, DifferentiableBandPassFilter, GaussianFilter, HighPassFilter, LowPassFilter
+from ._FreqGainChanger import FreqGainChanger
+from ._GaussianFilter import GaussianFilter
+from ._Hilbert import Hilbert
+from ._MNet_1000 import MNet1000, MNet_1000, MNet_config, ReshapeLayer, SwapLayer
+from ._ModulationIndex import ModulationIndex
+from ._PAC import PAC
+from ._PSD import PSD
+from ._ResNet1D import ResNet1D, ResNetBasicBlock
+from ._SpatialAttention import SpatialAttention
+from ._Spectrogram import Spectrogram, my_softmax, normalize, spectrograms, unbias
+from ._SwapChannels import SwapChannels
+from ._TransposeLayer import TransposeLayer
+from ._Wavelet import Wavelet
+
+__all__ = [
+    "AxiswiseDropout",
+    "BHead",
+    "BHead",
+    "BNet",
+    "BNet",
+    "BNet_config",
+    "BNet_config",
+    "BandPassFilter",
+    "BandStopFilter",
+    "BaseFilter1D",
+    "ChannelGainChanger",
+    "DifferentiableBandPassFilter",
+    "DropoutChannels",
+    "FreqGainChanger",
+    "GaussianFilter",
+    "GaussianFilter",
+    "HighPassFilter",
+    "Hilbert",
+    "LowPassFilter",
+    "MNet1000",
+    "MNet_1000",
+    "MNet_config",
+    "ModulationIndex",
+    "PAC",
+    "PSD",
+    "ResNet1D",
+    "ResNetBasicBlock",
+    "ReshapeLayer",
+    "SpatialAttention",
+    "Spectrogram",
+    "SwapChannels",
+    "SwapLayer",
+    "TransposeLayer",
+    "Wavelet",
+    "my_softmax",
+    "normalize",
+    "spectrograms",
+    "unbias",
+]

scitex/os/__init__.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Scitex os module."""
+
+from ._mv import mv
+
+__all__ = [
+    "mv",
+]

scitex/os/_mv.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-04-06 09:00:45 (ywatanabe)"
+
+# import os
+# import shutil
+
+# def mv(src, tgt):
+#     successful = True
+#     os.makedirs(tgt, exist_ok=True)
+
+#     if os.path.isdir(src):
+#         # Iterate over the items in the directory
+#         for item in os.listdir(src):
+#             item_path = os.path.join(src, item)
+#             # Check if the item is a file
+#             if os.path.isfile(item_path):
+#                 try:
+#                     shutil.move(item_path, tgt)
+#                     print(f"\nMoved file from {item_path} to {tgt}")
+#                 except OSError as e:
+#                     print(f"\nError: {e}")
+#                     successful = False
+#             else:
+#                 print(f"\nSkipped directory {item_path}")
+#     else:
+#         # If src is a file, just move it
+#         try:
+#             shutil.move(src, tgt)
+#             print(f"\nMoved from {src} to {tgt}")
+#         except OSError as e:
+#             print(f"\nError: {e}")
+#             successful = False
+
+#     return successful
+
+
+def mv(src, tgt):
+    import os
+    import shutil
+
+    successful = True
+    os.makedirs(tgt, exist_ok=True)
+
+    try:
+        shutil.move(src, tgt)
+        print(f"\nMoved from {src} to {tgt}")
+    except OSError as e:
+        print(f"\nError: {e}")
+        successful = False

scitex/parallel/__init__.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Scitex parallel module."""
+
+from ._run import run
+
+__all__ = [
+    "run",
+]

scitex/parallel/_run.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-14 23:12:20 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/parallel/_run.py
+
+"""
+1. Functionality:
+   - Runs functions in parallel using ProcessPoolExecutor
+   - Handles both single and multiple return values
+   - Supports automatic CPU core detection
+2. Input:
+   - Function to run
+   - List of items to process
+   - Optional parameters for execution control
+3. Output:
+   - List of results or concatenated DataFrame/tuple
+4. Prerequisites:
+   - concurrent.futures
+   - pandas
+   - tqdm
+"""
+
+import multiprocessing
+import warnings
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, List
+
+from tqdm import tqdm
+
+
+def run(
+    func: Callable,
+    args_list: List[tuple],
+    n_jobs: int = -1,
+    desc: str = "Processing",
+) -> List[Any]:
+    """Runs function in parallel using ThreadPoolExecutor with tuple arguments.
+
+    Parameters
+    ----------
+    func : Callable
+        Function to run in parallel
+    args_list : List[tuple]
+        List of argument tuples, each tuple contains arguments for one function call
+    n_jobs : int, optional
+        Number of jobs to run in parallel. -1 means using all processors
+    desc : str, optional
+        Description for progress bar
+
+    Returns
+    -------
+    List[Any]
+        Results of parallel execution
+
+    Examples
+    --------
+    >>> def add(x, y):
+    ...     return x + y
+    >>> args_list = [(1, 4), (2, 5), (3, 6)]
+    >>> run(add, args_list)
+    [5, 7, 9]
+    """
+    if not args_list:
+        raise ValueError("Args list cannot be empty")
+    if not callable(func):
+        raise ValueError("Func must be callable")
+
+    cpu_count = multiprocessing.cpu_count()
+    n_jobs = cpu_count if n_jobs < 0 else n_jobs
+
+    if n_jobs > cpu_count:
+        warnings.warn(f"n_jobs ({n_jobs}) is greater than CPU count ({cpu_count})")
+    if n_jobs < 1:
+        raise ValueError("n_jobs must be >= 1 or -1")
+
+    results = [None] * len(args_list)  # Pre-allocate list
+
+    with ThreadPoolExecutor(max_workers=n_jobs) as executor:
+        futures = {
+            executor.submit(func, *args): idx for idx, args in enumerate(args_list)
+        }
+        for future in tqdm(as_completed(futures), total=len(args_list), desc=desc):
+            idx = futures[future]
+            results[idx] = future.result()
+
+    # If results contain multiple values (tuples), transpose them
+    if results and isinstance(results[0], tuple):
+        n_vars = len(results[0])
+        return tuple([result[i] for result in results] for i in range(n_vars))
+
+    return results
+
+
+# def run(
+#     func: Callable,
+#     items: List[Any],
+#     n_jobs: int = -1,
+#     desc: str = "Processing",
+# ) -> List[Any]:
+#     """Runs function in parallel using ThreadPoolExecutor.
+
+#     Parameters
+#     ----------
+#     func : Callable
+#         Function to run in parallel
+#     items : List[Any]
+#         List of items to process
+#     n_jobs : int, optional
+#         Number of jobs to run in parallel. -1 means using all processors
+#     desc : str, optional
+#         Description for progress bar
+
+#     Returns
+#     -------
+#     List[Any]
+#         Results of parallel execution
+#     """
+#     if not items:
+#         raise ValueError("Items list cannot be empty")
+#     if not callable(func):
+#         raise ValueError("Func must be callable")
+#     if not isinstance(items, (list, tuple)):
+#         raise TypeError("Items must be a list or tuple")
+#     if not isinstance(n_jobs, int):
+#         raise TypeError("n_jobs must be an integer")
+
+#     cpu_count = multiprocessing.cpu_count()
+#     n_jobs = cpu_count if n_jobs < 0 else n_jobs
+
+#     if n_jobs > cpu_count:
+#         warnings.warn(f"n_jobs ({n_jobs}) is greater than CPU count ({cpu_count})")
+#     if n_jobs < 1:
+#         raise ValueError("n_jobs must be >= 1 or -1")
+
+#     results = [None] * len(items)  # Pre-allocate list
+#     with ThreadPoolExecutor(max_workers=n_jobs) as executor:
+#         futures = {executor.submit(func, item): idx
+#                   for idx, item in enumerate(items)}
+#         for future in tqdm(as_completed(futures), total=len(items), desc=desc):
+#             idx = futures[future]
+#             results[idx] = future.result()
+
+#     # If results contain multiple values (tuples), transpose them
+#     if results and isinstance(results[0], tuple):
+#         n_vars = len(results[0])
+#         return tuple([result[i] for result in results] for i in range(n_vars))
+
+#     return results
+
+
+# EOF