Fix usage of parallel in EEMD [Github] #145

laszukdawid · Oct 12, 2023 · db2a7b5 · db2a7b5
1 parent 7361a88
commit db2a7b5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/PyEMD/EEMD.py b/PyEMD/EEMD.py
@@ -103,8 +103,6 @@ def __call__(
 
     def __getstate__(self) -> Dict:
         self_dict = self.__dict__.copy()
-        if "pool" in self_dict:
-            del self_dict["pool"]
         return self_dict
 
     def generate_noise(self, scale: float, size: Union[int, Sequence[int]]) -> np.ndarray:
@@ -183,13 +181,14 @@ def eemd(
         # For trial number of iterations perform EMD on a signal
         # with added white noise
         if self.parallel:
-            map_pool = Pool(processes=self.processes)
+            pool = Pool(processes=self.processes)
+            map_pool = pool.map
         else:
             map_pool = map
         all_IMFs = map_pool(self._trial_update, range(self.trials))
 
         if self.parallel:
-            map_pool.close()
+            pool.close()
 
         self._all_imfs = defaultdict(list)
         it = iter if not progress else lambda x: tqdm(x, desc="EEMD", total=self.trials)

diff --git a/PyEMD/tests/test_eemd.py b/PyEMD/tests/test_eemd.py
@@ -139,7 +139,16 @@ def test_eemd_notParallel(self):
 
         self.assertTrue(eIMFs.shape[0] > 0)
         self.assertTrue(eIMFs.shape[1], len(S))
-        self.assertFalse("pool" in eemd.__dict__)
+
+    def test_eemd_yesParallel(self):
+        S = np.random.random(100)
+
+        eemd = EEMD(trials=5, max_imf=2, parallel=True)
+        eemd.EMD.FIXE_H = 2
+        eIMFs = eemd.eemd(S)
+
+        self.assertTrue(eIMFs.shape[0] > 0)
+        self.assertTrue(eIMFs.shape[1], len(S))
 
     def test_imfs_and_residue_accessor(self):
         S = np.random.random(100)