Skip to content

Feature/pool #100

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 2 commits into from
May 10, 2020
Merged

Feature/pool #100

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8419818
making pool branch
Jammy2211 May 8, 2020
fcf4fe5
parallelization of Dynesty implemented.
Jammy2211 May 9, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 51 additions & 9 deletions autofit/optimize/non_linear/emcee.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
import logging
import math
import os

import emcee
import numpy as np
import multiprocessing as mp

from autofit import conf
from autofit import exc
from autofit.text import samples_text
from autofit.optimize.non_linear import samples
Expand All @@ -30,6 +28,7 @@ def __init__(
auto_correlation_check_size=None,
auto_correlation_required_length=None,
auto_correlation_change_threshold=None,
number_of_cores=None,
):
"""
Class to setup and run an Emcee non-linear search.
Expand All @@ -40,11 +39,17 @@ def __init__(

https://emcee.readthedocs.io/en/stable/

**PyAutoFit** extends **emcee** by providing an option to check the auto-correlation length of the samples
during the run and terminating sampling early if these meet a specified threshold. See this page
(https://emcee.readthedocs.io/en/stable/tutorials/autocorr/#autocorr) for a description of how this is implemented.
Extensions:

**PyAutoFit** provides the option to check the auto-correlation length of the samples during the run and
terminating sampling early if these meet a specified threshold. See this page
(https://emcee.readthedocs.io/en/stable/tutorials/autocorr/#autocorr) for a description.

**PyAutoFit** also provides different options for walker initialization, with the default 'ball' method
starting all walkers close to one another in parameter space, as recommended in the Emcee documentation
(https://emcee.readthedocs.io/en/stable/user/faq/).

If you use *emcee* as part of a published work, please cite the package following the instructions under the
If you use *Emcee* as part of a published work, please cite the package following the instructions under the
*Attribution* section of the GitHub page.

Parameters
Expand All @@ -65,6 +70,20 @@ def __init__(
Whether the auto-correlation lengths of the MCMC samples should be checked to determine the stopping
criteria. If *True*, this option may terminate the Emcee run before the input number of steps, nsteps, has
been performed. If *False* nstep samples will be taken.
initialize_method : str
The method used to generate where walkers are initialized in parameter space, with options:
ball (default):
Walkers are initialized by randomly drawing unit values from a uniform distribution between the
initialize_ball_lower_limit and initialize_ball_upper_limit values. It is recommended these limits are
small, such that all walkers begin close to one another.
prior:
Walkers are initialized by randomly drawing unit values from a uniform distribution between 0 and 1,
thus being fully distributed over the prior.
initialize_ball_upper_limit : float
The lower limit of the uniform distribution unit values are drawn from when initializing walkers using the
ball method.
The upper limit of the uniform distribution unit values are drawn from when initializing walkers using the
ball method.
auto_correlation_check_size : int
The length of the samples used to check the auto-correlation lengths (from the latest sample backwards). For
convergence, the auto-correlations must not change over a certain range of samples. A longer check-size
Expand All @@ -76,11 +95,20 @@ def __init__(
auto_correlation_change_threshold : float
The threshold value by which if the change in auto_correlations is below sampling will be terminated early,
as it has been determined as converged.
number_of_cores : int
The number of cores Emcee sampling is performed using a Python multiprocessing Pool instance. If 1, a
pool instance is not created and the job runs in serial.

All remaining attributes are emcee parameters and described at the emcee API webpage:

https://emcee.readthedocs.io/en/stable/

Attributes
----------
sigma : float
The error-bound value that linked Gaussian prior withs are computed using. For example, if sigma=3.0,
parameters will use Gaussian Priors with widths coresponding to errors estimated at 3 sigma confidence.

"""

if paths is None:
Expand Down Expand Up @@ -130,6 +158,12 @@ def __init__(
else auto_correlation_change_threshold
)

self.number_of_cores = (
self.config("parallel", "number_of_cores", int)
if number_of_cores is None
else number_of_cores
)

logger.debug("Creating Emcee NLO")

def _fit(self, model, analysis):
Expand Down Expand Up @@ -161,10 +195,14 @@ def copy_with_name_extension(self, extension, remove_phase_tag=False):
copy.sigma = self.sigma
copy.nwalkers = self.nwalkers
copy.nsteps = self.nsteps
copy.initialize_method = self.initialize_method
copy.initialize_ball_lower_limit = self.initialize_ball_lower_limit
copy.initialize_ball_upper_limit = self.initialize_ball_upper_limit
copy.auto_correlation_check_for_convergence = self.auto_correlation_check_for_convergence
copy.auto_correlation_check_size = self.auto_correlation_check_size
copy.auto_correlation_required_length = self.auto_correlation_required_length
copy.auto_correlation_change_threshold = self.auto_correlation_change_threshold
copy.number_of_cores = self.number_of_cores

return copy

Expand All @@ -185,15 +223,18 @@ def __call__(self, params):

def _full_fit(self, model, analysis):

pool, pool_ids = self.make_pool()

fitness_function = self.fitness_function_from_model_and_analysis(
model=model, analysis=analysis
model=model, analysis=analysis, pool_ids=pool_ids,
)

emcee_sampler = emcee.EnsembleSampler(
nwalkers=self.nwalkers,
ndim=model.prior_count,
log_prob_fn=fitness_function.__call__,
backend=emcee.backends.HDFBackend(filename=self.paths.path + "/emcee.hdf"),
pool=pool,
)

try:
Expand Down Expand Up @@ -279,13 +320,14 @@ def backend(self) -> emcee.backends.HDFBackend:
"The file emcee.hdf does not exist at the path " + self.paths.path
)

def fitness_function_from_model_and_analysis(self, model, analysis):
def fitness_function_from_model_and_analysis(self, model, analysis, pool_ids=None):

return Emcee.Fitness(
paths=self.paths,
model=model,
analysis=analysis,
samples_from_model=self.samples_from_model,
pool_ids=pool_ids,
)

def samples_from_model(self, model):
Expand Down
35 changes: 26 additions & 9 deletions autofit/optimize/non_linear/nested_sampling/dynesty.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ def __init__(
self,
paths=None,
sigma=3,
iterations_per_update=None,
bound=None,
sample=None,
bootstrap=None,
Expand All @@ -37,6 +36,8 @@ def __init__(
max_move=None,
terminate_at_acceptance_ratio=None,
acceptance_ratio_threshold=None,
iterations_per_update=None,
number_of_cores=None,
):
"""
Class to setup and run a Dynesty non-linear search.
Expand Down Expand Up @@ -69,6 +70,9 @@ def __init__(
acceptance_ratio_threshold : float
The acceptance ratio threshold below which sampling terminates if *terminate_at_acceptance_ratio* is
*True* (see *NestedSampler* for a full description of this feature).
number_of_cores : int
The number of cores Emcee sampling is performed using a Python multiprocessing Pool instance. If 1, a
pool instance is not created and the job runs in serial.

Attributes
----------
Expand Down Expand Up @@ -125,6 +129,12 @@ def __init__(
else iterations_per_update
)

self.number_of_cores = (
self.config("parallel", "number_of_cores", int)
if number_of_cores is None
else number_of_cores
)

logger.debug("Creating DynestyStatic NLO")

def copy_with_name_extension(self, extension, remove_phase_tag=False):
Expand All @@ -147,6 +157,7 @@ def copy_with_name_extension(self, extension, remove_phase_tag=False):
copy.slices = self.slices
copy.fmove = self.fmove
copy.max_move = self.max_move
copy.number_of_cores = self.number_of_cores

return copy

Expand All @@ -171,6 +182,8 @@ def _fit(self, model: AbstractPriorModel, analysis) -> Result:
of the full samples used by the fit.
"""

pool, pool_ids = self.make_pool()

fitness_function = self.fitness_function_from_model_and_analysis(
model=model, analysis=analysis
)
Expand All @@ -193,12 +206,12 @@ def _fit(self, model: AbstractPriorModel, analysis) -> Result:
# These hacks are necessary to be able to pickle the sampler.

sampler.rstate = np.random
sampler.pool = None
sampler.M = map
sampler.pool = pool

# pool = Pool(processes=1)
# sampler.pool = None
# sampler.M = pool.map
if self.number_of_cores == 1:
sampler.M = map
else:
sampler.M = pool.map

dynesty_finished = False

Expand Down Expand Up @@ -281,7 +294,6 @@ def __init__(
self,
paths=None,
sigma=3,
iterations_per_update=None,
n_live_points=None,
bound=None,
sample=None,
Expand All @@ -297,6 +309,8 @@ def __init__(
max_move=None,
terminate_at_acceptance_ratio=None,
acceptance_ratio_threshold=None,
iterations_per_update=None,
number_of_cores=None,
):
"""
Class to setup and run a Dynesty non-linear search, using the static Dynesty nested sampler described at this
Expand Down Expand Up @@ -334,6 +348,7 @@ def __init__(
max_move=max_move,
terminate_at_acceptance_ratio=terminate_at_acceptance_ratio,
acceptance_ratio_threshold=acceptance_ratio_threshold,
number_of_cores=number_of_cores,
)

self.n_live_points = (
Expand Down Expand Up @@ -401,6 +416,7 @@ def __init__(
max_move=None,
terminate_at_acceptance_ratio=None,
acceptance_ratio_threshold=None,
number_of_cores=None,
):
"""
Class to setup and run a Dynesty non-linear search, using the dynamic Dynesty nested sampler described at this
Expand Down Expand Up @@ -438,6 +454,7 @@ def __init__(
max_move=max_move,
terminate_at_acceptance_ratio=terminate_at_acceptance_ratio,
acceptance_ratio_threshold=acceptance_ratio_threshold,
number_of_cores=number_of_cores,
)

logger.debug("Creating DynestyDynamic NLO")
Expand All @@ -446,8 +463,8 @@ def sampler_fom_model_and_fitness(self, model, fitness_function):
"""Get the dynamic Dynesty sampler which performs the non-linear search, passing it all associated input Dynesty
variables."""
return DynamicNestedSampler(
loglikelihood=fitness,
prior_transform=prior,
loglikelihood=fitness_function,
prior_transform=nl.NonLinearOptimizer.Fitness.prior,
ndim=model.prior_count,
logl_args=[model, fitness_function],
ptform_args=[model],
Expand Down
Loading