Repository: pints-team/pints

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 !

Comments:

Patch:
diff --git a/True b/True
new file mode 100644
index 00000000..844f2a6b
--- /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 d02d0ad3..0e089b08 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 51b3e350..dbe8f045 100644
--- a/pints/__init__.py
+++ b/pints/__init__.py
@@ -156,6 +156,7 @@ from ._evaluation import (
     Evaluator,
     ParallelEvaluator,
     SequentialEvaluator,
+    MultiSequentialEvaluator,
 )
 
 
diff --git a/pints/_evaluation.py b/pints/_evaluation.py
index 51d267ee..cb194e20 100644
--- a/pints/_evaluation.py
+++ b/pints/_evaluation.py
@@ -413,6 +413,41 @@ multiprocessing.html#all-platforms>`_ for details).
         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 @@ class _Worker(multiprocessing.Process):
         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 62071cab..f6514036 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 @@ class MCMCController(object):
             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 @@ class MCMCController(object):
             # 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 @@ class MCMCController(object):
         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 @@ class MCMCController(object):
         # 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 @@ class MCMCController(object):
         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 670823c9..b3fc7895 100755
--- a/pints/tests/test_evaluators.py
+++ b/pints/tests/test_evaluators.py
@@ -53,6 +53,43 @@ class TestEvaluators(unittest.TestCase):
         # 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 ceba9f04..fb03de62 100755
--- a/pints/tests/test_mcmc_controller.py
+++ b/pints/tests/test_mcmc_controller.py
@@ -11,7 +11,9 @@ import pints
 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 @@ class TestMCMCController(unittest.TestCase):
         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 @@ class TestMCMCController(unittest.TestCase):
         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 @@ class TestMCMCController(unittest.TestCase):
             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.