def test_exec_nuts_init(method):
with pm.Model() as model:
pm.Normal('a', mu=0, sd=1, shape=2)
with model:
start, _ = pm.init_nuts(init=method, n_init=10)
assert isinstance(start, dict)
start, _ = pm.init_nuts(init=method, n_init=10, njobs=2)
assert isinstance(start, list)
assert len(start) == 2
assert isinstance(start[0], dict)
def test_exec_nuts_init(method):
with pm.Model() as model:
pm.Normal('a', mu=0, sd=1, shape=2)
with model:
start, _ = pm.init_nuts(init=method, n_init=10)
assert isinstance(start, dict)
start, _ = pm.init_nuts(init=method, n_init=10, njobs=2)
assert isinstance(start, list)
assert len(start) == 2
assert isinstance(start[0], dict)
def test_exec_nuts_init(method):
with pm.Model() as model:
pm.Normal("a", mu=0, sigma=1, shape=2)
pm.HalfNormal("b", sigma=1)
with model:
start, _ = pm.init_nuts(init=method, n_init=10)
assert isinstance(start, list)
assert len(start) == 1
assert isinstance(start[0], dict)
assert "a" in start[0] and "b_log__" in start[0]
start, _ = pm.init_nuts(init=method, n_init=10, chains=2)
assert isinstance(start, list)
assert len(start) == 2
assert isinstance(start[0], dict)
assert "a" in start[0] and "b_log__" in start[0]
def test_exec_nuts_init(method):
with pm.Model() as model:
pm.Normal("a", mu=0, sigma=1, shape=2)
pm.HalfNormal("b", sigma=1)
with model:
start, _ = pm.init_nuts(init=method, n_init=10)
assert isinstance(start, list)
assert len(start) == 1
assert isinstance(start[0], dict)
assert "a" in start[0] and "b_log__" in start[0]
start, _ = pm.init_nuts(init=method, n_init=10, chains=2)
assert isinstance(start, list)
assert len(start) == 2
assert isinstance(start[0], dict)
assert "a" in start[0] and "b_log__" in start[0]
def test_exec_nuts_init(method):
with pm.Model() as model:
pm.Normal('a', mu=0, sd=1, shape=2)
pm.HalfNormal('b', sd=1)
with model:
start, _ = pm.init_nuts(init=method, n_init=10)
assert isinstance(start, list)
assert len(start) == 1
assert isinstance(start[0], dict)
assert 'a' in start[0] and 'b_log__' in start[0]
start, _ = pm.init_nuts(init=method, n_init=10, chains=2)
assert isinstance(start, list)
assert len(start) == 2
assert isinstance(start[0], dict)
assert 'a' in start[0] and 'b_log__' in start[0]
def setup(self, step, init):
"""Initialize model and get start position"""
np.random.seed(123)
self.chains = 4
data = pd.read_csv(pm.get_data('radon.csv'))
data['log_radon'] = data['log_radon'].astype(theano.config.floatX)
county_idx = data.county_code.values
n_counties = len(data.county.unique())
with pm.Model() as self.model:
mu_a = pm.Normal('mu_a', mu=0., sd=100**2)
sigma_a = pm.HalfCauchy('sigma_a', 5)
mu_b = pm.Normal('mu_b', mu=0., sd=100**2)
sigma_b = pm.HalfCauchy('sigma_b', 5)
a = pm.Normal('a', mu=mu_a, sd=sigma_a, shape=n_counties)
b = pm.Normal('b', mu=mu_b, sd=sigma_b, shape=n_counties)
eps = pm.HalfCauchy('eps', 5)
radon_est = a[county_idx] + b[county_idx] * data.floor.values
pm.Normal('radon_like',
mu=radon_est,
sd=eps,
observed=data.log_radon)
self.start, _ = pm.init_nuts(chains=self.chains, init=init)
def _mocked_init_nuts(*args, **kwargs):
if init == "adapt_diag":
start_ = [{"x": np.array(0.79788456)}]
else:
start_ = [{"x": np.array(-0.04949886)}]
_, step = pm.init_nuts(*args, **kwargs)
return start_, step
def check_exec_nuts_init(method):
with pm.Model() as model:
pm.Normal("a", mu=0, sigma=1, size=2)
pm.HalfNormal("b", sigma=1)
with model:
start, _ = pm.init_nuts(init=method, n_init=10)
assert isinstance(start, list)
assert len(start) == 1
assert isinstance(start[0], dict)
assert model.a.tag.value_var.name in start[0]
assert model.b.tag.value_var.name in start[0]
start, _ = pm.init_nuts(init=method, n_init=10, chains=2)
assert isinstance(start, list)
assert len(start) == 2
assert isinstance(start[0], dict)
assert model.a.tag.value_var.name in start[0]
assert model.b.tag.value_var.name in start[0]
def track_glm_hierarchical_ess(self, init):
with glm_hierarchical_model():
start, step = pm.init_nuts(init=init, chains=self.chains, progressbar=False, random_seed=123)
t0 = time.time()
trace = pm.sample(draws=self.draws, step=step, njobs=4, chains=self.chains,
start=start, random_seed=100)
tot = time.time() - t0
ess = pm.effective_n(trace, ('mu_a',))['mu_a']
return ess / tot
def track_glm_hierarchical_ess(self, init):
with glm_hierarchical_model():
start, step = pm.init_nuts(init=init, chains=self.chains, progressbar=False, random_seed=123)
t0 = time.time()
trace = pm.sample(draws=self.draws, step=step, cores=4, chains=self.chains,
start=start, random_seed=100, progressbar=False,
compute_convergence_checks=False)
tot = time.time() - t0
ess = float(pm.ess(trace, var_names=['mu_a'])['mu_a'].values)
return ess / tot
def sample_chain(model,
chain_i=0,
step=None,
num_samples=MAX_NUM_SAMPLES,
advi=False,
tune=5,
discard_tuned_samples=True,
num_scale1_iters=NUM_SCALE1_ITERS,
num_scale0_iters=NUM_SCALE0_ITERS):
"""Sample single chain from constructed Bayesian model"""
start = timer()
with model:
if not advi:
pm._log.info('Assigning NUTS sampler...')
if step is None:
start_, step = pm.init_nuts(init='advi',
njobs=1,
n_init=NUM_INIT_STEPS,
random_seed=-1,
progressbar=False)
discard = tune if discard_tuned_samples else 0
for i, trace in enumerate(
pm.iter_sample(num_samples + discard,
step,
start=start_,
chain=chain_i)):
if i == 0:
min_num_samples = get_min_samples_per_chain(
len(trace[0]), MIN_SAMPLES_CONSTANT, NUM_CHAINS)
elapsed = timer() - start
if elapsed > SOFT_MAX_TIME_IN_SECONDS / NUM_CHAINS:
print('exceeded soft time limit...')
if i + 1 - discard >= min_num_samples:
print('collected enough samples; stopping')
break
else:
print('but only collected {} of {}; continuing...'.
format(i + 1 - discard, min_num_samples))
if elapsed > HARD_MAX_TIME_IN_SECONDS / NUM_CHAINS:
print('exceeded HARD time limit; STOPPING')
break
return trace[discard:]
else: # ADVI for neural networks
scale = theano.shared(pm.floatX(1))
vi = pm.ADVI(cost_part_grad_scale=scale)
pm.fit(n=num_scale1_iters, method=vi)
scale.set_value(0)
approx = pm.fit(n=num_scale0_iters)
# one sample to get dimensions of trace
trace = approx.sample(draws=1)
min_num_samples = get_min_samples_per_chain(
len(trace.varnames), MIN_SAMPLES_CONSTANT, 1)
trace = approx.sample(draws=min_num_samples)
return trace
def track_marginal_mixture_model_ess(self, init):
model, start = mixture_model()
with model:
_, step = pm.init_nuts(init=init, chains=self.chains,
progressbar=False, random_seed=123)
start = [{k: v for k, v in start.items()} for _ in range(self.chains)]
t0 = time.time()
trace = pm.sample(draws=self.draws, step=step, njobs=4, chains=self.chains,
start=start, random_seed=100)
tot = time.time() - t0
ess = pm.effective_n(trace, ('mu',))['mu'].min() # worst case
return ess / tot
def track_marginal_mixture_model_ess(self, init):
model, start = mixture_model()
with model:
_, step = pm.init_nuts(init=init, chains=self.chains,
progressbar=False, random_seed=123)
start = [{k: v for k, v in start.items()} for _ in range(self.chains)]
t0 = time.time()
trace = pm.sample(draws=self.draws, step=step, cores=4, chains=self.chains,
start=start, random_seed=100, progressbar=False,
compute_convergence_checks=False)
tot = time.time() - t0
ess = pm.ess(trace, var_names=['mu'])['mu'].values.min() # worst case
return ess / tot
def track_glm_hierarchical_ess(self, init):
with glm_hierarchical_model():
start, step = pm.init_nuts(init=init,
chains=self.chains,
progressbar=False,
random_seed=123)
t0 = time.time()
trace = pm.sample(draws=self.draws,
step=step,
njobs=4,
chains=self.chains,
start=start,
random_seed=100)
tot = time.time() - t0
ess = pm.effective_n(trace, ('mu_a', ))['mu_a']
return ess / tot
def test_samples(self):
chains = 1
with pm.Model():
x = pm.Normal('x', 0, 1)
y = pm.Normal('y', x, 1)
start, step = pm.init_nuts(chains=chains)
l2hmc_step = L2HMC(potential=step.potential)
l2hmc_trace = pm.sample(2000,
step=l2hmc_step,
start=start,
chains=chains)
assert np.abs(l2hmc_trace['x'].mean()) < 0.02
assert np.abs(l2hmc_trace['x'].std() - 1) < 0.02
assert np.abs(l2hmc_trace['y'].mean()) < 0.05
def track_marginal_mixture_model_ess(self, init):
model, start = mixture_model()
with model:
_, step = pm.init_nuts(init=init,
chains=self.chains,
progressbar=False,
random_seed=123)
start = [{k: v
for k, v in start.items()} for _ in range(self.chains)]
t0 = time.time()
trace = pm.sample(draws=self.draws,
step=step,
njobs=4,
chains=self.chains,
start=start,
random_seed=100)
tot = time.time() - t0
ess = pm.effective_n(trace, ('mu', ))['mu'].min() # worst case
return ess / tot
def time_glm_hierarchical_init(self, init):
"""How long does it take to run the initialization."""
with glm_hierarchical_model():
pm.init_nuts(init=init, chains=self.chains, progressbar=False)
def time_glm_hierarchical_init(self, init):
"""How long does it take to run the initialization."""
with glm_hierarchical_model():
pm.init_nuts(init=init, chains=self.chains, progressbar=False)