dwelltime: allow parallel bootstraps

Author: JoepVanlier

lumicks/pylake/population/dwelltime.py

@@ -1,6 +1,7 @@
 import warnings
 from typing import Dict, Tuple, Union
 from dataclasses import field, dataclass
+from multiprocessing import Pool
 
 import numpy as np
 import scipy
@@ -292,7 +293,7 @@ def __post_init__(self):
             )
 
     @classmethod
-    def _from_dwelltime_model(cls, optimized, iterations):
+    def _from_dwelltime_model(cls, optimized, iterations, num_processes=None):
         """Construct bootstrap distributions for parameters from an optimized
         :class:`~lumicks.pylake.DwelltimeModel`.
 
@@ -306,9 +307,27 @@ def _from_dwelltime_model(cls, optimized, iterations):
             optimized model results
         iterations : int
             number of iterations (random samples) to use for the bootstrap
+        num_processes : int | None
+            number of processes to use for parallelization. If `None`, no parallelization is used.
         """
-        samples = DwelltimeBootstrap._sample(optimized, iterations)
-        return cls(optimized, samples[: optimized.n_components], samples[optimized.n_components :])
+
+        def from_samples(samples):
+            return cls(
+                optimized,
+                samples[: optimized.n_components],
+                samples[optimized.n_components :],
+            )
+
+        if num_processes is None:
+            return from_samples(DwelltimeBootstrap._sample(optimized, iterations))
+
+        with Pool(num_processes) as p:
+            result = p.starmap(
+                DwelltimeBootstrap._sample,
+                [(optimized, int(np.ceil(iterations / num_processes)))] * num_processes,
+            )
+
+        return from_samples(np.hstack(result)[:, :iterations])
 
     def extend(self, iterations):
         """Extend the distribution by additional sampling iterations.
@@ -721,15 +740,19 @@ def bic(self) -> float:
         n = self.dwelltimes.size  # number of observations
         return k * np.log(n) - 2 * self.log_likelihood
 
-    def calculate_bootstrap(self, iterations=500):
+    def calculate_bootstrap(self, iterations=500, *, num_processes=None):
         """Calculate a bootstrap distribution for the model.
 
         Parameters
         ----------
         iterations : int
             Number of iterations to sample for the distribution.
+        num_processes : int | None
+            Number of processes to use for parallelization. If `None`, no parallelization is used.
         """
-        bootstrap = DwelltimeBootstrap._from_dwelltime_model(self, iterations)
+        bootstrap = DwelltimeBootstrap._from_dwelltime_model(
+            self, iterations, num_processes=num_processes
+        )
         # TODO: remove with deprecation
         self._bootstrap = bootstrap
         return bootstrap

lumicks/pylake/population/tests/test_dwelltimes.py

@@ -158,6 +158,22 @@ def test_bootstrap_multi(min_obs, max_obs, ref_ci, time_step):
     np.testing.assert_allclose(ci, ref_ci, rtol=1e-5)
 
 
+@pytest.mark.parametrize("num_components", [1, 2])
+def test_bootstrap_parallel(monkeypatch, exponential_data, num_components):
+    import multiprocessing.dummy
+
+    dataset = exponential_data["dataset_2exp"]
+    fit = DwelltimeModel(
+        dataset["data"], num_components, **dataset["parameters"].observation_limits
+    )
+
+    with monkeypatch.context() as m:
+        m.setattr("multiprocessing.Pool", multiprocessing.dummy.Pool)
+
+        bootstrap = fit.calculate_bootstrap(iterations=3, num_processes=2)
+        assert bootstrap.n_samples == 3
+
+
 @pytest.mark.slow
 def test_bootstrap(exponential_data):
     # double exponential data