From e2dfd2698db2c96e34ef9f751e07b09603dc581d Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 15 Apr 2024 16:14:37 +0100 Subject: [PATCH 1/4] implemented batched, sneaky pool initialization --- autofit/non_linear/initializer.py | 69 +++++++++++++------ .../search/nest/dynesty/search/abstract.py | 4 +- .../search/optimize/lbfgs/search.py | 32 ++++----- 3 files changed, 64 insertions(+), 41 deletions(-) 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 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:] diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index fe7319a4c..28d0b4267 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 @@ def samples_via_internal_from(self, model, search_internal=None): 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 @@ def run_search_internal( } if iterations > 0: - with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -444,6 +443,7 @@ def live_points_init_from(self, model, fitness): total_points=self.number_live_points, model=model, fitness=fitness, + paths=self.paths, ) init_unit_parameters = np.zeros( diff --git a/autofit/non_linear/search/optimize/lbfgs/search.py b/autofit/non_linear/search/optimize/lbfgs/search.py index 6875f505a..60186f1b8 100644 --- a/autofit/non_linear/search/optimize/lbfgs/search.py +++ b/autofit/non_linear/search/optimize/lbfgs/search.py @@ -17,7 +17,6 @@ class LBFGS(AbstractOptimizer): - __identifier_fields__ = () def __init__( @@ -71,7 +70,6 @@ def __init__( @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 @@ def _fit( 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,15 @@ def _fit( ) 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, ) x0 = np.asarray(parameter_lists[0]) @@ -145,13 +144,11 @@ def _fit( 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 +162,9 @@ def _fit( 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 +179,16 @@ def _fit( 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 +205,6 @@ def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=N """ if search_internal is None: - search_internal = self.paths.load_search_internal() x0 = search_internal.x @@ -215,9 +215,7 @@ def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=N 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 +224,7 @@ def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=N 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 +236,4 @@ def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=N model=model, sample_list=sample_list, samples_info=samples_info, - ) \ No newline at end of file + ) From b9ddc12fd687d001cfa689fa32eae2553cb7c420 Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 15 Apr 2024 16:22:47 +0100 Subject: [PATCH 2/4] fixing tests --- test_autofit/non_linear/test_initializer.py | 76 +++++++++++++-------- 1 file changed, 48 insertions(+), 28 deletions(-) diff --git a/test_autofit/non_linear/test_initializer.py b/test_autofit/non_linear/test_initializer.py index 74a243556..ffb9788fd 100644 --- a/test_autofit/non_linear/test_initializer.py +++ b/test_autofit/non_linear/test_initializer.py @@ -5,14 +5,11 @@ class MockFitness: - - def __init__(self, figure_of_merit=0.0, increase_figure_of_merit = True): - + def __init__(self, figure_of_merit=0.0, increase_figure_of_merit=True): self.figure_of_merit = figure_of_merit self.increase_figure_of_merit = increase_figure_of_merit def __call__(self, parameters): - if self.increase_figure_of_merit: self.figure_of_merit += 1 @@ -28,8 +25,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 @@ -52,19 +56,21 @@ def test__priors__samples_from_model(): 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) 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(increase_figure_of_merit=False), ) -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,9 +81,11 @@ 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) assert 0.0 < unit_parameter_lists[0][0] < 1.0 assert 0.0 < unit_parameter_lists[1][0] < 1.0 @@ -101,8 +109,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,7 +119,11 @@ 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( + ( + unit_parameter_lists, + parameter_lists, + figure_of_merit_list, + ) = initializer.samples_from_model( total_points=2, model=model, fitness=MockFitness() ) @@ -135,7 +147,11 @@ 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( + ( + unit_parameter_lists, + parameter_lists, + figure_of_merit_list, + ) = initializer.samples_from_model( total_points=2, model=model, fitness=MockFitness() ) @@ -157,7 +173,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 +189,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 +205,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 +226,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 +241,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 From c88347c393d7479ab3d8fbd90d38c886dc4643cd Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 15 Apr 2024 16:27:15 +0100 Subject: [PATCH 3/4] fixing initializer tests... --- test_autofit/non_linear/test_initializer.py | 35 +++++++++++++-------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/test_autofit/non_linear/test_initializer.py b/test_autofit/non_linear/test_initializer.py index ffb9788fd..c63397306 100644 --- a/test_autofit/non_linear/test_initializer.py +++ b/test_autofit/non_linear/test_initializer.py @@ -1,18 +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 @@ -54,8 +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) @@ -66,7 +65,8 @@ def test__priors__samples_from_model__raise_exception_if_all_likelihoods_identic initializer.samples_from_model( total_points=2, model=model, - fitness=MockFitness(increase_figure_of_merit=False), + fitness=MockFitness(change_figure_of_merit=False), + paths=af.DirectoryPaths(), ) @@ -85,7 +85,12 @@ def test__priors__samples_in_test_mode(): unit_parameter_lists, parameter_lists, figure_of_merit_list, - ) = initializer.samples_from_model(total_points=2, model=model, fitness=None) + ) = initializer.samples_from_model( + total_points=2, + model=model, + fitness=None, + paths=af.DirectoryPaths(), + ) assert 0.0 < unit_parameter_lists[0][0] < 1.0 assert 0.0 < unit_parameter_lists[1][0] < 1.0 @@ -124,7 +129,10 @@ def test__ball__samples_sample_centre_of_priors(): parameter_lists, figure_of_merit_list, ) = initializer.samples_from_model( - total_points=2, model=model, fitness=MockFitness() + total_points=2, + model=model, + fitness=MockFitness(), + paths=af.DirectoryPaths(), ) assert 0.4999 < unit_parameter_lists[0][0] < 0.5001 @@ -152,7 +160,10 @@ def test__ball__samples_sample_centre_of_priors(): parameter_lists, figure_of_merit_list, ) = initializer.samples_from_model( - total_points=2, model=model, fitness=MockFitness() + total_points=2, + model=model, + fitness=MockFitness(), + paths=af.DirectoryPaths(), ) assert 0.799 < parameter_lists[0][0] < 0.801 @@ -164,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", From d296a7401e53217df2c8d978b282ee3ee387a741 Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 15 Apr 2024 16:30:38 +0100 Subject: [PATCH 4/4] ensure number of cores and paths are passed to initialiser --- .../non_linear/search/mcmc/emcee/search.py | 30 ++++--- autofit/non_linear/search/mcmc/zeus/search.py | 81 +++++++++---------- .../search/nest/dynesty/search/abstract.py | 1 + .../search/optimize/drawer/search.py | 16 ++-- .../search/optimize/lbfgs/search.py | 1 + .../optimize/pyswarms/search/abstract.py | 78 +++++++++--------- 6 files changed, 107 insertions(+), 100 deletions(-) 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 @@ 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) 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.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,13 +274,13 @@ 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() @@ -287,8 +288,7 @@ def samples_via_internal_from(self, model, search_internal=None): 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,7 +310,6 @@ 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() @@ -318,12 +317,12 @@ def auto_correlations_from(self, search_internal=None): 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, - } \ 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 28d0b4267..ccd87276f 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/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( 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 @@ 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] diff --git a/autofit/non_linear/search/optimize/lbfgs/search.py b/autofit/non_linear/search/optimize/lbfgs/search.py index 60186f1b8..963160d0e 100644 --- a/autofit/non_linear/search/optimize/lbfgs/search.py +++ b/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]) 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 @@ 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 @@ def __call__(self, parameters, *kwargs): 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 @@ def __call__(self, parameters, *kwargs): 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 @@ def _fit(self, model: AbstractPriorModel, analysis): 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 @@ def _fit(self, model: AbstractPriorModel, analysis): ) 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 @@ def _fit(self, model: AbstractPriorModel, analysis): ## 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 @@ def _fit(self, model: AbstractPriorModel, analysis): 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 @@ def _fit(self, model: AbstractPriorModel, analysis): 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 @@ def _fit(self, model: AbstractPriorModel, analysis): 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 @@ def samples_via_internal_from(self, model, search_internal=None): """ 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 @@ def samples_via_internal_from(self, model, search_internal=None): 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 @@ def samples_via_internal_from(self, model, search_internal=None): 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 @@ def config_dict_test_mode_from(self, config_dict: Dict) -> Dict: } def search_internal_from(self, model, fitness, bounds, init_pos): - raise NotImplementedError() \ No newline at end of file + raise NotImplementedError()