diff --git a/searches/pyswarms/abstract.py b/searches/pyswarms/abstract.py index 19c78bf..bd6e665 100755 --- a/searches/pyswarms/abstract.py +++ b/searches/pyswarms/abstract.py @@ -1,8 +1,7 @@ -from typing import Dict, Optional +from typing import Optional import numpy as np -from autofit import exc from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.fitness import Fitness @@ -10,6 +9,7 @@ from autofit.non_linear.search.mle.abstract_mle import AbstractMLE from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.samples import Samples +from autofit.non_linear.test_mode import is_test_mode class FitnessPySwarms(Fitness): @@ -18,13 +18,9 @@ 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 an `Analysis` class. - The interface is described in full in the `__init__` docstring. - `PySwarms` have a unique interface in that lists of parameters corresponding to multiple particles are - passed to the fitness function. A bespoke `__call__` method is therefore required to handle this. - - The figure of merit is the log posterior multiplied by -2.0, which is the chi-squared value that is minimized - by the `PySwarms` non-linear search. + passed to the fitness function. A bespoke `__call__` method is therefore required to handle this, + delegating per-particle evaluation to ``call_wrap``. Parameters ---------- @@ -42,26 +38,10 @@ def __call__(self, parameters, *kwargs): if isinstance(parameters[0], float): parameters = [parameters] - 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_posterior = log_likelihood + sum(log_prior) - figure_of_merit = -2.0 * log_posterior - except exc.FitException: - figure_of_merit = np.nan - - if np.isnan(figure_of_merit): - figure_of_merit = -2.0 * self.resample_figure_of_merit - - figure_of_merit_list.append(figure_of_merit) + figure_of_merit_list = [ + self.call_wrap(params_of_particle) + for params_of_particle in parameters + ] return np.asarray(figure_of_merit_list) @@ -72,10 +52,16 @@ def __init__( name: Optional[str] = None, path_prefix: Optional[str] = None, unique_tag: Optional[str] = None, + n_particles: int = 50, + cognitive: float = 0.5, + social: float = 0.3, + inertia: float = 0.9, + iters: int = 2000, initializer: Optional[AbstractInitializer] = None, iterations_per_quick_update: int = None, iterations_per_full_update: int = None, - number_of_cores: int = None, + number_of_cores: int = 1, + silence: bool = False, session: Optional[sa.orm.Session] = None, **kwargs ): @@ -96,20 +82,26 @@ def __init__( unique_tag The name of a unique tag for this model-fit, which will be given a unique entry in the sqlite database and also acts as the folder after the path prefix and before the search name. + n_particles + The number of particles in the swarm. + cognitive + The cognitive parameter controlling how much a particle is influenced by its own best position. + social + The social parameter controlling how much a particle is influenced by the swarm's best position. + inertia + The inertia weight controlling the momentum of the particles. + iters + The total number of iterations the swarm performs. initializer Generates the initialize samples of non-linear parameter space (see autofit.non_linear.initializer). number_of_cores The number of cores sampling is performed using a Python multiprocessing Pool instance. + silence + If True, the default print output of the non-linear search is silenced. session An SQLalchemy session instance so the results of the model-fit are written to an SQLite database. """ - number_of_cores = ( - self._config("parallel", "number_of_cores") - if number_of_cores is None - else number_of_cores - ) - super().__init__( name=name, path_prefix=path_prefix, @@ -118,10 +110,20 @@ def __init__( iterations_per_quick_update=iterations_per_quick_update, iterations_per_full_update=iterations_per_full_update, number_of_cores=number_of_cores, + silence=silence, session=session, **kwargs ) + self.n_particles = n_particles + self.cognitive = cognitive + self.social = social + self.inertia = inertia + self.iters = iters + + if is_test_mode(): + self.apply_test_mode() + self.logger.debug("Creating PySwarms Search") def _fit(self, model: AbstractPriorModel, analysis): @@ -146,6 +148,7 @@ def _fit(self, model: AbstractPriorModel, analysis): fitness = FitnessPySwarms( model=model, analysis=analysis, + paths=self.paths, fom_is_log_likelihood=False, resample_figure_of_merit=-np.inf, convert_to_chi_squared=True, @@ -167,7 +170,7 @@ def _fit(self, model: AbstractPriorModel, analysis): parameter_lists, log_posterior_list, ) = self.initializer.samples_from_model( - total_points=self.config_dict_search["n_particles"], + total_points=self.n_particles, model=model, fitness=fitness, paths=self.paths, @@ -175,7 +178,7 @@ def _fit(self, model: AbstractPriorModel, analysis): ) init_pos = np.zeros( - shape=(self.config_dict_search["n_particles"], model.prior_count) + shape=(self.n_particles, model.prior_count) ) for index, parameters in enumerate(parameter_lists): @@ -201,12 +204,12 @@ def _fit(self, model: AbstractPriorModel, analysis): bounds = (np.asarray(lower_bounds), np.asarray(upper_bounds)) - while total_iterations < self.config_dict_run["iters"]: + while total_iterations < self.iters: search_internal = self.search_internal_from( model=model, fitness=fitness, bounds=bounds, init_pos=init_pos ) - iterations_remaining = self.config_dict_run["iters"] - total_iterations + iterations_remaining = self.iters - total_iterations iterations = min(self.iterations_per_full_update, iterations_remaining) @@ -296,28 +299,8 @@ def samples_via_internal_from(self, model, search_internal=None): samples_info=search_internal_dict, ) - def config_dict_test_mode_from(self, config_dict: Dict) -> Dict: - """ - Returns a configuration dictionary for test mode meaning that the sampler terminates as quickly as possible. - - Entries which set the total number of samples of the sampler (e.g. maximum calls, maximum likelihood - evaluations) are reduced to low values meaning it terminates nearly immediately. - - Parameters - ---------- - config_dict - The original configuration dictionary for this sampler which includes entries controlling how fast the - sampler terminates. - - Returns - ------- - A configuration dictionary where settings which control the sampler's number of samples are reduced so it - terminates as quickly as possible. - """ - return { - **config_dict, - "iters": 1, - } + def apply_test_mode(self): + self.iters = 1 def search_internal_from(self, model, fitness, bounds, init_pos): raise NotImplementedError() diff --git a/searches/pyswarms/globe.py b/searches/pyswarms/globe.py index d837692..da5b67e 100755 --- a/searches/pyswarms/globe.py +++ b/searches/pyswarms/globe.py @@ -21,12 +21,13 @@ def __init__( initializer: Optional[AbstractInitializer] = None, iterations_per_full_update: int = None, iterations_per_quick_update: int = None, - number_of_cores: int = None, + number_of_cores: int = 1, + silence: bool = False, session: Optional[sa.orm.Session] = None, **kwargs ): """ - A PySwarms Particle Swarm MLE global non-linear search. + A PySwarms Particle Swarm MLE global-best non-linear search. For a full description of PySwarms, checkout its Github and readthedocs webpages: @@ -47,6 +48,8 @@ def __init__( Generates the initialize samples of non-linear parameter space (see autofit.non_linear.initializer). number_of_cores The number of cores sampling is performed using a Python multiprocessing Pool instance. + silence + If True, the default print output of the non-linear search is silenced. """ super().__init__( @@ -57,6 +60,7 @@ def __init__( iterations_per_quick_update=iterations_per_quick_update, iterations_per_full_update=iterations_per_full_update, number_of_cores=number_of_cores, + silence=silence, session=session, **kwargs ) @@ -64,24 +68,20 @@ def __init__( self.logger.debug("Creating PySwarms Search") def search_internal_from(self, model, fitness, bounds, init_pos): - """Get the static Dynesty sampler which performs the non-linear search, passing it all associated input Dynesty - variables.""" + """Get the PySwarms GlobalBestPSO sampler which performs the non-linear search.""" import pyswarms options = { - "c1": self.config_dict_search["cognitive"], - "c2": self.config_dict_search["social"], - "w": self.config_dict_search["inertia"] + "c1": self.cognitive, + "c2": self.social, + "w": self.inertia, } - filter_list = ["cognitive", "social", "inertia"] - config_dict = {key: value for key, value in self.config_dict_search.items() if key not in filter_list} - return pyswarms.global_best.GlobalBestPSO( + n_particles=self.n_particles, dimensions=model.prior_count, bounds=bounds, init_pos=init_pos, options=options, - **config_dict ) \ No newline at end of file diff --git a/searches/pyswarms/local.py b/searches/pyswarms/local.py index 3a9f5cb..8fa6bcc 100755 --- a/searches/pyswarms/local.py +++ b/searches/pyswarms/local.py @@ -19,14 +19,17 @@ def __init__( name: Optional[str] = None, path_prefix: Optional[str] = None, unique_tag: Optional[str] = None, + number_of_k_neighbors: int = 3, + minkowski_p_norm: int = 2, iterations_per_quick_update: int = None, iterations_per_full_update: int = None, - number_of_cores: int = None, + number_of_cores: int = 1, + silence: bool = False, session: Optional[sa.orm.Session] = None, **kwargs ): """ - A PySwarms Particle Swarm MLE global non-linear search. + A PySwarms Particle Swarm MLE local-best non-linear search. For a full description of PySwarms, checkout its Github and readthedocs webpages: @@ -43,10 +46,14 @@ def __init__( unique_tag The name of a unique tag for this model-fit, which will be given a unique entry in the sqlite database and also acts as the folder after the path prefix and before the search name. - initializer - Generates the initialize samples of non-linear parameter space (see autofit.non_linear.initializer). + number_of_k_neighbors + The number of neighbors each particle considers in the local-best topology. + minkowski_p_norm + The Minkowski p-norm used to compute distances between particles. number_of_cores The number of cores sampling is performed using a Python multiprocessing Pool instance. + silence + If True, the default print output of the non-linear search is silenced. """ super().__init__( @@ -56,35 +63,35 @@ def __init__( iterations_per_quick_update=iterations_per_quick_update, iterations_per_full_update=iterations_per_full_update, number_of_cores=number_of_cores, + silence=silence, session=session, **kwargs ) + self.number_of_k_neighbors = number_of_k_neighbors + self.minkowski_p_norm = minkowski_p_norm + self.logger.debug("Creating PySwarms Search") def search_internal_from(self, model, fitness, bounds, init_pos): """ - Get the static Dynesty sampler which performs the non-linear search, passing it all associated input Dynesty - variables. + Get the PySwarms LocalBestPSO sampler which performs the non-linear search. """ import pyswarms options = { - "c1": self.config_dict_search["cognitive"], - "c2": self.config_dict_search["social"], - "w": self.config_dict_search["inertia"], - "k": self.config_dict_search["number_of_k_neighbors"], - "p": self.config_dict_search["minkowski_p_norm"], + "c1": self.cognitive, + "c2": self.social, + "w": self.inertia, + "k": self.number_of_k_neighbors, + "p": self.minkowski_p_norm, } - filter_list = ["cognitive", "social", "inertia", "number_of_k_neighbors", "minkowski_p_norm"] - config_dict = {key: value for key, value in self.config_dict_search.items() if key not in filter_list} - return pyswarms.local_best.LocalBestPSO( + n_particles=self.n_particles, dimensions=model.prior_count, bounds=bounds, init_pos=init_pos, options=options, - **config_dict ) \ No newline at end of file diff --git a/searches/ultranest/search.py b/searches/ultranest/search.py index ee73415..107772a 100755 --- a/searches/ultranest/search.py +++ b/searches/ultranest/search.py @@ -1,340 +1,408 @@ -import os -from typing import Dict, Optional - -from autofit.database.sqlalchemy_ import sa - -from autofit.mapper.prior_model.abstract import AbstractPriorModel -from autofit.non_linear.search.nest import abstract_nest -from autofit.non_linear.fitness import Fitness -from autofit.non_linear.samples.sample import Sample -from autofit.non_linear.samples.nest import SamplesNest - -class UltraNest(abstract_nest.AbstractNest): - __identifier_fields__ = ( - "draw_multiple", - "ndraw_min", - "ndraw_max", - "min_num_live_points", - "cluster_num_live_points", - "insertion_test_zscore_threshold", - "stepsampler_cls", - "nsteps" - ) - - def __init__( - self, - name: Optional[str] = None, - path_prefix: Optional[str] = None, - unique_tag: Optional[str] = None, - iterations_per_quick_update: int = None, - iterations_per_full_update: int = None, - number_of_cores: int = None, - session: Optional[sa.orm.Session] = None, - **kwargs - ): - """ - An UltraNest non-linear search. - - UltraNest is an optional requirement and must be installed manually via the command `pip install ultranest`. - It is optional as it has certain dependencies which are generally straight forward to install (e.g. Cython). - - For a full description of UltraNest and its Python wrapper PyUltraNest, checkout its Github and documentation - webpages: - - https://github.com/JohannesBuchner/UltraNest - https://johannesbuchner.github.io/UltraNest/readme.html - - Parameters - ---------- - name - The name of the search, controlling the last folder results are output. - path_prefix - The path of folders prefixing the name folder where results are output. - unique_tag - The name of a unique tag for this model-fit, which will be given a unique entry in the sqlite database - and also acts as the folder after the path prefix and before the search name. - iterations_per_full_update - The number of iterations performed between update (e.g. output latest model to hard-disk, visualization). - number_of_cores - The number of cores sampling is performed using a Python multiprocessing Pool instance. - session - An SQLalchemy session instance so the results of the model-fit are written to an SQLite database. - """ - - number_of_cores = ( - self._config("parallel", "number_of_cores") - if number_of_cores is None - else number_of_cores - ) - - super().__init__( - name=name, - path_prefix=path_prefix, - unique_tag=unique_tag, - iterations_per_quick_update=iterations_per_quick_update, - iterations_per_full_update=iterations_per_full_update, - number_of_cores=number_of_cores, - session=session, - **kwargs - ) - - for key, value in self.config_dict_stepsampler.items(): - setattr(self, key, value) - if self.config_dict_stepsampler["stepsampler_cls"] is None: - self.nsteps = None - - self.logger.debug("Creating UltraNest Search") - - def _fit(self, model: AbstractPriorModel, analysis): - """ - Fit a model using the search and the Analysis class which contains the data and returns the log likelihood from - instances of the model, which the `NonLinearSearch` seeks to maximize. - - Parameters - ---------- - model : ModelMapper - The model which generates instances for different points in parameter space. - analysis : Analysis - Contains the data and the log likelihood function which fits an instance of the model to the data, returning - the log likelihood the `NonLinearSearch` maximizes. - - Returns - ------- - A result object comprising the Samples object that includes the maximum log likelihood instance and full - set of accepted ssamples of the fit. - """ - - try: - import ultranest - except ModuleNotFoundError: - raise ModuleNotFoundError( - "\n--------------------\n" - "You are attempting to perform a model-fit using UltraNest. \n\n" - "However, the optional library UltraNest (https://johannesbuchner.github.io/UltraNest/index.html) is " - "not installed.\n\n" - "Install it via the command `pip install ultranest==3.6.2`.\n\n" - "----------------------" - ) - - fitness = Fitness( - model=model, - analysis=analysis, - paths=self.paths, - fom_is_log_likelihood=True, - resample_figure_of_merit=-1.0e99 - ) - - def prior_transform(cube): - return model.vector_from_unit_vector( - unit_vector=cube, - ) - - log_dir = self.paths.search_internal_path - - try: - checkpoint_exists = os.path.exists(log_dir / "chains") - except TypeError: - checkpoint_exists = False - - if checkpoint_exists: - self.logger.info( - "Resuming UltraNest non-linear search (previous samples found)." - ) - else: - self.logger.info( - "Starting new UltraNest non-linear search (no previous samples found)." - ) - - search_internal = ultranest.ReactiveNestedSampler( - param_names=model.parameter_names, - loglike=fitness.call_wrap, - transform=prior_transform, - log_dir=log_dir, - **self.config_dict_search - ) - - search_internal.stepsampler = self.stepsampler - - finished = False - - while not finished: - - try: - total_iterations = search_internal.ncall - except AttributeError: - total_iterations = 0 - - if self.config_dict_run["max_ncalls"] is not None: - iterations = self.config_dict_run["max_ncalls"] - else: - iterations = total_iterations + self.iterations_per_full_update - - if iterations > 0: - - filter_list = ["max_ncalls", "dkl", "lepsilon"] - config_dict_run = { - key: value for key, value - in self.config_dict_run.items() - if key - not in filter_list - } - - config_dict_run["update_interval_ncall"] = iterations - - search_internal.run( - max_ncalls=iterations, - **config_dict_run - ) - - self.paths.save_search_internal( - obj=search_internal.results, - ) - - iterations_after_run = search_internal.ncall - - if ( - total_iterations == iterations_after_run - or iterations_after_run == self.config_dict_run["max_ncalls"] - ): - finished = True - - if not finished: - - self.perform_update( - model=model, - analysis=analysis, - during_analysis=True, - fitness=fitness, - search_internal=search_internal - ) - - return search_internal, fitness - - def output_search_internal(self, search_internal): - """ - Output the sampler results to hard-disk in their internal format. - - UltraNest uses a backend to store and load results, therefore the outputting of the search internal to a - dill file is disabled. - - However, a dictionary of the search results is output to dill above. - - Parameters - ---------- - sampler - The nautilus sampler object containing the results of the model-fit. - """ - pass - - def samples_info_from(self, search_internal=None): - - search_internal = search_internal or self.paths.load_search_internal() - - return { - "log_evidence": search_internal["logz"], - "total_samples": search_internal["ncall"], - "total_accepted_samples": len(search_internal["weighted_samples"]["logl"]), - "time": self.timer.time if self.timer else None, - "number_live_points": self.config_dict_run["min_num_live_points"] - } - - def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=None): - """ - Returns a `Samples` object from the ultranest internal results. - - The samples contain all information on the parameter space sampling (e.g. the parameters, - log likelihoods, etc.). - - The internal search results are converted from the native format used by the search to lists of values - (e.g. `parameter_lists`, `log_likelihood_list`). - - Parameters - ---------- - model - Maps input vectors of unit parameter values to physical values and model instances via priors. - """ - - search_internal = search_internal.results or self.paths.load_search_internal() - - parameters = search_internal["weighted_samples"]["points"] - log_likelihood_list = search_internal["weighted_samples"]["logl"] - log_prior_list = [ - sum(model.log_prior_list_from_vector(vector=vector)) for vector in parameters - ] - weight_list = search_internal["weighted_samples"]["weights"] - - sample_list = Sample.from_lists( - model=model, - parameter_lists=parameters, - log_likelihood_list=log_likelihood_list, - log_prior_list=log_prior_list, - weight_list=weight_list - ) - - return SamplesNest( - model=model, - sample_list=sample_list, - samples_info=self.samples_info_from(search_internal=search_internal), - ) - - def config_dict_test_mode_from(self, config_dict: Dict) -> Dict: - """ - Returns a configuration dictionary for test mode meaning that the sampler terminates as quickly as possible. - - Entries which set the total number of samples of the sampler (e.g. maximum calls, maximum likelihood - evaluations) are reduced to low values meaning it terminates nearly immediately. - - Parameters - ---------- - config_dict - The original configuration dictionary for this sampler which includes entries controlling how fast the - sampler terminates. - - Returns - ------- - A configuration dictionary where settings which control the sampler's number of samples are reduced so it - terminates as quickly as possible. - """ - return { - **config_dict, - "max_iters": 1, - "max_ncalls": 1, - } - - @property - def config_dict_stepsampler(self): - - config_dict = {} - - config_dict_step = self.config_type[self.__class__.__name__]["stepsampler"] - - for key, value in config_dict_step.items(): - try: - config_dict[key] = self.kwargs[key] - except KeyError: - config_dict[key] = value - - return config_dict - - @property - def stepsampler(self): - - from ultranest import stepsampler - - config_dict_stepsampler = self.config_dict_stepsampler - stepsampler_cls = config_dict_stepsampler["stepsampler_cls"] - config_dict_stepsampler.pop("stepsampler_cls") - - if stepsampler_cls is None: - return None - elif stepsampler_cls == "RegionMHSampler": - return stepsampler.RegionMHSampler(**config_dict_stepsampler) - elif stepsampler_cls == "AHARMSampler": - config_dict_stepsampler.pop("scale") - return stepsampler.AHARMSampler(**config_dict_stepsampler) - elif stepsampler_cls == "CubeMHSampler": - return stepsampler.CubeMHSampler(**config_dict_stepsampler) - elif stepsampler_cls == "CubeSliceSampler": - return stepsampler.CubeSliceSampler(**config_dict_stepsampler) - elif stepsampler_cls == "RegionSliceSampler": - return stepsampler.RegionSliceSampler(**config_dict_stepsampler) \ No newline at end of file +import os +from typing import Dict, Optional + +from autofit.database.sqlalchemy_ import sa + +from autofit.mapper.prior_model.abstract import AbstractPriorModel +from autofit.non_linear.search.nest import abstract_nest +from autofit.non_linear.fitness import Fitness +from autofit.non_linear.samples.sample import Sample +from autofit.non_linear.samples.nest import SamplesNest +from autofit.non_linear.test_mode import is_test_mode + +class UltraNest(abstract_nest.AbstractNest): + __identifier_fields__ = ( + "draw_multiple", + "ndraw_min", + "ndraw_max", + "min_num_live_points", + "cluster_num_live_points", + "insertion_test_zscore_threshold", + "stepsampler_cls", + "nsteps" + ) + + def __init__( + self, + name: Optional[str] = None, + path_prefix: Optional[str] = None, + unique_tag: Optional[str] = None, + draw_multiple: bool = True, + ndraw_min: int = 128, + ndraw_max: int = 65536, + num_bootstraps: int = 30, + num_test_samples: int = 2, + resume: bool = True, + run_num: Optional[int] = None, + storage_backend: str = "hdf5", + vectorized: bool = False, + warmstart_max_tau: float = -1.0, + min_num_live_points: int = 400, + cluster_num_live_points: int = 40, + insertion_test_window: int = 10, + insertion_test_zscore_threshold: int = 2, + dlogz: float = 0.5, + dkl: float = 0.5, + frac_remain: float = 0.01, + lepsilon: float = 0.001, + min_ess: int = 400, + max_iters: Optional[int] = None, + max_ncalls: Optional[int] = None, + max_num_improvement_loops: float = -1.0, + log_interval: Optional[int] = None, + show_status: bool = True, + update_interval_ncall: Optional[int] = None, + update_interval_volume_fraction: float = 0.8, + viz_callback: str = "auto", + stepsampler_cls: Optional[str] = None, + nsteps: int = 25, + adaptive_nsteps: bool = False, + log: bool = False, + max_nsteps: int = 1000, + region_filter: bool = False, + scale: float = 1.0, + iterations_per_quick_update: int = None, + iterations_per_full_update: int = None, + number_of_cores: int = 1, + silence: bool = False, + session: Optional[sa.orm.Session] = None, + **kwargs + ): + """ + An UltraNest non-linear search. + + UltraNest is an optional requirement and must be installed manually via the command `pip install ultranest`. + It is optional as it has certain dependencies which are generally straight forward to install (e.g. Cython). + + For a full description of UltraNest and its Python wrapper PyUltraNest, checkout its Github and documentation + webpages: + + https://github.com/JohannesBuchner/UltraNest + https://johannesbuchner.github.io/UltraNest/readme.html + + Parameters + ---------- + name + The name of the search, controlling the last folder results are output. + path_prefix + The path of folders prefixing the name folder where results are output. + unique_tag + The name of a unique tag for this model-fit, which will be given a unique entry in the sqlite database + and also acts as the folder after the path prefix and before the search name. + iterations_per_full_update + The number of iterations performed between update (e.g. output latest model to hard-disk, visualization). + number_of_cores + The number of cores sampling is performed using a Python multiprocessing Pool instance. + silence + If True, the default print output of the non-linear search is silenced. + session + An SQLalchemy session instance so the results of the model-fit are written to an SQLite database. + """ + + super().__init__( + name=name, + path_prefix=path_prefix, + unique_tag=unique_tag, + iterations_per_quick_update=iterations_per_quick_update, + iterations_per_full_update=iterations_per_full_update, + number_of_cores=number_of_cores, + silence=silence, + session=session, + **kwargs + ) + + self.draw_multiple = draw_multiple + self.ndraw_min = ndraw_min + self.ndraw_max = ndraw_max + self.num_bootstraps = num_bootstraps + self.num_test_samples = num_test_samples + self.resume = resume + self.run_num = run_num + self.storage_backend = storage_backend + self.vectorized = vectorized + self.warmstart_max_tau = warmstart_max_tau + + self.min_num_live_points = min_num_live_points + self.cluster_num_live_points = cluster_num_live_points + self.insertion_test_window = insertion_test_window + self.insertion_test_zscore_threshold = insertion_test_zscore_threshold + self.dlogz = dlogz + self.dkl = dkl + self.frac_remain = frac_remain + self.lepsilon = lepsilon + self.min_ess = min_ess + self.max_iters = max_iters + self.max_ncalls = max_ncalls + self.max_num_improvement_loops = max_num_improvement_loops + self.log_interval = log_interval + self.show_status = show_status + self.update_interval_ncall = update_interval_ncall + self.update_interval_volume_fraction = update_interval_volume_fraction + self.viz_callback = viz_callback + + self.stepsampler_cls = stepsampler_cls + self.nsteps = nsteps if stepsampler_cls is not None else None + self.adaptive_nsteps = adaptive_nsteps + self.log_stepsampler = log + self.max_nsteps = max_nsteps + self.region_filter = region_filter + self.scale = scale + + if is_test_mode(): + self.apply_test_mode() + + self.logger.debug("Creating UltraNest Search") + + def apply_test_mode(self): + self.max_iters = 1 + self.max_ncalls = 1 + + @property + def search_kwargs(self): + """Build the kwargs dict passed to ``ReactiveNestedSampler``.""" + return { + "draw_multiple": self.draw_multiple, + "ndraw_min": self.ndraw_min, + "ndraw_max": self.ndraw_max, + "num_bootstraps": self.num_bootstraps, + "num_test_samples": self.num_test_samples, + "resume": self.resume, + "run_num": self.run_num, + "storage_backend": self.storage_backend, + "vectorized": self.vectorized, + "warmstart_max_tau": self.warmstart_max_tau, + } + + @property + def run_kwargs(self): + """Build the kwargs dict passed to ``sampler.run()``.""" + return { + "min_num_live_points": self.min_num_live_points, + "cluster_num_live_points": self.cluster_num_live_points, + "insertion_test_window": self.insertion_test_window, + "insertion_test_zscore_threshold": self.insertion_test_zscore_threshold, + "frac_remain": self.frac_remain, + "min_ess": self.min_ess, + "max_iters": self.max_iters, + "max_num_improvement_loops": self.max_num_improvement_loops, + "log_interval": self.log_interval, + "show_status": self.show_status, + "update_interval_volume_fraction": self.update_interval_volume_fraction, + "viz_callback": self.viz_callback, + } + + def _fit(self, model: AbstractPriorModel, analysis): + """ + Fit a model using the search and the Analysis class which contains the data and returns the log likelihood from + instances of the model, which the `NonLinearSearch` seeks to maximize. + + Parameters + ---------- + model : ModelMapper + The model which generates instances for different points in parameter space. + analysis : Analysis + Contains the data and the log likelihood function which fits an instance of the model to the data, returning + the log likelihood the `NonLinearSearch` maximizes. + + Returns + ------- + A result object comprising the Samples object that includes the maximum log likelihood instance and full + set of accepted ssamples of the fit. + """ + + try: + import ultranest + except ModuleNotFoundError: + raise ModuleNotFoundError( + "\n--------------------\n" + "You are attempting to perform a model-fit using UltraNest. \n\n" + "However, the optional library UltraNest (https://johannesbuchner.github.io/UltraNest/index.html) is " + "not installed.\n\n" + "Install it via the command `pip install ultranest==3.6.2`.\n\n" + "----------------------" + ) + + fitness = Fitness( + model=model, + analysis=analysis, + paths=self.paths, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99 + ) + + def prior_transform(cube): + return model.vector_from_unit_vector( + unit_vector=cube, + ) + + log_dir = self.paths.search_internal_path + + try: + checkpoint_exists = os.path.exists(log_dir / "chains") + except TypeError: + checkpoint_exists = False + + if checkpoint_exists: + self.logger.info( + "Resuming UltraNest non-linear search (previous samples found)." + ) + else: + self.logger.info( + "Starting new UltraNest non-linear search (no previous samples found)." + ) + + search_internal = ultranest.ReactiveNestedSampler( + param_names=model.parameter_names, + loglike=fitness.call_wrap, + transform=prior_transform, + log_dir=log_dir, + **self.search_kwargs + ) + + search_internal.stepsampler = self.stepsampler + + finished = False + + while not finished: + + try: + total_iterations = search_internal.ncall + except AttributeError: + total_iterations = 0 + + if self.max_ncalls is not None: + iterations = self.max_ncalls + else: + iterations = total_iterations + self.iterations_per_full_update + + if iterations > 0: + + run_kwargs = self.run_kwargs + run_kwargs["update_interval_ncall"] = iterations + + search_internal.run( + max_ncalls=iterations, + **run_kwargs + ) + + self.paths.save_search_internal( + obj=search_internal.results, + ) + + iterations_after_run = search_internal.ncall + + if ( + total_iterations == iterations_after_run + or iterations_after_run == self.max_ncalls + ): + finished = True + + if not finished: + + self.perform_update( + model=model, + analysis=analysis, + during_analysis=True, + fitness=fitness, + search_internal=search_internal + ) + + return search_internal, fitness + + def output_search_internal(self, search_internal): + """ + Output the sampler results to hard-disk in their internal format. + + UltraNest uses a backend to store and load results, therefore the outputting of the search internal to a + dill file is disabled. + + However, a dictionary of the search results is output to dill above. + + Parameters + ---------- + sampler + The nautilus sampler object containing the results of the model-fit. + """ + pass + + def samples_info_from(self, search_internal=None): + + search_internal = search_internal or self.paths.load_search_internal() + + return { + "log_evidence": search_internal["logz"], + "total_samples": search_internal["ncall"], + "total_accepted_samples": len(search_internal["weighted_samples"]["logl"]), + "time": self.timer.time if self.timer else None, + "number_live_points": self.min_num_live_points + } + + def samples_via_internal_from(self, model: AbstractPriorModel, search_internal=None): + """ + Returns a `Samples` object from the ultranest internal results. + + The samples contain all information on the parameter space sampling (e.g. the parameters, + log likelihoods, etc.). + + The internal search results are converted from the native format used by the search to lists of values + (e.g. `parameter_lists`, `log_likelihood_list`). + + Parameters + ---------- + model + Maps input vectors of unit parameter values to physical values and model instances via priors. + """ + + search_internal = search_internal.results or self.paths.load_search_internal() + + parameters = search_internal["weighted_samples"]["points"] + log_likelihood_list = search_internal["weighted_samples"]["logl"] + log_prior_list = [ + sum(model.log_prior_list_from_vector(vector=vector)) for vector in parameters + ] + weight_list = search_internal["weighted_samples"]["weights"] + + sample_list = Sample.from_lists( + model=model, + parameter_lists=parameters, + log_likelihood_list=log_likelihood_list, + log_prior_list=log_prior_list, + weight_list=weight_list + ) + + return SamplesNest( + model=model, + sample_list=sample_list, + samples_info=self.samples_info_from(search_internal=search_internal), + ) + + @property + def stepsampler(self): + + from ultranest import stepsampler + + stepsampler_cls = self.stepsampler_cls + + if stepsampler_cls is None: + return None + + stepsampler_kwargs = { + "nsteps": self.nsteps, + "adaptive_nsteps": self.adaptive_nsteps, + "log": self.log_stepsampler, + "max_nsteps": self.max_nsteps, + "region_filter": self.region_filter, + "scale": self.scale, + } + + if stepsampler_cls == "RegionMHSampler": + return stepsampler.RegionMHSampler(**stepsampler_kwargs) + elif stepsampler_cls == "AHARMSampler": + stepsampler_kwargs.pop("scale") + return stepsampler.AHARMSampler(**stepsampler_kwargs) + elif stepsampler_cls == "CubeMHSampler": + return stepsampler.CubeMHSampler(**stepsampler_kwargs) + elif stepsampler_cls == "CubeSliceSampler": + return stepsampler.CubeSliceSampler(**stepsampler_kwargs) + elif stepsampler_cls == "RegionSliceSampler": + return stepsampler.RegionSliceSampler(**stepsampler_kwargs)