Update pool decorator, and multiprocessing (#539)

closes #535 #536

* 1. update pool_decorator 
* 2. convert lambda func to def func to avoid conflicts with pickle

Author: Yaoyx

cooltools/api/dotfinder.py

@@ -1262,6 +1262,9 @@ def cluster_filtering_hiccups(
 # large helper functions wrapping smaller step-specific ones
 ####################################################################
 
+def _compose_score_hist(tile, to_score, to_hist):
+    return to_hist(to_score(tile))
+
 @pool_decorator
 def scoring_and_histogramming_step(
     clr,
@@ -1300,7 +1303,7 @@ def scoring_and_histogramming_step(
     to_hist = partial(histogram_scored_pixels, kernels=kernels, ledges=ledges)
 
     # compose scoring and histogramming together :
-    job = lambda tile: to_hist(to_score(tile))
+    job = partial(_compose_score_hist, to_score=to_score, to_hist=to_hist)
 
     # standard multiprocessing implementation
     if nproc > 1:
@@ -1388,7 +1391,7 @@ def scoring_and_extraction_step(
     )
 
     # compose scoring and histogramming together
-    job = lambda tile: to_extract(to_score(tile))
+    job = partial(_compose_score_hist, to_score=to_score, to_hist=to_extract)
 
     # standard multiprocessing implementation
     if nproc > 1:

cooltools/api/expected.py

@@ -1035,6 +1035,10 @@ def per_region_smooth_cvd(
     )
 
     return cvd
+
+def _balance_transform(p, weight1, weight2):
+    return p["count"] * p[weight1] * p[weight2]
+
 # user-friendly wrapper for diagsum_symm and diagsum_pairwise - part of new "public" API
 @pool_decorator
 def expected_cis(
@@ -1179,7 +1183,7 @@ def expected_cis(
         # define balanced data transform:
         weight1 = clr_weight_name + "1"
         weight2 = clr_weight_name + "2"
-        transforms = {"balanced": lambda p: p["count"] * p[weight1] * p[weight2]}
+        transforms = {"balanced": partial(_balance_transform, weight1=weight1, weight2=weight2)}
     else:
         raise ValueError(
             "cooler is not balanced, or"
@@ -1317,8 +1321,7 @@ def expected_trans(
         # define balanced data transform:
         weight1 = clr_weight_name + "1"
         weight2 = clr_weight_name + "2"
-        transforms = {"balanced": lambda p: p["count"] * p[weight1] * p[weight2]}
-
+        transforms = {"balanced": partial(_balance_transform, weight1=weight1, weight2=weight2)}
     else:
         raise ValueError(
             "cooler is not balanced, or"

cooltools/lib/common.py

@@ -2,7 +2,7 @@
 import numpy as np
 import pandas as pd
 import bioframe
-from multiprocess import Pool
+from multiprocessing import Pool
 from functools import wraps
 import logging
 
@@ -526,7 +526,7 @@ def wrapper(*args, **kwargs):
         # If alternative or third party map functors are provided
         if "map_functor" in kwargs.keys():
             logging.info(f"using an alternative map functor: {kwargs['map_functor']}")
-            return func(*args, **kwargs, map_functor=kwargs["map_functor"])
+            return func(*args, **kwargs)
 
         pool = None
         if "nproc" in kwargs.keys():