def sample(
draws=1000,
tune=1000,
model=None,
warmup_window=50,
adapt_window=50,
cooldown_window=50,
target_accept=0.9,
gamma=0.05,
k=0.75,
t0=10,
step_kwargs=None,
**kwargs,
):
model = modelcontext(model)
if not all_continuous(model.vars):
raise ValueError("NUTS can only be used for models with only "
"continuous variables.")
warmup_window, update_steps = build_schedule(
tune,
warmup_window=warmup_window,
adapt_window=adapt_window,
cooldown_window=cooldown_window,
)
potential = WindowedQuadPotentialDenseAdapt(model.ndim, warmup_window,
update_steps)
if step_kwargs is None:
step_kwargs = {}
step = pm.NUTS(
potential=potential,
model=model,
target_accept=target_accept,
**step_kwargs,
)
step.step_adapt = WindowedDualAverageAdaptation(
np.append(warmup_window, update_steps),
step.step_size,
target_accept,
gamma,
k,
t0,
)
kwargs["step"] = step
return pm.sample(draws=draws, tune=tune, model=model, **kwargs)
def get_dense_nuts_step(start=None,
adaptation_window=101,
doubling=True,
model=None,
**kwargs):
"""Get a NUTS step function with a dense mass matrix
The entries in the mass matrix will be tuned based on the sample
covariances during tuning. All extra arguments are passed directly to
``pymc3.NUTS``.
Args:
start (dict, optional): A starting point in parameter space. If not
provided, the model's ``test_point`` is used.
adaptation_window (int, optional): The (initial) size of the window
used for sample covariance estimation.
doubling (bool, optional): If ``True`` (default) the adaptation window
is doubled each time the matrix is updated.
"""
model = modelcontext(model)
if not all_continuous(model.vars):
raise ValueError("NUTS can only be used for models with only "
"continuous variables.")
if start is None:
start = model.test_point
mean = model.dict_to_array(start)
var = np.eye(len(mean))
potential = QuadPotentialDenseAdapt(
model.ndim,
mean,
var,
10,
adaptation_window=adaptation_window,
doubling=doubling,
)
return pm.NUTS(potential=potential, model=model, **kwargs)
def get_dense_nuts_step(
start=None,
adaptation_window=101,
doubling=True,
initial_weight=10,
use_hessian=False,
use_hessian_diag=False,
hessian_regularization=1e-8,
model=None,
**kwargs,
):
"""Get a NUTS step function with a dense mass matrix
The entries in the mass matrix will be tuned based on the sample
covariances during tuning. All extra arguments are passed directly to
``pymc3.NUTS``.
Args:
start (dict, optional): A starting point in parameter space. If not
provided, the model's ``test_point`` is used.
adaptation_window (int, optional): The (initial) size of the window
used for sample covariance estimation.
doubling (bool, optional): If ``True`` (default) the adaptation window
is doubled each time the matrix is updated.
"""
model = modelcontext(model)
if not all_continuous(model.vars):
raise ValueError("NUTS can only be used for models with only "
"continuous variables.")
if start is None:
start = model.test_point
mean = model.dict_to_array(start)
if use_hessian or use_hessian_diag:
try:
import numdifftools as nd
except ImportError:
raise ImportError(
"The 'numdifftools' package is required for Hessian "
"computations")
logger.info("Numerically estimating Hessian matrix")
if use_hessian_diag:
hess = nd.Hessdiag(model.logp_array)(mean)
var = np.diag(-1.0 / hess)
else:
hess = nd.Hessian(model.logp_array)(mean)
var = -np.linalg.inv(hess)
factor = 1
success = False
while not success:
var[np.diag_indices_from(var)] += factor * hessian_regularization
try:
np.linalg.cholesky(var)
except np.linalg.LinAlgError:
factor *= 2
else:
success = True
else:
var = np.eye(len(mean))
potential = QuadPotentialDenseAdapt(
model.ndim,
mean,
var,
initial_weight,
adaptation_window=adaptation_window,
doubling=doubling,
)
return pm.NUTS(potential=potential, model=model, **kwargs)
def sample(
*,
draws=1000,
tune=1000,
model=None,
step_kwargs=None,
warmup_window=50,
adapt_window=50,
cooldown_window=100,
initial_accept=None,
target_accept=0.9,
gamma=0.05,
k=0.75,
t0=10,
**kwargs,
):
# Check that we're in a model context and that all the variables are
# continuous
model = modelcontext(model)
if not all_continuous(model.vars):
raise ValueError("NUTS can only be used for models with only "
"continuous variables.")
start = kwargs.get("start", None)
if start is None:
start = model.test_point
mean = model.dict_to_array(start)
update_steps = build_schedule(
tune,
warmup_window=warmup_window,
adapt_window=adapt_window,
cooldown_window=cooldown_window,
)
potential = QuadPotentialDenseAdapt(
model.ndim,
initial_mean=mean,
initial_weight=10,
update_steps=update_steps,
)
if "step" in kwargs:
step = kwargs["step"]
else:
if step_kwargs is None:
step_kwargs = {}
step = pm.NUTS(
potential=potential,
model=model,
target_accept=target_accept,
**step_kwargs,
)
if "target_accept" in step_kwargs and target_accept is not None:
raise ValueError(
"'target_accept' cannot be given as a keyword argument and in "
"'step_kwargs'")
target_accept = step_kwargs.pop("target_accept", target_accept)
if initial_accept is None:
target = target_accept
else:
if initial_accept > target_accept:
raise ValueError(
"initial_accept must be less than or equal to target_accept")
target = initial_accept + (target_accept - initial_accept) * np.sqrt(
np.arange(len(update_steps)) / (len(update_steps) - 1))
step.step_adapt = WindowedDualAverageAdaptation(update_steps,
step.step_size, target,
gamma, k, t0)
kwargs["step"] = step
return pm.sample(draws=draws, tune=tune, model=model, **kwargs)