Add option to only initialize some parameters

bilby/core/sampler/ptemcee.py

@@ -98,23 +98,35 @@ class Ptemcee(MCMCSampler):
                  for _ in range(self.sampler_init_kwargs["nwalkers"])]
                 for _ in range(self.kwargs['ntemps'])]
 
-    def get_pos0_from_minimize(self):
+    def get_pos0_from_minimize(self, minimize_list=None):
         logger.info("Attempting to set pos0 from minimize")
         from scipy.optimize import minimize
 
+        if minimize_list is None:
+            minimize_list = self.search_parameter_keys
+            pos0 = np.zeros((self.kwargs["ntemps"], self.kwargs["nwalkers"], self.ndim))
+        else:
+            pos0 = np.array(self.get_pos0_from_prior())
+
+        likelihood_copy = copy.copy(self.likelihood)
+
         def neg_log_like(params):
+            likelihood_copy.parameters.update(
+                {key: val for key, val in zip(minimize_list, params)})
             try:
-                return -self.log_likelihood(params)
+                return -likelihood_copy.log_likelihood()
             except RuntimeError:
                 return +np.inf
         bounds = [(self.priors[key].minimum, self.priors[key].maximum)
-                  for key in self.search_parameter_keys]
+                  for key in minimize_list]
         trials = 0
         success = []
         while True:
-            x0 = self.get_random_draw_from_prior()
+            draw = self.priors.sample()
+            likelihood_copy.parameters.update(draw)
+            x0 = [draw[key] for key in minimize_list]
             res = minimize(
-                neg_log_like, x0, bounds=bounds, method='L-BFGS-B', tol=1e-15)
+                neg_log_like, x0, bounds=bounds, method='L-BFGS-B')
             if res.success:
                 success.append(res.x)
             if trials > 100:
@@ -122,11 +134,16 @@ class Ptemcee(MCMCSampler):
             if len(success) >= 3:
                 break
 
-        pos0_min = np.min(success, axis=0)
-        pos0_max = np.max(success, axis=0)
-        pos0 = np.random.uniform(
-            pos0_min, pos0_max,
-            size=(self.kwargs["ntemps"], self.kwargs["nwalkers"], self.ndim))
+        success = np.array(success)
+        for i, key in enumerate(minimize_list):
+            pos0_min = np.min(success[:, i])
+            pos0_max = np.max(success[:, i])
+            logger.info("Initialize {} walkers from {}->{}"
+                        .format(key, pos0_min, pos0_max))
+            j = self.search_parameter_keys.index(key)
+            pos0[:, :, j] = np.random.uniform(
+                pos0_min, pos0_max,
+                size=(self.kwargs["ntemps"], self.kwargs["nwalkers"]))
         return pos0
 
     def setup_sampler(self):
@@ -169,10 +186,12 @@ class Ptemcee(MCMCSampler):
         return self.sampler, pos0
 
     def get_pos0(self):
-        if self.pos0.lower() == "prior":
+        if isinstance(self.pos0, str) and self.pos0.lower() == "prior":
             return self.get_pos0_from_prior()
-        elif self.pos0.lower() == "minimize":
+        elif isinstance(self.pos0, str) and self.pos0.lower() == "minimize":
             return self.get_pos0_from_minimize()
+        elif isinstance(self.pos0, list):
+            return self.get_pos0_from_minimize(minimize_list=self.pos0)
         else:
             raise SamplerError("pos0={} not implemented".format(self.pos0))