FEAT: Improve interactions of user pools

bilby/core/result.py

@@ -4,10 +4,10 @@
 import os
 from collections import namedtuple
 from copy import copy
+from functools import partial
 from importlib import import_module
 from itertools import product
-import multiprocessing
-from functools import partial
+
 import numpy as np
 import pandas as pd
 import scipy.stats
@@ -192,7 +192,7 @@ def read_in_result_list(filename_list, invalid="warning"):
 
 def get_weights_for_reweighting(
         result, new_likelihood=None, new_prior=None, old_likelihood=None,
-        old_prior=None, resume_file=None, n_checkpoint=5000, npool=1):
+        old_prior=None, resume_file=None, n_checkpoint=5000, npool=1, pool=None):
     """ Calculate the weights for reweight()
 
     See bilby.core.result.reweight() for help with the inputs
@@ -235,20 +235,27 @@ def get_weights_for_reweighting(
         basedir = os.path.split(resume_file)[0]
         check_directory_exists_and_if_not_mkdir(basedir)
 
-    dict_samples = [{key: sample[key] for key in result.posterior}
-                    for _, sample in result.posterior.iterrows()]
+    dict_samples = result.posterior.to_dict(orient="records")
     n = len(dict_samples) - starting_index
 
     # Helper function to compute likelihoods in parallel
     def eval_pool(this_logl):
-        with multiprocessing.Pool(processes=npool) as pool:
-            chunksize = max(100, n // (2 * npool))
-            return list(tqdm(
-                pool.imap(partial(_safe_likelihood_call, this_logl),
-                        dict_samples[starting_index:], chunksize=chunksize),
+        from .utils.parallel import bilby_pool
+
+        with bilby_pool(likelihood=this_logl, npool=npool) as my_pool:
+            if my_pool is None:
+                map_fn = map
+            else:
+                chunksize = max(100, n // (2 * npool))
+                map_fn = partial(my_pool.imap, chunksize=chunksize)
+            likelihood_fn = partial(_safe_likelihood_call, this_logl)
+
+            log_l = list(tqdm(
+                map_fn(likelihood_fn, dict_samples[starting_index:]),
                 desc='Computing likelihoods',
-                total=n)
-            )
+                total=n,
+            ))
+        return log_l
 
     if old_likelihood is None:
         old_log_likelihood_array[starting_index:] = \
@@ -319,7 +326,7 @@ def rejection_sample(posterior, weights):
 def reweight(result, label=None, new_likelihood=None, new_prior=None,
              old_likelihood=None, old_prior=None, conversion_function=None, npool=1,
              verbose_output=False, resume_file=None, n_checkpoint=5000,
-             use_nested_samples=False):
+             use_nested_samples=False, pool=None):
     """ Reweight a result to a new likelihood/prior using rejection sampling
 
     Parameters
@@ -382,7 +389,9 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None,
         get_weights_for_reweighting(
             result, new_likelihood=new_likelihood, new_prior=new_prior,
             old_likelihood=old_likelihood, old_prior=old_prior,
-            resume_file=resume_file, n_checkpoint=n_checkpoint, npool=npool)
+            resume_file=resume_file, n_checkpoint=n_checkpoint,
+            npool=npool, pool=pool,
+        )
 
     if use_nested_samples:
         ln_weights += np.log(result.posterior["weights"])
@@ -409,10 +418,14 @@ def reweight(result, label=None, new_likelihood=None, new_prior=None,
 
     if conversion_function is not None:
         data_frame = result.posterior
-        if "npool" in inspect.signature(conversion_function).parameters:
-            data_frame = conversion_function(data_frame, new_likelihood, new_prior, npool=npool)
-        else:
-            data_frame = conversion_function(data_frame, new_likelihood, new_prior)
+        parameters = inspect.signature(conversion_function).parameters
+        kwargs = dict()
+        for key, value in [
+            ("likelihood", new_likelihood), ("priors", new_prior), ("npool", npool), ("pool", pool)
+        ]:
+            if key in parameters:
+                kwargs[key] = value
+        data_frame = conversion_function(data_frame, **kwargs)
         result.posterior = data_frame
 
     if label:
@@ -765,6 +778,21 @@ def log_10_evidence_err(self):
     def log_10_noise_evidence(self):
         return self.log_noise_evidence / np.log(10)
 
+    @property
+    def sampler_kwargs(self):
+        return self._sampler_kwargs
+
+    @sampler_kwargs.setter
+    def sampler_kwargs(self, sampler_kwargs):
+        if sampler_kwargs is None:
+            sampler_kwargs = dict()
+        else:
+            sampler_kwargs = copy(sampler_kwargs)
+        if "pool" in sampler_kwargs:
+            # pool objects can't be neatly serialized
+            sampler_kwargs["pool"] = None
+        self._sampler_kwargs = sampler_kwargs
+
     @property
     def version(self):
         return self._version
@@ -1530,7 +1558,7 @@ def _add_prior_fixed_values_to_posterior(posterior, priors):
         return posterior
 
     def samples_to_posterior(self, likelihood=None, priors=None,
-                             conversion_function=None, npool=1):
+                             conversion_function=None, npool=1, pool=None):
         """
         Convert array of samples to posterior (a Pandas data frame)
 
@@ -1560,10 +1588,14 @@ def samples_to_posterior(self, likelihood=None, priors=None,
             data_frame['log_prior'] = self.log_prior_evaluations
 
         if conversion_function is not None:
-            if "npool" in inspect.signature(conversion_function).parameters:
-                data_frame = conversion_function(data_frame, likelihood, priors, npool=npool)
-            else:
-                data_frame = conversion_function(data_frame, likelihood, priors)
+            parameters = inspect.signature(conversion_function).parameters
+            kwargs = dict()
+            for key, value in [
+                ("likelihood", likelihood), ("priors", priors), ("npool", npool), ("pool", pool)
+            ]:
+                if key in parameters:
+                    kwargs[key] = value
+            data_frame = conversion_function(data_frame, **kwargs)
         self.posterior = data_frame
 
     def calculate_prior_values(self, priors):

bilby/core/sampler/__init__.py

@@ -11,6 +11,7 @@
     loaded_modules_dict,
     logger,
 )
+from ..utils.parallel import bilby_pool
 from . import proposal
 from .base_sampler import Sampler, SamplingMarginalisedParameterError
 
@@ -158,6 +159,7 @@ def run_sampler(
     gzip=False,
     result_class=None,
     npool=1,
+    pool=None,
     **kwargs,
 ):
     """
@@ -266,36 +268,27 @@ def run_sampler(
 
         likelihood = ZeroLikelihood(likelihood)
 
+    common_kwargs = dict(
+        likelihood=likelihood,
+        priors=priors,
+        outdir=outdir,
+        label=label,
+        injection_parameters=injection_parameters,
+        meta_data=meta_data,
+        use_ratio=use_ratio,
+        plot=plot,
+        result_class=result_class,
+        npool=npool,
+        pool=pool,
+    )
+
     if isinstance(sampler, Sampler):
         pass
     elif isinstance(sampler, str):
         sampler_class = get_sampler_class(sampler)
-        sampler = sampler_class(
-            likelihood,
-            priors=priors,
-            outdir=outdir,
-            label=label,
-            injection_parameters=injection_parameters,
-            meta_data=meta_data,
-            use_ratio=use_ratio,
-            plot=plot,
-            result_class=result_class,
-            npool=npool,
-            **kwargs,
-        )
+        sampler = sampler_class(**common_kwargs, **kwargs)
     elif inspect.isclass(sampler):
-        sampler = sampler.__init__(
-            likelihood,
-            priors=priors,
-            outdir=outdir,
-            label=label,
-            use_ratio=use_ratio,
-            plot=plot,
-            injection_parameters=injection_parameters,
-            meta_data=meta_data,
-            npool=npool,
-            **kwargs,
-        )
+        sampler = sampler.__init__(**common_kwargs, **kwargs)
     else:
         raise ValueError(
             "Provided sampler should be a Sampler object or name of a known "
@@ -305,42 +298,81 @@ def run_sampler(
     if sampler.cached_result:
         logger.warning("Using cached result")
         result = sampler.cached_result
+        result = apply_conversion_function(
+            result=result,
+            likelihood=likelihood,
+            conversion_function=conversion_function,
+            npool=npool,
+            pool=pool,
+        )
     else:
         # Run the sampler
-        start_time = datetime.datetime.now()
-        if command_line_args.bilby_test_mode:
-            result = sampler._run_test()
-        else:
-            result = sampler.run_sampler()
-        end_time = datetime.datetime.now()
-
-        # Some samplers calculate the sampling time internally
-        if result.sampling_time is None:
-            result.sampling_time = end_time - start_time
-        elif isinstance(result.sampling_time, (float, int)):
-            result.sampling_time = datetime.timedelta(result.sampling_time)
-
-        logger.info(f"Sampling time: {result.sampling_time}")
-        # Convert sampling time into seconds
-        result.sampling_time = result.sampling_time.total_seconds()
-
-        if sampler.use_ratio:
-            result.log_noise_evidence = likelihood.noise_log_likelihood()
-            result.log_bayes_factor = result.log_evidence
-            result.log_evidence = result.log_bayes_factor + result.log_noise_evidence
-        else:
-            result.log_noise_evidence = likelihood.noise_log_likelihood()
-            result.log_bayes_factor = result.log_evidence - result.log_noise_evidence
+        with bilby_pool(
+            likelihood,
+            priors,
+            use_ratio=sampler.use_ratio,
+            search_parameter_keys=sampler.search_parameter_keys,
+            npool=npool,
+            pool=pool,
+            parameters=priors.sample(),
+        ) as _pool:
+            start_time = datetime.datetime.now()
+            sampler.pool = _pool
+            if command_line_args.bilby_test_mode:
+                result = sampler._run_test()
+            else:
+                result = sampler.run_sampler()
+            end_time = datetime.datetime.now()
+            result = finalize_result(
+                result=result,
+                likelihood=likelihood,
+                use_ratio=sampler.use_ratio,
+                start_time=start_time,
+                end_time=end_time,
+            )
 
-        if None not in [result.injection_parameters, conversion_function]:
-            result.injection_parameters = conversion_function(
-                result.injection_parameters
+            # Initial save of the sampler in case of failure in samples_to_posterior
+            if save:
+                result.save_to_file(extension=save, gzip=gzip, outdir=outdir)
+
+            result = apply_conversion_function(
+                result=result,
+                likelihood=likelihood,
+                conversion_function=conversion_function,
+                npool=npool,
+                pool=_pool,
             )
 
-        # Initial save of the sampler in case of failure in samples_to_posterior
-        if save:
-            result.save_to_file(extension=save, gzip=gzip, outdir=outdir)
+    if save:
+        # The overwrite here ensures we overwrite the initially stored data
+        result.save_to_file(overwrite=True, extension=save, gzip=gzip, outdir=outdir)
+
+    if plot:
+        result.plot_corner()
+    logger.info(f"Summary of results:\n{result}")
+    return result
+
 
+def apply_conversion_function(
+    result, likelihood, conversion_function, npool=None, pool=None
+):
+    """
+    Apply the conversion function to the injected parameters and posterior if the
+    posterior has not already been created from the stored samples.
+
+    Parameters
+    ----------
+    result : bilby.core.result.Result
+        The result object from the sampler.
+    likelihood : bilby.Likelihood
+        The likelihood used during sampling.
+    conversion_function : function
+        The conversion function to apply.
+    npool : int, optional
+        The number of processes to use in a processing pool.
+    pool : multiprocessing.Pool, schwimmbad.MPIPool, optional
+        The pool to use for parallelisation, this overrides the :code:`npool` argument.
+    """
     if None not in [result.injection_parameters, conversion_function]:
         result.injection_parameters = conversion_function(
             result.injection_parameters,
@@ -354,15 +386,30 @@ def run_sampler(
             priors=result.priors,
             conversion_function=conversion_function,
             npool=npool,
+            pool=pool,
         )
+    return result
 
-    if save:
-        # The overwrite here ensures we overwrite the initially stored data
-        result.save_to_file(overwrite=True, extension=save, gzip=gzip, outdir=outdir)
 
-    if plot:
-        result.plot_corner()
-    logger.info(f"Summary of results:\n{result}")
+def finalize_result(result, likelihood, use_ratio, start_time=None, end_time=None):
+    # Some samplers calculate the sampling time internally
+    if result.sampling_time is None and None not in [start_time, end_time]:
+        result.sampling_time = end_time - start_time
+    elif isinstance(result.sampling_time, (float, int)):
+        result.sampling_time = datetime.timedelta(result.sampling_time)
+
+    logger.info(f"Sampling time: {result.sampling_time}")
+    # Convert sampling time into seconds
+    result.sampling_time = result.sampling_time.total_seconds()
+
+    if use_ratio:
+        result.log_noise_evidence = likelihood.noise_log_likelihood()
+        result.log_bayes_factor = result.log_evidence
+        result.log_evidence = result.log_bayes_factor + result.log_noise_evidence
+    else:
+        result.log_noise_evidence = likelihood.noise_log_likelihood()
+        result.log_bayes_factor = result.log_evidence - result.log_noise_evidence
+
     return result