Merge pull request #997 from rhayes777/feature/parallel_initializer

feature/parallel initializer

Author: Jammy2211

autofit/non_linear/initializer.py

@@ -3,13 +3,15 @@
 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 @@ def samples_from_model(
         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:]

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)