scitex 2.0.0__py2.py3-none-any.whl

scitex/nn/_Hilbert.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-04-10 12:46:06 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_repo/src/scitex/nn/_Hilbert.py
+# ----------------------------------------
+import os
+
+__FILE__ = "/home/ywatanabe/proj/scitex_repo/src/scitex/nn/_Hilbert.py"
+__DIR__ = os.path.dirname(__FILE__)
+# ----------------------------------------
+#!/usr/bin/env python
+
+import torch  # 1.7.1
+import torch.nn as nn
+from torch.fft import fft, ifft
+
+
+class Hilbert(nn.Module):
+    def __init__(self, seq_len, dim=-1, fp16=False, in_place=False):
+        super().__init__()
+        self.dim = dim
+        self.fp16 = fp16
+        self.in_place = in_place
+        self.n = seq_len
+        f = torch.cat(
+            [
+                torch.arange(0, (self.n - 1) // 2 + 1) / float(self.n),
+                torch.arange(-(self.n // 2), 0) / float(self.n),
+            ]
+        )
+        self.register_buffer("f", f)
+
+    def hilbert_transform(self, x):
+        # n = x.shape[self.dim]
+
+        # Create frequency dim
+        # f = torch.cat(
+        #     [
+        #         torch.arange(0, (n - 1) // 2 + 1, device=x.device) / float(n),
+        #         torch.arange(-(n // 2), 0, device=x.device) / float(n),
+        #     ]
+        # )
+
+        orig_dtype = x.dtype
+        x = x.float()
+        xf = fft(x, n=self.n, dim=self.dim)
+        x = x.to(orig_dtype)
+
+        # Create step function
+        steepness = 50  # This value can be adjusted
+        u = torch.sigmoid(
+            steepness * self.f.type_as(x)
+        )  # Soft step function for differentiability
+
+        transformed = ifft(xf * 2 * u, dim=self.dim)
+
+        return transformed
+
+    def forward(self, x):
+        if self.fp16:
+            x = x.half()
+
+        if not self.in_place:
+            x = x.clone()  # Ensure that we do not modify the input in-place
+
+        x_comp = self.hilbert_transform(x)
+
+        pha = torch.atan2(x_comp.imag, x_comp.real)
+        amp = x_comp.abs()
+
+        assert x.shape == pha.shape == amp.shape
+
+        out = torch.cat(
+            [
+                pha.unsqueeze(-1),
+                amp.unsqueeze(-1),
+            ],
+            dim=-1,
+        )
+
+        # if self.fp16:
+        #     out = (
+        #         out.float()
+        #     )
+        #     # Optionally cast back to float for stability in subsequent operations
+
+        if self.fp16:
+            out = out.float()
+
+        return out
+
+
+if __name__ == "__main__":
+    import scitex
+
+    xx, tt, fs = scitex.dsp.demo_sig()
+    xx = torch.tensor(xx)
+
+    # Parameters
+    device = "cuda"
+    fp16 = True
+    in_place = True
+
+    # Initialization
+    m = Hilbert(xx.shape[-1], fp16=fp16, in_place=in_place).to(device)
+
+    # Calculation
+    xx = xx.to(device)
+    y = m(xx)
+
+# EOF

scitex/nn/_MNet_1000.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2023-05-04 16:54:55 (ywatanabe)"
+
+#!/usr/bin/env python
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchsummary import summary
+import scitex
+
+MNet_config = {
+    "classes": ["class1", "class2"],
+    "n_chs": 270,
+    "n_fc1": 1024,
+    "d_ratio1": 0.85,
+    "n_fc2": 256,
+    "d_ratio2": 0.85,
+}
+
+
+class MNet1000(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        # basic
+        self.config = config
+        # fc
+        N_FC_IN = 15950
+
+        # conv
+        self.backborn = nn.Sequential(
+            *[
+                nn.Conv2d(1, 40, kernel_size=(config["n_chs"], 4)),
+                nn.Mish(),
+                nn.Conv2d(40, 40, kernel_size=(1, 4)),
+                nn.BatchNorm2d(40),
+                nn.MaxPool2d((1, 5)),
+                nn.Mish(),
+                SwapLayer(),
+                nn.Conv2d(1, 50, kernel_size=(8, 12)),
+                nn.BatchNorm2d(50),
+                nn.MaxPool2d((3, 3)),
+                nn.Mish(),
+                nn.Conv2d(50, 50, kernel_size=(1, 5)),
+                nn.BatchNorm2d(50),
+                nn.MaxPool2d((1, 2)),
+                nn.Mish(),
+                ReshapeLayer(),
+                nn.Linear(N_FC_IN, config["n_fc1"]),
+            ]
+        )
+
+        # # conv
+        # self.conv1 = nn.Conv2d(1, 40, kernel_size=(config["n_chs"], 4))
+        # self.act1 = nn.Mish()
+
+        # self.conv2 = nn.Conv2d(40, 40, kernel_size=(1, 4))
+        # self.bn2 = nn.BatchNorm2d(40)
+        # self.pool2 = nn.MaxPool2d((1, 5))
+        # self.act2 = nn.Mish()
+
+        # self.swap = SwapLayer()
+
+        # self.conv3 = nn.Conv2d(1, 50, kernel_size=(8, 12))
+        # self.bn3 = nn.BatchNorm2d(50)
+        # self.pool3 = nn.MaxPool2d((3, 3))
+        # self.act3 = nn.Mish()
+
+        # self.conv4 = nn.Conv2d(50, 50, kernel_size=(1, 5))
+        # self.bn4 = nn.BatchNorm2d(50)
+        # self.pool4 = nn.MaxPool2d((1, 2))
+        # self.act4 = nn.Mish()
+
+        self.fc = nn.Sequential(
+            # nn.Linear(N_FC_IN, config["n_fc1"]),
+            nn.Mish(),
+            nn.Dropout(config["d_ratio1"]),
+            nn.Linear(config["n_fc1"], config["n_fc2"]),
+            nn.Mish(),
+            nn.Dropout(config["d_ratio2"]),
+            nn.Linear(config["n_fc2"], len(config["classes"])),
+        )
+
+    @staticmethod
+    def _reshape_input(x, n_chs):
+        """
+        (batch, channel, time_length) -> (batch, channel, time_length, new_axis)
+        """
+        if x.ndim == 3:
+            x = x.unsqueeze(-1)
+        if x.shape[2] == n_chs:
+            x = x.transpose(1, 2)
+        x = x.transpose(1, 3).transpose(2, 3)
+        return x
+
+    @staticmethod
+    def _znorm_along_the_last_dim(x):
+        return (x - x.mean(dim=-1, keepdims=True)) / x.std(dim=-1, keepdims=True)
+
+    def forward(self, x):
+        # # time-wise normalization
+        # x = self._znorm_along_the_last_dim(x)
+        # x = self._reshape_input(x, self.config["n_chs"])
+
+        # x = self.backborn(x)
+        x = self.forward_bb(x)
+
+        # x = x.reshape(len(x), -1)
+
+        x = self.fc(x)
+
+        return x
+
+    def forward_bb(self, x):
+        # time-wise normalization
+        x = self._znorm_along_the_last_dim(x)
+        x = self._reshape_input(x, self.config["n_chs"])
+        x = self.backborn(x)
+        return x
+
+
+class SwapLayer(nn.Module):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def forward(self, x):
+        return x.transpose(1, 2)
+
+
+class ReshapeLayer(nn.Module):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def forward(self, x):
+        return x.reshape(len(x), -1)
+
+
+if __name__ == "__main__":
+    ## Demo data
+    BS, N_CHS, SEQ_LEN = 16, 270, 1000
+    x = torch.rand(BS, N_CHS, SEQ_LEN).cuda()
+
+    ## Config for the model
+    model = MNet_1000(MNet_config).cuda()
+
+    y = model(x)
+    summary(model, x)
+    print(y.shape)
+
+# Backward compatibility
+MNet_1000 = MNet1000  # Deprecated: use MNet1000 instead

scitex/nn/_ModulationIndex.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-11-04 02:08:01 (ywatanabe)"
+# File: ./scitex_repo/src/scitex/nn/_ModulationIndex.py
+
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-10-15 14:12:55 (ywatanabe)"
+
+"""
+This script defines the ModulationIndex module.
+"""
+
+# Imports
+import sys
+import warnings
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Functions
+class ModulationIndex(nn.Module):
+    def __init__(self, n_bins=18, fp16=False, amp_prob=False):
+        super(ModulationIndex, self).__init__()
+        self.n_bins = n_bins
+        self.fp16 = fp16
+        self.register_buffer(
+            "pha_bin_cutoffs", torch.linspace(-np.pi, np.pi, n_bins + 1)
+        )
+
+        self.amp_prob = amp_prob
+
+    @property
+    def pha_bin_centers(
+        self,
+    ):
+        return (
+            ((self.pha_bin_cutoffs[1:] + self.pha_bin_cutoffs[:-1]) / 2)
+            .detach()
+            .cpu()
+            .numpy()
+        )
+
+    def forward(self, pha, amp, epsilon=1e-9):
+        """
+        Compute the Modulation Index based on phase (pha) and amplitude (amp) tensors.
+
+        Parameters:
+        - pha (torch.Tensor): Tensor of phase values with shape
+                              (batch_size, n_channels, n_freqs_pha, n_segments, sequence_length).
+        - amp (torch.Tensor): Tensor of amplitude values with a similar shape as pha.
+                              (batch_size, n_channels, n_freqs_amp, n_segments, sequence_length).
+
+        Returns:
+        - MI (torch.Tensor): The Modulation Index for each batch and channel.
+        """
+        assert pha.ndim == amp.ndim == 5
+
+        if self.fp16:
+            pha, amp = pha.half().contiguous(), amp.half().contiguous()
+        else:
+            pha, amp = pha.float().contiguous(), amp.float().contiguous()
+
+        device = pha.device
+
+        pha_masks = self._phase_to_masks(pha, self.pha_bin_cutoffs.to(device))
+        # (batch_size, n_channels, n_freqs_pha, n_segments, sequence_length, n_bins)
+
+        # Expands amp and masks to utilize broadcasting
+        # i_batch = 0
+        # i_chs = 1
+        i_freqs_pha = 2
+        i_freqs_amp = 3
+        # i_segments = 4
+        i_time = 5
+        i_bins = 6
+
+        # Coupling
+        pha_masks = pha_masks.unsqueeze(i_freqs_amp)
+        amp = amp.unsqueeze(i_freqs_pha).unsqueeze(i_bins)
+
+        amp_bins = pha_masks * amp  # this is the most memory-consuming process
+
+        # # Batch processing to reduce maximum VRAM occupancy
+        # pha_masks = self.dh_pha.fit(pha_masks, keepdims=[2, 3, 5, 6])
+        # amp = self.dh_amp.fit(amp, keepdims=[2, 3, 5, 6])
+        # n_chunks = len(pha_masks) // self.chunk_size
+        # amp_bins = []
+        # for i_chunk in range(n_chunks):
+        #     start = i_chunk * self.chunk_size
+        #     end = (i_chunk + 1) * self.chunk_size
+        #     _amp_bins = pha_masks[start:end] * amp[start:end]
+        #     amp_bins.append(_amp_bins.cpu())
+        # amp_bins = torch.cat(amp_bins)
+        # amp_bins = self.dh_pha.unfit(amp_bins)
+        # pha_masks = self.dh_pha.unfit(pha_masks)
+        # Takes mean amplitude in each bin
+        amp_sums = amp_bins.sum(dim=i_time, keepdims=True).to(device)
+        counts = pha_masks.sum(dim=i_time, keepdims=True)
+        amp_means = amp_sums / (counts + epsilon)
+
+        amp_probs = amp_means / (amp_means.sum(dim=-1, keepdims=True) + epsilon)
+
+        if self.amp_prob:
+            return amp_probs.detach().cpu()
+
+        """
+        matplotlib.use("TkAgg")
+        fig, ax = scitex.plt.subplots(subplot_kw={'polar': True})
+        yy = amp_probs[0, 0, 0, 0, 0, 0, :].detach().cpu().numpy()
+        xx = ((self.pha_bin_cutoffs[1:] + self.pha_bin_cutoffs[:-1]) / 2).detach().cpu().numpy()
+        ax.bar(xx, yy, width=.1)
+        plt.show()
+        """
+
+        MI = (
+            torch.log(torch.tensor(self.n_bins, device=device) + epsilon)
+            + (amp_probs * (amp_probs + epsilon).log()).sum(dim=-1)
+        ) / torch.log(torch.tensor(self.n_bins, device=device))
+
+        # Squeeze the n_bin dimension
+        MI = MI.squeeze(-1)
+
+        # Takes mean along the n_segments dimension
+        i_segment = -1
+        MI = MI.mean(axis=i_segment)
+
+        if MI.isnan().any():
+            warnings.warn("NaN values detected in Modulation Index calculation.")
+            # raise ValueError(
+            #     "NaN values detected in Modulation Index calculation."
+            # )
+
+        return MI
+
+    @staticmethod
+    def _phase_to_masks(pha, phase_bin_cutoffs):
+        n_bins = int(len(phase_bin_cutoffs) - 1)
+        bin_indices = (
+            (torch.bucketize(pha, phase_bin_cutoffs, right=False) - 1).clamp(
+                0, n_bins - 1
+            )
+        ).long()
+        one_hot_masks = (
+            F.one_hot(
+                bin_indices,
+                num_classes=n_bins,
+            )
+            .bool()
+            .to(pha.device)
+        )
+        return one_hot_masks
+
+
+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 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 = 1
+    device = "cuda"
+
+    # 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)
+
+    m = ModulationIndex(n_bins=18, fp16=True).to(device)
+
+    pac_scitex = m(pha.to(device), amp.to(device))
+
+    # pac_scitex = scitex.dsp.modulation_index(pha, amp).cpu().numpy()
+    i_batch, i_ch = 0, 0
+    pac_scitex = pac_scitex[i_batch, i_ch].squeeze().numpy()
+
+    # Plots
+    fig = scitex.dsp.utils.pac.plot_PAC_scitex_vs_tensorpac(
+        pac_scitex, pac_tp, freqs_pha, freqs_amp
+    )
+    # fig = plot_PAC_scitex_vs_tensorpac(pac_scitex, pac_tp, freqs_pha, freqs_amp)
+    scitex.io.save(fig, CONFIG["SDIR"] + "modulation_index.png")  # plt.show()
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF
+
+"""
+/home/ywatanabe/proj/entrance/scitex/nn/_ModulationIndex.py
+"""
+
+
+# EOF