def track_1var_2par_ode_ess(self):
def freefall(y, t, p):
return 2.0 * p[1] - p[0] * y[0]
# Times for observation
times = np.arange(0, 10, 0.5)
y = np.array([
-2.01, 9.49, 15.58, 16.57, 27.58, 32.26, 35.13, 38.07, 37.36,
38.83, 44.86, 43.58, 44.59, 42.75, 46.9, 49.32, 44.06, 49.86,
46.48, 48.18
]).reshape(-1, 1)
ode_model = pm.ode.DifferentialEquation(func=freefall,
times=times,
n_states=1,
n_theta=2,
t0=0)
with pm.Model() as model:
# Specify prior distributions for some of our model parameters
sigma = pm.HalfCauchy("sigma", 1)
gamma = pm.Lognormal("gamma", 0, 1)
# If we know one of the parameter values, we can simply pass the value.
ode_solution = ode_model(y0=[0], theta=[gamma, 9.8])
# The ode_solution has a shape of (n_times, n_states)
Y = pm.Normal("Y", mu=ode_solution, sd=sigma, observed=y)
t0 = time.time()
trace = pm.sample(500, tune=1000, chains=2, cores=2, random_seed=0)
tot = time.time() - t0
ess = pm.ess(trace)
return np.mean([ess.sigma, ess.gamma]) / 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 track_glm_hierarchical_ess(self, step):
with glm_hierarchical_model():
if step is not None:
step = step()
t0 = time.time()
trace = pm.sample(draws=self.draws, step=step, cores=4, chains=4,
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 extract_mcmc_chain_from_pymc3_trace(trace, var_names, ndraws, nburn,
njobs):
nvars = len(var_names)
samples = np.empty((nvars, (ndraws - nburn) * njobs))
effective_sample_size = np.empty(nvars)
for ii in range(nvars):
samples[ii, :] = trace.get_values(var_names[ii],
burn=nburn,
chains=np.arange(njobs))
effective_sample_size[ii] = pm.ess(trace)[var_names[ii]].values
return samples, effective_sample_size
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 ess(self, dim=0, burn=50, clip=0, is_ess=True):
ess = pymc3.ess(self.fetch_samples()[burn:len(self.samples) - clip,
dim].detach().numpy())
if self.shadow and is_ess == True:
weights = torch.cat(self.radon_nikodym[burn:len(self.samples) -
clip],
dim=0).reshape(-1, 1).detach()
normed_weights = weights / weights.sum()
is_ess = torch.sum(normed_weights**2).reciprocal()
adjusted_ess = ess * is_ess / weights.shape[0]
return adjusted_ess
else:
return ess
def extract_mcmc_chain_from_pymc3_trace(trace,var_names,ndraws,nburn,njobs):
nvars = len(var_names)
samples = np.empty((nvars,(ndraws-nburn)*njobs))
effective_sample_size = -np.ones(nvars)
for ii in range(nvars):
samples[ii,:]=trace.get_values(
var_names[ii],burn=nburn,chains=np.arange(njobs))
try:
effective_sample_size[ii]=pm.ess(trace)[var_names[ii]].values
except:
print('could not compute ess. likely issue with theano')
return samples,effective_sample_size
def run(steppers, p):
steppers = set(steppers)
traces = {}
effn = {}
runtimes = {}
with pm.Model() as model:
if USE_XY:
x = pm.Flat("x")
y = pm.Flat("y")
mu = np.array([0.0, 0.0])
cov = np.array([[1.0, p], [p, 1.0]])
z = pm.MvNormal.dist(mu=mu, cov=cov,
shape=(2, )).logp(tt.stack([x, y]))
pot = pm.Potential("logp_xy", z)
start = {"x": 0, "y": 0}
else:
mu = np.array([0.0, 0.0])
cov = np.array([[1.0, p], [p, 1.0]])
z = pm.MvNormal("z", mu=mu, cov=cov, shape=(2, ))
start = {"z": [0, 0]}
for step_cls in steppers:
name = step_cls.__name__
t_start = time.time()
mt = pm.sample(draws=10000,
chains=16,
parallelize=False,
step=step_cls(),
start=start)
runtimes[name] = time.time() - t_start
print("{} samples across {} chains".format(
len(mt) * mt.nchains, mt.nchains))
traces[name] = mt
en = pm.ess(mt)
print(f"effective: {en}\r\n")
if USE_XY:
effn[name] = np.mean(en["x"]) / len(mt) / mt.nchains
else:
effn[name] = np.mean(en["z"]) / len(mt) / mt.nchains
return traces, effn, runtimes
def run(steppers, p):
steppers = set(steppers)
traces = {}
effn = {}
runtimes = {}
with pm.Model() as model:
if USE_XY:
x = pm.Flat('x')
y = pm.Flat('y')
mu = np.array([0., 0.])
cov = np.array([[1., p], [p, 1.]])
z = pm.MvNormal.dist(mu=mu, cov=cov,
shape=(2, )).logp(tt.stack([x, y]))
pot = pm.Potential('logp_xy', z)
start = {'x': 0, 'y': 0}
else:
mu = np.array([0., 0.])
cov = np.array([[1., p], [p, 1.]])
z = pm.MvNormal('z', mu=mu, cov=cov, shape=(2, ))
start = {'z': [0, 0]}
for step_cls in steppers:
name = step_cls.__name__
t_start = time.time()
mt = pm.sample(draws=10000,
chains=16,
parallelize=False,
step=step_cls(),
start=start)
runtimes[name] = time.time() - t_start
print('{} samples across {} chains'.format(
len(mt) * mt.nchains, mt.nchains))
traces[name] = mt
en = pm.ess(mt)
print('effective: {}\r\n'.format(en))
if USE_XY:
effn[name] = np.mean(en['x']) / len(mt) / mt.nchains
else:
effn[name] = np.mean(en['z']) / len(mt) / mt.nchains
return traces, effn, runtimes
def test_neff(self):
if hasattr(self, "min_n_eff"):
n_eff = pm.ess(self.trace[self.burn :])
for var in n_eff:
npt.assert_array_less(self.min_n_eff, n_eff[var])
def _run_convergence_checks(self, trace, model):
if trace.nchains == 1:
msg = ("Only one chain was sampled, this makes it impossible to "
"run some convergence checks")
warn = SamplerWarning(WarningType.BAD_PARAMS, msg, 'info', None,
None, None)
self._add_warnings([warn])
return
from pymc3 import rhat, ess
valid_name = [rv.name for rv in model.free_RVs + model.deterministics]
varnames = []
for rv in model.free_RVs:
rv_name = rv.name
if is_transformed_name(rv_name):
rv_name2 = get_untransformed_name(rv_name)
rv_name = rv_name2 if rv_name2 in valid_name else rv_name
if rv_name in trace.varnames:
varnames.append(rv_name)
self._ess = ess = ess(trace, var_names=varnames)
self._rhat = rhat = rhat(trace, var_names=varnames)
warnings = []
rhat_max = max(val.max() for val in rhat.values())
if rhat_max > 1.4:
msg = ("The rhat statistic is larger than 1.4 for some "
"parameters. The sampler did not converge.")
warn = SamplerWarning(WarningType.CONVERGENCE, msg, 'error', None,
None, rhat)
warnings.append(warn)
elif rhat_max > 1.2:
msg = ("The rhat statistic is larger than 1.2 for some "
"parameters.")
warn = SamplerWarning(WarningType.CONVERGENCE, msg, 'warn', None,
None, rhat)
warnings.append(warn)
elif rhat_max > 1.05:
msg = ("The rhat statistic is larger than 1.05 for some "
"parameters. This indicates slight problems during "
"sampling.")
warn = SamplerWarning(WarningType.CONVERGENCE, msg, 'info', None,
None, rhat)
warnings.append(warn)
eff_min = min(val.min() for val in ess.values())
n_samples = len(trace) * trace.nchains
if eff_min < 200 and n_samples >= 500:
msg = ("The estimated number of effective samples is smaller than "
"200 for some parameters.")
warn = SamplerWarning(WarningType.CONVERGENCE, msg, 'error', None,
None, ess)
warnings.append(warn)
elif eff_min / n_samples < 0.1:
msg = ("The number of effective samples is smaller than "
"10% for some parameters.")
warn = SamplerWarning(WarningType.CONVERGENCE, msg, 'warn', None,
None, ess)
warnings.append(warn)
elif eff_min / n_samples < 0.25:
msg = ("The number of effective samples is smaller than "
"25% for some parameters.")
warn = SamplerWarning(WarningType.CONVERGENCE, msg, 'info', None,
None, ess)
warnings.append(warn)
self._add_warnings(warnings)
