Repository: pints-team/pints (language: python)

Issue:
Optionally copy logpdfs before running mcmc
Currently, the MCMC controller calls `__call__` or `evaluateS1` on the same `log_pdf` object for all mcmc chains.

This causes problems for more complex LogPDF objects which have attributes which may be updated based on the parameters sent to the call method (in my case, the LogPDF object holds a grid of ODE solver time points which is updated adaptively as the chain moves through parameter space).

It's easy enough to fix by making one copy.deepcopy of the log_pdf for each chain, and adding an appropriate pints.Evaluator so each chain is just calling its own log_pdf object. @MichaelClerx let me know if this makes any sense as an optional argument to the MCMC controller !


Discussion / hints:
Hey @rcw5890 , yeah I can see why that would be useful.

Are there limitations to copy.deepcopy? I'm wondering if it would be better to let the user pass in a list of `log_pdf` objects of the same length as the number of chains?
That might help if, for example, you need to pass some special info to each log_pdf copy upon creation (i.e. a CPU to run on or something like that?)
Agreed, I'll implement it that way !

Fix patch:
diff --git a/True b/True
new file mode 100644
index 000000000..844f2a6b2
--- /dev/null
+++ b/True
@@ -0,0 +1,4 @@
+Iter. Eval. Acceptance rate Time m:s
+1     198    0.00505050505    0:00.1
+2     1029   0.0019436346     0:00.3
+3     3211   0.000934288384   0:01.0
diff --git a/docs/source/function_evaluation.rst b/docs/source/function_evaluation.rst
index d02d0ad32..0e089b085 100644
--- a/docs/source/function_evaluation.rst
+++ b/docs/source/function_evaluation.rst
@@ -24,6 +24,7 @@ Overview:
 - :class:`Evaluator`
 - :class:`ParallelEvaluator`
 - :class:`SequentialEvaluator`
+- :class:`MultiSequentialEvaluator`
 
 
 .. autofunction:: evaluate
@@ -34,3 +35,4 @@ Overview:
 
 .. autoclass:: SequentialEvaluator
 
+.. autoclass:: MultiSequentialEvaluator
diff --git a/pints/__init__.py b/pints/__init__.py
index 51b3e3506..dbe8f0450 100644
--- a/pints/__init__.py
+++ b/pints/__init__.py
@@ -156,6 +156,7 @@ def version(formatted=False):
     Evaluator,
     ParallelEvaluator,
     SequentialEvaluator,
+    MultiSequentialEvaluator,
 )
 
 
diff --git a/pints/_evaluation.py b/pints/_evaluation.py
index 51d267eea..cb194e20f 100644
--- a/pints/_evaluation.py
+++ b/pints/_evaluation.py
@@ -413,6 +413,41 @@ def clear(q):
         return errors
 
 
+class MultiSequentialEvaluator(Evaluator):
+    """
+    Evaluates a list of functions (or callable objects) for a list of input
+    values of the same length, and returns a list containing the calculated
+    function evaluations.
+
+    Extends :class:`Evaluator`.
+
+    Parameters
+    ----------
+    functions : list of callable
+        The functions to evaluate.
+    args : sequence
+        An optional tuple containing extra arguments to each element in
+        functions, ``f``. If ``args`` is specified, ``f`` will be called as
+        ``f(x, *args)``.
+    """
+    def __init__(self, functions, args=None):
+        super(MultiSequentialEvaluator, self).__init__(functions[0], args)
+
+        # Check functions
+        for function in functions:
+            if not callable(function):
+                raise ValueError('The given functions must be callable.')
+        self._functions = functions
+        self._n_functions = len(functions)
+
+    def _evaluate(self, positions):
+        if len(positions) != self._n_functions:
+            raise ValueError('Number of positions does not equal number of '
+                             'functions.')
+
+        return [f(x, *self._args) for f, x in zip(self._functions, positions)]
+
+
 class SequentialEvaluator(Evaluator):
     """
     Evaluates a function (or callable object) for a list of input values, and
@@ -522,4 +557,3 @@ def run(self):
         except (Exception, KeyboardInterrupt, SystemExit):
             self._errors.put((self.pid, traceback.format_exc()))
             self._error.set()
-
diff --git a/pints/_mcmc/__init__.py b/pints/_mcmc/__init__.py
index 62071cab1..f65140363 100644
--- a/pints/_mcmc/__init__.py
+++ b/pints/_mcmc/__init__.py
@@ -274,7 +274,10 @@ class MCMCController(object):
     ----------
     log_pdf : pints.LogPDF
         A :class:`LogPDF` function that evaluates points in the parameter
-        space.
+        space, or a list of :class:`LogPDF` of the same length as `chains`. If
+        multiple LogPDFs are provided, each chain will call only its
+        corresponding LogPDF. Note that if multiple LogPDFs are provided,
+        parallel running is not possible.
     chains : int
         The number of MCMC chains to generate.
     x0
@@ -301,16 +304,42 @@ def __init__(
             self, log_pdf, chains, x0, sigma0=None, transformation=None,
             method=None):
 
-        # Check function
-        if not isinstance(log_pdf, pints.LogPDF):
-            raise ValueError('Given function must extend pints.LogPDF')
+        if isinstance(log_pdf, pints.LogPDF):
+            self._multi_logpdf = False
+
+        else:
+            self._multi_logpdf = True
+            try:
+                if len(log_pdf) != chains:
+                    raise ValueError(
+                        '`log_pdf` must either extend pints.LogPDF, '
+                        'or be a list of objects which extend '
+                        'pints.LogPDF of the same length as `chains`')
+            except TypeError:
+                raise TypeError('`log_pdf` must either extend pints.LogPDF, '
+                                'or be a list of objects which extend '
+                                'pints.LogPDF')
+
+            first_n_params = log_pdf[0].n_parameters()
+            for pdf in log_pdf:
+                # Check function
+                if not isinstance(pdf, pints.LogPDF):
+                    raise ValueError('Elements of `log_pdf` must extend '
+                                     'pints.LogPDF')
+                if pdf.n_parameters() != first_n_params:
+                    raise ValueError('All log_pdfs must have the same number '
+                                     'of parameters.')
 
         # Apply a transformation (if given). From this point onward the MCMC
         # sampler will see only the transformed search space and will know
         # nothing about the model parameter space.
         if transformation is not None:
             # Convert log pdf
-            log_pdf = transformation.convert_log_pdf(log_pdf)
+            if self._multi_logpdf:
+                log_pdf = [transformation.convert_log_pdf(pdf)
+                           for pdf in log_pdf]
+            else:
+                log_pdf = transformation.convert_log_pdf(log_pdf)
 
             # Convert initial positions
             x0 = [transformation.to_search(x) for x in x0]
@@ -318,7 +347,10 @@ def __init__(
             # Convert sigma0, if provided
             if sigma0 is not None:
                 sigma0 = np.asarray(sigma0)
-                n_parameters = log_pdf.n_parameters()
+                if not self._multi_logpdf:
+                    n_parameters = log_pdf.n_parameters()
+                else:
+                    n_parameters = log_pdf[0].n_parameters()
                 # Make sure sigma0 is a (covariance) matrix
                 if np.product(sigma0.shape) == n_parameters:
                     # Convert from 1d array
@@ -341,7 +373,10 @@ def __init__(
         self._log_pdf = log_pdf
 
         # Get number of parameters
-        self._n_parameters = self._log_pdf.n_parameters()
+        if not self._multi_logpdf:
+            self._n_parameters = self._log_pdf.n_parameters()
+        else:
+            self._n_parameters = self._log_pdf[0].n_parameters()
 
         # Check number of chains
         self._n_chains = int(chains)
@@ -528,15 +563,24 @@ def run(self):
         # Choose method to evaluate
         f = self._log_pdf
         if self._needs_sensitivities:
-            f = f.evaluateS1
+            if not self._multi_logpdf:
+                f = f.evaluateS1
+            else:
+                f = [pdf.evaluateS1 for pdf in f]
 
         # Create evaluator object
         if self._parallel:
-            # Use at most n_workers workers
-            n_workers = min(self._n_workers, self._n_chains)
-            evaluator = pints.ParallelEvaluator(f, n_workers=n_workers)
+            if not self._multi_logpdf:
+                # Use at most n_workers workers
+                n_workers = min(self._n_workers, self._n_chains)
+                evaluator = pints.ParallelEvaluator(f, n_workers=n_workers)
+            else:
+                raise ValueError('Cannot run multiple logpdfs in parallel')
         else:
-            evaluator = pints.SequentialEvaluator(f)
+            if not self._multi_logpdf:
+                evaluator = pints.SequentialEvaluator(f)
+            else:
+                evaluator = pints.MultiSequentialEvaluator(f)
 
         # Initial phase
         if self._needs_initial_phase:
@@ -1041,6 +1085,9 @@ def set_parallel(self, parallel=False):
         than 0.
         Parallelisation can be disabled by setting ``parallel`` to ``0`` or
         ``False``.
+
+        Parallel evaluation is only supported when a single LogPDF has been
+        provided to the MCMC controller.
         """
         if parallel is True:
             self._parallel = True


Test patch:
diff --git a/pints/tests/test_evaluators.py b/pints/tests/test_evaluators.py
index 670823c94..b3fc7895e 100755
--- a/pints/tests/test_evaluators.py
+++ b/pints/tests/test_evaluators.py
@@ -53,6 +53,43 @@ def test_sequential(self):
         # Args must be a sequence
         self.assertRaises(ValueError, pints.SequentialEvaluator, f_args, 1)
 
+    def test_multi_sequential(self):
+
+        # Create test data
+        xs = np.random.normal(0, 10, 100)
+        ys = [f(x) for x in xs]
+
+        # Test sequential evaluator with multiple functions
+        e = pints.MultiSequentialEvaluator([f for _ in range(100)])
+        self.assertTrue(np.all(ys == e.evaluate(xs)))
+
+        # check errors
+
+        # not iterable
+        with self.assertRaises(TypeError):
+            e = pints.MultiSequentialEvaluator(3)
+
+        # not callable
+        with self.assertRaises(ValueError):
+            e = pints.MultiSequentialEvaluator([f, 4])
+
+        e = pints.MultiSequentialEvaluator([f for _ in range(100)])
+        # Argument must be sequence
+        with self.assertRaises(ValueError):
+            e.evaluate(1)
+
+        # wrong number of arguments
+        with self.assertRaises(ValueError):
+            e.evaluate([1 for _ in range(99)])
+
+        # Test args
+        e = pints.MultiSequentialEvaluator([f_args, f_args_plus1], [10, 20])
+        self.assertEqual(e.evaluate([1, 1]), [31, 32])
+
+        # Args must be a sequence
+        self.assertRaises(
+            ValueError, pints.MultiSequentialEvaluator, [f_args], 1)
+
     def test_parallel(self):
 
         # Create test data
@@ -212,6 +249,10 @@ def f_args(x, y, z):
     return x + y + z
 
 
+def f_args_plus1(x, y, z):
+    return x + y + z + 1
+
+
 def ioerror_on_five(x):
     if x == 5:
         raise IOError
diff --git a/pints/tests/test_mcmc_controller.py b/pints/tests/test_mcmc_controller.py
index ceba9f04b..fb03de62e 100755
--- a/pints/tests/test_mcmc_controller.py
+++ b/pints/tests/test_mcmc_controller.py
@@ -11,7 +11,9 @@
 import pints.io
 import pints.toy
 import unittest
+import unittest.mock
 import numpy as np
+import numpy.testing as npt
 
 from shared import StreamCapture, TemporaryDirectory
 
@@ -80,6 +82,17 @@ def setUpClass(cls):
         cls.log_posterior = pints.LogPosterior(
             cls.log_likelihood, cls.log_prior)
 
+        # Create another log-likelihood with two noise parameters
+        cls.log_likelihood_2 = pints.AR1LogLikelihood(problem)
+        cls.log_prior_2 = pints.UniformLogPrior(
+            [0.01, 400, 0.0, 0.0],
+            [0.02, 600, 100, 1]
+        )
+
+        # Create an un-normalised log-posterior (log-likelihood + log-prior)
+        cls.log_posterior_2 = pints.LogPosterior(
+            cls.log_likelihood_2, cls.log_prior_2)
+
     def test_single(self):
         # Test with a SingleChainMCMC method.
 
@@ -114,7 +127,7 @@ def test_single(self):
         def f(x):
             return x
         self.assertRaisesRegex(
-            ValueError, 'extend pints.LogPDF', pints.MCMCController,
+            TypeError, 'extend pints.LogPDF', pints.MCMCController,
             f, n_chains, xs)
 
         # Test x0 and chain argument
@@ -361,6 +374,100 @@ def test_multi(self):
             pints.MCMCController, self.log_posterior, n_chains, xs, sigma0,
             method=meth, transformation=logt)
 
+    def test_multi_logpdf(self):
+        # Test with multiple logpdfs
+
+        # 2 chains
+        x0 = np.array(self.real_parameters) * 1.1
+        x1 = np.array(self.real_parameters) * 1.15
+        xs = [x0, x1]
+
+        # Not iterable
+        with self.assertRaises(TypeError):
+            mcmc = pints.MCMCController(1, 3, xs)
+
+        # Wrong number of logpdfs
+        with self.assertRaises(ValueError):
+            mcmc = pints.MCMCController(
+                [self.log_posterior, self.log_posterior], 3, xs)
+
+        # List does not contain logpdfs
+        with self.assertRaises(ValueError):
+            mcmc = pints.MCMCController(
+                [self.log_posterior, 'abc'], 2, xs)
+
+        # Pdfs have different numbers of n_parameters
+        with self.assertRaises(ValueError):
+            mcmc = pints.MCMCController(
+                [self.log_posterior, self.log_posterior_2], 2, xs)
+
+        # Correctly configured inputs
+        n_chains = len(xs)
+        n_parameters = len(x0)
+        n_iterations = 10
+        mcmc = pints.MCMCController(
+            [self.log_posterior, self.log_posterior],
+            n_chains,
+            xs,
+            transformation=pints.LogTransformation(n_parameters),
+            sigma0=[1, 0.1, 0.01])
+        mcmc.set_max_iterations(n_iterations)
+        mcmc.set_log_to_screen(False)
+        chains = mcmc.run()
+        self.assertEqual(chains.shape[0], n_chains)
+        self.assertEqual(chains.shape[1], n_iterations)
+        self.assertEqual(chains.shape[2], n_parameters)
+        self.assertIs(chains, mcmc.chains())
+
+        # With sensitivities needed
+        mcmc = pints.MCMCController(
+            [self.log_posterior, self.log_posterior],
+            n_chains,
+            xs,
+            transformation=pints.LogTransformation(n_parameters),
+            sigma0=[1, 0.1, 0.01],
+            method=pints.HamiltonianMCMC)
+        mcmc.set_max_iterations(n_iterations)
+        mcmc.set_log_to_screen(False)
+        chains = mcmc.run()
+        self.assertEqual(chains.shape[0], n_chains)
+        self.assertEqual(chains.shape[1], n_iterations)
+        self.assertEqual(chains.shape[2], n_parameters)
+        self.assertIs(chains, mcmc.chains())
+
+        # Parallel (currently raises error)
+        mcmc = pints.MCMCController(
+            [self.log_posterior, self.log_posterior],
+            n_chains,
+            xs,
+            transformation=pints.LogTransformation(n_parameters),
+            sigma0=[1, 0.1, 0.01])
+        mcmc.set_parallel(True)
+        mcmc.set_max_iterations(n_iterations)
+        mcmc.set_log_to_screen(False)
+        with self.assertRaises(ValueError):
+            chains = mcmc.run()
+
+        # Test that both logpdfs are called
+        logpdf1 = unittest.mock.MagicMock(
+            return_value=-1.0, spec=self.log_posterior)
+        logpdf2 = unittest.mock.MagicMock(
+            return_value=-2.0, spec=self.log_posterior)
+        attrs = {'n_parameters.return_value': 3}
+        logpdf1.configure_mock(**attrs)
+        logpdf2.configure_mock(**attrs)
+        mcmc = pints.MCMCController([logpdf1, logpdf2], n_chains, xs)
+        mcmc.set_max_iterations(n_iterations)
+        mcmc.set_log_to_screen(False)
+        chains = mcmc.run()
+
+        logpdf1.assert_called()
+        logpdf2.assert_called()
+
+        # Check that they got called with the corresponding x0 at the start
+        npt.assert_allclose(logpdf1.call_args_list[0][0][0], xs[0])
+        npt.assert_allclose(logpdf2.call_args_list[0][0][0], xs[1])
+
     def test_stopping(self):
         # Test different stopping criteria.