Source code for sampler.metropolis

Metropolis algorithm module, wrapping the pymc3 implementation.

Provides the possibility to update the involved covariance matrixes within
the course of sampling the chain.

import logging
import os
import shutil
from copy import deepcopy
from time import time

import numpy as num
from pymc3.backends import text
from pymc3.model import Point, modelcontext
from pymc3.step_methods.metropolis import metrop_select
from pymc3.step_methods.metropolis import tune as step_tune
from pymc3.theanof import inputvars, make_shared_replacements
from pymc3.vartypes import discrete_types
from pyrocko import util
from theano import config as tconfig

from beat import backend, utility
from beat.covariance import init_proposal_covariance

from .base import (

__all__ = ["metropolis_sample", "get_trace_stats", "get_final_stage", "Metropolis"]

logger = logging.getLogger("metropolis")

[docs]class Metropolis(backend.ArrayStepSharedLLK): """ Metropolis-Hastings sampler Parameters ---------- vars : list List of variables for sampler out_vars : list List of output variables for trace recording. If empty unobserved_RVs are taken. n_chains : int Number of chains per stage has to be a large number of number of n_jobs (processors to be used) on the machine. scaling : float Factor applied to the proposal distribution i.e. the step size of the Markov Chain covariance : :class:`numpy.ndarray` (n_chains x n_chains) for MutlivariateNormal, otherwise (n_chains) Initial Covariance matrix for proposal distribution, if None - identity matrix taken likelihood_name : string name of the :class:`pymc3.determinsitic` variable that contains the model likelihood - defaults to 'like' backend : str type of backend to use for sample results storage, for alternatives see :class:`backend.backend:catalog` proposal_dist : :class:`pymc3.metropolis.Proposal` Type of proposal distribution, see :mod:`pymc3.step_methods.metropolis` for options tune : boolean Flag for adaptive scaling based on the acceptance rate model : :class:`pymc3.Model` Optional model for sampling step. Defaults to None (taken from context). """ default_blocked = True def __init__( self, vars=None, out_vars=None, covariance=None, scale=1.0, n_chains=100, tune=True, tune_interval=100, model=None, check_bound=True, likelihood_name="like", backend="csv", proposal_name="MultivariateNormal", **kwargs ): model = modelcontext(model) if vars is None: vars = model.vars vars = inputvars(vars) if out_vars is None: out_vars = model.unobserved_RVs out_varnames = [ for out_var in out_vars] self.scaling = utility.scalar2floatX(num.atleast_1d(scale)) self.tune = tune self.check_bound = check_bound self.tune_interval = tune_interval self.steps_until_tune = tune_interval self.stage_sample = 0 self.cumulative_samples = 0 self.accepted = 0 self.beta = 1.0 self.stage = 0 self.chain_index = 0 # needed to use the same parallel implementation function as for SMC self.resampling_indexes = num.arange(n_chains) self.n_chains = n_chains self.likelihood_name = likelihood_name self._llk_index = out_varnames.index(likelihood_name) self.backend = backend self.discrete = num.concatenate( [[v.dtype in discrete_types] * (v.dsize or 1) for v in vars] ) self.any_discrete = self.discrete.any() self.all_discrete = self.discrete.all() # create initial population self.population = [] self.array_population = num.zeros(n_chains) "Creating initial population for {}" " chains ...".format(self.n_chains) ) for i in range(self.n_chains): self.population.append( Point({ v.random() for v in vars}, model=model) ) self.population[0] = model.test_point shared = make_shared_replacements(vars, model) self.logp_forw = logp_forw(out_vars, vars, shared) self.check_bnd = logp_forw([model.varlogpt], vars, shared) super(Metropolis, self).__init__(vars, out_vars, shared) # init proposal if covariance is None and proposal_name in multivariate_proposals: t0 = time() self.covariance = init_proposal_covariance( bij=self.bij, vars=vars, model=model, pop_size=1000 ) t1 = time()"Time for proposal covariance init: %f" % (t1 - t0)) scale = self.covariance elif covariance is None: scale = num.ones(sum(v.dsize for v in vars)) else: scale = covariance self.proposal_name = proposal_name self.proposal_dist = choose_proposal(self.proposal_name, scale=scale) self.proposal_samples_array = self.proposal_dist(n_chains) self.chain_previous_lpoint = [[]] * self.n_chains self._tps = None def _sampler_state_blacklist(self): """ Returns sampler attributes that are not saved. """ bl = [ "check_bnd", "logp_forw", "proposal_samples_array", "vars", "bij", "lij", "ordering", "lordering", "_BlockedStep__newargs", ] return bl
[docs] def get_sampler_state(self): """ Return dictionary of sampler state. Returns ------- dict of sampler state """ blacklist = self._sampler_state_blacklist() return {k: v for k, v in self.__dict__.items() if k not in blacklist}
[docs] def apply_sampler_state(self, state): """ Update sampler state to given state (obtained by 'get_sampler_state') Parameters ---------- state : dict with sampler parameters """ for k, v in state.items(): setattr(self, k, v)
def time_per_sample(self, n_points=10): if not self._tps: tps = num.zeros((n_points)) for i in range(n_points): q =[i]) t0 = time() self.logp_forw(q) t1 = time() tps[i] = t1 - t0 self._tps = tps.mean() return self._tps def astep(self, q0): if self.stage == 0: l_new = self.logp_forw(q0) if not num.isfinite(l_new[self._llk_index]): raise ValueError( "Got NaN in likelihood evaluation! " "Invalid model definition? " "Or starting point outside prior bounds!" ) q_new = q0 else: if self.stage_sample == 0: self.proposal_samples_array = self.proposal_dist(self.n_steps).astype( tconfig.floatX ) if not self.steps_until_tune and self.tune: # Tune scaling parameter logger.debug( "Tuning: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) self.scaling = utility.scalar2floatX( step_tune(self.scaling, self.accepted / float(self.tune_interval)) ) # Reset counter self.steps_until_tune = self.tune_interval self.accepted = 0 logger.debug( "Get delta: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) delta = self.proposal_samples_array[self.stage_sample, :] * self.scaling if self.any_discrete: if self.all_discrete: delta = num.round(delta, 0) q0 = q0.astype(int) q = (q0 + delta).astype(int) else: delta[self.discrete] = num.round(delta[self.discrete], 0).astype( int ) q = q0 + delta q = q[self.discrete].astype(int) else: q = q0 + delta try: l0 = self.chain_previous_lpoint[self.chain_index] llk0 = l0[self._llk_index] except IndexError: l0 = self.logp_forw(q0) self.chain_previous_lpoint[self.chain_index] = l0 llk0 = l0[self._llk_index] if self.check_bound: logger.debug( "Checking bound: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) varlogp = self.check_bnd(q) if num.isfinite(varlogp): logger.debug( "Calc llk: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) lp = self.logp_forw(q) logger.debug( "Select llk: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) tempered_llk_ratio = self.beta * ( lp[self._llk_index] - l0[self._llk_index] ) q_new, accepted = metrop_select(tempered_llk_ratio, q, q0) if accepted: logger.debug( "Accepted: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) logger.debug( "proposed: %f previous: %f" % (lp[self._llk_index], llk0) ) self.accepted += 1 l_new = lp self.chain_previous_lpoint[self.chain_index] = l_new else: logger.debug( "Rejected: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) logger.debug( "proposed: %f previous: %f" % (lp[self._llk_index], l0[self._llk_index]) ) l_new = l0 else: q_new = q0 l_new = l0 else: logger.debug( "Calc llk: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) lp = self.logp_forw(q) logger.debug( "Select: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) q_new, accepted = metrop_select( self.beta * (lp[self._llk_index] - llk0), q, q0 ) if accepted: self.accepted += 1 l_new = lp self.chain_previous_lpoint[self.chain_index] = l_new else: l_new = l0 logger.debug( "Counters: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) self.steps_until_tune -= 1 self.stage_sample += 1 self.cumulative_samples += 1 # reset sample counter if self.stage_sample == self.n_steps: self.stage_sample = 0 logger.debug( "End step: Chain_%i step_%i" % (self.chain_index, self.stage_sample) ) return q_new, l_new
[docs]def get_final_stage(homepath, n_stages, model): """ Combine Metropolis results into final stage to get one single chain for plotting results. """ util.ensuredir(homepath) mtraces = [] for stage in range(n_stages):"Loading Metropolis stage %i" % stage) stage_outpath = os.path.join(homepath, "stage_%i" % stage) mtraces.append(backend.load(name=stage_outpath, model=model)) ctrace = backend.concatenate_traces(mtraces) outname = os.path.join(homepath, "stage_final") if os.path.exists(outname):"Removing existing previous final stage!") shutil.rmtree(outname) util.ensuredir(outname)"Creating final Metropolis stage") text.dump(name=outname, trace=ctrace)
[docs]def metropolis_sample( n_steps=10000, homepath=None, start=None, progressbar=False, rm_flag=False, buffer_size=5000, buffer_thinning=1, step=None, model=None, n_jobs=1, update=None, burn=0.5, thin=2, ): """ Execute Metropolis algorithm repeatedly depending on the number of chains. """ # hardcoded stage here as there are no stages stage = -1 model = modelcontext(model) step.n_steps = int(n_steps) if n_steps < 1: raise TypeError("Argument `n_steps` should be above 0.", exc_info=1) if step is None: raise TypeError("Argument `step` has to be a Metropolis step object.") if homepath is None: raise TypeError("Argument `homepath` should be path to result_directory.") if n_jobs > 1: if not (step.n_chains / float(n_jobs)).is_integer(): raise Exception("n_chains / n_jobs has to be a whole number!") if start is not None: if len(start) != step.n_chains: raise Exception( "Argument `start` should have dicts equal the " "number of chains (step.N-chains)" ) else: step.population = start if not any(step.likelihood_name in for var in model.deterministics): raise Exception( "Model (deterministic) variables need to contain " "a variable %s as defined in `step`." % step.likelihood_name ) stage_handler = backend.SampleStage(homepath, backend=step.backend) util.ensuredir(homepath) chains, step, update = init_stage( stage_handler=stage_handler, step=step, stage=stage, # needs zero otherwise tries to load stage_0 results progressbar=progressbar, update=update, model=model, rm_flag=rm_flag, ) with model: chains = stage_handler.clean_directory(step.stage, chains, rm_flag)"Sampling stage ...") draws = n_steps step.stage = stage sample_args = { "draws": draws, "step": step, "stage_path": stage_handler.stage_path(step.stage), "progressbar": progressbar, "model": model, "n_jobs": n_jobs, "buffer_size": buffer_size, "buffer_thinning": buffer_thinning, "chains": chains, } mtrace = iter_parallel_chains(**sample_args) if step.proposal_name == "MultivariateNormal": pdict, step.covariance = get_trace_stats(mtrace, step, burn, thin, n_jobs) step.proposal_dist = choose_proposal( step.proposal_name, scale=step.covariance ) if update is not None:"Updating Covariances ...") update.update_weights(pdict["dist_mean"], n_jobs=n_jobs) mtrace = update_last_samples( homepath, step, progressbar, model, n_jobs, rm_flag ) elif update is not None and stage == 0: update.engine.close_cashed_stores() outparam_list = [step.get_sampler_state(), update] stage_handler.dump_atmip_params(step.stage, outparam_list)
[docs]def get_trace_stats(mtrace, step, burn=0.5, thin=2, n_jobs=1): """ Get mean value of trace variables and return point. Parameters ---------- mtrace : :class:`pymc3.backends.base.MultiTrace` Multitrace sampling result step : initialised :class:`smc.SMC` sampler object burn : float Burn-in parameter to throw out samples from the beginning of the trace thin : int Thinning of samples in the trace Returns ------- dict with points, covariance matrix """ n_steps = len(mtrace) array_population = num.zeros( (n_jobs * int(num.ceil(n_steps * (1 - burn) / thin)), step.ordering.size) ) # collect end points of each chain and put into array for var, slc, shp, _ in step.ordering.vmap: samples = mtrace.get_values( varname=var, burn=int(burn * n_steps), thin=thin, combine=True ) if len(shp) == 0: array_population[:, slc] = num.atleast_2d(samples).T else: array_population[:, slc] = samples llks = mtrace.get_values( varname=step.likelihood_name, burn=int(burn * n_steps), thin=thin, combine=True ) posterior_idxs = utility.get_fit_indexes(llks) d = {} for k, v in posterior_idxs.items(): d[k] = step.bij.rmap(array_population[v, :]) d["dist_mean"] = step.bij.rmap(array_population.mean(axis=0)) avar = array_population.var(axis=0) if avar.sum() == 0.0: logger.warn("Trace std not valid not enough samples! Use 1.") avar = 1.0 cov = num.eye(step.ordering.size) * avar return d, cov