Add optional progress dots to EEMD and CEEMDAN (#125) (#132)

* Add optional progress dots to EEMD and CEEMDAN

* Linting lib was broken due to its dep

* Github testing

* Python3.6 had EOL 12/2021 - issues with github actions

* Typo for github workflow

* Use tqdm as progress bar

Co-authored-by: Dawid <git@dawid.lasz.uk>

.github/workflows/ci-test.yml

@@ -8,17 +8,35 @@ on:
       - 'PyEMD/**.py'
 
 jobs:
+  build-n-test-36:
+    name: Linux 3.6
+    runs-on: ubuntu-20.04
+    steps:
+    - uses: actions/checkout@v3
+    - name: Running Python 3.6
+      uses: actions/setup-python@v4
+      with:
+        python-version: 3.6
+        architecture: x64
+    - name: Install test dependencies
+      run: |
+        pip install --only-binary=numpy,scipy numpy scipy
+        pip install -e .[test]
+    - name: Run tests
+      run: |
+        python -m PyEMD.tests.test_all
+
   build-n-test:
     runs-on: ubuntu-latest
 
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.7, 3.8, 3.9]
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Running Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install test dependencies

PyEMD/CEEMDAN.py

@@ -9,12 +9,12 @@
 .. currentmodule:: CEEMDAN
 """
 
-import itertools
 import logging
 from multiprocessing import Pool
 from typing import Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+from tqdm import tqdm
 
 
 class CEEMDAN:
@@ -117,6 +117,8 @@ def __init__(self, trials: int = 100, epsilon: float = 0.005, ext_EMD=None, para
         self.beta_progress = bool(kwargs.get("beta_progress", True))  # Scale noise by std
         self.random = np.random.RandomState(seed=kwargs.get("seed"))
         self.noise_kind = kwargs.get("noise_kind", "normal")
+
+        self._max_imf = int(kwargs.get("max_imf", 100))
         self.parallel = parallel
         self.processes = kwargs.get("processes")  # Optional[int]
         if self.processes is not None and not self.parallel:
@@ -133,8 +135,10 @@ def __init__(self, trials: int = 100, epsilon: float = 0.005, ext_EMD=None, para
         self.C_IMF = None  # Optional[np.ndarray]
         self.residue = None  # Optional[np.ndarray]
 
-    def __call__(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1) -> np.ndarray:
-        return self.ceemdan(S, T=T, max_imf=max_imf)
+    def __call__(
+        self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1, progress: bool = False
+    ) -> np.ndarray:
+        return self.ceemdan(S, T=T, max_imf=max_imf, progress=progress)
 
     def __getstate__(self) -> Dict:
         self_dict = self.__dict__.copy()
@@ -182,7 +186,9 @@ def noise_seed(self, seed: int) -> None:
         """Set seed for noise generation."""
         self.random.seed(seed)
 
-    def ceemdan(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1) -> np.ndarray:
+    def ceemdan(
+        self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1, progress: bool = False
+    ) -> np.ndarray:
         """Perform CEEMDAN decomposition.
 
         Parameters
@@ -193,13 +199,14 @@ def ceemdan(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int =
             Time (x) values for the signal. If not passed, i.e. `T = None`, then assumes equidistant values.
         max_imf : int (default: -1)
             Maximum number of components to extract.
+        progress : bool (default: False)
+            Whether to print out '.' every 1s to indicate progress.
 
         Returns
         -------
         components : np.ndarray
             CEEMDAN components.
         """
-
         scale_s = np.std(S)
         S = S / scale_s
 
@@ -211,14 +218,16 @@ def ceemdan(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int =
         self.all_noise_EMD = self._decompose_noise()
 
         # Create first IMF
-        last_imf = self._eemd(S, T, 1)[0]
+        last_imf = self._eemd(S, T, 1, progress)[0]
         res = np.empty(S.size)
 
         all_cimfs = last_imf.reshape((-1, last_imf.size))
         prev_res = S - last_imf
 
         self.logger.debug("Starting CEEMDAN")
-        while True:
+        total = (max_imf - 1) if max_imf != -1 else None
+        it = iter if not progress else lambda x: tqdm(x, desc="cIMF decomposition", total=total)
+        for _ in it(range(self._max_imf)):
             # Check end condition in the beginning because we've already have 1 IMF
             if self.end_condition(S, all_cimfs, max_imf):
                 self.logger.debug("End Condition - Pass")
@@ -321,7 +330,7 @@ def _decompose_noise(self) -> List[np.ndarray]:
 
         return all_noise_EMD
 
-    def _eemd(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1) -> np.ndarray:
+    def _eemd(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1, progress=True) -> np.ndarray:
         if T is None:
             T = np.arange(len(S), dtype=S.dtype)
 
@@ -334,20 +343,22 @@ def _eemd(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1
         # with added white noise
         if self.parallel:
             pool = Pool(processes=self.processes)
-            all_IMFs = pool.map(self._trial_update, range(self.trials))
-            pool.close()
-
+            map_pool = pool.imap_unordered
         else:  # Not parallel
-            all_IMFs = map(self._trial_update, range(self.trials))
+            map_pool = map
 
-        all_IMFs_1, all_IMFs_2 = itertools.tee(all_IMFs, 2)
+        self.E_IMF = np.zeros((1, N))
+        it = iter if not progress else lambda x: tqdm(x, desc="Decomposing noise", total=self.trials)
 
-        max_imfNo = max([IMFs.shape[0] for IMFs in all_IMFs_1])
-
-        self.E_IMF = np.zeros((max_imfNo, N))
-        for IMFs in all_IMFs_2:
+        for IMFs in it(map_pool(self._trial_update, range(self.trials))):
+            if self.E_IMF.shape[0] < IMFs.shape[0]:
+                num_new_layers = IMFs.shape[0] - self.E_IMF.shape[0]
+                self.E_IMF = np.vstack((self.E_IMF, np.zeros(shape=(num_new_layers, N))))
             self.E_IMF[: IMFs.shape[0]] += IMFs
 
+        if self.parallel:
+            pool.close()
+
         return self.E_IMF / self.trials
 
     def _trial_update(self, trial: int) -> np.ndarray:

PyEMD/EEMD.py

@@ -15,6 +15,7 @@
 from typing import Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
+from tqdm import tqdm
 
 from PyEMD.utils import get_timeline
 
@@ -96,8 +97,10 @@ def __init__(self, trials: int = 100, noise_width: float = 0.05, ext_EMD=None, p
         self.residue = None  # Optional[np.ndarray]
         self._all_imfs = {}
 
-    def __call__(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1) -> np.ndarray:
-        return self.eemd(S, T=T, max_imf=max_imf)
+    def __call__(
+        self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1, progress: bool = False
+    ) -> np.ndarray:
+        return self.eemd(S, T=T, max_imf=max_imf, progress=progress)
 
     def __getstate__(self) -> Dict:
         self_dict = self.__dict__.copy()
@@ -141,7 +144,9 @@ def noise_seed(self, seed: int) -> None:
         """Set seed for noise generation."""
         self.random.seed(seed)
 
-    def eemd(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1) -> np.ndarray:
+    def eemd(
+        self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1, progress: bool = False
+    ) -> np.ndarray:
         """
         Performs EEMD on provided signal.
 
@@ -179,15 +184,17 @@ def eemd(self, S: np.ndarray, T: Optional[np.ndarray] = None, max_imf: int = -1)
         # For trial number of iterations perform EMD on a signal
         # with added white noise
         if self.parallel:
-            pool = Pool(processes=self.processes)
-            all_IMFs = pool.map(self._trial_update, range(self.trials))
-            pool.close()
+            map_pool = Pool(processes=self.processes)
+        else:
+            map_pool = map
+        all_IMFs = map_pool(self._trial_update, range(self.trials))
 
-        else:  # Not parallel
-            all_IMFs = map(self._trial_update, range(self.trials))
+        if self.parallel:
+            map_pool.close()
 
         self._all_imfs = defaultdict(list)
-        for (imfs, trend) in all_IMFs:
+        it = iter if not progress else lambda x: tqdm(x, desc="EEMD", total=self.trials)
+        for (imfs, trend) in it(all_IMFs):
 
             # A bit of explanation here.
             # If the `trend` is not None, that means it was intentionally separated in the decomp process.

PyEMD/__init__.py

@@ -1,6 +1,6 @@
 import logging
 
-__version__ = "1.2.3"
+__version__ = "1.3.0"
 logger = logging.getLogger("pyemd")
 
 from PyEMD.CEEMDAN import CEEMDAN  # noqa

setup.cfg

@@ -22,6 +22,7 @@ install_requires =
     numpy >= 1.12
     scipy >= 0.19
     pathos >= 0.2.1
+    tqdm == 4.64.*
 python_requires = >=3.6, <4
 test_suite = PyEMD.tests
 
@@ -34,6 +35,7 @@ dev =
     pycodestyle==2.8.*
     black==21.10b0
     isort==5.10.*
+    click==8.0.4
 test =
     pytest
     codecov