PKGOtrvae/__init__.py"""trVAE - Regularized Conditional Variational Autoencoders""" from . import models as archs from . import utils as tl from . import data_loader as dl from . import plotting as pl from . import metrics as mt __author__ = ', '.join([ 'Mohsen Naghipourfar', 'Mohammad Lotfollahi' ]) __email__ = ', '.join([ 'mohsen.naghipourfar@gmail.com', 'Mohammad.lotfollahi@helmholtz-muenchen.de', ]) from get_version import get_version __version__ = get_version(__file__) del get_version PKGO*#trvae/data_loader.pyimport os import tarfile import zipfile import anndata import cv2 import keras import numpy as np import pandas as pd import scanpy as sc from PIL import Image def prepare_and_load_celeba(file_path, attr_path, landmark_path, gender='Male', attribute='Smiling', max_n_images=None, restore=True, save=True, img_width=64, img_height=78, verbose=True): data_path = os.path.dirname(file_path) zip_filename = os.path.basename(file_path).split(".")[0] if restore and os.path.exists( os.path.join(data_path, f"celeba_{attribute}_{img_width}x{img_height}_{max_n_images}.h5ad")): return sc.read(os.path.join(data_path, f"celeba_{attribute}_{img_width}x{img_height}_{max_n_images}.h5ad")) def load_attr_list(file_path): indices = [] attributes = [] with open(file_path) as f: lines = f.read().splitlines() columns = lines[1].split(" ") columns.remove('') for i in range(2, len(lines)): elements = lines[i].split() indices.append(elements[0]) attributes.append(list(map(int, elements[1:]))) attr_df = pd.DataFrame(attributes) attr_df.index = indices attr_df.columns = columns if verbose: print(attr_df.shape[0]) return attr_df def load_landmark_list(file_path): indices = [] landmarks = [] with open(file_path) as f: lines = f.read().splitlines() columns = lines[1].split(" ") for i in range(2, len(lines)): elements = lines[i].split() indices.append(elements[0]) landmarks.append(list(map(int, elements[1:]))) landmarks_df = pd.DataFrame(landmarks) landmarks_df.index = indices landmarks_df.columns = columns print(landmarks_df.shape[0]) return landmarks_df images = [] zfile = zipfile.ZipFile(file_path) counter = 0 attr_df = load_attr_list(attr_path) landmarks = load_landmark_list(landmark_path) landmarks = landmarks[abs(landmarks['lefteye_x'] - landmarks['righteye_x']) > 30] landmarks = landmarks[abs(landmarks['lefteye_x'] - landmarks['nose_x']) > 15] landmarks = landmarks[abs(landmarks['righteye_x'] - landmarks['nose_x']) > 15] landmarks.head() attr_df = attr_df.loc[landmarks.index] print("# of images after preprocessing: ", attr_df.shape[0]) indices = [] for filename in attr_df.index.tolist(): ifile = zfile.open(os.path.join(f"{zip_filename}/", filename)) image = Image.open(ifile) image_landmarks = landmarks.loc[filename] most_left_x = max(0, min(image_landmarks['lefteye_x'], image_landmarks['leftmouth_x']) - 15) most_right_x = min(178, min(image_landmarks['righteye_x'], image_landmarks['rightmouth_x']) + 15) most_up_y = max(0, image_landmarks['lefteye_y'] - 35) most_down_y = min(218, image_landmarks['rightmouth_y'] + 25) image_cropped = image.crop((most_left_x, most_up_y, most_right_x, most_down_y)) image_cropped = image_cropped.resize((img_width, img_height), Image.NEAREST) image = image_cropped image = np.reshape(image, (img_width, img_height, 3)) if max_n_images is None: images.append(image) indices.append(filename) counter += 1 if verbose and counter % 1000 == 0: print(counter) else: if counter < max_n_images: images.append(image) indices.append(filename) counter += 1 if verbose and counter % 1000 == 0: print(counter) else: break images = np.array(images) if verbose: print(images.shape) images_df = pd.DataFrame(images.reshape(-1, np.prod(images.shape[1:]))) images_df.index = indices if save: data = anndata.AnnData(X=images_df.values) attr_df = attr_df.loc[images_df.index] print(data.shape, attr_df.shape) data.obs['labels'] = attr_df[gender].values data.obs['condition'] = attr_df[attribute].values sc.write(filename=os.path.join(data_path, f"celeba_{attribute}_{img_width}x{img_height}_{max_n_images}.h5ad"), adata=data) return data def prepare_and_load_edge2shoe(file_path, restore=True, save=True, img_width=64, img_height=64, verbose=True): data_path = os.path.dirname(file_path) if restore and os.path.exists(os.path.join(data_path, f"edges2shoes_{img_width}x{img_height}.h5ad")): return sc.read(os.path.join(data_path, f"edges2shoes_{img_width}x{img_height}.h5ad")) tar = tarfile.open(file_path) images, edges = [], [] counter = 0 for member in tar.getmembers(): if member.name.endswith(".jpg"): f = tar.extractfile(member) image = Image.open(f) edge, image = image.crop((0, 0, 256, 256)), image.crop((256, 0, 512, 256)) edge = edge.resize((64, 64), Image.BICUBIC) image = image.resize((64, 64), Image.NEAREST) edge = np.array(edge) image = np.array(image) images.append(image) edges.append(edge) counter += 1 if verbose and counter % 1000 == 0: print(counter) images = np.array(images) edges = np.array(edges) images = images.reshape(-1, np.prod(images.shape[1:])) edges = edges.reshape(-1, np.prod(edges.shape[1:])) data = np.concatenate([images, edges], axis=0) if save: data = anndata.AnnData(X=data) data.obs['id'] = np.concatenate([np.arange(images.shape[0]), np.arange(images.shape[0])]) data.obs['condition'] = ['shoe'] * images.shape[0] + ['edge'] * images.shape[0] sc.write(filename=os.path.join(data_path, f"edges2shoes_{img_width}x{img_height}.h5ad"), adata=data) return data def resize_image(images, img_width, img_height): images_list = [] for i in range(images.shape[0]): image = cv2.resize(images[i], (img_width, img_height), cv2.INTER_NEAREST) images_list.append(image) return np.array(images_list) class PairedDataSequence(keras.utils.Sequence): def __init__(self, image_paths, batch_size): self.image_paths = image_paths self.batch_size = batch_size def __len__(self): return len(self.image_paths) def __getitem__(self, idx): edges, images = [], [] batch_image_paths = self.image_paths[idx:idx + self.batch_size] for image_path in batch_image_paths: with Image.open(image_path) as image: edge = np.array(image.crop((0, 0, 256, 256)).resize((64, 64), Image.BICUBIC)) image = np.array(image.crop((256, 0, 512, 256)).resize((64, 64), Image.NEAREST)) edges.append(edge) images.append(image) edges = np.array(edges) images = np.array(images) edges = edges.astype(np.float32) images = images.astype(np.float32) # Pre-processing edges /= 255.0 images /= 255.0 x = np.concatenate([edges, edges, images, images], axis=0) y = np.concatenate([edges, images, images, edges], axis=0) encoder_labels_feed = np.concatenate([np.zeros(edges.shape[0]), np.zeros(edges.shape[0]), np.ones(images.shape[0]), np.ones(images.shape[0])]) decoder_labels_feed = np.concatenate([np.zeros(edges.shape[0]), np.ones(edges.shape[0]), np.ones(images.shape[0]), np.zeros(images.shape[0])]) x_feed = [x, encoder_labels_feed, decoder_labels_feed] y_feed = [y, encoder_labels_feed] return x_feed, y_feed PKGOg۽ trvae/metrics.pyimport numpy as np from scipy.stats import itemfreq, entropy from sklearn.cluster import KMeans from sklearn.metrics import silhouette_score, adjusted_rand_score, normalized_mutual_info_score from sklearn.neighbors import NearestNeighbors from sklearn.preprocessing import LabelEncoder from trvae.utils import remove_sparsity def __entropy_from_indices(indices): return entropy(np.array(itemfreq(indices)[:, 1].astype(np.int32))) def entropy_batch_mixing(adata, label_key='batch', n_neighbors=50, n_pools=50, n_samples_per_pool=100, subsample_frac=1.0): adata = remove_sparsity(adata) n_samples = adata.shape[0] keep_idx = np.random.choice(np.arange(n_samples), size=min(n_samples, int(subsample_frac * n_samples)), replace=False) adata = adata[keep_idx, :] neighbors = NearestNeighbors(n_neighbors=n_neighbors + 1).fit(adata.X) indices = neighbors.kneighbors(adata.X, return_distance=False)[:, 1:] batch_indices = np.vectorize(lambda i: adata.obs[label_key].values[i])(indices) entropies = np.apply_along_axis(__entropy_from_indices, axis=1, arr=batch_indices) # average n_pools entropy results where each result is an average of n_samples_per_pool random samples. if n_pools == 1: score = np.mean(entropies) else: score = np.mean([ np.mean(entropies[np.random.choice(len(entropies), size=n_samples_per_pool)]) for _ in range(n_pools) ]) return score def asw(adata, label_key): adata = remove_sparsity(adata) labels = adata.obs[label_key].values labels_encoded = LabelEncoder().fit_transform(labels) return silhouette_score(adata.X, labels_encoded) def ari(adata, label_key): adata = remove_sparsity(adata) n_labels = len(adata.obs[label_key].unique().tolist()) kmeans = KMeans(n_labels, n_init=200) labels_pred = kmeans.fit_predict(adata.X) labels = adata.obs[label_key].values labels_encoded = LabelEncoder().fit_transform(labels) return adjusted_rand_score(labels_encoded, labels_pred) def nmi(adata, label_key): adata = remove_sparsity(adata) n_labels = len(adata.obs[label_key].unique().tolist()) kmeans = KMeans(n_labels, n_init=200) labels_pred = kmeans.fit_predict(adata.X) labels = adata.obs[label_key].values labels_encoded = LabelEncoder().fit_transform(labels) return normalized_mutual_info_score(labels_encoded, labels_pred) PKGO0k66trvae/plotting.pyimport matplotlib import numpy import pandas as pd import scanpy as sc from adjustText import adjust_text from matplotlib import pyplot from scipy import stats, sparse font = {'family': 'Arial', # 'weight' : 'bold', 'size': 14} matplotlib.rc('font', **font) matplotlib.rc('ytick', labelsize=14) matplotlib.rc('xtick', labelsize=14) def reg_mean_plot(adata, condition_key, axis_keys, labels, path_to_save="./reg_mean.pdf", gene_list=None, top_100_genes=None, show=False, legend=True, title=None, x_coeff=0.30, y_coeff=0.8, fontsize=14, **kwargs): """ Plots mean matching figure for a set of specific genes. # Parameters adata: `~anndata.AnnData` Annotated Data Matrix. condition_key: basestring Condition state to be used. axis_keys: dict dictionary of axes labels. path_to_save: basestring path to save the plot. gene_list: list list of gene names to be plotted. show: bool if `True`: will show to the plot after saving it. # Example ```python import anndata import scgen import scanpy as sc train = sc.read("./tests/data/train.h5ad", backup_url="https://goo.gl/33HtVh") network = scgen.VAEArith(x_dimension=train.shape[1], model_path="../models/test") network.train(train_data=train, n_epochs=0) unperturbed_data = train[((train.obs["cell_type"] == "CD4T") & (train.obs["condition"] == "control"))] condition = {"ctrl": "control", "stim": "stimulated"} pred, delta = network.predict(adata=train, adata_to_predict=unperturbed_data, conditions=condition) pred_adata = anndata.AnnData(pred, obs={"condition": ["pred"] * len(pred)}, var={"var_names": train.var_names}) CD4T = train[train.obs["cell_type"] == "CD4T"] all_adata = CD4T.concatenate(pred_adata) scgen.plotting.reg_mean_plot(all_adata, condition_key="condition", axis_keys={"x": "control", "y": "pred", "y1": "stimulated"}, gene_list=["ISG15", "CD3D"], path_to_save="tests/reg_mean.pdf", show=False) network.sess.close() ``` """ import seaborn as sns sns.set() sns.set(color_codes=True) if sparse.issparse(adata.X): adata.X = adata.X.A diff_genes = top_100_genes stim = adata[adata.obs[condition_key] == axis_keys["y"]] ctrl = adata[adata.obs[condition_key] == axis_keys["x"]] if diff_genes is not None: if hasattr(diff_genes, "tolist"): diff_genes = diff_genes.tolist() adata_diff = adata[:, diff_genes] stim_diff = adata_diff[adata_diff.obs[condition_key] == axis_keys["y"]] ctrl_diff = adata_diff[adata_diff.obs[condition_key] == axis_keys["x"]] x_diff = numpy.average(ctrl_diff.X, axis=0) y_diff = numpy.average(stim_diff.X, axis=0) m, b, r_value_diff, p_value_diff, std_err_diff = stats.linregress(x_diff, y_diff) print('reg_mean_top100:', r_value_diff ** 2) if "y1" in axis_keys.keys(): real_stim = adata[adata.obs[condition_key] == axis_keys["y1"]] x = numpy.average(ctrl.X, axis=0) y = numpy.average(stim.X, axis=0) m, b, r_value, p_value, std_err = stats.linregress(x, y) print('reg_mean_all:', r_value ** 2) df = pd.DataFrame({axis_keys["x"]: x, axis_keys["y"]: y}) ax = sns.regplot(x=axis_keys["x"], y=axis_keys["y"], data=df, scatter_kws={'rasterized': True}) ax.tick_params(labelsize=fontsize) if "range" in kwargs: start, stop, step = kwargs.get("range") ax.set_xticks(numpy.arange(start, stop, step)) ax.set_yticks(numpy.arange(start, stop, step)) # _p1 = pyplot.scatter(x, y, marker=".", label=f"{axis_keys['x']}-{axis_keys['y']}") # pyplot.plot(x, m * x + b, "-", color="green") ax.set_xlabel(labels["x"], fontsize=fontsize) ax.set_ylabel(labels["y"], fontsize=fontsize) # if "y1" in axis_keys.keys(): # y1 = numpy.average(real_stim.X, axis=0) # _p2 = pyplot.scatter(x, y1, marker="*", c="red", alpha=.5, label=f"{axis_keys['x']}-{axis_keys['y1']}") if gene_list is not None: texts = [] for i in gene_list: j = adata.var_names.tolist().index(i) x_bar = x[j] y_bar = y[j] texts.append(pyplot.text(x_bar, y_bar, i, fontsize=11, color="black")) pyplot.plot(x_bar, y_bar, 'o', color="red", markersize=5) # if "y1" in axis_keys.keys(): # y1_bar = y1[j] # pyplot.text(x_bar, y1_bar, i, fontsize=11, color="black") if gene_list is not None: adjust_text(texts, x=x, y=y, arrowprops=dict(arrowstyle="->", color='grey', lw=0.5), force_points=(0.0, 0.0)) if legend: pyplot.legend(loc='center left', bbox_to_anchor=(1, 0.5)) if title is None: pyplot.title(f"", fontsize=fontsize) else: pyplot.title(title, fontsize=fontsize) ax.text(max(x) - max(x) * x_coeff, max(y) - y_coeff * max(y), r'$\mathrm{R^2_{\mathrm{\mathsf{all\ genes}}}}$= ' + f"{r_value ** 2:.2f}", fontsize=kwargs.get("textsize", fontsize)) if diff_genes is not None: ax.text(max(x) - max(x) * x_coeff, max(y) - (y_coeff + 0.15) * max(y), r'$\mathrm{R^2_{\mathrm{\mathsf{top\ ' + str(len(top_100_genes)) + '\ DEGs}}}}$= ' + f"{r_value_diff ** 2:.2f}", fontsize=kwargs.get("textsize", fontsize)) pyplot.savefig(f"{path_to_save}", bbox_inches='tight', dpi=100) if show: pyplot.show() pyplot.close() def reg_var_plot(adata, condition_key, axis_keys, labels, path_to_save="./reg_var.pdf", gene_list=None, top_100_genes=None, show=False, legend=True, title=None, x_coeff=0.30, y_coeff=0.8, fontsize=14, **kwargs): """ Plots variance matching figure for a set of specific genes. # Parameters adata: `~anndata.AnnData` Annotated Data Matrix. condition_key: basestring Condition state to be used. axis_keys: dict dictionary of axes labels. path_to_save: basestring path to save the plot. gene_list: list list of gene names to be plotted. show: bool if `True`: will show to the plot after saving it. # Example ```python import anndata import scgen import scanpy as sc train = sc.read("./tests/data/train.h5ad", backup_url="https://goo.gl/33HtVh") network = scgen.VAEArith(x_dimension=train.shape[1], model_path="../models/test") network.train(train_data=train, n_epochs=0) unperturbed_data = train[((train.obs["cell_type"] == "CD4T") & (train.obs["condition"] == "control"))] condition = {"ctrl": "control", "stim": "stimulated"} pred, delta = network.predict(adata=train, adata_to_predict=unperturbed_data, conditions=condition) pred_adata = anndata.AnnData(pred, obs={"condition": ["pred"] * len(pred)}, var={"var_names": train.var_names}) CD4T = train[train.obs["cell_type"] == "CD4T"] all_adata = CD4T.concatenate(pred_adata) scgen.plotting.reg_var_plot(all_adata, condition_key="condition", axis_keys={"x": "control", "y": "pred", "y1": "stimulated"}, gene_list=["ISG15", "CD3D"], path_to_save="tests/reg_var4.pdf", show=False) network.sess.close() ``` """ import seaborn as sns sns.set() sns.set(color_codes=True) if sparse.issparse(adata.X): adata.X = adata.X.A diff_genes = top_100_genes stim = adata[adata.obs[condition_key] == axis_keys["y"]] ctrl = adata[adata.obs[condition_key] == axis_keys["x"]] if diff_genes is not None: if hasattr(diff_genes, "tolist"): diff_genes = diff_genes.tolist() adata_diff = adata[:, diff_genes] stim_diff = adata_diff[adata_diff.obs[condition_key] == axis_keys["y"]] ctrl_diff = adata_diff[adata_diff.obs[condition_key] == axis_keys["x"]] x_diff = numpy.var(ctrl_diff.X, axis=0) y_diff = numpy.var(stim_diff.X, axis=0) m, b, r_value_diff, p_value_diff, std_err_diff = stats.linregress(x_diff, y_diff) print('reg_var_top100:', r_value_diff ** 2) if "y1" in axis_keys.keys(): real_stim = adata[adata.obs[condition_key] == axis_keys["y1"]] x = numpy.var(ctrl.X, axis=0) y = numpy.var(stim.X, axis=0) m, b, r_value, p_value, std_err = stats.linregress(x, y) print('reg_var_all:', r_value ** 2) df = pd.DataFrame({axis_keys["x"]: x, axis_keys["y"]: y}) ax = sns.regplot(x=axis_keys["x"], y=axis_keys["y"], data=df, scatter_kws={'rasterized': True}) ax.tick_params(labelsize=fontsize) if "range" in kwargs: start, stop, step = kwargs.get("range") ax.set_xticks(numpy.arange(start, stop, step)) ax.set_yticks(numpy.arange(start, stop, step)) # _p1 = pyplot.scatter(x, y, marker=".", label=f"{axis_keys['x']}-{axis_keys['y']}") # pyplot.plot(x, m * x + b, "-", color="green") ax.set_xlabel(labels['x'], fontsize=fontsize) ax.set_ylabel(labels['y'], fontsize=fontsize) if "y1" in axis_keys.keys(): y1 = numpy.var(real_stim.X, axis=0) _p2 = pyplot.scatter(x, y1, marker="*", c="grey", alpha=.5, label=f"{axis_keys['x']}-{axis_keys['y1']}") if gene_list is not None: for i in gene_list: j = adata.var_names.tolist().index(i) x_bar = x[j] y_bar = y[j] pyplot.text(x_bar, y_bar, i, fontsize=11, color="black") pyplot.plot(x_bar, y_bar, 'o', color="red", markersize=5) if "y1" in axis_keys.keys(): y1_bar = y1[j] pyplot.text(x_bar, y1_bar, '*', color="black", alpha=.5) if legend: pyplot.legend(loc='center left', bbox_to_anchor=(1, 0.5)) if title is None: pyplot.title(f"", fontsize=12) else: pyplot.title(title, fontsize=12) ax.text(max(x) - max(x) * x_coeff, max(y) - y_coeff * max(y), r'$\mathrm{R^2_{\mathrm{\mathsf{all\ genes}}}}$= ' + f"{r_value ** 2:.2f}", fontsize=kwargs.get("textsize", fontsize)) if diff_genes is not None: ax.text(max(x) - max(x) * x_coeff, max(y) - (y_coeff + 0.15) * max(y), r'$\mathrm{R^2_{\mathrm{\mathsf{top\ ' + str(len(top_100_genes)) + '\ DEGs}}}}$= ' + f"{r_value_diff ** 2:.2f}", fontsize=kwargs.get("textsize", fontsize)) pyplot.savefig(f"{path_to_save}", bbox_inches='tight', dpi=100) if show: pyplot.show() pyplot.close() def binary_classifier(scg_object, adata, delta, condition_key, conditions, path_to_save, fontsize=14): """ Builds a linear classifier based on the dot product between the difference vector and the latent representation of each cell and plots the dot product results between delta and latent representation. # Parameters scg_object: `~scgen.models.VAEArith` one of scGen models object. adata: `~anndata.AnnData` Annotated Data Matrix. delta: float Difference between stimulated and control cells in latent space condition_key: basestring Condition state to be used. conditions: dict dictionary of conditions. path_to_save: basestring path to save the plot. # Example ```python import anndata import scgen import scanpy as sc train = sc.read("./tests/data/train.h5ad", backup_url="https://goo.gl/33HtVh") network = scgen.VAEArith(x_dimension=train.shape[1], model_path="../models/test") network.train(train_data=train, n_epochs=0) unperturbed_data = train[((train.obs["cell_type"] == "CD4T") & (train.obs["condition"] == "control"))] condition = {"ctrl": "control", "stim": "stimulated"} pred, delta = network.predict(adata=train, adata_to_predict=unperturbed_data, conditions=condition) scgen.plotting.binary_classifier(network, train, delta, condtion_key="condition", conditions={"ctrl": "control", "stim": "stimulated"}, path_to_save="tests/binary_classifier.pdf") network.sess.close() ``` """ # matplotlib.rcParams.update(matplotlib.rcParamsDefault) pyplot.close("all") if sparse.issparse(adata.X): adata.X = adata.X.A cd = adata[adata.obs[condition_key] == conditions["ctrl"], :] stim = adata[adata.obs[condition_key] == conditions["stim"], :] all_latent_cd = scg_object.to_latent(cd.X) all_latent_stim = scg_object.to_latent(stim.X) dot_cd = numpy.zeros((len(all_latent_cd))) dot_sal = numpy.zeros((len(all_latent_stim))) for ind, vec in enumerate(all_latent_cd): dot_cd[ind] = numpy.dot(delta, vec) for ind, vec in enumerate(all_latent_stim): dot_sal[ind] = numpy.dot(delta, vec) pyplot.hist(dot_cd, label=conditions["ctrl"], bins=50, ) pyplot.hist(dot_sal, label=conditions["stim"], bins=50) # pyplot.legend(loc=1, prop={'size': 7}) pyplot.axvline(0, color='k', linestyle='dashed', linewidth=1) pyplot.title(" ", fontsize=fontsize) pyplot.xlabel(" ", fontsize=fontsize) pyplot.ylabel(" ", fontsize=fontsize) pyplot.xticks(fontsize=fontsize) pyplot.yticks(fontsize=fontsize) ax = pyplot.gca() ax.grid(False) pyplot.savefig(f"{path_to_save}", bbox_inches='tight', dpi=100) pyplot.show() PKGO-=trvae/utils.pyfrom random import shuffle import anndata import numpy as np import scanpy as sc from scipy import sparse from sklearn.preprocessing import LabelEncoder def normalize(adata, filter_min_counts=True, size_factors=True, normalize_input=True, logtrans_input=True, n_top_genes=2000): if filter_min_counts: sc.pp.filter_genes(adata, min_counts=1) sc.pp.filter_cells(adata, min_counts=1) adata_count = adata.copy() if size_factors: sc.pp.normalize_per_cell(adata) adata.obs['size_factors'] = adata.obs.n_counts / np.median(adata.obs.n_counts) else: adata.obs['size_factors'] = 1.0 if logtrans_input: sc.pp.log1p(adata) if n_top_genes > 0 and adata.shape[1] > n_top_genes: sc.pp.highly_variable_genes(adata, n_top_genes=n_top_genes) genes = adata.var['highly_variable'] adata = adata[:, genes] adata_count = adata_count[:, genes] if normalize_input: sc.pp.scale(adata) if sparse.issparse(adata_count.X): adata_count.X = adata_count.X.A if sparse.issparse(adata.X): adata.X = adata.X.A if size_factors or normalize_input or logtrans_input: adata.raw = adata_count.copy() else: adata.raw = adata_count return adata def train_test_split(adata, train_frac=0.85): train_size = int(adata.shape[0] * train_frac) indices = np.arange(adata.shape[0]) np.random.shuffle(indices) train_idx = indices[:train_size] test_idx = indices[train_size:] train_data = adata[train_idx, :] valid_data = adata[test_idx, :] return train_data, valid_data def label_encoder(adata, label_encoder=None, condition_key='condition'): """ Encode labels of Annotated `adata` matrix using sklearn.preprocessing.LabelEncoder class. Parameters ---------- adata: `~anndata.AnnData` Annotated data matrix. Returns ------- labels: numpy nd-array Array of encoded labels Example -------- >>> import trvae >>> import scanpy as sc >>> train_data = sc.read("./data/train.h5ad") >>> train_labels, label_encoder = label_encoder(train_data) """ if label_encoder is None: le = LabelEncoder() labels = le.fit_transform(adata.obs[condition_key].tolist()) else: le = label_encoder labels = np.zeros(adata.shape[0]) for condition, label in label_encoder.items(): labels[adata.obs[condition_key] == condition] = label return labels.reshape(-1, 1), le def remove_sparsity(adata): if sparse.issparse(adata.X): adata.X = adata.X.A return adata def create_dictionary(conditions, target_conditions=[]): if isinstance(target_conditions, list): target_conditions = [target_conditions] dictionary = {} conditions = [e for e in conditions if e not in target_conditions] for idx, condition in enumerate(conditions): dictionary[condition] = idx return dictionary def shuffle_adata(adata): """ Shuffles the `adata`. # Parameters adata: `~anndata.AnnData` Annotated data matrix. labels: numpy nd-array list of encoded labels # Returns adata: `~anndata.AnnData` Shuffled annotated data matrix. labels: numpy nd-array Array of shuffled labels if `labels` is not None. # Example ```python import scgen import anndata import pandas as pd train_data = anndata.read("./data/train.h5ad") train_labels = pd.read_csv("./data/train_labels.csv", header=None) train_data, train_labels = shuffle_data(train_data, train_labels) ``` """ adata = remove_sparsity(adata) ind_list = [i for i in range(adata.shape[0])] shuffle(ind_list) new_adata = adata[ind_list, :] return new_adata PK+GO9Htrvae/models/__init__.pyfrom ._cvae import CVAE from ._trae import trAE from ._dctrvae import DCtrVAE from ._trvae_multi import trVAEMulti from ._trvae_multi_tf import trVAEMultiTF from ._trvae_atac import trVAEATAC from ._vae import VAE from ._mmdcvae import MMDCVAE PKGOq Ltrvae/models/_activations.pyimport tensorflow as tf from keras import backend as K from keras.layers import Activation, Lambda from keras.layers.advanced_activations import LeakyReLU def mean_activation(x): return tf.clip_by_value(K.exp(x), 1e-5, 1e6) def disp_activation(x): return tf.clip_by_value(tf.nn.softplus(x), 1e-4, 1e4) ACTIVATIONS = { "relu": Activation("relu", name='reconstruction_output'), 'leaky_relu': LeakyReLU(name="reconstruction_output"), 'linear': Activation("linear", name='reconstruction_output'), 'mean_activation': Activation(mean_activation, name="decoder_mean"), 'disp_activation': Activation(disp_activation, name="decoder_disp"), 'sigmoid': Activation('sigmoid', name='decoder_pi'), } PKGOÈPMPMtrvae/models/_cvae.pyimport os import logging import tensorflow from scipy import sparse from trvae.utils import label_encoder log = logging.getLogger(__file__) class CVAE: """ C-VAE vector Network class. This class contains the implementation of Conditional Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, z_dimension=100, **kwargs): tensorflow.reset_default_graph() self.x_dim = x_dimension self.z_dim = z_dimension self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 0.01) self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_to_use = kwargs.get("model_path", "../models/cvae") self.is_training = tensorflow.placeholder(tensorflow.bool, name='training_flag') self.global_step = tensorflow.Variable(0, name='global_step', trainable=False, dtype=tensorflow.int32) self.x = tensorflow.placeholder(tensorflow.float32, shape=[None, self.x_dim], name="data") self.z = tensorflow.placeholder(tensorflow.float32, shape=[None, self.z_dim], name="latent") self.y = tensorflow.placeholder(tensorflow.float32, shape=[None, 1], name="labels") self.time_step = tensorflow.placeholder(tensorflow.int32) self.size = tensorflow.placeholder(tensorflow.int32) self.init_w = tensorflow.contrib.layers.xavier_initializer() self._create_network() self._loss_function() init = tensorflow.global_variables_initializer() self.sess = tensorflow.InteractiveSession() self.saver = tensorflow.train.Saver(max_to_keep=1) self.sess.run(init) def _encoder(self): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ with tensorflow.variable_scope("encoder", reuse=tensorflow.AUTO_REUSE): xy = tensorflow.concat([self.x, self.y], axis=1) h = tensorflow.layers.dense(inputs=xy, units=700, kernel_initializer=self.init_w, use_bias=False) h = tensorflow.layers.batch_normalization(h, axis=1, training=self.is_training) h = tensorflow.nn.leaky_relu(h) h = tensorflow.layers.dense(inputs=h, units=400, kernel_initializer=self.init_w, use_bias=False) h = tensorflow.layers.batch_normalization(h, axis=1, training=self.is_training) h = tensorflow.nn.leaky_relu(h) h = tensorflow.layers.dropout(h, self.dr_rate, training=self.is_training) mean = tensorflow.layers.dense(inputs=h, units=self.z_dim, kernel_initializer=self.init_w) log_var = tensorflow.layers.dense(inputs=h, units=self.z_dim, kernel_initializer=self.init_w) return mean, log_var def _decoder(self): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ with tensorflow.variable_scope("decoder", reuse=tensorflow.AUTO_REUSE): xy = tensorflow.concat([self.z_mean, self.y], axis=1) h = tensorflow.layers.dense(inputs=xy, units=400, kernel_initializer=self.init_w, use_bias=False) h = tensorflow.layers.batch_normalization(h, axis=1, training=self.is_training) h_mmd = tensorflow.nn.leaky_relu(h) h = tensorflow.layers.dense(inputs=h_mmd, units=700, kernel_initializer=self.init_w, use_bias=False) h = tensorflow.layers.batch_normalization(h, axis=1, training=self.is_training) h = tensorflow.nn.leaky_relu(h) h = tensorflow.layers.dropout(h, self.dr_rate, training=self.is_training) h = tensorflow.layers.dense(inputs=h, units=self.x_dim, kernel_initializer=self.init_w, use_bias=True) h = tensorflow.nn.relu(h) return h, h_mmd def _sample_z(self): """ Samples from standard Normal distribution with shape [size, z_dim] and applies re-parametrization trick. It is actually sampling from latent space distributions with N(mu, var) computed in `_encoder` function. # Parameters No parameters are needed. # Returns The computed Tensor of samples with shape [size, z_dim]. """ eps = tensorflow.random_normal(shape=[self.size, self.z_dim]) return self.mu + tensorflow.exp(self.log_var / 2) * eps def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ self.mu, self.log_var = self._encoder() self.z_mean = self._sample_z() self.x_hat, self.mmd_hl = self._decoder() @staticmethod def compute_kernel(x, y): """ Computes RBF kernel between x and y. # Parameters x: Tensor Tensor with shape [batch_size, z_dim] y: Tensor Tensor with shape [batch_size, z_dim] # Returns returns the computed RBF kernel between x and y """ x_size = tensorflow.shape(x)[0] y_size = tensorflow.shape(y)[0] dim = tensorflow.shape(x)[1] tiled_x = tensorflow.tile(tensorflow.reshape(x, tensorflow.stack([x_size, 1, dim])), tensorflow.stack([1, y_size, 1])) tiled_y = tensorflow.tile(tensorflow.reshape(y, tensorflow.stack([1, y_size, dim])), tensorflow.stack([x_size, 1, 1])) return tensorflow.exp( -tensorflow.reduce_mean(tensorflow.square(tiled_x - tiled_y), axis=2) / tensorflow.cast(dim, tensorflow.float32)) def _loss_function(self): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ self.kl_loss = 0.5 * tensorflow.reduce_sum( tensorflow.exp(self.log_var) + tensorflow.square(self.mu) - 1. - self.log_var, 1) self.recon_loss = 0.5 * tensorflow.reduce_sum(tensorflow.square((self.x - self.x_hat)), 1) self.vae_loss = tensorflow.reduce_mean(self.recon_loss + self.alpha * self.kl_loss) with tensorflow.control_dependencies(tensorflow.get_collection(tensorflow.GraphKeys.UPDATE_OPS)): self.solver = tensorflow.train.AdamOptimizer(learning_rate=self.lr).minimize(self.vae_loss) def to_latent(self, data, labels): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ if sparse.issparse(data): data = data.A latent = self.sess.run(self.z_mean, feed_dict={self.x: data, self.y: labels, self.size: data.shape[0], self.is_training: False}) return latent def to_mmd_layer(self, data, labels): """ Map `data` in to the pn layer after latent layer. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ latent = self.sess.run(self.mmd_hl, feed_dict={self.x: data, self.y: labels, self.size: data.shape[0], self.is_training: False}) return latent def _reconstruct(self, data, labels, use_data=False): """ Map back the latent space encoding via the decoder. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in latent space or primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. use_data: bool this flag determines whether the `data` is already in latent space or not. if `True`: The `data` is in latent space (`data.X` is in shape [n_obs, z_dim]). if `False`: The `data` is not in latent space (`data.X` is in shape [n_obs, n_vars]). # Returns rec_data: 'numpy nd-array' returns 'numpy nd-array` containing reconstructed 'data' in shape [n_obs, n_vars]. """ if use_data: latent = data else: latent = self.to_latent(data, labels) rec_data = self.sess.run(self.x_hat, feed_dict={self.z_mean: latent, self.y: labels.reshape(-1, 1), self.is_training: False}) return rec_data def predict(self, data, labels): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ if sparse.issparse(data.X): stim_pred = self._reconstruct(data.X.A, labels) else: stim_pred = self._reconstruct(data.X, labels) return stim_pred def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.saver.restore(self.sess, self.model_to_use) def train(self, train_data, use_validation=False, valid_data=None, n_epochs=25, batch_size=32, early_stop_limit=20, threshold=0.0025, initial_run=True, shuffle=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_katrain_kang.h5ad >>> validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ if initial_run: log.info("----Training----") assign_step_zero = tensorflow.assign(self.global_step, 0) _init_step = self.sess.run(assign_step_zero) if not initial_run: self.saver.restore(self.sess, self.model_to_use) train_labels, le = label_encoder(train_data) if use_validation and valid_data is None: raise Exception("valid_data is None but use_validation is True.") if use_validation: valid_labels, _ = label_encoder(valid_data) loss_hist = [] patience = early_stop_limit min_delta = threshold patience_cnt = 0 for it in range(n_epochs): increment_global_step_op = tensorflow.assign(self.global_step, self.global_step + 1) _step = self.sess.run(increment_global_step_op) current_step = self.sess.run(self.global_step) train_loss = 0 for lower in range(0, train_data.shape[0], batch_size): upper = min(lower + batch_size, train_data.shape[0]) if sparse.issparse(train_data.X): x_mb = train_data[lower:upper, :].X.A else: x_mb = train_data[lower:upper, :].X y_mb = train_labels[lower:upper] _, current_loss_train = self.sess.run([self.solver, self.vae_loss], feed_dict={self.x: x_mb, self.y: y_mb, self.time_step: current_step, self.size: len(x_mb), self.is_training: True}) train_loss += current_loss_train print(f"iteration {it}: {train_loss // (train_data.shape[0] // batch_size)}") if use_validation: valid_loss = 0 for lower in range(0, valid_data.shape[0], batch_size): upper = min(lower + batch_size, valid_data.shape[0]) if sparse.issparse(valid_data.X): x_mb = valid_data[lower:upper, :].X.A else: x_mb = valid_data[lower:upper, :].X y_mb = valid_labels[lower:upper] current_loss_valid = self.sess.run(self.vae_loss, feed_dict={self.x: x_mb, self.y: y_mb, self.time_step: current_step, self.size: len(x_mb), self.is_training: False}) valid_loss += current_loss_valid loss_hist.append(valid_loss / valid_data.shape[0]) if it > 0 and loss_hist[it - 1] - loss_hist[it] > min_delta: patience_cnt = 0 else: patience_cnt += 1 if patience_cnt > patience: os.makedirs(self.model_to_use, exist_ok=True) save_path = self.saver.save(self.sess, self.model_to_use) break else: os.makedirs(self.model_to_use, exist_ok=True) save_path = self.saver.save(self.sess, self.model_to_use) log.info(f"Model saved in file: {save_path}. Training finished")PKGO@ 6AmAmtrvae/models/_dctrvae.pyimport logging import os import anndata import keras import numpy as np from keras.callbacks import CSVLogger, History, EarlyStopping, ReduceLROnPlateau, LambdaCallback from keras.layers import Activation from keras.layers import Dense, BatchNormalization, Dropout, Input, concatenate, Lambda, Conv2D, \ Flatten, Reshape, Conv2DTranspose, UpSampling2D, MaxPooling2D from keras.layers.advanced_activations import LeakyReLU from keras.models import Model, load_model from keras.utils import multi_gpu_model, to_categorical from trvae.models._utils import sample_z, print_message from trvae.utils import label_encoder, remove_sparsity from ._losses import LOSSES log = logging.getLogger(__file__) class DCtrVAE: """ Regularized Convolutional C-VAE vector Network class. This class contains the implementation of Conditional Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, z_dimension=100, **kwargs): self.x_dim = x_dimension if isinstance(x_dimension, tuple) else (x_dimension,) self.z_dim = z_dimension self.image_shape = x_dimension self.n_conditions = kwargs.get("n_conditions", 2) self.mmd_dim = kwargs.get("mmd_dimension", 128) self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 0.001) self.beta = kwargs.get("beta", 100) self.eta = kwargs.get("eta", 1.0) self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_path = kwargs.get("model_path", "./") self.kernel_method = kwargs.get("kernel", "multi-scale-rbf") self.arch_style = kwargs.get("arch_style", 1) self.n_gpus = kwargs.get("gpus", 1) self.x = Input(shape=self.x_dim, name="data") self.encoder_labels = Input(shape=(self.n_conditions,), name="encoder_labels") self.decoder_labels = Input(shape=(self.n_conditions,), name="decoder_labels") self.z = Input(shape=(self.z_dim,), name="latent_data") self.init_w = keras.initializers.glorot_normal() self._create_network() self._loss_function(compile_gpu_model=True) self.cvae_model.summary() def _encoder(self, name="encoder"): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ if self.arch_style == 1: # Baseline CNN h = Dense(128, activation='relu')(self.encoder_labels) h = Dense(np.prod(self.x_dim[:-1]), activation='relu')(h) h = Reshape((*self.x_dim[:-1], 1))(h) h = concatenate([self.x, h]) h = Conv2D(64, kernel_size=(4, 4), strides=2, padding='same', activation=None)(h) h = LeakyReLU()(h) h = MaxPooling2D()(h) h = Conv2D(64, kernel_size=(4, 4), strides=2, padding='same', activation=None)(h) h = LeakyReLU()(h) h = MaxPooling2D()(h) h = Flatten()(h) h = Dense(1024, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) mean = Dense(self.z_dim, kernel_initializer=self.init_w)(h) log_var = Dense(self.z_dim, kernel_initializer=self.init_w)(h) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=[self.x, self.encoder_labels], outputs=[mean, log_var, z], name=name) model.summary() return mean, log_var, model elif self.arch_style == 2: # FCN x_reshaped = Reshape(target_shape=(np.prod(self.x_dim),))(self.x) xy = concatenate([x_reshaped, self.encoder_labels], axis=1) h = Dense(512, kernel_initializer=self.init_w, use_bias=False)(xy) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(512, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) mean = Dense(self.z_dim, kernel_initializer=self.init_w)(h) log_var = Dense(self.z_dim, kernel_initializer=self.init_w)(h) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=[self.x, self.encoder_labels], outputs=[mean, log_var, z], name=name) model.summary() return mean, log_var, model else: h = Dense(128, activation='relu')(self.encoder_labels) h = Dense(np.prod(self.x_dim[:-1]), activation='relu')(h) h = Reshape((*self.x_dim[:-1], 1))(h) h = concatenate([self.x, h]) conv1 = Conv2D(64, 3, activation='relu', padding='same', name='conv1_1')(h) conv1 = Conv2D(64, 3, activation='relu', padding='same', name='conv1_2')(conv1) pool1 = MaxPooling2D(pool_size=(2, 2))(conv1) conv2 = Conv2D(128, 3, activation='relu', padding='same', name='conv2_1')(pool1) conv2 = Conv2D(128, 3, activation='relu', padding='same', name='conv2_2')(conv2) pool2 = MaxPooling2D(pool_size=(2, 2))(conv2) conv3 = Conv2D(256, 3, activation='relu', padding='same', name='conv3_1')(pool2) conv3 = Conv2D(256, 3, activation='relu', padding='same', name='conv3_2')(conv3) conv3 = Conv2D(256, 3, activation='relu', padding='same', name='conv3_3')(conv3) pool3 = MaxPooling2D(pool_size=(2, 2))(conv3) conv4 = Conv2D(512, 3, activation='relu', padding='same', name='conv4_1')(pool3) conv4 = Conv2D(512, 3, activation='relu', padding='same', name='conv4_2')(conv4) conv4 = Conv2D(512, 3, activation='relu', padding='same', name='conv4_3')(conv4) pool4 = MaxPooling2D(pool_size=(2, 2))(conv4) conv5 = Conv2D(512, 3, activation='relu', padding='same', name='conv5_1')(pool4) conv5 = Conv2D(512, 3, activation='relu', padding='same', name='conv5_2')(conv5) conv5 = Conv2D(512, 3, activation='relu', padding='same', name='conv5_3')(conv5) pool5 = MaxPooling2D(pool_size=(2, 2))(conv5) flat = Flatten(name='flatten')(pool5) dense = Dense(1024, activation='linear', name='fc1')(flat) dense = Activation('relu')(dense) dense = Dense(256, activation='linear', name='fc2')(dense) dense = Activation('relu')(dense) self.enc_dense = Dropout(self.dr_rate)(dense) mean = Dense(self.z_dim, kernel_initializer=self.init_w)(self.enc_dense) log_var = Dense(self.z_dim, kernel_initializer=self.init_w)(self.enc_dense) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=[self.x, self.encoder_labels], outputs=[mean, log_var, z], name=name) model.summary() return mean, log_var, model def _mmd_decoder(self, name="decoder"): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ if self.arch_style == 1: # Baseline CNN for MNIST zy = concatenate([self.z, self.decoder_labels], axis=1) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, use_bias=False)(zy) h = BatchNormalization(axis=1)(h) h_mmd = LeakyReLU(name="mmd")(h) h = Dropout(self.dr_rate)(h_mmd) h = Dense(1024, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(np.prod(self.x_dim), kernel_initializer=self.init_w, use_bias=False)(h) h = LeakyReLU()(h) h = Reshape(target_shape=self.x_dim)(h) h = Conv2DTranspose(128, kernel_size=(4, 4), padding='same')(h) h = LeakyReLU()(h) h = Conv2DTranspose(64, kernel_size=(4, 4), padding='same')(h) h = LeakyReLU()(h) h = Conv2DTranspose(self.x_dim[-1], kernel_size=(4, 4), padding='same', activation="relu")(h) decoder_model = Model(inputs=[self.z, self.decoder_labels], outputs=h, name=name) decoder_mmd_model = Model(inputs=[self.z, self.decoder_labels], outputs=h_mmd, name='deocder_mmd') return decoder_model, decoder_mmd_model elif self.arch_style == 2: # FCN zy = concatenate([self.z, self.decoder_labels], axis=1) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, use_bias=False)(zy) h = BatchNormalization(axis=1)(h) h_mmd = LeakyReLU(name="mmd")(h) h = Dropout(self.dr_rate)(h) h = Dense(512, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(512, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(np.prod(self.x_dim), kernel_initializer=self.init_w, use_bias=True)(h) h = Activation('relu', name="reconstruction_output")(h) h = Reshape(target_shape=self.x_dim)(h) decoder_model = Model(inputs=[self.z, self.decoder_labels], outputs=h, name=name) decoder_mmd_model = Model(inputs=[self.z, self.decoder_labels], outputs=h_mmd, name='deocder_mmd') return decoder_model, decoder_mmd_model else: encode_y = Dense(64, activation='relu')(self.decoder_labels) zy = concatenate([self.z, encode_y], axis=1) zy = Activation('relu')(zy) h = Dense(self.mmd_dim, activation="linear", kernel_initializer='he_normal')(zy) h_mmd = Activation('relu', name="mmd")(h) h = Dense(1024, kernel_initializer='he_normal')(h_mmd) h = Activation('relu')(h) h = Dense(256 * 4 * 4, kernel_initializer='he_normal')(h) h = Activation('relu')(h) width = self.x_dim[0] // 16 height = self.x_dim[1] // 16 n_channels = 4096 // (width * height) h = Reshape(target_shape=(width, height, n_channels))(h) up6 = Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')( UpSampling2D(size=(2, 2))(h)) conv6 = Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up6) up7 = Conv2D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal')( UpSampling2D(size=(2, 2))(conv6)) conv7 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up7) up8 = Conv2D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal')( UpSampling2D(size=(2, 2))(conv7)) conv8 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up8) up9 = Conv2D(32, 2, activation='relu', padding='same', kernel_initializer='he_normal')( UpSampling2D(size=(2, 2))(conv8)) conv9 = Conv2D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(up9) conv10 = Conv2D(self.x_dim[-1], 1, activation='relu')(conv9) decoder_model = Model(inputs=[self.z, self.decoder_labels], outputs=conv10, name=name) decoder_mmd_model = Model(inputs=[self.z, self.decoder_labels], outputs=h_mmd, name='deocder_mmd') return decoder_model, decoder_mmd_model def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ inputs = [self.x, self.encoder_labels, self.decoder_labels] self.mu, self.log_var, self.encoder_model = self._encoder(name="encoder") self.decoder_model, self.decoder_mmd_model = self._mmd_decoder(name="decoder") decoder_output = self.decoder_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) mmd_output = self.decoder_mmd_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) reconstruction_output = Lambda(lambda x: x, name="kl_reconstruction")(decoder_output) mmd_output = Lambda(lambda x: x, name="mmd")(mmd_output) self.cvae_model = Model(inputs=inputs, outputs=[reconstruction_output, mmd_output], name="cvae") if self.n_gpus > 1: self.gpu_cvae_model = multi_gpu_model(self.cvae_model, gpus=self.n_gpus) self.gpu_encoder_model = multi_gpu_model(self.encoder_model, gpus=self.n_gpus) self.gpu_decoder_model = multi_gpu_model(self.decoder_model, gpus=self.n_gpus) else: self.gpu_cvae_model = self.cvae_model self.gpu_encoder_model = self.encoder_model self.gpu_decoder_model = self.decoder_model def _calculate_loss(self): loss = LOSSES['mse'](self.mu, self.log_var, self.alpha, self.eta) mmd_loss = LOSSES['mmd'](self.n_conditions, self.beta, self.kernel_method, "general") return loss, mmd_loss def _loss_function(self, compile_gpu_model=True): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ mse_loss, mmd_loss = self._calculate_loss() self.cvae_optimizer = keras.optimizers.Adam(lr=self.lr) if compile_gpu_model: self.gpu_cvae_model.compile(optimizer=self.cvae_optimizer, loss=[mse_loss, mmd_loss], metrics={self.cvae_model.outputs[0].name: mse_loss, self.cvae_model.outputs[1].name: mmd_loss}) else: self.cvae_model.compile(optimizer=self.cvae_optimizer, loss=[mse_loss, mmd_loss], metrics={self.cvae_model.outputs[0].name: mse_loss, self.cvae_model.outputs[1].name: mmd_loss}) def to_latent(self, adata, encoder_labels): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ adata = remove_sparsity(adata) images = np.reshape(adata.X, (-1, *self.x_dim)) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) latent = self.encoder_model.predict([images, encoder_labels])[2] latent_adata = anndata.AnnData(X=latent) latent_adata.obs = adata.obs.copy(deep=True) return latent_adata def to_mmd_layer(self, adata, encoder_labels, feed_fake=0): """ Map `data` in to the pn layer after latent layer. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ if feed_fake > 0: decoder_labels = np.zeros(shape=encoder_labels.shape) + feed_fake else: decoder_labels = encoder_labels adata = remove_sparsity(adata) images = np.reshape(adata.X, (-1, *self.x_dim)) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) mmd_latent = self.cvae_model.predict([images, encoder_labels, decoder_labels])[1] mmd_adata = anndata.AnnData(X=mmd_latent) mmd_adata.obs = adata.obs.copy(deep=True) return mmd_adata def predict(self, adata, encoder_labels, decoder_labels): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ adata = remove_sparsity(adata) images = np.reshape(adata.X, (-1, *self.x_dim)) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) reconstructed = self.cvae_model.predict([images, encoder_labels, decoder_labels])[0] reconstructed = np.reshape(reconstructed, (-1, np.prod(self.x_dim))) reconstructed_adata = anndata.AnnData(X=reconstructed) reconstructed_adata.obs = adata.obs.copy(deep=True) reconstructed_adata.var_names = adata.var_names return reconstructed_adata def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.cvae_model = load_model(os.path.join(self.model_path, 'mmd_cvae.h5'), compile=False) self.encoder_model = load_model(os.path.join(self.model_path, 'encoder.h5'), compile=False) self.decoder_model = load_model(os.path.join(self.model_path, 'decoder.h5'), compile=False) self.decoder_mmd_model = load_model(os.path.join(self.model_path, 'decoder_mmd.h5'), compile=False) self._loss_function(compile_gpu_model=False) def save_model(self): os.makedirs(self.model_path, exist_ok=True) self.cvae_model.save(os.path.join(self.model_path, "mmd_cvae.h5"), overwrite=True) self.encoder_model.save(os.path.join(self.model_path, "encoder.h5"), overwrite=True) self.decoder_model.save(os.path.join(self.model_path, "decoder.h5"), overwrite=True) self.decoder_model.save(os.path.join(self.model_path, "decoder_mmd.h5"), overwrite=True) log.info(f"Model saved in file: {self.model_path}. Training finished") def train(self, train_adata, valid_adata=None, condition_encoder=None, condition_key='condition', n_epochs=25, batch_size=32, early_stop_limit=20, lr_reducer=10, threshold=0.0025, monitor='val_loss', shuffle=True, verbose=2, save=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_katrain_kang.h5ad >>> validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ train_adata = remove_sparsity(train_adata) train_labels_encoded, self.condition_encoder = label_encoder(train_adata, condition_encoder, condition_key) train_labels_onehot = to_categorical(train_labels_encoded, num_classes=self.n_conditions) callbacks = [ History(), CSVLogger(filename="./csv_logger.log"), ] if early_stop_limit > 0: callbacks.append(EarlyStopping(patience=early_stop_limit, monitor=monitor, min_delta=threshold)) if lr_reducer > 0: callbacks.append(ReduceLROnPlateau(monitor=monitor, patience=lr_reducer, verbose=verbose)) if verbose > 2: callbacks.append( LambdaCallback(on_epoch_end=lambda epoch, logs: print_message(epoch, logs, n_epochs, verbose))) fit_verbose = 0 else: fit_verbose = verbose train_images = np.reshape(train_adata.X, (-1, *self.x_dim)) x = [train_images, train_labels_onehot, train_labels_onehot] y = [train_images, train_labels_encoded] if valid_adata is not None: valid_adata = remove_sparsity(valid_adata) valid_labels_encoded, _ = label_encoder(valid_adata, condition_encoder, condition_key) valid_labels_onehot = to_categorical(valid_labels_encoded, num_classes=self.n_conditions) valid_images = np.reshape(valid_adata.X, (-1, *self.x_dim)) x_valid = [valid_images, valid_labels_onehot, valid_labels_onehot] y_valid = [valid_images, valid_labels_encoded] self.cvae_model.fit(x=x, y=y, epochs=n_epochs, batch_size=batch_size, validation_data=(x_valid, y_valid), shuffle=shuffle, callbacks=callbacks, verbose=fit_verbose) else: self.cvae_model.fit(x=x, y=y, epochs=n_epochs, batch_size=batch_size, shuffle=shuffle, callbacks=callbacks, verbose=fit_verbose) if save: self.save_model() PKGO#..trvae/models/_losses.pyimport tensorflow as tf from keras import backend as K, Model from keras.applications.imagenet_utils import preprocess_input from keras_vggface import VGGFace from ._utils import compute_mmd, _nelem, _nan2zero, _nan2inf, _reduce_mean def kl_recon(mu, log_var, alpha=0.1, eta=1.0): def kl_recon_loss(y_true, y_pred): kl_loss = 0.5 * K.mean(K.exp(log_var) + K.square(mu) - 1. - log_var, 1) y_true_min, y_true_max = K.min(y_true), K.max(y_true) recon_loss = K.switch(K.equal(y_true_min, y_true_max), then_expression=lambda: 0.5 * K.sum(K.zeros_like(y_true), axis=1), else_expression=lambda: 0.5 * K.sum(K.square((y_true - y_pred)), axis=1) ) return _nan2inf(eta * recon_loss + alpha * kl_loss) return kl_recon_loss def kl_loss(mu, log_var, alpha=0.1): def kl_recon_loss(y_true, y_pred): kl_loss = 0.5 * K.mean(K.exp(log_var) + K.square(mu) - 1. - log_var, 1) return _nan2inf(alpha * kl_loss) return kl_recon_loss def mmd(n_conditions, beta, kernel_method='multi-scale-rbf', computation_way="general"): def mmd_loss(real_labels, y_pred): with tf.variable_scope("mmd_loss", reuse=tf.AUTO_REUSE): real_labels = K.reshape(K.cast(real_labels, 'int32'), (-1,)) conditions_mmd = tf.dynamic_partition(y_pred, real_labels, num_partitions=n_conditions) loss = 0.0 if computation_way.isdigit(): boundary = int(computation_way) for i in range(boundary): for j in range(boundary, n_conditions): loss += compute_mmd(conditions_mmd[i], conditions_mmd[j], kernel_method) else: for i in range(len(conditions_mmd)): for j in range(i): loss += compute_mmd(conditions_mmd[j], conditions_mmd[j + 1], kernel_method) return _nan2zero(beta * loss) return mmd_loss def categrical_crossentropy(gamma): def cce_loss(real_classes, pred_classes): return gamma * K.categorical_crossentropy(real_classes, pred_classes) return cce_loss def perceptual_loss(x_dim, gamma=1.0): def percept_loss(input_image, reconstructed_image): vggface = VGGFace(include_top=False, input_shape=x_dim, model='vgg16') vgg_layers = ['conv1_1'] outputs = [vggface.get_layer(l).output for l in vgg_layers] model = Model(inputs=vggface.input, outputs=outputs) for layer in model.layers: layer.trainable = False input_image *= 255.0 reconstructed_image *= 255.0 input_image = preprocess_input(input_image, mode='tf', data_format='channels_last') reconstructed_image = preprocess_input(reconstructed_image, mode='tf', data_format='channels_last') h1_list = model(input_image) h2_list = model(reconstructed_image) if not isinstance(h1_list, list): h1_list = [h1_list] h2_list = [h2_list] p_loss = 0.0 for h1, h2 in zip(h1_list, h2_list): h1 = K.batch_flatten(h1) h2 = K.batch_flatten(h2) p_loss += K.mean(K.square(h1 - h2), axis=-1) return gamma * p_loss return percept_loss class NB(object): def __init__(self, theta=None, masking=False, scope='nbinom_loss/', scale_factor=1.0): # for numerical stability self.eps = 1e-10 self.scale_factor = scale_factor self.scope = scope self.masking = masking self.theta = theta def loss(self, y_true, y_pred, mean=True): scale_factor = self.scale_factor eps = self.eps with tf.name_scope(self.scope): y_true = tf.cast(y_true, tf.float32) y_pred = tf.cast(y_pred, tf.float32) * scale_factor if self.masking: nelem = _nelem(y_true) y_true = _nan2zero(y_true) # Clip theta theta = tf.minimum(self.theta, 1e6) t1 = tf.lgamma(theta + eps) + tf.lgamma(y_true + 1.0) - tf.lgamma(y_true + theta + eps) t2 = (theta + y_true) * tf.log(1.0 + (y_pred / (theta + eps))) + ( y_true * (tf.log(theta + eps) - tf.log(y_pred + eps))) final = t1 + t2 final = _nan2inf(final) if mean: if self.masking: final = tf.divide(tf.reduce_sum(final), nelem) else: final = tf.reduce_mean(final) return final class ZINB(NB): def __init__(self, pi, ridge_lambda=0.0, scope='zinb_loss/', **kwargs): super().__init__(scope=scope, **kwargs) self.pi = pi self.ridge_lambda = ridge_lambda def loss(self, y_true, y_pred, mean=True): scale_factor = self.scale_factor eps = self.eps with tf.name_scope(self.scope): # reuse existing NB neg.log.lik. # mean is always False here, because everything is calculated # element-wise. we take the mean only in the end nb_case = super().loss(y_true, y_pred, mean=False) - tf.log(1.0 - self.pi + eps) y_true = tf.cast(y_true, tf.float32) y_pred = tf.cast(y_pred, tf.float32) * scale_factor theta = tf.minimum(self.theta, 1e6) zero_nb = tf.pow(theta / (theta + y_pred + eps), theta) zero_case = -tf.log(self.pi + ((1.0 - self.pi) * zero_nb) + eps) result = tf.where(tf.less(y_true, 1e-8), zero_case, nb_case) ridge = self.ridge_lambda * tf.square(self.pi) result += ridge if mean: if self.masking: result = _reduce_mean(result) else: result = tf.reduce_mean(result) result = _nan2inf(result) return result def nb_loss(disp, mu, log_var, scale_factor=1.0, alpha=0.1): kl = kl_loss(mu, log_var, alpha=alpha) def nb(y_true, y_pred): nb_obj = NB(theta=disp, masking=True, scale_factor=scale_factor) return nb_obj.loss(y_true, y_pred) + kl(y_true, y_pred) return nb def zinb_loss(pi, disp, mu, log_var, ridge=0.1, alpha=0.1): kl = kl_loss(mu, log_var, alpha=alpha) def zinb(y_true, y_pred): zinb_obj = ZINB(pi, theta=disp, ridge_lambda=ridge) return zinb_obj.loss(y_true, y_pred) + kl(y_true, y_pred) return zinb LOSSES = { "mse": kl_recon, "mmd": mmd, "cce": categrical_crossentropy, "nb": nb_loss, "zinb": zinb_loss, "perceptual": perceptual_loss, } PKGOض0J0Jtrvae/models/_mmdcvae.pyimport logging import os import anndata import keras import numpy as np from keras.callbacks import CSVLogger, History, EarlyStopping, ReduceLROnPlateau, LambdaCallback from keras.layers import Dense, BatchNormalization, Dropout, Input, concatenate, Lambda, Activation from keras.layers.advanced_activations import LeakyReLU from keras.models import Model, load_model from keras.utils import to_categorical from scipy import sparse from trvae.models._losses import LOSSES from trvae.models._activations import ACTIVATIONS from trvae.models._utils import sample_z, print_message from trvae.utils import label_encoder, remove_sparsity log = logging.getLogger(__file__) class MMDCVAE: """ Regularized C-VAE vector Network class. This class contains the implementation of Conditional Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, z_dimension=100, n_conditions=3, **kwargs): self.x_dim = x_dimension self.z_dim = z_dimension self.n_conditions = n_conditions self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 0.001) self.beta = kwargs.get("beta", 100) self.eta = kwargs.get("eta", 1.0) self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_to_use = kwargs.get("model_path", "./") self.kernel_method = kwargs.get("kernel", "multi-scale-rbf") self.output_activation = kwargs.get("output_activation", 'relu') self.mmd_computation_way = kwargs.get("mmd_computation_way", "general") self.ridge = kwargs.get('ridge', 0.1) self.scale_factor = kwargs.get('scale_factor', 1.0) self.clip_value = kwargs.get('clip_value', 3.0) self.lambda_l1 = kwargs.get('lambda_l1', 0.01) self.lambda_l2 = kwargs.get('lambda_l2', 0.1) self.x = Input(shape=(self.x_dim,), name="data") self.encoder_labels = Input(shape=(self.n_conditions,), name="encoder_labels") self.decoder_labels = Input(shape=(self.n_conditions,), name="decoder_labels") self.z = Input(shape=(self.z_dim,), name="latent_data") self.init_w = keras.initializers.glorot_normal() self.regularizer = keras.regularizers.l1_l2(self.lambda_l1, self.lambda_l2) self._create_network() self._loss_function() self.encoder_model.summary() self.decoder_model.summary() self.cvae_model.summary() def _encoder(self, name="encoder"): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ xy = concatenate([self.x, self.encoder_labels], axis=1) h = Dense(800, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(xy) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(800, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(h) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(128, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(h) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) mean = Dense(self.z_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer)(h) log_var = Dense(self.z_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer)(h) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=[self.x, self.encoder_labels], outputs=[mean, log_var, z], name=name) return mean, log_var, model def _mmd_decoder(self, name="decoder"): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ zy = concatenate([self.z, self.decoder_labels], axis=1) h = Dense(128, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(zy) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU(name="mmd")(h) h = Dropout(self.dr_rate)(h) h = Dense(800, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(h) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(800, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(h) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.x_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=True)(h) h = ACTIVATIONS[self.output_activation](h) decoder_model = Model(inputs=[self.z, self.decoder_labels], outputs=h, name=name) decoder_mmd_model = Model(inputs=[self.z, self.decoder_labels], outputs=self.z, name='decoder_mmd') return decoder_model, decoder_mmd_model def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ inputs = [self.x, self.encoder_labels, self.decoder_labels] self.mu, self.log_var, self.encoder_model = self._encoder(name="encoder") self.decoder_model, self.decoder_mmd_model = self._mmd_decoder(name="decoder") decoder_output = self.decoder_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) mmd_output = self.decoder_mmd_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) reconstruction_output = Lambda(lambda x: x, name="kl_mse")(decoder_output) mmd_output = Lambda(lambda x: x, name="mmd")(mmd_output) self.cvae_model = Model(inputs=inputs, outputs=[reconstruction_output, mmd_output], name="cvae") def _calculate_loss(self): loss = LOSSES['mse'](self.mu, self.log_var, self.alpha, self.eta) mmd_loss = LOSSES['mmd'](self.n_conditions, self.beta, self.kernel_method, self.mmd_computation_way) return loss, mmd_loss def _loss_function(self): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ loss, mmd_loss = self._calculate_loss() self.cvae_optimizer = keras.optimizers.Adam(lr=self.lr, clipvalue=self.clip_value) self.cvae_model.compile(optimizer=self.cvae_optimizer, loss=[loss, mmd_loss], metrics={self.cvae_model.outputs[0].name: loss, self.cvae_model.outputs[1].name: mmd_loss}) def to_latent(self, adata, encoder_labels): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ adata = remove_sparsity(adata) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) latent = self.encoder_model.predict([adata.X, encoder_labels])[2] latent = np.nan_to_num(latent) latent_adata = anndata.AnnData(X=latent) latent_adata.obs = adata.obs.copy(deep=True) return latent_adata def to_mmd_layer(self, adata, encoder_labels): """ Map `data` in to the pn layer after latent layer. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ return self.to_latent(adata, encoder_labels) def predict(self, adata, encoder_labels, decoder_labels): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ adata = remove_sparsity(adata) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) reconstructed = self.cvae_model.predict([adata.X, encoder_labels, decoder_labels])[0] reconstructed = np.nan_to_num(reconstructed) reconstructed_adata = anndata.AnnData(X=reconstructed) reconstructed_adata.obs = adata.obs.copy(deep=True) reconstructed_adata.var_names = adata.var_names return reconstructed_adata def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.cvae_model = load_model(os.path.join(self.model_to_use, 'mmd_cvae.h5'), compile=False) self.encoder_model = load_model(os.path.join(self.model_to_use, 'encoder.h5'), compile=False) self.decoder_model = load_model(os.path.join(self.model_to_use, 'decoder.h5'), compile=False) self.decoder_mmd_model = load_model(os.path.join(self.model_to_use, 'decoder_mmd.h5'), compile=False) self._loss_function() def save_model(self): os.makedirs(self.model_to_use, exist_ok=True) self.cvae_model.save(os.path.join(self.model_to_use, "mmd_cvae.h5"), overwrite=True) self.encoder_model.save(os.path.join(self.model_to_use, "encoder.h5"), overwrite=True) self.decoder_model.save(os.path.join(self.model_to_use, "decoder.h5"), overwrite=True) self.decoder_mmd_model.save(os.path.join(self.model_to_use, "decoder_mmd.h5"), overwrite=True) def train(self, train_adata, valid_adata=None, condition_encoder=None, condition_key='condition', n_epochs=25, batch_size=32, early_stop_limit=20, lr_reducer=10, threshold=0.0025, monitor='val_loss', shuffle=True, verbose=2, save=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_katrain_kang.h5ad >>> validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ train_labels_encoded, _ = label_encoder(train_adata, condition_encoder, condition_key) train_labels_onehot = to_categorical(train_labels_encoded, num_classes=self.n_conditions) callbacks = [ History(), CSVLogger(filename="./csv_logger.log"), ] if early_stop_limit > 0: callbacks.append(EarlyStopping(patience=early_stop_limit, monitor=monitor, min_delta=threshold)) if lr_reducer > 0: callbacks.append(ReduceLROnPlateau(monitor=monitor, patience=lr_reducer, verbose=verbose)) if verbose > 2: callbacks.append( LambdaCallback(on_epoch_end=lambda epoch, logs: print_message(epoch, logs, n_epochs, verbose))) fit_verbose = 0 else: fit_verbose = verbose if sparse.issparse(train_adata.X): train_adata.X = train_adata.X.A x = [train_adata.X, train_labels_onehot, train_labels_onehot] y = [train_adata.X, train_labels_encoded] if valid_adata is not None: if sparse.issparse(valid_adata.X): valid_adata.X = valid_adata.X.A valid_labels_encoded, _ = label_encoder(valid_adata, condition_encoder, condition_key) valid_labels_onehot = to_categorical(valid_labels_encoded, num_classes=self.n_conditions) x_valid = [valid_adata.X, valid_labels_onehot, valid_labels_onehot] y_valid = [valid_adata.X, valid_labels_encoded] self.cvae_model.fit(x=x, y=y, epochs=n_epochs, batch_size=batch_size, validation_data=(x_valid, y_valid), shuffle=shuffle, callbacks=callbacks, verbose=fit_verbose) else: self.cvae_model.fit(x=x, y=y, epochs=n_epochs, batch_size=batch_size, shuffle=shuffle, callbacks=callbacks, verbose=fit_verbose) if save: self.save_model() PKGOf q=q=trvae/models/_trae.pyimport logging import os import keras import tensorflow as tf from keras import backend as K from keras.callbacks import CSVLogger, History, EarlyStopping from keras.layers import Dense, BatchNormalization, Dropout, Input, Lambda, Activation from keras.layers.advanced_activations import LeakyReLU from keras.models import Model, load_model from scipy import sparse from trvae.models._utils import compute_mmd from trvae.utils import label_encoder log = logging.getLogger(__file__) class trAE: """ Regularized C-VAE vector Network class. This class contains the implementation of Conditional Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, z_dimension=100, **kwargs): self.x_dim = x_dimension self.z_dim = z_dimension self.lr = kwargs.get("learning_rate", 0.001) self.beta = kwargs.get("beta", 100) self.conditions = kwargs.get("condition_list") self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_to_use = kwargs.get("model_path", "./") self.kernel_method = kwargs.get("kernel", "multi-scale-rbf") self.x = Input(shape=(self.x_dim,), name="data") self.z = Input(shape=(self.z_dim,), name="latent_data") self.init_w = keras.initializers.glorot_normal() self._create_network() self._loss_function() self.encoder_model.summary() self.decoder_model.summary() self.rae_model.summary() def _encoder(self, x, name="encoder"): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ h = Dense(700, kernel_initializer=self.init_w, use_bias=False)(x) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(400, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) z = Dense(self.z_dim, kernel_initializer=self.init_w, name='mmd')(h) model = Model(inputs=x, outputs=z, name=name) return model def _decoder(self, z, name="decoder"): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ h = Dense(400, kernel_initializer=self.init_w, use_bias=False)(z) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dense(700, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.x_dim, kernel_initializer=self.init_w, use_bias=True)(h) h = Activation('relu', name="reconstruction_output")(h) model = Model(inputs=z, outputs=h, name=name) return model def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ self.encoder_model = self._encoder(self.x, name="encoder") self.decoder_model = self._decoder(self.z, name="decoder") decoder_outputs = self.decoder_model(self.encoder_model(self.x)) encoder_outputs = self.encoder_model(self.x) reconstruction_output = Lambda(lambda x: x, name="kl_reconstruction")(decoder_outputs) mmd_output = Lambda(lambda x: x, name="mmd")(encoder_outputs) self.rae_model = Model(inputs=self.x, outputs=[reconstruction_output, mmd_output], name="rae") def _loss_function(self): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ def batch_loss(): def reconstruction_loss(y_true, y_pred): recon_loss = K.sum(K.square((y_true - y_pred)), axis=1) return recon_loss def mmd_loss(real_labels, y_pred): with tf.variable_scope("mmd_loss", reuse=tf.AUTO_REUSE): real_labels = K.reshape(K.cast(real_labels, 'int32'), (-1,)) source_mmd, dest_mmd = tf.dynamic_partition(y_pred, real_labels, num_partitions=2) loss = compute_mmd(source_mmd, dest_mmd, self.kernel_method) return self.beta * loss self.cvae_optimizer = keras.optimizers.Adam(lr=self.lr) self.rae_model.compile(optimizer=self.cvae_optimizer, loss=[reconstruction_loss, mmd_loss], metrics={self.rae_model.outputs[0].name: reconstruction_loss, self.rae_model.outputs[1].name: mmd_loss}) batch_loss() def to_latent(self, data): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ latent = self.encoder_model.predict(data) return latent def _reconstruct(self, data, use_data=False): """ Map back the latent space encoding via the decoder. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in latent space or primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. use_data: bool this flag determines whether the `data` is already in latent space or not. if `True`: The `data` is in latent space (`data.X` is in shape [n_obs, z_dim]). if `False`: The `data` is not in latent space (`data.X` is in shape [n_obs, n_vars]). # Returns rec_data: 'numpy nd-array' returns 'numpy nd-array` containing reconstructed 'data' in shape [n_obs, n_vars]. """ if use_data: latent = data else: latent = self.to_latent(data) rec_data = self.decoder_model.predict(latent) return rec_data def predict(self, data, data_space='None'): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ if sparse.issparse(data.X): if data_space == 'latent': stim_pred = self._reconstruct(data.X.A, use_data=True) else: stim_pred = self._reconstruct(data.X.A) else: if data_space == 'latent': stim_pred = self._reconstruct(data.X, use_data=True) else: stim_pred = self._reconstruct(data.X) return stim_pred def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.rae_model = load_model(os.path.join(self.model_to_use, 'mmd_rae.h5'), compile=False) self.encoder_model = load_model(os.path.join(self.model_to_use, 'encoder.h5'), compile=False) self.decoder_model = load_model(os.path.join(self.model_to_use, 'decoder.h5'), compile=False) self._loss_function() def train(self, train_data, use_validation=False, valid_data=None, n_epochs=25, batch_size=32, early_stop_limit=20, threshold=0.0025, initial_run=True, shuffle=True, verbose=2, save=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_katrain_kang.h5ad >>> validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ if initial_run: log.info("----Training----") if use_validation and valid_data is None: raise Exception("valid_data is None but use_validation is True.") callbacks = [ History(), EarlyStopping(patience=early_stop_limit, monitor='val_loss', min_delta=threshold), CSVLogger(filename="./csv_logger.log") ] if sparse.issparse(train_data.X): train_data.X = train_data.X.A train_labels, _ = label_encoder(train_data) x = train_data.X y = [train_data.X, train_labels] if use_validation: if sparse.issparse(valid_data.X): valid_data.X = valid_data.X.A valid_labels, _ = label_encoder(valid_data) x_valid = valid_data.X y_valid = [valid_data.X, valid_labels] histories = self.rae_model.fit( x=x, y=y, epochs=n_epochs, batch_size=batch_size, validation_data=(x_valid, y_valid), shuffle=shuffle, callbacks=callbacks, verbose=verbose) else: histories = self.rae_model.fit( x=x, y=y, epochs=n_epochs, batch_size=batch_size, shuffle=shuffle, callbacks=callbacks, verbose=verbose) if save: os.makedirs(self.model_to_use, exist_ok=True) self.rae_model.save(os.path.join(self.model_to_use, "mmd_rae.h5"), overwrite=True) self.encoder_model.save(os.path.join(self.model_to_use, "encoder.h5"), overwrite=True) self.decoder_model.save(os.path.join(self.model_to_use, "decoder.h5"), overwrite=True) log.info(f"Model saved in file: {self.model_to_use}. Training finished") return histories PKGOIE좸]]trvae/models/_trvae_atac.pyimport logging import os import anndata import keras import numpy as np from keras.layers import Dense, BatchNormalization, Dropout, Input, concatenate, Lambda, Activation from keras.layers.advanced_activations import LeakyReLU from keras.models import Model, load_model from keras.utils import to_categorical from scipy import sparse from sklearn.preprocessing import LabelEncoder from trvae.models._losses import LOSSES from trvae.models._activations import ACTIVATIONS from trvae.utils import remove_sparsity, label_encoder from ._utils import sample_z log = logging.getLogger(__file__) class trVAEATAC: """ Regularized C-VAE vector Network class. This class contains the implementation of Conditional Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, n_labels, z_dimension=100, **kwargs): self.x_dim = x_dimension self.z_dim = z_dimension self.n_labels = n_labels self.n_domains = kwargs.get("n_domains", 2) self.mmd_dim = kwargs.get('mmd_dimension', 128) self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 0.001) self.beta = kwargs.get("beta", 100) self.prev_beta = kwargs.get("beta", 100) self.gamma = kwargs.get("gamma", 1.0) self.eta = kwargs.get("eta", 1.0) self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_path = kwargs.get("model_path", "./") self.kernel_method = kwargs.get("kernel", "multi-scale-rbf") self.output_activation = kwargs.get("output_activation", 'relu') self.mmd_computation_way = kwargs.get("mmd_computation_way", "general") self.print_summary = kwargs.get("print_summary", True) self.x = Input(shape=(self.x_dim,), name="data") self.encoder_labels = Input(shape=(self.n_domains,), name="encoder_labels") self.decoder_labels = Input(shape=(self.n_domains,), name="decoder_labels") self.z = Input(shape=(self.z_dim,), name="latent_data") self.mmd_latent = Input(shape=(self.mmd_dim,), name="MMD_Latent") self.label_enc = LabelEncoder() self.domain_enc = LabelEncoder() self.init_w = keras.initializers.glorot_normal() self._create_network() self._loss_function() if self.print_summary: self.encoder_model.summary() self.decoder_model.summary() self.cvae_model.summary() self.classifier_model.summary() def _encoder(self, x, y, name="encoder"): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ xy = concatenate([x, y], axis=1) h = Dense(800, kernel_initializer=self.init_w, use_bias=False)(xy) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(800, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) mean = Dense(self.z_dim, kernel_initializer=self.init_w)(h) log_var = Dense(self.z_dim, kernel_initializer=self.init_w)(h) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=[x, y], outputs=[mean, log_var, z], name=name) return mean, log_var, model def _mmd_decoder(self, z, y, name="decoder"): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ zy = concatenate([z, y], axis=1) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, use_bias=False)(zy) h = BatchNormalization()(h) h_mmd = LeakyReLU(name="mmd")(h) h = Dense(800, kernel_initializer=self.init_w, use_bias=False)(h_mmd) h = BatchNormalization(axis=1)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(800, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.x_dim, kernel_initializer=self.init_w, use_bias=True)(h) h = ACTIVATIONS[self.output_activation](h) decoder_model = Model(inputs=[z, y], outputs=h, name=name) mmd_model = Model(inputs=[z, y], outputs=h_mmd, name='decoder_mmd') return decoder_model, mmd_model def _classifier(self, name='classifier_from_mmd_latent'): mmd_latent = Input(shape=(self.mmd_dim,)) h = Dense(self.n_labels, activation='softmax', name='classifier_prob', kernel_initializer=self.init_w, use_bias=True)(mmd_latent) model = Model(inputs=mmd_latent, outputs=h, name=name) return model def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ inputs = [self.x, self.encoder_labels, self.decoder_labels] self.mu, self.log_var, self.encoder_model = self._encoder(*inputs[:2], name="encoder") self.decoder_model, self.decoder_mmd_model = self._mmd_decoder(self.z, self.decoder_labels, name="decoder") self.classifier_from_mmd_latent_model = self._classifier() decoder_output = self.decoder_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) decoder_mmd_output = self.decoder_mmd_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) reconstruction_output = Lambda(lambda x: x, name="kl_reconstruction")(decoder_output) mmd_output = Lambda(lambda x: x, name="mmd")(decoder_mmd_output) classifier_output = self.classifier_from_mmd_latent_model(decoder_mmd_output) self.cvae_model = Model(inputs=inputs, outputs=[reconstruction_output, mmd_output], name="cvae") self.classifier_model = Model(inputs=inputs, outputs=classifier_output, name='classifier') def _loss_function(self): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ def batch_loss(): self.cvae_optimizer = keras.optimizers.Adam(lr=self.lr) self.classifier_optimizer = keras.optimizers.Adam(lr=self.lr) self.cvae_model.compile(optimizer=self.cvae_optimizer, loss=[LOSSES['mse'](self.mu, self.log_var, self.alpha, self.eta), LOSSES['mmd'](self.n_domains, self.beta, self.kernel_method, self.mmd_computation_way)], metrics={self.cvae_model.outputs[0].name: LOSSES['mse'](self.mu, self.log_var, self.alpha, self.eta), self.cvae_model.outputs[1].name: LOSSES['mmd'](self.n_domains, self.beta, self.kernel_method, self.mmd_computation_way), }) self.classifier_model.compile(optimizer=self.classifier_optimizer, loss=LOSSES['cce'](self.gamma), metrics=['acc']) batch_loss() def to_latent(self, adata, encoder_labels): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ adata = remove_sparsity(adata) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_domains) latent = self.encoder_model.predict([adata.X, encoder_labels])[2] latent_adata = anndata.AnnData(X=latent) latent_adata.obs = adata.obs.copy(deep=True) latent_adata.var_names = adata.var_names return latent_adata def to_mmd_layer(self, adata, encoder_labels, feed_fake=0): """ Map `data` in to the pn layer after latent layer. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ if feed_fake > 0: decoder_labels = np.zeros(shape=encoder_labels.shape) + feed_fake else: decoder_labels = encoder_labels encoder_labels = to_categorical(encoder_labels, num_classes=self.n_domains) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_domains) adata = remove_sparsity(adata) mmd_latent = self.cvae_model.predict([adata.X, encoder_labels, decoder_labels])[1] mmd_latent_adata = anndata.AnnData(X=mmd_latent) mmd_latent_adata.obs = adata.obs.copy(deep=True) return mmd_latent_adata def predict(self, adata, encoder_labels, decoder_labels): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ adata = remove_sparsity(adata) reconstructed = self.cvae_model.predict([adata.X, encoder_labels, decoder_labels])[0] reconstructed_adata = anndata.AnnData(X=reconstructed) reconstructed_adata.obs = adata.obs.copy(deep=True) reconstructed_adata.var_names = adata.var_names return reconstructed_adata def _make_classifier_trainable(self, trainable): self.classifier_model.layers[-1].trainable = trainable def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.cvae_model = load_model(os.path.join(self.model_path, 'mmd_cvae.h5'), compile=False) self.classifier_model = load_model(os.path.join(self.model_path, 'classifier.h5'), compile=False) self.encoder_model = load_model(os.path.join(self.model_path, 'encoder.h5'), compile=False) self.decoder_model = load_model(os.path.join(self.model_path, 'decoder.h5'), compile=False) self.decoder_mmd_model = load_model(os.path.join(self.model_path, 'decoder_mmd.h5'), compile=False) self._loss_function() def save_model(self): os.makedirs(self.model_path, exist_ok=True) self.cvae_model.save(os.path.join(self.model_path, "mmd_cvae.h5"), overwrite=True) self.classifier_model.save(os.path.join(self.model_path, "classifier.h5"), overwrite=True) self.encoder_model.save(os.path.join(self.model_path, "encoder.h5"), overwrite=True) self.decoder_model.save(os.path.join(self.model_path, "decoder.h5"), overwrite=True) self.decoder_mmd_model.save(os.path.join(self.model_path, "decoder_mmd.h5"), overwrite=True) def train(self, train_adata, valid_adata, domain_key, label_key, source_key, target_key, domain_encoder, n_epochs=25, batch_size=32, early_stop_limit=20, lr_reducer=10, save=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_kang.h5ad) validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ if not isinstance(source_key, list): source_key = [source_key] if not isinstance(target_key, list): target_key = [target_key] train_labels_encoded, self.domain_enc = label_encoder(train_adata, label_encoder=domain_encoder, condition_key=domain_key) valid_labels_encoded, _ = label_encoder(valid_adata, label_encoder=domain_encoder, condition_key=domain_key) train_labels_onehot = to_categorical(train_labels_encoded, num_classes=self.n_domains) valid_labels_onehot = to_categorical(valid_labels_encoded, num_classes=self.n_domains) train_adata = remove_sparsity(train_adata) valid_adata = remove_sparsity(valid_adata) source_adata_train = train_adata.copy()[train_adata.obs[domain_key].isin(source_key)] source_classes_train = source_adata_train.obs[label_key].values source_classes_train = self.label_enc.fit_transform(source_classes_train) source_classes_train = to_categorical(source_classes_train, num_classes=self.n_labels) source_domains_train_encoded = np.zeros(source_adata_train.shape[0]) source_domains_train_onehot = to_categorical(source_domains_train_encoded, num_classes=self.n_domains) target_adata_train = train_adata.copy()[train_adata.obs[domain_key].isin(target_key)] target_classes_train = target_adata_train.obs[label_key].values target_classes_train = self.label_enc.transform(target_classes_train) target_classes_train = to_categorical(target_classes_train, num_classes=self.n_labels) target_domains_encoded = np.ones(target_adata_train.shape[0]) target_domains_onehot = to_categorical(target_domains_encoded, num_classes=self.n_domains) source_adata_valid = valid_adata.copy()[valid_adata.obs[domain_key].isin(source_key)] source_classes_valid = source_adata_valid.obs[label_key].values source_classes_valid = self.label_enc.transform(source_classes_valid) source_classes_valid = to_categorical(source_classes_valid, num_classes=self.n_labels) source_domains_valid_encoded = np.zeros(source_adata_valid.shape[0]) source_domains_valid_onehot = to_categorical(source_domains_valid_encoded, num_classes=self.n_domains) best_val_loss = 100000.0 patience = 0 for i in range(n_epochs): x_train = [train_adata.X, train_labels_onehot, train_labels_onehot] y_train = [train_adata.X, train_labels_encoded] x_valid = [valid_adata.X, valid_labels_onehot, valid_labels_onehot] y_valid = [valid_adata.X, valid_labels_encoded] cvae_history = self.cvae_model.fit( x=x_train, y=y_train, epochs=1, batch_size=batch_size, validation_data=(x_valid, y_valid), verbose=0, ) x_train = [source_adata_train.X, source_domains_train_onehot, source_domains_train_onehot] y_train = source_classes_train x_valid = [source_adata_valid.X, source_domains_valid_onehot, source_domains_valid_onehot] y_valid = source_classes_valid class_history = self.classifier_model.fit( x=x_train, y=y_train, epochs=1, batch_size=batch_size, validation_data=(x_valid, y_valid), verbose=0, ) x_target = [target_adata_train.X, target_domains_onehot, target_domains_onehot] y_target = target_classes_train _, target_acc = self.classifier_model.evaluate(x_target, y_target, verbose=0) cvae_loss = cvae_history.history['loss'][0] cvae_kl_recon_loss = cvae_history.history['kl_reconstruction_loss'][0] cvae_mmd_loss = cvae_history.history['mmd_loss'][0] cvae_loss_valid = cvae_history.history['val_loss'][0] cvae_kl_recon_loss_valid = cvae_history.history['val_kl_reconstruction_loss'][0] cvae_mmd_loss_valid = cvae_history.history['val_mmd_loss'][0] class_cce_loss = class_history.history['loss'][0] class_accuracy = class_history.history['acc'][0] class_cce_loss_valid = class_history.history['val_loss'][0] class_accuracy_valid = class_history.history['val_acc'][0] print(f"Epoch {i + 1}/{n_epochs}:") print(f'loss: {cvae_loss:.4f} - KL_Recon_loss: {cvae_kl_recon_loss:.4f}' f" - MMD_loss: {cvae_mmd_loss:.4f} - CCE_Loss: {class_cce_loss:.4f} - CCE_Acc: {class_accuracy:.3f}" f" - val_loss: {cvae_loss_valid:.4f} - val_KL_Recon_loss: {cvae_kl_recon_loss_valid:.4f}" f" - val_MMD_loss: {cvae_mmd_loss_valid:.4f} - val_CCE_Loss: {class_cce_loss_valid:.4f}" f" - val_CCE_Acc: {class_accuracy_valid:.3f} - Target_acc: {target_acc:.3f}") if class_cce_loss_valid > best_val_loss: patience += 1 if patience > early_stop_limit: break else: best_val_loss = class_cce_loss_valid patience = 0 if lr_reducer > 0 and (i + 1) % lr_reducer == 0: self.lr = self.lr / 10 self._loss_function() print(f"Epoch {i}: learning rate reduced to {self.lr}") if save: self.save_model() PKGO0__trvae/models/_trvae_multi.pyimport logging import os import anndata import keras import numpy as np from keras.callbacks import CSVLogger, History, EarlyStopping, ReduceLROnPlateau, LambdaCallback, ModelCheckpoint from keras.layers import Dense, BatchNormalization, Dropout, Input, concatenate, Lambda from keras.layers.advanced_activations import LeakyReLU from keras.models import Model, load_model from keras.utils import to_categorical from scipy import sparse from sklearn.preprocessing import LabelEncoder from trvae.models._activations import ACTIVATIONS from trvae.models._losses import LOSSES from trvae.models._utils import sample_z, print_message from trvae.utils import label_encoder, remove_sparsity log = logging.getLogger(__file__) class trVAEMulti: """ trVAE Network class. This class contains the implementation of Regularized Conditional Variational Auto-encoder network. # Parameters kwargs: key: `mmd_dimension`: int dimension of MMD layer in trVAE key: `kernel_method`: str kernel method for MMD loss calculation key: `output_activation`: str activation of output layer in trVAE key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for KL loss. key: `beta`: float beta coefficient for MMD loss. key: `eta`: float eta coefficient for reconstruction loss. key: `lambda_l1`: float lambda coefficient for L1 regularization key: `lambda_l2`: float lambda coefficient for L2 regularization key: `clip_value`: float clip_value for optimizer x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, n_conditions, z_dimension=20, **kwargs): self.x_dim = x_dimension self.z_dim = z_dimension self.mmd_dim = kwargs.get('mmd_dimension', 128) self.n_conditions = n_conditions self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 1e-5) self.beta = kwargs.get("beta", 10) self.eta = kwargs.get("eta", 100) self.dr_rate = kwargs.get("dropout_rate", 0.1) self.model_to_use = kwargs.get("model_path", "./") self.kernel_method = kwargs.get("kernel", "multi-scale-rbf") self.output_activation = kwargs.get("output_activation", 'relu') self.mmd_calc_mode = kwargs.get("mmd_computation_way", "general") self.clip_value = kwargs.get('clip_value', 5) self.lambda_l1 = kwargs.get('lambda_l1', 0.0) self.lambda_l2 = kwargs.get('lambda_l2', 0.0) self.x = Input(shape=(self.x_dim,), name="data") self.encoder_labels = Input(shape=(self.n_conditions,), name="encoder_labels") self.decoder_labels = Input(shape=(self.n_conditions,), name="decoder_labels") self.z = Input(shape=(self.z_dim,), name="latent_data") self.init_w = keras.initializers.glorot_normal() self.regularizer = keras.regularizers.l1_l2(self.lambda_l1, self.lambda_l2) self._create_network() self._loss_function() self.encoder_model.summary() self.decoder_model.summary() self.cvae_model.summary() def _encoder(self, name="encoder"): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. model: Model A keras Model object for Encoder subnetwork of trVAE """ xy = concatenate([self.x, self.encoder_labels], axis=1) h = Dense(512, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(xy) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(h) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) mean = Dense(self.z_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer)(h) log_var = Dense(self.z_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer)(h) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=[self.x, self.encoder_labels], outputs=[mean, log_var, z], name=name) return mean, log_var, model def _mmd_decoder(self, name="decoder"): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns decoder_model: Model A keras Model object for Decoder subnetwork of trVAE decoder_mmd_model: Model A keras Model object for MMD Decoder subnetwork of trVAE """ zy = concatenate([self.z, self.decoder_labels], axis=1) h = Dense(self.mmd_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(zy) h = BatchNormalization(axis=1, scale=True)(h) h_mmd = LeakyReLU(name="mmd")(h) h = Dropout(self.dr_rate)(h_mmd) h = Dense(512, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=False)(h) h = BatchNormalization(axis=1, scale=True)(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.x_dim, kernel_initializer=self.init_w, kernel_regularizer=self.regularizer, use_bias=True)(h) h = ACTIVATIONS[self.output_activation](h) decoder_model = Model(inputs=[self.z, self.decoder_labels], outputs=h, name=name) decoder_mmd_model = Model(inputs=[self.z, self.decoder_labels], outputs=h_mmd, name='decoder_mmd') return decoder_model, decoder_mmd_model def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ inputs = [self.x, self.encoder_labels, self.decoder_labels] self.mu, self.log_var, self.encoder_model = self._encoder(name="encoder") self.decoder_model, self.decoder_mmd_model = self._mmd_decoder(name="decoder") decoder_output = self.decoder_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) mmd_output = self.decoder_mmd_model([self.encoder_model(inputs[:2])[2], self.decoder_labels]) reconstruction_output = Lambda(lambda x: x, name="kl_mse")(decoder_output) mmd_output = Lambda(lambda x: x, name="mmd")(mmd_output) self.cvae_model = Model(inputs=inputs, outputs=[reconstruction_output, mmd_output], name="cvae") def _calculate_loss(self): """ Computes the MSE and MMD loss for trVAE using `LOSSES` dictionary # Parameters No parameters are needed. # Returns loss: function mse loss function mmd_loss: function mmd loss function """ loss = LOSSES['mse'](self.mu, self.log_var, self.alpha, self.eta) mmd_loss = LOSSES['mmd'](self.n_conditions, self.beta, self.kernel_method, self.mmd_calc_mode) return loss, mmd_loss def _loss_function(self): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ loss, mmd_loss = self._calculate_loss() self.cvae_optimizer = keras.optimizers.Adam(lr=self.lr) self.cvae_model.compile(optimizer=self.cvae_optimizer, loss=[loss, mmd_loss], metrics={self.cvae_model.outputs[0].name: loss, self.cvae_model.outputs[1].name: mmd_loss}) def to_latent(self, adata, encoder_labels, return_adata=True): """ Map `adata` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters adata: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. encoder_labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. return_adata: boolean if `True`, will output as an `anndata` object or put the results in the `obsm` attribute of `adata` # Returns output: `~anndata.AnnData` returns `anndata` object containing latent space encoding of 'adata' """ adata = remove_sparsity(adata) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) latent = self.encoder_model.predict([adata.X, encoder_labels])[2] latent = np.nan_to_num(latent) if return_adata: output = anndata.AnnData(X=latent) output.obs = adata.obs.copy(deep=True) else: output = latent return output def to_mmd_layer(self, adata, encoder_labels, feed_fake=0, return_adata=True): """ Map `adata` in to the MMD layer of trVAE network. This function will compute output activation of MMD layer in trVAE. # Parameters adata: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. encoder_labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. feed_fake: int if `feed_fake` is non-negative, `decoder_labels` will be identical to `encoder_labels`. if `feed_fake` is not non-negative, `decoder_labels` will be fed with `feed_fake` value. return_adata: boolean if `True`, will output as an `anndata` object or put the results in the `obsm` attribute of `adata` # Returns output: `~anndata.AnnData` returns `anndata` object containing MMD latent space encoding of 'adata' """ if feed_fake >= 0: decoder_labels = np.zeros(shape=encoder_labels.shape) + feed_fake else: decoder_labels = encoder_labels encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) adata = remove_sparsity(adata) x = [adata.X, encoder_labels, decoder_labels] mmd_latent = self.cvae_model.predict(x)[1] mmd_latent = np.nan_to_num(mmd_latent) if return_adata: output = anndata.AnnData(X=mmd_latent) output.obs = adata.obs.copy(deep=True) else: output = mmd_latent return output def predict(self, adata, encoder_labels, decoder_labels, return_adata=True): """ Predicts the cell type provided by the user in stimulated condition. # Parameters adata: `~anndata.AnnData` Annotated data matrix whether in primary space. encoder_labels: `numpy nd-array` `numpy nd-array` of labels to be fed as encoder's condition array. decoder_labels: `numpy nd-array` `numpy nd-array` of labels to be fed as decoder's condition array. return_adata: boolean if `True`, will output as an `anndata` object or put the results in the `obsm` attribute of `adata` # Returns output: `~anndata.AnnData` `anndata` object of predicted cells in primary space. # Example ```python import scanpy as sc import trvae train_data = sc.read("train.h5ad") valid_adata = sc.read("validation.h5ad") n_conditions = len(train_adata.obs['condition'].unique().tolist()) network = trvae.archs.trVAEMulti(train_adata.shape[1], n_conditions) network.train(train_adata, valid_adata, le=None, condition_key="condition", cell_type_key="cell_label", n_epochs=1000, batch_size=256 ) encoder_labels, _ = trvae.utils.label_encoder(train_adata, condition_key="condition") decoder_labels, _ = trvae.utils.label_encoder(train_adata, condition_key="condition") pred_adata = network.predict(train_adata, encoder_labels, decoder_labels) ``` """ adata = remove_sparsity(adata) encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) reconstructed = self.cvae_model.predict([adata.X, encoder_labels, decoder_labels])[0] reconstructed = np.nan_to_num(reconstructed) if return_adata: output = anndata.AnnData(X=reconstructed) output.obs = adata.obs.copy(deep=True) output.var_names = adata.var_names else: output = reconstructed return output def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import trvae train_data = sc.read("train.h5ad") valid_adata = sc.read("validation.h5ad") n_conditions = len(train_adata.obs['condition'].unique().tolist()) network = trvae.archs.trVAEMulti(train_adata.shape[1], n_conditions) network.restore_model() ``` """ self.cvae_model = load_model(filepath=os.path.join(self.model_to_use, "best_model.h5"), compile=False) self.encoder_model = self.cvae_model.get_layer("encoder") self.decoder_model = self.cvae_model.get_layer("decoder") self.decoder_mmd_model = self.cvae_model.get_layer("decoder_mmd") self._loss_function() def save_model(self): os.makedirs(self.model_to_use, exist_ok=True) self.cvae_model.save(os.path.join(self.model_to_use, "best_model.h5"), overwrite=True) def train(self, train_adata, valid_adata=None, condition_encoder=None, condition_key='condition', n_epochs=10000, batch_size=1024, early_stop_limit=100, lr_reducer=80, threshold=0.0, monitor='val_loss', shuffle=True, verbose=0, save=True, monitor_best=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters train_adata: `~anndata.AnnData` `AnnData` object for training trVAE valid_adata: `~anndata.AnnData` `AnnData` object for validating trVAE (if None, trVAE will automatically split the data with fraction of 80%/20%. condition_encoder: dict dictionary of encoded conditions (if None, trVAE will make one for data) condition_key: str name of conditions (domains) column in obs matrix cell_type_key: str name of cell_types (labels) column in obs matrix n_epochs: int number of epochs to iterate and optimize network weights batch_size: int number of samples to be used in each batch for network weights optimization early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. monitor: str metric to be monitored for early stopping. shuffle: boolean if `True`, `train_adata` will be shuffled before training. verbose: int level of verbosity save: boolean if `True`, the model will be saved in the specified path after training. # Returns Nothing will be returned # Example ```python import scanpy as sc import trvae train_data = sc.read("train.h5ad") valid_adata = sc.read("validation.h5ad") n_conditions = len(train_adata.obs['condition'].unique().tolist()) network = trvae.archs.trVAEMulti(train_adata.shape[1], n_conditions) network.train(train_adata, valid_adata, le=None, condition_key="condition", cell_type_key="cell_label", n_epochs=1000, batch_size=256 ) ``` """ train_labels_encoded, _ = label_encoder(train_adata, condition_encoder, condition_key) train_labels_onehot = to_categorical(train_labels_encoded, num_classes=self.n_conditions) callbacks = [ History(), CSVLogger(filename="./csv_logger.log"), ] if early_stop_limit > 0: callbacks.append(EarlyStopping(patience=early_stop_limit, monitor=monitor, min_delta=threshold)) if lr_reducer > 0: callbacks.append(ReduceLROnPlateau(monitor=monitor, patience=lr_reducer, verbose=verbose)) if verbose > 2: callbacks.append( LambdaCallback(on_epoch_end=lambda epoch, logs: print_message(epoch, logs, n_epochs, verbose))) fit_verbose = 0 else: fit_verbose = verbose if monitor_best: os.makedirs(self.model_to_use, exist_ok=True) callbacks.append(ModelCheckpoint(filepath=os.path.join(self.model_to_use, "best_model.h5"), save_best_only=True, monitor=monitor, period=50)) if sparse.issparse(train_adata.X): train_adata.X = train_adata.X.A x = [train_adata.X, train_labels_onehot, train_labels_onehot] y = [train_adata.X, train_labels_encoded] if valid_adata is not None: if sparse.issparse(valid_adata.X): valid_adata.X = valid_adata.X.A valid_labels_encoded, _ = label_encoder(valid_adata, condition_encoder, condition_key) valid_labels_onehot = to_categorical(valid_labels_encoded, num_classes=self.n_conditions) x_valid = [valid_adata.X, valid_labels_onehot, valid_labels_onehot] y_valid = [valid_adata.X, valid_labels_encoded] history = self.cvae_model.fit(x=x, y=y, epochs=n_epochs, batch_size=batch_size, validation_data=(x_valid, y_valid), shuffle=shuffle, callbacks=callbacks, verbose=fit_verbose) else: history = self.cvae_model.fit(x=x, y=y, epochs=n_epochs, batch_size=batch_size, validation_split=0.2, shuffle=shuffle, callbacks=callbacks, verbose=fit_verbose) if monitor_best: self.restore_model() elif save and not monitor_best: self.save_model() def get_corrected(self, adata, labels, return_z=False): """ Computes all trVAE's outputs (Latent (optional), MMD Latent, reconstruction) # Parameters adata: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: `numpy nd-array` `numpy nd-array` of labels to be fed as encoder's condition array. return_z: boolean if `True`, will also put z_latent in the `obsm` attribute of `adata` # Returns Nothing will be returned. # Example ```python import scanpy as sc import trvae train_data = sc.read("train.h5ad") valid_adata = sc.read("validation.h5ad") n_conditions = len(train_adata.obs['condition'].unique().tolist()) network = trvae.archs.trVAEMulti(train_adata.shape[1], n_conditions) network.restore_model() ``` """ reference_labels = np.zeros(adata.shape[0]) adata.obsm['mmd_latent'] = self.to_mmd_layer(adata, labels, -1, return_adata=False) adata.obsm['reconstructed'] = self.predict(adata, reference_labels, labels, return_adata=False) if return_z: adata.obsm['z_latent'] = self.to_latent(adata, labels, return_adata=False) def get_reconstruction_error(self, adata, condition_key): adata = remove_sparsity(adata) labels_encoded, _ = label_encoder(adata, None, condition_key) labels_onehot = to_categorical(labels_encoded, num_classes=self.n_conditions) x = [adata.X, labels_onehot, labels_onehot] y = [adata.X, labels_encoded] return self.cvae_model.evaluate(x, y, verbose=0) PKGOI> VVtrvae/models/_trvae_multi_tf.pyimport logging import os import numpy as np import tensorflow as tf from keras.utils import to_categorical from scipy import sparse from trvae.models._utils import compute_mmd from trvae.utils import label_encoder log = logging.getLogger(__file__) class trVAEMultiTF: """ C-VAE vector Network class. This class contains the implementation of Conditional Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, z_dimension=100, n_conditions=2, **kwargs): tf.reset_default_graph() self.x_dim = x_dimension self.z_dim = z_dimension self.mmd_dim = kwargs.get('mmd_dimension', 128) self.n_conditions = n_conditions self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 0.001) self.beta = kwargs.get("beta", 100) self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_to_use = kwargs.get("model_path", "./") self.kernel_method = kwargs.get("kernel", "multi-scale-rbf") self.output_activation = kwargs.get("output_activation", 'relu') self.loss_fn = kwargs.get("loss_fn", 'mse') self.ridge = kwargs.get('ridge', 0.1) self.clip_value = kwargs.get('clip_value', 3.0) self.is_training = tf.placeholder(tf.bool, name='training_flag') self.global_step = tf.Variable(0, name='global_step', trainable=False, dtype=tf.int32) self.x = tf.placeholder(tf.float32, shape=[None, self.x_dim], name="data") self.z = tf.placeholder(tf.float32, shape=[None, self.z_dim], name="latent") self.encoder_labels = tf.placeholder(tf.float32, shape=[None, self.n_conditions], name="encoder_labels") self.decoder_labels = tf.placeholder(tf.float32, shape=[None, self.n_conditions], name="decoder_labels") self.time_step = tf.placeholder(tf.int32) self.init_w = tf.contrib.layers.xavier_initializer() self._create_network() self._loss_function() init = tf.global_variables_initializer() self.sess = tf.InteractiveSession() self.saver = tf.train.Saver(max_to_keep=1) self.sess.run(init) def _encoder(self): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ with tf.variable_scope("encoder", reuse=tf.AUTO_REUSE): xy = tf.concat([self.x, self.encoder_labels], axis=1) h = tf.layers.dense(inputs=xy, units=800, kernel_initializer=self.init_w, use_bias=False) h = tf.layers.batch_normalization(h, axis=1, training=self.is_training) h = tf.nn.leaky_relu(h) h = tf.layers.dense(inputs=h, units=800, kernel_initializer=self.init_w, use_bias=False) h = tf.layers.batch_normalization(h, axis=1, training=self.is_training) h = tf.nn.leaky_relu(h) h = tf.layers.dropout(h, self.dr_rate, training=self.is_training) h = tf.layers.dense(inputs=h, units=self.mmd_dim, kernel_initializer=self.init_w, use_bias=False) h = tf.layers.batch_normalization(h, axis=1, training=self.is_training) h = tf.nn.leaky_relu(h) h = tf.layers.dropout(h, self.dr_rate, training=self.is_training) mean = tf.layers.dense(inputs=h, units=self.z_dim, kernel_initializer=self.init_w) log_var = tf.layers.dense(inputs=h, units=self.z_dim, kernel_initializer=self.init_w) return mean, log_var def _mmd_decoder(self): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE): xy = tf.concat([self.z_mean, self.decoder_labels], axis=1) h = tf.layers.dense(inputs=xy, units=self.mmd_dim, kernel_initializer=self.init_w, use_bias=False) h = tf.layers.batch_normalization(h, axis=1, training=self.is_training) h_mmd = tf.nn.leaky_relu(h) h = tf.layers.dropout(h_mmd, self.dr_rate, training=self.is_training) h = tf.layers.dense(inputs=h, units=800, kernel_initializer=self.init_w, use_bias=False) h = tf.layers.batch_normalization(h, axis=1, training=self.is_training) h = tf.nn.leaky_relu(h) h = tf.layers.dropout(h, self.dr_rate, training=self.is_training) h = tf.layers.dense(inputs=h, units=800, kernel_initializer=self.init_w, use_bias=False) h = tf.layers.batch_normalization(h, axis=1, training=self.is_training) h = tf.nn.leaky_relu(h) h = tf.layers.dropout(h, self.dr_rate, training=self.is_training) h = tf.layers.dense(inputs=h, units=self.x_dim, kernel_initializer=self.init_w, use_bias=True) if self.output_activation == 'relu': h = tf.nn.relu(h) return h, h_mmd def _sample_z(self): """ Samples from standard Normal distribution with shape [size, z_dim] and applies re-parametrization trick. It is actually sampling from latent space distributions with N(mu, var) computed in `_encoder` function. # Parameters No parameters are needed. # Returns The computed Tensor of samples with shape [size, z_dim]. """ batch_size = tf.shape(self.mu)[0] eps = tf.random_normal(shape=[batch_size, self.z_dim]) return self.mu + tf.exp(self.log_var / 2) * eps def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ self.mu, self.log_var = self._encoder() self.z_mean = self._sample_z() self.x_hat, self.mmd_hl = self._mmd_decoder() def _loss_function(self): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ self.kl_loss = 0.5 * tf.reduce_sum( tf.exp(self.log_var) + tf.square(self.mu) - 1. - self.log_var, 1) self.recon_loss = 0.5 * tf.reduce_sum(tf.square((self.x - self.x_hat)), 1) self.kl_recon_loss = tf.reduce_mean(self.recon_loss + self.alpha * self.kl_loss) # MMD loss computation with tf.variable_scope("mmd_loss", reuse=tf.AUTO_REUSE): real_labels = tf.argmax(self.decoder_labels, axis=1) real_labels = tf.reshape(tf.cast(real_labels, 'int32'), (-1,)) conditions_mmd = tf.dynamic_partition(self.mmd_hl, real_labels, num_partitions=self.n_conditions) loss = 0.0 for i in range(len(conditions_mmd)): for j in range(i): loss += compute_mmd(conditions_mmd[j], conditions_mmd[j + 1], self.kernel_method) self.mmd_loss = self.beta * loss self.trvae_loss = self.mmd_loss + self.kl_recon_loss with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)): self.solver = tf.train.AdamOptimizer(learning_rate=self.lr).minimize(self.trvae_loss) def to_latent(self, adata, labels): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ if sparse.issparse(adata.X): adata.X = adata.X.A labels = to_categorical(labels, num_classes=self.n_conditions) latent = self.sess.run(self.z_mean, feed_dict={self.x: adata.X, self.encoder_labels: labels, self.is_training: False}) return latent def to_mmd_layer(self, adata, encoder_labels, feed_fake=0): """ Map `data` in to the pn layer after latent layer. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ if feed_fake > 0: decoder_labels = np.zeros(shape=encoder_labels.shape) + feed_fake else: decoder_labels = encoder_labels encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) if sparse.issparse(adata.X): adata.X = adata.X.A latent = self.sess.run(self.mmd_hl, feed_dict={self.x: adata.X, self.encoder_labels: encoder_labels, self.decoder_labels: decoder_labels, self.is_training: False}) return latent def _reconstruct(self, data, encoder_labels, decoder_labels, size_factor=None): """ Map back the latent space encoding via the decoder. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in latent space or primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. use_data: bool this flag determines whether the `data` is already in latent space or not. if `True`: The `data` is in latent space (`data.X` is in shape [n_obs, z_dim]). if `False`: The `data` is not in latent space (`data.X` is in shape [n_obs, n_vars]). # Returns rec_data: 'numpy nd-array' returns 'numpy nd-array` containing reconstructed 'data' in shape [n_obs, n_vars]. """ rec_data = self.sess.run(self.x_hat, feed_dict={self.x: data, self.encoder_labels: encoder_labels, self.decoder_labels: decoder_labels, self.is_training: False}) return rec_data def predict(self, adata, encoder_labels, decoder_labels, size_factor=None): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ if sparse.issparse(adata.X): adata.X = adata.X.A encoder_labels = to_categorical(encoder_labels, num_classes=self.n_conditions) decoder_labels = to_categorical(decoder_labels, num_classes=self.n_conditions) stim_pred = self._reconstruct(adata.X, encoder_labels, decoder_labels, size_factor) return stim_pred def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.saver.restore(self.sess, self.model_to_use) def train(self, train_data, le=None, condition_key='condition', use_validation=False, valid_data=None, n_epochs=25, batch_size=32, early_stop_limit=20, lr_reducer=0, threshold=0.0025, initial_run=True, shuffle=True, verbose=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_katrain_kang.h5ad >>> validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ if initial_run: log.info("----Training----") assign_step_zero = tf.assign(self.global_step, 0) _init_step = self.sess.run(assign_step_zero) if not initial_run: self.saver.restore(self.sess, self.model_to_use) train_labels, le = label_encoder(train_data, label_encoder=le, condition_key=condition_key) if use_validation and valid_data is None: raise Exception("valid_data is None but use_validation is True.") if use_validation: valid_labels, _ = label_encoder(valid_data, label_encoder=le, condition_key=condition_key) loss_hist = [] patience = early_stop_limit min_delta = threshold patience_cnt = 0 min_loss = 1000000.0 for it in range(n_epochs): increment_global_step_op = tf.assign(self.global_step, self.global_step + 1) _step = self.sess.run(increment_global_step_op) current_step = self.sess.run(self.global_step) train_loss = 0 train_mmd_loss = 0 for lower in range(0, train_data.shape[0], batch_size): upper = min(lower + batch_size, train_data.shape[0]) if sparse.issparse(train_data.X): x_mb = train_data[lower:upper, :].X.A else: x_mb = train_data[lower:upper, :].X y_mb = train_labels[lower:upper] y_mb = to_categorical(y_mb, num_classes=self.n_conditions) _, current_loss_train, current_mmd_loss_train = self.sess.run( [self.solver, self.trvae_loss, self.mmd_loss], feed_dict={self.x: x_mb, self.encoder_labels: y_mb, self.decoder_labels: y_mb, self.time_step: current_step, self.is_training: True}) train_loss += current_loss_train train_mmd_loss += current_mmd_loss_train if use_validation: if sparse.issparse(valid_data.X): x_valid = valid_data.X.A else: x_valid = valid_data.X y_valid = to_categorical(valid_labels, num_classes=self.n_conditions) valid_loss, valid_mmd_loss = self.sess.run([self.trvae_loss, self.mmd_loss], feed_dict={self.x: x_valid, self.encoder_labels: y_valid, self.decoder_labels: y_valid, self.time_step: current_step, self.is_training: False}) if it > 0 and valid_loss - min_loss > min_delta: patience_cnt += 1 else: patience_cnt = 0 min_loss = valid_loss if patience_cnt > patience: os.makedirs(self.model_to_use, exist_ok=True) save_path = self.saver.save(self.sess, self.model_to_use) break if verbose: print(f"Epoch {it}/{n_epochs}: Loss: {train_loss / (train_data.shape[0] // batch_size):.4f} MMD Loss: {train_mmd_loss / (train_data.shape[0] // batch_size):.4f} Valid Loss: {valid_loss:.4f} Valid MMD Loss: {valid_mmd_loss:.4f}") else: if verbose: print(f"Epoch {it}/{n_epochs}: Loss: {train_loss / (train_data.shape[0] // batch_size):.4f} MMD Loss: {train_mmd_loss / (train_data.shape[0] // batch_size)}") os.makedirs(self.model_to_use, exist_ok=True) save_path = self.saver.save(self.sess, self.model_to_use) print(f"Model saved in file: {save_path}. Training finished") PKGOHzztrvae/models/_utils.pyimport numpy as np import tensorflow as tf from keras import backend as K def compute_kernel(x, y, kernel='rbf', **kwargs): """ Computes RBF kernel between x and y. # Parameters x: Tensor Tensor with shape [batch_size, z_dim] y: Tensor Tensor with shape [batch_size, z_dim] # Returns returns the computed RBF kernel between x and y """ scales = kwargs.get("scales", []) if kernel == "rbf": x_size = K.shape(x)[0] y_size = K.shape(y)[0] dim = K.shape(x)[1] tiled_x = K.tile(K.reshape(x, K.stack([x_size, 1, dim])), K.stack([1, y_size, 1])) tiled_y = K.tile(K.reshape(y, K.stack([1, y_size, dim])), K.stack([x_size, 1, 1])) return K.exp(-K.mean(K.square(tiled_x - tiled_y), axis=2) / K.cast(dim, tf.float32)) elif kernel == 'raphy': scales = K.variable(value=np.asarray(scales)) squared_dist = K.expand_dims(squared_distance(x, y), 0) scales = K.expand_dims(K.expand_dims(scales, -1), -1) weights = K.eval(K.shape(scales)[0]) weights = K.variable(value=np.asarray(weights)) weights = K.expand_dims(K.expand_dims(weights, -1), -1) return K.sum(weights * K.exp(-squared_dist / (K.pow(scales, 2))), 0) elif kernel == "multi-scale-rbf": sigmas = [1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1, 5, 10, 15, 20, 25, 30, 35, 100, 1e3, 1e4, 1e5, 1e6] beta = 1. / (2. * (K.expand_dims(sigmas, 1))) distances = squared_distance(x, y) s = K.dot(beta, K.reshape(distances, (1, -1))) return K.reshape(tf.reduce_sum(tf.exp(-s), 0), K.shape(distances)) / len(sigmas) def squared_distance(x, y): # returns the pairwise euclidean distance r = K.expand_dims(x, axis=1) return K.sum(K.square(r - y), axis=-1) def compute_mmd(x, y, kernel, **kwargs): # [batch_size, z_dim] [batch_size, z_dim] """ Computes Maximum Mean Discrepancy(MMD) between x and y. # Parameters x: Tensor Tensor with shape [batch_size, z_dim] y: Tensor Tensor with shape [batch_size, z_dim] # Returns returns the computed MMD between x and y """ x_kernel = compute_kernel(x, x, kernel=kernel, **kwargs) y_kernel = compute_kernel(y, y, kernel=kernel, **kwargs) xy_kernel = compute_kernel(x, y, kernel=kernel, **kwargs) return K.mean(x_kernel) + K.mean(y_kernel) - 2 * K.mean(xy_kernel) def sample_z(args): """ Samples from standard Normal distribution with shape [size, z_dim] and applies re-parametrization trick. It is actually sampling from latent space distributions with N(mu, var) computed in `_encoder` function. # Parameters No parameters are needed. # Returns The computed Tensor of samples with shape [size, z_dim]. """ mu, log_var = args batch_size = K.shape(mu)[0] z_dim = K.int_shape(mu)[1] eps = K.random_normal(shape=[batch_size, z_dim]) return mu + K.exp(log_var / 2) * eps def _nan2zero(x): return tf.where(tf.is_nan(x), tf.zeros_like(x), x) def _nan2inf(x): return tf.where(tf.is_nan(x), tf.zeros_like(x) + np.inf, x) def _nelem(x): nelem = tf.reduce_sum(tf.cast(~tf.is_nan(x), tf.float32)) return tf.cast(tf.where(tf.equal(nelem, 0.), 1., nelem), x.dtype) def _reduce_mean(x): nelem = _nelem(x) x = _nan2zero(x) return tf.divide(tf.reduce_sum(x), nelem) def print_message(epoch, logs, n_epochs=10000, duration=50): if epoch % duration == 0: print(f"Epoch {epoch + 1}/{n_epochs}:") print(f" - loss: {logs['loss']:.4f} - kl_sse_loss: {logs['kl_mse_loss']:.4f}" f" - mmd_loss: {logs['mmd_loss']:.4f} - val_loss: {logs['val_loss']:.4f}" f" - val_kl_sse_loss: {logs['val_kl_mse_loss']:.4f} - val_mmd_loss: {logs['val_mmd_loss']:.4f}") PKGO*OOtrvae/models/_vae.pyimport logging import os import keras import numpy as np from keras import backend as K from keras.callbacks import CSVLogger, History, EarlyStopping from keras.layers import Dense, BatchNormalization, Dropout, Input, Lambda, Reshape, Conv1D, \ MaxPooling1D, Flatten, Conv2DTranspose, UpSampling2D from keras.layers.advanced_activations import LeakyReLU from keras.models import Model, load_model from keras.utils import multi_gpu_model from scipy import sparse from trvae.models._utils import sample_z log = logging.getLogger(__file__) class VAE: """ VAE Network class. This class contains the implementation of Variational Auto-encoder network. # Parameters kwargs: key: `dropout_rate`: float dropout rate key: `learning_rate`: float learning rate of optimization algorithm key: `model_path`: basestring path to save the model after training key: `alpha`: float alpha coefficient for loss. key: `beta`: float beta coefficient for loss. x_dimension: integer number of gene expression space dimensions. z_dimension: integer number of latent space dimensions. """ def __init__(self, x_dimension, z_dimension=100, **kwargs): self.x_dim = x_dimension self.z_dim = z_dimension self.lr = kwargs.get("learning_rate", 0.001) self.alpha = kwargs.get("alpha", 0.001) self.conditions = kwargs.get("condition_list") self.dr_rate = kwargs.get("dropout_rate", 0.2) self.model_to_use = kwargs.get("model_path", "./") self.n_gpus = kwargs.get("gpus", 1) self.arch_style = kwargs.get("arch_style", 1) self.x = Input(shape=(self.x_dim,), name="data") self.z = Input(shape=(self.z_dim,), name="latent_data") self.init_w = keras.initializers.glorot_normal() self._create_network() self._loss_function(compile_gpu_model=True) self.encoder_model.summary() self.decoder_model.summary() self.vae_model.summary() def _encoder(self, x, name="encoder"): """ Constructs the encoder sub-network of C-VAE. This function implements the encoder part of Variational Auto-encoder. It will transform primary data in the `n_vars` dimension-space to the `z_dimension` latent space. # Parameters No parameters are needed. # Returns mean: Tensor A dense layer consists of means of gaussian distributions of latent space dimensions. log_var: Tensor A dense layer consists of log transformed variances of gaussian distributions of latent space dimensions. """ if self.arch_style == 1: h = Reshape((self.x_dim, 1))(x) h = Conv1D(32, kernel_size=256, activation='relu', padding='valid')(h) # h = Conv1D(32, kernel_size=256, activation='relu', padding='same')(h) h = MaxPooling1D(pool_size=100)(h) h = Flatten()(h) h = Dense(256, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) elif self.arch_style == 2: h = Dense(4096, kernel_initializer=self.init_w, use_bias=False)(x) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(512, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.z_dim, kernel_initializer=self.init_w)(h) mean = Dense(self.z_dim, kernel_initializer=self.init_w)(h) log_var = Dense(self.z_dim, kernel_initializer=self.init_w)(h) z = Lambda(sample_z, output_shape=(self.z_dim,))([mean, log_var]) model = Model(inputs=x, outputs=[mean, log_var, z], name=name) return mean, log_var, model def _decoder(self, z, name="decoder"): """ Constructs the decoder sub-network of C-VAE. This function implements the decoder part of Variational Auto-encoder. It will transform constructed latent space to the previous space of data with n_dimensions = n_vars. # Parameters No parameters are needed. # Returns h: Tensor A Tensor for last dense layer with the shape of [n_vars, ] to reconstruct data. """ if self.arch_style == 1: h = Dense(256, kernel_initializer=self.init_w, use_bias=False)(z) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Reshape((256, 1, 1))(h) h = UpSampling2D(size=(16, 1))(h) # h = Conv2DTranspose(32, kernel_size=(256, 1), activation='relu', padding='same', # kernel_initializer='he_normal')(h) # h = Conv2DTranspose(64, kernel_size=(512, 1), activation='relu', padding='same', kernel_initializer='he_normal')(h) # h = Conv2DTranspose(256, kernel_size=(1024, 1), activation='relu', padding='same', kernel_initializer='he_normal')(h) # h = UpSampling2D(size=(2, 1))(h) # h = Conv2DTranspose(64, kernel_size=(256, 1), activation='relu', padding='same', # kernel_initializer='he_normal')(h) # h = Conv2DTranspose(256, kernel_size=(1024, 1), activation='relu', padding='same', kernel_initializer='he_normal')(h) # h = Conv2DTranspose(256, kernel_size=(1024, 1), activation='relu', padding='same')(h) # h = UpSampling2D(size=(4, 1))(h) # h = Conv2DTranspose(4, kernel_size=(3152, 1), activation='relu', padding='valid')(h) h = Conv2DTranspose(32, kernel_size=(905, 1), activation='relu', padding='valid')(h) h = Conv2DTranspose(1, kernel_size=(256, 1), activation='relu', padding='same')(h) h = Reshape((self.x_dim,))(h) if self.arch_style == 2: h = Dense(512, kernel_initializer=self.init_w, use_bias=False)(z) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(4096, kernel_initializer=self.init_w, use_bias=False)(h) h = BatchNormalization()(h) h = LeakyReLU()(h) h = Dropout(self.dr_rate)(h) h = Dense(self.x_dim, activation='relu', kernel_initializer=self.init_w, use_bias=True)(h) model = Model(inputs=z, outputs=h, name=name) return h, model def _create_network(self): """ Constructs the whole C-VAE network. It is step-by-step constructing the C-VAE network. First, It will construct the encoder part and get mu, log_var of latent space. Second, It will sample from the latent space to feed the decoder part in next step. Finally, It will reconstruct the data by constructing decoder part of C-VAE. # Parameters No parameters are needed. # Returns Nothing will be returned. """ self.mu, self.log_var, self.encoder_model = self._encoder(self.x, name="encoder") self.x_hat, self.decoder_model = self._decoder(self.z, name="decoder") decoder_outputs = self.decoder_model(self.encoder_model(self.x)[2]) reconstruction_output = Lambda(lambda x: x, name="kl_reconstruction")(decoder_outputs) self.vae_model = Model(inputs=self.x, outputs=reconstruction_output, name="vae") if self.n_gpus > 1: self.gpu_vae_model = multi_gpu_model(self.vae_model, gpus=self.n_gpus) self.gpu_encoder_model = multi_gpu_model(self.encoder_model, gpus=self.n_gpus) self.gpu_decoder_model = multi_gpu_model(self.decoder_model, gpus=self.n_gpus) else: self.gpu_vae_model = self.vae_model self.gpu_encoder_model = self.encoder_model self.gpu_decoder_model = self.decoder_model def _loss_function(self, compile_gpu_model=True): """ Defines the loss function of C-VAE network after constructing the whole network. This will define the KL Divergence and Reconstruction loss for C-VAE and also defines the Optimization algorithm for network. The C-VAE Loss will be weighted sum of reconstruction loss and KL Divergence loss. # Parameters No parameters are needed. # Returns Nothing will be returned. """ def kl_recon_loss(y_true, y_pred): kl_loss = 0.5 * K.mean(K.exp(self.log_var) + K.square(self.mu) - 1. - self.log_var, 1) # recon_loss = 0.5 * K.sum(K.binary_crossentropy(y_true, y_pred), axis=-1) recon_loss = 0.5 * K.sum(K.square((y_true - y_pred)), axis=1) # output_shape = K.shape(y_pred) # # def get_column(tensor, col): # with tf.variable_scope("gather_cols", reuse=tf.AUTO_REUSE): # tensor = K.tf.convert_to_tensor(tensor, name='tensor') # col = K.tf.convert_to_tensor(col, name='column') # # return tf.transpose(tf.nn.embedding_lookup(tf.transpose(tensor), col)) # # def body(i): # y_true_i = get_column(y_true, i) # y_pred_i = get_column(y_pred, i) # bc_loss = K.binary_crossentropy(y_true_i, y_pred_i) # print(bc_loss) # return bc_loss # recon_loss = K.sum(K.map_fn(body, K.arange(0, self.x_dim)), axis=0) # recon_loss = K.cast(recon_loss, dtype='float32') # print(recon_loss) # print(kl_loss) return recon_loss + self.alpha * kl_loss self.vae_optimizer = keras.optimizers.Adam(lr=self.lr) if compile_gpu_model: self.gpu_vae_model.compile(optimizer=self.vae_optimizer, loss=kl_recon_loss, metrics={self.vae_model.outputs[0].name: kl_recon_loss}) else: self.vae_model.compile(optimizer=self.vae_optimizer, loss=kl_recon_loss, metrics={self.vae_model.outputs[0].name: kl_recon_loss}) def to_latent(self, data): """ Map `data` in to the latent space. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ latent = self.encoder_model.predict(data)[2] return latent def to_mmd_layer(self, model, data, encoder_labels, feed_fake=False): """ Map `data` in to the pn layer after latent layer. This function will feed data in encoder part of C-VAE and compute the latent space coordinates for each sample in data. # Parameters data: `~anndata.AnnData` Annotated data matrix to be mapped to latent space. `data.X` has to be in shape [n_obs, n_vars]. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns latent: numpy nd-array returns array containing latent space encoding of 'data' """ if feed_fake: decoder_labels = np.ones(shape=encoder_labels.shape) else: decoder_labels = encoder_labels mmd_latent = model.cvae_model.predict([data, encoder_labels, decoder_labels])[1] return mmd_latent def _reconstruct(self, data, use_data=False): """ Map back the latent space encoding via the decoder. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in latent space or primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. use_data: bool this flag determines whether the `data` is already in latent space or not. if `True`: The `data` is in latent space (`data.X` is in shape [n_obs, z_dim]). if `False`: The `data` is not in latent space (`data.X` is in shape [n_obs, n_vars]). # Returns rec_data: 'numpy nd-array' returns 'numpy nd-array` containing reconstructed 'data' in shape [n_obs, n_vars]. """ if use_data: latent = data else: latent = self.to_latent(data) rec_data = self.decoder_model.predict(latent) return rec_data def predict(self, data, data_space='None'): """ Predicts the cell type provided by the user in stimulated condition. # Parameters data: `~anndata.AnnData` Annotated data matrix whether in primary space. labels: numpy nd-array `numpy nd-array` of labels to be fed as CVAE's condition array. # Returns stim_pred: numpy nd-array `numpy nd-array` of predicted cells in primary space. # Example ```python import scanpy as sc import scgen train_data = sc.read("train_kang.h5ad") validation_data = sc.read("./data/validation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) prediction = network.predict('CD4T', obs_key={"cell_type": ["CD8T", "NK"]}) ``` """ if sparse.issparse(data.X): if data_space == 'latent': stim_pred = self._reconstruct(data.X.A, use_data=True) else: stim_pred = self._reconstruct(data.X.A) else: if data_space == 'latent': stim_pred = self._reconstruct(data.X, use_data=True) else: stim_pred = self._reconstruct(data.X) return stim_pred[0] def restore_model(self): """ restores model weights from `model_to_use`. # Parameters No parameters are needed. # Returns Nothing will be returned. # Example ```python import scanpy as sc import scgen train_data = sc.read("./data/train_kang.h5ad") validation_data = sc.read("./data/valiation.h5ad") network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.restore_model() ``` """ self.vae_model = load_model(os.path.join(self.model_to_use, 'vae.h5'), compile=False) self.encoder_model = load_model(os.path.join(self.model_to_use, 'encoder.h5'), compile=False) self.decoder_model = load_model(os.path.join(self.model_to_use, 'decoder.h5'), compile=False) self._loss_function() def train(self, train_data, use_validation=False, valid_data=None, n_epochs=25, batch_size=32, early_stop_limit=20, threshold=0.0025, initial_run=True, shuffle=True, verbose=2, save=True): """ Trains the network `n_epochs` times with given `train_data` and validates the model using validation_data if it was given in the constructor function. This function is using `early stopping` technique to prevent overfitting. # Parameters n_epochs: int number of epochs to iterate and optimize network weights early_stop_limit: int number of consecutive epochs in which network loss is not going lower. After this limit, the network will stop training. threshold: float Threshold for difference between consecutive validation loss values if the difference is upper than this `threshold`, this epoch will not considered as an epoch in early stopping. full_training: bool if `True`: Network will be trained with all batches of data in each epoch. if `False`: Network will be trained with a random batch of data in each epoch. initial_run: bool if `True`: The network will initiate training and log some useful initial messages. if `False`: Network will resume the training using `restore_model` function in order to restore last model which has been trained with some training dataset. # Returns Nothing will be returned # Example ```python import scanpy as sc import scgen train_data = sc.read(train_katrain_kang.h5ad >>> validation_data = sc.read(valid_kang.h5ad) network = scgen.CVAE(train_data=train_data, use_validation=True, validation_data=validation_data, model_path="./saved_models/", conditions={"ctrl": "control", "stim": "stimulated"}) network.train(n_epochs=20) ``` """ if initial_run: log.info("----Training----") if use_validation and valid_data is None: raise Exception("valid_data is None but use_validation is True.") callbacks = [ History(), EarlyStopping(patience=early_stop_limit, monitor='val_loss', min_delta=threshold), CSVLogger(filename="./csv_logger.log") ] if sparse.issparse(train_data.X): train_data.X = train_data.X.A x = train_data.X y = train_data.X if use_validation: if sparse.issparse(valid_data.X): valid_data.X = valid_data.X.A x_valid = valid_data.X y_valid = valid_data.X histories = self.gpu_vae_model.fit( x=x, y=y, epochs=n_epochs, batch_size=batch_size, validation_data=(x_valid, y_valid), shuffle=shuffle, callbacks=callbacks, verbose=verbose) else: histories = self.gpu_vae_model.fit( x=x, y=y, epochs=n_epochs, batch_size=batch_size, shuffle=shuffle, callbacks=callbacks, verbose=verbose) if save: os.makedirs(self.model_to_use, exist_ok=True) self.vae_model.save(os.path.join(self.model_to_use, "vae.h5"), overwrite=True) self.encoder_model.save(os.path.join(self.model_to_use, "encoder.h5"), overwrite=True) self.decoder_model.save(os.path.join(self.model_to_use, "decoder.h5"), overwrite=True) log.info(f"Model saved in file: {self.model_to_use}. Training finished") return histories PKGO܏ELLtrvae/models/layers.pyfrom keras import backend as K from keras.engine.topology import Layer class SliceLayer(Layer): def __init__(self, index=0, **kwargs): self.index = index super().__init__(**kwargs) def build(self, input_shape): if not isinstance(input_shape, list): raise ValueError('Input should be a list') super().build(input_shape) def call(self, x, **kwargs): assert isinstance(x, list), 'SliceLayer input is not a list' return x[self.index] def compute_output_shape(self, input_shape): return input_shape[self.index] class ColwiseMultLayer(Layer): def __init__(self, **kwargs): super().__init__(**kwargs) def build(self, input_shape): if not isinstance(input_shape, list): raise ValueError('Input should be a list') super().build(input_shape) def call(self, x, **kwargs): assert isinstance(x, list), 'SliceLayer input is not a list' return x[0] * K.reshape(x[1], (-1, 1)) def compute_output_shape(self, input_shape): return input_shape[0] PK!HMuSatrvae-1.0.dist-info/WHEEL HM K-*ϳR03rOK-J,/RH,szd&Y)r$[)T&UD"PK!Hͽtrvae-1.0.dist-info/METADATAUMs6Wl5))\2M<'$Q ,ȿ I{]}vIdr+fl0rDhKwe4-ҩ,i E+5\wdMHO`be5dR)yfB5*|gc`ڸ c 3YʚΕ tMHgרhzH l/P|D;eلS l)ZPA0V|{U*12zy:s t NEǙwsG- 윞ގ)>]n9w8M +YlZڮ>lql66| dI7`Nt =\jUI0žx3H旛?: $Ikg2OrWm-}~$kbHc!USԪj#n3wtN=R^?mN.:17p֒3Ec ; gm-Ơq!uABhY%Az9(PâVy k=Z̑St#CョFGڏQXN|:r/Yѕm|T4y9$Gc]#qdkD!_"p`pRqX#%qbc]Z,SbV<&TjSHц6X nwT6՛*ð# ޻W`8p G/qd Z8E֭#,}xZǪ'bW  {Jg8_yѱ4 } =- 7>lƽV` LjWBϡRp!tQwrԊRrmڐ ʲ+2F; ='(/V1"8!Lq6Nss-3V t1qy3b5+fb 8PK!HKUYtrvae-1.0.dist-info/RECORDuɒJ}? V3@AR@6  O튪wyX|$}h 樞2Y鉟ٝK؇|3>(G#P Niy:6:Km:wf=<}aM"_f![+֩z8֢#kqM/\PHNrt;6{N0_W[`! Yhl]N|Gh싚`Q&TXՊ’6l*:AZ+Qu#tw;w.NonjN r S,S5K B\TV$߼ k_C,P& o Kn`3whGr{˅(f^m϶=J+"<uWб~H~jB07.>QyL+cR'm2pa(uNxج7vS!8ܛXw㘼-,Au]O]7`mQK~/pO##4'l"1vxQw2u>1Ik`,pU|Yj&)t 겥sGCsdOF1{ZNn{%E J;-un VVtrvae/models/_trvae_multi_tf.pyPKGOHzzMtrvae/models/_utils.pyPKGO*OOtrvae/models/_vae.pyPKGO܏ELLGtrvae/models/layers.pyPK!HMuSaxLtrvae-1.0.dist-info/WHEELPK!HͽMtrvae-1.0.dist-info/METADATAPK!HKUY'Qtrvae-1.0.dist-info/RECORDPK%U