scitex 2.0.0__py2.py3-none-any.whl

scitex/dsp/_transform.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-04-08 12:41:59 (ywatanabe)"#!/usr/bin/env python3
+
+
+import warnings
+
+import numpy as np
+import pandas as pd
+import torch
+from ..decorators import torch_fn
+
+
+def to_sktime_df(arr):
+    """
+    Convert a 3D numpy array into a DataFrame suitable for sktime.
+
+    Parameters:
+    arr (numpy.ndarray): A 3D numpy array with shape (n_samples, n_channels, seq_len)
+
+    Returns:
+    pandas.DataFrame: A DataFrame in sktime format
+    """
+    if len(arr.shape) != 3:
+        raise ValueError("Input data must be a 3D array")
+
+    n_samples, seq_len, n_channels = arr.shape
+
+    # Initialize an empty DataFrame for sktime format
+    sktime_df = pd.DataFrame(index=range(n_samples), columns=["dim_0"])
+
+    # Iterate over each sample
+    for i in range(n_samples):
+        # Combine all channels into a single cell
+        combined_series = pd.Series(
+            {f"channel_{j}": pd.Series(arr[i, :, j]) for j in range(n_channels)}
+        )
+        sktime_df.iloc[i, 0] = combined_series
+
+    return sktime_df
+
+
+@torch_fn
+def to_segments(x, window_size, overlap_factor=1, dim=-1):
+    stride = window_size // overlap_factor
+    num_windows = (x.size(dim) - window_size) // stride + 1
+    windows = x.unfold(dim, window_size, stride)
+    return windows
+
+
+if __name__ == "__main__":
+    import scitex
+
+    x, t, f = scitex.dsp.demo_sig()
+
+    y = to_segments(x, 256)
+
+    x = 100 * np.random.rand(16, 160, 1000)
+    print(_normalize_time(x))
+
+    x = torch.randn(16, 160, 1000)
+    print(_normalize_time(x))
+
+    x = torch.randn(16, 160, 1000).cuda()
+    print(_normalize_time(x))
+
+
+# EOF

scitex/dsp/_wavelet.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-04 02:12:00 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/dsp/_wavelet.py
+
+"""scitex.dsp.wavelet function"""
+
+
+from ..decorators import batch_fn, signal_fn
+from ..nn._Wavelet import Wavelet
+
+
+# Functions
+@signal_fn
+@batch_fn
+def wavelet(
+    x,
+    fs,
+    freq_scale="linear",
+    out_scale="linear",
+    device="cuda",
+    batch_size=32,
+):
+    m = Wavelet(fs, freq_scale=freq_scale, out_scale="linear").to(device).eval()
+    pha, amp, freqs = m(x.to(device))
+
+    if out_scale == "log":
+        amp = (amp + 1e-5).log()
+        if amp.isnan().any():
+            print("NaN is detected while taking the lograrithm of amplitude.")
+
+    return pha, amp, freqs
+
+
+# @signal_fn
+# def wavelet(
+#     x,
+#     fs,
+#     freq_scale="linear",
+#     out_scale="linear",
+#     device="cuda",
+#     batch_size=32,
+# ):
+#     @signal_fn
+#     def _wavelet(
+#         x,
+#         fs,
+#         freq_scale="linear",
+#         out_scale="linear",
+#         device="cuda",
+#     ):
+#         m = (
+#             Wavelet(fs, freq_scale=freq_scale, out_scale=out_scale)
+#             .to(device)
+#             .eval()
+#         )
+#         pha, amp, freqs = m(x.to(device))
+
+#         if out_scale == "log":
+#             amp = (amp + 1e-5).log()
+#             if amp.isnan().any():
+#                 print(
+#                     "NaN is detected while taking the lograrithm of amplitude."
+#                 )
+
+#         return pha, amp, freqs
+
+#     if len(x) <= batch_size:
+#         try:
+#             pha, amp, freqs = _wavelet(
+#                 x,
+#                 fs,
+#                 freq_scale=freq_scale,
+#                 out_scale=out_scale,
+#                 device=device,
+#             )
+#             torch.cuda.empty_cache()
+#             return pha, amp, freqs
+
+#         except Exception as e:
+#             print(e)
+#             print("\nTrying Batch Mode...")
+
+#     n_batches = (len(x) + batch_size - 1) // batch_size
+#     device_orig = x.device
+#     pha, amp, freqs = [], [], []
+#     for i_batch in tqdm(range(n_batches)):
+#         start = i_batch * batch_size
+#         end = (i_batch + 1) * batch_size
+#         _pha, _amp, _freqs = _wavelet(
+#             x[start:end],
+#             fs,
+#             freq_scale=freq_scale,
+#             out_scale=out_scale,
+#             device=device,
+#         )
+#         torch.cuda.empty_cache()
+#         # to CPU
+#         pha.append(_pha.cpu())
+#         amp.append(_amp.cpu())
+#         freqs.append(_freqs.cpu())
+
+#     pha = torch.vstack(pha)
+#     amp = torch.vstack(amp)
+#     freqs = freqs[0]
+
+#     try:
+#         pha = pha.to(device_orig)
+#         amp = amp.to(device_orig)
+#         freqs = freqs.to(device_orig)
+#     except Exception as e:
+#         print(
+#             f"\nError occurred while transferring wavelet outputs back to the original device. Proceeding with CPU tensor. \n\n({e})"
+#         )
+
+#     sleep(0.5)
+#     torch.cuda.empty_cache()
+#     return pha, amp, freqs
+
+
+if __name__ == "__main__":
+    import sys
+
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    # Start
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt, agg=True)
+
+    # Parameters
+    FS = 512
+    SIG_TYPE = "chirp"
+    T_SEC = 4
+
+    # Demo signal
+    xx, tt, fs = scitex.dsp.demo_sig(
+        batch_size=64,
+        n_chs=19,
+        n_segments=2,
+        t_sec=T_SEC,
+        fs=FS,
+        sig_type=SIG_TYPE,
+    )
+
+    if SIG_TYPE in ["tensorpac", "pac"]:
+        i_segment = 0
+        xx = xx[:, :, i_segment, :]
+
+    # Main
+    pha, amp, freqs = wavelet(xx, fs, device="cuda")
+    freqs = freqs[0, 0]
+
+    # Plots
+    i_batch, i_ch = 0, 0
+    fig, axes = scitex.plt.subplots(nrows=3)
+
+    # # Time vector for x-axis extents
+    # time_extent = [tt.min(), tt.max()]
+
+    # Trace
+    axes[0].plot(tt, xx[i_batch, i_ch], label=SIG_TYPE)
+    axes[0].set_ylabel("Amplitude [?V]")
+    axes[0].legend(loc="upper left")
+    axes[0].set_title("Signal")
+
+    # Amplitude
+    # extent = [time_extent[0], time_extent[1], freqs.min(), freqs.max()]
+    axes[1].imshow2d(
+        np.log(amp[i_batch, i_ch] + 1e-5).T,
+        cbar_label="Log(amplitude [?V]) [a.u.]",
+        aspect="auto",
+        # extent=extent,
+        # origin="lower",
+    )
+    axes[1] = scitex.plt.ax.set_ticks(axes[1], x_ticks=tt, y_ticks=freqs)
+    axes[1].set_ylabel("Frequency [Hz]")
+    axes[1].set_title("Amplitude")
+
+    # Phase
+    axes[2].imshow2d(
+        pha[i_batch, i_ch].T,
+        cbar_label="Phase [rad]",
+        aspect="auto",
+        # extent=extent,
+        # origin="lower",
+    )
+    axes[2] = scitex.plt.ax.set_ticks(axes[2], x_ticks=tt, y_ticks=freqs)
+    axes[2].set_ylabel("Frequency [Hz]")
+    axes[2].set_title("Phase")
+
+    fig.suptitle("Wavelet Transformation")
+    fig.supxlabel("Time [s]")
+
+    for ax in axes:
+        ax = scitex.plt.ax.set_n_ticks(ax)
+        # ax.set_xlim(time_extent[0], time_extent[1])
+
+    fig.tight_layout(rect=[0, 0.03, 1, 0.95])
+
+    scitex.io.save(fig, "wavelet.png")
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF
+
+"""
+/home/ywatanabe/proj/entrance/scitex/dsp/_wavelet.py
+"""
+
+
+# EOF

scitex/dsp/add_noise.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "ywatanabe (2024-11-02 23:09:49)"
+# File: ./scitex_repo/src/scitex/dsp/add_noise.py
+
+import torch
+from ..decorators import signal_fn
+
+
+def _uniform(shape, amp=1.0):
+    a, b = -amp, amp
+    return -amp + (2 * amp) * torch.rand(shape)
+
+
+@signal_fn
+def gauss(x, amp=1.0):
+    noise = amp * torch.randn(x.shape)
+    return x + noise.to(x.device)
+
+
+@signal_fn
+def white(x, amp=1.0):
+    return x + _uniform(x.shape, amp=amp).to(x.device)
+
+
+@signal_fn
+def pink(x, amp=1.0, dim=-1):
+    """
+    Adds pink noise to a given tensor along a specified dimension.
+
+    Parameters:
+    - x (torch.Tensor): The input tensor to which pink noise will be added.
+    - amp (float, optional): The amplitude of the pink noise. Defaults to 1.0.
+    - dim (int, optional): The dimension along which to add pink noise. Defaults to -1.
+
+    Returns:
+    - torch.Tensor: The input tensor with added pink noise.
+    """
+    cols = x.size(dim)
+    noise = torch.randn(cols, dtype=x.dtype, device=x.device)
+    noise = torch.fft.rfft(noise)
+    indices = torch.arange(1, noise.size(0), dtype=x.dtype, device=x.device)
+    noise[1:] /= torch.sqrt(indices)
+    noise = torch.fft.irfft(noise, n=cols)
+    noise = noise - noise.mean()
+    noise_amp = torch.sqrt(torch.mean(noise**2))
+    noise = noise * (amp / noise_amp)
+    return x + noise.to(x.device)
+
+
+@signal_fn
+def brown(x, amp=1.0, dim=-1):
+    noise = _uniform(x.shape, amp=amp)
+    noise = torch.cumsum(noise, dim=dim)
+    noise = scitex.dsp.norm.minmax(noise, amp=amp, dim=dim)
+    return x + noise.to(x.device)
+
+
+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)
+
+    # Parameters
+    T_SEC = 1
+    FS = 128
+
+    # Demo signal
+    xx, tt, fs = scitex.dsp.demo_sig(t_sec=T_SEC, fs=FS)
+
+    funcs = {
+        "orig": lambda x: x,
+        "gauss": gauss,
+        "white": white,
+        "pink": pink,
+        "brown": brown,
+    }
+
+    # Plots
+    fig, axes = scitex.plt.subplots(nrows=len(funcs), ncols=2, sharex=True, sharey=True)
+    count = 0
+    for (k, fn), axes_row in zip(funcs.items(), axes):
+        for ax in axes_row:
+            if count % 2 == 0:
+                ax.plot(tt, fn(xx)[0, 0], label=k, c="blue")
+            else:
+                ax.plot(tt, (fn(xx) - xx)[0, 0], label=f"{k} - orig", c="red")
+            count += 1
+            ax.legend(loc="upper right")
+
+    fig.supxlabel("Time [s]")
+    fig.supylabel("Amplitude [?V]")
+    axes[0, 0].set_title("Signal + Noise")
+    axes[0, 1].set_title("Noise")
+
+    scitex.io.save(fig, "traces.png")
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF
+
+"""
+/home/ywatanabe/proj/entrance/scitex/dsp/add_noise.py
+"""
+
+# EOF

scitex/dsp/example.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-04-06 01:36:18 (ywatanabe)"
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import pandas as pd
+
+
+# Functions
+def calc_norm_resample_filt_hilbert(xx, tt, fs, sig_type, verbose=True):
+    sigs = {"index": ("signal", "time", "fs")}  # Collector
+
+    if sig_type == "tensorpac":
+        xx = xx[:, :, 0]
+
+    sigs[f"orig"] = (xx, tt, fs)
+
+    # Normalization
+    sigs["z_normed"] = (scitex.dsp.norm.z(xx), tt, fs)
+    sigs["minmax_normed"] = (scitex.dsp.norm.minmax(xx), tt, fs)
+
+    # Resampling
+    sigs["resampled"] = (
+        scitex.dsp.resample(xx, fs, TGT_FS),
+        tt[:: int(fs / TGT_FS)],
+        TGT_FS,
+    )
+
+    # Noise injection
+    sigs["gaussian_noise_added"] = (scitex.dsp.add_noise.gauss(xx), tt, fs)
+    sigs["white_noise_added"] = (scitex.dsp.add_noise.white(xx), tt, fs)
+    sigs["pink_noise_added"] = (scitex.dsp.add_noise.pink(xx), tt, fs)
+    sigs["brown_noise_added"] = (scitex.dsp.add_noise.brown(xx), tt, fs)
+
+    # Filtering
+    sigs[f"bandpass_filted ({LOW_HZ} - {HIGH_HZ} Hz)"] = (
+        scitex.dsp.filt.bandpass(xx, fs, low_hz=LOW_HZ, high_hz=HIGH_HZ),
+        tt,
+        fs,
+    )
+
+    sigs[f"bandstop_filted ({LOW_HZ} - {HIGH_HZ} Hz)"] = (
+        scitex.dsp.filt.bandstop(xx, fs, low_hz=LOW_HZ, high_hz=HIGH_HZ),
+        tt,
+        fs,
+    )
+    sigs[f"bandstop_gauss (sigma = {SIGMA})"] = (
+        scitex.dsp.filt.gauss(xx, sigma=SIGMA),
+        tt,
+        fs,
+    )
+
+    # Hilbert Transformation
+    pha, amp = scitex.dsp.hilbert(xx)
+    sigs["hilbert_amp"] = (amp, tt, fs)
+    sigs["hilbert_pha"] = (pha, tt, fs)
+
+    sigs = pd.DataFrame(sigs).set_index("index")
+
+    if verbose:
+        print(sigs.index)
+        print(sigs.columns)
+
+    return sigs
+
+
+def plot_signals(plt, sigs, sig_type):
+    fig, axes = plt.subplots(nrows=len(sigs.columns), sharex=True)
+
+    i_batch = 0
+    i_ch = 0
+    for ax, (i_col, col) in zip(axes, enumerate(sigs.columns)):
+
+        if col == "hilbert_amp":  # add the original signal to the ax
+            _col = "orig"
+            (
+                _xx,
+                _tt,
+                _fs,
+            ) = sigs[_col]
+            ax.plot(_tt, _xx[i_batch, i_ch], label=_col, c=CC["blue"])
+
+        # Main
+        xx, tt, fs = sigs[col]
+        # if sig_type == "tensorpac":
+        #     xx = xx[:, :, 0]
+
+        try:
+            ax.plot(
+                tt,
+                xx[i_batch, i_ch],
+                label=col,
+                c=CC["red"] if col == "hilbert_amp" else CC["blue"],
+            )
+        except Exception as e:
+            print(e)
+            import ipdb
+
+            ipdb.set_trace()
+
+        # Adjustments
+        ax.legend(loc="upper left")
+        ax.set_xlim(tt[0], tt[-1])
+
+        ax = scitex.plt.ax.set_n_ticks(ax)
+
+    fig.supxlabel("Time [s]")
+    fig.supylabel("Voltage")
+    fig.suptitle(sig_type)
+    return fig
+
+
+def plot_wavelet(plt, sigs, sig_col, sig_type):
+
+    xx, tt, fs = sigs[sig_col]
+    # if sig_type == "tensorpac":
+    #     xx = xx[:, :, 0]
+
+    # Wavelet Transformation
+    wavelet_coef, ff_ww = scitex.dsp.wavelet(xx, fs)
+
+    i_batch = 0
+    i_ch = 0
+
+    # Main
+    fig, axes = plt.subplots(nrows=2, sharex=True)
+    # Signal
+    axes[0].plot(
+        tt,
+        xx[i_batch, i_ch],
+        label=sig_col,
+        c=CC["blue"],
+    )
+    # Adjusts
+    axes[0].legend(loc="upper left")
+    axes[0].set_xlim(tt[0], tt[-1])
+    axes[0].set_ylabel("Voltage")
+    axes[0] = scitex.plt.ax.set_n_ticks(axes[0])
+
+    # Wavelet Spectrogram
+    axes[1].imshow(
+        wavelet_coef[i_batch, i_ch],
+        aspect="auto",
+        extent=[tt[0], tt[-1], 512, 1],
+        label="wavelet_coefficient",
+    )
+    # axes[1].set_xlabel("Time [s]")
+    axes[1].set_ylabel("Frequency [Hz]")
+    # axes[1].legend(loc="upper left")
+    axes[1].invert_yaxis()
+
+    fig.supxlabel("Time [s]")
+    fig.suptitle(sig_type)
+
+    return fig
+
+
+def plot_psd(plt, sigs, sig_col, sig_type):
+
+    xx, tt, fs = sigs[sig_col]
+
+    # if sig_type == "tensorpac":
+    #     xx = xx[:, :, 0]
+
+    # Power Spetrum Density
+    psd, ff_pp = scitex.dsp.psd(xx, fs)
+
+    # Main
+    i_batch = 0
+    i_ch = 0
+    fig, axes = plt.subplots(nrows=2, sharex=False)
+
+    # Signal
+    axes[0].plot(
+        tt,
+        xx[i_batch, i_ch],
+        label=sig_col,
+        c=CC["blue"],
+    )
+    # Adjustments
+    axes[0].legend(loc="upper left")
+    axes[0].set_xlim(tt[0], tt[-1])
+    axes[0].set_xlabel("Time [s]")
+    axes[0].set_ylabel("Voltage")
+    axes[0] = scitex.plt.ax.set_n_ticks(axes[0])
+
+    # PSD
+    axes[1].plot(ff_pp, psd[i_batch, i_ch], label="PSD")
+    axes[1].set_yscale("log")
+    axes[1].set_ylabel("Power [uV^2 / Hz]")
+    axes[1].set_xlabel("Frequency [Hz]")
+
+    fig.suptitle(sig_type)
+
+    return fig
+
+
+if __name__ == "__main__":
+
+    # Parameters
+    T_SEC = 4
+    SIG_TYPES = [
+        # "uniform",
+        # "gauss",
+        # "periodic",
+        # "chirp",
+        # "ripple",
+        # "meg",
+        "tensorpac",
+    ]
+    SRC_FS = 1024
+    TGT_FS = 512
+    FREQS_HZ = [10, 30, 100]
+    LOW_HZ = 20
+    HIGH_HZ = 50
+    SIGMA = 10
+
+    plt, CC = scitex.plt.configure_mpl(plt, fig_scale=10)
+    sdir = "/home/ywatanabe/proj/entrance/scitex/dsp/example/"
+
+    for sig_type in SIG_TYPES:
+        # Demo Signal
+        xx, tt, fs = scitex.dsp.demo_sig(
+            t_sec=T_SEC, fs=SRC_FS, freqs_hz=FREQS_HZ, sig_type=sig_type
+        )
+
+        # Apply calculations on the original signal
+        sigs = calc_norm_resample_filt_hilbert(xx, tt, fs, sig_type)
+
+        # Plots signals
+        fig = plot_signals(plt, sigs, sig_type)
+        scitex.io.save(fig, sdir + f"{sig_type}/1_signals.png")
+
+        # Plots wavelet coefficients and PSD
+        for sig_col in sigs.columns:
+
+            if "hilbert" in sig_col:
+                continue
+
+            fig = plot_wavelet(plt, sigs, sig_col, sig_type)
+            scitex.io.save(fig, sdir + f"{sig_type}/2_wavelet_{sig_col}.png")
+
+            fig = plot_psd(plt, sigs, sig_col, sig_type)
+            scitex.io.save(fig, sdir + f"{sig_type}/3_psd_{sig_col}.png")
+
+    # plt.show()
+
+    """
+    python ./dsp/example.py
+    """