Merge pull request #1700 from pints-team/1691-pints-logpdf-storage

LogPDF storage issue with multi-chain methods

Author: martinjrobins

pints/_mcmc/__init__.py

@@ -735,6 +735,8 @@ def run(self):
 
                     if reply is not None:
                         # Unpack reply into position, evaluation, and status
+                        # Note that evaluation is (fy, dfy) with sensitivities
+                        # enabled.
                         y, fy, accepted = reply
 
                         # Inverse transform to model space if transformation is
@@ -754,7 +756,10 @@ def run(self):
                         if store_evaluations:
                             # If accepted, update log_pdf and prior for logging
                             if accepted:
-                                current_logpdf[i] = fy
+                                if self._needs_sensitivities:
+                                    current_logpdf[i] = fy[0]
+                                else:
+                                    current_logpdf[i] = fy
                                 if prior is not None:
                                     current_prior[i] = prior(y)
 
@@ -818,7 +823,10 @@ def run(self):
                             # Check if accepted, if so, update log_pdf and
                             # prior to be logged
                             if accepted[i]:
-                                current_logpdf[i] = fys[i]
+                                if self._needs_sensitivities:
+                                    current_logpdf[i] = fys[i][0]
+                                else:
+                                    current_logpdf[i] = fys[i]
                                 if prior is not None:
                                     current_prior[i] = prior(ys[i])
 
@@ -1034,9 +1042,10 @@ def set_log_pdf_storage(self, store_in_memory=False):
         Store :class:`LogPDF` evaluations in memory as they are generated.
 
         By default, evaluations of the :class:`LogPDF` are not stored. This
-        method can be used to enable storage of the evaluations for the
-        accepted samples.
-        After running, evaluations can be obtained using :meth:`evaluations()`.
+        method can be used to enable (or disable) storage for the accepted
+        samples, in memory.
+
+        After running, evaluations can be obtained using :meth:`log_pdfs()`.
         """
         self._evaluations_in_memory = bool(store_in_memory)
 

pints/_mcmc/_adaptive_covariance.py

@@ -23,7 +23,7 @@ class AdaptiveCovarianceMC(pints.SingleChainMCMC):
     A basic implementation is provided for each, which extending methods can
     choose to override.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, does not use sensitivities.
     """
 
     def __init__(self, x0, sigma0=None):

pints/_mcmc/_differential_evolution.py

@@ -29,7 +29,7 @@ class DifferentialEvolutionMCMC(pints.MultiChainMCMC):
     If ``x_proposed / x[i,r] > u ~ U(0,1)``, then
     ``x[i+1,r] = x_proposed``; otherwise, ``x[i+1,r] = x[i]``.
 
-    Extends :class:`MultiChainMCMC`.
+    Extends :class:`MultiChainMCMC`, does not use sensitivities.
 
     .. note::
         This sampler requires a number of chains :math:`n \ge 3`, and

pints/_mcmc/_dram_ac.py

@@ -50,7 +50,7 @@ class DramACMC(pints.AdaptiveCovarianceMC):
     are then adapted as ``[scale_1, scale_2] * sigma``, where the
     scale factors are set using ``set_sigma_scale``.
 
-    *Extends:* :class:`GlobalAdaptiveCovarianceMC`
+    *Extends:* :class:`GlobalAdaptiveCovarianceMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_dream.py

@@ -53,7 +53,7 @@ class DreamMCMC(pints.MultiChainMCMC):
 
     Here b > 0, b* > 0,  1 >= p_g >= 0, 1 >= CR >= 0.
 
-    Extends :class:`MultiChainMCMC`.
+    Extends :class:`MultiChainMCMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_emcee_hammer.py

@@ -34,6 +34,8 @@ class EmceeHammerMCMC(pints.MultiChainMCMC):
     ``1 / sqrt(z)`` if ``z`` is in  ``[1 / a, a]`` or to 0, otherwise (where
     ``a`` is a parameter with default value ``2``).
 
+    Extends :class:`MultiChainMCMC`, does not use sensitivities.
+
     References
     ----------
     .. [1] "emcee: The MCMC Hammer", Daniel Foreman-Mackey, David W. Hogg,

pints/_mcmc/_haario_ac.py

@@ -44,7 +44,7 @@ class HaarioACMC(pints.AdaptiveCovarianceMC):
         log lambda = log lambda + gamma (alpha - self._target_acceptance)
         gamma = adaptation_count^-eta
 
-    Extends :class:`AdaptiveCovarianceMC`.
+    Extends :class:`AdaptiveCovarianceMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_haario_bardenet_ac.py

@@ -42,7 +42,7 @@ class HaarioBardenetACMC(pints.AdaptiveCovarianceMC):
         log lambda = log lambda + gamma (alpha - self._target_acceptance)
         gamma = adaptation_count^-eta
 
-    Extends :class:`AdaptiveCovarianceMC`.
+    Extends :class:`AdaptiveCovarianceMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_hamiltonian.py

@@ -46,7 +46,7 @@ class HamiltonianMCMC(pints.SingleChainMCMC):
     In particular, the algorithm we implement follows eqs. (4.14)-(4.16) in
     [1]_, since we allow different epsilon according to dimension.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, this method uses sensitivities.
 
     References
     ----------
@@ -190,7 +190,10 @@ def name(self):
         return 'Hamiltonian Monte Carlo'
 
     def needs_sensitivities(self):
-        """ See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
+        """
+        This method requires sensitivities.
+
+        See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
         return True
 
     def scaled_epsilon(self):

pints/_mcmc/_mala.py

@@ -67,7 +67,7 @@ class MALAMCMC(pints.SingleChainMCMC):
     where :math:`\epsilon' = \epsilon \sqrt{M}` is given by the initial value
     of ``sigma0``.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, this method uses sensitivities.
 
     References
     ----------
@@ -181,7 +181,11 @@ def name(self):
         return 'Metropolis-Adjusted Langevin Algorithm (MALA)'
 
     def needs_sensitivities(self):
-        """ See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
+        """
+        This method requires sensitivities.
+
+        See :meth:`pints.MCMCSampler.needs_sensitivities()`.
+        """
         return True
 
     def n_hyper_parameters(self):

pints/_mcmc/_metropolis.py

@@ -24,7 +24,7 @@ class MetropolisRandomWalkMCMC(pints.SingleChainMCMC):
 
     here Sigma is the covariance matrix of the proposal.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_monomial_gamma_hamiltonian.py

@@ -63,7 +63,7 @@ class MonomialGammaHamiltonianMCMC(pints.SingleChainMCMC):
     In particular, the algorithm we implement follows eqs. (4.14)-(4.16) in
     [2]_, since we allow different epsilon according to dimension.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, this method uses sensitivities.
 
     References
     ----------
@@ -297,7 +297,11 @@ def name(self):
         return 'Monomial-Gamma Hamiltonian Monte Carlo'
 
     def needs_sensitivities(self):
-        """ See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
+        """
+        This method requires sensitivities.
+
+        See also :meth:`pints.MCMCSampler.needs_sensitivities()`.
+        """
         return True
 
     def n_hyper_parameters(self):

pints/_mcmc/_nuts.py

@@ -499,7 +499,7 @@ class NoUTurnMCMC(pints.SingleChainMCMC):
 
     Note: This sampler is only supported on Python versions 3.3 and newer.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, this method uses sensitivities.
 
     References
     ----------
@@ -664,7 +664,11 @@ def name(self):
         return 'No-U-Turn MCMC'
 
     def needs_sensitivities(self):
-        """ See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
+        """
+        This method requires sensitivities.
+
+        See :meth:`pints.MCMCSampler.needs_sensitivities()`.
+        """
         return True
 
     def number_adaption_steps(self):

pints/_mcmc/_population.py

@@ -39,7 +39,7 @@ class PopulationMCMC(pints.SingleChainMCMC):
     ``delta_T = 1 / num_temperatures``, and the chain with ``T_i = 0`` is the
     one whose target distribution we want to sample.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_rao_blackwell_ac.py

@@ -37,7 +37,7 @@ class RaoBlackwellACMC(pints.AdaptiveCovarianceMC):
         Y_t+1 ~ N(theta_t, lambda * sigma0) rather than
             Y_t+1 ~ N(theta_t, sigma0)
 
-    Extends :class:`AdaptiveCovarianceMC`.
+    Extends :class:`AdaptiveCovarianceMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_relativistic.py

@@ -57,7 +57,7 @@ class RelativisticMCMC(pints.SingleChainMCMC):
     In particular, the algorithm we implement follows eqs. in section 2.1 of
     [1]_.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, this method uses sensitivities.
 
     References
     ----------
@@ -325,7 +325,11 @@ def name(self):
         return 'Relativistic MCMC'
 
     def needs_sensitivities(self):
-        """ See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
+        """
+        This method requires sensitivities.
+
+        See :meth:`pints.MCMCSampler.needs_sensitivities()`.
+        """
         return True
 
     def _sample_momentum(self):

pints/_mcmc/_slice_doubling.py

@@ -109,7 +109,7 @@ class SliceDoublingMCMC(pints.SingleChainMCMC):
     :math:`\epsilon \sim \text{exp}(1)` and define the slice as
     :math:`S = {x : z < g(x)}`.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, does not use sensitivities.
 
     References
     ----------

pints/_mcmc/_slice_rank_shrinking.py

@@ -63,7 +63,7 @@ class SliceRankShrinkingMCMC(pints.SingleChainMCMC):
     :math:`\epsilon \sim \text{exp}(1)` and define the slice as
     :math:`S = {x : z < log f(x)}`.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, this method uses sensitivities.
 
     References
     ----------
@@ -156,7 +156,11 @@ def name(self):
         return 'Slice Sampling - Covariance-Adaptive: Rank Shrinking.'
 
     def needs_sensitivities(self):
-        """ See :meth:`pints.MCMCSampler.needs_sensitivities()`. """
+        """
+        This method requires sensitivities.
+
+        See :meth:`pints.MCMCSampler.needs_sensitivities()`.
+        """
         return True
 
     def n_hyper_parameters(self):

pints/_mcmc/_slice_stepout.py

@@ -120,7 +120,7 @@ class SliceStepoutMCMC(pints.SingleChainMCMC):
     :math:`\epsilon \sim \text{exp}(1)` and define the slice as
     :math:`S = {x : z < g(x)}`.
 
-    Extends :class:`SingleChainMCMC`.
+    Extends :class:`SingleChainMCMC`, does not use sensitivities.
 
     References
     ----------

pints/tests/test_mcmc_controller.py

@@ -690,6 +690,26 @@ def test_log_pdf_storage_in_memory_single(self):
             likelihoods = [self.log_likelihood(x) for x in chain]
             self.assertTrue(np.all(evals[i] == likelihoods))
 
+        # Test with a sensitivity-using method
+        mcmc = pints.MCMCController(
+            self.log_posterior, n_chains, xs, method=pints.MALAMCMC)
+        mcmc.set_max_iterations(n_iterations)
+        mcmc.set_log_to_screen(False)
+        mcmc.set_log_pdf_storage(True)
+        chains = mcmc.run()
+        evals = mcmc.log_pdfs()
+        self.assertEqual(len(evals.shape), 3)
+        self.assertEqual(evals.shape[0], n_chains)
+        self.assertEqual(evals.shape[1], n_iterations)
+        self.assertEqual(evals.shape[2], 3)
+        for i, chain in enumerate(chains):
+            posteriors = [self.log_posterior(x) for x in chain]
+            self.assertTrue(np.all(evals[i, :, 0] == posteriors))
+            likelihoods = [self.log_likelihood(x) for x in chain]
+            self.assertTrue(np.all(evals[i, :, 1] == likelihoods))
+            priors = [self.log_prior(x) for x in chain]
+            self.assertTrue(np.all(evals[i, :, 2] == priors))
+
         # Test disabling again
         mcmc = pints.MCMCController(self.log_posterior, n_chains, xs)
         mcmc.set_max_iterations(n_iterations)
@@ -706,7 +726,9 @@ def test_log_pdf_storage_in_memory_multi(self):
         x0 = np.array(self.real_parameters) * 1.05
         x1 = np.array(self.real_parameters) * 1.15
         x2 = np.array(self.real_parameters) * 0.95
-        xs = [x0, x1, x2]
+        x3 = np.array(self.real_parameters) * 0.95
+        x4 = np.array(self.real_parameters) * 0.95
+        xs = [x0, x1, x2, x3, x4]
         n_chains = len(xs)
         n_iterations = 100
         meth = pints.DifferentialEvolutionMCMC
@@ -737,6 +759,28 @@ def test_log_pdf_storage_in_memory_multi(self):
             priors = [self.log_prior(x) for x in chain]
             self.assertTrue(np.all(evals[i, :, 2] == priors))
 
+        # Test with a sensitivity-using multi-chain method
+        # We don't have any of these!
+        mcmc = pints.MCMCController(
+            self.log_posterior, n_chains, xs,
+            method=FakeMultiChainSamplerWithSensitivities)
+        mcmc.set_max_iterations(n_iterations)
+        mcmc.set_log_to_screen(False)
+        mcmc.set_log_pdf_storage(True)
+        chains = mcmc.run()
+        evals = mcmc.log_pdfs()
+        self.assertEqual(len(evals.shape), 3)
+        self.assertEqual(evals.shape[0], n_chains)
+        self.assertEqual(evals.shape[1], n_iterations)
+        self.assertEqual(evals.shape[2], 3)
+        for i, chain in enumerate(chains):
+            posteriors = [self.log_posterior(x) for x in chain]
+            self.assertTrue(np.all(evals[i, :, 0] == posteriors))
+            likelihoods = [self.log_likelihood(x) for x in chain]
+            self.assertTrue(np.all(evals[i, :, 1] == likelihoods))
+            priors = [self.log_prior(x) for x in chain]
+            self.assertTrue(np.all(evals[i, :, 2] == priors))
+
         # Test with a loglikelihood
         mcmc = pints.MCMCController(
             self.log_likelihood, n_chains, xs, method=meth)
@@ -1651,6 +1695,29 @@ def tell(self, fx):
         return None if x is None else (x, fx, np.array([True] * self._n))
 
 
+class FakeMultiChainSamplerWithSensitivities(pints.MultiChainMCMC):
+    """
+    Fake sampler that pretends to be a multi-chain method that uses
+    sensitivities. At the moment (2024-12-18) we don't have these in PINTS, but
+    we need to check that code potentially handling this type of sampler exists
+    or raises not-implemented errors.
+    """
+    def ask(self):
+        self._xs = self._x0 + 1e-3 * np.random.uniform(
+            size=(self._n_chains, self._n_parameters))
+        return self._xs
+
+    def current_log_pdfs(self):
+        return self._fxs
+
+    def tell(self, fxs):
+        self._fxs = fxs
+        return self._xs, self._fxs, [True] * self._n_chains
+
+    def needs_sensitivities(self):
+        return True
+
+
 class TestMCMCControllerSingleChainStorage(unittest.TestCase):
     """
     Tests storage of samples and evaluations to disk, running with a