Update pool decorator, and multiprocessing

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 _transformer(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(_transformer, 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(_transformer, 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():

tests/test_expected.py

@@ -732,4 +732,4 @@ def test_diagsum_from_array():
     exp = _diagsum_symm_dense(ar, bad_bins=list(range(3, 5)))
     exp1 = diagsum_from_array(ar, ignore_diags=0)
     exp1["balanced.avg"] = exp1["balanced.sum"] / exp1["n_valid"]
-    assert np.allclose(exp, exp1["balanced.avg"].values, equal_nan=True)
+    assert np.allclose(exp, exp1["balanced.avg"].values, equal_nan=True)