def main(args: argparse.Namespace) -> None:
numpyro.set_platform(args.device)
numpyro.set_host_device_count(args.num_chains)
_, fetch_train = load_dataset(BASEBALL, split="train", shuffle=False)
train, player_names = fetch_train()
_, fetch_test = load_dataset(BASEBALL, split="test", shuffle=False)
test, _ = fetch_test()
at_bats = train[:, 0]
hits = train[:, 1]
at_bats_test = test[:, 0]
hist_test = test[:, 1]
model_list = [fully_pooled, not_pooled, partially_pooled, partially_pooled_with_logit]
for i, model in enumerate(model_list, 1):
rng_key, rng_key_predict = random.split(random.PRNGKey(i))
posterior = run_inference(
model,
at_bats,
hits,
rng_key,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
algo_name=args.algo_name,
)
predictions = predict(model, at_bats, posterior, rng_key_predict)["obs"]
print_results(model.__name__, predictions, at_bats, hits, player_names, is_train=True)
print_results(
model.__name__, predictions, at_bats_test, hist_test, player_names, is_train=False
)
def MIPS_24(prior,num_samples=500,num_warmup=500,num_chains=4,chain_method='parallel'):
numpyro.set_host_device_count(4)
nuts_kernel = NUTS(mips_model)
mcmc = MCMC(nuts_kernel, num_samples=num_samples, num_warmup=num_warmup,num_chains=num_chains,chain_method=chain_method)
rng_key = random.PRNGKey(0)
mcmc.run(rng_key, [prior])
return mcmc
def main(args: argparse.Namespace) -> None:
numpyro.set_platform(args.device)
numpyro.set_host_device_count(args.num_chains)
(
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=args.num_supervised,
num_unsupservised=args.num_unsupervised,
)
rng_key = random.PRNGKey(2)
posterior = inference(
semi_supervised_hmm,
args.num_categories,
args.num_words,
supervised_categories,
supervised_words,
unsupervised_words,
rng_key,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
)
print_results(posterior, transition_prob, emission_prob)
plot_results(posterior, transition_prob)
def __init__(self, random_state=0, num_bins=3, tot_init_items=1000, bin_values=[1.0,1.0,1.0], num_users=50):
self.num_bins = num_bins # The number of bins is roughly equivalent to the number of arms.
self.bin_values = bin_values # The bin values correspond to the potential reward if the user is provided the item and they purchase it.
self.rng_key = random.PRNGKey(random_state)
# Create initial allotments of bins where the lowest value items will be plentiful while the high value items may be scarce
bin_allotments = sorted(random.randint(self.rng_key,(self.num_bins,),25,75))[::-1]
bin_allotments /= sum(bin_allotments) # Normalize to get the percentage of items that are assigned to each bin initially
self.init_bins = onp.round(tot_init_items*bin_allotments) # Place most items in first, low value bin and decrease from there.
self.num_users = num_users
# Set initialization for BNN parameters
self.bnn_dx = num_bins+1 # Input dimensions
self.bnn_dh = 5 # Size of hidden layer
self.bnn_dy = 1 # Size of output dimensions
self.bnn_warm_up = 500 # Number of warmup runs for MCMC
self.bnn_num_samples = 2500 # Number of BNN Samples
self.bnn_num_chains = 1 # Number of MCMC chains
self.bnn_device = 'cpu'
numpyro.set_platform(self.bnn_device)
numpyro.set_host_device_count(self.bnn_num_chains)
self.restock() # Fill bins
self._generate_user_preferences()
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 select_device(use_gpu, num_chains):
if use_gpu:
try:
numpyro.set_platform('gpu')
numpyro.set_host_device_count(1)
except RuntimeError as e:
warnings.warn(f'No GPU found: {e}')
numpyro.set_platform('cpu')
else:
numpyro.set_platform('cpu')
numpyro.set_host_device_count(min((num_chains, os.cpu_count())))
# Sanity check
jax.lib.xla_bridge.get_backend().platform
def all_bands(priors,
num_samples=500,
num_warmup=500,
num_chains=4,
chain_method='parallel'):
numpyro.set_host_device_count(4)
nuts_kernel = NUTS(spire_model)
mcmc = MCMC(nuts_kernel,
num_samples=num_samples,
num_warmup=num_warmup,
num_chains=num_chains,
chain_method=chain_method)
rng_key = random.PRNGKey(0)
mcmc.run(rng_key, priors)
return mcmc
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"],
)
ax.plot(X_test[:, 1], mean_prediction, 'blue', ls='solid', lw=2.0)
ax.set(xlabel="X", ylabel="Y", title="Mean predictions with 90% CI")
plt.savefig('bnn_plot.pdf')
plt.close()
if __name__ == "__main__":
assert numpyro.__version__.startswith('0.2.0')
parser = argparse.ArgumentParser(
description="Bayesian neural network example")
parser.add_argument("-n",
"--num-samples",
nargs="?",
default=2000,
type=int)
parser.add_argument("--num-warmup", nargs='?', default=1000, type=int)
parser.add_argument("--num-chains", nargs='?', default=1, type=int)
parser.add_argument("--num-data", nargs='?', default=100, type=int)
parser.add_argument("--num-hidden", nargs='?', default=5, type=int)
parser.add_argument("--device",
default='cpu',
type=str,
help='use "cpu" or "gpu".')
args = parser.parse_args()
numpyro.set_platform(args.device)
numpyro.set_host_device_count(args.num_chains)
main(args)
about things in a gaussian model.
"""
import json
import jax.numpy as np
import numpy as onp
import numpyro
import pandas as pd
import pytest
from numpyro import distributions as dist
import shabadoo.exceptions as exceptions
from shabadoo import Normal
try:
numpyro.set_host_device_count(2)
except Exception:
print("Failed to set device count")
def test_single_coef_is_about_right():
"""Test that a single coef model which has a known value gets there.
This is a good round trip test that fitting works.
"""
# going to make a coef be the mean of y
df = pd.DataFrame(dict(y=[1, 2, 2, 2, 2, 3] * 100))
class Model(Normal):
dv = "y"
features = dict(mu=dict(transformer=1, prior=dist.Normal(0, 5)))
_author__ = 'pdh21'
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
import jax.numpy as jnp
from jax import random
import numpy as np
import jax
import os
from xidplus.numpyro_fit.misc import sp_matmul
numpyro.set_host_device_count(os.cpu_count())
def pacs_model(priors):
pointing_matrices = [([p.amat_row, p.amat_col], p.amat_data)
for p in priors]
flux_lower = np.asarray([p.prior_flux_lower for p in priors]).T
flux_upper = np.asarray([p.prior_flux_upper for p in priors]).T
bkg_mu = np.asarray([p.bkg[0] for p in priors]).T
bkg_sig = np.asarray([p.bkg[1] for p in priors]).T
with numpyro.plate('bands', len(priors)):
sigma_conf = numpyro.sample('sigma_conf', dist.HalfCauchy(1.0, 0.5))
bkg = numpyro.sample('bkg', dist.Normal(bkg_mu, bkg_sig))
with numpyro.plate('nsrc', priors[0].nsrc):
src_f = numpyro.sample('src_f',
dist.Uniform(flux_lower, flux_upper))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="numpyro model")
parser.add_argument(
"-n",
"--num-samples",
nargs="?",
default=10000,
type=int,
)
parser.add_argument("--num-warmup", nargs="?", default=5000, type=int)
parser.add_argument("--num-chains", nargs="?", default=2, type=int)
parser.add_argument(
"--device",
default="cpu",
type=str,
help='"cpu" or "gpu"?',
)
parser.add_argument("--num-cores", nargs="?", default=2, type=int)
parser.add_argument("--rng_seed",
default=1,
type=int,
help="random number generator seed")
args = parser.parse_args()
numpyro.set_platform(args.device)
numpyro.set_host_device_count(args.num_cores)
main(args)
def generate_posterior(self,
num_dimensions=20,
known_active_dimensions=3,
num_data=100,
labels=None,
num_samples=1000,
num_warmup=500,
num_chains=1,
device='gpu'):
X = self.X
Y = self.Y
hypers = self.hypers
sigma = hypers['sigma']
# args -- needed: num-chains
# num_dimensions
# active dimensions
if labels is None:
labs = [str(_) for _ in range(X.shape[1])]
else:
labs = labels
# TODO:set up Numpyro device (should maybe do this in main script ?)
numpyro.set_platform(device)
numpyro.set_host_device_count(num_chains)
# do inference
rng_key = random.PRNGKey(self.seed)
samples = self.run_inference(self.model,
rng_key,
X,
Y,
hypers,
num_warmup=num_warmup,
num_chains=num_chains,
num_samples=num_samples)
# compute the mean and square root variance of each coefficient theta_i
means, stds = vmap(
lambda dim: self.analyze_dimension(samples, X, Y, dim, hypers))(
np.arange(X.shape[1]))
num_dims = len(means)
active_dimensions = []
for dim, (mean, std) in enumerate(zip(means, stds)):
# we mark the dimension as inactive if the interval [mean - 3 * std, mean + 3 * std] contains zero
lower, upper = mean - sigma * std, mean + sigma * std
inactive = "inactive" if lower < 0.0 and upper > 0.0 else "active"
if inactive == "active":
active_dimensions.append(dim)
print("[dimension %02d/%02d] %s:\t%.2e +- %.2e" %
(dim + 1, X.shape[1], inactive, mean, std))
print("Identified a total of %d active dimensions." %
(len(active_dimensions)))
# Compute the mean and square root variance of coefficients theta_ij for i,j active dimensions.
# Note that the resulting numbers are only meaningful for i != j.
if len(active_dimensions) > 0:
dim_pairs = np.array(
list(itertools.product(active_dimensions, active_dimensions)))
means, stds = vmap(
lambda dim_pair: self.analyze_pair_of_dimensions(
samples, X, Y, dim_pair[0], dim_pair[1], hypers))(
dim_pairs)
# print(dim_pairs)
dim_pair_arr = []
dim_pair_index = num_dims - 1
dim_pair_name = []
pair_labs = []
for dim_pair, mean, std in zip(dim_pairs, means, stds):
dim1, dim2 = dim_pair
if dim1 >= dim2:
continue
dim_pair_index += 1
lower, upper = mean - sigma * std, mean + sigma * std
if not (lower < 0.0 and upper > 0.0):
dim_pair_arr.append(dim_pair_index)
dim_pair_name.append('%d and %d' % (dim1 + 1, dim2 + 1))
format_str = "Identified pairwise interaction between dimensions %d and %d: %.2e +- %.2e"
print(format_str % (dim1 + 1, dim2 + 1, mean, std))
pair_labs.append(str(labs[dim1]) + ' + ' + str(labs[dim2]))
# Draw a single sample of coefficients theta from the posterior, where we return all singleton
# coefficients theta_i and pairwise coefficients theta_ij for i, j active dimensions. We use the
# final MCMC sample obtained from the HMC sampler.
## Get posterior samples from the sample_theta_space_modified() method
thetas = self.sample_theta_posterior(
X, Y, active_dimensions, samples['msq'][-1],
samples['lambda'][-1], samples['eta1'][-1],
samples['xisq'][-1], hypers['c'], samples['var_obs'][-1],
num_samples, dim_pair_arr)
print("Active dimensions: " + str(active_dimensions))
## Visualize the posterior from the example with corner
labels = ['dim ' + str(i) for i in active_dimensions]
active_dimensions = active_dimensions + dim_pair_arr
if len(dim_pair_name) != 0:
for n in range(len(dim_pair_name)):
labels.append('dim ' + dim_pair_name[n])
#fig = corner.corner(thetas, labels = labels);
return active_dimensions, thetas, labels, pair_labs
else:
return active_dimensions, [], []
import jax.numpy as jnp
import jax.numpy as np # it is the same anyways
from jax import random, vmap
import numpyro
import numpyro.distributions as dist
from numpyro.infer import SVI, Trace_ELBO, MCMC, NUTS
numpyro.set_host_device_count(4) # MCMC chain
def mcmc_sample(data, num_samples=2000, num_chains=4, num_warmup=1000):
# data = dict(W=W, L=L)
mcmc = MCMC(NUTS(model), num_warmup=1000, num_samples=2000, num_chains=4)
mcmc.run(random.PRNGKey(0), **data)
post = mcmc.get_samples()
return post
class Sampler():
def __init__(self, model, data=None):
self.data = data
self.num_warmup = 1000
self.num_samples = 2000
self.num_chains = 4
self.mcmc = MCMC(NUTS(model),
num_warmup=self.num_warmup,
num_samples=self.num_samples,
num_chains=self.num_chains)
self.data = data
def main(params):
#X, Y, expected_thetas, expected_pairwise = get_data(N=args.num_data, P=args.num_dimensions,
# S=args.active_dimensions)
fxdata5 = pd.read_csv('fxdata5.csv', header=0)
df = fxdata5.copy()
df = df.iloc[-64:,:]
drop_cols = list(df.columns[df.columns.str.contains('Dummy')])
drop_cols = drop_cols + list(df.columns[df.columns.str.contains('LogReturn')])
drop_cols = drop_cols + list(df.columns[df.columns.str.contains('Spot')])
# drop_cols = drop_cols + list(df.columns[df.columns.str.contains('HAR')])
df.drop(drop_cols, axis=1, inplace=True)
target_col = 'EURGBP_log_RealVol'
df['Target_' + target_col] = df[target_col].shift(1).copy()
cols = df.columns.tolist()
cols.insert(0, cols.pop(cols.index('Target_' + target_col)))
df = df[cols].set_index('Date').dropna()
df[df.columns] = StandardScaler().fit_transform(df) # scale
Y = np.array(df.iloc[:,0])
X = np.array(df.iloc[:,1:])
parser = argparse.ArgumentParser(description="Gaussian Process example")
parser.add_argument("-n", "--num-samples", nargs="?", default=8000, type=int)
parser.add_argument("--num-warmup", nargs='?', default=500, type=int)
parser.add_argument("--num-chains", nargs='?', default=1, type=int)
parser.add_argument("--num-data", nargs='?', default=len(X), type=int)
parser.add_argument("--num-dimensions", nargs='?', default=len(X.T), type=int)
parser.add_argument("--active-dimensions", nargs='?', default=3, type=int)
parser.add_argument("--device", default='cpu', type=str, help='use "cpu" or "gpu".')
args = parser.parse_args(args=[])
numpyro.set_platform(args.device)
numpyro.set_host_device_count(args.num_chains)
# setup hyperparameters
hypers = {'expected_sparsity': max(1.0, args.num_dimensions / 2),
}
# add in hyperopt params
hypers.update(params)
# do inference
rng_key = random.PRNGKey(0)
samples = run_inference(model, args, rng_key, X, Y, hypers)
# compute the mean and square root variance of each coefficient theta_i
means, stds = vmap(lambda dim: analyze_dimension(samples, X, Y, dim, hypers))(np.arange(args.num_dimensions))
num_dims = len(means)
# print("Coefficients theta_1 to theta_%d used to generate the data:" % args.active_dimensions, expected_thetas)
# print("The single quadratic coefficient theta_{1,2} used to generate the data:", expected_pairwise)
active_dimensions = []
for dim, (mean, std) in enumerate(zip(means, stds)):
# we mark the dimension as inactive if the interval [mean - 3 * std, mean + 3 * std] contains zero
lower, upper = mean - 1.5 * std, mean + 1.5 * std
inactive = "inactive" if lower < 0.0 and upper > 0.0 else "active"
if inactive == "active":
active_dimensions.append(dim)
print("[dimension %02d/%02d] %s:\t%.2e +- %.2e" % (dim + 1, args.num_dimensions, inactive, mean, std))
print("Identified a total of %d active dimensions; expected %d." % (len(active_dimensions),
args.active_dimensions))
# Compute the mean and square root variance of coefficients theta_ij for i,j active dimensions.
# Note that the resulting numbers are only meaningful for i != j.
if len(active_dimensions) > 0:
dim_pairs = np.array(list(itertools.product(active_dimensions, active_dimensions)))
means, stds = vmap(lambda dim_pair: analyze_pair_of_dimensions(samples, X, Y,
dim_pair[0], dim_pair[1], hypers))(dim_pairs)
# print(dim_pairs)
dim_pair_arr = []
dim_pair_index = num_dims -1
dim_pair_name = []
for dim_pair, mean, std in zip(dim_pairs, means, stds):
dim1, dim2 = dim_pair
if dim1 >= dim2:
continue
dim_pair_index += 1
lower, upper = mean - 1.5 * std, mean + 1.5 * std
if not (lower < 0.0 and upper > 0.0):
dim_pair_arr.append(dim_pair_index)
dim_pair_name.append('%d and %d'%(dim1 + 1, dim2 + 1))
format_str = "Identified pairwise interaction between dimensions %d and %d: %.2e +- %.2e"
print(format_str % (dim1 + 1, dim2 + 1, mean, std))
active_dimensions = active_dimensions + dim_pair_arr
return np.exp(-len(active_dimensions))
