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 forward(self):
def RP(weights, distances, d):
return 1e4 * (weights * torch.pow(0.5, distances/(1e3 * d))).sum(-1)
with pyro.plate(self.name +"_regions", 3):
a = pyro.sample(self.name +"_a", dist.HalfNormal(12.))
with pyro.plate(self.name +"_upstream-downstream", 2):
d = torch.exp(pyro.sample(self.name +'_logdistance', dist.Normal(np.e, 2.)))
b = pyro.sample(self.name +"_b", dist.Normal(-10.,3.))
theta = pyro.sample(self.name +"_theta", dist.Gamma(2., 0.5))
psi = pyro.sample(self.name +"_dropout", dist.Beta(1., 10.))
with pyro.plate(self.name +"_data", self.N, subsample_size=64) as ind:
expr_rate = a[0] * RP(self.upstream_weights.index_select(0, ind), self.upstream_distances, d[0])\
+ a[1] * RP(self.downstream_weights.index_select(0, ind), self.downstream_distances, d[1]) \
+ a[2] * 1e4 * self.promoter_weights.index_select(0, ind).sum(-1) \
+ b
mu = torch.multiply(self.read_depth.index_select(0, ind), torch.exp(expr_rate))
p = torch.minimum(mu / (mu + theta), torch.tensor([0.99999]))
pyro.sample(self.name +'_obs',
dist.ZeroInflatedNegativeBinomial(total_count=theta, probs=p, gate = psi),
obs= self.gene_expr.index_select(0, ind))
def model(self, raw_expr, encoded_expr, read_depth):
pyro.module("decoder", self.decoder)
with pyro.plate("genes", self.num_genes):
dispersion = pyro.sample(
"dispersion",
dist.Gamma(
torch.tensor(2.).to(self.device),
torch.tensor(0.5).to(self.device)))
psi = pyro.sample(
"dropout",
dist.Beta(
torch.tensor(1.).to(self.device),
torch.tensor(10.).to(self.device)))
#pyro.module("decoder", self.decoder)
with pyro.plate("cells", encoded_expr.shape[0]):
# Dirichlet prior π(π|πΌ) is replaced by a log-normal distribution
theta_loc = self.prior_mu * encoded_expr.new_ones(
(encoded_expr.shape[0], self.num_topics))
theta_scale = self.prior_std * encoded_expr.new_ones(
(encoded_expr.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(π(π½π))
expr_rate = pyro.deterministic("expr_rate", self.decoder(theta))
mu = torch.multiply(read_depth, expr_rate)
p = torch.minimum(mu / (mu + dispersion), self.max_prob)
pyro.sample(
'obs',
dist.ZeroInflatedNegativeBinomial(total_count=dispersion,
probs=p,
gate=psi).to_event(1),
obs=raw_expr)
def model(self, x, log_library):
# register PyTorch module `decoder` with Pyro
pyro.module("scvi", self)
with pyro.plate("data", x.shape[0]):
# setup hyperparameters for prior p(z)
z_loc = x.new_zeros(torch.Size((x.shape[0], self.n_latent)))
z_scale = x.new_ones(torch.Size((x.shape[0], self.n_latent)))
# sample from prior (value will be sampled by guide when computing the ELBO)
z = pyro.sample("latent", dist.Normal(z_loc, z_scale).to_event(1))
# decode the latent code z
px_scale, _, px_rate, px_dropout = self.decoder("gene", z, log_library)
# build count distribution
nb_logits = (px_rate + self.epsilon).log() - (
self.px_r.exp() + self.epsilon
).log()
x_dist = dist.ZeroInflatedNegativeBinomial(
gate_logits=px_dropout, total_count=self.px_r.exp(), logits=nb_logits
)
# score against actual counts
pyro.sample("obs", x_dist.to_event(1), obs=x)
def model(self, raw_expr, encoded_expr, read_depth):
pyro.module("decoder", self.decoder)
dispersion = pyro.param("dispersion",
torch.tensor(5.).to(self.device) *
torch.ones(self.num_genes).to(self.device),
constraint=constraints.positive)
with pyro.plate("cells", encoded_expr.shape[0]):
# Dirichlet prior π(π|πΌ) is replaced by a log-normal distribution
theta_loc = self.prior_mu * encoded_expr.new_ones(
(encoded_expr.shape[0], self.num_topics))
theta_scale = self.prior_std * encoded_expr.new_ones(
(encoded_expr.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)
read_scale = pyro.sample(
'read_depth',
dist.LogNormal(torch.log(read_depth), 1.).to_event(1))
#read_scale = torch.minimum(read_scale, self.max_scale)
# conditional distribution of π€π is defined as
# π€π|π½,π ~ Categorical(π(π½π))
expr_rate, dropout = self.decoder(theta)
mu = torch.multiply(read_scale, expr_rate)
p = torch.minimum(mu / (mu + dispersion), self.max_prob)
pyro.sample('obs',
dist.ZeroInflatedNegativeBinomial(
total_count=dispersion,
probs=p,
gate_logits=dropout).to_event(1),
obs=raw_expr)
def forward(self):
with pyro.plate("gene_weights", self.G):
b = pyro.sample("b", dist.Normal(-10.,3.))
theta = pyro.sample("theta", dist.Gamma(2., 0.5))
psi = pyro.sample("dropout", dist.Beta(1., 10.))
with pyro.plate("topic-gene_weights", self.K):
beta = pyro.sample("beta", dist.Gamma(1., 5.))
with pyro.plate("gene", self.G) as gene:
with pyro.plate("data", self.N, subsample_size=64) as ind:
expr_rate = pyro.deterministic("rate", torch.matmul(self.cell_topics.index_select(0, ind), beta) + b)
mu = torch.reshape(self.read_depth, (-1,1)).index_select(0, ind) * torch.exp(expr_rate)
p = torch.minimum(mu / (mu + theta), torch.tensor([0.99999]))
pyro.sample("obs",
dist.ZeroInflatedNegativeBinomial(total_count=theta,
probs=p, gate = psi),
obs= self.gene_expr.index_select(0, ind))