From 4b1e9be2407517046dcbf9d0eff029fa51d9c043 Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 12:10:10 +0000 Subject: [PATCH 1/8] refactor unittest --- .../{draft-pdf.yml => draft_pdf.yml} | 0 .../{python-app.yml => selfeeg_test.yml} | 9 +- selfeeg/models/encoders.py | 6 +- test/EEGself/_testtools.py | 86 ++ test/EEGself/augmentation/compose_test.py | 163 +-- test/EEGself/augmentation/functional_test.py | 1142 ++++------------- test/EEGself/dataloading/load_test.py | 418 +++--- test/EEGself/losses/losses_test.py | 353 ++--- test/EEGself/models/layers_test.py | 351 ++--- test/EEGself/models/zoo_test.py | 775 ++++------- test/EEGself/ssl/ssl_test.py | 365 ++---- test/EEGself/utils/utils_test.py | 192 ++- test/README.md | 3 +- 13 files changed, 1318 insertions(+), 2545 deletions(-) rename .github/workflows/{draft-pdf.yml => draft_pdf.yml} (100%) rename .github/workflows/{python-app.yml => selfeeg_test.yml} (80%) create mode 100644 test/EEGself/_testtools.py diff --git a/.github/workflows/draft-pdf.yml b/.github/workflows/draft_pdf.yml similarity index 100% rename from .github/workflows/draft-pdf.yml rename to .github/workflows/draft_pdf.yml diff --git a/.github/workflows/python-app.yml b/.github/workflows/selfeeg_test.yml similarity index 80% rename from .github/workflows/python-app.yml rename to .github/workflows/selfeeg_test.yml index 4a8720f..e149e48 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/selfeeg_test.yml @@ -7,12 +7,12 @@ on: paths: - 'selfeeg/**' - 'test/**' - - '.github/workflows/python-app.yml' + - '.github/workflows/selfeeg_test.yml' pull_request: paths: - 'selfeeg/**' - 'test/**' - - '.github/workflows/python-app.yml' + - '.github/workflows/selfeeg_test.yml' concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} @@ -24,7 +24,7 @@ jobs: strategy: matrix: os: [ubuntu-latest, windows-latest, macos-latest] - python-version: ["3.10", "3.11"] + python-version: ["3.11", "3.12"] runs-on: ${{ matrix.os }} steps: - name: Checkout @@ -39,7 +39,8 @@ jobs: - name: Install dependencies run: | python3 -m pip install --upgrade pip - pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu + pip install torch torchvision + pip install torchaudio pip install -r requirements.txt - name: Test with unittest run: python3 -m unittest discover test "*_test.py" diff --git a/selfeeg/models/encoders.py b/selfeeg/models/encoders.py index 3f035d6..b0c11b0 100644 --- a/selfeeg/models/encoders.py +++ b/selfeeg/models/encoders.py @@ -135,13 +135,13 @@ def __init__( # Layer 1 self.conv1 = nn.Conv2d(1, F1, (1, kernLength), padding="same", bias=False) - self.batchnorm1 = nn.BatchNorm2d(F1, False) + self.batchnorm1 = nn.BatchNorm2d(F1) # Layer 2 self.conv2 = DepthwiseConv2d( F1, D, (Chans, 1), padding="valid", bias=False, max_norm=depthwise_max_norm ) - self.batchnorm2 = nn.BatchNorm2d(D * F1, False) + self.batchnorm2 = nn.BatchNorm2d(D * F1) self.elu2 = nn.ELU(alpha=ELUalpha) self.pooling2 = nn.AvgPool2d((1, pool1)) if dropType.lower() == "dropout": @@ -153,7 +153,7 @@ def __init__( self.sepconv3 = SeparableConv2d( D * F1, F2, (1, separable_kernel), bias=False, padding="same" ) - self.batchnorm3 = nn.BatchNorm2d(F2, False) + self.batchnorm3 = nn.BatchNorm2d(F2) self.elu3 = nn.ELU(alpha=ELUalpha) self.pooling3 = nn.AvgPool2d((1, pool2)) if dropType.lower() == "dropout": diff --git a/test/EEGself/_testtools.py b/test/EEGself/_testtools.py new file mode 100644 index 0000000..f962f1e --- /dev/null +++ b/test/EEGself/_testtools.py @@ -0,0 +1,86 @@ +"""Shared helpers for the selfEEG unittest suite.""" + +import itertools +import random + +import torch + +__all__ = ["get_device", "make_grid"] + + +def get_device(probe=None): + """Return the best available torch device (mps > cuda > cpu). + + If a ``probe`` callable is given it is executed on the selected accelerator; + should it fail (e.g. an operator is unsupported on that backend) the function + silently falls back to the cpu. This mirrors the behaviour previously + duplicated in every ``setUpClass``. + """ + if torch.backends.mps.is_available(): + device = torch.device("mps") + elif torch.cuda.is_available(): + device = torch.device("cuda") + else: + return torch.device("cpu") + + try: + if probe is None: + _ = torch.zeros(8, device=device) + 1 + else: + probe(device) + except Exception: + device = torch.device("cpu") + return device + + +def make_grid(pars_dict, max_comb=None, seed=0): + """Build a list of keyword-argument dictionaries from a dict of value lists. + + The full Cartesian product is returned when it contains at most ``max_comb`` + elements (or when ``max_comb`` is ``None``). Otherwise a deterministic subset + is selected that still exercises **every individual value of every parameter + at least once**, then topped up with random combinations up to ``max_comb``. + This keeps the "no crash / no NaN" smoke tests meaningful while avoiding the + combinatorial explosion of testing the whole product. + """ + keys = list(pars_dict) + values = [list(v) for v in pars_dict.values()] + index_product = [range(len(v)) for v in values] + + if max_comb is None: + selected = list(itertools.product(*index_product)) + return [dict(zip(keys, [values[k][ci] for k, ci in enumerate(c)])) for c in selected] + + # Materialize index combinations lazily-aware: only build the full list when + # it is small enough, otherwise sample without enumerating everything. + total = 1 + for v in values: + total *= len(v) + + if total <= max_comb: + selected = list(itertools.product(*index_product)) + else: + rng = random.Random(seed) + needed = {(k, vi) for k, v in enumerate(values) for vi in range(len(v))} + covered = set() + selected = [] + seen = set() + + # Coverage pass: greedily add random combinations that introduce at least + # one not-yet-tested parameter value, guaranteeing full single-value cover. + while needed - covered: + combo = tuple(rng.randrange(len(v)) for v in values) + new = {(k, combo[k]) for k in range(len(keys))} - covered + if new and combo not in seen: + selected.append(combo) + seen.add(combo) + covered |= new + + # Top-up pass: add random unique combinations until reaching max_comb. + while len(selected) < max_comb: + combo = tuple(rng.randrange(len(v)) for v in values) + if combo not in seen: + selected.append(combo) + seen.add(combo) + + return [dict(zip(keys, [values[k][ci] for k, ci in enumerate(c)])) for c in selected] diff --git a/test/EEGself/augmentation/compose_test.py b/test/EEGself/augmentation/compose_test.py index e35b8d6..8068e20 100644 --- a/test/EEGself/augmentation/compose_test.py +++ b/test/EEGself/augmentation/compose_test.py @@ -10,43 +10,35 @@ class TestAugmentationCompose(unittest.TestCase): @classmethod def setUpClass(cls): - print("\n---------------------------") - print("TESTING AUGMENTATION.COMPOSE MODULE") - print("---------------------------") cls.BatchEEG = torch.zeros(16, 32, 1024) cls.BatchEEG += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) cls.Fs = 128 def test_StaticSingleAug(self): - print("Testing static single augmentation...", end="", flush=True) - Aug_scal = aug.StaticSingleAug(aug.scaling, {"value": 2, "batch_equal": True}) + aug_scal = aug.StaticSingleAug(aug.scaling, {"value": 2, "batch_equal": True}) + batch_aug = aug_scal(self.BatchEEG) + self.assertAlmostEqual(batch_aug.min().item(), -2.0, delta=1e-5) + self.assertAlmostEqual(batch_aug.max().item(), 2.0, delta=1e-5) - BatchEEGaug = Aug_scal(self.BatchEEG) - self.assertTrue(((BatchEEGaug.min() + 2.0) < 1e-5).item()) - self.assertTrue(((BatchEEGaug.max() - 2.0) < 1e-5).item()) - - Aug_scal = aug.StaticSingleAug( + aug_scal = aug.StaticSingleAug( aug.scaling, [{"value": 1.5, "batch_equal": False}, {"value": 2.0, "batch_equal": False}], ) + batch_aug1 = aug_scal(self.BatchEEG) + self.assertAlmostEqual(batch_aug1.min().item(), -1.5, delta=1e-4) + self.assertAlmostEqual(batch_aug1.max().item(), 1.5, delta=1e-4) - BatchEEGaug1 = Aug_scal(self.BatchEEG) - self.assertTrue(((BatchEEGaug1.min() + 1.5) < 1e-5).item()) - self.assertTrue(((BatchEEGaug1.max() - 1.49999) < 1e-5).item()) - - BatchEEGaug2 = Aug_scal(self.BatchEEG) - self.assertTrue(((BatchEEGaug2.min() + 2.0) < 1e-5).item()) - self.assertTrue(((BatchEEGaug2.max() - 2.0) < 1e-5).item()) + batch_aug2 = aug_scal(self.BatchEEG) + self.assertAlmostEqual(batch_aug2.min().item(), -2.0, delta=1e-4) + self.assertAlmostEqual(batch_aug2.max().item(), 2.0, delta=1e-4) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug)) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug1)) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug2)) - self.assertFalse(torch.equal(BatchEEGaug1, BatchEEGaug2)) - print(" static single augmentation OK") + self.assertFalse(torch.equal(self.BatchEEG, batch_aug)) + self.assertFalse(torch.equal(self.BatchEEG, batch_aug1)) + self.assertFalse(torch.equal(self.BatchEEG, batch_aug2)) + self.assertFalse(torch.equal(batch_aug1, batch_aug2)) def test_DynamicSingleAug(self): - print("Testing dynamic single augmentation...", end="", flush=True) - Aug_warp = aug.DynamicSingleAug( + aug_warp = aug.DynamicSingleAug( aug.warp_signal, discrete_arg={"batch_equal": [True, False]}, range_arg={ @@ -56,77 +48,53 @@ def test_DynamicSingleAug(self): }, range_type={"segments": True, "stretch_strength": False, "squeeze_strength": False}, ) - BatchEEGaug1 = Aug_warp(self.BatchEEG) - BatchEEGaug2 = Aug_warp(self.BatchEEG) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug1)) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug2)) - self.assertFalse(torch.equal(BatchEEGaug1, BatchEEGaug2)) - print(" dynamic single augmentation OK") + batch_aug1 = aug_warp(self.BatchEEG) + batch_aug2 = aug_warp(self.BatchEEG) + self.assertFalse(torch.equal(self.BatchEEG, batch_aug1)) + self.assertFalse(torch.equal(self.BatchEEG, batch_aug2)) + self.assertFalse(torch.equal(batch_aug1, batch_aug2)) def test_SequentialAug(self): - print("Testing Sequential augmentation...", end="", flush=True) - Aug_scal = aug.StaticSingleAug(aug.scaling, {"value": 2, "batch_equal": True}) - Sequence1 = aug.SequentialAug(Aug_scal, aug.flip_vertical) - BatchEEGaug1 = Sequence1(self.BatchEEG) - BatchEEGaug2 = aug.scaling(self.BatchEEG, 2) - - # check that augmentation has been performed - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug1)) - - # check that scaling has been performed - self.assertTrue(((BatchEEGaug1.min() + 2.0) < 1e-5).item()) - self.assertTrue(((BatchEEGaug1.max() - 2.0) < 1e-5).item()) - - # check that flip has been performed - self.assertTrue(torch.equal(BatchEEGaug1, BatchEEGaug2 * (-1))) - print(" Sequential augmentation OK") + aug_scal = aug.StaticSingleAug(aug.scaling, {"value": 2, "batch_equal": True}) + sequence = aug.SequentialAug(aug_scal, aug.flip_vertical) + batch_aug1 = sequence(self.BatchEEG) + batch_aug2 = aug.scaling(self.BatchEEG, 2) + + self.assertFalse(torch.equal(self.BatchEEG, batch_aug1)) + # scaling has been applied + self.assertAlmostEqual(batch_aug1.min().item(), -2.0, delta=1e-5) + self.assertAlmostEqual(batch_aug1.max().item(), 2.0, delta=1e-5) + # vertical flip has been applied on top of the scaling + self.assertTrue(torch.equal(batch_aug1, batch_aug2 * (-1))) def test_CircularAug(self): - print("Testing Circular augmentation...", end="", flush=True) - Circular = aug.CircularAug(aug.flip_vertical, aug.identity) - BatchEEGaug1 = Circular(self.BatchEEG) - self.assertTrue(torch.equal(BatchEEGaug1, self.BatchEEG * (-1))) - BatchEEGaug1 = Circular(self.BatchEEG) - self.assertTrue(torch.equal(BatchEEGaug1, self.BatchEEG)) - - # repeat to check Circular calls - BatchEEGaug1 = Circular(self.BatchEEG) - self.assertTrue(torch.equal(BatchEEGaug1, self.BatchEEG * (-1))) - BatchEEGaug1 = Circular(self.BatchEEG) - self.assertTrue(torch.equal(BatchEEGaug1, self.BatchEEG)) - print(" Circular augmentation OK") + circular = aug.CircularAug(aug.flip_vertical, aug.identity) + # two full cycles to check the round-robin behaviour + for _ in range(2): + self.assertTrue(torch.equal(circular(self.BatchEEG), self.BatchEEG * (-1))) + self.assertTrue(torch.equal(circular(self.BatchEEG), self.BatchEEG)) def test_RandomAug(self): - print("Testing Random augmentation...", end="", flush=True) - Aug_scal = aug.StaticSingleAug(aug.scaling, {"value": 2, "batch_equal": True}) - Sequence2 = aug.RandomAug(Aug_scal, aug.flip_vertical, p=[0.7, 0.3], return_index=True) + aug_scal = aug.StaticSingleAug(aug.scaling, {"value": 2, "batch_equal": True}) + sequence = aug.RandomAug(aug_scal, aug.flip_vertical, p=[0.7, 0.3], return_index=True) counter = [0, 0] - N = 10000 + N = 3000 np.random.seed(1234) - for i in range(N): - BatchEEGaug, idx = Sequence2(self.BatchEEG) + for _ in range(N): + _, idx = sequence(self.BatchEEG) counter[idx] += 1 - counter[0] /= N - counter[1] /= N - self.assertTrue(abs(counter[0] - 0.7) < 1e-2) - self.assertTrue(abs(counter[1] - 0.3) < 1e-2) - print(" Random augmentation OK") + self.assertAlmostEqual(counter[0] / N, 0.7, delta=3e-2) + self.assertAlmostEqual(counter[1] / N, 0.3, delta=3e-2) def test_AugmentationComposition(self): - print( - "Testing final augmentation composition based on all previous classes...", - end="", - flush=True, - ) - # DEFINE AUGMENTER - # FIRST RANDOM SELECTION: APPLY FLIP OR CHANGE REFERENCE OR NOTHING - AUG_flipv = aug.StaticSingleAug(aug.flip_vertical) - AUG_flipr = aug.StaticSingleAug(aug.flip_horizontal) - AUG_id = aug.StaticSingleAug(aug.identity) - Sequence1 = aug.RandomAug(AUG_id, AUG_flipv, AUG_flipr, p=[0.5, 0.25, 0.25]) - - # SECOND RANDOM SELECTION: ADD SOME NOISE - AUG_band = aug.DynamicSingleAug( + # First random selection: flip vertical / flip horizontal / identity + flipv = aug.StaticSingleAug(aug.flip_vertical) + flipr = aug.StaticSingleAug(aug.flip_horizontal) + identity = aug.StaticSingleAug(aug.identity) + block1 = aug.RandomAug(identity, flipv, flipr, p=[0.5, 0.25, 0.25]) + + # Second random selection: add band noise or an eye artifact + band = aug.DynamicSingleAug( aug.add_band_noise, discrete_arg={ "bandwidth": ["delta", "theta", "alpha", "beta", (30, 49)], @@ -134,21 +102,22 @@ def test_AugmentationComposition(self): "noise_range": 0.1, }, ) - Aug_eye = aug.DynamicSingleAug( + eye = aug.DynamicSingleAug( aug.add_eeg_artifact, discrete_arg={"Fs": self.Fs, "artifact": "eye", "batch_equal": False}, range_arg={"amplitude": [0.1, 0.5]}, range_type={"amplitude": False}, ) - Sequence2 = aug.RandomAug(AUG_band, Aug_eye, return_index=True) - # THIRD RANDOM SELECTION: CROP OR RANDOM PERMUTATION - AUG_crop = aug.DynamicSingleAug( + block2 = aug.RandomAug(band, eye, return_index=True) + + # Third random selection: crop-and-resize or warp + crop = aug.DynamicSingleAug( aug.crop_and_resize, discrete_arg={"batch_equal": False}, range_arg={"N_cut": [1, 4], "segments": [10, 15]}, range_type={"N_cut": True, "segments": True}, ) - Aug_warp = aug.DynamicSingleAug( + warp = aug.DynamicSingleAug( aug.warp_signal, discrete_arg={"batch_equal": [True, False]}, range_arg={ @@ -158,15 +127,13 @@ def test_AugmentationComposition(self): }, range_type={"segments": True, "stretch_strength": False, "squeeze_strength": False}, ) - Sequence3 = aug.RandomAug(AUG_crop, Aug_warp, return_index=True) - - # FINAL AUGMENTER: SEQUENCE OF THE THREE RANDOM LISTS - Augmenter = aug.SequentialAug(Sequence1, Sequence2, Sequence3) - BatchEEGaug1 = Augmenter(self.BatchEEG) - BatchEEGaug2 = Augmenter(self.BatchEEG) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug1)) - self.assertFalse(torch.equal(self.BatchEEG, BatchEEGaug2)) - print(" final augmentation composition OK") + block3 = aug.RandomAug(crop, warp, return_index=True) + + augmenter = aug.SequentialAug(block1, block2, block3) + batch_aug1 = augmenter(self.BatchEEG) + batch_aug2 = augmenter(self.BatchEEG) + self.assertFalse(torch.equal(self.BatchEEG, batch_aug1)) + self.assertFalse(torch.equal(self.BatchEEG, batch_aug2)) if __name__ == "__main__": diff --git a/test/EEGself/augmentation/functional_test.py b/test/EEGself/augmentation/functional_test.py index 2df53bb..7ee702d 100644 --- a/test/EEGself/augmentation/functional_test.py +++ b/test/EEGself/augmentation/functional_test.py @@ -1,382 +1,177 @@ +import math import os import sys import unittest -sys.path.append(os.getcwd().split("/test")[0]) -import itertools -import math - import numpy as np import torch from scipy.signal import periodogram +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from EEGself._testtools import get_device, make_grid + from selfeeg import augmentation as aug -class TestAugmentationFunctional(unittest.TestCase): +def _sine(*shape): + return torch.zeros(*shape) + torch.sin(torch.linspace(0, 8 * torch.pi, shape[-1])) - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds + +class TestAugmentationFunctional(unittest.TestCase): @classmethod def setUpClass(cls): - if torch.backends.mps.is_available(): - cls.device = torch.device("mps") - elif torch.cuda.is_available(): - cls.device = torch.device("cuda") - else: - cls.device = torch.device("cpu") - - if cls.device.type != "cpu": - try: - xx = torch.randn(1024).to(device=cls.device) - xx = aug.add_band_noise(xx, "theta", 128) - except Exception: - cls.device = torch.device("cpu") - - device = cls.device - print("\n---------------------------") - print("TESTING AUGMENTATION.FUNCTIONAL MODULE") + cls.device = get_device( + probe=lambda dev: aug.add_band_noise(torch.randn(1024, device=dev), "theta", 128) + ) + # Tensors with 1, 2, 3 and 4 dimensions, plus numpy / gpu counterparts. + shapes = {1: (1024,), 2: (32, 1024), 3: (2, 32, 1024), 4: (32, 2, 32, 1024)} + cls._t = {k: _sine(*s) for k, s in shapes.items()} + cls._np = {k: v.numpy() for k, v in cls._t.items()} + cls._g = {} if cls.device.type != "cpu": - print("Found gpu device: testing module on it") + cls._g = {k: v.clone().to(cls.device) for k, v in cls._t.items()} + + def _inputs(self, dims=(1, 2, 3, 4), numpy=True, gpu=True, mps_ok=True): + xs = [self._t[d] for d in dims] + if numpy: + xs += [self._np[d] for d in dims] + if gpu and self._g and (mps_ok or self.device.type != "mps"): + xs += [self._g[d] for d in dims] + return xs + + def _assert_valid(self, xin, xout, changed=True): + if isinstance(xout, torch.Tensor): + self.assertEqual(torch.isnan(xout).sum().item(), 0) + same = torch.equal(xin, xout) else: - print("Didn't found cuda device: testing module on cpu") - print("---------------------------") - dims = (32, 2, 32, 1024) - pi = torch.pi - cls.x1 = torch.zeros(*dims[-1:]) + torch.sin(torch.linspace(0, 8 * pi, 1024)) - cls.x2 = torch.zeros(*dims[-2:]) + torch.sin(torch.linspace(0, 8 * pi, 1024)) - cls.x3 = torch.zeros(*dims[-3:]) + torch.sin(torch.linspace(0, 8 * pi, 1024)) - cls.x4 = torch.zeros(*dims) + torch.sin(torch.linspace(0, 8 * pi, 1024)) - cls.x1np = cls.x1.numpy() - cls.x2np = cls.x2.numpy() - cls.x3np = cls.x3.numpy() - cls.x4np = cls.x4.numpy() - if device.type != "cpu": - cls.x1gpu = torch.clone(cls.x1).to(device=device) - cls.x2gpu = torch.clone(cls.x2).to(device=device) - cls.x3gpu = torch.clone(cls.x3).to(device=device) - cls.x4gpu = torch.clone(cls.x4).to(device=device) + self.assertEqual(np.isnan(xout).sum(), 0) + same = np.array_equal(xin, xout) + self.assertEqual(same, not changed) + + @staticmethod + def _descr(args): + out = {} + for k, v in args.items(): + if torch.is_tensor(v): + out["x"] = f"T{tuple(v.shape)}/{v.device.type}" + elif isinstance(v, np.ndarray): + out["x"] = f"np{tuple(v.shape)}" + else: + out[k] = v + return out + + def _smoke(self, fn, params, x_list, changed=True, max_comb=32, prep=None, extra=None): + grid = make_grid({**params, "x": list(x_list)}, max_comb=max_comb) + for args in grid: + if prep is not None: + prep(args) + with self.subTest(fn=fn.__name__, **self._descr(args)): + xaug = fn(**args) + self._assert_valid(args["x"], xaug, changed) + if extra is not None: + extra(args, xaug) + + @staticmethod + def _bounded(args, xaug): + assert int((xaug > 1e2).sum()) == 0 and int((xaug < -1e2).sum()) == 0 + + # ------------------------------------------------------------------ tests def test_identity(self): - print("Testing identity...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np] - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.identity(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertTrue(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertTrue(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = {"x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.identity(**i) - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertTrue(torch.equal(i["x"], xaug)) - print(" identity OK: tested", N + len(aug_args), "combinations of input arguments") + self._smoke(aug.identity, {}, self._inputs(), changed=False) def test_shift_vertical(self): - print("Testing shift vertical...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "value": [1, 2.0, 4], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.shift_vertical(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = {"x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], "value": [1, 2.0, 4]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.shift_vertical(**i) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) - xaug = aug.shift_vertical(x, 4) - self.assertTrue(torch.equal(x + 4, xaug)) # should return True - print(" shift vertical OK: tested", N + len(aug_args), "combinations of input arguments") + self._smoke(aug.shift_vertical, {"value": [1, 2.0, 4]}, self._inputs()) + xaug = aug.shift_vertical(self._t[4], 4) + self.assertTrue(torch.equal(self._t[4] + 4, xaug)) def test_shift_horizontal(self): - print("Testing shift horizontal...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128], - "shift_time": [0.5, 1, 2.0], - "forward": [None, True, False], - "random_shift": [False, True], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - # change batch equal to avoid function print - if not (i["batch_equal"]): - if not (i["random_shift"] or (i["forward"] is None)): - i["batch_equal"] = True - xaug = aug.shift_horizontal(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + def prep(a): + if not a["batch_equal"] and not (a["random_shift"] or a["forward"] is None): + a["batch_equal"] = True + + self._smoke( + aug.shift_horizontal, + { "Fs": [128], "shift_time": [0.5, 1, 2.0], "forward": [None, True, False], "random_shift": [False, True], "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - if not (i["batch_equal"]): - if not (i["random_shift"] or (i["forward"] is None)): - i["batch_equal"] = True - xaug = aug.shift_horizontal(**i) - - xaug = aug.shift_horizontal(self.x4, 64, 1, True) - self.assertTrue(xaug[..., 0:64].sum() == 0) - self.assertFalse(xaug[..., 65].sum() == 0) - print( - " shift vertical OK: tested", N + len(aug_args) + 1, "combinations of input arguments" + }, + self._inputs(), + prep=prep, + max_comb=48, ) + xaug = aug.shift_horizontal(self._t[4], 64, 1, True) + self.assertEqual(xaug[..., 0:64].sum(), 0) + self.assertNotEqual(xaug[..., 65].sum(), 0) def test_shift_frequency(self): - print("Testing shift frequency...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128], - "shift_freq": [1.35, 2, 4.12], - "forward": [None, True, False], - "random_shift": [False, True], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - if not (i["batch_equal"]): - if not (i["random_shift"] or (i["forward"] is None)): - i["batch_equal"] = True - xaug = aug.shift_frequency(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if (self.device.type != "cpu") and (self.device.type != "mps"): - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + def prep(a): + if not a["batch_equal"] and not (a["random_shift"] or a["forward"] is None): + a["batch_equal"] = True + + self._smoke( + aug.shift_frequency, + { "Fs": [128], "shift_freq": [1.35, 2, 4.12], "forward": [None, True, False], "random_shift": [False, True], "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - if not (i["batch_equal"]): - if not (i["random_shift"] or (i["forward"] is None)): - i["batch_equal"] = True - xaug = aug.shift_frequency(**i) - + }, + self._inputs(mps_ok=False), + prep=prep, + max_comb=48, + ) Fs = 128 - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 48 * torch.pi, 1024)) - x = x + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + x = _sine(16, 32, 1024) + torch.sin(torch.linspace(0, 48 * torch.pi, 1024)) xaug = aug.shift_frequency(x, 10, Fs, True) - f, per1 = periodogram(x[0, 0], fs=Fs) + per1 = periodogram(x[0, 0], fs=Fs)[1] per2 = periodogram(xaug[0, 0], fs=Fs)[1] self.assertTrue(math.isclose(per1[4], per2[84], rel_tol=1e-5)) self.assertTrue(math.isclose(per1[24], per2[104], rel_tol=1e-5)) - print(" shift frequency OK: tested", N + len(aug_args), "combinations of input arguments") def test_phase_swap(self): - print("Testing phase swap...", end="", flush=True) - aug_args = {"x": [self.x3, self.x3np]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.phase_swap(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = {"x": [self.x3gpu]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.phase_swap(**i) - - print(" phase_swap OK: tested", N + len(aug_args), "combinations of input arguments") + self._smoke(aug.phase_swap, {}, self._inputs(dims=(3,))) def test_flip_vertical(self): - print("Testing flip vertical...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np] - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.flip_vertical(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = {"x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.flip_vertical(**i) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * np.pi, 1024)) - xaug = aug.flip_vertical(x) - self.assertTrue(torch.equal(xaug, x * (-1))) # should return True - print(" flip vertical OK: tested", N + len(aug_args), "combinations of input arguments") + self._smoke(aug.flip_vertical, {}, self._inputs()) + xaug = aug.flip_vertical(self._t[4]) + self.assertTrue(torch.equal(xaug, self._t[4] * (-1))) def test_flip_horizontal(self): - print("Testing flip horizontal...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np] - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.flip_horizontal(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = {"x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.flip_horizontal(**i) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) - xaug = aug.flip_horizontal(x) - self.assertTrue(torch.equal(xaug, torch.flip(x, [len(x.shape) - 1]))) - print(" flip horizontal OK: tested", N + len(aug_args), "combinations of input arguments") + self._smoke(aug.flip_horizontal, {}, self._inputs()) + xaug = aug.flip_horizontal(self._t[4]) + self.assertTrue(torch.equal(xaug, torch.flip(self._t[4], [self._t[4].ndim - 1]))) def test_gaussian_noise(self): - print("Testing gaussian noise...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "mean": [0, 1, 2.5], - "std": [1.35, 2, 0.72], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_gaussian_noise(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], - "mean": [0, 1, 2.5], - "std": [1.35, 2, 0.72], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_gaussian_noise(**i) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + self._smoke( + aug.add_gaussian_noise, + {"mean": [0, 1, 2.5], "std": [1.35, 2, 0.72]}, + self._inputs(), + max_comb=24, + ) + x = _sine(16, 32, 1024) xaug, noise = aug.add_gaussian_noise(x, 0.1, get_noise=True) self.assertTrue(math.isclose(noise.std(), 0.1, rel_tol=1e-2)) self.assertTrue(math.isclose(xaug.mean(), 0, rel_tol=1e-4, abs_tol=1e-3)) - print(" gaussian noise OK: tested", N + len(aug_args), "combinations of input arguments") def test_add_noise_SNR(self): - print("Testing noise SNR...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "target_snr": [1, 2, 5], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_noise_SNR(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], - "target_snr": [1, 2, 5], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_noise_SNR(**i) - - x = torch.zeros(16, 32, 512) + torch.sin(torch.linspace(0, 8 * np.pi, 512)) + self._smoke(aug.add_noise_SNR, {"target_snr": [1, 2, 5]}, self._inputs(), max_comb=24) + x = _sine(16, 32, 512) xaug, noise = aug.add_noise_SNR(x, 10, get_noise=True) - SNR = 10 * torch.log10(((x**2).sum().mean()) / ((noise**2).sum().mean())) - self.assertTrue(math.isclose(SNR, 10, rel_tol=1e-1)) - print(" noise SNR OK: tested", N + len(aug_args), "combinations of input arguments") + snr = 10 * torch.log10(((x**2).sum().mean()) / ((noise**2).sum().mean())) + self.assertTrue(math.isclose(snr, 10, rel_tol=1e-1)) def test_add_band_noise(self): - print("Testing band noise...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "bandwidth": [ - ["theta", "gamma"], - [(1, 10), (15, 18)], - [4, 50], - 50, - ["theta", (10, 20), 50], - ], - "samplerate": [128], - "noise_range": [None, 2, 1.5], - "std": [None, 1.4, 1.23], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_band_noise(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.add_band_noise, + { "bandwidth": [ ["theta", "gamma"], [(1, 10), (15, 18)], @@ -387,431 +182,185 @@ def test_add_band_noise(self): "samplerate": [128], "noise_range": [None, 2, 1.5], "std": [None, 1.4, 1.23], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_band_noise(**i) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(), + max_comb=40, + ) + x = _sine(16, 32, 1024) xaug, noise = aug.add_band_noise(x, "beta", 128, noise_range=0.2, get_noise=True) f, per = periodogram(noise, 128) index = np.where(per > 1e-12)[0] - self.assertTrue(len(np.where(((f[index] < 13) | (f[index] > 30)))[0]) == 0) - print(" band noise OK: tested", N + len(aug_args), "combinations of input arguments") + self.assertEqual(len(np.where(((f[index] < 13) | (f[index] > 30)))[0]), 0) def test_scaling(self): - print("Testing scaling...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "value": [None, 1.5, 2, 0.5], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.scaling(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], - "value": [None, 1.5, 2, 0.5], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.scaling(**i) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + self._smoke( + aug.scaling, + {"value": [None, 1.5, 2, 0.5], "batch_equal": [True, False]}, + self._inputs(), + max_comb=24, + ) + x = _sine(16, 32, 1024) xaug = aug.scaling(x, 1.5) - self.assertTrue(xaug.max() == x.max() * 1.5) # should return True - self.assertTrue(xaug.min() == x.min() * 1.5) # should return True - print(" scaling OK: tested", N + len(aug_args), "combinations of input arguments") + self.assertEqual(xaug.max(), x.max() * 1.5) + self.assertEqual(xaug.min(), x.min() * 1.5) def test_random_slope_scale(self): - print("Testing random slope scale...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "min_scale": [0.7, 0.9], - "max_scale": [1.2, 1.5], - "batch_equal": [True, False], - "keep_memory": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - if i["batch_equal"] and len(i["x"].shape) < 2: - i["batch_equal"] = False - xaug = aug.random_slope_scale(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + def prep(a): + if a["batch_equal"] and len(a["x"].shape) < 2: + a["batch_equal"] = False + + self._smoke( + aug.random_slope_scale, + { "min_scale": [0.7, 0.9], "max_scale": [1.2, 1.5], "batch_equal": [True, False], "keep_memory": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - if i["batch_equal"] and len(i["x"].shape) < 2: - i["batch_equal"] = False - xaug = aug.random_slope_scale(**i) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(), + prep=prep, + max_comb=24, + ) + x = _sine(16, 32, 1024) xaug = aug.random_slope_scale(x) diff1 = torch.abs(xaug[0, 0, 1:] - xaug[0, 0, :-1]) diff2 = torch.abs(x[0, 0, 1:] - x[0, 0, :-1]) self.assertEqual( torch.logical_or(diff1 <= (diff2 * 1.2), diff1 >= (diff2 * 0.9)).sum(), 1023 ) - print( - " random slope scale OK: tested", N + len(aug_args), "combinations of input arguments" - ) def test_random_FT_phase(self): - print("Testing random FT phase...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "value": [0.2, 0.5, 0.75], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.random_FT_phase(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if (self.device.type != "cpu") and (self.device.type != "mps"): - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], - "value": [0.2, 0.5, 0.75], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.random_FT_phase(**i) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + self._smoke( + aug.random_FT_phase, + {"value": [0.2, 0.5, 0.75], "batch_equal": [True, False]}, + self._inputs(mps_ok=False), + max_comb=24, + ) + x = _sine(16, 32, 1024) xaug = aug.random_FT_phase(x, 0.8) phase_shift = torch.arccos(2 * ((x[0, 0, 0:512] * xaug[0, 0, :512]).mean())) a = torch.sin(torch.linspace(0, 8 * torch.pi, 1024) + phase_shift) if (a[0] - xaug[0, 0, 0]).abs() > 0.1: a = torch.sin(torch.linspace(0, 8 * torch.pi, 1024) - phase_shift) self.assertTrue((a - xaug[0, 0]).mean() < 1e-3) - print(" random FT phase OK: tested", N + len(aug_args), "combinations of input arguments") def test_moving_average(self): - print("Testing moving average...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "order": [3, 5, 9], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.moving_avg(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = {"x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], "order": [3, 5, 9]} - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.moving_avg(**i) - + self._smoke(aug.moving_avg, {"order": [3, 5, 9]}, self._inputs(), max_comb=24) x = torch.randn(16, 32, 1024) xaug = aug.moving_avg(x, 5) self.assertTrue(math.isclose(x[0, 0, 5 : 5 + 5].sum() / 5, xaug[0, 0, 7], rel_tol=1e-5)) - print(" moving average OK: tested", N + len(aug_args), "combinations of input arguments") def test_filter_lowpass(self): - print("Testing lowpass filter...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128, 256], - "Wp": [30], - "Ws": [50], - "rp": [-20 * np.log10(0.90)], - "rs": [-20 * np.log10(0.15)], - "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_lowpass(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.filter_lowpass, + { "Fs": [128, 256], "Wp": [30], "Ws": [50], - "rp": [-20 * np.log10(0.95)], + "rp": [-20 * np.log10(0.90)], "rs": [-20 * np.log10(0.15)], "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_lowpass(**i) - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(), + max_comb=32, + extra=self._bounded, + ) + x = _sine(16, 32, 1024) x += torch.sin(torch.linspace(0, 48 * 2 * torch.pi, 1024)) x += torch.sin(torch.linspace(0, 256 * 2 * torch.pi, 1024)) - f, per1 = periodogram(x[0, 0], 128) xaug = aug.filter_lowpass(x, 128, 20, 30) f, per2 = periodogram(xaug[0, 0], 128) self.assertTrue(np.isclose(np.max(per2[f > 30]), 0, rtol=1e-04, atol=1e-04)) - print(" lowpass filter OK: tested", N + len(aug_args), "combinations of input arguments") def test_filter_highpass(self): - print("Testing highpass filter...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128, 256], - "Wp": [40], - "Ws": [20], - "rp": [-20 * np.log10(0.9)], - "rs": [-20 * np.log10(0.15)], - "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_highpass(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.filter_highpass, + { "Fs": [128, 256], "Wp": [40], "Ws": [20], - "rp": [-20 * np.log10(0.95)], + "rp": [-20 * np.log10(0.9)], "rs": [-20 * np.log10(0.15)], "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_highpass(**i) - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(), + max_comb=32, + extra=self._bounded, + ) + x = _sine(16, 32, 1024) x += torch.sin(torch.linspace(0, 48 * 2 * torch.pi, 1024)) x += torch.sin(torch.linspace(0, 256 * 2 * torch.pi, 1024)) - f, per1 = periodogram(x[0, 0], 128) xaug = aug.filter_highpass(x, 128, 20, 30) f, per2 = periodogram(xaug[0, 0], 128) self.assertTrue(np.isclose(np.max(per2[f < 20]), 0, rtol=1e-04, atol=1e-04)) - print(" highpass filter OK: tested", N + len(aug_args), "combinations of input arguments") def test_filter_bandpass(self): - print("Testing bandpass filter...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128, 256], - "eeg_band": ["delta", "alpha", "beta", "gamma", "gamma_low"], - "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_bandpass(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.filter_bandpass, + { "Fs": [128, 256], - "eeg_band": [None, "delta", "alpha", "beta", "gamma", "gamma_low"], + "eeg_band": ["delta", "alpha", "beta", "gamma", "gamma_low"], "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_bandpass(**i) - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(), + max_comb=32, + extra=self._bounded, + ) + x = _sine(16, 32, 1024) x += torch.sin(torch.linspace(0, 48 * 2 * torch.pi, 1024)) x += torch.sin(torch.linspace(0, 256 * 2 * torch.pi, 1024)) - f, per1 = periodogram(x[0, 0], 128) xaug = aug.filter_bandpass(x, 128, [13, 22], [5, 27]) f, per2 = periodogram(xaug[0, 0], 128) self.assertTrue(np.isclose(np.max(per2[f < 5]), 0, rtol=1e-04, atol=1e-04)) self.assertTrue(np.isclose(np.max(per2[f > 27]), 0, rtol=1e-04, atol=1e-04)) - print(" bandpass filter OK: tested", N + len(aug_args), "combinations of input arguments") def test_filter_bandstop(self): - print("Testing bandstop filter...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128, 256], - "eeg_band": [None, "delta", "theta", "alpha", "beta", "gamma", "gamma_low"], - "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_bandstop(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.filter_bandstop, + { "Fs": [128, 256], "eeg_band": [None, "delta", "theta", "alpha", "beta", "gamma", "gamma_low"], "filter_type": ["butter", "ellip", "cheby1", "cheby2"], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.filter_bandstop(**i) - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - self.assertTrue((xaug > 1e2).sum() == 0) - self.assertTrue((xaug < -1e2).sum() == 0) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(), + max_comb=32, + extra=self._bounded, + ) + x = _sine(16, 32, 1024) x += torch.sin(torch.linspace(0, 48 * 2 * torch.pi, 1024)) x += torch.sin(torch.linspace(0, 256 * 2 * torch.pi, 1024)) - f, per1 = periodogram(x[0, 0], 128) xaug = aug.filter_bandpass(x, 128, [13, 22], [5, 27]) f, per2 = periodogram(xaug[0, 0], 128) self.assertTrue( np.isclose(np.max(per2[np.logical_and(f > 5, f < 27)]), 0, rtol=1e-04, atol=1e-04) ) - print(" bandstop filter OK: tested", N + len(aug_args), "combinations of input arguments") def test_permute_channels(self): - print("Testing permute channels...", end="", flush=True) channel_map = [ - "FP1", - "AF3", - "F1", - "F3", - "FC5", - "FC3", - "FC1", - "C1", - "C5", - "TP7", - "CP5", - "CP3", - "CP1", - "P7", - "PO7", - "POZ", - "PZ", - "FPZ", - "FP2", - "AFZ", - "FZ", - "F2", - "F4", - "F6", - "FT8", - "C4", - "T8", - "TP8", - "CP6", - "CP4", - "CP2", - "PO8", + "FP1", "AF3", "F1", "F3", "FC5", "FC3", "FC1", "C1", "C5", "TP7", "CP5", + "CP3", "CP1", "P7", "PO7", "POZ", "PZ", "FPZ", "FP2", "AFZ", "FZ", "F2", + "F4", "F6", "FT8", "C4", "T8", "TP8", "CP6", "CP4", "CP2", "PO8", ] - aug_args = { - "x": [self.x2, self.x3, self.x4, self.x2np, self.x3np, self.x4np], - "chan2shuf": [-1, 5, 10], - "mode": ["random", "network"], - "chan_net": ["DMN", "FPN", ["DMN", "FPN"], "all"], - "batch_equal": [True, False], - "channel_map": [channel_map], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - i["x"] = i["x"] - if isinstance(i["x"], torch.Tensor): - i["x"] += torch.randn(32, 1) - else: - i["x"] += np.random.randn(32, 1) - xaug = aug.permute_channels(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - if torch.equal(i["x"], xaug): - print(i["x"][:, 10]) - print(xaug[:, 10]) - print(i) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x2gpu, self.x3gpu, self.x4gpu], + # Channels must be distinct for a permutation to change the signal. + torch.manual_seed(0) + offset = torch.randn(32, 1) + inputs = [self._t[d] + offset for d in (2, 3, 4)] + inputs += [(self._t[d] + offset).numpy() for d in (2, 3, 4)] + if self._g: + inputs += [self._g[d] + offset.to(self.device) for d in (2, 3, 4)] + + self._smoke( + aug.permute_channels, + { "chan2shuf": [-1, 5, 10], "mode": ["random", "network"], "chan_net": ["DMN", "FPN", ["DMN", "FPN"], "all"], "batch_equal": [True, False], "channel_map": [channel_map], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - i["x"] = i["x"] + torch.randn(32, 1).to(device=self.device) - xaug = aug.permute_channels(**i) + }, + inputs, + max_comb=40, + ) x = torch.zeros(61, 4) + torch.arange(61).reshape(61, 1) xaug = aug.permute_channels(x, 10) self.assertEqual((x[:, 0] != xaug[:, 0]).sum(), 10) @@ -820,258 +369,100 @@ def test_permute_channels(self): a = np.intersect1d(eeg1010, chan2per, return_indices=True)[1] b = torch.from_numpy(np.delete(np.arange(61), a)) xaug2 = aug.permute_channels(x, 50, mode="network", chan_net=["DMN", "VFN"]) - self.assertTrue(((x[:, 0] != xaug2[:, 0]).sum()) == 50) - self.assertTrue(((x[b, 0] == xaug2[b, 0]).sum()) == len(b)) - print( - " permute channels OK: tested", N + len(aug_args), "combinations of input arguments" - ) + self.assertEqual((x[:, 0] != xaug2[:, 0]).sum(), 50) + self.assertEqual((x[b, 0] == xaug2[b, 0]).sum(), len(b)) def test_permutation_signal(self): - print("Testing permute signal...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "segments": [10, 15, 20], - "seg_to_per": [-1, 2, 5, 8], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.permutation_signal(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], - "segments": [10, 15, 20], - "seg_to_per": [-1, 2, 5, 8], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.permutation_signal(**i) - torch.manual_seed(1234) - x = torch.ones(16, 32, 1024) * 2 - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + self._smoke( + aug.permutation_signal, + {"segments": [10, 15, 20], "seg_to_per": [-1, 2, 5, 8], "batch_equal": [True, False]}, + self._inputs(), + max_comb=32, + ) + x = torch.ones(16, 32, 1024) * 2 + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) xaug = aug.masking(x, 3, 0.5) self.assertTrue( torch.isclose( - ((xaug[0, 0] == 0).sum() / len(xaug[0, 0])), + (xaug[0, 0] == 0).sum() / len(xaug[0, 0]), torch.tensor([0.5]), rtol=1e-8, atol=1e-8, ) ) a = xaug[0, 0] == 0 - self.assertTrue((a[:-1].ne(a[1:])).sum() == 6) # should return True - print(" permute signal OK: tested", N + len(aug_args), "combinations of input arguments") + self.assertEqual((a[:-1].ne(a[1:])).sum(), 6) def test_warp_signal(self): - print("Testing warp signal...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "segments": [15], - "stretch_strength": [2, 1.5], - "squeeze_strength": [0.4, 0.8], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.warp_signal(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.warp_signal, + { "segments": [15], "stretch_strength": [2, 1.5], "squeeze_strength": [0.4, 0.8], "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.warp_signal(**i) - print(" warp signal OK: tested", N + len(aug_args), "combinations of input arguments") + }, + self._inputs(), + max_comb=24, + ) def test_crop_and_resize(self): - print("Testing crop and resize...", end="", flush=True) - aug_args = { - "x": [self.x2, self.x3, self.x4, self.x2np, self.x3np, self.x4np], - "segments": [15], - "N_cut": [1, 5], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.crop_and_resize(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x2gpu, self.x3gpu, self.x4gpu], - "segments": [15], - "N_cut": [1, 5], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.crop_and_resize(**i) - print(" crop and resize OK: tested", N + len(aug_args), "combinations of input arguments") + self._smoke( + aug.crop_and_resize, + {"segments": [15], "N_cut": [1, 5], "batch_equal": [True, False]}, + self._inputs(dims=(2, 3, 4)), + max_comb=24, + ) def test_change_ref(self): - print("Testing change reference...", end="", flush=True) - aug_args = { - "x": [self.x2, self.x3, self.x4, self.x2np, self.x3np, self.x4np], - "mode": ["chan", "avg"], - "reference": [None, 5], - "exclude_from_ref": [None, 9, [9, 10]], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.change_ref(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.change_ref, + { "mode": ["chan", "avg"], "reference": [None, 5], "exclude_from_ref": [None, 9, [9, 10]], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.change_ref(**i) - - torch.manual_seed(1234) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + }, + self._inputs(dims=(2, 3, 4)), + max_comb=24, + ) + x = _sine(16, 32, 1024) x[:, 0, :] = 0.0 xaug = aug.change_ref(x, "channel", 5) - self.assertFalse(x[0, 0].max() != 0 and x[0, 0].min() != 0) # should return False - self.assertTrue( - (xaug[0, [i for i in range(1, 32)]].min().item() == 0 and xaug[0, 0].min().item() != 0) - ) # should return True - print( - " change refeference OK: tested", N + len(aug_args), "combinations of input arguments" - ) + self.assertEqual(xaug[0, list(range(1, 32))].min().item(), 0) + self.assertNotEqual(xaug[0, 0].min().item(), 0) def test_masking(self): - print("Testing masking...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "mask_number": [1, 2, 4], - "masked_ratio": [0.1, 0.2, 0.4], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.masking(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], - "mask_number": [1, 2, 4], - "masked_ratio": [0.1, 0.2, 0.4], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.masking(**i) - - x = torch.ones(16, 32, 1024) * 2 - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) + self._smoke( + aug.masking, + {"mask_number": [1, 2, 4], "masked_ratio": [0.1, 0.2, 0.4], "batch_equal": [True, False]}, + self._inputs(), + max_comb=24, + ) + x = torch.ones(16, 32, 1024) * 2 + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) xaug = aug.masking(x, 3, 0.5) self.assertTrue( torch.isclose( - ((xaug[0, 0] == 0).sum() / len(xaug[0, 0])), + (xaug[0, 0] == 0).sum() / len(xaug[0, 0]), torch.tensor([0.5]), rtol=1e-6, atol=1e-8, ) ) - print(" masking OK: tested", N + len(aug_args), "combinations of input arguments") def test_channel_dropout(self): - print("Testing channel dropout...", end="", flush=True) - aug_args = { - "x": [self.x2, self.x3, self.x4, self.x2np, self.x3np, self.x4np], - "Nchan": [None, 2, 3], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.channel_dropout(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x2gpu, self.x3gpu, self.x4gpu], - "Nchan": [None, 2, 3], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.channel_dropout(**i) + self._smoke( + aug.channel_dropout, + {"Nchan": [None, 2, 3], "batch_equal": [True, False]}, + self._inputs(dims=(2, 3, 4)), + max_comb=24, + ) x = torch.ones(32, 1024) * 2 + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) xaug = aug.channel_dropout(x, 3) - self.assertTrue((xaug[:, 10] == 0).sum() == 3) - print(" channel dropout OK: tested", N + len(aug_args), "combinations of input arguments") + self.assertEqual((xaug[:, 10] == 0).sum(), 3) def test_eeg_artifact(self): - print("Testing eeg artifact...", end="", flush=True) - aug_args = { - "x": [self.x1, self.x2, self.x3, self.x4, self.x1np, self.x2np, self.x3np, self.x4np], - "Fs": [128], - "artifact": [None, "white", "line", "eye", "muscle", "drift", "lost"], - "amplitude": [None, 1], - "line_at_60Hz": [True, False], - "lost_time": [0.5, None], - "drift_slope": [None, 0.2], - "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_eeg_artifact(**i) - if isinstance(xaug, torch.Tensor): - self.assertTrue(torch.isnan(xaug).sum() == 0) - self.assertFalse(torch.equal(i["x"], xaug)) - else: - self.assertTrue(np.isnan(xaug).sum() == 0) - self.assertFalse(np.array_equal(i["x"], xaug)) - N = len(aug_args) - if self.device.type != "cpu": - aug_args = { - "x": [self.x1gpu, self.x2gpu, self.x3gpu, self.x4gpu], + self._smoke( + aug.add_eeg_artifact, + { "Fs": [128], "artifact": [None, "white", "line", "eye", "muscle", "drift", "lost"], "amplitude": [None, 1], @@ -1079,11 +470,10 @@ def test_eeg_artifact(self): "lost_time": [0.5, None], "drift_slope": [None, 0.2], "batch_equal": [True, False], - } - aug_args = self.makeGrid(aug_args) - for i in aug_args: - xaug = aug.add_eeg_artifact(**i) - print(" eeg artifact OK: tested", N + len(aug_args), "combinations of input arguments") + }, + self._inputs(), + max_comb=48, + ) if __name__ == "__main__": diff --git a/test/EEGself/dataloading/load_test.py b/test/EEGself/dataloading/load_test.py index e71c1b9..614edfb 100644 --- a/test/EEGself/dataloading/load_test.py +++ b/test/EEGself/dataloading/load_test.py @@ -1,13 +1,17 @@ -import itertools import os import pickle -import platform -import random +import shutil +import sys import unittest import numpy as np import torch +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from EEGself._testtools import make_grid + from selfeeg import dataloading as dl @@ -39,28 +43,15 @@ def loadEEG(self, path, return_label=False, return_label_array=False): # just to complicate things y = np.array([[y] * 100, [y] * 100]).T return x, y - else: - return x, y - else: - return x + return x, y + return x def transformEEG(self, EEG, value=64): - EEG = EEG[:, :-value] - return EEG - - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds + return EEG[:, :-value] @classmethod def setUpClass(cls): - print("\n---------------------------") - print("TESTING DATALOADING MODULE") - print("the process will generetate two additional folders") - print("---------------------------") - if not (os.path.isdir("tmpsave")): + if not os.path.isdir("tmpsave"): os.mkdir("tmpsave") cls.eegpath = "Simulated_EEG" cls.create_dataset(cls) @@ -70,44 +61,10 @@ def setUpClass(cls): def setUp(self): self.seed = 1234 - random.seed(self.seed) np.random.seed(self.seed) - def test_get_eeg_partition_number(self): - # checks = Table with length 1000 and results in specific setting is correct - print("Testing get_eeg_partition_number...", end="", flush=True) - input_grid = { - "EEGpath": [self.eegpath], - "freq": [self.freq], - "window": [2], - "overlap": [0, 0.15], - "includePartial": [True, False], - "file_format": ["*.pickle"], - "load_function": [self.loadEEG], - "optional_load_fun_args": [[False]], - "transform_function": [None, self.transformEEG], - "optional_transform_fun_args": [None, [32]], - "keep_zero_sample": [True, False], - "save": [True, False], - "save_path": ["tmpsave/results1.csv"], - "verbose": [False], - } - input_grid = self.makeGrid(input_grid) - for i in input_grid: - EEGlen = dl.get_eeg_partition_number(**i) - self.assertEqual(EEGlen.shape[0], 1000) - print( - " get_eeg_partition_number OK: tested", - len(input_grid), - "combination of input arguments", - ) - - def test_get_eeg_split_table(self): - # checks: table has length 1000, - # ratio = 0 means empty set, - # ids when given are splitted corretly - print("Testing get_eeg_split_table (this may take some time)...", end="", flush=True) - EEGlen = dl.get_eeg_partition_number( + def _partition_table(self): + return dl.get_eeg_partition_number( self.eegpath, self.freq, self.window, @@ -117,9 +74,41 @@ def test_get_eeg_split_table(self): optional_load_fun_args=[False], transform_function=self.transformEEG, ) - Labels = np.zeros(EEGlen.shape[0], dtype=int) + + def _labels(self, EEGlen): + labels = np.zeros(EEGlen.shape[0], dtype=int) for i in range(EEGlen.shape[0]): - _, Labels[i] = self.loadEEG(EEGlen.iloc[i]["full_path"], True) + _, labels[i] = self.loadEEG(EEGlen.iloc[i]["full_path"], True) + return labels + + def test_get_eeg_partition_number(self): + grid = make_grid( + { + "EEGpath": [self.eegpath], + "freq": [self.freq], + "window": [2], + "overlap": [0, 0.15], + "includePartial": [True, False], + "file_format": ["*.pickle"], + "load_function": [self.loadEEG], + "optional_load_fun_args": [[False]], + "transform_function": [None, self.transformEEG], + "optional_transform_fun_args": [None, [32]], + "keep_zero_sample": [True, False], + "save": [True, False], + "save_path": ["tmpsave/results1.csv"], + "verbose": [False], + }, + max_comb=16, + ) + for args in grid: + with self.subTest(overlap=args["overlap"], includePartial=args["includePartial"]): + EEGlen = dl.get_eeg_partition_number(**args) + self.assertEqual(EEGlen.shape[0], 1000) + + def test_get_eeg_split_table(self): + EEGlen = self._partition_table() + Labels = self._labels(EEGlen) # fmt: off val_dict_id = [ @@ -142,91 +131,95 @@ def test_get_eeg_split_table(self): ] # fmt: on - input_grid = { - "partition_table": [EEGlen], - "test_ratio": [0, 0.2], - "val_ratio": [0, 0.2], - "test_split_mode": [0, 1, 2], - "val_split_mode": [1, 2], - "exclude_data_id": [None, {x: [13, 23] for x in range(1, 6)}], - "test_data_id": [None, {x: [14, 22] for x in range(1, 6)}, 4], - "val_data_id": [None, {x: [15, 21] for x in range(1, 6)}, [3]], - "val_ratio_on_all_data": [True, False], - "stratified": [False, True], - "labels": [Labels], - "dataset_id_extractor": [lambda x: int(x.split("_")[0])], - "subject_id_extractor": [None], - "save": [True], - "split_tolerance": [0.001], - "perseverance": [1000], - "save_path": ["tmpsave/results1.csv"], - "seed": [self.seed], - } - input_grid = self.makeGrid(input_grid) - for n, i in enumerate(input_grid): - if n == 500: - print("proceding...", end="", flush=True) - elif n == 1000: - print("a little more...", end="", flush=True) - elif n == 1500: - print("almost done...", end="", flush=True) - - EEGsplit = dl.get_eeg_split_table(**i) - total_list = EEGsplit[EEGsplit["split_set"] != -1].index.tolist() - tot = EEGlen.iloc[total_list]["N_samples"].sum() - self.assertEqual(EEGsplit.shape[0], 1000) - - if isinstance(i["exclude_data_id"], dict): - check = EEGsplit.iloc[excl_dict_id]["split_set"].unique() - self.assertTrue(len(check) == 1) - self.assertTrue(check[0] == -1) - - if i["test_ratio"] == 0: - if isinstance(i["test_data_id"], dict): - check = EEGsplit.iloc[test_dict_id]["split_set"].unique() - self.assertTrue(len(check) == 1) - self.assertTrue(check[0] == 2) - elif isinstance(i["test_data_id"], int): - if not (isinstance(i["exclude_data_id"], dict)): - cond = (EEGsplit["file_name"].str[0] == "4").values - check = EEGsplit[cond]["split_set"].unique() - self.assertTrue(len(check) == 1) - self.assertTrue(check[0] == 2) + grid = make_grid( + { + "partition_table": [EEGlen], + "test_ratio": [0, 0.2], + "val_ratio": [0, 0.2], + "test_split_mode": [0, 1, 2], + "val_split_mode": [1, 2], + "exclude_data_id": [None, {x: [13, 23] for x in range(1, 6)}], + "test_data_id": [None, {x: [14, 22] for x in range(1, 6)}, 4], + "val_data_id": [None, {x: [15, 21] for x in range(1, 6)}, [3]], + "val_ratio_on_all_data": [True, False], + "stratified": [False, True], + "labels": [Labels], + "dataset_id_extractor": [lambda x: int(x.split("_")[0])], + "subject_id_extractor": [None], + "save": [True], + "split_tolerance": [0.001], + "perseverance": [1000], + "save_path": ["tmpsave/results1.csv"], + "seed": [self.seed], + }, + max_comb=48, + ) + for args in grid: + with self.subTest( + test_ratio=args["test_ratio"], + val_ratio=args["val_ratio"], + test_split_mode=args["test_split_mode"], + stratified=args["stratified"], + ): + EEGsplit = dl.get_eeg_split_table(**args) + total_list = EEGsplit[EEGsplit["split_set"] != -1].index.tolist() + tot = EEGlen.iloc[total_list]["N_samples"].sum() + self.assertEqual(EEGsplit.shape[0], 1000) + + if isinstance(args["exclude_data_id"], dict): + check = EEGsplit.iloc[excl_dict_id]["split_set"].unique() + self.assertEqual(len(check), 1) + self.assertEqual(check[0], -1) + + if args["test_ratio"] == 0: + if isinstance(args["test_data_id"], dict): + check = EEGsplit.iloc[test_dict_id]["split_set"].unique() + self.assertEqual(len(check), 1) + self.assertEqual(check[0], 2) + elif isinstance(args["test_data_id"], int): + if not isinstance(args["exclude_data_id"], dict): + cond = (EEGsplit["file_name"].str[0] == "4").values + check = EEGsplit[cond]["split_set"].unique() + self.assertEqual(len(check), 1) + self.assertEqual(check[0], 2) + else: + self.assertEqual(EEGlen["N_samples"][EEGsplit["split_set"] == 2].sum(), 0) else: - self.assertEqual(EEGlen["N_samples"][EEGsplit["split_set"] == 2].sum(), 0) - else: - if i["test_data_id"] is None: - ratio = abs(0.2 - EEGlen["N_samples"][EEGsplit["split_set"] == 2].sum() / tot) - self.assertTrue(ratio < 2.5e-2) - if not (i["stratified"]) and i["test_split_mode"] == 0: - EEGsplit["dataid"] = EEGsplit["file_name"].str[0] - group = EEGsplit.groupby(["dataid", "split_set"]) - lst = list(group.split_set.groups.keys()) - result = [t for t in lst if t[1] == 2] - self.assertEqual(len(result), 1) - check = 200 if i["exclude_data_id"] is None else 190 - self.assertEqual(group.get_group(result[0]).shape[0], check) - - if i["val_ratio"] == 0: - if isinstance(i["val_data_id"], dict): - if isinstance(i["test_data_id"], dict): - check = EEGsplit.iloc[val_dict_id]["split_set"].unique() - self.assertTrue(len(check) == 1) - self.assertTrue(check[0] == 1) - elif i["val_data_id"] is None: - self.assertEqual(EEGlen["N_samples"][EEGsplit["split_set"] == 1].sum(), 0) - else: - if i["val_data_id"] is None: - thresh = 0.2 - if not (i["val_ratio_on_all_data"]): - test_list = EEGsplit[EEGsplit["split_set"] == 2].index.tolist() - test_ratio = EEGlen.iloc[test_list]["N_samples"].sum() / tot - thresh = 0.2 * (1 - test_ratio) - ratio = abs( - thresh - EEGlen["N_samples"][EEGsplit["split_set"] == 1].sum() / tot - ) - self.assertTrue(ratio < 2.5e-2) - + if args["test_data_id"] is None: + ratio = abs( + 0.2 - EEGlen["N_samples"][EEGsplit["split_set"] == 2].sum() / tot + ) + self.assertLess(ratio, 2.5e-2) + if not args["stratified"] and args["test_split_mode"] == 0: + EEGsplit["dataid"] = EEGsplit["file_name"].str[0] + group = EEGsplit.groupby(["dataid", "split_set"]) + lst = list(group.split_set.groups.keys()) + result = [t for t in lst if t[1] == 2] + self.assertEqual(len(result), 1) + expected = 200 if args["exclude_data_id"] is None else 190 + self.assertEqual(group.get_group(result[0]).shape[0], expected) + + if args["val_ratio"] == 0: + if isinstance(args["val_data_id"], dict): + if isinstance(args["test_data_id"], dict): + check = EEGsplit.iloc[val_dict_id]["split_set"].unique() + self.assertEqual(len(check), 1) + self.assertEqual(check[0], 1) + elif args["val_data_id"] is None: + self.assertEqual(EEGlen["N_samples"][EEGsplit["split_set"] == 1].sum(), 0) + else: + if args["val_data_id"] is None: + thresh = 0.2 + if not args["val_ratio_on_all_data"]: + test_list = EEGsplit[EEGsplit["split_set"] == 2].index.tolist() + test_ratio = EEGlen.iloc[test_list]["N_samples"].sum() / tot + thresh = 0.2 * (1 - test_ratio) + ratio = abs( + thresh - EEGlen["N_samples"][EEGsplit["split_set"] == 1].sum() / tot + ) + self.assertLess(ratio, 2.5e-2) + + # stratified split preserves the per-class ratio across sets EEGsplit = dl.get_eeg_split_table( EEGlen, 0.2, @@ -241,84 +234,45 @@ def test_get_eeg_split_table(self): seed=1234, ) ratio = dl.check_split(EEGlen, EEGsplit, Labels, True, False)["class_ratio"] - self.assertTrue(np.abs(ratio - ratio.mean(0)).max() < 1e-3) - print( - " get_eeg_split_table OK: tested", len(input_grid), "combination of input arguments" - ) + self.assertLess(np.abs(ratio - ratio.mean(0)).max(), 1e-3) def test_get_eeg_split_table_kfold(self): - # check: since this function is based on multiple calls of the previous one, - # we have already verified the quality of the single splits, so checks - # will be done on the size of the table and if each file is placed only - # ones in validation set, excluding those placed in test or excluded - print("Testing get_eeg_split_table_kfold...", end="", flush=True) - EEGlen = dl.get_eeg_partition_number( - self.eegpath, - self.freq, - self.window, - self.overlap, - file_format="*.pickle", - load_function=self.loadEEG, - optional_load_fun_args=[False], - transform_function=self.transformEEG, - ) - Labels = np.zeros(EEGlen.shape[0], dtype=int) - for i in range(EEGlen.shape[0]): - _, Labels[i] = self.loadEEG(EEGlen.iloc[i]["full_path"], True) - input_grid = { - "partition_table": [EEGlen], - "test_ratio": [0, 0.2], - "kfold": [5, 10], - "test_split_mode": [1, 2], - "val_split_mode": [1, 2], - "exclude_data_id": [None, {x: [13, 23] for x in range(1, 6)}], - "test_data_id": [None, {x: [14, 22] for x in range(1, 6)}, 4], - "stratified": [False, True], - "labels": [Labels], - "save": [True], - "split_tolerance": [0.01], - "perseverance": [1000], - "save_path": ["tmpsave/results1.csv"], - } - input_grid = self.makeGrid(input_grid) - for i in input_grid: - EEGsplit = dl.get_eeg_split_table_kfold(**i) - self.assertEqual(EEGsplit.shape[0], 1000) - self.assertEqual(EEGsplit.shape[1], i["kfold"] + 1) - sums = set(EEGsplit.sum(axis=1, numeric_only=True).unique().tolist()) - self.assertTrue(sums.issubset(set([-1 * i["kfold"], 1, 2 * i["kfold"]]))) - - print( - " get_eeg_split_table_kfold OK: tested", - len(input_grid), - "combination of input arguments", + EEGlen = self._partition_table() + Labels = self._labels(EEGlen) + grid = make_grid( + { + "partition_table": [EEGlen], + "test_ratio": [0, 0.2], + "kfold": [5, 10], + "test_split_mode": [1, 2], + "val_split_mode": [1, 2], + "exclude_data_id": [None, {x: [13, 23] for x in range(1, 6)}], + "test_data_id": [None, {x: [14, 22] for x in range(1, 6)}, 4], + "stratified": [False, True], + "labels": [Labels], + "save": [True], + "split_tolerance": [0.01], + "perseverance": [1000], + "save_path": ["tmpsave/results1.csv"], + }, + max_comb=24, ) + for args in grid: + with self.subTest(test_ratio=args["test_ratio"], kfold=args["kfold"], stratified=args["stratified"]): + EEGsplit = dl.get_eeg_split_table_kfold(**args) + self.assertEqual(EEGsplit.shape[0], 1000) + self.assertEqual(EEGsplit.shape[1], args["kfold"] + 1) + sums = set(EEGsplit.sum(axis=1, numeric_only=True).unique().tolist()) + self.assertTrue(sums.issubset({-1 * args["kfold"], 1, 2 * args["kfold"]})) def test_EEGDataset(self): - # checks: extraction is performed correctly - print("Testing EEGDataset...", end="", flush=True) - - EEGlen = dl.get_eeg_partition_number( - self.eegpath, - self.freq, - self.window, - self.overlap, - file_format="*.pickle", - load_function=self.loadEEG, - optional_load_fun_args=[False], - transform_function=self.transformEEG, - ) - Labels = np.zeros(EEGlen.shape[0], dtype=int) - for i in range(EEGlen.shape[0]): - _, Labels[i] = self.loadEEG(EEGlen.iloc[i]["full_path"], True) - + EEGlen = self._partition_table() EEGsplit = dl.get_eeg_split_table( EEGlen, test_ratio=0.1, val_ratio=0.1, test_split_mode="file", val_split_mode="file", - # stratified=False, labels=Labels, perseverance=5000, split_tolerance=0.005, ) @@ -331,8 +285,8 @@ def test_EEGDataset(self): load_function=self.loadEEG, transform_function=self.transformEEG, ) - sample_1 = dataset_pretrain.__getitem__(0) - self.assertTrue(isinstance(sample_1, torch.Tensor)) + sample_1 = dataset_pretrain[0] + self.assertIsInstance(sample_1, torch.Tensor) self.assertEqual(sample_1.shape[-1], 256) dataset_finetune = dl.EEGDataset( @@ -346,10 +300,10 @@ def test_EEGDataset(self): transform_function=self.transformEEG, label_on_load=True, ) - sample_2, label_2 = dataset_finetune.__getitem__(0) - self.assertTrue(isinstance(sample_2, torch.Tensor)) + sample_2, label_2 = dataset_finetune[0] + self.assertIsInstance(sample_2, torch.Tensor) self.assertEqual(sample_2.shape[-1], 256) - self.assertTrue(isinstance(label_2, int)) + self.assertIsInstance(label_2, int) # try again with an array of labels dataset_finetune = dl.EEGDataset( @@ -363,19 +317,16 @@ def test_EEGDataset(self): transform_function=self.transformEEG, multilabel_on_load=True, ) - sample_2, label_2 = dataset_finetune.__getitem__(0) - self.assertTrue(isinstance(sample_2, torch.Tensor)) + sample_2, label_2 = dataset_finetune[0] + self.assertIsInstance(sample_2, torch.Tensor) self.assertEqual(sample_2.shape[-1], 256) - # self.assertTrue(len(label_2.shape)==0) dataset_finetune.preload_dataset() - sample_3, label_3 = dataset_finetune.__getitem__(0) - self.assertTrue(isinstance(sample_3, torch.Tensor)) + sample_3, label_3 = dataset_finetune[0] + self.assertIsInstance(sample_3, torch.Tensor) self.assertEqual(sample_3.shape[-1], 256) - print(" EEGDataset OK") def test_EEGSampler(self): - print("Testing Sampler on both modalities...", end="", flush=True) EEGlen = dl.get_eeg_partition_number( self.eegpath, self.freq, @@ -387,9 +338,9 @@ def test_EEGSampler(self): save=True, save_path="tmpsave/results1.csv", ) - Labels = np.zeros(EEGlen.shape[0]) # , dtype=in) + Labels = np.zeros(EEGlen.shape[0]) for i in range(EEGlen.shape[0]): - EEG, Labels[i] = self.loadEEG(EEGlen["full_path"][i], True) + _, Labels[i] = self.loadEEG(EEGlen["full_path"][i], True) Labels = Labels.astype(int) EEGsplit = dl.get_eeg_split_table( EEGlen, @@ -410,25 +361,14 @@ def test_EEGSampler(self): load_function=self.loadEEG, transform_function=self.transformEEG, ) - sampler_linear = dl.EEGSampler(dataset_pretrain, Mode=0) - sampler_custom = dl.EEGSampler(dataset_pretrain, 16, 4) - print(" EEGDataset OK") + # both sampling modalities must build without error + dl.EEGSampler(dataset_pretrain, Mode=0) + dl.EEGSampler(dataset_pretrain, 16, 4) @classmethod def tearDownClass(cls): - print("removing generated residual directories (Simulated_EEG, tmpsave)") - try: - if platform.system() == "Windows": - os.system("rmdir /Q /S Simulated_EEG") # nosec - os.system("rmdir /Q /S tmpsave") # nosec - else: - os.system("rm -r Simulated_EEG") # nosec - os.system("rm -r tmpsave") # nosec - except: - print( - 'Failed to delete "Simulated_EEG" and "tmpsave" folders.' - " Please don't hate me and do it manually" - ) + for folder in ("Simulated_EEG", "tmpsave"): + shutil.rmtree(folder, ignore_errors=True) if __name__ == "__main__": diff --git a/test/EEGself/losses/losses_test.py b/test/EEGself/losses/losses_test.py index dcfd19d..272497f 100644 --- a/test/EEGself/losses/losses_test.py +++ b/test/EEGself/losses/losses_test.py @@ -1,284 +1,127 @@ -import itertools import os import sys +import types import unittest -import numpy as np import torch -from selfeeg import losses +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from EEGself._testtools import get_device, make_grid -class TestLoss(unittest.TestCase): +from selfeeg import losses - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds +N, FEAT = 64, 128 - @classmethod - def setUpClass(cls): - if torch.backends.mps.is_available(): - cls.device = torch.device("mps") - elif torch.cuda.is_available(): - cls.device = torch.device("cuda") - else: - cls.device = torch.device("cpu") - if cls.device.type != "cpu": - try: - xx = torch.randn(64, 128).to(device=cls.device) - yy = torch.randn(64, 128).to(device=cls.device) - xx = losses.barlow_loss(xx, yy) - except Exception: - cls.device = torch.device("cpu") +def _phase_signal(phase=0.0, length=FEAT, rows=N): + """Deterministic (seed-free) signal used for the numeric regression checks.""" + base = torch.sin(torch.linspace(0, 8 * torch.pi, length) + phase) + return base + (torch.arange(rows) / rows).unsqueeze(1) - device = cls.device - print("\n---------------------") - print("TESTING LOSSES MODULE") - if cls.device.type != "cpu": - print("Found other device: testing module with both cpu and gpu") - else: - print("Didn't found cuda device: testing module with only cpu") - print("---------------------") - N, Feat = 64, 128 - cls.N = 64 - cls.Feat = 128 - cls.x = torch.randn(N, Feat) - cls.y = torch.randn(N, Feat) - cls.p = torch.randn(N, Feat) - cls.z = torch.randn(N, Feat) - cls.u = torch.randn(Feat, 1024) - if device.type != "cpu": - cls.x2 = torch.randn(N, Feat).to(device=device) - cls.y2 = torch.randn(N, Feat).to(device=device) - cls.p2 = torch.randn(N, Feat).to(device=device) - cls.z2 = torch.randn(N, Feat).to(device=device) - cls.u2 = torch.randn(Feat, 1024).to(device=device) +class TestLoss(unittest.TestCase): - def setUp(self): - self.seed = 1234 - np.random.seed(self.seed) - torch.manual_seed(self.seed) + @classmethod + def setUpClass(cls): + cls.device = get_device( + probe=lambda dev: losses.barlow_loss( + torch.randn(N, FEAT, device=dev), torch.randn(N, FEAT, device=dev) + ) + ) + + def _tensor_sets(self): + """Yield a namespace of random tensors per device to be tested.""" + torch.manual_seed(1234) + devices = ["cpu"] + ([self.device.type] if self.device.type != "cpu" else []) + for dev in devices: + yield types.SimpleNamespace( + x=torch.randn(N, FEAT, device=dev), + y=torch.randn(N, FEAT, device=dev), + p=torch.randn(N, FEAT, device=dev), + z=torch.randn(N, FEAT, device=dev), + u=torch.randn(FEAT, 1024, device=dev), + ) + + def _assert_no_nan(self, grid, loss_fn): + for args in grid: + with self.subTest(**{k: v for k, v in args.items() if not torch.is_tensor(v)}): + self.assertEqual(torch.isnan(loss_fn(**args)).sum().item(), 0) def test_barlow_loss(self): - print("Testing Barlow Loss...", end="", flush=True) - Barlow_args = {"z1": [self.x], "z2": [self.y, None], "lambda_coeff": [0.005, 0.05, 0.5, 1]} - Barlow_args = self.makeGrid(Barlow_args) - for i in Barlow_args: - loss = losses.barlow_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - if self.device.type != "cpu": - Barlow_args = { - "z1": [self.x2], - "z2": [self.y2, None], - "lambda_coeff": [0.005, 0.05, 0.5, 1], - } - Barlow_args = self.makeGrid(Barlow_args) - for i in Barlow_args: - loss = losses.barlow_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - x = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat)) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - loss = losses.barlow_loss(x) - self.assertTrue((loss - 76.4043) < 1e-4) - print(" Barlow Loss OK: tested", len(Barlow_args), "combinations of input arguments") + for t in self._tensor_sets(): + grid = make_grid( + {"z1": [t.x], "z2": [t.y, None], "lambda_coeff": [0.005, 0.05, 0.5, 1]} + ) + self._assert_no_nan(grid, losses.barlow_loss) + loss = losses.barlow_loss(_phase_signal()) + self.assertAlmostEqual(loss.item(), 76.4043, delta=1e-3) def test_byol_loss(self): - print("Testing BYOL Loss...", end="", flush=True) - BYOL_args = { - "z1": [self.x], - "z2": [self.y], - "p1": [self.p], - "p2": [self.z], - "projections_norm": [True, False], - } - BYOL_args = self.makeGrid(BYOL_args) - for i in BYOL_args: - loss = losses.byol_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - if self.device.type != "cpu": - BYOL_args = { - "z1": [self.x2], - "z2": [self.y2], - "p1": [self.p2], - "p2": [self.z2], - "projections_norm": [True, False], - } - BYOL_args = self.makeGrid(BYOL_args) - for i in BYOL_args: - loss = losses.byol_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - x = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 6) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - y = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 4) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - p = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 3) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - z = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 5) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - loss = losses.byol_loss(x, y, p, z) - self.assertTrue((loss - 0.0534) < 1e-4) - print(" BYOL Loss OK: tested", len(BYOL_args), "combinations of input arguments") + for t in self._tensor_sets(): + grid = make_grid( + {"z1": [t.x], "z2": [t.y], "p1": [t.p], "p2": [t.z], "projections_norm": [True, False]} + ) + self._assert_no_nan(grid, losses.byol_loss) + loss = losses.byol_loss( + _phase_signal(torch.pi / 6), + _phase_signal(torch.pi / 4), + _phase_signal(torch.pi / 3), + _phase_signal(torch.pi / 5), + ) + self.assertAlmostEqual(loss.item(), 0.0534, delta=1e-3) def test_simclr_loss(self): - print("Testing SimCLR Loss...", end="", flush=True) - SimCLR_args = { - "projections": [self.x], - "temperature": [0.15, 0.5, 0.7], - "projections_norm": [True, False], - } - SimCLR_args = self.makeGrid(SimCLR_args) - for i in SimCLR_args: - loss = losses.simclr_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - if self.device.type != "cpu": - SimCLR_args = { - "projections": [self.x], - "temperature": [0.15, 0.5, 0.7], - "projections_norm": [True, False], - } - SimCLR_args = self.makeGrid(SimCLR_args) - for i in SimCLR_args: - loss = losses.simclr_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - x = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 6) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - loss = losses.simclr_loss(x) - self.assertTrue((loss - 9.04887) < 1e-4) - print(" SimCLR Loss OK: tested", len(SimCLR_args), "combinations of input arguments") + for t in self._tensor_sets(): + grid = make_grid( + {"projections": [t.x], "temperature": [0.15, 0.5, 0.7], "projections_norm": [True, False]} + ) + self._assert_no_nan(grid, losses.simclr_loss) + loss = losses.simclr_loss(_phase_signal(torch.pi / 6)) + self.assertAlmostEqual(loss.item(), 9.04887, delta=1e-3) def test_simsiam_loss(self): - print("Testing SimSiam Loss...", end="", flush=True) - Siam_args = { - "z1": [self.x], - "z2": [self.y], - "p1": [self.p], - "p2": [self.z], - "projections_norm": [True, False], - } - Siam_args = self.makeGrid(Siam_args) - for i in Siam_args: - loss = losses.simsiam_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - if self.device.type != "cpu": - Siam_args = { - "z1": [self.x2], - "z2": [self.y2], - "p1": [self.p2], - "p2": [self.z2], - "projections_norm": [True, False], - } - Siam_args = self.makeGrid(Siam_args) - for i in Siam_args: - loss = losses.simsiam_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - x = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 6) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - y = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 4) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - p = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 3) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - z = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 5) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - loss = losses.simsiam_loss(x, y, p, z) - self.assertTrue((loss - (-0.9867)) < 1e-4) - print(" SimSiam Loss OK: tested", len(Siam_args), "combinations of input arguments") + for t in self._tensor_sets(): + grid = make_grid( + {"z1": [t.x], "z2": [t.y], "p1": [t.p], "p2": [t.z], "projections_norm": [True, False]} + ) + self._assert_no_nan(grid, losses.simsiam_loss) + loss = losses.simsiam_loss( + _phase_signal(torch.pi / 6), + _phase_signal(torch.pi / 4), + _phase_signal(torch.pi / 3), + _phase_signal(torch.pi / 5), + ) + self.assertAlmostEqual(loss.item(), -0.9867, delta=1e-3) def test_vicreg_loss(self): - print("Testing VICReg Loss...", end="", flush=True) - Vicreg_args = { - "z1": [self.x], - "z2": [self.y, None], - "Lambda": [25, 10, 50], - "Mu": [25, 5, 50], - "Nu": [2, 1, 0.5], - } - Vicreg_args = self.makeGrid(Vicreg_args) - for i in Vicreg_args: - loss = losses.vicreg_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - if self.device.type != "cpu": - Vicreg_args = { - "z1": [self.x2], - "z2": [self.y2, None], - "Lambda": [25, 10, 50], - "Mu": [25, 5, 50], - "Nu": [2, 1, 0.5], - } - Vicreg_args = self.makeGrid(Vicreg_args) - for i in Vicreg_args: - loss = losses.vicreg_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - x = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat)) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - loss = losses.vicreg_loss(x) - self.assertTrue((loss - 21.4443) < 1e-4) - print(" VICReg Loss OK: tested", len(Vicreg_args), "combinations of input arguments") + for t in self._tensor_sets(): + grid = make_grid( + {"z1": [t.x], "z2": [t.y, None], "Lambda": [25, 10, 50], "Mu": [25, 5, 50], "Nu": [2, 1, 0.5]}, + max_comb=24, + ) + self._assert_no_nan(grid, losses.vicreg_loss) + loss = losses.vicreg_loss(_phase_signal()) + self.assertAlmostEqual(loss.item(), 21.4443, delta=1e-3) def test_moco_loss(self): - print("Testing MoCo Loss...", end="", flush=True) - Moco_args = { - "q": [self.x], - "k": [self.y], - "queue": [None, self.u], - "projections_norm": [True, False], - "temperature": [0.15, 0.5, 0.9], - } - Moco_args = self.makeGrid(Moco_args) - for i in Moco_args: - loss = losses.moco_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - if self.device.type != "cpu": - Moco_args = { - "q": [self.x], - "k": [self.y], - "queue": [None, self.u], - "projections_norm": [True, False], - "temperature": [0.15, 0.5, 0.9], - } - Moco_args = self.makeGrid(Moco_args) - for i in Moco_args: - loss = losses.moco_loss(**i) - self.assertTrue(torch.isnan(loss).sum() == 0) - - x = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 6) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - y = torch.sin(torch.linspace(0, 8 * torch.pi, self.Feat) + torch.pi / 4) + ( - torch.arange(self.N) / self.N - ).unsqueeze(1) - u = torch.sin(torch.linspace(0, 8 * torch.pi, 1024) + torch.pi / 3) + ( - torch.arange(self.Feat) / self.Feat - ).unsqueeze(1) - loss = losses.moco_loss(x, y, u) - self.assertTrue((loss - 108.6667) < 1e-4) - loss = losses.moco_loss(x, y) - self.assertTrue((loss - 0.4952) < 1e-4) - print(" MoCo Loss OK: tested", len(Moco_args), "combinations of input arguments") + for t in self._tensor_sets(): + grid = make_grid( + { + "q": [t.x], + "k": [t.y], + "queue": [None, t.u], + "projections_norm": [True, False], + "temperature": [0.15, 0.5, 0.9], + } + ) + self._assert_no_nan(grid, losses.moco_loss) + x = _phase_signal(torch.pi / 6) + y = _phase_signal(torch.pi / 4) + u = _phase_signal(torch.pi / 3, length=1024, rows=FEAT) + self.assertAlmostEqual(losses.moco_loss(x, y, u).item(), 108.6667, delta=1e-3) + self.assertAlmostEqual(losses.moco_loss(x, y).item(), 0.4952, delta=1e-3) if __name__ == "__main__": diff --git a/test/EEGself/models/layers_test.py b/test/EEGself/models/layers_test.py index ebbd3b7..1503e12 100644 --- a/test/EEGself/models/layers_test.py +++ b/test/EEGself/models/layers_test.py @@ -1,250 +1,167 @@ -import itertools import os import sys import unittest -import numpy as np import torch +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from EEGself._testtools import get_device, make_grid + from selfeeg import models +N, CHAN, SAMPLES = 2, 8, 2048 -class TestModels(unittest.TestCase): - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds +class TestModels(unittest.TestCase): @classmethod def setUpClass(cls): - if torch.backends.mps.is_available(): - cls.device = torch.device("mps") - elif torch.cuda.is_available(): - cls.device = torch.device("cuda") - else: - cls.device = torch.device("cpu") - - if cls.device.type != "cpu": - try: - xx = torch.randn(2, 8, 2048).to(device=cls.device) - lay = models.ConstrainedConv1d(8, 4, 16).to(device=cls.device) - xx = lay(xx) - except Exception: - cls.device = torch.device("cpu") - - print("\n----------------------------") - print("TESTING MODELS.LAYERS MODULE") - if cls.device.type != "cpu": - print("Found gpu device: testing module with both cpu and gpu") - else: - print("Didn't found cuda device: testing module with only cpu") - print("----------------------------") - cls.N = 2 - cls.Chan = 8 - cls.Samples = 2048 - cls.x = torch.randn(cls.N, cls.Chan, cls.Samples) - cls.xl = torch.randn(cls.N, 1, 16, cls.Samples) - cls.xd = torch.randn(cls.N, 128) - if cls.device.type != "cpu": - cls.x2 = torch.randn(cls.N, cls.Chan, cls.Samples).to(device=cls.device) - cls.xl2 = torch.randn(cls.N, 1, 16, cls.Samples).to(device=cls.device) - cls.xd2 = torch.randn(cls.N, 128).to(device=cls.device) + cls.device = get_device( + probe=lambda dev: models.ConstrainedConv1d(8, 4, 16).to(dev)( + torch.randn(2, 8, 2048, device=dev) + ) + ) + cls.x = torch.randn(N, CHAN, SAMPLES) + cls.xl = torch.randn(N, 1, 16, SAMPLES) + cls.xd = torch.randn(N, 128) def setUp(self): torch.manual_seed(1234) + def _devices(self): + return ["cpu"] + ([self.device.type] if self.device.type != "cpu" else []) + + def _assert_constraints(self, weight, out, max_norm, min_norm, norm_axes): + self.assertEqual(torch.isnan(out).sum().item(), 0) + if max_norm is not None: + norms = torch.sqrt(torch.sum(torch.square(weight), axis=norm_axes)) + self.assertEqual(torch.sum(norms > (max_norm + 1e-3)).item(), 0) + if min_norm is not None: + self.assertEqual(torch.sum(norms < (min_norm - 1e-3)).item(), 0) + def test_ConstrainedConv1d(self): - print("Testing conv1d with norm constraint...", end="", flush=True) - Conv_args = { - "in_channels": [8], - "out_channels": [4, 16], - "kernel_size": [16], - "stride": [1, 2, 3], - "dilation": [1, 2], - "bias": [True, False], - "max_norm": [None, 1, 2], - "min_norm": [None, 1], - "padding": ["valid", "causal"], - } - Conv_args = self.makeGrid(Conv_args) - for i in Conv_args: - model = models.ConstrainedConv1d(**i) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.x) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=[1, 2])) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - - if self.device.type != "cpu": - for i in Conv_args: - model = models.ConstrainedConv1d(**i).to(device=self.device) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.x2) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=[1, 2])) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - print( - " Constrained conv1d OK: tested", len(Conv_args), " combinations of input arguments" + grid = make_grid( + { + "in_channels": [8], + "out_channels": [4, 16], + "kernel_size": [16], + "stride": [1, 2, 3], + "dilation": [1, 2], + "bias": [True, False], + "max_norm": [None, 1, 2], + "min_norm": [None, 1], + "padding": ["valid", "causal"], + }, + max_comb=40, ) + for dev in self._devices(): + x = self.x.to(dev) + for args in grid: + with self.subTest(device=dev, **{k: args[k] for k in ("stride", "max_norm", "padding")}): + model = models.ConstrainedConv1d(**args).to(dev) + model.weight = torch.nn.Parameter(model.weight * 10) + out = model(x) + self._assert_constraints(model.weight, out, args["max_norm"], args["min_norm"], [1, 2]) def test_ConstrainedConv2d(self): - print("Testing conv2d with norm constraint...", end="", flush=True) - Conv_args = { - "in_channels": [1], - "out_channels": [5, 16], - "kernel_size": [(1, 64), (5, 1), (5, 64)], - "stride": [1, 2, 3], - "dilation": [1, 2], - "bias": [True, False], - "max_norm": [None, 1, 2, 3], - "min_norm": [None, 1], - "padding": ["valid"], - } - Conv_args = self.makeGrid(Conv_args) - for i in Conv_args: - model = models.ConstrainedConv2d(**i) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.xl) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=[1, 2, 3])) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - - if self.device.type != "cpu": - for i in Conv_args: - model = models.ConstrainedConv2d(**i).to(device=self.device) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.xl2) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=[1, 2, 3])) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - print( - " Constrained conv2d OK: tested", len(Conv_args), " combinations of input arguments" + grid = make_grid( + { + "in_channels": [1], + "out_channels": [5, 16], + "kernel_size": [(1, 64), (5, 1), (5, 64)], + "stride": [1, 2, 3], + "dilation": [1, 2], + "bias": [True, False], + "max_norm": [None, 1, 2, 3], + "min_norm": [None, 1], + "padding": ["valid"], + }, + max_comb=40, ) + for dev in self._devices(): + xl = self.xl.to(dev) + for args in grid: + with self.subTest(device=dev, kernel=args["kernel_size"], max_norm=args["max_norm"]): + model = models.ConstrainedConv2d(**args).to(dev) + model.weight = torch.nn.Parameter(model.weight * 10) + out = model(xl) + self._assert_constraints( + model.weight, out, args["max_norm"], args["min_norm"], [1, 2, 3] + ) def test_ConstrainedDense(self): - print("Testing Dense layer with max norm constraint...", end="", flush=True) - Dense_args = { - "in_features": [128], - "out_features": [32], - "bias": [True, False], - "max_norm": [None, 1, 3], - "min_norm": [None, 1], - } - Dense_args = self.makeGrid(Dense_args) - for i in Dense_args: - model = models.ConstrainedDense(**i) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.xd) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=1)) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], 32) - - if self.device.type != "cpu": - for i in Dense_args: - model = models.ConstrainedDense(**i).to(device=self.device) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.xd2) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=1)) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], 32) - print(" Dense layer OK: tested", len(Dense_args), " combinations of input arguments") + grid = make_grid( + { + "in_features": [128], + "out_features": [32], + "bias": [True, False], + "max_norm": [None, 1, 3], + "min_norm": [None, 1], + } + ) + for dev in self._devices(): + xd = self.xd.to(dev) + for args in grid: + with self.subTest(device=dev, max_norm=args["max_norm"], min_norm=args["min_norm"]): + model = models.ConstrainedDense(**args).to(dev) + model.weight = torch.nn.Parameter(model.weight * 10) + out = model(xd) + self._assert_constraints(model.weight, out, args["max_norm"], args["min_norm"], 1) + self.assertEqual(out.shape[1], 32) def test_DepthwiseConv2d(self): - print("Testing Depthwise conv2d with norm constraint...", end="", flush=True) - Depthwise_args = { - "in_channels": [1], - "depth_multiplier": [2, 3, 4], - "kernel_size": [(1, 64), (5, 1), (5, 64)], - "stride": [1, 2, 3], - "dilation": [1, 2], - "bias": [True, False], - "max_norm": [None, 1, 3], - "min_norm": [None, 1], - "padding": ["valid"], - } - Depthwise_args = self.makeGrid(Depthwise_args) - for i in Depthwise_args: - model = models.DepthwiseConv2d(**i) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.xl) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=[1, 2, 3])) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["depth_multiplier"]) - - if self.device.type != "cpu": - for i in Depthwise_args: - model = models.DepthwiseConv2d(**i).to(device=self.device) - model.weight = torch.nn.Parameter(model.weight * 10) - out = model(self.xl2) - if i["max_norm"] is not None: - norms = torch.sqrt(torch.sum(torch.square(model.weight), axis=[1, 2, 3])) - self.assertTrue(torch.sum(norms > (i["max_norm"] + 1e-3)).item() == 0) - if i["min_norm"] is not None: - self.assertTrue(torch.sum(norms < (i["min_norm"] - 1e-3)).item() == 0) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["depth_multiplier"]) - print( - " Depthwise conv2d OK: tested", - len(Depthwise_args), - " combinations of input arguments", + grid = make_grid( + { + "in_channels": [1], + "depth_multiplier": [2, 3, 4], + "kernel_size": [(1, 64), (5, 1), (5, 64)], + "stride": [1, 2, 3], + "dilation": [1, 2], + "bias": [True, False], + "max_norm": [None, 1, 3], + "min_norm": [None, 1], + "padding": ["valid"], + }, + max_comb=40, ) + for dev in self._devices(): + xl = self.xl.to(dev) + for args in grid: + with self.subTest(device=dev, depth=args["depth_multiplier"], max_norm=args["max_norm"]): + model = models.DepthwiseConv2d(**args).to(dev) + model.weight = torch.nn.Parameter(model.weight * 10) + out = model(xl) + self._assert_constraints(model.weight, out, args["max_norm"], args["min_norm"], [1, 2, 3]) + self.assertEqual(out.shape[1], args["depth_multiplier"]) def test_SeparableConv2d(self): - print("Testing Separable conv2d with norm constraint...", end="", flush=True) - Separable_args = { - "in_channels": [1], - "out_channels": [5, 16], - "depth_multiplier": [1, 3], - "kernel_size": [(1, 64), (5, 1), (5, 64)], - "stride": [1, 2, 3], - "dilation": [1, 2], - "bias": [True, False], - "depth_max_norm": [None, 1, 2], - "depth_min_norm": [None, 1], - "point_max_norm": [None, 1, 2], - "point_min_norm": [None, 1], - "padding": ["valid"], - } - Separable_args = self.makeGrid(Separable_args) - for i in Separable_args: - model = models.SeparableConv2d(**i) - out = model(self.xl) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["out_channels"]) - - if self.device.type != "cpu": - for i in Separable_args: - model = models.SeparableConv2d(**i).to(device=self.device) - out = model(self.xl2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["out_channels"]) - print( - " Separable conv2d OK: tested", len(Separable_args), "combinations of input arguments" + grid = make_grid( + { + "in_channels": [1], + "out_channels": [5, 16], + "depth_multiplier": [1, 3], + "kernel_size": [(1, 64), (5, 1), (5, 64)], + "stride": [1, 2, 3], + "dilation": [1, 2], + "bias": [True, False], + "depth_max_norm": [None, 1, 2], + "depth_min_norm": [None, 1], + "point_max_norm": [None, 1, 2], + "point_min_norm": [None, 1], + "padding": ["valid"], + }, + max_comb=40, ) + for dev in self._devices(): + xl = self.xl.to(dev) + for args in grid: + with self.subTest(device=dev, out_channels=args["out_channels"], kernel=args["kernel_size"]): + model = models.SeparableConv2d(**args).to(dev) + out = model(xl) + self.assertEqual(torch.isnan(out).sum().item(), 0) + self.assertEqual(out.shape[1], args["out_channels"]) if __name__ == "__main__": diff --git a/test/EEGself/models/zoo_test.py b/test/EEGself/models/zoo_test.py index 2775c98..83308ec 100644 --- a/test/EEGself/models/zoo_test.py +++ b/test/EEGself/models/zoo_test.py @@ -1,540 +1,325 @@ -import itertools import os import sys import unittest import warnings -import numpy as np + import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from EEGself._testtools import get_device, make_grid + from selfeeg import models +N, CHAN, SAMPLES = 2, 8, 2048 -class TestModels(unittest.TestCase): - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds +class TestModels(unittest.TestCase): @classmethod def setUpClass(cls): warnings.filterwarnings("ignore", message="Using padding='same'", category=UserWarning) - if torch.backends.mps.is_available(): - cls.device = torch.device("mps") - elif torch.cuda.is_available(): - cls.device = torch.device("cuda") - else: - cls.device = torch.device("cpu") - - if cls.device.type != "cpu": - try: - xx = torch.randn(2, 8, 2048).to(device=cls.device) - model = models.EEGNet(2, 8, 2048).to(device=cls.device) - xx = model(xx) - except Exception: - cls.device = torch.device("cpu") - - print("\n-------------------------") - print("TESTING MODELS.ZOO MODULE") - if cls.device.type != "cpu": - print("Found gpu device: testing module with both cpu and gpu") - else: - print("Didn't found cuda device: testing module with only cpu") - print("-------------------------") - cls.N = 2 - cls.Chan = 8 - cls.Samples = 2048 - cls.x = torch.randn(cls.N, cls.Chan, cls.Samples) - cls.xl = torch.randn(cls.N, 1, 16, cls.Samples) - cls.xd = torch.randn(cls.N, 128) - if cls.device.type != "cpu": - cls.x2 = torch.randn(cls.N, cls.Chan, cls.Samples).to(device=cls.device) - cls.xl2 = torch.randn(cls.N, 1, 16, cls.Samples).to(device=cls.device) - cls.xd2 = torch.randn(cls.N, 128).to(device=cls.device) + cls.device = get_device( + probe=lambda dev: models.EEGNet(2, 8, 2048).to(dev)(torch.randn(2, 8, 2048, device=dev)) + ) + cls.x = torch.randn(N, CHAN, SAMPLES) def setUp(self): torch.manual_seed(1234) - def test_ATCNet(self): - print("Testing ATCNet...", end="", flush=True) + def _devices(self, allow_mps=True): + if self.device.type == "cpu": + return ["cpu"] + if self.device.type == "mps" and not allow_mps: + return ["cpu"] + return ["cpu", self.device.type] + + def _check_classifier(self, model_cls, grid, input_fn=None, skip=None, allow_mps=True, label_keys=()): + """Build ``model_cls`` for every combination and check its output. + + The output must be NaN-free, have ``nb_classes`` columns (or 1 when + binary), and, unless ``return_logits`` is set, lie in the [0, 1] range. + """ + for dev in self._devices(allow_mps): + for args in grid: + if skip is not None and skip(args): + continue + with self.subTest(device=dev, **{k: args[k] for k in label_keys}): + model = model_cls(**args).to(dev) + x = input_fn(args, dev) if input_fn else self.x.to(dev) + out = model(x) + self.assertEqual(torch.isnan(out).sum().item(), 0) + expected = args["nb_classes"] if args["nb_classes"] > 2 else 1 + self.assertEqual(out.shape[1], expected) + if not args["return_logits"]: + self.assertLessEqual(out.max().item(), 1) + self.assertGreaterEqual(out.min().item(), 0) - DCN_args = { - "nb_classes": [2, 4], - "Chans": [self.Chan], - "Samples": [self.Samples], - "Fs": [128], - "num_windows": [4], - "mha_heads": [2, 4], - "tcn_depth": [2, 3], - "F1": [12, 8], - "D": [2, 3], - "return_logits": [False], - "seed": [42], - } - DCN_grid = self.makeGrid(DCN_args) - for i in DCN_grid: - model = models.ATCNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - if self.device.type != "cpu": - for i in DCN_grid: - model = models.ATCNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" ATCNet OK: tested ", len(DCN_grid), " combinations of input arguments") + def test_ATCNet(self): + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "Samples": [SAMPLES], + "Fs": [128], + "num_windows": [4], + "mha_heads": [2, 4], + "tcn_depth": [2, 3], + "F1": [12, 8], + "D": [2, 3], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.ATCNet, grid, label_keys=("nb_classes", "mha_heads", "tcn_depth")) def test_DeepConvNet(self): - print("Testing DeepConvNet...", end="", flush=True) - DCN_args = { - "nb_classes": [2, 4], - "Chans": [self.Chan], - "Samples": [self.Samples], - "kernLength": [10, 20], - "F": [12, 25], - "Pool": [3, 4], - "stride": [3, 4], - "max_norm": [2.0], - "batch_momentum": [0.9], - "ELUalpha": [1], - "dropRate": [0.5], - "max_dense_norm": [1.0], - "return_logits": [False], - "seed": [42], - } - DCN_grid = self.makeGrid(DCN_args) - for i in DCN_grid: - model = models.DeepConvNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - if self.device.type != "cpu": - for i in DCN_grid: - model = models.DeepConvNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" DeepConvNet OK: tested ", len(DCN_grid), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "Samples": [SAMPLES], + "kernLength": [10, 20], + "F": [12, 25], + "Pool": [3, 4], + "stride": [3, 4], + "max_norm": [2.0], + "batch_momentum": [0.9], + "ELUalpha": [1], + "dropRate": [0.5], + "max_dense_norm": [1.0], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.DeepConvNet, grid, label_keys=("nb_classes", "kernLength", "F")) def test_EEGConformer(self): - print("Testing EEGConformer...", end="", flush=True) - EEGcon_args = { - "nb_classes": [2, 4], - "Samples": [2048], - "Chans": [self.Chan], - "F": [40], - "K1": [25, 12], - "Pool": [75, 50], - "stride_pool": [20], - "nlayers": [4], - "d_model": [40, 80], - "nheads": [8, 10], - "dim_feedforward": [80], - "activation_transformer": ["gelu"], - "mlp_dim": [[128, 32], [64, 32]], - "return_logits": [False], - "seed": [42], - } - EEGcon_args = self.makeGrid(EEGcon_args) - for i in EEGcon_args: - model = models.EEGConformer(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGcon_args: - model = models.EEGConformer(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" EEGConformer OK: tested", len(EEGcon_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Samples": [2048], + "Chans": [CHAN], + "F": [40], + "K1": [25, 12], + "Pool": [75, 50], + "stride_pool": [20], + "nlayers": [4], + "d_model": [40, 80], + "nheads": [8, 10], + "dim_feedforward": [80], + "activation_transformer": ["gelu"], + "mlp_dim": [[128, 32], [64, 32]], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.EEGConformer, grid, label_keys=("nb_classes", "d_model", "nheads")) def test_EEGInception(self): - print("Testing EEGInception...", end="", flush=True) - EEGin_args = { - "nb_classes": [2, 4], - "Chans": [self.Chan], - "Samples": [self.Samples], - "kernel_size": [32, 128], - "F1": [4, 16], - "D": [2, 4], - "pool": [4, 8], - "batch_momentum": [0.9], - "dropRate": [0.5], - "max_depth_norm": [1.0], - "return_logits": [False], - "bias": [True, False], - "seed": [42], - } - EEGin_args = self.makeGrid(EEGin_args) - for i in EEGin_args: - model = models.EEGInception(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGin_args: - model = models.EEGInception(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" EEGInception OK: tested", len(EEGin_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "Samples": [SAMPLES], + "kernel_size": [32, 128], + "F1": [4, 16], + "D": [2, 4], + "pool": [4, 8], + "batch_momentum": [0.9], + "dropRate": [0.5], + "max_depth_norm": [1.0], + "return_logits": [False], + "bias": [True, False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.EEGInception, grid, label_keys=("nb_classes", "kernel_size", "bias")) def test_EEGNet(self): - print("Testing EEGnet...", end="", flush=True) - EEGnet_args = { - "nb_classes": [2, 4], - "Chans": [self.Chan], - "Samples": [self.Samples], - "kernLength": [32, 64], - "F1": [4, 8], - "D": [2, 4], - "F2": [8, 16], - "pool1": [4, 8], - "pool2": [8, 16], - "separable_kernel": [16, 32], - "return_logits": [False], - "seed": [42], - } - EEGnet_args = self.makeGrid(EEGnet_args) - for i in EEGnet_args: - model = models.EEGNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGnet_args: - model = models.EEGNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" EEGnet OK: tested", len(EEGnet_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "Samples": [SAMPLES], + "kernLength": [32, 64], + "F1": [4, 8], + "D": [2, 4], + "F2": [8, 16], + "pool1": [4, 8], + "pool2": [8, 16], + "separable_kernel": [16, 32], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.EEGNet, grid, label_keys=("nb_classes", "kernLength", "F1")) def test_EEGSym(self): - print("Testing EEGsym...", end="", flush=True) - EEGsym_args = { - "nb_classes": [2, 4], - "Samples": [2048], - "Chans": [self.Chan], - "Fs": [64], - "scales_time": [(500, 250, 125), (250, 183, 95)], - "lateral_chans": [2, 3], - "first_left": [True, False], - "F": [8, 24], - "pool": [2, 3], - "bias": [True, False], - "return_logits": [False], - "seed": [42], - } - EEGsym_args = self.makeGrid(EEGsym_args) - for i in EEGsym_args: - model = models.EEGSym(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type not in ["cpu", "mps"]: - for i in EEGsym_args: - model = models.EEGSym(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" EEGsym OK: tested", len(EEGsym_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Samples": [2048], + "Chans": [CHAN], + "Fs": [64], + "scales_time": [(500, 250, 125), (250, 183, 95)], + "lateral_chans": [2, 3], + "first_left": [True, False], + "F": [8, 24], + "pool": [2, 3], + "bias": [True, False], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier( + models.EEGSym, grid, allow_mps=False, label_keys=("nb_classes", "lateral_chans", "F") + ) def test_FBCNet(self): - print("Testing FBCNet...", end="", flush=True) - - DCN_args = { - "nb_classes": [2, 4], - "Chans": [self.Chan], - "Samples": [self.Samples], - "Fs": [128], - "FilterBands": [4, 8], - "FilterRange": [4, 5], - "FilterType": ["Cheby2", "ellip"], - "TemporalType": ["var", "max", "mean", "std", "logvar"], - "return_logits": [False], - "seed": [42], - } - DCN_grid = self.makeGrid(DCN_args) - for i in DCN_grid: - model = models.FBCNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - if self.device.type != "cpu": - for i in DCN_grid: - model = models.FBCNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" FBCNet OK: tested ", len(DCN_grid), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "Samples": [SAMPLES], + "Fs": [128], + "FilterBands": [4, 8], + "FilterRange": [4, 5], + "FilterType": ["Cheby2", "ellip"], + "TemporalType": ["var", "max", "mean", "std", "logvar"], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.FBCNet, grid, label_keys=("nb_classes", "FilterType", "TemporalType")) def test_ResNet(self): - print("Testing ResNet...", end="", flush=True) - EEGres_args = { - "nb_classes": [2, 4], - "Samples": [2048], - "Chans": [self.Chan], - "block": [models.BasicBlock1], - "Layers": [[1, 1, 1, 1], [1, 2, 4, 3]], - "inplane": [8, 16], - "kernLength": [7, 13], - "addConnection": [True, False], - "return_logits": [False], - "seed": [42], - } - EEGres_args = self.makeGrid(EEGres_args) - for i in EEGres_args: - model = models.ResNet1D(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGres_args: - model = models.ResNet1D(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" ResNet OK: tested", len(EEGres_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Samples": [2048], + "Chans": [CHAN], + "block": [models.BasicBlock1], + "Layers": [[1, 1, 1, 1], [1, 2, 4, 3]], + "inplane": [8, 16], + "kernLength": [7, 13], + "addConnection": [True, False], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.ResNet1D, grid, label_keys=("nb_classes", "inplane", "addConnection")) def test_ShallowNet(self): - print("Testing ShallowNet...", end="", flush=True) - EEGsha_args = { - "nb_classes": [2, 4], - "Samples": [2048], - "Chans": [self.Chan], - "F": [20, 40], - "K1": [25, 12], - "Pool": [75, 50], - "return_logits": [False], - "seed": [42], - } - EEGsha_args = self.makeGrid(EEGsha_args) - for i in EEGsha_args: - model = models.ShallowNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGsha_args: - model = models.ShallowNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" ShallowNet OK: tested", len(EEGsha_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Samples": [2048], + "Chans": [CHAN], + "F": [20, 40], + "K1": [25, 12], + "Pool": [75, 50], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.ShallowNet, grid, label_keys=("nb_classes", "F", "K1")) def test_StagerNet(self): - print("Testing StageRNet...", end="", flush=True) - EEGsta_args = { - "nb_classes": [2, 4], - "Samples": [2048], - "Chans": [self.Chan], - "F": [8, 16], - "kernLength": [64, 120], - "Pool": [16, 8], - "return_logits": [False], - "seed": [42], - } - EEGsta_args = self.makeGrid(EEGsta_args) - for i in EEGsta_args: - model = models.StagerNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGsta_args: - model = models.StagerNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" StageRNet OK: tested", len(EEGsta_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Samples": [2048], + "Chans": [CHAN], + "F": [8, 16], + "kernLength": [64, 120], + "Pool": [16, 8], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.StagerNet, grid, label_keys=("nb_classes", "F", "kernLength")) def test_STNet(self): - print("Testing STNet...", end="", flush=True) - EEGstn_args = { - "nb_classes": [2, 4], - "Samples": [2048], - "grid_size": [5, 9], - "F": [256, 64], - "kernlength": [5, 7], - "dense_size": [1024, 512], - "return_logits": [False], - "seed": [42], - } - EEGstn_args = self.makeGrid(EEGstn_args) - for i in EEGstn_args: - model = models.STNet(**i) - xst = torch.randn(self.N, i["Samples"], i["grid_size"], i["grid_size"]) - out = model(xst) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGstn_args: - model = models.STNet(**i).to(device=self.device) - xst2 = torch.randn(self.N, i["Samples"], i["grid_size"], i["grid_size"]).to( - device=self.device - ) - out = model(xst2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" STNet OK: tested", len(EEGstn_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Samples": [2048], + "grid_size": [5, 9], + "F": [256, 64], + "kernlength": [5, 7], + "dense_size": [1024, 512], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + + def stnet_input(args, dev): + return torch.randn(N, args["Samples"], args["grid_size"], args["grid_size"], device=dev) + + self._check_classifier( + models.STNet, grid, input_fn=stnet_input, label_keys=("nb_classes", "grid_size", "F") + ) def test_TinySleepNet(self): - print("Testing TinySleepNet...", end="", flush=True) - EEGsleep_args = { - "nb_classes": [2, 4], - "Chans": [self.Chan], - "Fs": [64], - "F": [128, 32], - "kernlength": [8, 30], - "pool": [16, 5], - "hidden_lstm": [128, 50], - "return_logits": [False], - "seed": [42], - } - EEGsleep_args = self.makeGrid(EEGsleep_args) - for i in EEGsleep_args: - model = models.TinySleepNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGsleep_args: - model = models.TinySleepNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" TinySleepNet OK: tested", len(EEGsleep_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "Fs": [64], + "F": [128, 32], + "kernlength": [8, 30], + "pool": [16, 5], + "hidden_lstm": [128, 50], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier(models.TinySleepNet, grid, label_keys=("nb_classes", "F", "hidden_lstm")) def test_xEEGNet(self): - print("Testing xEEGNet...", end="", flush=True) - EEGxeg_args = { - "nb_classes": [4], - "Samples": [2048], - "Chans": [self.Chan], - "Fs": [125], - "F1": [7, 126], - "K1": [125, 75], - "F2": [7, 126], - "Pool": [75, 50], - "random_temporal_filter": [True, False], - "freeze_temporal": [0, 1e12], - "spatial_depthwise": [True, False], - "log_activation_base": ["dB"], - "norm_type": ["batchnorm"], - "global_pooling": [True, False], - "bias": [[False] * 3], - "dense_hidden": [-1, 32], - "return_logits": [False], - "seed": [42], - } - - EEGxeg_args = self.makeGrid(EEGxeg_args) - for i in EEGxeg_args: - if i["F1"] > i["F2"] and i["spatial_depthwise"]: - continue - model = models.xEEGNet(**i) - out = model(self.x) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - - if self.device.type != "cpu": - for i in EEGxeg_args: - if i["F1"] > i["F2"] and i["spatial_depthwise"]: - continue - model = models.xEEGNet(**i).to(device=self.device) - out = model(self.x2) - self.assertEqual(torch.isnan(out).sum(), 0) - self.assertEqual(out.shape[1], i["nb_classes"] if i["nb_classes"] > 2 else 1) - if not (i["return_logits"]): - self.assertLessEqual(out.max(), 1) - self.assertGreaterEqual(out.min(), 0) - print(" xEEGNet OK: tested", len(EEGxeg_args), " combinations of input arguments") + grid = make_grid( + { + "nb_classes": [4], + "Samples": [2048], + "Chans": [CHAN], + "Fs": [125], + "F1": [7, 126], + "K1": [125, 75], + "F2": [7, 126], + "Pool": [75, 50], + "random_temporal_filter": [True, False], + "freeze_temporal": [0, 1e12], + "spatial_depthwise": [True, False], + "log_activation_base": ["dB"], + "norm_type": ["batchnorm"], + "global_pooling": [True, False], + "bias": [[False] * 3], + "dense_hidden": [-1, 32], + "return_logits": [False], + "seed": [42], + }, + max_comb=16, + ) + self._check_classifier( + models.xEEGNet, + grid, + skip=lambda a: a["F1"] > a["F2"] and a["spatial_depthwise"], + label_keys=("F1", "F2", "spatial_depthwise"), + ) if __name__ == "__main__": diff --git a/test/EEGself/ssl/ssl_test.py b/test/EEGself/ssl/ssl_test.py index de6be7b..9e7d2bb 100644 --- a/test/EEGself/ssl/ssl_test.py +++ b/test/EEGself/ssl/ssl_test.py @@ -1,15 +1,23 @@ -import itertools +import glob import os import pickle -import platform import random +import shutil +import sys import unittest import warnings + import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from EEGself._testtools import get_device + import selfeeg from selfeeg import augmentation as aug from selfeeg import dataloading as dl @@ -22,8 +30,7 @@ def loadEEG(path, return_label=False): y = EEG["label"] if return_label: return x, y - else: - return x + return x class Decoder(nn.Module): @@ -47,47 +54,30 @@ def loss_finetuning(self, yhat, ytrue): def loss_finetuning_val(self, yhat, ytrue): return torch.sum((yhat > 0.5) == ytrue) / ytrue.shape[0] - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds - @classmethod def setUpClass(cls): warnings.filterwarnings("ignore", message="Using padding='same'", category=UserWarning) cls.seed = 1234 + cls.device = get_device( + probe=lambda dev: aug.add_band_noise(torch.randn(1024, device=dev), "theta", 128) + ) - if torch.backends.mps.is_available(): - cls.device = torch.device("mps") - elif torch.cuda.is_available(): - cls.device = torch.device("cuda") - else: - cls.device = torch.device("cpu") - - if cls.device.type != "cpu": - try: - xx = torch.randn(1024).to(device=cls.device) - xx = aug.add_band_noise(xx, "theta", 128) - except Exception: - cls.device = torch.device("cpu") - - print("\n---------------------------") - print("TESTING SSL MODULE") - if cls.device.type != "cpu": - print("Found cuda device: testing module on it") - else: - print("Didn't found cuda device: testing module on cpu") - print("---------------------------") cls.eegpath = "Simulated_EEG" - cls.freq = 128 # sampling frequency in [Hz] - cls.overlap = 0.3 # overlap between partitions - cls.window = 1 # window length in [seconds] + cls.freq = 128 + cls.overlap = 0.3 + cls.window = 1 cls.workers = 0 cls.batchsize = 16 cls.Chan = 16 + # A trimmed simulated dataset keeps the training-based tests fast while + # still exercising the full data-loading / SSL pipeline. selfeeg.utils.create_dataset() + keep = sorted(glob.glob(os.path.join(cls.eegpath, "*.pickle")))[:200] + keep = set(keep) + for f in glob.glob(os.path.join(cls.eegpath, "*.pickle")): + if f not in keep: + os.remove(f) cls.EEGlen = dl.get_eeg_partition_number( cls.eegpath, cls.freq, cls.window, cls.overlap, load_function=loadEEG @@ -113,8 +103,7 @@ def setUpClass(cls): ) cls.valloader = DataLoader(dataset=valset, batch_size=cls.batchsize, shuffle=False) - # DEFINE AUGMENTER - AUG_band = aug.DynamicSingleAug( + band = aug.DynamicSingleAug( aug.add_band_noise, discrete_arg={ "bandwidth": ["delta", "theta", "alpha", "beta", (30, 49)], @@ -122,12 +111,12 @@ def setUpClass(cls): "noise_range": 0.5, }, ) - AUG_mask = aug.DynamicSingleAug( + mask = aug.DynamicSingleAug( aug.masking, discrete_arg={"mask_number": [1, 2, 3, 4], "masked_ratio": 0.25} ) - Block1 = aug.RandomAug(AUG_band, AUG_mask, p=[0.7, 0.3]) - Block2 = lambda x: selfeeg.utils.scale_range_soft_clip(x, 500, 1.5, "uV", True) - cls.Augmenter = aug.SequentialAug(Block1, Block2) + block1 = aug.RandomAug(band, mask, p=[0.7, 0.3]) + block2 = lambda x: selfeeg.utils.scale_range_soft_clip(x, 500, 1.5, "uV", True) + cls.Augmenter = aug.SequentialAug(block1, block2) cls.enc = selfeeg.models.ShallowNetEncoder(8, 8) cls.head_size = [16, 32, 32] @@ -139,41 +128,56 @@ def setUp(self): np.random.seed(self.seed) torch.manual_seed(self.seed) + def _fit_and_check(self, model, epochs=2, out_dim=32, **fit_kwargs): + """Fit an SSL model for a few epochs and run the shared sanity checks.""" + model = model.to(self.device) + loss_train = model.fit( + train_dataloader=self.trainloader, + augmenter=self.Augmenter, + epochs=epochs, + validation_dataloader=self.valloader, + verbose=False, + device=self.device, + return_loss_info=True, + **fit_kwargs, + ) + model = model.to("cpu") + self.assertIsInstance(loss_train, dict) + self.assertEqual(model(torch.randn(32, 8, 128)).shape, torch.Size([32, out_dim])) + model.test(self.valloader, augmenter=self.Augmenter, verbose=False) + return loss_train + def test_evaluate_loss(self): - print("testng evaluate loss function...", end="", flush=True) y1 = torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) y2 = torch.sin(torch.linspace(0, 8 * torch.pi, 1024) + torch.pi / 6) loss = selfeeg.ssl.evaluate_loss(torch.nn.functional.mse_loss, [y1, y2]) - self.assertTrue(torch.abs(loss - 0.1341).item() < 1e-4) - print(" evaluate loss ok") + self.assertAlmostEqual(loss.item(), 0.1341, delta=1e-4) def test_EarlyStopping(self): - print("testng EarlyStopper...") - - # first set of assertion on device and preservation of best model - Stopper = selfeeg.ssl.EarlyStopping(device=self.device) + # device placement and preservation of the best model + stopper = selfeeg.ssl.EarlyStopping(device=self.device) eegnet = selfeeg.models.EEGNet(2, 16, 256) optimizer = torch.optim.SGD(eegnet.parameters(), 0.5) ytrue = torch.randn(16) yhat = eegnet(torch.randn(16, 16, 256)).squeeze() - Stopper.rec_best_weights(eegnet) + stopper.rec_best_weights(eegnet) - self.assertEqual(eegnet.Dense.bias.item(), Stopper.best_model["Dense.bias"].item()) - self.assertEqual(Stopper.best_model["Dense.bias"].device.type, self.device.type) + self.assertEqual(eegnet.Dense.bias.item(), stopper.best_model["Dense.bias"].item()) + self.assertEqual(stopper.best_model["Dense.bias"].device.type, self.device.type) loss = torch.nn.functional.binary_cross_entropy_with_logits(yhat, ytrue) loss.backward() optimizer.step() - self.assertNotEqual(eegnet.Dense.bias.item(), Stopper.best_model["Dense.bias"].item()) + self.assertNotEqual(eegnet.Dense.bias.item(), stopper.best_model["Dense.bias"].item()) eegnet = selfeeg.models.EEGNet(2, 16, 256) - Stopper.restore_best_weights(eegnet) - self.assertEqual(eegnet.Dense.bias.item(), Stopper.best_model["Dense.bias"].item()) + stopper.restore_best_weights(eegnet) + self.assertEqual(eegnet.Dense.bias.item(), stopper.best_model["Dense.bias"].item()) self.assertEqual( - Stopper.best_model["Dense.bias"].device.type, eegnet.Dense.bias.device.type + stopper.best_model["Dense.bias"].device.type, eegnet.Dense.bias.device.type ) - TrainSet = dl.EEGDataset( + trainset = dl.EEGDataset( self.EEGlen, self.EEGsplit, [128, 2, 0.3], @@ -183,216 +187,96 @@ def test_EarlyStopping(self): optional_load_fun_args=[True], label_on_load=True, ) - TrainLoader = torch.utils.data.DataLoader(TrainSet, batch_size=32) + trainloader = torch.utils.data.DataLoader(trainset, batch_size=32) eegnet = selfeeg.models.EEGNet(2, 8, 256) - Stopper = selfeeg.ssl.EarlyStopping(patience=1, monitored="train", device=self.device) - Stopper.rec_best_weights(eegnet) # little hack to force early stop correctly - self.assertEqual(Stopper.best_model["Dense.bias"].device.type, self.device.type) + stopper = selfeeg.ssl.EarlyStopping(patience=1, monitored="train", device=self.device) + stopper.rec_best_weights(eegnet) # little hack to force early stop correctly + self.assertEqual(stopper.best_model["Dense.bias"].device.type, self.device.type) eegnet = selfeeg.models.EEGNet(2, 8, 256) - Stopper.best_loss = 0 # little hack to force early stop correctly - loss_info = selfeeg.ssl.fine_tune( + stopper.best_loss = 0 # little hack to force early stop correctly + selfeeg.ssl.fine_tune( eegnet, - TrainLoader, + trainloader, 2, - EarlyStopper=Stopper, + EarlyStopper=stopper, loss_func=self.loss_finetuning, verbose=False, ) - self.assertTrue(Stopper.earlystop) + self.assertTrue(stopper.earlystop) self.assertEqual(eegnet.Dense.bias.device.type, self.device.type) - self.assertEqual(eegnet.Dense.bias.item(), Stopper.best_model["Dense.bias"].item()) - print("testng EarlyStopper... EarlyStopper OK") + self.assertEqual(eegnet.Dense.bias.item(), stopper.best_model["Dense.bias"].item()) def test_SimCLR(self): - print("Testing SimCLR (2 epochs)...", end="", flush=True) - - SelfMdl = selfeeg.ssl.SimCLR(encoder=self.enc, projection_head=self.head_size).to( - device=self.device - ) - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, - epochs=2, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, - ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test(self.valloader, augmenter=self.Augmenter, verbose=False) - print(" SimCLR OK") + model = selfeeg.ssl.SimCLR(encoder=self.enc, projection_head=self.head_size) + self._fit_and_check(model) def test_MoCo(self): - print("Testing MoCo v2 (2 epochs)...", end="", flush=True) - - SelfMdl = selfeeg.ssl.MoCo( + # MoCo v2: momentum encoder + memory bank, custom optimizer/scheduler/loss + model = selfeeg.ssl.MoCo( encoder=self.enc, projection_head=self.head_size, bank_size=1024, m=0.9995 - ).to(device=self.device) - loss = selfeeg.losses.moco_loss - loss_arg = {"temperature": 0.5} - optimizer = torch.optim.SGD(SelfMdl.parameters(), lr=1e-3) + ) + optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.98) - - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, - epochs=2, + self._fit_and_check( + model, optimizer=optimizer, - loss_func=loss, - loss_args=loss_arg, + loss_func=selfeeg.losses.moco_loss, + loss_args={"temperature": 0.5}, lr_scheduler=scheduler, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test(self.valloader, augmenter=self.Augmenter, verbose=False) - print(" MoCo v2 OK") - - print("Testing MoCo v3 (2 epochs)...", end="", flush=True) - SelfMdl = selfeeg.ssl.MoCo( + + # MoCo v3: predictor instead of memory bank + model = selfeeg.ssl.MoCo( encoder=self.enc, projection_head=self.head_size, predictor=self.predictor_size, m=0.9995, - ).to(device=self.device) - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, - epochs=2, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test(self.valloader, augmenter=self.Augmenter, verbose=False) - print(" MoCo v3 OK") + self._fit_and_check(model) def test_BYOL(self): - print("Testing BYOL (10 epochs, Earlystop, Scheduler)...", end="", flush=True) - SelfMdl = selfeeg.ssl.BYOL( + model = selfeeg.ssl.BYOL( encoder=self.enc, projection_head=self.head_size, predictor=self.predictor_size, m=0.9995, - ).to(device=self.device) - - loss = selfeeg.losses.byol_loss - optimizer = torch.optim.Adam(SelfMdl.parameters(), lr=1e-4) + ) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.98) earlystop = selfeeg.ssl.EarlyStopping( - patience=2, - min_delta=1e-05, - record_best_weights=True, - device=self.device, + patience=2, min_delta=1e-05, record_best_weights=True, device=self.device ) - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, + loss_train = self._fit_and_check( + model, epochs=10, EarlyStopper=earlystop, optimizer=optimizer, - loss_func=loss, + loss_func=selfeeg.losses.byol_loss, lr_scheduler=scheduler, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(loss_train[9][1] == earlystop.best_loss) - self.assertTrue( - (optimizer.param_groups[-1]["lr"] - (0.98 ** len(loss_train)) * 1e-4) < 1e-4 + self.assertEqual(loss_train[9][1], earlystop.best_loss) + self.assertAlmostEqual( + optimizer.param_groups[-1]["lr"], (0.98 ** len(loss_train)) * 1e-4, delta=1e-4 ) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test(self.valloader, augmenter=self.Augmenter, verbose=False) - earlystop.restore_best_weights(SelfMdl) - print(" BYOL OK") + earlystop.restore_best_weights(model) def test_SimSiam(self): - print("Testing SimSiam (2 epochs)...", end="", flush=True) - - SelfMdl = selfeeg.ssl.SimSiam( + model = selfeeg.ssl.SimSiam( encoder=self.enc, projection_head=self.head_size, predictor=self.predictor_size - ).to(device=self.device) - - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, - epochs=2, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test(self.valloader, augmenter=self.Augmenter, verbose=False) - print(" SimSiam OK") + self._fit_and_check(model) def test_VICReg(self): - print("Testing VICReg (2 epochs)...", end="", flush=True) - - SelfMdl = selfeeg.ssl.VICReg(encoder=self.enc, projection_head=self.head_size).to( - device=self.device - ) - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, - epochs=2, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, - ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test( - self.valloader, augmenter=self.Augmenter, verbose=False - ) # just to show it works - print(" VICReg OK") + model = selfeeg.ssl.VICReg(encoder=self.enc, projection_head=self.head_size) + self._fit_and_check(model) def test_BarlowTwins(self): - print("Testing BarlowTwins (2 epochs)...", end="", flush=True) - - SelfMdl = selfeeg.ssl.BarlowTwins(encoder=self.enc, projection_head=self.head_size).to( - device=self.device - ) - - loss_train = SelfMdl.fit( - train_dataloader=self.trainloader, - augmenter=self.Augmenter, - epochs=2, - validation_dataloader=self.valloader, - verbose=False, - device=self.device, - return_loss_info=True, - ) - SelfMdl = SelfMdl.to(device="cpu") - self.assertTrue(isinstance(loss_train, dict)) - self.assertTrue(SelfMdl(torch.randn(32, 8, 128)).shape == torch.Size([32, 32])) - loss_test = SelfMdl.test( - self.valloader, augmenter=self.Augmenter, verbose=False - ) # just to show it works - print(" BarlowTwins OK") + model = selfeeg.ssl.BarlowTwins(encoder=self.enc, projection_head=self.head_size) + self._fit_and_check(model) def test_PredictiveSSL(self): - print("Testing Predictive SSL (2 epochs)...", end="", flush=True) - - SelfMdl = selfeeg.ssl.PredictiveSSL(self.enc, [16, 1]) - - AUG_band = aug.DynamicSingleAug( + model = selfeeg.ssl.PredictiveSSL(self.enc, [16, 1]) + band = aug.DynamicSingleAug( aug.add_band_noise, discrete_arg={ "bandwidth": ["delta", "theta", "alpha", "beta", (30, 49)], @@ -400,11 +284,11 @@ def test_PredictiveSSL(self): "noise_range": 0.5, }, ) - AUG_mask = aug.DynamicSingleAug( + mask = aug.DynamicSingleAug( aug.masking, discrete_arg={"mask_number": [1, 2, 3, 4], "masked_ratio": 0.25} ) - augment = aug.RandomAug(AUG_band, AUG_mask, return_index=True) - loss_train = SelfMdl.fit( + augment = aug.RandomAug(band, mask, return_index=True) + loss_train = model.fit( train_dataloader=self.trainloader, epochs=2, loss_func=self.loss_finetuning, @@ -415,7 +299,8 @@ def test_PredictiveSSL(self): device=self.device, return_loss_info=True, ) - loss_test = SelfMdl.test( + self.assertIsInstance(loss_train, dict) + model.test( self.valloader, loss_func=self.loss_finetuning, augmenter=augment, @@ -423,11 +308,8 @@ def test_PredictiveSSL(self): verbose=False, device=self.device, ) - print(" Predictive SSL OK") def test_ReconstructiveSSL(self): - print("Testing Reconstructive SSL (2 epochs)...", end="", flush=True) - dec = Decoder(16, 8, 128) gen = selfeeg.ssl.ReconstructiveSSL(self.enc, dec) loss_train = gen.fit( @@ -439,15 +321,11 @@ def test_ReconstructiveSSL(self): device=self.device, return_loss_info=True, ) - loss_train = gen.test( - self.valloader, augmenter=self.Augmenter, verbose=False, device=self.device - ) - print(" Reconstructive SSL OK") + self.assertIsInstance(loss_train, dict) + gen.test(self.valloader, augmenter=self.Augmenter, verbose=False, device=self.device) def test_finetuning(self): - - print("testing fine-tuning phase (10 epochs)...", end="", flush=True) - TrainSet = dl.EEGDataset( + trainset = dl.EEGDataset( self.EEGlen, self.EEGsplit, [self.freq, self.window, self.overlap], @@ -457,8 +335,8 @@ def test_finetuning(self): optional_load_fun_args=[True], label_on_load=True, ) - TrainLoader = torch.utils.data.DataLoader(TrainSet, batch_size=32) - ValSet = dl.EEGDataset( + trainloader = torch.utils.data.DataLoader(trainset, batch_size=32) + valset = dl.EEGDataset( self.EEGlen, self.EEGsplit, [self.freq, self.window, self.overlap], @@ -468,32 +346,25 @@ def test_finetuning(self): optional_load_fun_args=[True], label_on_load=True, ) - ValLoader = torch.utils.data.DataLoader(ValSet, batch_size=32) + valloader = torch.utils.data.DataLoader(valset, batch_size=32) shanet = selfeeg.models.ShallowNet(2, 8, 128).to(device=self.device) loss_info = selfeeg.ssl.fine_tune( shanet, - TrainLoader, + trainloader, device=self.device, epochs=10, - validation_dataloader=ValLoader, + validation_dataloader=valloader, loss_func=self.loss_finetuning, validation_loss_func=self.loss_finetuning_val, verbose=False, return_loss_info=True, ) - self.assertTrue(loss_info[9][0] < 0.015) - print(" fine-tuning OK") + # training must actually reduce the loss (robust to dataset size) + self.assertLess(loss_info[9][0], loss_info[0][0]) @classmethod def tearDownClass(cls): - print("removing generated residual directory (Simulated_EEG)") - try: - if platform.system() == "Windows": - os.system("rmdir /Q /S Simulated_EEG") # nosec - else: - os.system("rm -r Simulated_EEG") # nosec - except: - print('Failed to delete "Simulated_EEG" folder' " Please do it manually") + shutil.rmtree(cls.eegpath, ignore_errors=True) if __name__ == "__main__": diff --git a/test/EEGself/utils/utils_test.py b/test/EEGself/utils/utils_test.py index 5d784d2..8a83b10 100644 --- a/test/EEGself/utils/utils_test.py +++ b/test/EEGself/utils/utils_test.py @@ -1,151 +1,123 @@ -import itertools +import os import random +import sys import unittest import numpy as np import torch +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from EEGself._testtools import get_device, make_grid + from selfeeg import models, utils +def _sine_batch(scale=1.0): + return torch.zeros(16, 32, 1024) + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) * scale + + class TestUtils(unittest.TestCase): @classmethod def setUpClass(cls): - print("\n--------------------") - print("TESTING UTILS MODULE") - if torch.backends.mps.is_available(): - cls.device = torch.device("mps") - elif torch.cuda.is_available(): - cls.device = torch.device("cuda") - else: - cls.device = torch.device("cpu") - - if cls.device.type != "cpu": - try: - xx = torch.zeros(16, 32, 1024) - xx = xx + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) * 500 - xx = xx.to(device=cls.device) - xx = utils.scale_range_soft_clip(xx, "scale", "uV") - except Exception: - cls.device = torch.device("cpu") - - if cls.device.type != "cpu": - print("Found gpu device: testing module with both cpu and gpu") - else: - print("Didn't found cuda device: testing module with only cpu") - print("--------------------") + cls.device = get_device( + probe=lambda dev: utils.scale_range_soft_clip(_sine_batch(500).to(dev), "scale", "uV") + ) + + def _inputs(self, x): + """Return the input variants (cpu tensor, numpy, gpu tensor) to test.""" + variants = [x, x.numpy()] + if self.device.type != "cpu": + variants.append(x.clone().to(self.device)) + return variants - def makeGrid(self, pars_dict): - keys = pars_dict.keys() - combinations = itertools.product(*pars_dict.values()) - ds = [dict(zip(keys, cc)) for cc in combinations] - return ds + def _assert_scaled(self, x, x_scaled, asintote): + if isinstance(x_scaled, torch.Tensor): + self.assertEqual(torch.isnan(x_scaled).sum().item(), 0) + else: + self.assertEqual(np.isnan(x_scaled).sum(), 0) + self.assertFalse(x.max() <= asintote and x.min() >= -asintote) + self.assertTrue(x_scaled.max() <= asintote and x_scaled.min() >= -asintote) def test_scale_range_soft_clip(self): - print("testing scale range with soft clip function...", end="", flush=True) random.seed(1234) - x = torch.zeros(16, 32, 1024) + torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) * 500 - xnp = x.numpy() - inplist = [x, xnp] - if self.device.type != "cpu": - xgpu = torch.clone(x).to(device=self.device) - inplist = [x, xnp, xgpu] - input_args = { - "x": inplist, - "Range": [200, 300], - "asintote": [1.0, 2.3, 3.5], - "scale": ["uV"], - "exact": [True, False], - } - input_args = self.makeGrid(input_args) - for i in input_args: - x_scaled = utils.scale_range_soft_clip(**i) - if isinstance(i["x"], torch.Tensor): - self.assertTrue(torch.isnan(x_scaled).sum() == 0) - else: - self.assertTrue(np.isnan(x_scaled).sum() == 0) - self.assertFalse(x.max() <= i["asintote"] and x.min() >= -i["asintote"]) - self.assertTrue(x_scaled.max() <= i["asintote"] and x_scaled.min() >= -i["asintote"]) - print(" scale range with soft clip OK") + x = _sine_batch(500) + grid = make_grid( + { + "x": self._inputs(x), + "Range": [200, 300], + "asintote": [1.0, 2.3, 3.5], + "scale": ["uV"], + "exact": [True, False], + }, + max_comb=24, + ) + for args in grid: + with self.subTest(asintote=args["asintote"], exact=args["exact"]): + x_scaled = utils.scale_range_soft_clip(**args) + self._assert_scaled(x, x_scaled, args["asintote"]) def test_RangeScaler(self): - print("testing Range Scaler...", end="", flush=True) random.seed(1234) - x = torch.zeros(16, 32, 1024) - x += torch.sin(torch.linspace(0, 8 * torch.pi, 1024)) * 500 - xnp = x.numpy() - inplist = [x, xnp] - if self.device.type != "cpu": - xgpu = torch.clone(x).to(device=self.device) - inplist = [x, xnp, xgpu] - input_args = { - "x": inplist, - "Range": [200, 300], - "asintote": [1.0, 2.3, 3.5], - "scale": ["uV"], - "exact": [True, False], - } - input_args = self.makeGrid(input_args) - for i in input_args: - Scaler = utils.RangeScaler(i["Range"], i["asintote"], i["scale"], i["exact"]) - x_scaled = Scaler(i["x"]) - if isinstance(i["x"], torch.Tensor): - self.assertTrue(torch.isnan(x_scaled).sum() == 0) - else: - self.assertTrue(np.isnan(x_scaled).sum() == 0) - self.assertFalse(x.max() <= i["asintote"] and x.min() >= -i["asintote"]) - self.assertTrue(x_scaled.max() <= i["asintote"] and x_scaled.min() >= -i["asintote"]) - print(" Range Scaler OK") + x = _sine_batch(500) + grid = make_grid( + { + "x": self._inputs(x), + "Range": [200, 300], + "asintote": [1.0, 2.3, 3.5], + "scale": ["uV"], + "exact": [True, False], + }, + max_comb=24, + ) + for args in grid: + with self.subTest(asintote=args["asintote"], exact=args["exact"]): + scaler = utils.RangeScaler( + args["Range"], args["asintote"], args["scale"], args["exact"] + ) + x_scaled = scaler(args["x"]) + self._assert_scaled(x, x_scaled, args["asintote"]) def test_torch_zscore(self): - print("testing Zscore Scaler...", end="", flush=True) x = torch.ones(2, 16, 16) x[0] += torch.randn(16, 16) - if self.device.type != "cpu": - x = x.to(device=self.device) + x = x.to(self.device) + xz = utils.torch_zscore(x, -2, 1) - self.assertTrue(torch.isnan(xz).sum() == 16 * 16) - self.assertTrue(xz.device.type == self.device.type) + self.assertEqual(torch.isnan(xz).sum().item(), 16 * 16) + self.assertEqual(xz.device.type, self.device.type) - x[1] += torch.randn(16, 16).to(device=self.device) + x[1] += torch.randn(16, 16).to(self.device) xz = utils.torch_zscore(x, -2, 1) - self.assertTrue(torch.isnan(xz).sum() == 0) - self.assertTrue((xz.mean(-2).abs() > 1e-6).sum().item() == 0) - self.assertTrue(((xz.std(-2, correction=0).abs() - 1) > 1e-6).sum().item() == 0) - print(" Zscore Scaler OK") + self.assertEqual(torch.isnan(xz).sum().item(), 0) + self.assertEqual((xz.mean(-2).abs() > 1e-6).sum().item(), 0) + self.assertEqual(((xz.std(-2, correction=0).abs() - 1) > 1e-6).sum().item(), 0) def test_get_subarray_closest_sum(self): - print("testing subarray closest sum function...", end="", flush=True) random.seed(1235) - arr = [i for i in range(1, 100)] - _, best_sub_arr = utils.get_subarray_closest_sum( - arr, 3251, tolerance=1e-4, perseverance=10000 - ) - self.assertEqual(sum(best_sub_arr), 3251) - _, best_sub_arr = utils.get_subarray_closest_sum( - arr, 2497, tolerance=1e-4, perseverance=10000 - ) - self.assertEqual(sum(best_sub_arr), 2497) - print(" subarray closest sum OK") + arr = list(range(1, 100)) + for target in (3251, 2497): + with self.subTest(target=target): + _, best_sub_arr = utils.get_subarray_closest_sum( + arr, target, tolerance=1e-4, perseverance=10000 + ) + self.assertEqual(sum(best_sub_arr), target) def test_check_models(self): - print("testing check models function...", end="", flush=True) model1 = models.EEGNet(4, 8, 512) model2 = models.EEGNet(4, 8, 512) - self.assertFalse(utils.check_models(model1, model2)) # Should return False + self.assertFalse(utils.check_models(model1, model2)) model2.load_state_dict(model1.state_dict()) - self.assertTrue(utils.check_models(model1, model2)) # Should return True - print(" check models OK") + self.assertTrue(utils.check_models(model1, model2)) def test_count_parameters(self): - print("testng count parameters function...\n") mdl = models.ShallowNet(4, 8, 1024) - for n, i in enumerate(mdl.parameters()): # bias require grad put to False - i.requires_grad = False if n in [1, 3, 5, 7] else True - a, b = utils.count_parameters(mdl, True, True, True) - self.assertEqual(b, 23760) # should return True - print("\n count parameters OK") + for n, param in enumerate(mdl.parameters()): # freeze biases + param.requires_grad = n not in [1, 3, 5, 7] + _, trainable = utils.count_parameters(mdl, True, True, True) + self.assertEqual(trainable, 23760) if __name__ == "__main__": diff --git a/test/README.md b/test/README.md index 9eca8dc..1b1cc0c 100644 --- a/test/README.md +++ b/test/README.md @@ -23,7 +23,7 @@ Here is reported a basic list of assertion you must include in your tests based ### General -All the functions must check that any allowed combination of input arguments will not raise an unexpected error. To speed up the procedure, you can create a dictionary of possible values per arguments and create the input iterator with the `makegrid` method included in each unittest.TestCase class. See already implemented testing pipelines. +All the functions must check that any allowed combination of input arguments will not raise an unexpected error. To speed up the procedure, you can create a dictionary of possible values per argument and build the input iterator with the shared `make_grid` helper in `test/EEGself/_testtools.py`. Passing its `max_comb` argument caps the number of tested combinations while still exercising every individual parameter value at least once, which keeps the suite fast. See the already implemented testing pipelines and wrap each combination in a `self.subTest(...)` block so failures report the exact arguments involved. ### Dataloading @@ -76,4 +76,5 @@ All the functions must check that any allowed combination of input arguments wil 1. MacBook Pro 14 inch M2-Pro, MacOS 13.6.2, mps backend 2. Padova Neuroscience Center Server, Ubuntu 18.04.1, Tesla V100 GPU +2. Department of Neuroscience Server, Ubuntu 22.04.5, NVIDIA A30 GPU 3. Custom built PC, Windows 10 22H2, NVIDIA RTX 2080 super From b64c68308de439795b99136ba4db6ed70d46bb5d Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 12:18:38 +0000 Subject: [PATCH 2/8] update github actions --- .github/workflows/doc_build_test.yml | 1 + .github/workflows/selfeeg_test.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/doc_build_test.yml b/.github/workflows/doc_build_test.yml index d2c45a8..dcc3a4a 100644 --- a/.github/workflows/doc_build_test.yml +++ b/.github/workflows/doc_build_test.yml @@ -4,6 +4,7 @@ on: pull_request: branches: - main + - develop paths: - 'selfeeg/**' - 'docs/**' diff --git a/.github/workflows/selfeeg_test.yml b/.github/workflows/selfeeg_test.yml index e149e48..7a3116c 100644 --- a/.github/workflows/selfeeg_test.yml +++ b/.github/workflows/selfeeg_test.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - develop paths: - 'selfeeg/**' - 'test/**' From dccd0f8ae2dfb15ec9345f03d039054f6a8ad44d Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 13:02:00 +0000 Subject: [PATCH 3/8] add develop branch to actions --- .github/workflows/doc_build_test.yml | 4 ++-- .github/workflows/selfeeg_test.yml | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/doc_build_test.yml b/.github/workflows/doc_build_test.yml index dcc3a4a..506c6c4 100644 --- a/.github/workflows/doc_build_test.yml +++ b/.github/workflows/doc_build_test.yml @@ -16,9 +16,9 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v7 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v7 with: python-version: '3.11' cache: 'pip' diff --git a/.github/workflows/selfeeg_test.yml b/.github/workflows/selfeeg_test.yml index 7a3116c..6955b0a 100644 --- a/.github/workflows/selfeeg_test.yml +++ b/.github/workflows/selfeeg_test.yml @@ -24,14 +24,14 @@ jobs: strategy: matrix: - os: [ubuntu-latest, windows-latest, macos-latest] - python-version: ["3.11", "3.12"] + os: [ubuntu-latest] + python-version: ["3.10", "3.11", "3.12"] runs-on: ${{ matrix.os }} steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v7 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + uses: actions/setup-python@v7 with: python-version: ${{ matrix.python-version }} cache: 'pip' From 344e27bb6fc76dd0e76b7b6f0e722ae199d3d823 Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 13:03:55 +0000 Subject: [PATCH 4/8] fix useless import and pylint warnings --- selfeeg/augmentation/compose.py | 6 ++--- selfeeg/augmentation/functional.py | 8 +++---- selfeeg/dataloading/load.py | 38 +++++++++++++++--------------- selfeeg/ssl/base.py | 21 ++++++++--------- selfeeg/ssl/contrastive.py | 32 ++++++++++++------------- selfeeg/ssl/generative.py | 15 +++++------- selfeeg/ssl/predictive.py | 16 ++++++------- selfeeg/utils/utils.py | 14 +++++------ 8 files changed, 71 insertions(+), 79 deletions(-) diff --git a/selfeeg/augmentation/compose.py b/selfeeg/augmentation/compose.py index 772ff9d..8030741 100644 --- a/selfeeg/augmentation/compose.py +++ b/selfeeg/augmentation/compose.py @@ -98,7 +98,7 @@ class StaticSingleAug: """ def __init__( - self, augmentation: "function", arguments: list or dict or list[list or dict] = None + self, augmentation: "function", arguments: list | dict | list[list | dict] = None ): if not (inspect.isfunction(augmentation) or inspect.isbuiltin(augmentation)): @@ -252,8 +252,8 @@ def __init__( self, augmentation, discrete_arg: Dict[str, Any] = None, - range_arg: Dict[str, list[int or float, int or float]] = None, - range_type: Dict[str, str or bool] or list[str or bool] = None, + range_arg: Dict[str, list[int | float, int | float]] = None, + range_type: Dict[str, str | bool] | list[str | bool] = None, ): # set augmentation function diff --git a/selfeeg/augmentation/functional.py b/selfeeg/augmentation/functional.py index cf3ee04..e4b9eb4 100755 --- a/selfeeg/augmentation/functional.py +++ b/selfeeg/augmentation/functional.py @@ -726,7 +726,7 @@ def add_band_noise( x: ArrayLike, bandwidth: list[tuple[float, float], str, float], samplerate: float = 256, - noise_range: float or list[float, float] = None, + noise_range: float | list[float, float] = None, std: float = None, get_noise: bool = False, ) -> tuple[ArrayLike, Optional[ArrayLike]]: @@ -1419,7 +1419,7 @@ def get_filter_coeff( if btype.lower() == "bandstop": btype = "lowpass" else: - message = 'Brainwave "', bandwidth[i], '" not exist. \n' + message = f'Brainwave "{eeg_band}" not exist. \n' message += "Choose between delta, theta, alpha, beta, " message += "gamma, gamma_low, gamma_high" raise ValueError(message) @@ -2829,9 +2829,9 @@ def crop_and_resize( # RE-REFERENCING def change_ref( x: ArrayLike, - mode: str or int = "avg", + mode: str | int = "avg", reference: int = None, - exclude_from_ref: int or list[int] = None, + exclude_from_ref: int | list[int] = None, ) -> ArrayLike: """ changes the reference of all EEG record in the ArrayLike object. diff --git a/selfeeg/dataloading/load.py b/selfeeg/dataloading/load.py index 73bde9f..1ef27ba 100644 --- a/selfeeg/dataloading/load.py +++ b/selfeeg/dataloading/load.py @@ -32,15 +32,15 @@ # get_eeg_partition_number def get_eeg_partition_number( EEGpath: str, - freq: int or float = 250, - window: int or float = 2, + freq: int | float = 250, + window: int | float = 2, overlap: float = 0.10, includePartial: bool = True, - file_format: str or list[str] = "*", + file_format: str | list[str] = "*", load_function: "function" = None, - optional_load_fun_args: list or dict = None, + optional_load_fun_args: list | dict = None, transform_function: "function" = None, - optional_transform_fun_args: list or dict = None, + optional_transform_fun_args: list | dict = None, keep_zero_sample: bool = True, save: bool = False, save_path: str = None, @@ -335,11 +335,11 @@ def get_eeg_split_table( partition_table: pd.DataFrame, test_ratio: float = 0.2, val_ratio: float = 0.2, - test_split_mode: str or int = 2, - val_split_mode: str or int = 2, - exclude_data_id: list or dict = None, - test_data_id: list or dict = None, - val_data_id: list or dict = None, + test_split_mode: str | int = 2, + val_split_mode: str | int = 2, + exclude_data_id: list | dict = None, + test_data_id: list | dict = None, + val_data_id: list | dict = None, val_ratio_on_all_data: bool = True, stratified: bool = False, labels: ArrayLike = None, @@ -862,12 +862,12 @@ def get_eeg_split_table_kfold( partition_table: pd.DataFrame, kfold: int = 10, test_ratio: float = 0.2, - test_split_mode: str or int = 2, - val_split_mode: str or int = 2, - exclude_data_id: list or dict = None, - test_data_id: list or dict = None, + test_split_mode: str | int = 2, + val_split_mode: str | int = 2, + exclude_data_id: list | dict = None, + test_data_id: list | dict = None, stratified: bool = False, - labels: "array like" = None, + labels: ArrayLike = None, dataset_id_extractor: "function" = None, subject_id_extractor: "function" = None, split_tolerance=0.01, @@ -1496,9 +1496,9 @@ def __init__( load_function: "function" = None, transform_function: "function" = None, label_function: "function" = None, - optional_load_fun_args: list or dict = None, - optional_transform_fun_args: list or dict = None, - optional_label_fun_args: list or dict = None, + optional_load_fun_args: list | dict = None, + optional_transform_fun_args: list | dict = None, + optional_label_fun_args: list | dict = None, multilabel_on_load: bool = False, label_on_load: bool = False, label_key: list = None, @@ -1780,7 +1780,7 @@ def __getitem__(self, index): ( *dim_idx, slice(None), - slice(self.currEEG.shape[-1] - Nsample, self.currEEG.shape[-1]), + slice(self.currEEG.shape[-1] - self.Nsample, self.currEEG.shape[-1]), ) ] else: diff --git a/selfeeg/ssl/base.py b/selfeeg/ssl/base.py index 642272a..27b0172 100644 --- a/selfeeg/ssl/base.py +++ b/selfeeg/ssl/base.py @@ -4,7 +4,6 @@ from collections.abc import Iterable, Callable import copy import datetime -from itertools import zip_longest import math import os import random @@ -48,7 +47,7 @@ def _default_augmentation(x): def evaluate_loss( loss_fun: Callable, - arguments: torch.Tensor or list[torch.Tensor], + arguments: torch.Tensor | list[torch.Tensor], loss_arg: Union[list, dict] = None, ): """ @@ -119,15 +118,15 @@ def fine_tune( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], validation_loss_func: Callable = None, - validation_loss_args: list or dict = [], - label_encoder: Callable or list[Callable] = None, + validation_loss_args: list | dict = [], + label_encoder: Callable | list[Callable] = None, lr_scheduler=None, EarlyStopper=None, validation_dataloader: torch.utils.data.DataLoader = None, verbose=True, - device: str or torch.device = None, + device: str | torch.device = None, return_loss_info: bool = False, ) -> Optional[dict]: """ @@ -595,7 +594,7 @@ def __init__( improvement: str = "decrease", monitored: str = "validation", record_best_weights: bool = True, - device: str or torch.device = None, + device: str | torch.device = None, ): if device is None: self.device = torch.device("cpu") @@ -790,7 +789,7 @@ def forward(self, x): def evaluate_loss( self, loss_fun: Callable, - arguments: torch.Tensor or list[torch.Tensors], + arguments: torch.Tensor | list[torch.Tensors], loss_arg: Union[list, dict] = None, ) -> torch.Tensor: """ @@ -906,10 +905,10 @@ def _set_fit_args( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], EarlyStopper=None, validation_dataloader=None, - device: str or torch.device = None, + device: str | torch.device = None, ): # Various checks on input parameters. # If some arguments weren't given they will be automatically set @@ -1015,7 +1014,7 @@ def _set_test_args( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], device: str = None, ): if device == None: diff --git a/selfeeg/ssl/contrastive.py b/selfeeg/ssl/contrastive.py index b9ecb2a..a1586ac 100644 --- a/selfeeg/ssl/contrastive.py +++ b/selfeeg/ssl/contrastive.py @@ -1,17 +1,15 @@ from __future__ import annotations -from collections import OrderedDict from collections.abc import Callable import copy -import os import sys -from typing import Optional, Union +from typing import Union import torch import torch.nn as nn import tqdm -from .base import EarlyStopping, SSLBase +from .base import SSLBase __all__ = [ "BarlowTwins", @@ -139,12 +137,12 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, verbose=True, - device: str or torch.device = None, + device: str | torch.device = None, cat_augmentations: bool = False, return_loss_info: bool = False, ): @@ -373,7 +371,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): @@ -553,7 +551,7 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, @@ -780,7 +778,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): @@ -1087,7 +1085,7 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, @@ -1355,7 +1353,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): @@ -1587,7 +1585,7 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, @@ -1814,7 +1812,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): @@ -1950,7 +1948,7 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, @@ -2076,7 +2074,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): @@ -2165,7 +2163,7 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, @@ -2291,7 +2289,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): diff --git a/selfeeg/ssl/generative.py b/selfeeg/ssl/generative.py index 639f2b0..b6bc516 100644 --- a/selfeeg/ssl/generative.py +++ b/selfeeg/ssl/generative.py @@ -1,17 +1,14 @@ from __future__ import annotations -from collections import OrderedDict -from collections.abc import Iterable, Callable -import copy -import os +from collections.abc import Callable import sys -from typing import Optional, Union +from typing import Union import torch import torch.nn as nn import tqdm -from .base import EarlyStopping, SSLBase +from .base import SSLBase __all__ = ["ReconstructiveSSL"] @@ -84,12 +81,12 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, verbose=True, - device: str or torch.device = None, + device: str | torch.device = None, return_loss_info: bool = False, ): """ @@ -304,7 +301,7 @@ def test( test_dataloader, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], verbose: bool = True, device: str = None, ): diff --git a/selfeeg/ssl/predictive.py b/selfeeg/ssl/predictive.py index efd5170..6f8c43a 100644 --- a/selfeeg/ssl/predictive.py +++ b/selfeeg/ssl/predictive.py @@ -1,17 +1,15 @@ from __future__ import annotations -from collections import OrderedDict -from collections.abc import Iterable, Callable -import copy -import os + +from collections.abc import Callable import sys -from typing import Optional, Union +from typing import Union import torch import torch.nn as nn import tqdm -from .base import EarlyStopping, SSLBase +from .base import SSLBase __all__ = ["PredictiveSSL"] @@ -127,14 +125,14 @@ def fit( optimizer=None, augmenter=None, loss_func: Callable = None, - loss_args: list or dict = [], + loss_args: list | dict = [], lr_scheduler=None, EarlyStopper=None, validation_dataloader=None, augmenter_batch_calls=2, labels_on_dataloader=False, verbose=True, - device: str or torch.device = None, + device: str | torch.device = None, return_loss_info: bool = False, ): """ @@ -465,7 +463,7 @@ def test( test_dataloader, augmenter=None, loss_func=None, - loss_args: list or dict = [], + loss_args: list | dict = [], augmenter_batch_calls=2, labels_on_dataloader=False, verbose: bool = True, diff --git a/selfeeg/utils/utils.py b/selfeeg/utils/utils.py index bd6fa4e..3606683 100644 --- a/selfeeg/utils/utils.py +++ b/selfeeg/utils/utils.py @@ -1,11 +1,11 @@ from __future__ import annotations -import copy import os import pickle import random -from typing import Optional, Sequence, Union +from typing import Optional +from scipy.stats import zscore import numpy as np import pandas as pd import torch @@ -24,7 +24,7 @@ ] -def subarray_closest_sum(arr: ArrayLike, n: int, k: float) -> tuple(ArrayLike, float, float, float): +def subarray_closest_sum(arr: ArrayLike, n: int, k: float) -> tuple[ArrayLike, float, float, float]: """ returns a subarray whose element sum is closest to k. @@ -508,9 +508,9 @@ def __call__(self, x): def torch_pchip( - x: "1D Tensor", - y: "ND Tensor", - xv: "1D Tensor", + x: torch.Tensor, + y: torch.Tensor, + xv: torch.Tensor, save_memory: bool = True, new_y_max_numel: int = 4194304, ) -> torch.Tensor: @@ -835,7 +835,7 @@ def count_parameters( return_table: bool = False, print_table: bool = False, add_not_trainable=False, -) -> [int, Optional[pd.DataFrame]]: +) -> tuple[int, Optional[pd.DataFrame]]: """ counts the number of **trainable parameters** of a Pytorch's nn.Module. From 13395b66e7c443a70c97bc017dbd1b25bb3dc682 Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 13:04:32 +0000 Subject: [PATCH 5/8] add transformeeg model --- selfeeg/models/__init__.py | 2 + selfeeg/models/encoders.py | 259 ++++++++++++++++++++++++++++++-- selfeeg/models/zoo.py | 216 +++++++++++++++++++++++++- test/EEGself/models/zoo_test.py | 28 ++++ 4 files changed, 487 insertions(+), 18 deletions(-) diff --git a/selfeeg/models/__init__.py b/selfeeg/models/__init__.py index 5a336c8..cdd1415 100644 --- a/selfeeg/models/__init__.py +++ b/selfeeg/models/__init__.py @@ -18,6 +18,7 @@ StagerNetEncoder, STNetEncoder, TinySleepNetEncoder, + TransformEEGEncoder, xEEGNetEncoder, ) @@ -34,5 +35,6 @@ StagerNet, STNet, TinySleepNet, + TransformEEG, xEEGNet, ) diff --git a/selfeeg/models/encoders.py b/selfeeg/models/encoders.py index b0c11b0..81edf17 100644 --- a/selfeeg/models/encoders.py +++ b/selfeeg/models/encoders.py @@ -1,12 +1,12 @@ from itertools import chain, combinations +import math from scipy.signal import firwin import torch import torch.nn as nn import torch.nn.functional as F +from typing import Callable from .layers import ( - ConstrainedConv1d, ConstrainedConv2d, - ConstrainedDense, DepthwiseConv2d, SeparableConv2d, FilterBank, @@ -26,6 +26,7 @@ "StagerNetEncoder", "STNetEncoder", "TinySleepNetEncoder", + "TransformEEGEncoder", "xEEGNetEncoder", ] @@ -956,7 +957,7 @@ def __init__( self, Chans: int, block: nn.Module = BasicBlock1, - Layers: "list of 4 ints" = [2, 2, 2, 2], + Layers: list = [2, 2, 2, 2], inplane: int = 16, kernLength: int = 7, addConnection: bool = False, @@ -1856,13 +1857,13 @@ def __init__( self, Chans: int, Samples: int, - Fs: int or float, + Fs: int | float, FilterBands: int = 9, - FilterRange: float or int = 4, + FilterRange: float | int = 4, FilterType: str = "Cheby2", - FilterStopRippple: int or float = 30, - FilterPassRipple: int or float = 3, - FilterRangeTol: int or float = 2, + FilterStopRipple: int | float = 30, + FilterPassRipple: int | float = 3, + FilterRangeTol: int | float = 2, FilterSkipFirst=True, D: int = 32, TemporalType: str = "logvar", @@ -1882,7 +1883,7 @@ def __init__( FilterBands, FilterRange, FilterType, - FilterStopRippple, + FilterStopRipple, FilterPassRipple, FilterRangeTol, FilterSkipFirst, @@ -2015,7 +2016,7 @@ def __init__( nlayers: int = 6, nheads: int = 10, dim_feedforward: int = 160, - activation_transformer: str or Callable = "gelu", + activation_transformer: str | Callable = "gelu", p: float = 0.2, p_transformer: float = 0.5, seed: int = None, @@ -2315,3 +2316,241 @@ def _combinatorial_op(self, N, k): :meta private: """ return int((math.factorial(N)) / (math.factorial(k) * math.factorial(N - k))) + + +# ------------------------------ +# TransformEEG +# ------------------------------ +def _resolve_activation(activation): + """ + Resolve a transformer activation given either a callable or a string alias. + + :meta private: + """ + if callable(activation): + return activation + act_map = { + "relu": F.relu, + "gelu": F.gelu, + "elu": F.elu, + "leaky_relu": F.leaky_relu, + "hardswish": F.hardswish, + } + if isinstance(activation, str) and activation.lower() in act_map: + return act_map[activation.lower()] + raise ValueError( + "activation must be a callable or one of " + "'relu', 'gelu', 'elu', 'leaky_relu', 'hardswish'" + ) + + +class TransformEEGEncoder(nn.Module): + """ + Pytorch implementation of the TransformEEG Encoder. + + See TransformEEG for some references. + The expected **input** is a **3D tensor** with size + (Batch x Channels x Samples). + + The encoder is composed of two stages. First, a depthwise convolutional + tokenizer expands the channel dimension and produces a sequence of + ``Chans * D1 * D2`` feature tokens (the token embedding size). Then, a + stack of standard transformer encoder layers processes the token sequence. + The output is obtained by average pooling the transformer output along the + temporal (token) dimension, returning a 2D tensor of size + (Batch x ``Chans * D1 * D2``). + + Parameters + ---------- + Chans: int + The number of EEG channels. + D1: int, optional + The depth multiplier of the first depthwise convolutional block. + The number of intermediate feature maps will be ``Chans * D1``. + + Default = 2 + D2: int, optional + The depth multiplier of the second depthwise convolutional block. + The token embedding size will be ``Chans * D1 * D2``. + + Default = 2 + kernLength1: int, optional + The kernel length of the first convolutional block. + + Default = 5 + kernLength2: int, optional + The kernel length of the second, third and fourth convolutional blocks. + + Default = 5 + pool: int, optional + The kernel size of the two average pooling layers of the tokenizer. + + Default = 4 + stridePool: int, optional + The stride of the two average pooling layers of the tokenizer. + + Default = 2 + dropRate: float, optional + The dropout probability of the tokenizer in range [0,1]. + + Default = 0.2 + ELUalpha: float, optional + The alpha value of the ELU activation function used in the tokenizer. + + Default = 0.1 + batchMomentum: float, optional + The momentum of the batch normalization layers of the tokenizer. + + Default = 0.25 + num_heads: int, optional + The number of heads of each transformer encoder layer. The token + embedding size (``Chans * D1 * D2``) must be divisible by ``num_heads``. + + Default = 1 + dim_feedforward: int, optional + The dimension of the feedforward block of each transformer encoder layer. + If None, it is set equal to the token embedding size. + + Default = None + num_layers: int, optional + The number of stacked transformer encoder layers. + + Default = 2 + transformer_dropout: float, optional + The dropout probability applied inside the transformer encoder layers. + + Default = 0.2 + activation: str or Callable, optional + The activation function of the transformer encoder layers. It can be a + callable (e.g., ``torch.nn.functional.gelu``) or one of the strings + 'relu', 'gelu', 'elu', 'leaky_relu', 'hardswish'. + + Default = 'hardswish' + seed: int, optional + A custom seed for model initialization. It must be a nonnegative number. + If None is passed, no custom seed will be set. + + Default = None + + Note + ---- + Since the temporal dimension is collapsed with an adaptive average pooling, + the encoder can process inputs with an arbitrary number of samples, provided + that they are long enough not to be entirely consumed by the pooling layers. + + Example + ------- + >>> import selfeeg.models + >>> import torch + >>> x = torch.randn(4, 8, 512) + >>> mdl = models.TransformEEGEncoder(8) + >>> out = mdl(x) + >>> print(out.shape) # shoud return torch.Size([4, 32]) + >>> print(torch.isnan(out).sum()) # shoud return 0 + + """ + + def __init__( + self, + Chans: int, + D1: int = 2, + D2: int = 2, + kernLength1: int = 5, + kernLength2: int = 5, + pool: int = 4, + stridePool: int = 2, + dropRate: float = 0.2, + ELUalpha: float = 0.1, + batchMomentum: float = 0.25, + num_heads: int = 1, + dim_feedforward: int = None, + num_layers: int = 2, + transformer_dropout: float = 0.2, + activation="hardswish", + seed: int = None, + ): + + if D1 < 1 or D2 < 1: + raise ValueError("D1 and D2 must be positive integers") + + super(TransformEEGEncoder, self).__init__() + _reset_seed(seed) + + F1 = Chans * D1 + self.embed_dim = Chans * D1 * D2 + if self.embed_dim % num_heads != 0: + raise ValueError( + "the token embedding size (Chans * D1 * D2 = " + f"{self.embed_dim}) must be divisible by num_heads ({num_heads})" + ) + if dim_feedforward is None: + dim_feedforward = self.embed_dim + + # Depthwise convolutional tokenizer + self.block1 = nn.Sequential( + nn.Conv1d(Chans, F1, kernLength1, padding="same", groups=Chans), + nn.BatchNorm1d(F1, momentum=batchMomentum), + nn.ELU(ELUalpha), + ) + self.pool1 = nn.AvgPool1d(pool, stridePool) + self.drop1 = nn.Dropout1d(dropRate) + + self.block2 = nn.Sequential( + nn.Conv1d(F1, F1, kernLength2, padding="same", groups=F1), + nn.BatchNorm1d(F1, momentum=batchMomentum), + nn.ELU(ELUalpha), + ) + + self.block3 = nn.Sequential( + nn.Conv1d(F1, self.embed_dim, kernLength2, padding="same", groups=F1), + nn.BatchNorm1d(self.embed_dim, momentum=batchMomentum), + nn.ELU(ELUalpha), + ) + self.pool2 = nn.AvgPool1d(pool, stridePool) + self.drop2 = nn.Dropout1d(dropRate) + + self.block4 = nn.Sequential( + nn.Conv1d( + self.embed_dim, self.embed_dim, kernLength2, padding="same", groups=self.embed_dim + ), + nn.BatchNorm1d(self.embed_dim, momentum=batchMomentum), + nn.ELU(ELUalpha), + ) + + # Transformer encoder + _reset_seed(seed) + self.transformer = nn.TransformerEncoder( + nn.TransformerEncoderLayer( + self.embed_dim, + nhead=num_heads, + dim_feedforward=dim_feedforward, + dropout=transformer_dropout, + activation=_resolve_activation(activation), + batch_first=True, + ), + num_layers=num_layers, + enable_nested_tensor=False, + ) + self.pool_lay = nn.AdaptiveAvgPool1d(1) + + def forward(self, x): + """ + :meta private: + """ + # Depthwise convolutional tokenizer with residual connections + x1 = self.block1(x) + x1 = self.pool1(x1) + x1 = self.drop1(x1) + x2 = x1 + self.block2(x1) + x3 = self.block3(x2) + x3 = self.pool2(x3) + x3 = self.drop2(x3) + x4 = x3 + self.block4(x3) + + # Transformer expects (Batch x Tokens x Features) + x = torch.permute(x4, [0, 2, 1]) + x = self.transformer(x) + x = torch.permute(x, [0, 2, 1]) + x = self.pool_lay(x) + x = x.squeeze(-1) + return x diff --git a/selfeeg/models/zoo.py b/selfeeg/models/zoo.py index 717ae90..fd94285 100755 --- a/selfeeg/models/zoo.py +++ b/selfeeg/models/zoo.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from typing import Callable from .layers import ConstrainedDense, ConstrainedConv1d, ConstrainedConv2d from .encoders import ( BasicBlock1, @@ -15,6 +16,7 @@ StagerNetEncoder, STNetEncoder, TinySleepNetEncoder, + TransformEEGEncoder, xEEGNetEncoder, ) from ..utils.utils import _reset_seed @@ -32,6 +34,7 @@ "StagerNet", "STNet", "TinySleepNet", + "TransformEEG", "xEEGNet", ] @@ -960,7 +963,7 @@ def __init__( Chans: int, Samples: int, block: nn.Module = BasicBlock1, - Layers: "list of 4 int" = [2, 2, 2, 2], + Layers: list = [2, 2, 2, 2], inplane: int = 16, kernLength: int = 7, addConnection: bool = False, @@ -1420,13 +1423,13 @@ def __init__( nb_classes: int, Chans: int, Samples: int, - Fs: int or float, + Fs: int | float, FilterBands: int = 9, - FilterRange: float or int = 4, + FilterRange: float | int = 4, FilterType: str = "Cheby2", - FilterStopRippple: int or float = 30, - FilterPassRipple: int or float = 3, - FilterRangeTol: int or float = 2, + FilterStopRipple: int | float = 30, + FilterPassRipple: int | float = 3, + FilterRangeTol: int | float = 2, FilterSkipFirst=True, D: int = 32, TemporalType: str = "logvar", @@ -1451,7 +1454,7 @@ def __init__( FilterBands, FilterRange, FilterType, - FilterStopRippple, + FilterStopRipple, FilterPassRipple, FilterRangeTol, FilterSkipFirst, @@ -1997,7 +2000,7 @@ def __init__( nlayers: int = 6, nheads: int = 10, dim_feedforward: int = 160, - activation_transformer: str or Callable = "gelu", + activation_transformer: str | Callable = "gelu", p: float = 0.2, p_transformer: float = 0.5, mlp_dim: list[int, int] = [256, 32], @@ -2234,3 +2237,200 @@ def forward(self, x): else: x = F.softmax(x, dim=1) return x + + +# ------------------------------ +# TransformEEG +# ------------------------------ +class TransformEEG(nn.Module): + """ + Pytorch implementation of the TransformEEG model. + + TransformEEG is a hybrid convolution-transformer architecture for EEG + decoding. A depthwise convolutional tokenizer turns the multichannel EEG + into a sequence of feature tokens, which is then processed by a stack of + transformer encoder layers. The temporal dimension is finally collapsed + with an adaptive average pooling and mapped to the class scores by a small + multilayer perceptron. + + For more information see the following paper [tfeeg]_ . + The expected **input** is a **3D tensor** with size: + (Batch x Channels x Samples). + + Parameters + ---------- + nb_classes: int + The number of classes. If less than 2, a binary classification problem + is considered (output dimensions will be [batch, 1] in this case). + Chans: int + The number of EEG channels. + D1: int, optional + The depth multiplier of the first depthwise convolutional block. + + Default = 2 + D2: int, optional + The depth multiplier of the second depthwise convolutional block. + The token embedding size will be ``Chans * D1 * D2``. + + Default = 2 + kernLength1: int, optional + The kernel length of the first convolutional block. + + Default = 5 + kernLength2: int, optional + The kernel length of the remaining convolutional blocks. + + Default = 5 + pool: int, optional + The kernel size of the two average pooling layers of the tokenizer. + + Default = 4 + stridePool: int, optional + The stride of the two average pooling layers of the tokenizer. + + Default = 2 + dropRate: float, optional + The dropout probability of the tokenizer in range [0,1]. + + Default = 0.2 + ELUalpha: float, optional + The alpha value of the ELU activation function used in the tokenizer. + + Default = 0.1 + batchMomentum: float, optional + The momentum of the batch normalization layers of the tokenizer. + + Default = 0.25 + num_heads: int, optional + The number of heads of each transformer encoder layer. The token + embedding size (``Chans * D1 * D2``) must be divisible by ``num_heads``. + + Default = 1 + dim_feedforward: int, optional + The dimension of the feedforward block of each transformer encoder layer. + If None, it is set equal to the token embedding size. + + Default = None + num_layers: int, optional + The number of stacked transformer encoder layers. + + Default = 2 + transformer_dropout: float, optional + The dropout probability applied inside the transformer encoder layers. + + Default = 0.2 + activation: str or Callable, optional + The activation function of the transformer encoder layers. It can be a + callable (e.g., ``torch.nn.functional.gelu``) or one of the strings + 'relu', 'gelu', 'elu', 'leaky_relu', 'hardswish'. + + Default = 'hardswish' + dense_hidden: int, optional + The number of hidden units of the classification MLP. If None, it is set + to ``max(embedding_size // 2, 64)``. If set to a value lower than 1, a + single linear layer is used as the classification head. + + Default = None + return_logits: bool, optional + If True, return the output as logit. It is suggested to not use False as + the pytorch crossentropy loss function applies the softmax internally. + + Default = True + seed: int, optional + A custom seed for model initialization. It must be a nonnegative number. + If None is passed, no custom seed will be set. + + Default = None + + References + ---------- + .. [tfeeg] Del Pup et al., TransformEEG: Towards Improving Model + Generalizability in Deep Learning-based EEG Parkinson's Disease Detection. + arXiv preprint. 2025. https://doi.org/10.48550/arXiv.2506.17188 + + Example + ------- + >>> import selfeeg.models + >>> import torch + >>> x = torch.randn(4, 8, 512) + >>> mdl = models.TransformEEG(4, 8) + >>> out = mdl(x) + >>> print(out.shape) # shoud return torch.Size([4, 4]) + >>> print(torch.isnan(out).sum()) # shoud return 0 + + """ + + def __init__( + self, + nb_classes: int, + Chans: int, + D1: int = 2, + D2: int = 2, + kernLength1: int = 5, + kernLength2: int = 5, + pool: int = 4, + stridePool: int = 2, + dropRate: float = 0.2, + ELUalpha: float = 0.1, + batchMomentum: float = 0.25, + num_heads: int = 1, + dim_feedforward: int = None, + num_layers: int = 2, + transformer_dropout: float = 0.2, + activation="hardswish", + dense_hidden: int = None, + return_logits: bool = True, + seed: int = None, + ): + + super(TransformEEG, self).__init__() + _reset_seed(seed) + + self.nb_classes = nb_classes + self.return_logits = return_logits + self.encoder = TransformEEGEncoder( + Chans, + D1, + D2, + kernLength1, + kernLength2, + pool, + stridePool, + dropRate, + ELUalpha, + batchMomentum, + num_heads, + dim_feedforward, + num_layers, + transformer_dropout, + activation, + seed, + ) + self.emb_size = self.encoder.embed_dim + + if dense_hidden is None: + dense_hidden = max(self.emb_size // 2, 64) + + _reset_seed(seed) + out_features = 1 if nb_classes <= 2 else nb_classes + if dense_hidden < 1: + self.Dense = nn.Linear(self.emb_size, out_features) + else: + self.Dense = nn.Sequential( + nn.Linear(self.emb_size, dense_hidden), + nn.LeakyReLU(), + nn.Linear(dense_hidden, out_features), + ) + + def forward(self, x): + """ + :meta private: + """ + x = self.encoder(x) + x = self.Dense(x) + if not (self.return_logits): + if self.nb_classes <= 2: + x = torch.sigmoid(x) + else: + x = F.softmax(x, dim=1) + return x diff --git a/test/EEGself/models/zoo_test.py b/test/EEGself/models/zoo_test.py index 83308ec..f62acc8 100644 --- a/test/EEGself/models/zoo_test.py +++ b/test/EEGself/models/zoo_test.py @@ -290,6 +290,34 @@ def test_TinySleepNet(self): ) self._check_classifier(models.TinySleepNet, grid, label_keys=("nb_classes", "F", "hidden_lstm")) + def test_TransformEEG(self): + # embedding size is Chans * D1 * D2 (a multiple of Chans=8), so + # num_heads in {1, 4} always divides it. + grid = make_grid( + { + "nb_classes": [2, 4], + "Chans": [CHAN], + "D1": [2, 4], + "D2": [2], + "kernLength1": [5, 7], + "kernLength2": [5], + "pool": [4, 3], + "stridePool": [2], + "num_heads": [1, 4], + "num_layers": [2], + "dim_feedforward": [None, 64], + "transformer_dropout": [0.2], + "activation": ["hardswish", "gelu"], + "dense_hidden": [None, 0, 32], + "return_logits": [False], + "seed": [42], + }, + max_comb=12, + ) + self._check_classifier( + models.TransformEEG, grid, label_keys=("nb_classes", "D1", "num_heads", "activation") + ) + def test_xEEGNet(self): grid = make_grid( { From 7e7fa59226365a41b2b3481e64b91508b4e2ce40 Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 13:08:50 +0000 Subject: [PATCH 6/8] update doc build test --- .github/workflows/doc_build_test.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/doc_build_test.yml b/.github/workflows/doc_build_test.yml index 506c6c4..61b71c6 100644 --- a/.github/workflows/doc_build_test.yml +++ b/.github/workflows/doc_build_test.yml @@ -1,6 +1,13 @@ name: Sphinx_build_test on: + push: + branches: + - develop + paths: + - 'selfeeg/**' + - 'docs/**' + - '.github/workflows/doc_build_test.yml' pull_request: branches: - main From 33d739b72ef4fa069398613389fbf45cffd16c7c Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 13:15:21 +0000 Subject: [PATCH 7/8] update zoo.py --- selfeeg/models/zoo.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/selfeeg/models/zoo.py b/selfeeg/models/zoo.py index fd94285..787d0ca 100755 --- a/selfeeg/models/zoo.py +++ b/selfeeg/models/zoo.py @@ -2062,8 +2062,7 @@ class xEEGNet(nn.Module): For more information see the following paper [xEEG]_ . The original implementation of xEEGNet can be found here [xEEGgit]_ . - The expected **input** is a **3D tensor** with size: - (Batch x Channels x Samples). + The expected **input** is a **3D tensor** with size: (Batch x Channels x Samples). Parameters ---------- @@ -2346,7 +2345,7 @@ class TransformEEG(nn.Module): ---------- .. [tfeeg] Del Pup et al., TransformEEG: Towards Improving Model Generalizability in Deep Learning-based EEG Parkinson's Disease Detection. - arXiv preprint. 2025. https://doi.org/10.48550/arXiv.2506.17188 + neurocomputing. 2025. https://doi.org/10.1016/j.neucom.2025.132075 Example ------- From 1b452982f2e5cc816effb3229a74f6405b92bd2d Mon Sep 17 00:00:00 2001 From: fedepup Date: Tue, 4 Aug 2026 13:32:08 +0000 Subject: [PATCH 8/8] version update --- README.md | 2 +- RELEASE.md | 10 ++++++++-- pyproject.toml | 2 +- selfeeg/VERSION.txt | 2 +- selfeeg/models/zoo.py | 3 +-- setup.py | 11 +++-------- 6 files changed, 15 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index 0c300a1..7fbf093 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ [![PyPI](https://img.shields.io/pypi/v/selfeeg?label=PyPI&color=blue)](https://pypi.org/project/selfeeg/) [![Conda](https://img.shields.io/conda/vn/conda-forge/selfeeg.svg?color=blue)](https://anaconda.org/conda-forge/selfeeg) [![Docs](https://img.shields.io/readthedocs/selfeeg)](https://readthedocs.org/projects/selfeeg/) -[![Unittest](https://github.com/MedMaxLab/selfEEG/actions/workflows/python-app.yml/badge.svg)](https://github.com/MedMaxLab/selfEEG/actions/workflows/python-app.yml) +[![Unittest](https://github.com/MedMaxLab/selfEEG/actions/workflows/python-app.yml/badge.svg)](https://github.com/MedMaxLab/selfEEG/actions/workflows/selfeeg_test.yml) [![DOI](https://joss.theoj.org/papers/10.21105/joss.06224/status.svg)](https://doi.org/10.21105/joss.06224) [![License](https://img.shields.io/badge/License-MIT-violet.svg)](https://github.com/MedMaxLab/selfEEG/blob/main/LICENSE.md) diff --git a/RELEASE.md b/RELEASE.md index f28dc99..c023d28 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,11 +1,17 @@ # Version X.X.X (only via git install) +# Version 0.2.2 (latest) + **Functionality** - **models module**: - - Fix issue in EEGConformer (projection from F to d_model missing in the forward). + - Add TransformEEG model. + +**maintenance** + +* Improve unittest time -# Version 0.2.1 (latest) +# Version 0.2.1 **Functionality** diff --git a/pyproject.toml b/pyproject.toml index 1eb7190..6a9dd79 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,4 +4,4 @@ build-backend = "setuptools.build_meta" [tool.black] line-length = 100 -target-version = ['py38', 'py39', 'py310', 'py311'] +target-version = ['py310', 'py311', 'py312', 'py313'] diff --git a/selfeeg/VERSION.txt b/selfeeg/VERSION.txt index 0c62199..ee1372d 100644 --- a/selfeeg/VERSION.txt +++ b/selfeeg/VERSION.txt @@ -1 +1 @@ -0.2.1 +0.2.2 diff --git a/selfeeg/models/zoo.py b/selfeeg/models/zoo.py index 787d0ca..e363e7a 100755 --- a/selfeeg/models/zoo.py +++ b/selfeeg/models/zoo.py @@ -2253,8 +2253,7 @@ class TransformEEG(nn.Module): multilayer perceptron. For more information see the following paper [tfeeg]_ . - The expected **input** is a **3D tensor** with size: - (Batch x Channels x Samples). + The expected **input** is a **3D tensor** with size: (Batch x Channels x Samples). Parameters ---------- diff --git a/setup.py b/setup.py index 67bc455..3eb8264 100644 --- a/setup.py +++ b/setup.py @@ -1,7 +1,4 @@ -# license info - import os - import setuptools requirements = [ @@ -43,14 +40,12 @@ "Environment :: Console", "Environment :: GPU", "Intended Audience :: Science/Research", - "Operating System :: Microsoft :: Windows", - "Operating System :: MacOS", "Operating System :: POSIX :: Linux", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering :: Artificial Intelligence", ], keywords=[ @@ -67,7 +62,7 @@ }, install_requires=requirements, include_package_data=True, - python_requires=">=3.8", + python_requires=">=3.10", extras_require=extra_require, zip_safe=False, )