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