def evaluate(self, adata, batch_key):
adata = remove_sparsity(adata)
encoder_labels, _ = label_encoder(adata, self.condition_encoder, batch_key)
decoder_labels, _ = label_encoder(adata, self.condition_encoder, batch_key)
encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions)
decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions)
cvae_inputs = [adata.X, encoder_labels, decoder_labels]
encoded_labels = self.cvae_model.predict(cvae_inputs)[2].argmax(axis=1)
self._reverse_cell_type_encoder()
labels = []
for encoded_label in encoded_labels:
labels.append(self.inv_cell_type_encoder[encoded_label])
labels = np.array(labels)
true_labels = adata.obs[batch_key].values
accuracy = np.mean(labels == true_labels)
print(classification_report(true_labels, labels))
return accuracy, confusion_matrix(true_labels, labels)
def get_latent(self, adata, batch_key):
""" Transforms `adata` in latent space of CVAE and returns the latent
coordinates in the annotated (adata) format.
Parameters
----------
adata: :class:`~anndata.AnnData`
Annotated dataset matrix in Primary space.
"""
if set(self.gene_names).issubset(set(adata.var_names)):
adata = adata[:, self.gene_names]
else:
raise Exception(
"set of gene names in train adata are inconsistent with scNet's gene_names"
)
encoder_labels, _ = label_encoder(adata, self.condition_encoder,
batch_key)
encoder_labels = to_categorical(encoder_labels,
num_classes=self.n_conditions)
return self.get_z_latent(adata, encoder_labels)
def get_latent(self, adata, batch_key, return_z=False):
""" Transforms `adata` in latent space of scNet and returns the latent
coordinates in the annotated (adata) format.
Parameters
----------
adata: :class:`~anndata.AnnData`
Annotated dataset matrix in Primary space.
batch_key: str
Name of the column containing the study (batch) names for each sample.
return_z: bool
``False`` by defaul. if ``True``, the output of bottleneck layer of network will be computed.
Returns
-------
adata_pred: `~anndata.AnnData`
Annotated data of transformed ``adata`` into latent space.
"""
if set(self.gene_names).issubset(set(adata.var_names)):
adata = adata[:, self.gene_names]
else:
raise Exception("set of gene names in train adata are inconsistent with scNet's gene_names")
if self.beta == 0:
return_z = True
encoder_labels, _ = label_encoder(adata, self.condition_encoder, batch_key)
encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions)
if return_z or self.beta == 0:
return self.get_z_latent(adata, encoder_labels)
else:
return self.to_mmd_layer(adata, batch_key)
def __init__(self,
filename: str,
adata: anndata.AnnData,
batch_key: str,
cell_type_key: str,
encoder_model: Model,
n_per_epoch: int = 5,
n_batch_labels: int = 0,
n_celltype_labels: int = 0,
clustering_scores: list_or_str = 'all'):
super(ScoreCallback, self).__init__()
self.adata = remove_sparsity(adata)
self.batch_labels, _ = label_encoder(adata,
le=None,
condition_key=batch_key)
self.batch_labels = np.reshape(self.batch_labels, (-1, ))
self.batch_labels_onehot = to_categorical(self.batch_labels,
num_classes=n_batch_labels)
self.celltype_labels, _ = label_encoder(adata,
le=None,
condition_key=cell_type_key)
self.celltype_labels = np.reshape(self.celltype_labels, (-1, ))
self.celltype_labels_onehot = to_categorical(
self.celltype_labels, num_classes=n_celltype_labels)
self.filename = filename
self.encoder_model = encoder_model
self.n_per_epoch = n_per_epoch
self.n_batch_labels = n_batch_labels
self.n_celltype_labels = n_celltype_labels
self.clustering_scores = clustering_scores
self.score_computers = {
"asw": self.asw,
"ari": self.ari,
"nmi": self.nmi,
"ebm": self.entropy_of_batch_mixing,
"knn": self.knn_purity
}
self.kmeans_batch = KMeans(self.n_batch_labels, n_init=200)
self.kmeans_celltype = KMeans(self.n_celltype_labels, n_init=200)
def annotate(self, adata, batch_key, cell_type_key):
adata = remove_sparsity(adata)
encoder_labels, _ = label_encoder(adata, self.condition_encoder, batch_key)
decoder_labels, _ = label_encoder(adata, self.condition_encoder, batch_key)
encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions)
decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions)
cvae_inputs = [adata.X, encoder_labels, decoder_labels]
encoded_labels = self.cvae_model.predict(cvae_inputs)[2].argmax(axis=1)
self._reverse_cell_type_encoder()
labels = []
for encoded_label in encoded_labels:
labels.append(self.inv_cell_type_encoder[encoded_label])
adata.obs[f'pred_{cell_type_key}'] = np.array(labels)
def to_mmd_layer(self, adata, batch_key):
"""
Map ``adata`` in to the MMD space. This function will feed data
in ``mmd_model`` of scArches and compute the MMD space coordinates
for each sample in data.
Parameters
----------
adata: :class:`~anndata.AnnData`
Annotated data matrix to be mapped to MMD latent space.
Please note that ``adata.X`` has to be in shape [n_obs, x_dimension]
encoder_labels: :class:`~numpy.ndarray`
:class:`~numpy.ndarray` of labels to be fed as scArches' encoder condition array.
decoder_labels: :class:`~numpy.ndarray`
:class:`~numpy.ndarray` of labels to be fed as scArches' decoder condition array.
Returns
-------
adata_mmd: :class:`~anndata.AnnData`
returns Annotated data containing MMD latent space encoding of ``adata``
"""
adata = remove_sparsity(adata)
encoder_labels, _ = label_encoder(adata, self.condition_encoder,
batch_key)
decoder_labels, _ = label_encoder(adata, self.condition_encoder,
batch_key)
encoder_labels = to_categorical(encoder_labels,
num_classes=self.n_conditions)
decoder_labels = to_categorical(decoder_labels,
num_classes=self.n_conditions)
cvae_inputs = [adata.X, encoder_labels, decoder_labels]
mmd = self.cvae_model.predict(cvae_inputs)[1]
mmd = np.nan_to_num(mmd, nan=0.0, posinf=0.0, neginf=0.0)
adata_mmd = anndata.AnnData(X=mmd)
adata_mmd.obs = adata.obs.copy(deep=True)
return adata_mmd
def get_latent(self, adata, batch_key, return_mean=False):
""" Transforms `adata` in latent space of CVAE and returns the latent
coordinates in the annotated (adata) format.
Parameters
----------
adata: :class:`~anndata.AnnData`
Annotated dataset matrix in Primary space.
batch_key: str
key for the observation that has batch labels in adata.obs.
return_mean: bool
if False, z will be sampled. Set to `True` if want a fix z value every time you call
get_latent.
Returns
-------
latent_adata: :class:`~anndata.AnnData`
Annotated dataset matrix in Latent space.
"""
if set(self.gene_names).issubset(set(adata.var_names)):
adata = adata[:, self.gene_names]
else:
raise Exception(
"set of gene names in train adata are inconsistent with scNet's gene_names"
)
encoder_labels, _ = label_encoder(adata, self.condition_encoder,
batch_key)
encoder_labels = to_categorical(encoder_labels,
num_classes=self.n_conditions)
return self.get_z_latent(adata, encoder_labels, return_mean)
def _train_on_batch(self,
adata,
condition_key,
train_size=0.8,
cell_type_key='cell_type',
n_epochs=100,
batch_size=128,
early_stop_limit=10,
lr_reducer=8,
n_per_epoch=0,
score_filename=None,
save=True,
retrain=True,
verbose=3):
train_adata, valid_adata = train_test_split(adata, train_size)
if self.gene_names is None:
self.gene_names = train_adata.var_names.tolist()
else:
if set(self.gene_names).issubset(set(train_adata.var_names)):
train_adata = train_adata[:, self.gene_names]
else:
raise Exception(
"set of gene names in train adata are inconsistent with class' gene_names"
)
if set(self.gene_names).issubset(set(valid_adata.var_names)):
valid_adata = valid_adata[:, self.gene_names]
else:
raise Exception(
"set of gene names in valid adata are inconsistent with class' gene_names"
)
train_conditions_encoded, self.condition_encoder = label_encoder(
train_adata,
le=self.condition_encoder,
condition_key=condition_key)
valid_conditions_encoded, self.condition_encoder = label_encoder(
valid_adata,
le=self.condition_encoder,
condition_key=condition_key)
if not retrain and os.path.exists(
os.path.join(self.model_path, f"{self.model_name}.h5")):
self.restore_model_weights()
return
train_conditions_onehot = to_categorical(train_conditions_encoded,
num_classes=self.n_conditions)
valid_conditions_onehot = to_categorical(valid_conditions_encoded,
num_classes=self.n_conditions)
if sparse.issparse(train_adata.X):
is_sparse = True
else:
is_sparse = False
train_expr = train_adata.X
valid_expr = valid_adata.X.A if is_sparse else valid_adata.X
x_valid = [
valid_expr, valid_conditions_onehot, valid_conditions_onehot
]
if self.loss_fn in ['nb', 'zinb']:
x_valid.append(valid_adata.obs[self.size_factor_key].values)
y_valid = valid_adata.raw.X.A if sparse.issparse(
valid_adata.raw.X) else valid_adata.raw.X
else:
y_valid = valid_expr
es_patience, best_val_loss = 0, 1e10
for i in range(n_epochs):
train_loss = train_recon_loss = train_kl_loss = 0.0
for j in range(min(500, train_adata.shape[0] // batch_size)):
batch_indices = np.random.choice(train_adata.shape[0],
batch_size)
batch_expr = train_expr[
batch_indices, :].A if is_sparse else train_expr[
batch_indices, :]
x_train = [
batch_expr, train_conditions_onehot[batch_indices],
train_conditions_onehot[batch_indices]
]
if self.loss_fn in ['nb', 'zinb']:
x_train.append(train_adata.obs[
self.size_factor_key].values[batch_indices])
y_train = train_adata.raw.X[
batch_indices].A if sparse.issparse(
train_adata.raw.X[batch_indices]
) else train_adata.raw.X[batch_indices]
else:
y_train = batch_expr
batch_loss, batch_recon_loss, batch_kl_loss = self.cvae_model.train_on_batch(
x_train, y_train)
train_loss += batch_loss / batch_size
train_recon_loss += batch_recon_loss / batch_size
train_kl_loss += batch_kl_loss / batch_size
valid_loss, valid_recon_loss, valid_kl_loss = self.cvae_model.evaluate(
x_valid, y_valid, verbose=0)
if valid_loss < best_val_loss:
best_val_loss = valid_loss
es_patience = 0
else:
es_patience += 1
if es_patience == early_stop_limit:
print("Training stopped with Early Stopping")
break
logs = {
"loss": train_loss,
"recon_loss": train_recon_loss,
"kl_loss": train_kl_loss,
"val_loss": valid_loss,
"val_recon_loss": valid_recon_loss,
"val_kl_loss": valid_kl_loss
}
print_progress(i, logs, n_epochs)
if save:
self.update_kwargs()
self.save(make_dir=True)
def _fit(self,
adata,
condition_key,
train_size=0.8,
cell_type_key='cell_type',
n_epochs=100,
batch_size=128,
early_stop_limit=10,
lr_reducer=8,
n_per_epoch=0,
score_filename=None,
save=True,
retrain=True,
verbose=3):
train_adata, valid_adata = train_test_split(adata, train_size)
if self.gene_names is None:
self.gene_names = train_adata.var_names.tolist()
else:
if set(self.gene_names).issubset(set(train_adata.var_names)):
train_adata = train_adata[:, self.gene_names]
else:
raise Exception(
"set of gene names in train adata are inconsistent with class' gene_names"
)
if set(self.gene_names).issubset(set(valid_adata.var_names)):
valid_adata = valid_adata[:, self.gene_names]
else:
raise Exception(
"set of gene names in valid adata are inconsistent with class' gene_names"
)
if self.loss_fn in ['nb', 'zinb']:
train_raw_expr = train_adata.raw.X.A if sparse.issparse(
train_adata.raw.X) else train_adata.raw.X
valid_raw_expr = valid_adata.raw.X.A if sparse.issparse(
valid_adata.raw.X) else valid_adata.raw.X
train_expr = train_adata.X.A if sparse.issparse(
train_adata.X) else train_adata.X
valid_expr = valid_adata.X.A if sparse.issparse(
valid_adata.X) else valid_adata.X
train_conditions_encoded, self.condition_encoder = label_encoder(
train_adata,
le=self.condition_encoder,
condition_key=condition_key)
valid_conditions_encoded, self.condition_encoder = label_encoder(
valid_adata,
le=self.condition_encoder,
condition_key=condition_key)
if not retrain and os.path.exists(
os.path.join(self.model_path, f"{self.model_name}.h5")):
self.restore_model_weights()
return
callbacks = [
History(),
]
if verbose > 2:
callbacks.append(
LambdaCallback(on_epoch_end=lambda epoch, logs: print_progress(
epoch, logs, n_epochs)))
fit_verbose = 0
else:
fit_verbose = verbose
if (n_per_epoch > 0 or n_per_epoch == -1) and not score_filename:
adata = train_adata.concatenate(valid_adata)
train_celltypes_encoded, _ = label_encoder(
train_adata, le=None, condition_key=cell_type_key)
valid_celltypes_encoded, _ = label_encoder(
valid_adata, le=None, condition_key=cell_type_key)
celltype_labels = np.concatenate(
[train_celltypes_encoded, valid_celltypes_encoded], axis=0)
callbacks.append(
ScoreCallback(score_filename,
adata,
condition_key,
cell_type_key,
self.cvae_model,
n_per_epoch=n_per_epoch,
n_batch_labels=self.n_conditions,
n_celltype_labels=len(
np.unique(celltype_labels))))
if early_stop_limit > 0:
callbacks.append(
EarlyStopping(patience=early_stop_limit, monitor='val_loss'))
if lr_reducer > 0:
callbacks.append(
ReduceLROnPlateau(monitor='val_loss', patience=lr_reducer))
train_conditions_onehot = to_categorical(train_conditions_encoded,
num_classes=self.n_conditions)
valid_conditions_onehot = to_categorical(valid_conditions_encoded,
num_classes=self.n_conditions)
x_train = [
train_expr, train_conditions_onehot, train_conditions_onehot
]
x_valid = [
valid_expr, valid_conditions_onehot, valid_conditions_onehot
]
if self.loss_fn in ['nb', 'zinb']:
x_train.append(train_adata.obs[self.size_factor_key].values)
y_train = train_raw_expr
x_valid.append(valid_adata.obs[self.size_factor_key].values)
y_valid = valid_raw_expr
else:
y_train = train_expr
y_valid = valid_expr
self.cvae_model.fit(
x=x_train,
y=y_train,
validation_data=(x_valid, y_valid),
epochs=n_epochs,
batch_size=batch_size,
verbose=fit_verbose,
callbacks=callbacks,
)
if save:
self.update_kwargs()
self.save(make_dir=True)
def _fit(self, adata, condition_key, cell_type_key,
train_size=0.8,
n_epochs=25, batch_size=32,
early_stop_limit=20, lr_reducer=10,
n_per_epoch=0, score_filename=None,
save=True, retrain=True, verbose=3):
"""
Trains scNet with ``n_epochs`` times given ``train_adata``
and validates the model using ``valid_adata``
This function is using ``early stopping`` and ``learning rate reduce on plateau``
techniques to prevent over-fitting.
Parameters
----------
adata: :class:`~anndata.AnnData`
Annotated dataset used to train & evaluate scNet.
condition_key: str
column name for conditions in the `obs` matrix of `train_adata` and `valid_adata`.
train_size: float
fraction of samples used to train scNet.
n_epochs: int
number of epochs.
batch_size: int
number of samples in the mini-batches used to optimize scNet.
early_stop_limit: int
patience of EarlyStopping
lr_reducer: int
patience of LearningRateReduceOnPlateau.
save: bool
Whether to save scNet after the training or not.
verbose: int
Verbose level
retrain: bool
``True`` by default. if ``True`` scNet will be trained regardless of existance of pre-trained scNet in ``model_path``. if ``False`` scNet will not be trained if pre-trained scNet exists in ``model_path``.
"""
train_adata, valid_adata = train_test_split(adata, train_size)
if self.gene_names is None:
self.gene_names = train_adata.var_names.tolist()
else:
if set(self.gene_names).issubset(set(train_adata.var_names)):
train_adata = train_adata[:, self.gene_names]
else:
raise Exception("set of gene names in train adata are inconsistent with scNet's gene_names")
if set(self.gene_names).issubset(set(valid_adata.var_names)):
valid_adata = valid_adata[:, self.gene_names]
else:
raise Exception("set of gene names in valid adata are inconsistent with scNet's gene_names")
train_conditions_encoded, self.condition_encoder = label_encoder(train_adata, le=self.condition_encoder,
condition_key=condition_key)
valid_conditions_encoded, self.condition_encoder = label_encoder(valid_adata, le=self.condition_encoder,
condition_key=condition_key)
train_cell_types_encoded, encoder = label_encoder(train_adata, le=self.cell_type_encoder,
condition_key=cell_type_key)
if self.cell_type_encoder is None:
self.cell_type_encoder = encoder
valid_cell_types_encoded, self.cell_type_encoder = label_encoder(valid_adata, le=self.cell_type_encoder,
condition_key=cell_type_key)
if not retrain and self.restore_model_weights():
return
train_conditions_onehot = to_categorical(train_conditions_encoded, num_classes=self.n_conditions)
valid_conditions_onehot = to_categorical(valid_conditions_encoded, num_classes=self.n_conditions)
train_cell_types_onehot = to_categorical(train_cell_types_encoded, num_classes=self.n_classes)
valid_cell_types_onehot = to_categorical(valid_cell_types_encoded, num_classes=self.n_classes)
if self.loss_fn in ['nb', 'zinb']:
train_raw_expr = train_adata.raw.X.A if sparse.issparse(train_adata.raw.X) else train_adata.raw.X
valid_raw_expr = valid_adata.raw.X.A if sparse.issparse(valid_adata.raw.X) else valid_adata.raw.X
train_expr = train_adata.X.A if sparse.issparse(train_adata.X) else train_adata.X
valid_expr = valid_adata.X.A if sparse.issparse(valid_adata.X) else valid_adata.X
x_train = [train_expr, train_conditions_onehot, train_conditions_onehot]
y_train = [train_expr, train_conditions_encoded, train_cell_types_onehot]
x_valid = [valid_expr, valid_conditions_onehot, valid_conditions_onehot]
y_valid = [valid_expr, valid_conditions_encoded, valid_cell_types_onehot]
callbacks = [
History(),
]
if verbose > 2:
callbacks.append(
LambdaCallback(on_epoch_end=lambda epoch, logs: print_progress(epoch, logs, n_epochs)))
fit_verbose = 0
else:
fit_verbose = verbose
if (n_per_epoch > 0 or n_per_epoch == -1) and not score_filename:
adata = train_adata.concatenate(valid_adata)
celltype_labels = np.concatenate([train_cell_types_encoded, valid_cell_types_encoded], axis=0)
callbacks.append(ScoreCallback(score_filename, adata, condition_key, cell_type_key, self.cvae_model,
n_per_epoch=n_per_epoch, n_batch_labels=self.n_conditions,
n_celltype_labels=len(np.unique(celltype_labels))))
if early_stop_limit > 0:
callbacks.append(EarlyStopping(patience=early_stop_limit, monitor='val_loss'))
if lr_reducer > 0:
callbacks.append(ReduceLROnPlateau(monitor='val_loss', patience=lr_reducer))
self.cvae_model.fit(x=x_train,
y=y_train,
validation_data=(x_valid, y_valid),
epochs=n_epochs,
batch_size=batch_size,
verbose=fit_verbose,
callbacks=callbacks,
)
if save:
self.update_kwargs()
self.save(make_dir=True)
def make_dataset(adata,
condition_key,
le,
batch_size,
n_epochs,
is_training,
loss_fn,
n_conditions,
size_factor_key=None,
use_mmd=False):
if sparse.issparse(adata.X):
expressions = adata.X.A
else:
expressions = adata.X
encoded_conditions, le = label_encoder(adata, le, condition_key)
if loss_fn == 'nb':
if sparse.issparse(adata.raw.X):
raw_expressions = adata.raw.X.A
else:
raw_expressions = adata.raw.X
dataset = tf.data.Dataset.from_tensor_slices(({
"expression":
expressions,
"encoder_label":
encoded_conditions,
"decoder_label":
encoded_conditions,
"size_factor":
adata.obs[size_factor_key].values
}, {
"reconstruction":
raw_expressions
}))
elif loss_fn == 'zinb':
if sparse.issparse(adata.raw.X):
raw_expressions = adata.raw.X.A
else:
raw_expressions = adata.raw.X
dataset = tf.data.Dataset.from_tensor_slices(({
"expression":
expressions,
"encoder_label":
encoded_conditions,
"decoder_label":
encoded_conditions,
'size_factor':
adata.obs[size_factor_key].values
}, {
"reconstruction":
raw_expressions
}))
else:
if use_mmd:
dataset = tf.data.Dataset.from_tensor_slices(({
"expression":
expressions,
"encoder_label":
encoded_conditions,
"decoder_label":
encoded_conditions
}, {
"reconstruction":
expressions,
"mmd":
encoded_conditions
}))
else:
dataset = tf.data.Dataset.from_tensor_slices(({
"expression":
expressions,
"encoder_label":
encoded_conditions,
"decoder_label":
encoded_conditions
}, {
"reconstruction":
expressions
}))
if is_training:
dataset = dataset.shuffle(1000)
dataset = dataset.map(preprocess_cvae_input(n_conditions),
num_parallel_calls=4,
deterministic=None)
if is_training:
dataset = dataset.repeat(n_epochs)
else:
dataset = dataset.repeat()
dataset = dataset.batch(batch_size,
drop_remainder=True if is_training else False)
dataset = dataset.prefetch(buffer_size=5)
return dataset, le