def generative(
self, z, library, batch_index, cont_covs=None, cat_covs=None, y=None
):
"""Runs the generative model."""
# TODO: refactor forward function to not rely on y
decoder_input = z if cont_covs is None else torch.cat([z, cont_covs], dim=-1)
if cat_covs is not None:
categorical_input = torch.split(cat_covs, 1, dim=1)
else:
categorical_input = tuple()
px_scale, px_r, px_rate, px_dropout = self.decoder(
self.dispersion, decoder_input, library, batch_index, *categorical_input, y
)
if self.dispersion == "gene-label":
px_r = F.linear(
one_hot(y, self.n_labels), self.px_r
) # px_r gets transposed - last dimension is nb genes
elif self.dispersion == "gene-batch":
px_r = F.linear(one_hot(batch_index, self.n_batch), self.px_r)
elif self.dispersion == "gene":
px_r = self.px_r
px_r = torch.exp(px_r)
return dict(
px_scale=px_scale, px_r=px_r, px_rate=px_rate, px_dropout=px_dropout
)
def reshape_bernoulli(
self,
bernoulli_params: torch.Tensor,
batch_index: Optional[torch.Tensor] = None,
y: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if self.zero_inflation == "gene-label":
one_hot_label = one_hot(y, self.n_labels)
# If we sampled several random Bernoulli parameters
if len(bernoulli_params.shape) == 2:
bernoulli_params = F.linear(one_hot_label, bernoulli_params)
else:
bernoulli_params_res = []
for sample in range(bernoulli_params.shape[0]):
bernoulli_params_res.append(
F.linear(one_hot_label, bernoulli_params[sample]))
bernoulli_params = torch.stack(bernoulli_params_res)
elif self.zero_inflation == "gene-batch":
one_hot_batch = one_hot(batch_index, self.n_batch)
if len(bernoulli_params.shape) == 2:
bernoulli_params = F.linear(one_hot_batch, bernoulli_params)
# If we sampled several random Bernoulli parameters
else:
bernoulli_params_res = []
for sample in range(bernoulli_params.shape[0]):
bernoulli_params_res.append(
F.linear(one_hot_batch, bernoulli_params[sample]))
bernoulli_params = torch.stack(bernoulli_params_res)
return bernoulli_params
def generative(
self,
z: torch.Tensor,
library: torch.Tensor,
batch_index: Optional[torch.Tensor] = None,
y: Optional[torch.Tensor] = None,
mode: Optional[int] = None,
) -> dict:
px_scale, px_r, px_rate, px_dropout = self.decoder(
z, mode, library, self.dispersion, batch_index, y
)
if self.dispersion == "gene-label":
px_r = F.linear(one_hot(y, self.n_labels), self.px_r)
elif self.dispersion == "gene-batch":
px_r = F.linear(one_hot(batch_index, self.n_batch), self.px_r)
elif self.dispersion == "gene":
px_r = self.px_r.view(1, self.px_r.size(0))
px_r = torch.exp(px_r)
px_scale = px_scale / torch.sum(
px_scale[:, self.indices_mappings[mode]], dim=1
).view(-1, 1)
px_rate = px_scale * torch.exp(library)
return dict(
px_scale=px_scale, px_r=px_r, px_rate=px_rate, px_dropout=px_dropout
)
def marginal_ll(self, tensors, n_mc_samples):
x = tensors[_CONSTANTS.X_KEY]
batch_index = tensors[_CONSTANTS.BATCH_KEY]
to_sum = torch.zeros(x.size()[0], n_mc_samples)
for i in range(n_mc_samples):
# Distribution parameters and sampled variables
inference_outputs, generative_outputs, losses = self.forward(tensors)
# outputs = self.module.inference(x, y, batch_index, labels)
qz_m = inference_outputs["qz_m"]
qz_v = inference_outputs["qz_v"]
ql_m = inference_outputs["ql_m"]
ql_v = inference_outputs["ql_v"]
py_ = generative_outputs["py_"]
log_library = inference_outputs["untran_l"]
# really need not softmax transformed random variable
z = inference_outputs["untran_z"]
log_pro_back_mean = generative_outputs["log_pro_back_mean"]
# Reconstruction Loss
reconst_loss = losses._reconstruction_loss
reconst_loss_gene = reconst_loss["reconst_loss_gene"]
reconst_loss_protein = reconst_loss["reconst_loss_protein"]
# Log-probabilities
log_prob_sum = torch.zeros(qz_m.shape[0]).to(self.device)
if not self.use_observed_lib_size:
n_batch = self.library_log_means.shape[1]
local_library_log_means = F.linear(
one_hot(batch_index, n_batch), self.library_log_means
)
local_library_log_vars = F.linear(
one_hot(batch_index, n_batch), self.library_log_vars
)
p_l_gene = (
Normal(local_library_log_means, local_library_log_vars.sqrt())
.log_prob(log_library)
.sum(dim=-1)
)
q_l_x = Normal(ql_m, ql_v.sqrt()).log_prob(log_library).sum(dim=-1)
log_prob_sum += p_l_gene - q_l_x
p_z = Normal(0, 1).log_prob(z).sum(dim=-1)
p_mu_back = self.back_mean_prior.log_prob(log_pro_back_mean).sum(dim=-1)
p_xy_zl = -(reconst_loss_gene + reconst_loss_protein)
q_z_x = Normal(qz_m, qz_v.sqrt()).log_prob(z).sum(dim=-1)
q_mu_back = (
Normal(py_["back_alpha"], py_["back_beta"])
.log_prob(log_pro_back_mean)
.sum(dim=-1)
)
log_prob_sum += p_z + p_mu_back + p_xy_zl - q_z_x - q_mu_back
to_sum[:, i] = log_prob_sum
batch_log_lkl = torch.logsumexp(to_sum, dim=-1) - np.log(n_mc_samples)
log_lkl = torch.sum(batch_log_lkl).item()
return log_lkl
def generative(
self,
z: torch.Tensor,
library_gene: torch.Tensor,
batch_index: torch.Tensor,
label: torch.Tensor,
cont_covs=None,
cat_covs=None,
size_factor=None,
transform_batch: Optional[int] = None,
) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]:
decoder_input = z if cont_covs is None else torch.cat([z, cont_covs],
dim=-1)
if cat_covs is not None:
categorical_input = torch.split(cat_covs, 1, dim=1)
else:
categorical_input = tuple()
if transform_batch is not None:
batch_index = torch.ones_like(batch_index) * transform_batch
if not self.use_size_factor_key:
size_factor = library_gene
px_, py_, log_pro_back_mean = self.decoder(decoder_input, size_factor,
batch_index,
*categorical_input)
if self.gene_dispersion == "gene-label":
# px_r gets transposed - last dimension is nb genes
px_r = F.linear(one_hot(label, self.n_labels), self.px_r)
elif self.gene_dispersion == "gene-batch":
px_r = F.linear(one_hot(batch_index, self.n_batch), self.px_r)
elif self.gene_dispersion == "gene":
px_r = self.px_r
px_r = torch.exp(px_r)
if self.protein_dispersion == "protein-label":
# py_r gets transposed - last dimension is n_proteins
py_r = F.linear(one_hot(label, self.n_labels), self.py_r)
elif self.protein_dispersion == "protein-batch":
py_r = F.linear(one_hot(batch_index, self.n_batch), self.py_r)
elif self.protein_dispersion == "protein":
py_r = self.py_r
py_r = torch.exp(py_r)
px_["r"] = px_r
py_["r"] = py_r
return dict(
px_=px_,
py_=py_,
log_pro_back_mean=log_pro_back_mean,
)
def _compute_local_library_params(self, batch_index):
"""
Computes local library parameters.
Compute two tensors of shape (batch_index.shape[0], 1) where each
element corresponds to the mean and variances, respectively, of the
log library sizes in the batch the cell corresponds to.
"""
n_batch = self.library_log_means.shape[1]
local_library_log_means = F.linear(
one_hot(batch_index, n_batch), self.library_log_means
)
local_library_log_vars = F.linear(
one_hot(batch_index, n_batch), self.library_log_vars
)
return local_library_log_means, local_library_log_vars
def loss_adversarial_classifier(self, z, batch_index, predict_true_class=True):
n_classes = self.n_output_classifier
cls_logits = torch.nn.LogSoftmax(dim=1)(self.adversarial_classifier(z))
if predict_true_class:
cls_target = one_hot(batch_index, n_classes)
else:
one_hot_batch = one_hot(batch_index, n_classes)
cls_target = torch.zeros_like(one_hot_batch)
# place zeroes where true label is
cls_target.masked_scatter_(
~one_hot_batch.bool(), torch.ones_like(one_hot_batch) / (n_classes - 1)
)
l_soft = cls_logits * cls_target
loss = -l_soft.sum(dim=1).mean()
return loss
def forward(self, x: torch.Tensor, *cat_list: int):
"""
Forward computation on x for sparse layer.
Parameters
----------
x
tensor of values with shape ``(n_in,)``
cat_list
list of category membership(s) for this sample
x: torch.Tensor
Returns
-------
py:class:`torch.Tensor`
tensor of shape ``(n_out,)``
"""
one_hot_cat_list = [
] # for generality in this list many indices useless.
if len(self.n_cat_list) > len(cat_list):
raise ValueError(
"nb. categorical args provided doesn't match init. params.")
for n_cat, cat in zip(self.n_cat_list, cat_list):
if n_cat and cat is None:
raise ValueError(
"cat not provided while n_cat != 0 in init. params.")
if n_cat > 1: # n_cat = 1 will be ignored - no additional information
if cat.size(1) != n_cat:
one_hot_cat = one_hot(cat, n_cat)
else:
one_hot_cat = cat # cat has already been one_hot encoded
one_hot_cat_list += [one_hot_cat]
for layer in self.sparse_layer:
if layer is not None:
if isinstance(layer, nn.BatchNorm1d):
if x.dim() == 3:
x = torch.cat([(layer(slice_x)).unsqueeze(0)
for slice_x in x],
dim=0)
else:
x = layer(x)
else:
if isinstance(layer, CustomizedLinear):
if x.dim() == 3:
one_hot_cat_list_layer = [
o.unsqueeze(0).expand(
(x.size(0), o.size(0), o.size(1)))
for o in one_hot_cat_list
]
else:
one_hot_cat_list_layer = one_hot_cat_list
x = torch.cat((x, *one_hot_cat_list_layer), dim=-1)
x = layer(x)
return x
def generative(self,
z,
library_gene,
batch_index,
label,
cont_covs=None,
cat_covs=None):
decoder_input = z if cont_covs is None else torch.cat([z, cont_covs],
dim=-1)
if cat_covs is not None:
categorical_input = torch.split(cat_covs, 1, dim=1)
else:
categorical_input = tuple()
px_, py_, log_pro_back_mean = self.decoder(decoder_input, library_gene,
batch_index,
*categorical_input)
if self.gene_dispersion == "gene-label":
# px_r gets transposed - last dimension is nb genes
px_r = F.linear(one_hot(label, self.n_labels), self.px_r)
elif self.gene_dispersion == "gene-batch":
px_r = F.linear(one_hot(batch_index, self.n_batch), self.px_r)
elif self.gene_dispersion == "gene":
px_r = self.px_r
px_r = torch.exp(px_r)
if self.protein_dispersion == "protein-label":
# py_r gets transposed - last dimension is n_proteins
py_r = F.linear(one_hot(label, self.n_labels), self.py_r)
elif self.protein_dispersion == "protein-batch":
py_r = F.linear(one_hot(batch_index, self.n_batch), self.py_r)
elif self.protein_dispersion == "protein":
py_r = self.py_r
py_r = torch.exp(py_r)
px_["r"] = px_r
py_["r"] = py_r
return dict(
px_=px_,
py_=py_,
log_pro_back_mean=log_pro_back_mean,
)
def broadcast_labels(y, *o, n_broadcast=-1):
"""
Utility for the semi-supervised setting.
If y is defined(labelled batch) then one-hot encode the labels (no broadcasting needed)
If y is undefined (unlabelled batch) then generate all possible labels (and broadcast other arguments if not None)
"""
if not len(o):
raise ValueError("Broadcast must have at least one reference argument")
if y is None:
ys = enumerate_discrete(o[0], n_broadcast)
new_o = iterate(
o,
lambda x: x.repeat(n_broadcast, 1)
if len(x.size()) == 2 else x.repeat(n_broadcast),
)
else:
ys = one_hot(y, n_broadcast)
new_o = o
return (ys, ) + new_o
def forward(self, x_data, idx, batch_index):
obs2sample = one_hot(batch_index, self.n_batch)
obs_plate = self.create_plates(x_data, idx, batch_index)
# =====================Cell abundances w_sf======================= #
# factorisation prior on w_sf models similarity in locations
# across cell types f and reflects the absolute scale of w_sf
with obs_plate:
n_s_cells_per_location = pyro.sample(
"n_s_cells_per_location",
dist.Gamma(
self.N_cells_per_location * self.N_cells_mean_var_ratio,
self.N_cells_mean_var_ratio,
),
)
y_s_groups_per_location = pyro.sample(
"y_s_groups_per_location",
dist.Gamma(self.Y_groups_per_location, self.ones),
)
# cell group loadings
shape = self.ones_1_n_groups * y_s_groups_per_location / self.n_groups_tensor
rate = self.ones_1_n_groups / (n_s_cells_per_location /
y_s_groups_per_location)
with obs_plate:
z_sr_groups_factors = pyro.sample(
"z_sr_groups_factors",
dist.Gamma(
shape,
rate), # .to_event(1)#.expand([self.n_groups]).to_event(1)
) # (n_obs, n_groups)
k_r_factors_per_groups = pyro.sample(
"k_r_factors_per_groups",
dist.Gamma(self.factors_per_groups,
self.ones).expand([self.n_groups, 1]).to_event(2),
) # (self.n_groups, 1)
c2f_shape = k_r_factors_per_groups / self.n_factors_tensor
x_fr_group2fact = pyro.sample(
"x_fr_group2fact",
dist.Gamma(c2f_shape, k_r_factors_per_groups).expand(
[self.n_groups, self.n_factors]).to_event(2),
) # (self.n_groups, self.n_factors)
with obs_plate:
w_sf_mu = z_sr_groups_factors @ x_fr_group2fact
w_sf = pyro.sample(
"w_sf",
dist.Gamma(
w_sf_mu * self.w_sf_mean_var_ratio_tensor,
self.w_sf_mean_var_ratio_tensor,
),
) # (self.n_obs, self.n_factors)
# =====================Location-specific detection efficiency ======================= #
# y_s with hierarchical mean prior
detection_mean_y_e = pyro.sample(
"detection_mean_y_e",
dist.Gamma(
self.ones * self.detection_mean_hyp_prior_alpha,
self.ones * self.detection_mean_hyp_prior_beta,
).expand([self.n_batch, 1]).to_event(2),
)
detection_hyp_prior_alpha = pyro.deterministic(
"detection_hyp_prior_alpha",
self.ones_n_batch_1 * self.detection_hyp_prior_alpha,
)
beta = (obs2sample @ detection_hyp_prior_alpha) / (
obs2sample @ detection_mean_y_e)
with obs_plate:
detection_y_s = pyro.sample(
"detection_y_s",
dist.Gamma(obs2sample @ detection_hyp_prior_alpha, beta),
) # (self.n_obs, 1)
# =====================Gene-specific additive component ======================= #
# per gene molecule contribution that cannot be explained by
# cell state signatures (e.g. background, free-floating RNA)
s_g_gene_add_alpha_hyp = pyro.sample(
"s_g_gene_add_alpha_hyp",
dist.Gamma(self.gene_add_alpha_hyp_prior_alpha,
self.gene_add_alpha_hyp_prior_beta),
)
s_g_gene_add_mean = pyro.sample(
"s_g_gene_add_mean",
dist.Gamma(
self.gene_add_mean_hyp_prior_alpha,
self.gene_add_mean_hyp_prior_beta,
).expand([self.n_batch, 1]).to_event(2),
) # (self.n_batch)
s_g_gene_add_alpha_e_inv = pyro.sample(
"s_g_gene_add_alpha_e_inv",
dist.Exponential(s_g_gene_add_alpha_hyp).expand([self.n_batch,
1]).to_event(2),
) # (self.n_batch)
s_g_gene_add_alpha_e = self.ones / s_g_gene_add_alpha_e_inv.pow(2)
s_g_gene_add = pyro.sample(
"s_g_gene_add",
dist.Gamma(s_g_gene_add_alpha_e, s_g_gene_add_alpha_e /
s_g_gene_add_mean).expand([self.n_batch,
self.n_vars]).to_event(2),
) # (self.n_batch, n_vars)
# =====================Gene-specific overdispersion ======================= #
alpha_g_phi_hyp = pyro.sample(
"alpha_g_phi_hyp",
dist.Gamma(self.alpha_g_phi_hyp_prior_alpha,
self.alpha_g_phi_hyp_prior_beta),
)
alpha_g_inverse = pyro.sample(
"alpha_g_inverse",
dist.Exponential(alpha_g_phi_hyp).expand(
[self.n_batch, self.n_vars]).to_event(2),
) # (self.n_batch, self.n_vars)
# =====================Expected expression ======================= #
# expected expression
mu = ((w_sf @ self.cell_state) +
(obs2sample @ s_g_gene_add)) * detection_y_s
alpha = obs2sample @ (self.ones / alpha_g_inverse.pow(2))
# convert mean and overdispersion to total count and logits
# total_count, logits = _convert_mean_disp_to_counts_logits(
# mu, alpha, eps=self.eps
# )
# =====================DATA likelihood ======================= #
# Likelihood (sampling distribution) of data_target & add overdispersion via NegativeBinomial
with obs_plate:
pyro.sample(
"data_target",
dist.GammaPoisson(concentration=alpha, rate=alpha / mu),
# dist.NegativeBinomial(total_count=total_count, logits=logits),
obs=x_data,
)
# =====================Compute mRNA count from each factor in locations ======================= #
with obs_plate:
mRNA = w_sf * (self.cell_state).sum(-1)
pyro.deterministic("u_sf_mRNA_factors", mRNA)
def inference(
self,
x: torch.Tensor,
y: torch.Tensor,
batch_index: Optional[torch.Tensor] = None,
label: Optional[torch.Tensor] = None,
n_samples=1,
cont_covs=None,
cat_covs=None,
) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]:
"""
Internal helper function to compute necessary inference quantities.
We use the dictionary ``px_`` to contain the parameters of the ZINB/NB for genes.
The rate refers to the mean of the NB, dropout refers to Bernoulli mixing parameters.
`scale` refers to the quanity upon which differential expression is performed. For genes,
this can be viewed as the mean of the underlying gamma distribution.
We use the dictionary ``py_`` to contain the parameters of the Mixture NB distribution for proteins.
`rate_fore` refers to foreground mean, while `rate_back` refers to background mean. ``scale`` refers to
foreground mean adjusted for background probability and scaled to reside in simplex.
``back_alpha`` and ``back_beta`` are the posterior parameters for ``rate_back``. ``fore_scale`` is the scaling
factor that enforces `rate_fore` > `rate_back`.
``px_["r"]`` and ``py_["r"]`` are the inverse dispersion parameters for genes and protein, respectively.
Parameters
----------
x
tensor of values with shape ``(batch_size, n_input_genes)``
y
tensor of values with shape ``(batch_size, n_input_proteins)``
batch_index
array that indicates which batch the cells belong to with shape ``batch_size``
label
tensor of cell-types labels with shape (batch_size, n_labels)
n_samples
Number of samples to sample from approximate posterior
cont_covs
Continuous covariates to condition on
cat_covs
Categorical covariates to condition on
"""
x_ = x
y_ = y
if self.use_observed_lib_size:
library_gene = x.sum(1).unsqueeze(1)
if self.log_variational:
x_ = torch.log(1 + x_)
y_ = torch.log(1 + y_)
if cont_covs is not None and self.encode_covariates is True:
encoder_input = torch.cat((x_, y_, cont_covs), dim=-1)
else:
encoder_input = torch.cat((x_, y_), dim=-1)
if cat_covs is not None and self.encode_covariates is True:
categorical_input = torch.split(cat_covs, 1, dim=1)
else:
categorical_input = tuple()
qz_m, qz_v, ql_m, ql_v, latent, untran_latent = self.encoder(
encoder_input, batch_index, *categorical_input
)
z = latent["z"]
untran_z = untran_latent["z"]
untran_l = untran_latent["l"]
if not self.use_observed_lib_size:
library_gene = latent["l"]
if n_samples > 1:
qz_m = qz_m.unsqueeze(0).expand((n_samples, qz_m.size(0), qz_m.size(1)))
qz_v = qz_v.unsqueeze(0).expand((n_samples, qz_v.size(0), qz_v.size(1)))
untran_z = Normal(qz_m, qz_v.sqrt()).sample()
z = self.encoder.z_transformation(untran_z)
ql_m = ql_m.unsqueeze(0).expand((n_samples, ql_m.size(0), ql_m.size(1)))
ql_v = ql_v.unsqueeze(0).expand((n_samples, ql_v.size(0), ql_v.size(1)))
untran_l = Normal(ql_m, ql_v.sqrt()).sample()
if self.use_observed_lib_size:
library_gene = library_gene.unsqueeze(0).expand(
(n_samples, library_gene.size(0), library_gene.size(1))
)
else:
library_gene = self.encoder.l_transformation(untran_l)
# Background regularization
if self.gene_dispersion == "gene-label":
# px_r gets transposed - last dimension is nb genes
px_r = F.linear(one_hot(label, self.n_labels), self.px_r)
elif self.gene_dispersion == "gene-batch":
px_r = F.linear(one_hot(batch_index, self.n_batch), self.px_r)
elif self.gene_dispersion == "gene":
px_r = self.px_r
px_r = torch.exp(px_r)
if self.protein_dispersion == "protein-label":
# py_r gets transposed - last dimension is n_proteins
py_r = F.linear(one_hot(label, self.n_labels), self.py_r)
elif self.protein_dispersion == "protein-batch":
py_r = F.linear(one_hot(batch_index, self.n_batch), self.py_r)
elif self.protein_dispersion == "protein":
py_r = self.py_r
py_r = torch.exp(py_r)
if self.n_batch > 0:
py_back_alpha_prior = F.linear(
one_hot(batch_index, self.n_batch), self.background_pro_alpha
)
py_back_beta_prior = F.linear(
one_hot(batch_index, self.n_batch),
torch.exp(self.background_pro_log_beta),
)
else:
py_back_alpha_prior = self.background_pro_alpha
py_back_beta_prior = torch.exp(self.background_pro_log_beta)
self.back_mean_prior = Normal(py_back_alpha_prior, py_back_beta_prior)
return dict(
qz_m=qz_m,
qz_v=qz_v,
z=z,
untran_z=untran_z,
ql_m=ql_m,
ql_v=ql_v,
library_gene=library_gene,
untran_l=untran_l,
)
def batch(batch_size, label):
labels = torch.ones(batch_size, 1, device=x.device,
dtype=torch.long) * label
return one_hot(labels, y_dim)
def loss(
self,
tensors,
inference_outputs,
generative_outputs,
mode: Optional[int] = None,
kl_weight=1.0,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Return the reconstruction loss and the Kullback divergences.
Parameters
----------
x
tensor of values with shape ``(batch_size, n_input)``
or ``(batch_size, n_input_fish)`` depending on the mode
batch_index
array that indicates which batch the cells belong to with shape ``batch_size``
y
tensor of cell-types labels with shape (batch_size, n_labels)
mode
indicates which head/tail to use in the joint network
Returns
-------
the reconstruction loss and the Kullback divergences
"""
if mode is None:
if len(self.n_input_list) == 1:
mode = 0
else:
raise Exception("Must provide a mode")
x = tensors[_CONSTANTS.X_KEY]
batch_index = tensors[_CONSTANTS.BATCH_KEY]
qz_m = inference_outputs["qz_m"]
qz_v = inference_outputs["qz_v"]
ql_m = inference_outputs["ql_m"]
ql_v = inference_outputs["ql_v"]
px_rate = generative_outputs["px_rate"]
px_r = generative_outputs["px_r"]
px_dropout = generative_outputs["px_dropout"]
# mask loss to observed genes
mapping_indices = self.indices_mappings[mode]
reconstruction_loss = self.reconstruction_loss(
x,
px_rate[:, mapping_indices],
px_r[:, mapping_indices],
px_dropout[:, mapping_indices],
mode,
)
# KL Divergence
mean = torch.zeros_like(qz_m)
scale = torch.ones_like(qz_v)
kl_divergence_z = kl(Normal(qz_m, torch.sqrt(qz_v)), Normal(mean, scale)).sum(
dim=1
)
if self.model_library_bools[mode]:
library_log_means = getattr(self, f"library_log_means_{mode}")
library_log_vars = getattr(self, f"library_log_vars_{mode}")
local_library_log_means = F.linear(
one_hot(batch_index, self.n_batch), library_log_means
)
local_library_log_vars = F.linear(
one_hot(batch_index, self.n_batch), library_log_vars
)
kl_divergence_l = kl(
Normal(ql_m, torch.sqrt(ql_v)),
Normal(local_library_log_means, local_library_log_vars.sqrt()),
).sum(dim=1)
else:
kl_divergence_l = torch.zeros_like(kl_divergence_z)
kl_local = kl_divergence_l + kl_divergence_z
kl_global = torch.tensor(0.0)
loss = torch.mean(reconstruction_loss + kl_weight * kl_local) * x.size(0)
return LossRecorder(loss, reconstruction_loss, kl_local, kl_global)
def loss(
self,
tensors,
inference_outputs,
generative_outputs,
pro_recons_weight=1.0, # double check these defaults
kl_weight=1.0,
) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor,
torch.FloatTensor]:
"""
Returns the reconstruction loss and the Kullback divergences.
Parameters
----------
x
tensor of values with shape ``(batch_size, n_input_genes)``
y
tensor of values with shape ``(batch_size, n_input_proteins)``
batch_index
array that indicates which batch the cells belong to with shape ``batch_size``
label
tensor of cell-types labels with shape (batch_size, n_labels)
Returns
-------
type
the reconstruction loss and the Kullback divergences
"""
qz_m = inference_outputs["qz_m"]
qz_v = inference_outputs["qz_v"]
ql_m = inference_outputs["ql_m"]
ql_v = inference_outputs["ql_v"]
px_ = generative_outputs["px_"]
py_ = generative_outputs["py_"]
x = tensors[REGISTRY_KEYS.X_KEY]
batch_index = tensors[REGISTRY_KEYS.BATCH_KEY]
y = tensors[REGISTRY_KEYS.PROTEIN_EXP_KEY]
if self.protein_batch_mask is not None:
pro_batch_mask_minibatch = torch.zeros_like(y)
for b in torch.unique(batch_index):
b_indices = (batch_index == b).reshape(-1)
pro_batch_mask_minibatch[b_indices] = torch.tensor(
self.protein_batch_mask[b.item()].astype(np.float32),
device=y.device,
)
else:
pro_batch_mask_minibatch = None
reconst_loss_gene, reconst_loss_protein = self.get_reconstruction_loss(
x, y, px_, py_, pro_batch_mask_minibatch)
# KL Divergence
kl_div_z = kl(Normal(qz_m, torch.sqrt(qz_v)), Normal(0, 1)).sum(dim=1)
if not self.use_observed_lib_size:
n_batch = self.library_log_means.shape[1]
local_library_log_means = F.linear(one_hot(batch_index, n_batch),
self.library_log_means)
local_library_log_vars = F.linear(one_hot(batch_index, n_batch),
self.library_log_vars)
kl_div_l_gene = kl(
Normal(ql_m, torch.sqrt(ql_v)),
Normal(local_library_log_means,
torch.sqrt(local_library_log_vars)),
).sum(dim=1)
else:
kl_div_l_gene = 0.0
kl_div_back_pro_full = kl(Normal(py_["back_alpha"], py_["back_beta"]),
self.back_mean_prior)
if pro_batch_mask_minibatch is not None:
kl_div_back_pro = torch.zeros_like(kl_div_back_pro_full)
kl_div_back_pro.masked_scatter_(pro_batch_mask_minibatch.bool(),
kl_div_back_pro_full)
kl_div_back_pro = kl_div_back_pro.sum(dim=1)
else:
kl_div_back_pro = kl_div_back_pro_full.sum(dim=1)
loss = torch.mean(reconst_loss_gene +
pro_recons_weight * reconst_loss_protein +
kl_weight * kl_div_z + kl_div_l_gene +
kl_weight * kl_div_back_pro)
reconst_losses = dict(
reconst_loss_gene=reconst_loss_gene,
reconst_loss_protein=reconst_loss_protein,
)
kl_local = dict(
kl_div_z=kl_div_z,
kl_div_l_gene=kl_div_l_gene,
kl_div_back_pro=kl_div_back_pro,
)
return LossRecorder(loss,
reconst_losses,
kl_local,
kl_global=torch.tensor(0.0))
def forward(self, x_data, idx, batch_index, label_index,
extra_categoricals):
obs2sample = one_hot(batch_index, self.n_batch)
obs2label = one_hot(label_index, self.n_factors)
if self.n_extra_categoricals is not None:
obs2extra_categoricals = torch.cat(
[
one_hot(
extra_categoricals[:, i].view(
(extra_categoricals.shape[0], 1)),
n_cat,
) for i, n_cat in enumerate(self.n_extra_categoricals)
],
dim=1,
)
obs_plate = self.create_plates(x_data, idx, batch_index, label_index,
extra_categoricals)
# =====================Per-cluster average mRNA count ======================= #
# \mu_{f,g}
per_cluster_mu_fg = pyro.sample(
"per_cluster_mu_fg",
dist.Gamma(self.ones,
self.ones).expand([self.n_factors,
self.n_vars]).to_event(2),
)
# =====================Gene-specific multiplicative component ======================= #
# `y_{t, g}` per gene multiplicative effect that explains the difference
# in sensitivity between genes in each technology or covariate effect
if self.n_extra_categoricals is not None:
detection_tech_gene_tg = pyro.sample(
"detection_tech_gene_tg",
dist.Gamma(
self.ones * self.gene_tech_prior_alpha,
self.ones * self.gene_tech_prior_beta,
).expand([np.sum(self.n_extra_categoricals),
self.n_vars]).to_event(2),
)
# =====================Cell-specific detection efficiency ======================= #
# y_c with hierarchical mean prior
detection_mean_y_e = pyro.sample(
"detection_mean_y_e",
dist.Gamma(
self.ones * self.detection_mean_hyp_prior_alpha,
self.ones * self.detection_mean_hyp_prior_beta,
).expand([self.n_batch, 1]).to_event(2),
)
detection_y_c = obs2sample @ detection_mean_y_e # (self.n_obs, 1)
# =====================Gene-specific additive component ======================= #
# s_{e,g} accounting for background, free-floating RNA
s_g_gene_add_alpha_hyp = pyro.sample(
"s_g_gene_add_alpha_hyp",
dist.Gamma(self.ones * self.gene_add_alpha_hyp_prior_alpha,
self.ones * self.gene_add_alpha_hyp_prior_beta),
)
s_g_gene_add_mean = pyro.sample(
"s_g_gene_add_mean",
dist.Gamma(
self.gene_add_mean_hyp_prior_alpha,
self.gene_add_mean_hyp_prior_beta,
).expand([self.n_batch, 1]).to_event(2),
) # (self.n_batch)
s_g_gene_add_alpha_e_inv = pyro.sample(
"s_g_gene_add_alpha_e_inv",
dist.Exponential(s_g_gene_add_alpha_hyp).expand([self.n_batch,
1]).to_event(2),
) # (self.n_batch)
s_g_gene_add_alpha_e = self.ones / s_g_gene_add_alpha_e_inv.pow(2)
s_g_gene_add = pyro.sample(
"s_g_gene_add",
dist.Gamma(s_g_gene_add_alpha_e, s_g_gene_add_alpha_e /
s_g_gene_add_mean).expand([self.n_batch,
self.n_vars]).to_event(2),
) # (self.n_batch, n_vars)
# =====================Gene-specific overdispersion ======================= #
alpha_g_phi_hyp = pyro.sample(
"alpha_g_phi_hyp",
dist.Gamma(self.ones * self.alpha_g_phi_hyp_prior_alpha,
self.ones * self.alpha_g_phi_hyp_prior_beta),
)
alpha_g_inverse = pyro.sample(
"alpha_g_inverse",
dist.Exponential(alpha_g_phi_hyp).expand([1, self.n_vars
]).to_event(2),
) # (self.n_batch or 1, self.n_vars)
# =====================Expected expression ======================= #
# overdispersion
alpha = self.ones / alpha_g_inverse.pow(2)
# biological expression
mu = (
obs2label @ per_cluster_mu_fg +
obs2sample @ s_g_gene_add # contaminating RNA
) * detection_y_c # cell-specific normalisation
if self.n_extra_categoricals is not None:
# gene-specific normalisation for covatiates
mu = mu * (obs2extra_categoricals @ detection_tech_gene_tg)
# total_count, logits = _convert_mean_disp_to_counts_logits(
# mu, alpha, eps=self.eps
# )
# =====================DATA likelihood ======================= #
# Likelihood (sampling distribution) of data_target & add overdispersion via NegativeBinomial
with obs_plate:
pyro.sample(
"data_target",
dist.GammaPoisson(concentration=alpha, rate=alpha / mu),
# dist.NegativeBinomial(total_count=total_count, logits=logits),
obs=x_data,
)
