Include proposals that jump along eigenvectors of covariance matrix of past samples.

bilby/bilby_mcmc/proposals.py

@@ -346,6 +346,84 @@ class UniformProposal(BaseProposal):
         return sample, log_factor
 
 
+class CovarianceAMProposal(BaseProposal):
+    def __init__(
+        self,
+        priors,
+        weight=1,
+        subset=None,
+        n_samples_for_covariance=10000,
+        min_samples=100,
+    ):
+        super(CovarianceAMProposal, self).__init__(priors, weight, subset)
+        """
+        Jump along eigenvectors of most recent samples. Based on proposals
+        in PTMCMCSampler: https://github.com/jellis18/PTMCMCSampler
+
+        Parameters
+        ----------
+        priors: bilby.core.prior.PriorDict
+            The set of priors
+        weight: float
+            Weighting factor
+        subset: list
+            A list of keys for which to restrict the proposal to (other parameters
+        will be kept fixed)
+        n_samples_for_covariance: int, optional, default=10000
+            Max number of samples to use for calculating covariance
+        min_samples: int, optional, default=100
+            Min number of samples before calculating covariance matrix
+        """
+        self.U = None
+        self.S = None
+        self.n_samples_for_covariance = n_samples_for_covariance
+        self.min_samples = min_samples
+        if subset is not None:
+            if len(subset) == 1:
+                raise ValueError(
+                    "Don't use CovarianceAMProposal on only one parameter."
+                )
+            self.parameter_indexes = np.array(
+                [self.parameters.index(key) for key in subset]
+            )
+        else:
+            self.parameter_indexes = np.arange(len(self.parameters))
+
+    def propose(self, chain):
+        prob = np.random.rand()
+        Nparams = len(self.parameter_indexes)
+        if prob > 0.97:
+            scale = 10
+        elif prob > 0.9:
+            scale = 0.2
+        else:
+            scale = 1.0
+
+        current_sample = chain.current_sample
+        if chain.position < self.min_samples:
+            covmat = np.eye(Nparams)
+        else:
+            lookback = min(chain.position, self.n_samples_for_covariance)
+            idxs = np.arange(-lookback, 0) + chain.position
+            recent_samples = chain._chain_array[idxs][:, self.parameter_indexes]
+            covmat = np.cov(recent_samples.T)
+        self.U, self.S, _ = np.linalg.svd(covmat)
+
+        # put recent sample into eigenbasis
+        y = self.U.T @ np.array(current_sample.list)[self.parameter_indexes]
+
+        cd = 2.4 / np.sqrt(2 * y.size) * scale
+
+        y += np.random.randn(y.size) * cd * np.sqrt(self.S)
+
+        # change back to original basis
+        q = self.U @ y
+
+        for jj, paridx in enumerate(self.parameter_indexes):
+            current_sample[self.parameters[paridx]] = q[jj]
+        return current_sample, 0
+
+
 class PriorProposal(BaseProposal):
     """A proposal using draws from the prior distribution