Implement comprehensive Registry Pattern for system-wide modularity and decoupling

.github/workflows/test-databases.yml

@@ -18,7 +18,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.12", "3.13"]
       fail-fast: false
 
     steps:

.github/workflows/test-models.yml

@@ -18,7 +18,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.12", "3.13"]
       fail-fast: false
 
     steps:

.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
     hooks:
       - id: flake8
         additional_dependencies: [pycodestyle>=2.11.0]
-        args: [--max-line-length=128, '--exclude=./.*,build,dist', '--ignore=E501,W503,E203,F841,E231,W604,E226', --count, --statistics, --show-source]
+        args: [--max-line-length=128, '--exclude=./.*,build,dist', '--ignore=E501,W503,E203,F841,E231,W604,E226,E225', --count, --statistics, --show-source]
   - repo: https://github.com/pycqa/isort
     rev: 8.0.1
     hooks:

torch_ecg/augmenters/__init__.py

@@ -15,6 +15,7 @@
     :toctree: generated/
     :recursive:
 
+    AUGMENTERS
     AugmenterManager
     Augmenter
     BaselineWanderAugmenter
@@ -38,9 +39,11 @@
 from .random_flip import RandomFlip
 from .random_masking import RandomMasking
 from .random_renormalize import RandomRenormalize
+from .registry import AUGMENTERS
 from .stretch_compress import StretchCompress, StretchCompressOffline
 
 __all__ = [
+    "AUGMENTERS",
     "Augmenter",
     "BaselineWanderAugmenter",
     "CutMix",

torch_ecg/augmenters/augmenter_manager.py

@@ -9,13 +9,7 @@
 
 from ..utils.misc import add_docstring, default_class_repr
 from .base import Augmenter, _augmenter_forward_doc
-from .baseline_wander import BaselineWanderAugmenter
-from .label_smooth import LabelSmooth
-from .mixup import Mixup
-from .random_flip import RandomFlip
-from .random_masking import RandomMasking
-from .random_renormalize import RandomRenormalize
-from .stretch_compress import StretchCompress
+from .registry import AUGMENTERS
 
 __all__ = [
     "AugmenterManager",
@@ -65,84 +59,34 @@ class AugmenterManager(torch.nn.Module):
     def __init__(self, *augs: Optional[Tuple[Augmenter, ...]], random: bool = False) -> None:
         super().__init__()
         self.random = random
-        self._augmenters = list(augs)
+        self._augmenters = torch.nn.ModuleList(list(augs))
 
-    def _add_baseline_wander(self, **config: dict) -> None:
-        """Add the baseline wander augmenter to the manager.
+    def add_(self, aug: Union[Augmenter, str, dict], pos: int = -1, **kwargs: Any) -> None:
+        """Add an augmenter to the manager.
 
         Parameters
         ----------
-        **config : dict
-            The configuration for the baseline wander augmenter.
+        aug : Augmenter or str or dict
+            The augmenter to be added.
+            If it's a string or a dict, it will be built using the registry.
+        pos : int, default -1
+            The position to insert the augmenter.
+            Should be >= -1, with -1 being the indicator of the end.
+        **kwargs : Any
+            Additional keyword arguments for building the augmenter.
 
         """
-        self._augmenters.append(BaselineWanderAugmenter(**config))
-
-    def _add_label_smooth(self, **config: dict) -> None:
-        """Add the label smooth augmenter to the manager.
-
-        Parameters
-        ----------
-        **config : dict
-            The configuration for the label smooth augmenter.
-
-        """
-        self._augmenters.append(LabelSmooth(**config))
-
-    def _add_mixup(self, **config: dict) -> None:
-        """Add the mixup augmenter to the manager.
-
-        Parameters
-        ----------
-        **config : dict
-            The configuration for the mixup augmenter.
-
-        """
-        self._augmenters.append(Mixup(**config))
-
-    def _add_random_flip(self, **config: dict) -> None:
-        """Add the random flip augmenter to the manager.
-
-        Parameters
-        ----------
-        **config : dict
-            The configuration for the random flip augmenter.
-
-        """
-        self._augmenters.append(RandomFlip(**config))
-
-    def _add_random_masking(self, **config: dict) -> None:
-        """Add the random masking augmenter to the manager.
-
-        Parameters
-        ----------
-        **config : dict
-            The configuration for the random masking augmenter.
-
-        """
-        self._augmenters.append(RandomMasking(**config))
-
-    def _add_random_renormalize(self, **config: dict) -> None:
-        """Add the random renormalize augmenter to the manager.
-
-        Parameters
-        ----------
-        **config : dict
-            The configuration for the random renormalize augmenter.
-
-        """
-        self._augmenters.append(RandomRenormalize(**config))
-
-    def _add_stretch_compress(self, **config: dict) -> None:
-        """Add the stretch compress augmenter to the manager.
-
-        Parameters
-        ----------
-        **config : dict
-            The configuration for the stretch compress augmenter.
-
-        """
-        self._augmenters.append(StretchCompress(**config))
+        if isinstance(aug, (str, dict)):
+            aug = AUGMENTERS.build(aug, **kwargs)
+        assert isinstance(aug, Augmenter)
+        assert aug.__class__.__name__ not in [
+            a.__class__.__name__ for a in self.augmenters
+        ], f"Augmenter {aug.__class__.__name__} already exists."
+        assert isinstance(pos, int) and pos >= -1, f"pos must be an integer >= -1, but got {pos}."
+        if pos == -1:
+            self._augmenters.append(aug)
+        else:
+            self._augmenters.insert(pos, aug)
 
     @add_docstring(
         _augmenter_forward_doc.replace(
@@ -168,7 +112,7 @@ def forward(
         return (sig, label, *extra_tensors)
 
     @property
-    def augmenters(self) -> List[Augmenter]:
+    def augmenters(self) -> torch.nn.ModuleList:
         """The list of augmenters in the manager."""
         return self._augmenters
 
@@ -199,23 +143,24 @@ def from_config(cls, config: dict) -> "AugmenterManager":
 
         """
         am = cls(random=config.get("random", False))
-        _mapping = {
-            "baseline_wander": am._add_baseline_wander,
-            "label_smooth": am._add_label_smooth,
-            "mixup": am._add_mixup,
-            "random_flip": am._add_random_flip,
-            "random_masking": am._add_random_masking,
-            "random_renormalize": am._add_random_renormalize,
-            "stretch_compress": am._add_stretch_compress,
-        }
         for aug_name, aug_config in config.items():
             if aug_name in [
                 "fs",
                 "random",
             ]:
                 continue
-            elif aug_name in _mapping and isinstance(aug_config, dict):
-                _mapping[aug_name](fs=config["fs"], **aug_config)
+            if aug_name in AUGMENTERS or aug_name in [
+                "".join([w.capitalize() for w in k.split("_")]) for k in AUGMENTERS.list_all()
+            ]:
+                if aug_config is False:
+                    continue
+                if isinstance(aug_config, dict):
+                    # add default fs from config if not specified in aug_config
+                    if "fs" not in aug_config and "fs" in config:
+                        aug_config["fs"] = config["fs"]
+                    am.add_(aug_name, **aug_config)
+                else:
+                    am.add_(aug_name)
             else:
                 # just ignore the other items
                 pass
@@ -236,9 +181,10 @@ def rearrange(self, new_ordering: List[str]) -> None:
                 "The augmenters are applied in random order, " "rearranging the augmenters will not take effect.",
                 RuntimeWarning,
             )
+        # TODO: use a more robust way to map names to classes
         _mapping = {  # built-in augmenters
             "".join([w.capitalize() for w in k.split("_")]): k
-            for k in "label_smooth,mixup,random_flip,random_masking,random_renormalize,stretch_compress".split(",")
+            for k in "label_smooth,mixup,random_flip,random_masking,random_renormalize,stretch_compress,cutmix".split(",")
         }
         _mapping.update({"BaselineWanderAugmenter": "baseline_wander"})
         _mapping.update({v: k for k, v in _mapping.items()})
@@ -250,4 +196,6 @@ def rearrange(self, new_ordering: List[str]) -> None:
         assert len(new_ordering) == len(set(new_ordering)), "Duplicate augmenter names."
         assert len(new_ordering) == len(self._augmenters), "Number of augmenters mismatch."
 
-        self._augmenters.sort(key=lambda aug: new_ordering.index(_mapping[aug.__class__.__name__]))
+        augs = list(self._augmenters)
+        augs.sort(key=lambda aug: new_ordering.index(_mapping[aug.__class__.__name__]))
+        self._augmenters = torch.nn.ModuleList(augs)

torch_ecg/augmenters/baseline_wander.py

@@ -13,12 +13,15 @@
 from ..cfg import DEFAULTS
 from ..utils.utils_signal import get_ampl
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "BaselineWanderAugmenter",
 ]
 
 
+@AUGMENTERS.register(name="baseline_wander")
+@AUGMENTERS.register()
 class BaselineWanderAugmenter(Augmenter):
     """Generate baseline wander composed of
     sinusoidal and Gaussian noise.

torch_ecg/augmenters/cutmix.py

@@ -11,12 +11,15 @@
 
 from ..cfg import DEFAULTS
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "CutMix",
 ]
 
 
+@AUGMENTERS.register(name="cutmix")
+@AUGMENTERS.register()
 class CutMix(Augmenter):
     """CutMix augmentation.
 

torch_ecg/augmenters/label_smooth.py

@@ -6,12 +6,15 @@
 from torch import Tensor
 
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "LabelSmooth",
 ]
 
 
+@AUGMENTERS.register(name="label_smooth")
+@AUGMENTERS.register()
 class LabelSmooth(Augmenter):
     """Label smoothing augmentation.
 

torch_ecg/augmenters/mixup.py

@@ -11,12 +11,15 @@
 
 from ..cfg import DEFAULTS
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "Mixup",
 ]
 
 
+@AUGMENTERS.register(name="mixup")
+@AUGMENTERS.register()
 class Mixup(Augmenter):
     """Mixup augmentor.
 

torch_ecg/augmenters/random_flip.py

@@ -8,12 +8,15 @@
 from torch import Tensor
 
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "RandomFlip",
 ]
 
 
+@AUGMENTERS.register(name="random_flip")
+@AUGMENTERS.register()
 class RandomFlip(Augmenter):
     """Randomly flip the ECGs along the voltage axis.
 

torch_ecg/augmenters/random_masking.py

@@ -9,12 +9,15 @@
 from torch import Tensor
 
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "RandomMasking",
 ]
 
 
+@AUGMENTERS.register(name="random_masking")
+@AUGMENTERS.register()
 class RandomMasking(Augmenter):
     """Randomly mask ECGs with a probability.
 

torch_ecg/augmenters/random_renormalize.py

@@ -9,12 +9,15 @@
 from ..cfg import DEFAULTS
 from ..utils.utils_signal_t import normalize as normalize_t
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "RandomRenormalize",
 ]
 
 
+@AUGMENTERS.register(name="random_renormalize")
+@AUGMENTERS.register()
 class RandomRenormalize(Augmenter):
     """Randomly re-normalize the ECG tensor,
     using the Z-score normalization method.

torch_ecg/augmenters/registry.py

@@ -0,0 +1,11 @@
+"""
+Registry for augmenters.
+"""
+
+from ..utils.registry import Registry
+
+__all__ = [
+    "AUGMENTERS",
+]
+
+AUGMENTERS = Registry("augmenters")

torch_ecg/augmenters/stretch_compress.py

@@ -14,13 +14,16 @@
 from ..cfg import DEFAULTS
 from ..utils.misc import ReprMixin, add_docstring
 from .base import Augmenter
+from .registry import AUGMENTERS
 
 __all__ = [
     "StretchCompress",
     "StretchCompressOffline",
 ]
 
 
+@AUGMENTERS.register(name="stretch_compress")
+@AUGMENTERS.register()
 class StretchCompress(Augmenter):
     """Stretch-or-compress augmenter on ECG tensors.
 

torch_ecg/components/registry.py

@@ -0,0 +1,15 @@
+"""
+Registries for components (optimizers, schedulers, losses, etc.).
+"""
+
+from ..utils.registry import Registry
+
+__all__ = [
+    "OPTIMIZERS",
+    "SCHEDULERS",
+    "LOSSES",
+]
+
+OPTIMIZERS = Registry("optimizers")
+SCHEDULERS = Registry("schedulers")
+LOSSES = Registry("losses")