Repository: rhayes777/PyAutoFit
Pull request #997 (2024-04-15 15:28:37)

feature/parallel initializer
Performs initialisation in parallel batches for efficiency



Patch:
diff --git a/autofit/non_linear/initializer.py b/autofit/non_linear/initializer.py
index 4624d9e23..5747d6109 100644
--- a/autofit/non_linear/initializer.py
+++ b/autofit/non_linear/initializer.py
@@ -3,13 +3,15 @@ import logging
 import os
 import random
 from abc import ABC, abstractmethod
-from typing import Dict, Tuple, List
+from typing import Dict, Tuple, List, Optional
 
 import numpy as np
 
 from autofit import exc
+from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.mapper.prior.abstract import Prior
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
+from autofit.non_linear.parallel import SneakyPool
 
 logger = logging.getLogger(__name__)
 
@@ -23,13 +25,28 @@ class AbstractInitializer(ABC):
     def _generate_unit_parameter_list(self, model):
         pass
 
+    @staticmethod
+    def figure_of_metric(args) -> Optional[float]:
+        fitness, parameter_list = args
+        try:
+            figure_of_merit = fitness(parameters=parameter_list)
+
+            if np.isnan(figure_of_merit) or figure_of_merit < -1e98:
+                return None
+
+            return figure_of_merit
+        except exc.FitException:
+            return None
+
     def samples_from_model(
         self,
         total_points: int,
         model: AbstractPriorModel,
         fitness,
+        paths: AbstractPaths,
         use_prior_medians: bool = False,
         test_mode_samples: bool = True,
+        n_cores: int = 1,
     ):
         """
         Generate the initial points of the non-linear search, by randomly drawing unit values from a uniform
@@ -55,31 +72,39 @@ class AbstractInitializer(ABC):
         parameter_lists = []
         figures_of_merit_list = []
 
-        point_index = 0
-
-        while point_index < total_points:
-            if not use_prior_medians:
-                unit_parameter_list = self._generate_unit_parameter_list(model)
-            else:
-                unit_parameter_list = [0.5] * model.prior_count
-
-            parameter_list = model.vector_from_unit_vector(
-                unit_vector=unit_parameter_list
-            )
+        sneaky_pool = SneakyPool(n_cores, fitness, paths)
 
-            try:
-                figure_of_merit = fitness(parameters=parameter_list)
+        while len(figures_of_merit_list) < total_points:
+            remaining_points = total_points - len(figures_of_merit_list)
+            batch_size = min(remaining_points, n_cores)
+            parameter_lists_ = []
+            unit_parameter_lists_ = []
 
-                if np.isnan(figure_of_merit) or figure_of_merit < -1e98:
-                    raise exc.FitException
+            for _ in range(batch_size):
+                if not use_prior_medians:
+                    unit_parameter_list = self._generate_unit_parameter_list(model)
+                else:
+                    unit_parameter_list = [0.5] * model.prior_count
 
-                unit_parameter_lists.append(unit_parameter_list)
-                parameter_lists.append(parameter_list)
-                figures_of_merit_list.append(figure_of_merit)
-                point_index += 1
-            except exc.FitException:
-                pass
+                parameter_list = model.vector_from_unit_vector(
+                    unit_vector=unit_parameter_list
+                )
 
+                parameter_lists_.append(parameter_list)
+                unit_parameter_lists_.append(unit_parameter_list)
+
+            for figure_of_merit, unit_parameter_list, parameter_list in zip(
+                sneaky_pool.map(
+                    self.figure_of_metric,
+                    [(fitness, parameter_list) for parameter_list in parameter_lists_],
+                ),
+                unit_parameter_lists_,
+                parameter_lists_,
+            ):
+                if figure_of_merit is not None:
+                    unit_parameter_lists.append(unit_parameter_list)
+                    parameter_lists.append(parameter_list)
+                    figures_of_merit_list.append(figure_of_merit)
 
         if total_points > 1 and np.allclose(
             a=figures_of_merit_list[0], b=figures_of_merit_list[1:]
diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py
index 0f9e762da..e8c2d81a1 100644
--- a/autofit/non_linear/search/mcmc/emcee/search.py
+++ b/autofit/non_linear/search/mcmc/emcee/search.py
@@ -105,7 +105,7 @@ class Emcee(AbstractMCMC):
             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 @@ class Emcee(AbstractMCMC):
                 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 @@ class Emcee(AbstractMCMC):
             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 @@ class Emcee(AbstractMCMC):
         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 @@ class Emcee(AbstractMCMC):
             "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 @@ class Emcee(AbstractMCMC):
         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 @@ class Emcee(AbstractMCMC):
             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)
diff --git a/autofit/non_linear/search/mcmc/zeus/search.py b/autofit/non_linear/search/mcmc/zeus/search.py
index 18e014152..ea1d3b11f 100644
--- a/autofit/non_linear/search/mcmc/zeus/search.py
+++ b/autofit/non_linear/search/mcmc/zeus/search.py
@@ -13,6 +13,7 @@ from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelations
 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 @@ class Zeus(AbstractMCMC):
             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 @@ class Zeus(AbstractMCMC):
                 )
 
         except (FileNotFoundError, AttributeError):
-
             search_internal = zeus.EnsembleSampler(
                 nwalkers=self.config_dict_search["nwalkers"],
                 ndim=model.prior_count,
@@ -168,7 +167,9 @@ class Zeus(AbstractMCMC):
                 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 @@ class Zeus(AbstractMCMC):
             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 @@ class Zeus(AbstractMCMC):
             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 @@ class Zeus(AbstractMCMC):
                     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,13 +274,13 @@ class Zeus(AbstractMCMC):
 
         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()
@@ -287,8 +288,7 @@ class Zeus(AbstractMCMC):
 
         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 @@ class Zeus(AbstractMCMC):
             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,7 +310,6 @@ class Zeus(AbstractMCMC):
         )
 
     def auto_correlations_from(self, search_internal=None):
-
         import zeus
 
         search_internal = search_internal or self.paths.load_search_internal()
@@ -318,12 +317,12 @@ class Zeus(AbstractMCMC):
         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 @@ class Zeus(AbstractMCMC):
             **config_dict,
             "nwalkers": 20,
             "nsteps": 10,
-        }
\ No newline at end of file
+        }
diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py
index fe7319a4c..ccd87276f 100644
--- a/autofit/non_linear/search/nest/dynesty/search/abstract.py
+++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py
@@ -253,7 +253,7 @@ class AbstractDynesty(AbstractNest, ABC):
         return SamplesNest(
             model=model,
             sample_list=sample_list,
-            samples_info=self.samples_info_from(search_internal=search_internal)
+            samples_info=self.samples_info_from(search_internal=search_internal),
         )
 
     @property
@@ -334,7 +334,6 @@ class AbstractDynesty(AbstractNest, ABC):
         }
 
         if iterations > 0:
-
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore")
 
@@ -444,6 +443,8 @@ class AbstractDynesty(AbstractNest, ABC):
             total_points=self.number_live_points,
             model=model,
             fitness=fitness,
+            paths=self.paths,
+            n_cores=self.number_of_cores,
         )
 
         init_unit_parameters = np.zeros(
diff --git a/autofit/non_linear/search/optimize/drawer/search.py b/autofit/non_linear/search/optimize/drawer/search.py
index f4b40aba1..43c481740 100644
--- a/autofit/non_linear/search/optimize/drawer/search.py
+++ b/autofit/non_linear/search/optimize/drawer/search.py
@@ -103,7 +103,7 @@ class Drawer(AbstractOptimizer):
             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 @@ class Drawer(AbstractOptimizer):
             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 @@ class Drawer(AbstractOptimizer):
         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 @@ class Drawer(AbstractOptimizer):
             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]
diff --git a/autofit/non_linear/search/optimize/lbfgs/search.py b/autofit/non_linear/search/optimize/lbfgs/search.py
index 6875f505a..963160d0e 100644
--- a/autofit/non_linear/search/optimize/lbfgs/search.py
+++ b/autofit/non_linear/search/optimize/lbfgs/search.py
@@ -17,7 +17,6 @@ import numpy as np
 
 
 class LBFGS(AbstractOptimizer):
-
     __identifier_fields__ = ()
 
     def __init__(
@@ -71,7 +70,6 @@ class LBFGS(AbstractOptimizer):
 
     @cached_property
     def config_dict_options(self):
-
         config_dict = copy.deepcopy(self._class_config["options"])
 
         for key, value in config_dict.items():
@@ -110,11 +108,10 @@ class LBFGS(AbstractOptimizer):
             paths=self.paths,
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
-            convert_to_chi_squared=True
+            convert_to_chi_squared=True,
         )
 
         try:
-
             search_internal_dict = self.paths.load_search_internal()
 
             x0 = search_internal_dict["x0"]
@@ -125,13 +122,16 @@ class LBFGS(AbstractOptimizer):
             )
 
         except (FileNotFoundError, TypeError):
-
             (
                 unit_parameter_lists,
                 parameter_lists,
                 log_posterior_list,
             ) = self.initializer.samples_from_model(
-                total_points=1, model=model, fitness=fitness,
+                total_points=1,
+                model=model,
+                fitness=fitness,
+                paths=self.paths,
+                n_cores=self.number_of_cores,
             )
 
             x0 = np.asarray(parameter_lists[0])
@@ -145,13 +145,11 @@ class LBFGS(AbstractOptimizer):
         maxiter = self.config_dict_options.get("maxiter", 1e8)
 
         while total_iterations < maxiter:
-
             iterations_remaining = maxiter - total_iterations
 
             iterations = min(self.iterations_per_update, iterations_remaining)
 
             if iterations > 0:
-
                 config_dict_options = self.config_dict_options
                 config_dict_options["maxiter"] = iterations
 
@@ -165,7 +163,9 @@ class LBFGS(AbstractOptimizer):
 
                 total_iterations += search_internal.nit
 
-                search_internal.log_posterior_list = -0.5 * fitness(parameters=search_internal.x)
+                search_internal.log_posterior_list = -0.5 * fitness(
+                    parameters=search_internal.x
+                )
 
                 self.paths.save_search_internal(
                     obj=search_internal,
@@ -180,14 +180,16 @@ class LBFGS(AbstractOptimizer):
                     model=model,
                     analysis=analysis,
                     during_analysis=True,
-                    search_internal=search_internal
+                    search_internal=search_internal,
                 )
 
         self.logger.info("L-BFGS sampling complete.")
 
         return search_internal
 
-    def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=None):
+    def samples_via_internal_from(
+        self, model: AbstractPriorModel, search_internal=None
+    ):
         """
         Returns a `Samples` object from the LBFGS internal results.
 
@@ -204,7 +206,6 @@ class LBFGS(AbstractOptimizer):
         """
 
         if search_internal is None:
-
             search_internal = self.paths.load_search_internal()
 
         x0 = search_internal.x
@@ -215,9 +216,7 @@ class LBFGS(AbstractOptimizer):
         log_prior_list = model.log_prior_list_from(parameter_lists=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]
 
@@ -226,7 +225,7 @@ class LBFGS(AbstractOptimizer):
             parameter_lists=parameter_lists,
             log_likelihood_list=log_likelihood_list,
             log_prior_list=log_prior_list,
-            weight_list=weight_list
+            weight_list=weight_list,
         )
 
         samples_info = {
@@ -238,4 +237,4 @@ class LBFGS(AbstractOptimizer):
             model=model,
             sample_list=sample_list,
             samples_info=samples_info,
-        )
\ No newline at end of file
+        )
diff --git a/autofit/non_linear/search/optimize/pyswarms/search/abstract.py b/autofit/non_linear/search/optimize/pyswarms/search/abstract.py
index 5b7af1b2d..7c0acb7cd 100644
--- a/autofit/non_linear/search/optimize/pyswarms/search/abstract.py
+++ b/autofit/non_linear/search/optimize/pyswarms/search/abstract.py
@@ -13,7 +13,6 @@ from autofit.non_linear.samples.samples import Samples
 
 
 class FitnessPySwarms(Fitness):
-
     def __call__(self, parameters, *kwargs):
         """
         Interfaces with any non-linear in order to fit a model to the data and return a log likelihood via
@@ -46,11 +45,14 @@ class FitnessPySwarms(Fitness):
         figure_of_merit_list = []
 
         for params_of_particle in parameters:
-
             try:
                 instance = self.model.instance_from_vector(vector=params_of_particle)
-                log_likelihood = self.analysis.log_likelihood_function(instance=instance)
-                log_prior = self.model.log_prior_list_from_vector(vector=params_of_particle)
+                log_likelihood = self.analysis.log_likelihood_function(
+                    instance=instance
+                )
+                log_prior = self.model.log_prior_list_from_vector(
+                    vector=params_of_particle
+                )
                 log_posterior = log_likelihood + sum(log_prior)
                 figure_of_merit = -2.0 * log_posterior
             except exc.FitException:
@@ -64,18 +66,17 @@ class FitnessPySwarms(Fitness):
         return np.asarray(figure_of_merit_list)
 
 
-
 class AbstractPySwarms(AbstractOptimizer):
     def __init__(
-            self,
-            name: Optional[str] = None,
-            path_prefix: Optional[str] = None,
-            unique_tag: Optional[str] = None,
-            initializer: Optional[AbstractInitializer] = None,
-            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[AbstractInitializer] = None,
+        iterations_per_update: int = None,
+        number_of_cores: int = None,
+        session: Optional[sa.orm.Session] = None,
+        **kwargs
     ):
         """
         A PySwarms Particle Swarm Optimizer global non-linear search.
@@ -145,11 +146,10 @@ class AbstractPySwarms(AbstractOptimizer):
             analysis=analysis,
             fom_is_log_likelihood=False,
             resample_figure_of_merit=-np.inf,
-            convert_to_chi_squared=True
+            convert_to_chi_squared=True,
         )
 
         try:
-
             search_internal = self.paths.load_search_internal()
 
             init_pos = search_internal.pos_history[-1]
@@ -160,17 +160,23 @@ class AbstractPySwarms(AbstractOptimizer):
             )
 
         except (FileNotFoundError, TypeError, AttributeError):
-
-            unit_parameter_lists, parameter_lists, log_posterior_list = self.initializer.samples_from_model(
+            (
+                unit_parameter_lists,
+                parameter_lists,
+                log_posterior_list,
+            ) = self.initializer.samples_from_model(
                 total_points=self.config_dict_search["n_particles"],
                 model=model,
                 fitness=fitness,
+                paths=self.paths,
+                n_cores=self.number_of_cores,
             )
 
-            init_pos = np.zeros(shape=(self.config_dict_search["n_particles"], model.prior_count))
+            init_pos = np.zeros(
+                shape=(self.config_dict_search["n_particles"], model.prior_count)
+            )
 
             for index, parameters in enumerate(parameter_lists):
-
                 init_pos[index, :] = np.asarray(parameters)
 
             total_iterations = 0
@@ -182,12 +188,10 @@ class AbstractPySwarms(AbstractOptimizer):
         ## TODO : Use actual limits
 
         vector_lower = model.vector_from_unit_vector(
-            unit_vector=[1e-6] * model.prior_count,
-            ignore_prior_limits=True
+            unit_vector=[1e-6] * model.prior_count, ignore_prior_limits=True
         )
         vector_upper = model.vector_from_unit_vector(
-            unit_vector=[0.9999999] * model.prior_count,
-            ignore_prior_limits=True
+            unit_vector=[0.9999999] * model.prior_count, ignore_prior_limits=True
         )
 
         lower_bounds = [lower for lower in vector_lower]
@@ -196,12 +200,8 @@ class AbstractPySwarms(AbstractOptimizer):
         bounds = (np.asarray(lower_bounds), np.asarray(upper_bounds))
 
         while total_iterations < self.config_dict_run["iters"]:
-
             search_internal = self.search_internal_from(
-                model=model,
-                fitness=fitness,
-                bounds=bounds,
-                init_pos=init_pos
+                model=model, fitness=fitness, bounds=bounds, init_pos=init_pos
             )
 
             iterations_remaining = self.config_dict_run["iters"] - total_iterations
@@ -209,8 +209,9 @@ class AbstractPySwarms(AbstractOptimizer):
             iterations = min(self.iterations_per_update, iterations_remaining)
 
             if iterations > 0:
-
-                search_internal.optimize(objective_func=fitness.__call__, iters=iterations)
+                search_internal.optimize(
+                    objective_func=fitness.__call__, iters=iterations
+                )
 
                 total_iterations += iterations
 
@@ -222,7 +223,7 @@ class AbstractPySwarms(AbstractOptimizer):
                     model=model,
                     analysis=analysis,
                     during_analysis=True,
-                    search_internal=search_internal
+                    search_internal=search_internal,
                 )
 
                 init_pos = search_internal.pos_history[-1]
@@ -246,12 +247,13 @@ class AbstractPySwarms(AbstractOptimizer):
         """
 
         if search_internal is None:
-
             search_internal = self.paths.load_search_internal()
 
         search_internal_dict = {
             "total_iterations": None,
-            "log_posterior_list": [-0.5 * cost for cost in search_internal.cost_history],
+            "log_posterior_list": [
+                -0.5 * cost for cost in search_internal.cost_history
+            ],
             "time": self.timer.time if self.timer else None,
         }
         pos_history = search_internal.pos_history
@@ -263,7 +265,9 @@ class AbstractPySwarms(AbstractOptimizer):
 
         log_posterior_list = search_internal_dict["log_posterior_list"]
         log_prior_list = model.log_prior_list_from(parameter_lists=parameter_lists)
-        log_likelihood_list = [lp - prior for lp, prior in zip(log_posterior_list, log_prior_list)]
+        log_likelihood_list = [
+            lp - prior for lp, prior in zip(log_posterior_list, log_prior_list)
+        ]
         weight_list = len(log_likelihood_list) * [1.0]
 
         sample_list = Sample.from_lists(
@@ -271,7 +275,7 @@ class AbstractPySwarms(AbstractOptimizer):
             parameter_lists=parameter_lists_2,
             log_likelihood_list=log_likelihood_list,
             log_prior_list=log_prior_list,
-            weight_list=weight_list
+            weight_list=weight_list,
         )
 
         return Samples(
@@ -304,4 +308,4 @@ class AbstractPySwarms(AbstractOptimizer):
         }
 
     def search_internal_from(self, model, fitness, bounds, init_pos):
-        raise NotImplementedError()
\ No newline at end of file
+        raise NotImplementedError()


Test patch:
diff --git a/test_autofit/non_linear/test_initializer.py b/test_autofit/non_linear/test_initializer.py
index 74a243556..c63397306 100644
--- a/test_autofit/non_linear/test_initializer.py
+++ b/test_autofit/non_linear/test_initializer.py
@@ -1,21 +1,19 @@
 import os
+from random import random
+
 import pytest
 
 import autofit as af
 
 
 class MockFitness:
-
-    def __init__(self, figure_of_merit=0.0, increase_figure_of_merit = True):
-
+    def __init__(self, figure_of_merit=0.0, change_figure_of_merit=True):
         self.figure_of_merit = figure_of_merit
-        self.increase_figure_of_merit = increase_figure_of_merit
+        self.change_figure_of_merit = change_figure_of_merit
 
     def __call__(self, parameters):
-
-        if self.increase_figure_of_merit:
-            self.figure_of_merit += 1
-
+        if self.change_figure_of_merit:
+            return -random() * 10
         return self.figure_of_merit
 
 
@@ -28,8 +26,15 @@ def test__priors__samples_from_model():
 
     initializer = af.InitializerPrior()
 
-    unit_parameter_lists, parameter_lists, figure_of_merit_list = initializer.samples_from_model(
-        total_points=2, model=model, fitness=MockFitness()
+    (
+        unit_parameter_lists,
+        parameter_lists,
+        figure_of_merit_list,
+    ) = initializer.samples_from_model(
+        total_points=2,
+        model=model,
+        fitness=MockFitness(),
+        paths=af.DirectoryPaths(),
     )
 
     assert 0.0 < unit_parameter_lists[0][0] < 1.0
@@ -50,7 +55,6 @@ def test__priors__samples_from_model():
     assert 0.399 < parameter_lists[0][3] < 0.401
     assert 0.399 < parameter_lists[1][3] < 0.401
 
-    assert figure_of_merit_list == [1.0, 2.0]
 
 def test__priors__samples_from_model__raise_exception_if_all_likelihoods_identical():
     model = af.Model(af.m.MockClassx4)
@@ -58,13 +62,15 @@ def test__priors__samples_from_model__raise_exception_if_all_likelihoods_identic
     initializer = af.InitializerPrior()
 
     with pytest.raises(af.exc.InitializerException):
-
         initializer.samples_from_model(
-            total_points=2, model=model, fitness=MockFitness(increase_figure_of_merit=False)
+            total_points=2,
+            model=model,
+            fitness=MockFitness(change_figure_of_merit=False),
+            paths=af.DirectoryPaths(),
         )
 
-def test__priors__samples_in_test_mode():
 
+def test__priors__samples_in_test_mode():
     os.environ["PYAUTOFIT_TEST_MODE"] = "1"
 
     model = af.Model(af.m.MockClassx4)
@@ -75,8 +81,15 @@ def test__priors__samples_in_test_mode():
 
     initializer = af.InitializerPrior()
 
-    unit_parameter_lists, parameter_lists, figure_of_merit_list = initializer.samples_from_model(
-        total_points=2, model=model, fitness=None
+    (
+        unit_parameter_lists,
+        parameter_lists,
+        figure_of_merit_list,
+    ) = initializer.samples_from_model(
+        total_points=2,
+        model=model,
+        fitness=None,
+        paths=af.DirectoryPaths(),
     )
 
     assert 0.0 < unit_parameter_lists[0][0] < 1.0
@@ -101,8 +114,8 @@ def test__priors__samples_in_test_mode():
 
     os.environ["PYAUTOFIT_TEST_MODE"] = "0"
 
-def test__ball__samples_sample_centre_of_priors():
 
+def test__ball__samples_sample_centre_of_priors():
     model = af.Model(af.m.MockClassx4)
     model.one = af.UniformPrior(lower_limit=0.0, upper_limit=1.0)
     model.two = af.UniformPrior(lower_limit=0.0, upper_limit=2.0)
@@ -111,8 +124,15 @@ def test__ball__samples_sample_centre_of_priors():
 
     initializer = af.InitializerBall(lower_limit=0.4999, upper_limit=0.5001)
 
-    unit_parameter_lists, parameter_lists, figure_of_merit_list = initializer.samples_from_model(
-        total_points=2, model=model, fitness=MockFitness()
+    (
+        unit_parameter_lists,
+        parameter_lists,
+        figure_of_merit_list,
+    ) = initializer.samples_from_model(
+        total_points=2,
+        model=model,
+        fitness=MockFitness(),
+        paths=af.DirectoryPaths(),
     )
 
     assert 0.4999 < unit_parameter_lists[0][0] < 0.5001
@@ -135,8 +155,15 @@ def test__ball__samples_sample_centre_of_priors():
 
     initializer = af.InitializerBall(lower_limit=0.7999, upper_limit=0.8001)
 
-    unit_parameter_lists, parameter_lists, figure_of_merit_list = initializer.samples_from_model(
-        total_points=2, model=model, fitness=MockFitness()
+    (
+        unit_parameter_lists,
+        parameter_lists,
+        figure_of_merit_list,
+    ) = initializer.samples_from_model(
+        total_points=2,
+        model=model,
+        fitness=MockFitness(),
+        paths=af.DirectoryPaths(),
     )
 
     assert 0.799 < parameter_lists[0][0] < 0.801
@@ -148,8 +175,6 @@ def test__ball__samples_sample_centre_of_priors():
     assert 3.199 < parameter_lists[0][3] < 3.201
     assert 3.199 < parameter_lists[1][3] < 3.201
 
-    assert figure_of_merit_list == [1.0, 2.0]
-
 
 @pytest.mark.parametrize(
     "unit_value, physical_value",
@@ -157,7 +182,7 @@ def test__ball__samples_sample_centre_of_priors():
         (0.0, 0.0),
         (0.5, 0.5),
         (1.0, 1.0),
-    ]
+    ],
 )
 def test_invert_physical(unit_value, physical_value):
     prior = af.UniformPrior(
@@ -173,7 +198,7 @@ def test_invert_physical(unit_value, physical_value):
         (1.0, 0.0),
         (2.0, 0.5),
         (3.0, 1.0),
-    ]
+    ],
 )
 def test_invert_physical_offset(unit_value, physical_value):
     prior = af.UniformPrior(
@@ -189,7 +214,7 @@ def test_invert_physical_offset(unit_value, physical_value):
         (-float("inf"), 0.0),
         (0.0, 0.5),
         (float("inf"), 1.0),
-    ]
+    ],
 )
 def test_invert_gaussian(unit_value, physical_value):
     prior = af.GaussianPrior(
@@ -210,11 +235,13 @@ def make_model():
 
 
 def test_starting_point_initializer(model):
-    initializer = af.SpecificRangeInitializer({
-        model.centre: (1.0, 2.0),
-        model.normalization: (2.0, 3.0),
-        model.sigma: (-2.0, -1.0),
-    })
+    initializer = af.SpecificRangeInitializer(
+        {
+            model.centre: (1.0, 2.0),
+            model.normalization: (2.0, 3.0),
+            model.sigma: (-2.0, -1.0),
+        }
+    )
 
     parameter_list = initializer._generate_unit_parameter_list(model)
     assert len(parameter_list) == 3
@@ -223,11 +250,13 @@ def test_starting_point_initializer(model):
 
 
 def test_offset(model):
-    initializer = af.SpecificRangeInitializer({
-        model.centre: (1.5, 2.0),
-        model.normalization: (2.5, 3.0),
-        model.sigma: (-1.5, -1.0),
-    })
+    initializer = af.SpecificRangeInitializer(
+        {
+            model.centre: (1.5, 2.0),
+            model.normalization: (2.5, 3.0),
+            model.sigma: (-1.5, -1.0),
+        }
+    )
 
     parameter_list = initializer._generate_unit_parameter_list(model)
     assert len(parameter_list) == 3