def test_predictive(parallel):
model, data, true_probs = beta_bernoulli()
mcmc = MCMC(NUTS(model), num_warmup=100, num_samples=100)
mcmc.run(random.PRNGKey(0), data)
samples = mcmc.get_samples()
predictive = Predictive(model, samples, parallel=parallel)
predictive_samples = predictive.get_samples(random.PRNGKey(1))
assert predictive_samples.keys() == {"obs"}
predictive.return_sites = ["beta", "obs"]
predictive_samples = predictive.get_samples(random.PRNGKey(1))
# check shapes
assert predictive_samples["beta"].shape == (100, ) + true_probs.shape
assert predictive_samples["obs"].shape == (100, ) + data.shape
# check sample mean
assert_allclose(
predictive_samples["obs"].reshape((-1, ) + true_probs.shape).mean(0),
true_probs,
rtol=0.1)
def test_prior_with_sample_shape():
data = {
"J": 8,
"y": jnp.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0]),
"sigma": jnp.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0]),
}
def schools_model():
mu = numpyro.sample('mu', dist.Normal(0, 5))
tau = numpyro.sample('tau', dist.HalfCauchy(5))
theta = numpyro.sample('theta',
dist.Normal(mu, tau),
sample_shape=(data['J'], ))
numpyro.sample('obs', dist.Normal(theta, data['sigma']), obs=data['y'])
num_samples = 500
mcmc = MCMC(NUTS(schools_model), num_warmup=500, num_samples=num_samples)
mcmc.run(random.PRNGKey(0))
assert mcmc.get_samples()['theta'].shape == (num_samples, data['J'])
def main(args):
print('Simulating data...')
(transition_prior, emission_prior, transition_prob, emission_prob,
supervised_categories, supervised_words,
unsupervised_words) = simulate_data(
random.PRNGKey(1),
num_categories=args.num_categories,
num_words=args.num_words,
num_supervised_data=args.num_supervised,
num_unsupervised_data=args.num_unsupervised,
)
print('Starting inference...')
rng_key = random.PRNGKey(2)
start = time.time()
kernel = NUTS(semi_supervised_hmm)
mcmc = MCMC(
kernel,
args.num_warmup,
args.num_samples,
num_chains=args.num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
mcmc.run(rng_key, transition_prior, emission_prior, supervised_categories,
supervised_words, unsupervised_words, args.unroll_loop)
samples = mcmc.get_samples()
print_results(samples, transition_prob, emission_prob)
print('\nMCMC elapsed time:', time.time() - start)
# make plots
fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True)
x = np.linspace(0, 1, 101)
for i in range(transition_prob.shape[0]):
for j in range(transition_prob.shape[1]):
ax.plot(x,
gaussian_kde(samples['transition_prob'][:, i, j])(x),
label="trans_prob[{}, {}], true value = {:.2f}".format(
i, j, transition_prob[i, j]))
ax.set(xlabel="Probability",
ylabel="Frequency",
title="Transition probability posterior")
ax.legend()
plt.savefig("hmm_plot.pdf")
def run_inference(design_matrix: jnp.ndarray, outcome: jnp.ndarray,
rng_key: jnp.ndarray,
num_warmup: int,
num_samples: int, num_chains: int,
interval_size: float = 0.95) -> None:
"""
Estimate the effect size.
"""
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup, num_samples, num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
mcmc.run(rng_key, design_matrix, outcome)
# 0th column is intercept (not getting called)
# 1st column is effect of getting called
# 2nd column is effect of gender (should be none since assigned at random)
coef = mcmc.get_samples()['coefficients']
print_results(coef, interval_size)
def test_bernoulli_latent_model():
def model(data):
y_prob = numpyro.sample("y_prob", dist.Beta(1., 1.))
with numpyro.plate("data", data.shape[0]):
y = numpyro.sample("y", dist.Bernoulli(y_prob))
z = numpyro.sample("z", dist.Bernoulli(0.65 * y + 0.1))
numpyro.sample("obs", dist.Normal(2. * z, 1.), obs=data)
N = 2000
y_prob = 0.3
y = dist.Bernoulli(y_prob).sample(random.PRNGKey(0), (N, ))
z = dist.Bernoulli(0.65 * y + 0.1).sample(random.PRNGKey(1))
data = dist.Normal(2. * z, 1.0).sample(random.PRNGKey(2))
nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=500)
mcmc.run(random.PRNGKey(3), data)
samples = mcmc.get_samples()
assert_allclose(samples["y_prob"].mean(0), y_prob, atol=0.05)
def test_chain_smoke(chain_method, compile_args):
def model(data):
concentration = jnp.array([1.0, 1.0, 1.0])
p_latent = numpyro.sample('p_latent', dist.Dirichlet(concentration))
numpyro.sample('obs', dist.Categorical(p_latent), obs=data)
return p_latent
data = dist.Categorical(jnp.array([0.1, 0.6,
0.3])).sample(random.PRNGKey(1),
(2000, ))
kernel = NUTS(model)
mcmc = MCMC(kernel,
2,
5,
num_chains=2,
chain_method=chain_method,
jit_model_args=compile_args)
mcmc.warmup(random.PRNGKey(0), data)
mcmc.run(random.PRNGKey(1), data)
def test_improper_normal(max_tree_depth):
true_coef = 0.9
def model(data):
alpha = numpyro.sample("alpha", dist.Uniform(0, 1))
with numpyro.handlers.reparam(config={"loc": TransformReparam()}):
loc = numpyro.sample(
"loc",
dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False), AffineTransform(0, alpha)),
)
numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
data = true_coef + random.normal(random.PRNGKey(0), (1000, ))
kernel = NUTS(model=model, max_tree_depth=max_tree_depth)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(0), data)
samples = mcmc.get_samples()
assert_allclose(jnp.mean(samples["loc"], 0), true_coef, atol=0.05)
def test_initial_inverse_mass_matrix(dense_mass):
def model():
numpyro.sample("x", dist.Normal(0, 1).expand([3]))
numpyro.sample("z", dist.Normal(0, 1).expand([2]))
expected_mm = jnp.arange(1, 4.0)
kernel = NUTS(
model,
dense_mass=dense_mass,
inverse_mass_matrix={("x", ): expected_mm},
adapt_mass_matrix=False,
)
mcmc = MCMC(kernel, num_warmup=1, num_samples=1)
mcmc.run(random.PRNGKey(0))
inverse_mass_matrix = mcmc.last_state.adapt_state.inverse_mass_matrix
assert set(inverse_mass_matrix.keys()) == {("x", ), ("z", )}
expected_mm = jnp.diag(expected_mm) if dense_mass else expected_mm
assert_allclose(inverse_mass_matrix[("x", )], expected_mm)
assert_allclose(inverse_mass_matrix[("z", )], jnp.ones(2))
def test_predictive_with_improper():
true_coef = 0.9
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
with handlers.reparam(config={'loc': TransformReparam()}):
loc = numpyro.sample(
'loc',
dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False), AffineTransform(0, alpha)))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
data = true_coef + random.normal(random.PRNGKey(0), (1000, ))
kernel = NUTS(model=model)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(0), data)
samples = mcmc.get_samples()
obs_pred = Predictive(model, samples)(random.PRNGKey(1), data=None)["obs"]
assert_allclose(jnp.mean(obs_pred), true_coef, atol=0.05)
def mcmc_inference(model, num_warmup, num_samples, num_chains, rng_key, X, Y):
""""
Helper function for doing NUTS inference.
:param model: a parametric function proportional to the posterior (see gp_regression.likelihood).
:param num_warmup: warmup steps.
:param num_samples: number of samples.
:param num_chains: number of Markov chains used for MCMC sampling.
:param rng_key: random seed.
:param X: X data.
:param Y: Y data.
:return: Dictionary key: name of parameter (from defined in model), value: list of samples.
"""
start = time.time()
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup, num_samples, num_chains=num_chains)
mcmc.run(rng_key, X, Y)
print('\nMCMC time:', time.time() - start)
print(mcmc.print_summary())
return mcmc.get_samples()
def test_scan():
def model(T=10, q=1, r=1, phi=0., beta=0.):
def transition(state, i):
x0, mu0 = state
x1 = numpyro.sample('x', dist.Normal(phi * x0, q))
mu1 = beta * mu0 + x1
y1 = numpyro.sample('y', dist.Normal(mu1, r))
numpyro.deterministic('y2', y1 * 2)
return (x1, mu1), (x1, y1)
mu0 = x0 = numpyro.sample('x_0', dist.Normal(0, q))
y0 = numpyro.sample('y_0', dist.Normal(mu0, r))
_, xy = scan(transition, (x0, mu0), jnp.arange(T))
x, y = xy
return jnp.append(x0, x), jnp.append(y0, y)
T = 10
num_samples = 100
kernel = NUTS(model)
mcmc = MCMC(kernel, 100, num_samples)
mcmc.run(jax.random.PRNGKey(0), T=T)
assert set(mcmc.get_samples()) == {'x', 'y', 'y2', 'x_0', 'y_0'}
mcmc.print_summary()
samples = mcmc.get_samples()
x = samples.pop('x')[0] # take 1 sample of x
# this tests for the composition of condition and substitute
# this also tests if we can use `vmap` for predictive.
future = 5
predictive = Predictive(numpyro.handlers.condition(model, {'x': x}),
samples,
return_sites=['x', 'y', 'y2'],
parallel=True)
result = predictive(jax.random.PRNGKey(1), T=T + future)
expected_shape = (num_samples, T + future)
assert result['x'].shape == expected_shape
assert result['y'].shape == expected_shape
assert result['y2'].shape == expected_shape
assert_allclose(result['x'][:, :T], jnp.broadcast_to(x, (num_samples, T)))
assert_allclose(result['y'][:, :T], samples['y'])
def run_inference(model, args, rng_key, X, Y):
start = time.time()
# demonstrate how to use different HMC initialization strategies
if args.init_strategy == "value":
init_strategy = init_to_value(values={"kernel_var": 1.0, "kernel_noise": 0.05, "kernel_length": 0.5})
elif args.init_strategy == "median":
init_strategy = init_to_median(num_samples=10)
elif args.init_strategy == "feasible":
init_strategy = init_to_feasible()
elif args.init_strategy == "sample":
init_strategy = init_to_sample()
elif args.init_strategy == "uniform":
init_strategy = init_to_uniform(radius=1)
kernel = NUTS(model, init_strategy=init_strategy)
mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains, thinning=args.thinning,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
mcmc.run(rng_key, X, Y)
mcmc.print_summary()
print('\nMCMC elapsed time:', time.time() - start)
return mcmc.get_samples()
def test_reuse_mcmc_run(jit_args, shape):
y1 = np.random.normal(3, 0.1, (100, ))
y2 = np.random.normal(-3, 0.1, (shape, ))
def model(y_obs):
mu = numpyro.sample("mu", dist.Normal(0.0, 1.0))
sigma = numpyro.sample("sigma", dist.HalfCauchy(3.0))
numpyro.sample("y", dist.Normal(mu, sigma), obs=y_obs)
# Run MCMC on zero observations.
kernel = NUTS(model)
mcmc = MCMC(kernel,
num_warmup=300,
num_samples=500,
jit_model_args=jit_args)
mcmc.run(random.PRNGKey(32), y1)
# Re-run on new data - should be much faster.
mcmc.run(random.PRNGKey(32), y2)
assert_allclose(mcmc.get_samples()["mu"].mean(), -3.0, atol=0.1)
def main(args):
print('Simulating data...')
(transition_prior, emission_prior, transition_prob, emission_prob,
supervised_categories, supervised_words, unsupervised_words) = simulate_data(
random.PRNGKey(1),
num_categories=args.num_categories,
num_words=args.num_words,
num_supervised_data=args.num_supervised,
num_unsupervised_data=args.num_unsupervised,
)
print('Starting inference...')
rng_key = random.PRNGKey(2)
start = time.time()
kernel = NUTS(semi_supervised_hmm)
mcmc = MCMC(kernel, args.num_warmup, args.num_samples)
mcmc.run(rng_key, transition_prior, emission_prior, supervised_categories,
supervised_words, unsupervised_words)
samples = mcmc.get_samples()
print_results(samples, transition_prob, emission_prob)
print('\nMCMC elapsed time:', time.time() - start)
def test_uniform_normal():
true_coef = 0.9
num_warmup, num_samples = 1000, 1000
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
with reparam(config={'loc': TransformReparam()}):
loc = numpyro.sample('loc', dist.Uniform(0, alpha))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
data = true_coef + random.normal(random.PRNGKey(0), (1000,))
kernel = NUTS(model=model)
mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
mcmc.warmup(random.PRNGKey(2), data, collect_warmup=True)
warmup_samples = mcmc.get_samples()
mcmc.run(random.PRNGKey(3), data)
samples = mcmc.get_samples()
assert len(warmup_samples['loc']) == num_warmup
assert len(samples['loc']) == num_samples
assert_allclose(jnp.mean(samples['loc'], 0), true_coef, atol=0.05)
def test_beta_bernoulli():
from numpyro.contrib.tfp import distributions as dist
warmup_steps, num_samples = (500, 2000)
def model(data):
alpha = jnp.array([1.1, 1.1])
beta = jnp.array([1.1, 1.1])
p_latent = numpyro.sample('p_latent', dist.Beta(alpha, beta))
numpyro.sample('obs', dist.Bernoulli(p_latent), obs=data)
return p_latent
true_probs = jnp.array([0.9, 0.1])
data = dist.Bernoulli(true_probs)(rng_key=random.PRNGKey(1), sample_shape=(1000, 2))
kernel = NUTS(model=model, trajectory_length=0.1)
mcmc = MCMC(kernel, num_warmup=warmup_steps, num_samples=num_samples)
mcmc.run(random.PRNGKey(2), data)
mcmc.print_summary()
samples = mcmc.get_samples()
assert_allclose(jnp.mean(samples['p_latent'], 0), true_probs, atol=0.05)
def run_inference(model, args, rng_key, X, Y):
start = time.time()
kernel = NUTS(model)
mcmc = MCMC(
kernel,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
)
mcmc.run(rng_key, X, Y)
mcmc.print_summary(exclude_deterministic=False)
samples = mcmc.get_samples()
summary_dict = summary(samples, group_by_chain=False)
print("\nMCMC elapsed time:", time.time() - start)
return summary_dict
def sample(
model,
num_samples,
num_warmup,
num_chains=2,
seed=0,
chain_method="parallel",
summary=True,
**kwargs,
):
"""Run the No-U-Turn sampler
Args:
model: an NumPyro model function
num_samples: number of samples to draw in each chain
num_warmup: number of samples to use for tuning in each chain
num_chains: number of chains to draw (default: {2})
**kwargs: other arguments to be passed to the model function
seed: random seed (default: {0})
chain_method: one of NumPyro's sampling methods — "parallel" / "sequential" /
"vectorized" (default: {"parallel"})
summary: print diagnostics, including the Effective sample size and the
Gelman-Rubin test (default: {True})
Returns:
mcmc: A fitted MCMC object
"""
rng_key = random.PRNGKey(seed)
kernel = NUTS(model)
# Note: sampling more than one chain doesn't show a progress bar
mcmc = MCMC(kernel,
num_warmup,
num_samples,
num_chains,
chain_method=chain_method)
mcmc.run(rng_key, **kwargs)
if summary:
mcmc.print_summary()
return mcmc
def test_correlated_mvn():
# This requires dense mass matrix estimation.
D = 5
warmup_steps, num_samples = 5000, 8000
true_mean = 0.
a = jnp.tril(0.5 * jnp.fliplr(jnp.eye(D)) + 0.1 * jnp.exp(random.normal(random.PRNGKey(0), shape=(D, D))))
true_cov = jnp.dot(a, a.T)
true_prec = jnp.linalg.inv(true_cov)
def potential_fn(z):
return 0.5 * jnp.dot(z.T, jnp.dot(true_prec, z))
init_params = jnp.zeros(D)
kernel = NUTS(potential_fn=potential_fn, dense_mass=True)
mcmc = MCMC(kernel, warmup_steps, num_samples)
mcmc.run(random.PRNGKey(0), init_params=init_params)
samples = mcmc.get_samples()
assert_allclose(jnp.mean(samples), true_mean, atol=0.02)
assert np.sum(np.abs(np.cov(samples.T) - true_cov)) / D**2 < 0.02
def test_binomial_stable_x64(with_logits):
# Ref: https://github.com/pyro-ppl/pyro/issues/1706
warmup_steps, num_samples = 200, 200
def model(data):
p = numpyro.sample('p', dist.Beta(1., 1.))
if with_logits:
logits = logit(p)
numpyro.sample('obs', dist.Binomial(data['n'], logits=logits), obs=data['x'])
else:
numpyro.sample('obs', dist.Binomial(data['n'], probs=p), obs=data['x'])
data = {'n': 5000000, 'x': 3849}
kernel = NUTS(model=model)
mcmc = MCMC(kernel, warmup_steps, num_samples)
mcmc.run(random.PRNGKey(2), data)
samples = mcmc.get_samples()
assert_allclose(jnp.mean(samples['p'], 0), data['x'] / data['n'], rtol=0.05)
if 'JAX_ENABLE_X64' in os.environ:
assert samples['p'].dtype == jnp.float64
def test_beta_bernoulli():
from tensorflow_probability.substrates.jax import distributions as tfd
num_warmup, num_samples = (500, 2000)
def model(data):
alpha = jnp.array([1.1, 1.1])
beta = jnp.array([1.1, 1.1])
p_latent = numpyro.sample("p_latent", tfd.Beta(alpha, beta))
numpyro.sample("obs", tfd.Bernoulli(p_latent), obs=data)
return p_latent
true_probs = jnp.array([0.9, 0.1])
data = tfd.Bernoulli(true_probs).sample(seed=random.PRNGKey(1),
sample_shape=(1000, 2))
kernel = NUTS(model=model, trajectory_length=0.1)
mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
mcmc.run(random.PRNGKey(2), data)
mcmc.print_summary()
samples = mcmc.get_samples()
assert_allclose(jnp.mean(samples["p_latent"], 0), true_probs, atol=0.05)
def main(args):
_, fetch = load_dataset(LYNXHARE, shuffle=False)
year, data = fetch() # data is in hare -> lynx order
# use dense_mass for better mixing rate
mcmc = MCMC(
NUTS(model, dense_mass=True),
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
)
mcmc.run(PRNGKey(1), N=data.shape[0], y=data)
mcmc.print_summary()
# predict populations
pop_pred = Predictive(model, mcmc.get_samples())(PRNGKey(2),
data.shape[0])["y"]
mu = jnp.mean(pop_pred, 0)
pi = jnp.percentile(pop_pred, jnp.array([10, 90]), 0)
plt.figure(figsize=(8, 6), constrained_layout=True)
plt.plot(year,
data[:, 0],
"ko",
mfc="none",
ms=4,
label="true hare",
alpha=0.67)
plt.plot(year, data[:, 1], "bx", label="true lynx")
plt.plot(year, mu[:, 0], "k-.", label="pred hare", lw=1, alpha=0.67)
plt.plot(year, mu[:, 1], "b--", label="pred lynx")
plt.fill_between(year, pi[0, :, 0], pi[1, :, 0], color="k", alpha=0.2)
plt.fill_between(year, pi[0, :, 1], pi[1, :, 1], color="b", alpha=0.3)
plt.gca().set(ylim=(0, 160),
xlabel="year",
ylabel="population (in thousands)")
plt.title("Posterior predictive (80% CI) with predator-prey pattern.")
plt.legend()
plt.savefig("ode_plot.pdf")
def test_inference_data_constant_data(self):
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
x1 = 10
x2 = 12
y1 = np.random.randn(10)
def model_constant_data(x, y1=None):
_x = numpyro.sample("x", dist.Normal(1, 3))
numpyro.sample("y1", dist.Normal(x * _x, 1), obs=y1)
nuts_kernel = NUTS(model_constant_data)
mcmc = MCMC(nuts_kernel, num_samples=10, num_warmup=2)
mcmc.run(PRNGKey(0), x=x1, y1=y1)
posterior = mcmc.get_samples()
posterior_predictive = Predictive(model_constant_data,
posterior)(PRNGKey(1), x1)
predictions = Predictive(model_constant_data, posterior)(PRNGKey(2),
x2)
inference_data = from_numpyro(
mcmc,
posterior_predictive=posterior_predictive,
predictions=predictions,
constant_data={"x1": x1},
predictions_constant_data={"x2": x2},
)
test_dict = {
"posterior": ["x"],
"posterior_predictive": ["y1"],
"sample_stats": ["diverging"],
"log_likelihood": ["y1"],
"predictions": ["y1"],
"observed_data": ["y1"],
"constant_data": ["x1"],
"predictions_constant_data": ["x2"],
}
fails = check_multiple_attrs(test_dict, inference_data)
assert not fails
def test_hmcecs_multiple_plates():
true_loc = jnp.array([0.3, 0.1, 0.9])
num_warmup, num_samples = 2, 2
data = true_loc + dist.Normal(jnp.zeros(3), jnp.ones(3)).sample(
random.PRNGKey(1), (1000,)
)
def model(data):
mean = numpyro.sample("mean", dist.Normal().expand((3,)).to_event(1))
with numpyro.plate("batch", data.shape[0], dim=-2, subsample_size=10):
sub_data = numpyro.subsample(data, 0)
with numpyro.plate("dim", 3):
numpyro.sample("obs", dist.Normal(mean, 1), obs=sub_data)
ref_params = {
"mean": true_loc + dist.Normal(true_loc, 5e-2).sample(random.PRNGKey(0))
}
proxy_fn = HMCECS.taylor_proxy(ref_params)
kernel = HMCECS(NUTS(model), proxy=proxy_fn)
mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
mcmc.run(random.PRNGKey(0), data)
def __init__(
self,
model,
num_warmup=2000,
num_samples=10000,
num_chains=1,
rng_key=0,
to_numpy: bool = True,
*args,
**kwargs,
):
self.model = model
self.num_warmup = num_warmup
self.num_samples = num_samples
self.num_chains = num_chains
self.rng_key, self.rng_key_ = random.split(random.PRNGKey(rng_key))
self.to_numpy = to_numpy
self.kernel = NUTS(model, **kwargs)
self.mcmc = MCMC(self.kernel,
num_warmup,
num_samples,
num_chains=num_chains)
def test_logistic_regression():
from numpyro.contrib.tfp import distributions as dist
N, dim = 3000, 3
num_warmup, num_samples = (1000, 1000)
data = random.normal(random.PRNGKey(0), (N, dim))
true_coefs = jnp.arange(1., dim + 1.)
logits = jnp.sum(true_coefs * data, axis=-1)
labels = dist.Bernoulli(logits=logits)(rng_key=random.PRNGKey(1))
def model(labels):
coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
logits = numpyro.deterministic('logits', jnp.sum(coefs * data, axis=-1))
return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup, num_samples)
mcmc.run(random.PRNGKey(2), labels)
mcmc.print_summary()
samples = mcmc.get_samples()
assert samples['logits'].shape == (num_samples, N)
assert_allclose(jnp.mean(samples['coefs'], 0), true_coefs, atol=0.22)
def test_gaussian_hmm():
dim = 4
num_steps = 10
@config_enumerate
def model(data):
with numpyro.plate("states", dim):
transition = numpyro.sample("transition",
dist.Dirichlet(jnp.ones(dim)))
emission_loc = numpyro.sample("emission_loc", dist.Normal(0, 1))
emission_scale = numpyro.sample("emission_scale",
dist.LogNormal(0, 1))
trans_prob = numpyro.sample("initialize",
dist.Dirichlet(jnp.ones(dim)))
for t, y in markov(enumerate(data)):
x = numpyro.sample("x_{}".format(t), dist.Categorical(trans_prob))
numpyro.sample("y_{}".format(t),
dist.Normal(emission_loc[x], emission_scale[x]),
obs=y)
trans_prob = transition[x]
def _generate_data():
transition_probs = np.random.rand(dim, dim)
transition_probs = transition_probs / transition_probs.sum(
-1, keepdims=True)
emissions_loc = np.arange(dim)
emissions_scale = 1.0
state = np.random.choice(3)
obs = [np.random.normal(emissions_loc[state], emissions_scale)]
for _ in range(num_steps - 1):
state = np.random.choice(dim, p=transition_probs[state])
obs.append(np.random.normal(emissions_loc[state], emissions_scale))
return np.stack(obs)
data = _generate_data()
nuts_kernel = NUTS(model)
mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=500)
mcmc.run(random.PRNGKey(0), data)
def run_inference(self,
model,
rng_key,
X,
Y,
hypers,
num_warmup=500,
num_chains=1,
num_samples=1000):
start = time.time()
kernel = NUTS(model)
mcmc = MCMC(kernel,
num_warmup,
num_samples,
num_chains=num_chains,
progress_bar=False
if "NUMPYRO_SPHINXBUILD" in os.environ else True)
mcmc.run(rng_key, X, Y, hypers)
mcmc.print_summary()
print('\nMCMC elapsed time:', time.time() - start)
return mcmc.get_samples()
def test_gaussian_mixture_model():
K, N = 3, 1000
def gmm(data):
mix_proportions = numpyro.sample("phi", dist.Dirichlet(jnp.ones(K)))
with numpyro.plate("num_clusters", K, dim=-1):
cluster_means = numpyro.sample("cluster_means", dist.Normal(jnp.arange(K), 1.))
with numpyro.plate("data", data.shape[0], dim=-1):
assignments = numpyro.sample("assignments", dist.Categorical(mix_proportions))
numpyro.sample("obs", dist.Normal(cluster_means[assignments], 1.), obs=data)
true_cluster_means = jnp.array([1., 5., 10.])
true_mix_proportions = jnp.array([0.1, 0.3, 0.6])
cluster_assignments = dist.Categorical(true_mix_proportions).sample(random.PRNGKey(0), (N,))
data = dist.Normal(true_cluster_means[cluster_assignments], 1.0).sample(random.PRNGKey(1))
nuts_kernel = NUTS(gmm)
mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=500)
mcmc.run(random.PRNGKey(2), data)
samples = mcmc.get_samples()
assert_allclose(samples["phi"].mean(0).sort(), true_mix_proportions, atol=0.05)
assert_allclose(samples["cluster_means"].mean(0).sort(), true_cluster_means, atol=0.2)
def __init__(
self,
model,
posterior=None,
num_warmup=2000,
num_samples=10000,
num_chains=1,
key=0,
*args,
**kwargs,
):
self.model = model
self.rng_key, self.rng_key_ = random.split(random.PRNGKey(key))
if posterior is not None:
self.mcmc = posterior
self.posterior = self.mcmc.get_samples()
else:
self.kernel = NUTS(model, **kwargs)
self.mcmc = MCMC(
self.kernel, num_warmup, num_samples, num_chains=num_chains
)
