Ensemble
The eryn.ensemble.Ensemble
is the main driver and container for MCMC runs.
- class eryn.ensemble.EnsembleSampler(nwalkers, ndims, log_like_fn, priors, provide_groups=False, provide_supplemental=False, tempering_kwargs={}, branch_names=None, nbranches=1, nleaves_max=1, nleaves_min=0, pool=None, moves=None, rj_moves=None, dr_moves=None, dr_max_iter=5, args=None, kwargs=None, backend=None, vectorize=False, blobs_dtype=None, plot_iterations=-1, plot_generator=None, plot_name=None, periodic=None, update_fn=None, update_iterations=-1, stopping_fn=None, stopping_iterations=-1, fill_zero_leaves_val=-1e+300, num_repeats_in_model=1, num_repeats_rj=1, track_moves=True, info={})
Bases:
object
An ensemble MCMC sampler
The class controls the entire sampling run. It can handle everything from a basic non-tempered MCMC to a parallel-tempered, global fit containing multiple branches (models) and a variable number of leaves (sources) per branch. See here for a basic explainer.
Parameters related to parallelization can be controlled via the
pool
argument.- Parameters:
nwalkers (int) – The number of walkers in the ensemble per temperature.
ndims (int, list of ints, or dict) – The number of dimensions for each branch. If
dict
, keys should be the branch names and values the associated dimensionality.log_like_fn (callable) –
A function that returns the natural logarithm of the likelihood for that position. The inputs to
log_like_fn
depend on whether the function is vectorized (kwargvectorize
below), if you are using reversible jump, and how many branches you have.In the simplest case where
vectorize == False
, no reversible jump, and only one type of model, the inputs are just the array of parameters for one walker, so shape is(ndim,)
.If
vectorize == True
, no reversible jimp, and only one type of model, the inputs will be a 2D array of parameters of all the walkers going in. Shape:(num positions, ndim)
.If using reversible jump, the leaves that go together in the same Likelihood will be grouped together into a single function call. If
vectorize == False
, then each group is sent as an individual computation. WithN
different branches (N > 1
), inputs would be a list of 2D arrays of the coordinates for all leaves within each branch:([x0, x1,...,xN])
wherexi
is 2D with shape(number of leaves in this branch, ndim)
. IfN == 1
, then a list is not provided, just x0, the 2D array of coordinates for the one branch considered.If using reversible jump and
vectorize == True
, then the arrays of parameters will be output with information as regards the grouping of branch and leaf set. Inputs will be([X0, X1,..XN], [group0, group1,...,groupN])
whereXi
is a 2D array of all leaves in the sampler for branchi
.groupi
is an index indicating which unique group that sources belongs. For example, if we have 3 walkers with (1, 2, 1) leaves for modeli
, respectively, we wil have anXi = array([params0, params1, params2, params3])
andgroupsi = array([0, 1, 1, 2])
. IfN == 1
, then the lists are removed and the inputs become(X0, group0)
.Extra
args
andkwargs
for the Likelihood function can be added with the kwargsargs
andkwargs
below.Please see the tutorial for more information.
priors (dict) – The prior dictionary can take four forms. 1) A dictionary with keys as int or tuple containing the int or tuple of int that describe the parameter number over which to assess the prior, and values that are prior probability distributions that must have a
logpdf
class method. 2) Aeryn.prior.ProbDistContainer
object. 3) A dictionary with keys that arebranch_names
and values that are dictionaries for each branch as described for (1). 4) A dictionary with keys that arebranch_names
and values areeryn.prior.ProbDistContainer
objects. If the priors dictionary has the specific key"all_models_together"
in it, then a special prior must be input by the user that can produce the priorlogpdf
information that can depend on all of the models together rather than handling them separately as is the default. In this case, the user must input as the item attached to this special key a class object with alogpdf
method that takes as input two arguments:(coords, inds)
, which are the coordinate and index dictionaries across all models with shapes of(ntemps, nwalkers, nleaves_max, ndim)
and(ntemps, nwalkers, nleaves_max)
for each individual model, respectively. This function must then return the numpy array of logpdf values for the prior value with shape(ntemps, nwalkers)
.provide_groups (bool, optional) – If
True
, provide groups as described inlog_like_fn
above. A group parameter is added for each branch. (default:False
)provide_supplemental (bool, optional) –
If
True
, it will provide keyword arguments to the Likelihood function:supps
andbranch_supps
. Please see the Tutorial anderyn.state.BranchSupplemental
for more information.tempering_kwargs (dict, optional) – Keyword arguments for initialization of the tempering class:
eryn.moves.tempering.TemperatureControl
. (default:{}
)branch_names (list, optional) – List of branch names. If
None
, models will be assigned names asf"model_{index}"
based onnbranches
. (default:None
)nbranches (int, optional) – Number of branches (models) tested. Only used if
branch_names is None
. (default:1
)nleaves_max (int, list of ints, or dict, optional) – Maximum allowable leaf count for each branch. It should have the same length as the number of branches. If
dict
, keys should be the branch names and values the associated maximal leaf value. (default:1
)nleaves_min (int, list of ints, or dict, optional) – Minimum allowable leaf count for each branch. It should have the same length as the number of branches. Only used with Reversible Jump. If
dict
, keys should be the branch names and values the associated maximal leaf value. IfNone
and using Reversible Jump, will fill all branches with zero. (default:0
)pool (object, optional) – An object with a
map
method that follows the same calling sequence as the built-inmap
function. This is generally used to compute the log-probabilities for the ensemble in parallel.moves (list or object, optional) – This can be a single move object, a list of moves, or a “weighted” list of the form
[(eryn.moves.StretchMove(), 0.1), ...]
. When running, the sampler will randomly select a move from this list (optionally with weights) for each proposal. IfNone
, the default will beStretchMove
. (default:None
)rj_moves (list or object or bool or str, optional) – If
None
orFalse
, reversible jump will not be included in the run. This can be a single move object, a list of moves, or a “weighted” list of the form[(eryn.moves.DistributionGenerateRJ(), 0.1), ...]
. When running, the sampler will randomly select a move from this list (optionally with weights) for each proposal. IfTrue
, it defaults toDistributionGenerateRJ
. When running just theDistributionGenerateRJ
with multiple branches, it will propose changes to all branhces simultaneously. When running with more than one branch, useful options forrj_moves
are"iterate_branches"
,"separate_branches"
, or"together"
. Ifrj_moves == "iterate_branches"
, sample one branch by one branch in order of the branch names. This occurs within one RJ proposal, for each RJ proposal. Ifrj_moves == "separate_branches"
, there will be one RJ move per branch. During each individual RJ move, one of these proposals is chosen at random with equal propability. This is generally recommended when using multiple branches. Ifrj_moves == "together"
, this is equivalent torj_moves == True
. (default:None
)dr_moves (bool, optional) – If
None
otFalse
, delayed rejection when proposing “birth” of new components/models will be switched off for this run. Requiresrj_moves
set toTrue
. Not implemented yet. Working on it. (default:None
)dr_max_iter (int, optional) – Maximum number of iterations used with delayed rejection. (default: 5)
args (optional) – A list of extra positional arguments for
log_like_fn
.log_like_fn
will be called aslog_like_fn(sampler added args, *args, sampler added kwargs, **kwargs)
.kwargs (optional) – A dict of extra keyword arguments for
log_like_fn
.log_like_fn
will be called aslog_like_fn(sampler added args, *args, sampler added kwargs, **kwargs)
.backend (optional) – Either a
backends.Backend
or a subclass (likebackends.HDFBackend
) that is used to store and serialize the state of the chain. By default, the chain is stored as a set of numpy arrays in memory, but new backends can be written to support other mediums.vectorize (bool, optional) – If
True
,log_like_fn
is expected to accept an array of position vectors instead of just one. Note thatpool
will be ignored if this isTrue
. Seelog_like_fn
information above to understand the arguments oflog_like_fn
based on whethervectorize
isTrue
. (default:False
)periodic (dict, optional) – Keys are
branch_names
. Values are dictionaries that have (key: value) pairs as (index to parameter: period). Periodic parameters are treated as having periodic boundary conditions in proposals.update_fn (callable, optional) –
eryn.utils.updates.AdjustStretchProposalScale
object that allows the user to update the sampler in any preferred way everyupdate_iterations
sampler iterations. The callable must have signature:(sampler iteration, last sample state object, EnsembleSampler object)
.update_iterations (int, optional) – Number of iterations between sampler updates using
update_fn
. Updates are only performed at the thinning rate. Ifthin_by>1
whenEnsembleSampler.run_mcmc()
is used, the sampler is updated everythin_by * update_iterations
iterations.stopping_fn (callable, optional) –
eryn.utils.stopping.Stopping
object that allows the user to end the sampler if specified criteria are met. The callable must have signature:(sampler iteration, last sample state object, EnsembleSampler object)
.stopping_iterations (int, optional) – Number of iterations between sampler attempts to evaluate the
stopping_fn
. Stopping checks are only performed at the thinning rate. Ifthin_by>1
whenEnsembleSampler.run_mcmc()
is used, the sampler is checked for the stopping criterion everythin_by * stopping_iterations
iterations.fill_zero_leaves_val (double, optional) – When there are zero leaves in a given walker (across all branches), fill the likelihood value with
fill_zero_leaves_val
. If wanting to keep zero leaves as a possible model, this should be set to the value of the contribution to the Likelihood from the data. (Default:-1e300
).num_repeats_in_model (int, optional) – Number of times to repeat the in-model step within in one sampler iteration. When analyzing the acceptance fraction, you must include the value of
num_repeats_in_model
to get the proper denominator.num_repeats_rj (int, optional) – Number of time to repeat the reversible jump step within in one sampler iteration. When analyzing the acceptance fraction, you must include the value of
num_repeats_rj
to get the proper denominator.track_moves (bool, optional) – If
True
, track acceptance fraction of each move in the backend. IfFalse
, no tracking is done. IfTrue
and run is interrupted, it will check that the move configuration has not changed. It will not allow the run to go on if it is changed. In this case, the user should declare a new backend and use the last state from the previous backend. Warning: If the order of moves of the same move class is changed, the check may not catch it, so the tracking may mix move acceptance fractions together.info (dict, optional) – Key and value pairs reprenting any information the user wants to add to the backend if the user is not inputing their own backend.
- Raises:
ValueError – Any startup issues.
- property random_state
The state of the internal random number generator. In practice, it’s the result of calling
get_state()
on anumpy.random.mtrand.RandomState
object. You can try to set this property but be warned that if you do this and it fails, it will do so silently.
- property priors
Return the priors in the sampler.
- reset(**info)
Reset the backend.
- Parameters:
**info (dict, optional) – information to pass to backend reset method.
- get_model()
Get
Model
object from samplerThe model object is used to pass necessary information to the proposals. This method can be used to retrieve the
model
used in the sampler from outside the sampler.- Returns:
Model
object used by sampler.- Return type:
Model
- sample(initial_state, iterations=1, tune=False, skip_initial_state_check=True, thin_by=1, store=True, progress=False)
Advance the chain as a generator
- Parameters:
initial_state (State or ndarray[ntemps, nwalkers, nleaves_max, ndim] or dict) – The initial
State
or positions of the walkers in the parameter space. If multiple branches used, must be dict with keys as thebranch_names
and values as the positions. Ifbetas
are provided in the state object, they will be loaded into thetemperature_control
.iterations (int or None, optional) – The number of steps to generate.
None
generates an infinite stream (requiresstore=False
). (default: 1)tune (bool, optional) – If
True
, the parameters of some moves will be automatically tuned. (default:False
)thin_by (int, optional) – If you only want to store and yield every
thin_by
samples in the chain, setthin_by
to an integer greater than 1. When this is set,iterations * thin_by
proposals will be made. (default: 1)store (bool, optional) – By default, the sampler stores in the backend the positions (and other information) of the samples in the chain. If you are using another method to store the samples to a file or if you don’t need to analyze the samples after the fact (for burn-in for example) set
store
toFalse
. (default:True
)progress (bool or str, optional) – If
True
, a progress bar will be shown as the sampler progresses. If a string, will select a specifictqdm
progress bar - most notable is'notebook'
, which shows a progress bar suitable for Jupyter notebooks. IfFalse
, no progress bar will be shown. (default:False
)skip_initial_state_check (bool, optional) – If
True
, a check that the initial_state can fully explore the space will be skipped. If using reversible jump, the user needs to ensure this on their own (skip_initial_state_check``is set to ``False
in this case. (default:True
)
- Returns:
Every
thin_by
steps, this generator yields theState
of the ensemble.- Return type:
- Raises:
ValueError – Improper initialization.
- run_mcmc(initial_state, nsteps, burn=None, post_burn_update=False, **kwargs)
Iterate
sample()
fornsteps
iterations and return the result.- Parameters:
initial_state (State or ndarray[ntemps, nwalkers, nleaves_max, ndim] or dict) – The initial
State
or positions of the walkers in the parameter space. If multiple branches used, must be dict with keys as thebranch_names
and values as the positions. Ifbetas
are provided in the state object, they will be loaded into thetemperature_control
.nsteps (int) – The number of steps to generate. The total number of proposals is
nsteps * thin_by
.burn (int, optional) – Number of burn steps to run before storing information. The
thin_by
kwarg is ignored when counting burn steps since there is no storage (equivalent tothin_by=1
).post_burn_update (bool, optional) – If
True
, runupdate_fn
after burn in.
Other parameters are directly passed to
sample()
.
- compute_log_prior(coords, inds=None, supps=None, branch_supps=None)
Calculate the vector of log-prior for the walkers
- Parameters:
coords (dict) – Keys are
branch_names
and values are the position np.arrays[ntemps, nwalkers, nleaves_max, ndim]. This dictionary is created with thebranches_coords
attribute fromState
.inds (dict, optional) – Keys are
branch_names
and values are theinds
np.arrays[ntemps, nwalkers, nleaves_max] that indicates which leaves are being used. This dictionary is created with thebranches_inds
attribute fromState
. (default:None
)
- Returns:
Prior Values
- Return type:
np.ndarray[ntemps, nwalkers]
- compute_log_like(coords, inds=None, logp=None, supps=None, branch_supps=None)
Calculate the vector of log-likelihood for the walkers
- Parameters:
coords (dict) – Keys are
branch_names
and values are the position np.arrays[ntemps, nwalkers, nleaves_max, ndim]. This dictionary is created with thebranches_coords
attribute fromState
.inds (dict, optional) – Keys are
branch_names
and values are the inds np.arrays[ntemps, nwalkers, nleaves_max] that indicates which leaves are being used. This dictionary is created with thebranches_inds
attribute fromState
. (default:None
)logp (np.ndarray[ntemps, nwalkers], optional) – Log prior values associated with all walkers. If not provided, it will be calculated because if a walker has logp = -inf, its likelihood is not calculated. This prevents evaluting likelihood outside the prior. (default:
None
)
- Returns:
- Carries log-likelihood and blob information.
First entry is np.ndarray[ntemps, nwalkers] with values corresponding to the log likelihood of each walker. Second entry is
blobs
.
- Return type:
tuple
- Raises:
ValueError: Infinite or NaN values in parameters.
- property acceptance_fraction
The fraction of proposed steps that were accepted
- property rj_acceptance_fraction
The fraction of proposed reversible jump steps that were accepted
- property swap_acceptance_fraction
The fraction of proposed temperature swaps that were accepted
- get_chain(**kwargs)
Get the stored chain of MCMC samples
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
- MCMC samples
The dictionary contains np.ndarrays of samples across the branches.
- Return type:
dict
- get_blobs(**kwargs)
Get the chain of blobs for each sample in the chain
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
The chain of blobs.
- Return type:
double np.ndarray[nsteps, ntemps, nwalkers, nblobs]
- get_log_like(**kwargs)
Get the chain of log Prior evaluated at the MCMC samples
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
The chain of log prior values.
- Return type:
double np.ndarray[nsteps, ntemps, nwalkers]
- get_log_prior(**kwargs)
Get the chain of log Prior evaluated at the MCMC samples
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
The chain of log prior values.
- Return type:
double np.ndarray[nsteps, ntemps, nwalkers]
- get_log_posterior(**kwargs)
Get the chain of log posterior values evaluated at the MCMC samples
- Parameters:
temper (bool, optional) – Apply tempering to the posterior values. (default:
False
)thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
The chain of log prior values.
- Return type:
double np.ndarray[nsteps, ntemps, nwalkers]
- get_inds(**kwargs)
Get the stored chain of MCMC samples
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
- The
inds
associated with the MCMC samples. The dictionary contains np.ndarrays of
inds
across the branches indicated which leaves were used at each step.
- The
- Return type:
dict
- get_nleaves(**kwargs)
Get the number of leaves for each walker
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
- nleaves on each branch.
- The number of leaves on each branch associated with the MCMC samples
within each branch.
- Return type:
dict
- get_last_sample(**kwargs)
Access the most recent sample in the chain
- Returns:
eryn.state.State
object containing the last sample from the chain.- Return type:
- get_betas(**kwargs)
Get the chain of inverse temperatures
- Parameters:
thin (int, optional) – Take only every
thin
steps from the chain. (default:1
)discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)slice_vals (indexing np.ndarray or slice, optional) – This is only available in
eryn.backends.hdfbackend
. If provided, slice the array directly from the HDF5 file with slice =slice_vals
.thin
anddiscard
will be ignored if slice_vals is notNone
. This is particularly useful if files are very large and the user only wants a small subset of the overall array. (default:None
)
- Returns:
The chain of temperatures.
- Return type:
double np.ndarray[nsteps, ntemps]
- get_value(name, **kwargs)
Get a specific value
- get_autocorr_time(**kwargs)
Compute an estimate of the autocorrelation time for each parameter
- Parameters:
discard (int, optional) – Discard the first
discard
steps in the chain as burn-in. (default:0
)thin (int, optional) – Use only every
thin
steps from the chain. The returned estimate is multiplied bythin
so the estimated time is in units of steps, not thinned steps. (default:1
)all_temps (bool, optional) – If True, calculate autocorrelation across all temperatures. If False, calculate autocorrelation across the minumum temperature chain (usually
T=1
). (default:False
)multiply_thin (bool, optional) – into the autocorrelation length. (default:
True
)
Other arguments are passed directly to
emcee.autocorr.integrated_time()
.- Returns:
- autocorrelation times
The dictionary contains autocorrelation times for all parameters as 1D double np.ndarrays as values with associated
branch_names
as keys.
- Return type:
dict