ensure number of cores and paths are passed to initialiser

autofit/non_linear/search/mcmc/emcee/search.py

@@ -105,7 +105,7 @@ def _fit(self, model: AbstractPriorModel, analysis):
             analysis=analysis,
             paths=self.paths,
             fom_is_log_likelihood=False,
-            resample_figure_of_merit=-np.inf
+            resample_figure_of_merit=-np.inf,
         )
 
         pool = self.make_sneaky_pool(fitness)
@@ -147,6 +147,8 @@ def _fit(self, model: AbstractPriorModel, analysis):
                 total_points=search_internal.nwalkers,
                 model=model,
                 fitness=fitness,
+                paths=self.paths,
+                n_cores=self.number_of_cores,
             )
 
             state = np.zeros(shape=(search_internal.nwalkers, model.prior_count))
@@ -184,17 +186,19 @@ def _fit(self, model: AbstractPriorModel, analysis):
             samples = self.samples_from(model=model, search_internal=search_internal)
 
             if self.auto_correlation_settings.check_for_convergence:
-                if search_internal.iteration > self.auto_correlation_settings.check_size:
+                if (
+                    search_internal.iteration
+                    > self.auto_correlation_settings.check_size
+                ):
                     if samples.converged:
                         iterations_remaining = 0
 
             if iterations_remaining > 0:
-
                 self.perform_update(
                     model=model,
                     analysis=analysis,
                     search_internal=search_internal,
-                    during_analysis=True
+                    during_analysis=True,
                 )
 
         return search_internal
@@ -214,7 +218,6 @@ def output_search_internal(self, search_internal):
         pass
 
     def samples_info_from(self, search_internal=None):
-
         search_internal = search_internal or self.backend
 
         auto_correlations = self.auto_correlations_from(search_internal=search_internal)
@@ -225,7 +228,7 @@ def samples_info_from(self, search_internal=None):
             "change_threshold": auto_correlations.change_threshold,
             "total_walkers": len(search_internal.get_chain()[0, :, 0]),
             "total_steps": len(search_internal.get_log_prob()),
-            "time": self.timer.time if self.timer else None
+            "time": self.timer.time if self.timer else None,
         }
 
     def samples_via_internal_from(self, model, search_internal=None):
@@ -247,14 +250,14 @@ def samples_via_internal_from(self, model, search_internal=None):
         search_internal = search_internal or self.backend
 
         if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
-
             samples_after_burn_in = search_internal.get_chain(
-            discard=5, thin=5, flat=True
-        )
+                discard=5, thin=5, flat=True
+            )
 
         else:
-
-            auto_correlations = self.auto_correlations_from(search_internal=search_internal)
+            auto_correlations = self.auto_correlations_from(
+                search_internal=search_internal
+            )
 
             discard = int(3.0 * np.max(auto_correlations.times))
             thin = int(np.max(auto_correlations.times) / 2.0)
@@ -292,11 +295,12 @@ def samples_via_internal_from(self, model, search_internal=None):
             sample_list=sample_list,
             samples_info=self.samples_info_from(search_internal=search_internal),
             auto_correlation_settings=self.auto_correlation_settings,
-            auto_correlations=self.auto_correlations_from(search_internal=search_internal),
+            auto_correlations=self.auto_correlations_from(
+                search_internal=search_internal
+            ),
         )
 
     def auto_correlations_from(self, search_internal=None):
-
         search_internal = search_internal or self.backend
 
         times = search_internal.get_autocorr_time(tol=0)

autofit/non_linear/search/mcmc/zeus/search.py

@@ -13,6 +13,7 @@
 from autofit.non_linear.samples.sample import Sample
 from autofit.non_linear.samples.mcmc import SamplesMCMC
 
+
 class Zeus(AbstractMCMC):
     __identifier_fields__ = (
         "nwalkers",
@@ -24,16 +25,16 @@ class Zeus(AbstractMCMC):
     )
 
     def __init__(
-            self,
-            name: Optional[str] = None,
-            path_prefix: Optional[str] = None,
-            unique_tag: Optional[str] = None,
-            initializer: Optional[Initializer] = None,
-            auto_correlation_settings=AutoCorrelationsSettings(),
-            iterations_per_update: int = None,
-            number_of_cores: int = None,
-            session: Optional[sa.orm.Session] = None,
-            **kwargs
+        self,
+        name: Optional[str] = None,
+        path_prefix: Optional[str] = None,
+        unique_tag: Optional[str] = None,
+        initializer: Optional[Initializer] = None,
+        auto_correlation_settings=AutoCorrelationsSettings(),
+        iterations_per_update: int = None,
+        number_of_cores: int = None,
+        session: Optional[sa.orm.Session] = None,
+        **kwargs
     ):
         """
         An Zeus non-linear search.
@@ -126,11 +127,10 @@ def _fit(self, model: AbstractPriorModel, analysis):
             analysis=analysis,
             paths=self.paths,
             fom_is_log_likelihood=False,
-            resample_figure_of_merit=-np.inf
+            resample_figure_of_merit=-np.inf,
         )
 
         try:
-
             search_internal = self.paths.load_search_internal()
 
             state = search_internal.get_last_sample()
@@ -150,7 +150,6 @@ def _fit(self, model: AbstractPriorModel, analysis):
                 )
 
         except (FileNotFoundError, AttributeError):
-
             search_internal = zeus.EnsembleSampler(
                 nwalkers=self.config_dict_search["nwalkers"],
                 ndim=model.prior_count,
@@ -168,7 +167,9 @@ def _fit(self, model: AbstractPriorModel, analysis):
                 total_points=search_internal.nwalkers,
                 model=model,
                 fitness=fitness,
-                test_mode_samples=False
+                test_mode_samples=False,
+                paths=self.paths,
+                n_cores=self.number_of_cores,
             )
 
             state = np.zeros(shape=(search_internal.nwalkers, model.prior_count))
@@ -184,17 +185,16 @@ def _fit(self, model: AbstractPriorModel, analysis):
             iterations_remaining = self.config_dict_run["nsteps"]
 
         while iterations_remaining > 0:
-
             if self.iterations_per_update > iterations_remaining:
                 iterations = iterations_remaining
             else:
                 iterations = self.iterations_per_update
 
             for sample in search_internal.sample(
-                    start=state,
-                    log_prob0=log_posterior_list,
-                    iterations=iterations,
-                    progress=True,
+                start=state,
+                log_prob0=log_posterior_list,
+                iterations=iterations,
+                progress=True,
             ):
                 pass
 
@@ -213,11 +213,16 @@ def _fit(self, model: AbstractPriorModel, analysis):
             samples = self.samples_from(model=model, search_internal=search_internal)
 
             if self.auto_correlation_settings.check_for_convergence:
-                if search_internal.iteration > self.auto_correlation_settings.check_size:
+                if (
+                    search_internal.iteration
+                    > self.auto_correlation_settings.check_size
+                ):
                     if samples.converged:
                         iterations_remaining = 0
 
-            auto_correlation_time = zeus.AutoCorrTime(samples=search_internal.get_chain())
+            auto_correlation_time = zeus.AutoCorrTime(
+                samples=search_internal.get_chain()
+            )
 
             discard = int(3.0 * np.max(auto_correlation_time))
             thin = int(np.max(auto_correlation_time) / 2.0)
@@ -228,23 +233,19 @@ def _fit(self, model: AbstractPriorModel, analysis):
                     iterations_remaining = 0
 
             if iterations_remaining > 0:
-
                 self.perform_update(
                     model=model,
                     analysis=analysis,
                     search_internal=search_internal,
-                    during_analysis=True
+                    during_analysis=True,
                 )
 
         return search_internal
 
-    def samples_info_from(self, search_internal = None):
-
+    def samples_info_from(self, search_internal=None):
         search_internal = search_internal or self.paths.load_search_internal()
 
-        auto_correlations = self.auto_correlations_from(
-            search_internal=search_internal
-        )
+        auto_correlations = self.auto_correlations_from(search_internal=search_internal)
 
         return {
             "check_size": auto_correlations.check_size,
@@ -273,22 +274,21 @@ def samples_via_internal_from(self, model, search_internal=None):
 
         search_internal = search_internal or self.paths.load_search_internal()
 
-        auto_correlations = self.auto_correlations_from(
-            search_internal=search_internal
-        )
+        auto_correlations = self.auto_correlations_from(search_internal=search_internal)
 
         discard = int(3.0 * np.max(auto_correlations.times))
         thin = int(np.max(auto_correlations.times) / 2.0)
-        samples_after_burn_in =  search_internal.get_chain(discard=discard, thin=thin, flat=True)
+        samples_after_burn_in = search_internal.get_chain(
+            discard=discard, thin=thin, flat=True
+        )
 
         parameter_lists = samples_after_burn_in.tolist()
         log_posterior_list = search_internal.get_log_prob(flat=True).tolist()
         log_prior_list = model.log_prior_list_from(parameter_lists=parameter_lists)
 
         log_likelihood_list = [
             log_posterior - log_prior
-            for log_posterior, log_prior
-            in zip(log_posterior_list, log_prior_list)
+            for log_posterior, log_prior in zip(log_posterior_list, log_prior_list)
         ]
 
         weight_list = len(log_likelihood_list) * [1.0]
@@ -298,7 +298,7 @@ def samples_via_internal_from(self, model, search_internal=None):
             parameter_lists=parameter_lists,
             log_likelihood_list=log_likelihood_list,
             log_prior_list=log_prior_list,
-            weight_list=weight_list
+            weight_list=weight_list,
         )
 
         return SamplesMCMC(
@@ -310,20 +310,19 @@ def samples_via_internal_from(self, model, search_internal=None):
         )
 
     def auto_correlations_from(self, search_internal=None):
-
         import zeus
 
         search_internal = search_internal or self.paths.load_search_internal()
 
         times = zeus.AutoCorrTime(samples=search_internal.get_chain())
         try:
             previous_auto_correlation_times = zeus.AutoCorrTime(
-                samples=search_internal.get_chain()[: - self.auto_correlation_settings.check_size, :, :],
+                samples=search_internal.get_chain()[
+                    : -self.auto_correlation_settings.check_size, :, :
+                ],
             )
         except IndexError:
-            self.logger.debug(
-                "Unable to compute previous auto correlation times."
-            )
+            self.logger.debug("Unable to compute previous auto correlation times.")
             previous_auto_correlation_times = None
 
         return AutoCorrelations(
@@ -356,4 +355,4 @@ def config_dict_test_mode_from(self, config_dict: Dict) -> Dict:
             **config_dict,
             "nwalkers": 20,
             "nsteps": 10,
-        }
+        }

autofit/non_linear/search/nest/dynesty/search/abstract.py

@@ -444,6 +444,7 @@ def live_points_init_from(self, model, fitness):
             model=model,
             fitness=fitness,
             paths=self.paths,
+            n_cores=self.number_of_cores,
         )
 
         init_unit_parameters = np.zeros(

autofit/non_linear/search/optimize/drawer/search.py

@@ -103,7 +103,7 @@ def _fit(self, model: AbstractPriorModel, analysis):
             paths=self.paths,
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
-            convert_to_chi_squared=False
+            convert_to_chi_squared=False,
         )
 
         total_draws = self.config_dict_search["total_draws"]
@@ -120,12 +120,14 @@ def _fit(self, model: AbstractPriorModel, analysis):
             total_points=self.config_dict_search["total_draws"],
             model=model,
             fitness=fitness,
+            paths=self.paths,
+            n_cores=self.number_of_cores,
         )
 
         search_internal = {
-            "parameter_lists" : parameter_lists,
-            "log_posterior_list" : log_posterior_list,
-            "time": self.timer.time
+            "parameter_lists": parameter_lists,
+            "log_posterior_list": log_posterior_list,
+            "time": self.timer.time,
         }
 
         self.paths.save_search_internal(
@@ -135,7 +137,6 @@ def _fit(self, model: AbstractPriorModel, analysis):
         self.logger.info("Drawer complete")
 
     def samples_from(self, model, search_internal):
-
         search_internal_dict = self.paths.load_search_internal()
 
         parameter_lists = search_internal_dict["parameter_lists"]
@@ -146,10 +147,7 @@ def samples_from(self, model, search_internal):
             for vector in parameter_lists
         ]
         log_likelihood_list = [
-            lp - prior
-            for lp, prior in zip(
-                log_posterior_list, log_prior_list
-            )
+            lp - prior for lp, prior in zip(log_posterior_list, log_prior_list)
         ]
 
         weight_list = len(log_likelihood_list) * [1.0]

autofit/non_linear/search/optimize/lbfgs/search.py

@@ -131,6 +131,7 @@ def _fit(
                 model=model,
                 fitness=fitness,
                 paths=self.paths,
+                n_cores=self.number_of_cores,
             )
 
             x0 = np.asarray(parameter_lists[0])