def init_nuts(init='advi', n_init=500000, model=None, **kwargs):
"""Initialize and sample from posterior of a continuous model.
This is a convenience function. NUTS convergence and sampling speed is extremely
dependent on the choice of mass/scaling matrix. In our experience, using ADVI
to estimate a diagonal covariance matrix and using this as the scaling matrix
produces robust results over a wide class of continuous models.
Parameters
----------
init : str {'advi', 'advi_map', 'map', 'nuts'}
Initialization method to use.
* advi : Run ADVI to estimate posterior mean and diagonal covariance matrix.
* advi_map: Initialize ADVI with MAP and use MAP as starting point.
* map : Use the MAP as starting point.
* nuts : Run NUTS and estimate posterior mean and covariance matrix.
n_init : int
Number of iterations of initializer
If 'advi', number of iterations, if 'metropolis', number of draws.
model : Model (optional if in `with` context)
**kwargs : keyword arguments
Extra keyword arguments are forwarded to pymc3.NUTS.
Returns
-------
start, nuts_sampler
start : pymc3.model.Point
Starting point for sampler
nuts_sampler : pymc3.step_methods.NUTS
Instantiated and initialized NUTS sampler object
"""
model = pm.modelcontext(model)
pm._log.info('Initializing NUTS using {}...'.format(init))
if init == 'advi':
v_params = pm.variational.advi(n=n_init)
start = pm.variational.sample_vp(v_params, 1, progressbar=False)[0]
cov = np.power(model.dict_to_array(v_params.stds), 2)
elif init == 'advi_map':
start = pm.find_MAP()
v_params = pm.variational.advi(n=n_init, start=start)
cov = np.power(model.dict_to_array(v_params.stds), 2)
elif init == 'map':
start = pm.find_MAP()
cov = pm.find_hessian(point=start)
elif init == 'nuts':
init_trace = pm.sample(step=pm.NUTS(), draws=n_init)
cov = pm.trace_cov(init_trace[n_init//2:])
start = {varname: np.mean(init_trace[varname]) for varname in init_trace.varnames}
else:
raise NotImplemented('Initializer {} is not supported.'.format(init))
step = pm.NUTS(scaling=cov, is_cov=True, **kwargs)
return start, step
def test_hh_flow():
cov = pm.floatX([[2, -1], [-1, 3]])
with pm.Model():
pm.MvNormal('mvN', mu=pm.floatX([0, 1]), cov=cov, shape=2)
nf = NFVI('scale-hh*2-loc')
nf.fit(25000, obj_optimizer=pm.adam(learning_rate=0.001))
trace = nf.approx.sample(10000)
cov2 = pm.trace_cov(trace)
np.testing.assert_allclose(cov, cov2, rtol=0.07)
def test_hh_flow():
cov = pm.floatX([[2, -1], [-1, 3]])
with pm.Model():
pm.MvNormal('mvN', mu=pm.floatX([0, 1]), cov=cov, shape=2)
nf = NFVI('scale-hh*2-loc')
nf.fit(25000, obj_optimizer=pm.adam(learning_rate=0.001))
trace = nf.approx.sample(10000)
cov2 = pm.trace_cov(trace)
np.testing.assert_allclose(cov, cov2, rtol=0.07)
def init_nuts(init='advi', n_init=500000, model=None):
"""Initialize and sample from posterior of a continuous model.
This is a convenience function. NUTS convergence and sampling speed is extremely
dependent on the choice of mass/scaling matrix. In our experience, using ADVI
to estimate a diagonal covariance matrix and using this as the scaling matrix
produces robust results over a wide class of continuous models.
Parameters
----------
init : str {'advi', 'advi_map', 'map', 'nuts'}
Initialization method to use.
* advi : Run ADVI to estimate posterior mean and diagonal covariance matrix.
* advi_map: Initialize ADVI with MAP and use MAP as starting point.
* map : Use the MAP as starting point.
* nuts : Run NUTS and estimate posterior mean and covariance matrix.
n_init : int
Number of iterations of initializer
If 'advi', number of iterations, if 'metropolis', number of draws.
model : Model (optional if in `with` context)
Returns
-------
start, nuts_sampler
start : pymc3.model.Point
Starting point for sampler
nuts_sampler : pymc3.step_methods.NUTS
Instantiated and initialized NUTS sampler object
"""
model = pm.modelcontext(model)
pm._log.info('Initializing NUTS using {}...'.format(init))
if init == 'advi':
v_params = pm.variational.advi(n=n_init)
start = pm.variational.sample_vp(v_params, 1, progressbar=False)[0]
cov = np.power(model.dict_to_array(v_params.stds), 2)
elif init == 'advi_map':
start = pm.find_MAP()
v_params = pm.variational.advi(n=n_init, start=start)
cov = np.power(model.dict_to_array(v_params.stds), 2)
elif init == 'map':
start = pm.find_MAP()
cov = pm.find_hessian(point=start)
elif init == 'nuts':
init_trace = pm.sample(step=pm.NUTS(), draws=n_init)
cov = pm.trace_cov(init_trace[n_init // 2:])
start = {
varname: np.mean(init_trace[varname])
for varname in init_trace.varnames
}
else:
raise NotImplemented('Initializer {} is not supported.'.format(init))
step = pm.NUTS(scaling=cov, is_cov=True)
return start, step
def init_nuts(init='ADVI',
njobs=1,
n_init=500000,
model=None,
random_seed=-1,
**kwargs):
"""Initialize and sample from posterior of a continuous model.
This is a convenience function. NUTS convergence and sampling speed is extremely
dependent on the choice of mass/scaling matrix. In our experience, using ADVI
to estimate a diagonal covariance matrix and using this as the scaling matrix
produces robust results over a wide class of continuous models.
Parameters
----------
init : str {'ADVI', 'ADVI_MAP', 'MAP', 'NUTS'}
Initialization method to use.
* ADVI : Run ADVI to estimate posterior mean and diagonal covariance matrix.
* ADVI_MAP: Initialize ADVI with MAP and use MAP as starting point.
* MAP : Use the MAP as starting point.
* NUTS : Run NUTS and estimate posterior mean and covariance matrix.
njobs : int
Number of parallel jobs to start.
n_init : int
Number of iterations of initializer
If 'ADVI', number of iterations, if 'metropolis', number of draws.
model : Model (optional if in `with` context)
**kwargs : keyword arguments
Extra keyword arguments are forwarded to pymc3.NUTS.
Returns
-------
start, nuts_sampler
start : pymc3.model.Point
Starting point for sampler
nuts_sampler : pymc3.step_methods.NUTS
Instantiated and initialized NUTS sampler object
"""
model = pm.modelcontext(model)
pm._log.info('Initializing NUTS using {}...'.format(init))
random_seed = int(np.atleast_1d(random_seed)[0])
if init is not None:
init = init.lower()
if init == 'advi':
v_params = pm.variational.advi(n=n_init, random_seed=random_seed)
start = pm.variational.sample_vp(v_params,
njobs,
progressbar=False,
hide_transformed=False,
random_seed=random_seed)
if njobs == 1:
start = start[0]
cov = np.power(model.dict_to_array(v_params.stds), 2)
elif init == 'advi_map':
start = pm.find_MAP()
v_params = pm.variational.advi(n=n_init,
start=start,
random_seed=random_seed)
cov = np.power(model.dict_to_array(v_params.stds), 2)
elif init == 'map':
start = pm.find_MAP()
cov = pm.find_hessian(point=start)
elif init == 'nuts':
init_trace = pm.sample(step=pm.NUTS(),
draws=n_init,
random_seed=random_seed)[n_init // 2:]
cov = np.atleast_1d(pm.trace_cov(init_trace))
start = np.random.choice(init_trace, njobs)
if njobs == 1:
start = start[0]
else:
raise NotImplementedError(
'Initializer {} is not supported.'.format(init))
step = pm.NUTS(scaling=cov, is_cov=True, **kwargs)
return start, step
def init_nuts(init='ADVI',
njobs=1,
n_init=500000,
model=None,
random_seed=-1,
progressbar=True,
**kwargs):
"""Initialize and sample from posterior of a continuous model.
This is a convenience function. NUTS convergence and sampling speed is extremely
dependent on the choice of mass/scaling matrix. In our experience, using ADVI
to estimate a diagonal covariance matrix and using this as the scaling matrix
produces robust results over a wide class of continuous models.
Parameters
----------
init : str {'ADVI', 'ADVI_MAP', 'MAP', 'NUTS'}
Initialization method to use.
* ADVI : Run ADVI to estimate posterior mean and diagonal covariance matrix.
* ADVI_MAP: Initialize ADVI with MAP and use MAP as starting point.
* MAP : Use the MAP as starting point.
* NUTS : Run NUTS and estimate posterior mean and covariance matrix.
njobs : int
Number of parallel jobs to start.
n_init : int
Number of iterations of initializer
If 'ADVI', number of iterations, if 'metropolis', number of draws.
model : Model (optional if in `with` context)
progressbar : bool
Whether or not to display a progressbar for advi sampling.
**kwargs : keyword arguments
Extra keyword arguments are forwarded to pymc3.NUTS.
Returns
-------
start : pymc3.model.Point
Starting point for sampler
nuts_sampler : pymc3.step_methods.NUTS
Instantiated and initialized NUTS sampler object
"""
model = pm.modelcontext(model)
pm._log.info('Initializing NUTS using {}...'.format(init))
random_seed = int(np.atleast_1d(random_seed)[0])
if init is not None:
init = init.lower()
cb = [
pm.callbacks.CheckParametersConvergence(tolerance=1e-2,
diff='absolute'),
pm.callbacks.CheckParametersConvergence(tolerance=1e-2,
diff='relative'),
]
if init == 'advi':
approx = pm.fit(random_seed=random_seed,
n=n_init,
method='advi',
model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window) # type: pm.MeanField
start = approx.sample(draws=njobs)
stds = approx.gbij.rmap(approx.std.eval())
cov = model.dict_to_array(stds)**2
if njobs == 1:
start = start[0]
elif init == 'advi_map':
start = pm.find_MAP()
approx = pm.MeanField(model=model, start=start)
pm.fit(random_seed=random_seed,
n=n_init,
method=pm.ADVI.from_mean_field(approx),
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window)
start = approx.sample(draws=njobs)
stds = approx.gbij.rmap(approx.std.eval())
cov = model.dict_to_array(stds)**2
if njobs == 1:
start = start[0]
elif init == 'map':
start = pm.find_MAP()
cov = pm.find_hessian(point=start)
elif init == 'nuts':
init_trace = pm.sample(draws=n_init,
step=pm.NUTS(),
tune=n_init // 2,
random_seed=random_seed)
cov = np.atleast_1d(pm.trace_cov(init_trace))
start = np.random.choice(init_trace, njobs)
if njobs == 1:
start = start[0]
else:
raise NotImplementedError(
'Initializer {} is not supported.'.format(init))
step = pm.NUTS(scaling=cov, is_cov=True, **kwargs)
return start, step
def train_model(self, method, prior, res=22.):
test_model = self.get_model(prior, test=True)
with test_model:
# True step size is scaled down
# by (1/n)^(1/4), so that a res of 20º -> 15º sampling in 3-d
# therefore, we scale UP by (1/n)^(1/4) so when we specify
# 20º, we get 20º!
step_scale = res * np.pi / 180 / (1 / self.n_d)**0.25
if method == 'MH':
step = pm.Metropolis()
test_trace = pm.sample(step=step,
cores=1,
chains=1,
draws=1,
tune=8000,
discard_tuned_samples=False)
cov = pm.trace_cov(test_trace)
self.write_(cov, method, prior)
return
elif method == 'HMC':
step = pm.HamiltonianMC(step_scale=step_scale)
test_trace = pm.sample(step=step,
cores=1,
chains=1,
draws=1,
tune=8000,
discard_tuned_samples=False)
cov = pm.trace_cov(test_trace)
self.write_(cov, method, prior)
return
elif method == 'NUTS':
step = pm.NUTS(step_scale=step_scale,
adapt_step_size=False,
target_accept=0.8)
test_trace = pm.sample(step=step,
cores=1,
chains=1,
draws=1,
tune=8000,
discard_tuned_samples=False)
cov = pm.trace_cov(test_trace)
accept = test_trace.get_sampler_stats("mean_tree_accept")
print("Mean acceptance for method", method, "and prior", prior,
"is", accept.mean())
print("Step size is",
test_trace.get_sampler_stats("step_size_bar"))
t_a = accept.mean() if accept.mean() > 0.5 else 0.8
training_model = self.get_model(prior)
# Method is NUTS
# NUTS tuning will be a little more intense:
with training_model:
step = pm.NUTS(scaling=cov,
is_cov=True,
step_scale=step_scale,
adapt_step_size=False,
target_accept=t_a)
train_trace = pm.sample(step=step,
cores=1,
chains=1,
draws=1,
tune=500,
discard_tuned_samples=False)
cov = pm.trace_cov(train_trace)
self.write_(cov, method, prior)
# print(f"Map Estimate:{map_estimate}")
times_consumed = []
N = 100
alphas = norm.rvs(size=N)
betas = halfnorm.rvs(size=(N, 2))
sigmas = halfnorm.rvs(size=N)
print("=========================================")
print(f"Tuning NUTS")
print("=========================================")
n_chains = 4
init_trace = pm.sample(draws=1000, tune=1000, cores=n_chains)
cov = np.atleast_1d(pm.trace_cov(init_trace))
start = list(np.random.choice(init_trace, n_chains))
potential = quadpotential.QuadPotentialFull(cov)
step_size = init_trace.get_sampler_stats("step_size_bar")[-1]
size = m.bijection.ordering.size
step_scale = step_size * (size**0.25)
# with pm.Model() as model_new: # reset model. If you use theano.shared you can also update the value of model1 above
for i in range(N):
# No-U-Turn Sampler NUTS
print("=========================================")
print(f"Turn {i}")
print("=========================================")
# start = {"alpha": alphas[i], "beta": betas[i], "sigma": sigmas[i]}
time_zero = default_timer()
step = pm.NUTS(potential=potential,
def init_nuts(init='auto',
chains=1,
n_init=500000,
model=None,
random_seed=None,
progressbar=True,
**kwargs):
"""Set up the mass matrix initialization for NUTS.
NUTS convergence and sampling speed is extremely dependent on the
choice of mass/scaling matrix. This function implements different
methods for choosing or adapting the mass matrix.
Parameters
----------
init : str
Initialization method to use.
* auto : Choose a default initialization method automatically.
Currently, this is `'jitter+adapt_diag'`, but this can change in
the future. If you depend on the exact behaviour, choose an
initialization method explicitly.
* adapt_diag : Start with a identity mass matrix and then adapt
a diagonal based on the variance of the tuning samples. All
chains use the test value (usually the prior mean) as starting
point.
* jitter+adapt_diag : Same as `adapt_diag`, but add uniform jitter
in [-1, 1] to the starting point in each chain.
* advi+adapt_diag : Run ADVI and then adapt the resulting diagonal
mass matrix based on the sample variance of the tuning samples.
* advi+adapt_diag_grad : Run ADVI and then adapt the resulting
diagonal mass matrix based on the variance of the gradients
during tuning. This is **experimental** and might be removed
in a future release.
* advi : Run ADVI to estimate posterior mean and diagonal mass
matrix.
* advi_map: Initialize ADVI with MAP and use MAP as starting point.
* map : Use the MAP as starting point. This is discouraged.
* nuts : Run NUTS and estimate posterior mean and mass matrix from
the trace.
chains : int
Number of jobs to start.
n_init : int
Number of iterations of initializer
If 'ADVI', number of iterations, if 'nuts', number of draws.
model : Model (optional if in `with` context)
progressbar : bool
Whether or not to display a progressbar for advi sampling.
**kwargs : keyword arguments
Extra keyword arguments are forwarded to pymc3.NUTS.
Returns
-------
start : pymc3.model.Point
Starting point for sampler
nuts_sampler : pymc3.step_methods.NUTS
Instantiated and initialized NUTS sampler object
"""
model = pm.modelcontext(model)
vars = kwargs.get('vars', model.vars)
if set(vars) != set(model.vars):
raise ValueError('Must use init_nuts on all variables of a model.')
if not pm.model.all_continuous(vars):
raise ValueError('init_nuts can only be used for models with only '
'continuous variables.')
if not isinstance(init, str):
raise TypeError('init must be a string.')
if init is not None:
init = init.lower()
if init == 'auto':
init = 'jitter+adapt_diag'
pm._log.info('Initializing NUTS using {}...'.format(init))
if random_seed is not None:
random_seed = int(np.atleast_1d(random_seed)[0])
np.random.seed(random_seed)
cb = [
pm.callbacks.CheckParametersConvergence(tolerance=1e-2,
diff='absolute'),
pm.callbacks.CheckParametersConvergence(tolerance=1e-2,
diff='relative'),
]
if init == 'adapt_diag':
start = [model.test_point] * chains
mean = np.mean([model.dict_to_array(vals) for vals in start], axis=0)
var = np.ones_like(mean)
potential = quadpotential.QuadPotentialDiagAdapt(
model.ndim, mean, var, 10)
elif init == 'jitter+adapt_diag':
start = []
for _ in range(chains):
mean = {var: val.copy() for var, val in model.test_point.items()}
for val in mean.values():
val[...] += 2 * np.random.rand(*val.shape) - 1
start.append(mean)
mean = np.mean([model.dict_to_array(vals) for vals in start], axis=0)
var = np.ones_like(mean)
potential = quadpotential.QuadPotentialDiagAdapt(
model.ndim, mean, var, 10)
elif init == 'advi+adapt_diag_grad':
approx = pm.fit(
random_seed=random_seed,
n=n_init,
method='advi',
model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window,
) # type: pm.MeanField
start = approx.sample(draws=chains)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds)**2
mean = approx.bij.rmap(approx.mean.get_value())
mean = model.dict_to_array(mean)
weight = 50
potential = quadpotential.QuadPotentialDiagAdaptGrad(
model.ndim, mean, cov, weight)
elif init == 'advi+adapt_diag':
approx = pm.fit(
random_seed=random_seed,
n=n_init,
method='advi',
model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window,
) # type: pm.MeanField
start = approx.sample(draws=chains)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds)**2
mean = approx.bij.rmap(approx.mean.get_value())
mean = model.dict_to_array(mean)
weight = 50
potential = quadpotential.QuadPotentialDiagAdapt(
model.ndim, mean, cov, weight)
elif init == 'advi':
approx = pm.fit(random_seed=random_seed,
n=n_init,
method='advi',
model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window) # type: pm.MeanField
start = approx.sample(draws=chains)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds)**2
potential = quadpotential.QuadPotentialDiag(cov)
elif init == 'advi_map':
start = pm.find_MAP(include_transformed=True)
approx = pm.MeanField(model=model, start=start)
pm.fit(random_seed=random_seed,
n=n_init,
method=pm.KLqp(approx),
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window)
start = approx.sample(draws=chains)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds)**2
potential = quadpotential.QuadPotentialDiag(cov)
elif init == 'map':
start = pm.find_MAP(include_transformed=True)
cov = pm.find_hessian(point=start)
start = [start] * chains
potential = quadpotential.QuadPotentialFull(cov)
elif init == 'nuts':
init_trace = pm.sample(draws=n_init,
step=pm.NUTS(),
tune=n_init // 2,
random_seed=random_seed)
cov = np.atleast_1d(pm.trace_cov(init_trace))
start = list(np.random.choice(init_trace, chains))
potential = quadpotential.QuadPotentialFull(cov)
else:
raise NotImplementedError(
'Initializer {} is not supported.'.format(init))
step = pm.NUTS(potential=potential, **kwargs)
return start, step
def init_nuts(init='auto', njobs=1, n_init=500000, model=None,
random_seed=-1, progressbar=True, **kwargs):
"""Set up the mass matrix initialization for NUTS.
NUTS convergence and sampling speed is extremely dependent on the
choice of mass/scaling matrix. This function implements different
methods for choosing or adapting the mass matrix.
Parameters
----------
init : str
Initialization method to use.
* auto : Choose a default initialization method automatically.
Currently, this is `'jitter+adapt_diag'`, but this can change in
the future. If you depend on the exact behaviour, choose an
initialization method explicitly.
* adapt_diag : Start with a identity mass matrix and then adapt
a diagonal based on the variance of the tuning samples. All
chains use the test value (usually the prior mean) as starting
point.
* jitter+adapt_diag : Same as `adapt_diag`, but add uniform jitter
in [-1, 1] to the starting point in each chain.
* advi+adapt_diag : Run ADVI and then adapt the resulting diagonal
mass matrix based on the sample variance of the tuning samples.
* advi+adapt_diag_grad : Run ADVI and then adapt the resulting
diagonal mass matrix based on the variance of the gradients
during tuning. This is **experimental** and might be removed
in a future release.
* advi : Run ADVI to estimate posterior mean and diagonal mass
matrix.
* advi_map: Initialize ADVI with MAP and use MAP as starting point.
* map : Use the MAP as starting point. This is discouraged.
* nuts : Run NUTS and estimate posterior mean and mass matrix from
the trace.
njobs : int
Number of parallel jobs to start.
n_init : int
Number of iterations of initializer
If 'ADVI', number of iterations, if 'nuts', number of draws.
model : Model (optional if in `with` context)
progressbar : bool
Whether or not to display a progressbar for advi sampling.
**kwargs : keyword arguments
Extra keyword arguments are forwarded to pymc3.NUTS.
Returns
-------
start : pymc3.model.Point
Starting point for sampler
nuts_sampler : pymc3.step_methods.NUTS
Instantiated and initialized NUTS sampler object
"""
model = pm.modelcontext(model)
vars = kwargs.get('vars', model.vars)
if set(vars) != set(model.vars):
raise ValueError('Must use init_nuts on all variables of a model.')
if not pm.model.all_continuous(vars):
raise ValueError('init_nuts can only be used for models with only '
'continuous variables.')
if not isinstance(init, str):
raise TypeError('init must be a string.')
if init is not None:
init = init.lower()
if init == 'auto':
init = 'jitter+adapt_diag'
pm._log.info('Initializing NUTS using {}...'.format(init))
random_seed = int(np.atleast_1d(random_seed)[0])
cb = [
pm.callbacks.CheckParametersConvergence(
tolerance=1e-2, diff='absolute'),
pm.callbacks.CheckParametersConvergence(
tolerance=1e-2, diff='relative'),
]
if init == 'adapt_diag':
start = [model.test_point] * njobs
mean = np.mean([model.dict_to_array(vals) for vals in start], axis=0)
var = np.ones_like(mean)
potential = quadpotential.QuadPotentialDiagAdapt(
model.ndim, mean, var, 10)
if njobs == 1:
start = start[0]
elif init == 'jitter+adapt_diag':
start = []
for _ in range(njobs):
mean = {var: val.copy() for var, val in model.test_point.items()}
for val in mean.values():
val[...] += 2 * np.random.rand(*val.shape) - 1
start.append(mean)
mean = np.mean([model.dict_to_array(vals) for vals in start], axis=0)
var = np.ones_like(mean)
potential = quadpotential.QuadPotentialDiagAdapt(
model.ndim, mean, var, 10)
if njobs == 1:
start = start[0]
elif init == 'advi+adapt_diag_grad':
approx = pm.fit(
random_seed=random_seed,
n=n_init, method='advi', model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window,
) # type: pm.MeanField
start = approx.sample(draws=njobs)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds) ** 2
mean = approx.bij.rmap(approx.mean.get_value())
mean = model.dict_to_array(mean)
weight = 50
potential = quadpotential.QuadPotentialDiagAdaptGrad(
model.ndim, mean, cov, weight)
if njobs == 1:
start = start[0]
elif init == 'advi+adapt_diag':
approx = pm.fit(
random_seed=random_seed,
n=n_init, method='advi', model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window,
) # type: pm.MeanField
start = approx.sample(draws=njobs)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds) ** 2
mean = approx.bij.rmap(approx.mean.get_value())
mean = model.dict_to_array(mean)
weight = 50
potential = quadpotential.QuadPotentialDiagAdapt(
model.ndim, mean, cov, weight)
if njobs == 1:
start = start[0]
elif init == 'advi':
approx = pm.fit(
random_seed=random_seed,
n=n_init, method='advi', model=model,
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window
) # type: pm.MeanField
start = approx.sample(draws=njobs)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds) ** 2
potential = quadpotential.QuadPotentialDiag(cov)
if njobs == 1:
start = start[0]
elif init == 'advi_map':
start = pm.find_MAP()
approx = pm.MeanField(model=model, start=start)
pm.fit(
random_seed=random_seed,
n=n_init, method=pm.KLqp(approx),
callbacks=cb,
progressbar=progressbar,
obj_optimizer=pm.adagrad_window
)
start = approx.sample(draws=njobs)
start = list(start)
stds = approx.bij.rmap(approx.std.eval())
cov = model.dict_to_array(stds) ** 2
potential = quadpotential.QuadPotentialDiag(cov)
if njobs == 1:
start = start[0]
elif init == 'map':
start = pm.find_MAP()
cov = pm.find_hessian(point=start)
start = [start] * njobs
potential = quadpotential.QuadPotentialFull(cov)
if njobs == 1:
start = start[0]
elif init == 'nuts':
init_trace = pm.sample(draws=n_init, step=pm.NUTS(),
tune=n_init // 2,
random_seed=random_seed)
cov = np.atleast_1d(pm.trace_cov(init_trace))
start = list(np.random.choice(init_trace, njobs))
potential = quadpotential.QuadPotentialFull(cov)
if njobs == 1:
start = start[0]
else:
raise NotImplementedError('Initializer {} is not supported.'.format(init))
step = pm.NUTS(potential=potential, **kwargs)
return start, step
total_tools = shared(dk.total_tools.values)
# %%
with pm.Model() as m_10_10:
a = pm.Normal("a", 0, 100)
b = pm.Normal("b", 0, 1, shape=3)
lam = pm.math.exp(a + b[0] * log_pop + b[1] * contact_high +
b[2] * contact_high * log_pop)
obs = pm.Poisson("total_tools", lam, observed=total_tools)
trace_10_10 = pm.sample(1000, tune=1000)
# %%
summary = az.summary(trace_10_10, credible_interval=0.89)[[
"mean", "sd", "hpd_5.5%", "hpd_94.5%"
]]
trace_cov = pm.trace_cov(trace_10_10, model=m_10_10)
invD = (np.sqrt(np.diag(trace_cov))**-1)[:, None]
trace_corr = pd.DataFrame(invD * trace_cov * invD.T,
index=summary.index,
columns=summary.index)
summary.join(trace_corr).round(2)
# %%
az.plot_forest(trace_10_10)
# %%
lambda_high = np.exp(trace_10_10["a"] + trace_10_10["b"][:, 1] +
(trace_10_10["b"][:, 0] + trace_10_10["b"][:, 2]) * 8)
lambda_low = np.exp(trace_10_10["a"] + trace_10_10["b"][:, 0] * 8)
