def main() -> None:
# Data
x, t = _load_data()
num_train = int(len(x) * 0.8)
x_train = x[:num_train]
t_train = t[:num_train]
rng_key = random.PRNGKey(0)
rng_key, rng_key_prior, rng_key_infer, rng_key_posterior = random.split(
rng_key, 4)
# prior
predictive = infer.Predictive(model, num_samples=10)
prior_samples = predictive(rng_key_prior, t)
# Inference
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel, 100, 100)
mcmc.run(rng_key_infer, t_train, x_train)
posterior_samples = mcmc.get_samples()
# Posterior prediction
predictive = infer.Predictive(
model,
posterior_samples=posterior_samples,
return_sites=["x", "s0", "z", "trend", "weight"],
)
posterior_predictive = predictive(rng_key_posterior, t)
_save_results(x, prior_samples, posterior_samples, posterior_predictive,
num_train)
def sample(
self,
objective: Union[Objective, StochasticObjective],
prng_key: jnp.ndarray,
num_samples: int = 1,
parallel_chains: int = 1,
init_state: Optional[jnp.ndarray] = None,
) -> jnp.ndarray:
# Make sure stochastic objectives are evaluated on the full data.
if isinstance(objective, StochasticObjective):
objective = functools.partial(objective,
prng_key=prng_key,
deterministic=True)
if init_state is None:
init_state = jnp.squeeze(
jnp.zeros((
parallel_chains,
objective.dim,
)))
# Instantiate and run MCMC.
self.mcmc = infer.MCMC(
self.kernel(potential_fn=objective),
num_warmup=self.num_warmup,
num_samples=num_samples,
num_chains=parallel_chains,
progress_bar=False,
)
self.mcmc.run(prng_key, init_params=init_state)
return self.mcmc.get_samples()
def main() -> None:
x = _load_dataset()
rng_key = random.PRNGKey(0)
rng_key, rng_key_prior, rng_key_infer, rng_key_posterior = random.split(
rng_key, 4)
# prior
predictive = infer.Predictive(model, num_samples=10)
prior_samples = predictive(rng_key_prior, None, *x.shape, future_steps=20)
# Inference
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel, 100, 100)
mcmc.run(rng_key_infer, x)
posterior_samples = mcmc.get_samples()
# Posterior prediction
predictive = infer.Predictive(model, posterior_samples=posterior_samples)
posterior_predictive = predictive(rng_key_posterior,
None,
*x.shape,
future_steps=20)
_save_results(x, prior_samples, posterior_samples, posterior_predictive,
len(x))
def infer(
self,
data: xr.Dataset,
iterations: int,
num_warmup: int,
seed: int,
algorithm: str = "NUTS",
**infer_args,
) -> xr.Dataset:
"""
See :class:`numpyro.infer.mcmc.MCMC` for the MCMC inference API.
"""
if algorithm == "NUTS":
kernel = infer.NUTS(self.impl.model)
else:
raise ValueError(
f"{algorithm} algorithm not registered for NumPyro.")
mcmc = infer.MCMC(
kernel,
num_warmup=num_warmup,
num_samples=iterations - num_warmup,
**infer_args,
)
# Note that we need to run warmup separately to collect samples
# from the warmup phase.
rng_key_1, rng_key_2 = random.split(random.PRNGKey(seed))
mcmc.warmup(rng_key_1, data, collect_warmup=True)
warmup_samples = mcmc.get_samples()
mcmc.run(rng_key_2, data)
# merge samples from the warmup phase
samples = {
k: np.concatenate([warmup_samples[k], v])
for k, v in mcmc.get_samples().items()
}
return self.impl.extract_data_from_numpyro(samples)
def run_inference(model: Callable,
rng_key: np.ndarray) -> Dict[str, jnp.ndarray]:
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel, 1000, 1000, 1)
mcmc.run(rng_key)
return mcmc.get_samples()
def main() -> None:
x = jnp.concatenate([
np.random.randn(10),
np.random.randn(10) + 2,
np.random.randn(10) - 1,
np.random.randn(10) + 1,
])
x = x[:, None, None]
rng_key = random.PRNGKey(0)
rng_key, rng_key_prior, rng_key_infer, rng_key_posterior = random.split(
rng_key, 4)
# prior
predictive = infer.Predictive(model, num_samples=10)
prior_samples = predictive(rng_key_prior,
future_steps=20,
batch=10,
x_dim=1)
# Inference
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel, 100, 100)
mcmc.run(rng_key_infer, x)
posterior_samples = mcmc.get_samples()
# Posterior prediction
predictive = infer.Predictive(model, posterior_samples=posterior_samples)
posterior_predictive = predictive(rng_key_posterior, None, *x.shape, 10)
_save_results(x, prior_samples, posterior_samples, posterior_predictive)
def main() -> None:
# Data
x, betas, covariates = _load_data()
num_train = int(len(x) * 0.8)
x_train = x[:num_train]
c_train = covariates[:num_train]
rng_key = random.PRNGKey(0)
rng_key, rng_key_prior, rng_key_infer, rng_key_posterior = random.split(
rng_key, 4)
# prior
predictive = infer.Predictive(model, num_samples=10)
prior_samples = predictive(rng_key_prior, c_train)
# Inference
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel, 100, 100)
mcmc.run(rng_key_infer, c_train, x_train)
posterior_samples = mcmc.get_samples()
# Posterior prediction
posterior_given = posterior_samples.copy()
posterior_given.pop("weight")
predictive = infer.Predictive(model,
posterior_samples=posterior_given,
return_sites=["x", "weight"])
posterior_predictive = predictive(rng_key_posterior, covariates)
_save_results(x, betas, prior_samples, posterior_samples,
posterior_predictive, num_train)
def main() -> None:
_, y, x_missing = _load_dataset()
train_len = int(len(y) * 0.8)
x_train = x_missing[:train_len]
y_train = y[:train_len]
num_chains = 1
numpyro.set_platform("cpu")
numpyro.set_host_device_count(num_chains)
rng_key = random.PRNGKey(0)
rng_key, rng_key_posterior, rng_key_prior = random.split(rng_key, 3)
predictive = infer.Predictive(bayesian_regression, num_samples=500)
prior = predictive(rng_key_prior, x_train)
kernel = infer.NUTS(bayesian_regression)
mcmc = infer.MCMC(kernel,
num_warmup=1000,
num_samples=1000,
num_chains=num_chains)
mcmc.run(rng_key, x_train, y_train)
posterior_samples = mcmc.get_samples()
predictive = infer.Predictive(bayesian_regression,
posterior_samples=posterior_samples)
posterior_predictive = predictive(rng_key_posterior, x_missing)
_save_results(y, mcmc, prior, posterior_samples, posterior_predictive)
def fit(
self,
df: pd.DataFrame,
sampler: str = "NUTS",
rng_key: np.ndarray = None,
sampler_kwargs: typing.Dict[str, typing.Any] = None,
**mcmc_kwargs,
):
"""Fit the model to a DataFrame.
Parameters
----------
df : pd.DataFrame
Source dataframe.
sampler : str
Numpyro sampler name. Default NUTS
rng_key : two-element ndarray.
Optional rng key, will be randomly splitted if not provided.
sampler_kwargs :
Passed to the numpyro sampler selected.
**mcmc_kwargs :
Passed to numpyro.infer.MCMC
Returns
-------
Model
The fitted model.
"""
if self.fitted:
raise exceptions.AlreadyFittedError(self)
if sampler.upper() not in ("NUTS", "HMC"):
raise ValueError("Invalid sampler, try NUTS or HMC.")
sampler = getattr(infer, sampler.upper())
# store fit df
self.df = df
# set up mcmc
_mcmc_kwargs = dict(num_warmup=500, num_samples=1000)
_mcmc_kwargs.update(mcmc_kwargs)
_sampler_kwargs = dict(model=self.model)
_sampler_kwargs.update(sampler_kwargs or {})
mcmc = infer.MCMC(sampler(**_sampler_kwargs), **_mcmc_kwargs)
# do it
rng_key_ = (self.split_rand_key()
if rng_key is None else rng_key.astype("uint32"))
mcmc.run(rng_key_, df=df)
# store results
self.samples = mcmc.get_samples(group_by_chain=True)
self.fitted = True
return self
def main(args: argparse.Namespace) -> None:
_, y, x_missing = _load_dataset()
batch, x_dim = x_missing.shape
train_len = int(len(y) * 0.8)
x_train = x_missing[:train_len]
y_train = y[:train_len]
numpyro.set_platform("cpu")
numpyro.set_host_device_count(args.num_chains)
rng_key = random.PRNGKey(1)
rng_key, rng_key_prior, rng_key_posterior, rng_key_pca_pred = random.split(
rng_key, 4)
predictive = infer.Predictive(pca_regression, num_samples=500)
prior = predictive(rng_key_prior, batch=batch, x_dim=x_dim)
kernel = infer.NUTS(pca_regression)
mcmc = infer.MCMC(
kernel,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
)
mcmc.run(rng_key_posterior, x_train, y_train)
posterior_samples = mcmc.get_samples()
posterior_without_z = posterior_samples.copy()
posterior_without_z.pop("z")
predictive = infer.Predictive(pca_regression,
posterior_samples=posterior_without_z)
posterior_predictive = predictive(rng_key_pca_pred,
batch=batch,
x_dim=x_dim)
_save_results(
y,
mcmc,
prior,
posterior_samples,
posterior_predictive,
var_names=["phi", "eta", "theta", "sigma"],
)
def main(args: argparse.Namespace) -> None:
model = model_dict[args.model]
_, fetch = load_dataset(JSB_CHORALES, split="train", shuffle=False)
lengths, sequences = fetch()
# Remove never used data dimension to reduce computation time
present_notes = (sequences == 1).sum(0).sum(0) > 0
sequences = sequences[..., present_notes]
batch, seq_len, data_dim = sequences.shape
rng_key = random.PRNGKey(0)
rng_key, rng_key_prior, rng_key_pred = random.split(rng_key, 3)
predictive = infer.Predictive(model, num_samples=10)
prior_samples = predictive(rng_key_prior,
batch=batch,
seq_len=seq_len,
data_dim=data_dim,
future_steps=20)
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel, args.num_warmup, args.num_samples,
args.num_chains)
mcmc.run(rng_key, sequences, lengths)
posterior_samples = mcmc.get_samples()
predictive = infer.Predictive(model, posterior_samples)
predictive_samples = predictive(rng_key_pred,
sequences,
lengths,
future_steps=10)
path = pathlib.Path("./data/hmm_enum")
path.mkdir(exist_ok=True)
jnp.savez(path / "prior_samples.npz", **prior_samples)
jnp.savez(path / "posterior_samples.npz", **posterior_samples)
jnp.savez(path / "predictive_samples.npz", **predictive_samples)
def test_mcmc_model_side_enumeration(model, temperature):
mcmc = infer.MCMC(infer.NUTS(model), 0, 1)
mcmc.run(random.PRNGKey(0))
mcmc_data = {
k: v[0] for k, v in mcmc.get_samples().items() if k in ["loc", "scale"]
}
# MAP estimate discretes, conditioned on posterior sampled continous latents.
model = handlers.seed(model, rng_seed=1)
actual_trace = handlers.trace(
infer_discrete(
# TODO support replayed sites in infer_discrete.
# handlers.replay(config_enumerate(model), mcmc_trace),
handlers.condition(config_enumerate(model), mcmc_data),
temperature=temperature,
rng_key=random.PRNGKey(1),
)
).get_trace()
# Check site names and shapes.
expected_trace = handlers.trace(model).get_trace()
assert set(actual_trace) == set(expected_trace)
def inference(
model: Callable,
sigma: np.ndarray,
y: np.ndarray,
rng_key: np.ndarray,
*,
num_warmup: int = 500,
num_samples: int = 1000,
num_chains: int = 1,
verbose: bool = True,
) -> infer.MCMC:
kernel = infer.NUTS(model)
mcmc = infer.MCMC(kernel,
num_warmup=num_warmup,
num_samples=num_samples,
num_chains=num_chains)
mcmc.run(rng_key, sigma, y)
if verbose:
mcmc.print_summary()
return mcmc
