def guide(home_id, away_id, score1_obs=None, score2_obs=None):
mu_locs = pyro.param("mu_loc", torch.tensor(0.0).expand(4))
mu_scales = pyro.param(
"mu_scale",
torch.tensor(0.1).expand(4),
constraint=constraints.positive,
)
pyro.sample("alpha", dist.Normal(mu_locs[0], mu_scales[0]))
pyro.sample("sd_att", dist.LogNormal(mu_locs[1], mu_scales[1]))
pyro.sample("sd_def", dist.LogNormal(mu_locs[2], mu_scales[2]))
pyro.sample("home", dist.Normal(mu_locs[3], mu_scales[3]))
nt = len(np.unique(home_id))
mu_team_locs = pyro.param("mu_team_loc", torch.tensor(0.0).expand(2, nt))
mu_team_scales = pyro.param(
"mu_team_scale",
torch.tensor(0.1).expand(2, nt),
constraint=constraints.positive,
)
with pyro.plate("plate_teams", nt):
pyro.sample("attack", dist.Normal(mu_team_locs[0], mu_team_scales[0]))
pyro.sample("defend", dist.Normal(mu_team_locs[1], mu_team_scales[1]))
def guide_ret(*args, **kwargs):
I, N = self._data['data'].shape
seg_mean = torch.mean(self._data['data'] / self._data['pld'].reshape([I, 1]), axis=1)
batch = N if self._params['batch_size'] else self._params['batch_size']
param_weights = pyro.param("param_mixture_weights", lambda: torch.ones(self._params['K']) / self._params['K'],
constraint=constraints.simplex)
cnv_mean = pyro.param("param_cnv_probs", lambda: self.create_gaussian_init_values(),
constraint=constraints.positive)
cnv_var = pyro.param("param_cnv_var", lambda: torch.ones(I) * self._params['cnv_sd'],
constraint=constraints.positive)
seg_var = pyro.param("param_seg_var", lambda: torch.ones(I) * self._params['cnv_sd'],
constraint=constraints.positive)
seg_mean = pyro.param("param_seg_mean", lambda: seg_mean)
gamma_scale = pyro.param("param_norm_factor", lambda: torch.sum(self._data['data'], axis = 0) / torch.sum(seg_mean),
constraint=constraints.positive)
pyro.sample('mixture_weights', dist.Dirichlet(param_weights))
with pyro.plate('segments', I):
pyro.sample('segment_mean', dist.LogNormal(torch.log(seg_mean) - seg_var ** 2 / 2, seg_var))
with pyro.plate('components', self._params['K']):
pyro.sample('cnv_probs', dist.LogNormal(torch.log(cnv_mean) - cnv_var ** 2 / 2, cnv_var))
with pyro.plate("data2", N, batch):
pyro.sample('norm_factor', dist.Delta(gamma_scale))
def model(self, zero_data, covariates):
data_dim = zero_data.size(-1)
feature_dim = covariates.size(-1)
bias = pyro.sample("bias", dist.Normal(5.0, 10.0).expand((data_dim,)).to_event(1))
weight = pyro.sample("weight", dist.Normal(0.0, 5).expand((feature_dim,)).to_event(1))
# We'll sample a time-global scale parameter outside the time plate,
# then time-local iid noise inside the time plate.
drift_scale = pyro.sample("drift_scale",
dist.LogNormal(-20, 5.0).expand((1,)).to_event(1))
with self.time_plate:
# We'll use a reparameterizer to improve variational fit. The model would still be
# correct if you removed this context manager, but the fit appears to be worse.
with poutine.reparam(config={"drift": LocScaleReparam()}):
drift = pyro.sample("drift", dist.Normal(zero_data.double(), drift_scale.double()).to_event(1))
# After we sample the iid "drift" noise we can combine it in any time-dependent way.
# It is important to keep everything inside the plate independent and apply dependent
# transforms outside the plate.
motion = drift.cumsum(-2) # A Brownian motion.
# The prediction now includes three terms.
prediction = motion + bias + (weight * covariates).sum(-1, keepdim=True)
assert prediction.shape[-2:] == zero_data.shape
# Construct the noise distribution and predict.
noise_scale = pyro.sample("noise_scale", dist.LogNormal(0.0, 1.0).expand((1,)).to_event(1))
noise_dist = dist.Normal(0, noise_scale)
self.predict(noise_dist, prediction)
def state_space_model(data, N=1, T=2, prior_drift=0., verbose=False):
# global rvs
drift = pyro.sample('drift', dist.Normal(prior_drift, 1))
vol = pyro.sample('vol', dist.LogNormal(0, 1))
uncert = pyro.sample('uncert', dist.LogNormal(-5, 1))
if verbose:
print(f"Using drift = {drift}, vol = {vol}, uncert = {uncert}")
# the latent time series you want to infer
# since you want to output this, we initialize a vector where you'll
# save the inferred values
latent = torch.empty((T, N)) # +1 comes from hidden initial condition
# I think you want to plate out the same state space model for N different obs
with pyro.plate('data_plate', N) as n:
x0 = pyro.sample('x0',
dist.Normal(drift,
vol)) # or whatever your IC might be
latent[0, n] = x0
# now comes the markov part, as you correctly noted
for t in pyro.markov(range(1, T)):
x_t = pyro.sample(f"x_{t}",
dist.Normal(latent[t - 1, n] + drift, vol))
y_t = pyro.sample(f"y_{t}",
dist.Normal(x_t, uncert),
obs=data[t - 1, n] if data is not None else None)
latent[t, n] = x_t
return pyro.deterministic('latent', latent)
def guide(self, X):
_id = self._id
K = self.K
N, D = X.shape
loc_locinit = self.param_init[f'loc_loc_init_{_id}']
loc_scaleinit = self.param_init[f'loc_scale_init_{_id}']
cov_diag_loc_init = self.param_init[f'cov_diag_loc_init_{_id}']
cov_diag_scale_init = self.param_init[f'cov_diag_scale_init_{_id}']
cov_factor_loc = self.param_init[f'cov_factor_loc_init_{_id}']
cov_factor_scale = self.param_init[f'cov_factor_scale_init_{_id}']
with pyro.plate(f'D_{_id}', D, dim=-1):
loc_loc = pyro.param(f'loc_loc_{_id}', loc_locinit)
loc_scale = pyro.param(f'loc_scale_{_id}', loc_scaleinit, constraint=constraints.positive)
loc = pyro.sample(f'loc_{_id}', dist.LogNormal(loc_loc, loc_scale))
cov_diag_loc = pyro.param(f'cov_diag_loc_{_id}', cov_diag_loc_init)
cov_diag_scale = pyro.param(f'cov_diag_scale_{_id}', cov_diag_scale_init, constraint=constraints.positive)
cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale))
cov_diag = cov_diag*torch.ones(D)
cov_diag = cov_diag + jitter
# sample variables
with pyro.plate(f'K_{_id}', K):
cov_factor_loc = pyro.param(f'cov_factor_loc_{_id}', cov_factor_loc_init)
cov_factor_scale = pyro.param(f'cov_factor_scale_{_id}', cov_factor_scale_init, constraint=constraints.positive)
cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(cov_factor_loc, cov_factor_scale))
cov_factor = cov_factor.transpose(-2,-1)
return loc, cov_factor, cov_diag
def guide(self, X):
_id = self._id
N, D = X.shape
global_shrinkage_loc_init = self.param_init[f'global_shrinkage_loc_init_{_id}']
global_shrinkage_scale_init = self.param_init[f'global_shrinkage_scale_init_{_id}']
local_shrinkage_loc_init = self.param_init[f'local_shrinkage_loc_init_{_id}']
local_shrinkage_scale_init = self.param_init[f'local_shrinkage_scale_init_{_id}']
cov_diag_loc_init = self.param_init[f'cov_diag_loc_init_{_id}']
cov_diag_scale_init = self.param_init[f'cov_diag_scale_init_{_id}']
cov_factor_loc_init = self.param_init[f'cov_factor_loc_init_{_id}']
global_shrinkage_loc = pyro.param(f'global_shrinkage_loc_{_id}', global_shrinkage_loc_init)
global_shrinkage_scale = pyro.param(f'global_shrinkage_scale_{_id}', global_shrinkage_scale_init, constraint=constraints.positive)
tau = pyro.sample(f'global_shrinkage_{_id}', dist.LogNormal(global_shrinkage_loc,global_shrinkage_scale))
b = pyro.sample('b', dist.InverseGamma(0.5,1./torch.ones(D)**2).to_event(1))
local_shrinkage_loc = pyro.param(f'local_shrinkage_loc_{_id}', local_shrinkage_loc_init)
local_shrinkage_scale = pyro.param(f'local_shrinkage_scale_{_id}', local_shrinkage_scale_init, constraint=constraints.positive)
lambdasquared = pyro.sample(f'local_shrinkage_{_id}', dist.LogNormal(local_shrinkage_loc,local_shrinkage_scale).to_event(1))
cov_diag_loc = pyro.param(f'cov_diag_loc_{_id}', cov_diag_loc_init)
cov_diag_scale = pyro.param(f'cov_diag_scale_{_id}', cov_diag_scale_init, constraint=constraints.positive)
cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale).to_event(1))
cov_diag = cov_diag*torch.ones(D)
lambdasquared = lambdasquared.squeeze()
if lambdasquared.dim() == 1:
cov_factor_scale = torch.ger(torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,)))
else:
cov_factor_scale = torch.einsum('bp, br->bpr', torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,)))
cov_factor_loc = pyro.param(f'cov_factor_loc_{_id}', cov_factor_loc_init)
cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(cov_factor_loc, cov_factor_scale).to_event(2))
return tau, lambdasquared, cov_factor, cov_diag
def model(self, X):
_id = self._id
K = self.K
N, D = X.shape
loc_locinit = self.param_init[f'loc_prior_loc_init_{_id}']
loc_scaleinit = self.param_init[f'loc_prior_scale_init_{_id}']
cov_diag_loc_init = self.param_init[f'cov_diag_prior_loc_init_{_id}']
cov_diag_scale_init = self.param_init[f'cov_diag_prior_scale_init_{_id}']
cov_factor_loc_init = self.param_init[f'cov_factor_prior_loc_init_{_id}']
cov_factor_scale_init = self.param_init[f'cov_factor_prior_scale_init_{_id}']
with pyro.plate(f'D_{_id}', D):
loc_loc = pyro.param(f'loc_loc_prior_{_id}', loc_locinit)
loc_scale = pyro.param(f'loc_scale_prior_{_id}', loc_scaleinit, constraint=constraints.positive)
loc = pyro.sample(f'loc_{_id}', dist.LogNormal(loc_loc, loc_scale))
cov_diag_loc = pyro.param(f'cov_diag_prior_loc_{_id}', cov_diag_loc_init)
cov_diag_scale = pyro.param(f'cov_diag_prior_scale_{_id}', cov_diag_scale_init, constraint=constraints.positive)
cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale))
cov_diag = cov_diag + jitter
cov_factor = None
with pyro.plate(f'K_{_id}', K):
cov_factor_loc = pyro.param(f'cov_factor_prior_loc_{_id}', cov_factor_loc_init)
cov_factor_scale = pyro.param(f'cov_factor_prior_scale_{_id}', cov_factor_scale_init, constraint=constraints.positive)
cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(cov_factor_loc,cov_factor_scale))
cov_factor = cov_factor.transpose(-2,-1)
with pyro.plate(f'N_{_id}', size=N, subsample_size=self.batch_size) as ind:
X = pyro.sample('obs', dist.LowRankMultivariateNormal(loc, cov_factor=cov_factor, cov_diag=cov_diag), obs=X.index_select(0, ind))
return X
def model(self, zero_data, covariates):
duration,data_dim = zero_data.shape
feature_dim = covariates.size(-1)
drift_stability = pyro.sample("drift_stability", dist.Uniform(1, 2))
drift_scale = pyro.sample("drift_scale", dist.LogNormal(-20, 5))
with pyro.plate("item", data_dim, dim=-2):
bias = pyro.sample("bias", dist.Normal(5.0, 10.0))
with pyro.plate('features',feature_dim,dim=-1):
weight = pyro.sample("weight", dist.Normal(0.0, 5))
# We'll sample a time-global scale parameter outside the time plate,
# then time-local iid noise inside the time plate.
with self.time_plate:
# We combine two different reparameterizers: the inner SymmetricStableReparam
# is needed for the Stable site, and the outer LocScaleReparam is optional but
# appears to improve inference.
with poutine.reparam(config={"drift": LocScaleReparam()}):
with poutine.reparam(config={"drift": SymmetricStableReparam()}):
drift = pyro.sample("drift",
dist.Stable(drift_stability, 0, drift_scale))
motion = drift.cumsum(-1) # A Brownian motion.
# The prediction now includes three terms.
regression = torch.matmul(covariates, weight[...,None])
prediction = motion + bias + regression.sum(axis=-1)
prediction = prediction.unsqueeze(-1).transpose(-1, -3)
# Finally we can construct a noise distribution.
# We will share parameters across all time series.
obs_scale = pyro.sample("obs_scale", dist.LogNormal(-5, 5))
noise_dist = dist.Normal(loc=0.0, scale=obs_scale.unsqueeze(-1))
self.predict(noise_dist, prediction)
def model(self, zero_data, covariates):
assert zero_data.size(-1) == 1 # univariate
duration = zero_data.size(-2)
time, feature = covariates[..., 0], covariates[..., 1:]
bias = pyro.sample("bias", dist.Normal(0, 10))
# construct a linear trend; we know that the sales are increasing
# through years, so a positive-support prior should be used here
trend_coef = pyro.sample("trend", dist.LogNormal(-2, 1))
trend = trend_coef * time
# set prior of weights of the remaining covariates
weight = pyro.sample(
"weight",
dist.Normal(0, 1).expand([feature.size(-1)]).to_event(1))
regressor = (weight * feature).sum(-1)
# encode the additive weekly seasonality
with pyro.plate("day_of_week", 7, dim=-1):
seasonal = pyro.sample("seasonal", dist.Normal(0, 5))
seasonal = periodic_repeat(seasonal, duration, dim=-1)
# make prediction
prediction = bias + trend + seasonal + regressor
# because Pyro forecasting framework is multivariate,
# for univariate timeseries we need to make sure that
# the last dimension is 1
prediction = prediction.unsqueeze(-1)
# Now, we will use heavy tail noise because the data has some outliers
# (such as Christmas day)
dof = pyro.sample("dof", dist.Uniform(1, 10))
noise_scale = pyro.sample("noise_scale", dist.LogNormal(-2, 1))
noise_dist = dist.StudentT(dof.unsqueeze(-1), 0,
noise_scale.unsqueeze(-1))
self.predict(noise_dist, prediction)
def _param_map_estimates(
self, data: torch.Tensor,
chi_ambient: torch.Tensor) -> Dict[str, torch.Tensor]:
"""Calculate MAP estimates of mu, the mean of the true count matrix, and
lambda, the rate parameter of the Poisson background counts.
Args:
data: Dense tensor minibatch of cell by gene count data.
chi_ambient: Point estimate of inferred ambient gene expression.
Returns:
mu_map: Dense tensor of Negative Binomial means for true counts.
lambda_map: Dense tensor of Poisson rate params for noise counts.
alpha_map: Dense tensor of Dirichlet concentration params that
inform the overdispersion of the Negative Binomial.
"""
# Encode latents.
enc = self.vi_model.encoder.forward(x=data, chi_ambient=chi_ambient)
z_map = enc['z']['loc']
chi_map = self.vi_model.decoder.forward(z_map)
phi_loc = pyro.param('phi_loc')
phi_scale = pyro.param('phi_scale')
phi_conc = phi_loc.pow(2) / phi_scale.pow(2)
phi_rate = phi_loc / phi_scale.pow(2)
alpha_map = 1. / dist.Gamma(phi_conc, phi_rate).mean
y = (enc['p_y'] > 0).float()
d_empty = dist.LogNormal(loc=pyro.param('d_empty_loc'),
scale=pyro.param('d_empty_scale')).mean
d_cell = dist.LogNormal(loc=enc['d_loc'],
scale=pyro.param('d_cell_scale')).mean
epsilon = dist.Gamma(enc['epsilon'] * self.vi_model.epsilon_prior,
self.vi_model.epsilon_prior).mean
if self.vi_model.include_rho:
rho = pyro.param("rho_alpha") / (pyro.param("rho_alpha") +
pyro.param("rho_beta"))
else:
rho = None
# Calculate MAP estimates of mu and lambda.
mu_map = self.vi_model.calculate_mu(epsilon=epsilon,
d_cell=d_cell,
chi=chi_map,
y=y,
rho=rho)
lambda_map = self.vi_model.calculate_lambda(
epsilon=epsilon,
chi_ambient=chi_ambient,
d_empty=d_empty,
y=y,
d_cell=d_cell,
rho=rho,
chi_bar=self.vi_model.avg_gene_expression)
return {'mu': mu_map, 'lam': lambda_map, 'alpha': alpha_map}
def run_bo_pyro(params_bo, params_data, outer_loop_steps=10, q_size=5):
"""
run the bo
"""
X = params_data['X']
y = params_data['y']
#import ipdb; ipdb.set_trace()
gpmodel = gp.models.GPRegression(
X,
y,
gp.kernels.Matern52(input_dim=params_data['dim']),
noise=torch.tensor(0.1),
jitter=10**(-1))
gpmodel.kernel.set_prior("lengthscale", dist.LogNormal(0.0, 1.0))
gpmodel.kernel.set_prior("variance", dist.LogNormal(0.0, 1.0))
sampling_type = params_bo['sampling_type']
sample_size = params_bo['sample_size']
f_target = params_data['f_target']
print('run model with %s and %s samples' % (sampling_type, sample_size))
start_time = time.time()
optimizer = torch.optim.Adam(gpmodel.parameters(), lr=0.001)
gp.util.train(gpmodel, optimizer)
#gpmodel.optimize()
y_list_min = []
x_list_min = []
for i in range(outer_loop_steps):
print('approach %s, sample size %s, outer loop step %s of %s' %
(sampling_type, sample_size, i, outer_loop_steps))
xmin = next_x(gpmodel,
sampling_type=sampling_type,
sample_size=sample_size,
q_size=q_size)
print("next points evaluated:")
#import ipdb; ipdb.set_trace()
print(xmin, f_target(xmin))
# update the posterior with the new point found
#TODO: do we evaluate the target function at all points??
update_posterior(xmin, f_target, gpmodel)
# add to the list of best sampled points
val_y, ind = gpmodel.y.min(0)
val_x = gpmodel.X[ind, :]
y_list_min.append(val_y.detach().numpy())
x_list_min.append(val_x.detach().numpy())
print("best points so far:")
print(x_list_min[-1], y_list_min[-1])
print("run time %s seconds" % (time.time() - start_time))
res_dict = {
'X_exp': gpmodel.X.detach().numpy(),
'y_exp': gpmodel.y.numpy(),
'X_min': np.array(x_list_min),
'y_min': np.array(y_list_min)
}
return gpmodel, res_dict
def module_ppca_gm_means_sigma(self, input_batch, epsilon):
max_hidden_dim = max(self.hidden_dims)
batch_size = input_batch.shape[0]
ppca_means = input_batch.new_zeros(
[self.K, batch_size, self.output_dim])
if self.likelihood == 'normal':
#with pyro.plate('ppca_sigma_plate', self.K):
if self.group_isotropic:
ppca_gm_sigma = pyro.sample(
f'ppca_gm_sigma',
dist.LogNormal(0, input_batch.new_ones(1,
1)).independent(1))
else:
ppca_gm_sigma_list = []
ppca_gm_sigma = input_batch.new_ones(1, 1)
for i in range(self.d):
ppca_gm_sigma_list.append(
pyro.sample(
f'ppca_gm_sigma_{i}',
dist.LogNormal(0, input_batch.new_ones(
1, self.n[i])).independent(1)))
ppca_gm_sigma = torch_utils.krp_cw_torch(
ppca_gm_sigma_list[i], ppca_gm_sigma, column=False)
#with pyro.plate('alpha_plate', self.K):
alpha_gm = pyro.sample(
f'alpha_gm',
dist.LogNormal(0, input_batch.new_ones([1, self.group_hidden_dim
])).independent(1))
if self.group_iterm is None:
ppca_gm_means = self.group_term.multi_project(
input_batch,
remove_bias=False,
tensorize=False #fast_svd=False
)
else:
ppca_gm_means = self.group_term(
self.group_iterm(
torch_utils.flatten_torch(input_batch),
T=True #, fast_svd=False
))
ppca_gm_means += self.group_term(epsilon[:, :self.group_hidden_dim] *
alpha_gm[:, :self.group_hidden_dim])
if self.likelihood == 'bernoulli':
return ppca_gm_means.sigmoid()
if self.likelihood == 'normal':
return ppca_gm_means, ppca_gm_sigma
raise ValueError
def normal_model():
zero = torch.zeros(2)
a = pyro.sample("a", dist.Normal(0, 1))
b = pyro.sample("b", dist.LogNormal(0, 1))
c = pyro.sample("c", dist.Normal(a, b))
d = pyro.sample("d", dist.LogNormal(a, b))
e = pyro.sample("e", dist.Normal(zero, b).to_event(1))
f = pyro.sample("f", dist.LogNormal(zero, b).to_event(1))
g = pyro.sample("g", dist.Normal(0, 1), obs=a)
h = pyro.sample("h", dist.LogNormal(0, 1), obs=b)
with pyro.plate("plate", 5):
i = pyro.sample("i", dist.Normal(a, b))
j = pyro.sample("j", dist.LogNormal(a, b))
return a, b, c, d, e, f, g, h, i, j
def get_model(self, train_data, **kwargs):
X_train, y_train = train_data
# initialize the kernel and model
pyro.clear_param_store()
kernel = self.kernel(input_dim=X_train.shape[1])
gpr = gp.models.GPRegression(X_train,
y_train,
kernel,
noise=torch.tensor(1.).double())
# optional: fit the model using MAP
gpr.kernel.lengthscale = pyro.nn.PyroSample(dist.LogNormal(0.0, 1.0))
gpr.kernel.variance = pyro.nn.PyroSample(dist.LogNormal(0.0, 1.0))
return gpr
def guide(self, x, y=None):
pyro.module("scanvi", self)
with pyro.plate("batch",
len(x)), poutine.scale(scale=self.scale_factor):
z2_loc, z2_scale, l_loc, l_scale = self.z2l_encoder(x)
pyro.sample("l", dist.LogNormal(l_loc, l_scale).to_event(1))
z2 = pyro.sample("z2", dist.Normal(z2_loc, z2_scale).to_event(1))
y_logits = self.classifier(z2)
y_dist = dist.OneHotCategorical(logits=y_logits)
if y is None:
# x is unlabeled so sample y using q(y|z2)
y = pyro.sample("y", y_dist)
else:
# x is labeled so add a classification loss term
# (this way q(y|z2) learns from both labeled and unlabeled data)
classification_loss = y_dist.log_prob(y)
# Note that the negative sign appears because we're adding this term in the guide
# and the guide log_prob appears in the ELBO as -log q
pyro.factor(
"classification_loss",
-self.alpha * classification_loss,
has_rsample=False,
)
z1_loc, z1_scale = self.z1_encoder(z2, y)
pyro.sample("z1", dist.Normal(z1_loc, z1_scale).to_event(1))
def model(self, peak_idx, read_depth, onehot_obs=None):
pyro.module("decoder", self.decoder)
#pyro.module("decoder", self.decoder)
with pyro.plate("cells", peak_idx.shape[0]):
# Dirichlet prior 𝑝(𝜃|𝛼) is replaced by a log-normal distribution
theta_loc = self.prior_mu * peak_idx.new_ones(
(peak_idx.shape[0], self.num_topics))
theta_scale = self.prior_std * peak_idx.new_ones(
(peak_idx.shape[0], self.num_topics))
theta = pyro.sample(
"theta",
dist.LogNormal(theta_loc, theta_scale).to_event(1))
theta = theta / theta.sum(-1, keepdim=True)
# conditional distribution of 𝑤𝑛 is defined as
# 𝑤𝑛|𝛽,𝜃 ~ Categorical(𝜎(𝛽𝜃))
peak_probs = self.decoder(theta)
pyro.sample(
'obs',
dist.Multinomial(
total_count=read_depth if onehot_obs is None else 1,
probs=peak_probs).to_event(1),
obs=onehot_obs)
def stable_model():
zero = torch.zeros(2)
a = pyro.sample("a", dist.Normal(0, 1))
b = pyro.sample("b", dist.LogNormal(0, 1))
c = pyro.sample("c", dist.Stable(1.5, 0.0, b, a))
d = pyro.sample("d", dist.Stable(1.5, 0.0, b, 0.0), obs=a)
e = pyro.sample("e", dist.Stable(1.5, 0.1, b, a))
f = pyro.sample("f", dist.Stable(1.5, 0.1, b, 0.0), obs=a)
g = pyro.sample("g", dist.Stable(1.5, zero, b, a).to_event(1))
h = pyro.sample("h", dist.Stable(1.5, zero, b, 0).to_event(1), obs=a)
i = pyro.sample(
"i",
dist.TransformedDistribution(dist.Stable(1.5, 0, b, a),
dist.transforms.ExpTransform()),
)
j = pyro.sample(
"j",
dist.TransformedDistribution(dist.Stable(1.5, 0, b, a),
dist.transforms.ExpTransform()),
obs=a.exp(),
)
k = pyro.sample(
"k",
dist.TransformedDistribution(dist.Stable(
1.5, zero, b, a), dist.transforms.ExpTransform()).to_event(1),
)
l = pyro.sample(
"l",
dist.TransformedDistribution(dist.Stable(
1.5, zero, b, a), dist.transforms.ExpTransform()).to_event(1),
obs=a.exp() + zero,
)
return a, b, c, d, e, f, g, h, i, j, k, l
def guide(self, bows, _):
pyro.module("prop_inference_net", self.prop_inference_net)
with pyro.plate("articles", bows.shape[0]):
prop_mu, prop_sigma = self.prop_inference_net(bows)
props = pyro.sample(
"theta",
dist.LogNormal(prop_mu, prop_sigma).to_event(1))
def model(self, bows, embeddings, article_ids):
pyro.module("topic_recognition_net", self.topic_recognition_net)
with pyro.plate("articles", bows.shape[0]):
# instead of a Dirichlet prior, we use a log-normal distribution
prop_mu = bows.new_zeros((bows.shape[0], self.nav_topics))
prop_sigma = bows.new_ones((bows.shape[0], self.nav_topics))
props = pyro.sample(
"theta",
dist.LogNormal(prop_mu, prop_sigma).to_event(1))
topics_mu, topics_sigma = self.topic_recognition_net(props)
for batch_article_id, article_id in enumerate(article_ids):
nav_embeddings = torch.tensor(
embeddings[self.article_navs[article_id]],
dtype=torch.float32).to(device)
for article_nav_id in pyro.plate(
"navs_{}".format(article_id),
len(self.article_navs[article_id])):
pyro.sample("nav_{}_{}".format(article_id, article_nav_id),
dist.MultivariateNormal(
topics_mu[batch_article_id],
scale_tril=torch.diag(
topics_sigma[batch_article_id])),
obs=nav_embeddings[article_nav_id])
def gp_model(X, y):
N = X.shape[0]
# Priors for kernel parameters.
alpha = pyro.sample("alpha", dist.LogNormal(0., 0.1))
rho = pyro.sample("rho", dist.LogNormal(0., 1.))
sigma = pyro.sample("sigma", dist.LogNormal(0., 1.))
# Covariance matrix.
K = sq_exp_kernel_matrix(X, alpha, rho) + torch.eye(N) * sigma * sigma
L = torch.cholesky(K)
# Marginal likelihood.
pyro.sample('obs',
dist.MultivariateNormal(torch.zeros(N), scale_tril=L),
obs=y)
def gpc(X, y):
N = y.shape[0]
# Priors.
alpha = pyro.sample('alpha', dist.LogNormal(0, 1))
rho = pyro.sample('rho', dist.LogNormal(0, 1))
beta = pyro.sample('beta', dist.Normal(0, 1))
with pyro.plate('latent_response', N):
eta = pyro.sample('eta', dist.Normal(0, 1))
# Latent function.
f = compute_f(alpha, rho, beta, eta, X)
with pyro.plate('response', N):
pyro.sample('obs', dist.Bernoulli(logits=f), obs=y)
def model(self, X):
_id = self._id
N, D = X.shape
global_shrinkage_prior_scale_init = self.param_init[f'global_shrinkage_prior_scale_init_{_id}']
cov_diag_prior_loc_init = self.param_init[f'cov_diag_prior_loc_init_{_id}']
cov_diag_prior_scale_init = self.param_init[f'cov_diag_prior_scale_init_{_id}']
global_shrinkage_prior_scale = pyro.param(f'global_shrinkage_scale_prior_{_id}', global_shrinkage_prior_scale_init, constraint=constraints.positive)
tau = pyro.sample(f'global_shrinkage_{_id}', dist.HalfNormal(global_shrinkage_prior_scale))
b = pyro.sample('b', dist.InverseGamma(0.5,1./torch.ones(D)**2).to_event(1))
lambdasquared = pyro.sample(f'local_shrinkage_{_id}', dist.InverseGamma(0.5,1./b).to_event(1))
cov_diag_loc = pyro.param(f'cov_diag_prior_loc_{_id}', cov_diag_prior_loc_init)
cov_diag_scale = pyro.param(f'cov_diag_prior_scale_{_id}', cov_diag_prior_scale_init, constraint=constraints.positive)
cov_diag = pyro.sample(f'cov_diag_{_id}', dist.LogNormal(cov_diag_loc, cov_diag_scale).to_event(1))
#cov_diag = cov_diag*torch.ones(D)
cov_diag = cov_diag + jitter
lambdasquared = lambdasquared.squeeze()
if lambdasquared.dim() == 1:
# outer product
cov_factor_scale = torch.ger(torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,)))
else:
# batch outer product
cov_factor_scale = torch.einsum('bp, br->bpr', torch.sqrt(lambdasquared),tau.repeat((tau.dim()-1)*(1,)+(D,)))
cov_factor = pyro.sample(f'cov_factor_{_id}', dist.Normal(0., cov_factor_scale).to_event(2))
cov_factor = cov_factor.transpose(-2,-1)
with pyro.plate(f'N_{_id}', size=N, subsample_size=self.batch_size, dim=-1) as ind:
X = pyro.sample('obs', dist.LowRankMultivariateNormal(torch.zeros(D), cov_factor=cov_factor, cov_diag=cov_diag), obs=X.index_select(0, ind))
return X
def factorGuide(X, batch_size, variational_parameter_initialization):
N, D = X.shape
K, locloc, locscale, scaleloc, scalescale, cov_factor_loc_init, cov_factor_scale_init = variational_parameter_initialization[1]
with pyro.plate('D', D, dim=-1):
cov_diag_loc = pyro.param('scale_loc', scaleloc)
cov_diag_scale = pyro.param('scale_scale', scalescale, constraint=constraints.positive)
cov_diag = pyro.sample('scale', dist.LogNormal(cov_diag_loc, cov_diag_scale))
cov_diag = cov_diag*torch.ones(D)
loc_loc = pyro.param('loc_loc', locloc)
loc_scale = pyro.param('loc_scale', locscale, constraint=constraints.positive)
# sample variables
loc = pyro.sample('loc', dist.Normal(loc_loc,loc_scale))
cov_factor = None
if K > 1:
with pyro.plate('K', K-1, dim=-2):
cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init[:K-1,:])
cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init[:K-1,:], constraint=constraints.positive)
cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc, cov_factor_scale))
cov_factor_new_loc = pyro.param('cov_factor_new_loc_{}'.format(K), cov_factor_loc_init[-1,:])
cov_factor_new_scale = pyro.param('cov_factor_new_scale_{}'.format(K), cov_factor_scale_init[-1,:], constraint=constraints.positive)
cov_factor_new = pyro.sample('cov_factor_new', dist.Normal(cov_factor_new_loc,cov_factor_new_scale))
cov_factor = torch.cat([cov_factor, torch.unsqueeze(cov_factor_new, dim=0)])
else:
with pyro.plate('K', K):
cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init)
cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init, constraint=constraints.positive)
cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc,cov_factor_scale))
cov_factor = cov_factor.transpose(0,1)
return loc, cov_factor, cov_diag
def model(self, zero_data, covariates):
with pyro.plate("batch", len(zero_data), dim=-2):
zero_data = pyro.subsample(zero_data, event_dim=1)
covariates = pyro.subsample(covariates, event_dim=1)
loc = zero_data[..., :1, :]
scale = pyro.sample("scale", dist.LogNormal(loc, 1).to_event(1))
with self.time_plate:
jumps = pyro.sample("jumps", dist.Normal(0, scale).to_event(1))
prediction = jumps.cumsum(-2)
duration, obs_dim = zero_data.shape[-2:]
noise_dist = dist.LinearHMM(
dist.Stable(1.9, 0).expand([obs_dim]).to_event(1),
torch.eye(obs_dim),
dist.Stable(1.9, 0).expand([obs_dim]).to_event(1),
torch.eye(obs_dim),
dist.Stable(1.9, 0).expand([obs_dim]).to_event(1),
duration=duration,
)
rep = StableReparam()
with poutine.reparam(
config={"residual": LinearHMMReparam(rep, rep, rep)}):
self.predict(noise_dist, prediction)
def projectedMixtureGuide(X, batch_size, variational_parameter_initialization):
"""
Covariances of all clusters are locked, we're just learning mixture weights and means
"""
N, D = X.shape
loc_loc, loc_scale, scale_loc, scale_scale, component_logits_concentration, cov_factor_loc_init,cov_factor_scale_init = variational_parameter_initialization[1]
C = loc_loc.shape[0]
K = cov_factor_loc_init.shape[0]
component_logits_concentration = pyro.param('component_logits_concentration', component_logits_concentration, constraint=constraints.positive)
component_logits = pyro.sample('component_logits', dist.Dirichlet(component_logits_concentration))
with pyro.plate('D', D):
cov_diag_loc = pyro.param('scale_loc', scale_loc)
cov_diag_scale = pyro.param('scale_scale', scale_scale, constraint=constraints.positive)
cov_diag = pyro.sample('scale', dist.LogNormal(scale_loc, scale_scale))
cov_diag = cov_diag*torch.ones(D)
with pyro.plate('K', K):
cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init)
cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init, constraint=constraints.positive)
cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc,cov_factor_scale))
cov_factor = cov_factor.transpose(0,1)
with pyro.plate('C', C):
loc_loc = pyro.param('loc_loc', loc_loc)
loc_scale = pyro.param('loc_scale', loc_scale, constraint=constraints.positive)
locs = pyro.sample('locs', dist.Normal(loc_loc,loc_scale))
return component_logits, cov_diag, cov_factor, locs
def zeroMeanFactorGuide(X, batch_size, variational_parameter_initialization):
N, D = X.shape
K, scaleloc, scalescale, cov_factor_loc_init, cov_factor_scale_init = variational_parameter_initialization[1]
with pyro.plate('D', D, dim=-1):
cov_diag_loc = pyro.param('scale_loc', scaleloc, constraint=constraints.positive)
cov_diag_scale = pyro.param('scale_scale', scalescale, constraint=constraints.positive)
cov_diag = pyro.sample('scale', dist.LogNormal(cov_diag_loc, cov_diag_scale))
cov_diag = cov_diag*torch.ones(D)
# sample variables
cov_factor = None
if K > 1:
with pyro.plate('K', K-1, dim=-2):
cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init[:K-1,:])
cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init[:K-1,:], constraint=constraints.positive)
cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc, cov_factor_scale))
cov_factor_new_loc = pyro.param('cov_factor_new_loc_{}'.format(K), cov_factor_loc_init[-1,:])
cov_factor_new_scale = pyro.param('cov_factor_new_scale_{}'.format(K), cov_factor_scale_init[-1,:], constraint=constraints.positive)
cov_factor_new = pyro.sample('cov_factor_new', dist.Normal(cov_factor_new_loc,cov_factor_new_scale))
# when using pyro.infer.Predictive, cov_factor_new is somehow sampled as 2-d tensors instead of 1-d
if cov_factor_new.dim() == cov_factor.dim():
cov_factor = torch.cat([cov_factor, cov_factor_new], dim=-2)
else:
cov_factor = torch.cat([cov_factor, torch.unsqueeze(cov_factor_new, dim=-2)], dim=-2)
else:
with pyro.plate('K', K):
cov_factor_loc = pyro.param('cov_factor_loc_{}'.format(K), cov_factor_loc_init)
cov_factor_scale = pyro.param('cov_factor_scale_{}'.format(K), cov_factor_scale_init, constraint=constraints.positive)
cov_factor = pyro.sample('cov_factor', dist.Normal(cov_factor_loc,cov_factor_scale))
cov_factor = cov_factor.transpose(-2,-1)
return cov_factor, cov_diag
def sphericalMixture(X, batch_size, prior_parameters):
"""
Covariances of all clusters are diagonal, only params are mixture weights, means and shared variable variances
...so it's actually an ellipsoidal mixture, if you wanna be anal about it. C**t.
"""
N, D = X.shape
locloc, locscale, scaleloc, scalescale, component_logits_concentration = prior_parameters[0]
# get number of clusters from first dimension of means
C = locloc.shape[0]
component_logits = pyro.sample('component_logits', dist.Dirichlet(component_logits_concentration))
with pyro.plate('D', D):
cov_diag = pyro.sample('scale', dist.LogNormal(scaleloc, scalescale))
with pyro.plate('C', C):
locs = pyro.sample('locs', dist.Normal(locloc,locscale))
with pyro.plate('N', size=N, subsample_size=batch_size) as ind:
assignment = pyro.sample('assignment', dist.Categorical(component_logits), infer={"enumerate": "parallel"})
#X = pyro.sample('obs', dist.MultivariateNormal(locs.index_select(-2, assignment), torch.diag(cov_diag)), obs=X.index_select(0, ind))
# use index_select instead of locs[assignment] so batching works correctly
# have to select the right index both in a parallel enumeration context and a batch context
# leading to this ugly hack
indexed_locs = locs.index_select(-2, assignment.squeeze()).view(*locs.shape[:-2],*assignment.shape,locs.shape[-1])
#print(locs.shape)
#print(assignment.shape)
#print(indexed_locs.shape)
X = pyro.sample('obs', dist.MultivariateNormal(indexed_locs, torch.diag_embed(cov_diag)), obs=X.index_select(0, ind))
return X
def test_transformed_hmm_shape(batch_shape, duration, hidden_dim, obs_dim):
init_dist = random_mvn(batch_shape, hidden_dim)
trans_mat = torch.randn(batch_shape + (duration, hidden_dim, hidden_dim))
trans_dist = random_mvn(batch_shape + (duration, ), hidden_dim)
obs_mat = torch.randn(batch_shape + (duration, hidden_dim, obs_dim))
obs_dist = dist.LogNormal(
torch.randn(batch_shape + (duration, obs_dim)),
torch.rand(batch_shape + (duration, obs_dim)).exp()).to_event(1)
hmm = dist.LinearHMM(init_dist,
trans_mat,
trans_dist,
obs_mat,
obs_dist,
duration=duration)
def model(data=None):
with pyro.plate_stack("plates", batch_shape):
return pyro.sample("x", hmm, obs=data)
data = model()
with poutine.trace() as tr:
with poutine.reparam(config={"x": LinearHMMReparam()}):
model(data)
fn = tr.trace.nodes["x"]["fn"]
assert isinstance(fn, dist.TransformedDistribution)
assert isinstance(fn.base_dist, dist.GaussianHMM)
tr.trace.compute_log_prob() # smoke test only
def model(self, x, y=None):
# Register various nn.Modules with Pyro
pyro.module("scanvi", self)
# This gene-level parameter modulates the variance of the observation distribution
theta = pyro.param("inverse_dispersion", 10.0 * x.new_ones(self.num_genes),
constraint=constraints.positive)
# We scale all sample statements by scale_factor so that the ELBO is normalized
# wrt the number of datapoints and genes
with pyro.plate("batch", len(x)), poutine.scale(scale=self.scale_factor):
z1 = pyro.sample("z1", dist.Normal(0, x.new_ones(self.latent_dim)).to_event(1))
# Note that if y is None (i.e. y is unobserved) then y will be sampled;
# otherwise y will be treated as observed.
y = pyro.sample("y", dist.OneHotCategorical(logits=x.new_zeros(self.num_labels)),
obs=y)
z2_loc, z2_scale = self.z2_decoder(z1, y)
z2 = pyro.sample("z2", dist.Normal(z2_loc, z2_scale).to_event(1))
l_scale = self.l_scale * x.new_ones(1)
l = pyro.sample("l", dist.LogNormal(self.l_loc, l_scale).to_event(1))
# Note that by construction mu is normalized (i.e. mu.sum(-1) == 1) and the
# total scale of counts for each cell is determined by `l`
gate_logits, mu = self.x_decoder(z2)
# TODO revisit this parameterization if torch.distributions.NegativeBinomial changes
# from failure to success parametrization;
# see https://github.com/pytorch/pytorch/issues/42449
nb_logits = (l * mu + self.epsilon).log() - (theta + self.epsilon).log()
x_dist = dist.ZeroInflatedNegativeBinomial(gate_logits=gate_logits, total_count=theta,
logits=nb_logits)
# Observe the datapoint x using the observation distribution x_dist
pyro.sample("x", x_dist.to_event(1), obs=x)
def setUp(self):
self.mu = Variable(torch.Tensor([1.4]))
self.sigma = Variable(torch.Tensor([0.4]))
self.test_data = Variable(torch.Tensor([5.5]))
self.batch_mu = Variable(torch.Tensor([[1.4], [2.6]]))
self.batch_sigma = Variable(torch.Tensor([[0.4], [0.5]]))
self.batch_test_data = Variable(torch.Tensor([[5.5], [6.4]]))
self.dist = dist.LogNormal(self.mu, self.sigma)
self.batch_dist = dist.LogNormal(self.batch_mu, self.batch_sigma)
self.analytic_mean = torch.exp(
self.mu + 0.5 * torch.pow(self.sigma, 2.0)).data.cpu().numpy()[0]
var_factor = torch.exp(torch.pow(self.sigma, 2.0)) - Variable(
torch.ones(1))
self.analytic_var = var_factor.data.cpu().numpy()[0] * np.square(
self.analytic_mean)
self.n_samples = 70000