def test_new_10x(): """ Test new 10X data format, which is a bit different than newer ones :return: """ data = Dataset10X("pbmc_1k_v2") data.subsample_genes(new_n_genes=100) assert data.X.shape[1] == 100
def main(): usage = 'solo' parser = ArgumentParser(usage, formatter_class=ArgumentDefaultsHelpFormatter) parser.add_argument(dest='model_json_file', help='json file to pass VAE parameters') parser.add_argument( dest='data_path', help= 'path to h5ad, loom or 10x directory containing cell by genes counts') parser.add_argument('-d', dest='doublet_depth', default=2., type=float, help='Depth multiplier for a doublet relative to the \ average of its constituents') parser.add_argument('-g', dest='gpu', default=True, action='store_true', help='Run on GPU') parser.add_argument('-a', dest='anndata_output', default=False, action='store_true', help='output modified anndata object with solo scores \ Only works for anndata') parser.add_argument('-o', dest='out_dir', default='solo_out') parser.add_argument('-r', dest='doublet_ratio', default=2., type=float, help='Ratio of doublets to true \ cells') parser.add_argument('-s', dest='seed', default=None, help='Path to previous solo output \ directory. Seed VAE models with previously \ trained solo model. Directory structure is assumed to \ be the same as solo output directory structure. \ should at least have a vae.pt a pickled object of \ vae weights and a latent.npy an np.ndarray of the \ latents of your cells.') parser.add_argument('-k', dest='known_doublets', help='Experimentally defined doublets tsv file. \ Should be a single column of True/False. True \ indicates the cell is a doublet. No header.', type=str) parser.add_argument('-t', dest='doublet_type', help='Please enter \ multinomial, average, or sum', default='multinomial', choices=['multinomial', 'average', 'sum']) parser.add_argument('-e', dest='expected_number_of_doublets', help='Experimentally expected number of doublets', type=int, default=None) parser.add_argument('-p', dest='plot', default=False, action='store_true', help='Plot outputs for solo') parser.add_argument('-l', dest='normal_logging', default=False, action='store_true', help='Logging level set to normal (aka not debug)') parser.add_argument('--random_size', dest='randomize_doublet_size', default=False, action='store_true', help='Sample depth multipliers from Unif(1, \ DoubletDepth) \ to provide a diversity of possible doublet depths.') args = parser.parse_args() if not args.normal_logging: scvi._settings.set_verbosity(10) model_json_file = args.model_json_file data_path = args.data_path if args.gpu and not torch.cuda.is_available(): args.gpu = torch.cuda.is_available() print('Cuda is not available, switching to cpu running!') if not os.path.isdir(args.out_dir): os.mkdir(args.out_dir) ################################################## # data # read loom/anndata data_ext = os.path.splitext(data_path)[-1] if data_ext == '.loom': scvi_data = LoomDataset(data_path) elif data_ext == '.h5ad': adata = anndata.read(data_path) if issparse(adata.X): adata.X = adata.X.todense() scvi_data = AnnDatasetFromAnnData(adata) elif os.path.isdir(data_path): scvi_data = Dataset10X(save_path=data_path, measurement_names_column=1, dense=True) cell_umi_depth = scvi_data.X.sum(axis=1) fifth, ninetyfifth = np.percentile(cell_umi_depth, [5, 95]) min_cell_umi_depth = np.min(cell_umi_depth) max_cell_umi_depth = np.max(cell_umi_depth) if fifth * 10 < ninetyfifth: print("""WARNING YOUR DATA HAS A WIDE RANGE OF CELL DEPTHS. PLEASE MANUALLY REVIEW YOUR DATA""") print( f"Min cell depth: {min_cell_umi_depth}, Max cell depth: {max_cell_umi_depth}" ) else: msg = f'{data_path} is not a recognized format.\n' msg += 'must be one of {h5ad, loom, 10x directory}' raise TypeError(msg) num_cells, num_genes = scvi_data.X.shape if args.known_doublets is not None: print('Removing known doublets for in silico doublet generation') print('Make sure known doublets are in the same order as your data') known_doublets = np.loadtxt(args.known_doublets, dtype=str) == 'True' assert len(known_doublets) == scvi_data.X.shape[0] known_doublet_data = make_gene_expression_dataset( scvi_data.X[known_doublets], scvi_data.gene_names) known_doublet_data.labels = np.ones(known_doublet_data.X.shape[0]) singlet_scvi_data = make_gene_expression_dataset( scvi_data.X[~known_doublets], scvi_data.gene_names) singlet_num_cells, _ = singlet_scvi_data.X.shape else: known_doublet_data = None singlet_num_cells = num_cells known_doublets = np.zeros(num_cells, dtype=bool) singlet_scvi_data = scvi_data singlet_scvi_data.labels = np.zeros(singlet_scvi_data.X.shape[0]) scvi_data.labels = known_doublets.astype(int) ################################################## # parameters # check for parameters if not os.path.exists(model_json_file): raise FileNotFoundError(f'{model_json_file} does not exist.') # read parameters with open(model_json_file, 'r') as model_json_open: params = json.load(model_json_open) # set VAE params vae_params = {} for par in [ 'n_hidden', 'n_latent', 'n_layers', 'dropout_rate', 'ignore_batch' ]: if par in params: vae_params[par] = params[par] vae_params['n_batch'] = 0 if params.get('ignore_batch', False) else scvi_data.n_batches # training parameters batch_size = params.get('batch_size', 128) valid_pct = params.get('valid_pct', 0.1) learning_rate = params.get('learning_rate', 1e-3) stopping_params = {'patience': params.get('patience', 10), 'threshold': 0} # protect against single example batch while num_cells % batch_size == 1: batch_size = int(np.round(1.25 * batch_size)) print('Increasing batch_size to %d to avoid single example batch.' % batch_size) ################################################## # VAE vae = VAE(n_input=singlet_scvi_data.nb_genes, n_labels=2, reconstruction_loss='nb', log_variational=True, **vae_params) if args.seed: if args.gpu: device = torch.device('cuda') vae.load_state_dict(torch.load(os.path.join(args.seed, 'vae.pt'))) vae.to(device) else: map_loc = 'cpu' vae.load_state_dict( torch.load(os.path.join(args.seed, 'vae.pt'), map_location=map_loc)) # save latent representation utrainer = \ UnsupervisedTrainer(vae, singlet_scvi_data, train_size=(1. - valid_pct), frequency=2, metrics_to_monitor=['reconstruction_error'], use_cuda=args.gpu, early_stopping_kwargs=stopping_params, batch_size=batch_size) full_posterior = utrainer.create_posterior(utrainer.model, singlet_scvi_data, indices=np.arange( len(singlet_scvi_data))) latent, _, _ = full_posterior.sequential(batch_size).get_latent() np.save(os.path.join(args.out_dir, 'latent.npy'), latent.astype('float32')) else: stopping_params['early_stopping_metric'] = 'reconstruction_error' stopping_params['save_best_state_metric'] = 'reconstruction_error' # initialize unsupervised trainer utrainer = \ UnsupervisedTrainer(vae, singlet_scvi_data, train_size=(1. - valid_pct), frequency=2, metrics_to_monitor=['reconstruction_error'], use_cuda=args.gpu, early_stopping_kwargs=stopping_params, batch_size=batch_size) utrainer.history['reconstruction_error_test_set'].append(0) # initial epoch utrainer.train(n_epochs=2000, lr=learning_rate) # drop learning rate and continue utrainer.early_stopping.wait = 0 utrainer.train(n_epochs=500, lr=0.5 * learning_rate) # save VAE torch.save(vae.state_dict(), os.path.join(args.out_dir, 'vae.pt')) # save latent representation full_posterior = utrainer.create_posterior(utrainer.model, singlet_scvi_data, indices=np.arange( len(singlet_scvi_data))) latent, _, _ = full_posterior.sequential(batch_size).get_latent() np.save(os.path.join(args.out_dir, 'latent.npy'), latent.astype('float32')) ################################################## # simulate doublets non_zero_indexes = np.where(singlet_scvi_data.X > 0) cells = non_zero_indexes[0] genes = non_zero_indexes[1] cells_ids = defaultdict(list) for cell_id, gene in zip(cells, genes): cells_ids[cell_id].append(gene) # choose doublets function type if args.doublet_type == 'average': doublet_function = create_average_doublet elif args.doublet_type == 'sum': doublet_function = create_summed_doublet else: doublet_function = create_multinomial_doublet cell_depths = singlet_scvi_data.X.sum(axis=1) num_doublets = int(args.doublet_ratio * singlet_num_cells) if known_doublet_data is not None: num_doublets -= known_doublet_data.X.shape[0] # make sure we are making a non negative amount of doublets assert num_doublets >= 0 in_silico_doublets = np.zeros((num_doublets, num_genes), dtype='float32') # for desired # doublets for di in range(num_doublets): # sample two cells i, j = np.random.choice(singlet_num_cells, size=2) # generate doublets in_silico_doublets[di, :] = \ doublet_function(singlet_scvi_data.X, i, j, doublet_depth=args.doublet_depth, cell_depths=cell_depths, cells_ids=cells_ids, randomize_doublet_size=args.randomize_doublet_size) # merge datasets # we can maybe up sample the known doublets # concatentate classifier_data = GeneExpressionDataset() classifier_data.populate_from_data( X=np.vstack([scvi_data.X, in_silico_doublets]), labels=np.hstack( [np.ravel(scvi_data.labels), np.ones(in_silico_doublets.shape[0])]), remap_attributes=False) assert (len(np.unique(classifier_data.labels.flatten())) == 2) ################################################## # classifier # model classifier = Classifier(n_input=(vae.n_latent + 1), n_hidden=params['cl_hidden'], n_layers=params['cl_layers'], n_labels=2, dropout_rate=params['dropout_rate']) # trainer stopping_params['early_stopping_metric'] = 'accuracy' stopping_params['save_best_state_metric'] = 'accuracy' strainer = ClassifierTrainer(classifier, classifier_data, train_size=(1. - valid_pct), frequency=2, metrics_to_monitor=['accuracy'], use_cuda=args.gpu, sampling_model=vae, sampling_zl=True, early_stopping_kwargs=stopping_params, batch_size=batch_size) # initial strainer.train(n_epochs=1000, lr=learning_rate) # drop learning rate and continue strainer.early_stopping.wait = 0 strainer.train(n_epochs=300, lr=0.1 * learning_rate) torch.save(classifier.state_dict(), os.path.join(args.out_dir, 'classifier.pt')) ################################################## # post-processing # use logits for predictions for better results logits_classifier = Classifier(n_input=(vae.n_latent + 1), n_hidden=params['cl_hidden'], n_layers=params['cl_layers'], n_labels=2, dropout_rate=params['dropout_rate'], logits=True) logits_classifier.load_state_dict(classifier.state_dict()) # using logits leads to better performance in for ranking logits_strainer = ClassifierTrainer(logits_classifier, classifier_data, train_size=(1. - valid_pct), frequency=2, metrics_to_monitor=['accuracy'], use_cuda=args.gpu, sampling_model=vae, sampling_zl=True, early_stopping_kwargs=stopping_params, batch_size=batch_size) # models evaluation mode vae.eval() classifier.eval() logits_classifier.eval() print('Train accuracy: %.4f' % strainer.train_set.accuracy()) print('Test accuracy: %.4f' % strainer.test_set.accuracy()) # compute predictions manually # output logits train_y, train_score = strainer.train_set.compute_predictions(soft=True) test_y, test_score = strainer.test_set.compute_predictions(soft=True) # train_y == true label # train_score[:, 0] == singlet score; train_score[:, 1] == doublet score train_score = train_score[:, 1] train_y = train_y.astype('bool') test_score = test_score[:, 1] test_y = test_y.astype('bool') train_auroc = roc_auc_score(train_y, train_score) test_auroc = roc_auc_score(test_y, test_score) print('Train AUROC: %.4f' % train_auroc) print('Test AUROC: %.4f' % test_auroc) train_fpr, train_tpr, train_t = roc_curve(train_y, train_score) test_fpr, test_tpr, test_t = roc_curve(test_y, test_score) train_t = np.minimum(train_t, 1 + 1e-9) test_t = np.minimum(test_t, 1 + 1e-9) train_acc = np.zeros(len(train_t)) for i in range(len(train_t)): train_acc[i] = np.mean(train_y == (train_score > train_t[i])) test_acc = np.zeros(len(test_t)) for i in range(len(test_t)): test_acc[i] = np.mean(test_y == (test_score > test_t[i])) # write predictions # softmax predictions order_y, order_score = strainer.compute_predictions(soft=True) _, order_pred = strainer.compute_predictions() doublet_score = order_score[:, 1] np.save(os.path.join(args.out_dir, 'no_updates_softmax_scores.npy'), doublet_score[:num_cells]) np.save(os.path.join(args.out_dir, 'no_updates_softmax_scores_sim.npy'), doublet_score[num_cells:]) # logit predictions logit_y, logit_score = logits_strainer.compute_predictions(soft=True) logit_doublet_score = logit_score[:, 1] np.save(os.path.join(args.out_dir, 'logit_scores.npy'), logit_doublet_score[:num_cells]) np.save(os.path.join(args.out_dir, 'logit_scores_sim.npy'), logit_doublet_score[num_cells:]) # update threshold as a function of Solo's estimate of the number of # doublets # essentially a log odds update # TODO put in a function diff = np.inf counter_update = 0 solo_scores = doublet_score[:num_cells] logit_scores = logit_doublet_score[:num_cells] d_s = (args.doublet_ratio / (args.doublet_ratio + 1)) while (diff > .01) | (counter_update < 5): # calculate log odss calibration for logits d_o = np.mean(solo_scores) c = np.log(d_o / (1 - d_o)) - np.log(d_s / (1 - d_s)) # update soloe scores solo_scores = 1 / (1 + np.exp(-(logit_scores + c))) # update while conditions diff = np.abs(d_o - np.mean(solo_scores)) counter_update += 1 np.save(os.path.join(args.out_dir, 'softmax_scores.npy'), solo_scores) if args.expected_number_of_doublets is not None: k = len(solo_scores) - args.expected_number_of_doublets if args.expected_number_of_doublets / len(solo_scores) > .5: print('''Make sure you actually expect more than half your cells to be doublets. If not change your -e parameter value''') assert k > 0 idx = np.argpartition(solo_scores, k) threshold = np.max(solo_scores[idx[:k]]) is_solo_doublet = solo_scores > threshold else: is_solo_doublet = solo_scores > .5 is_doublet = known_doublets new_doublets_idx = np.where(~(is_doublet) & is_solo_doublet[:num_cells])[0] is_doublet[new_doublets_idx] = True np.save(os.path.join(args.out_dir, 'is_doublet.npy'), is_doublet[:num_cells]) np.save(os.path.join(args.out_dir, 'is_doublet_sim.npy'), is_doublet[num_cells:]) np.save(os.path.join(args.out_dir, 'preds.npy'), order_pred[:num_cells]) np.save(os.path.join(args.out_dir, 'preds_sim.npy'), order_pred[num_cells:]) smoothed_preds = knn_smooth_pred_class(X=latent, pred_class=is_doublet[:num_cells]) np.save(os.path.join(args.out_dir, 'smoothed_preds.npy'), smoothed_preds) if args.anndata_output and data_ext == '.h5ad': adata.obs['is_doublet'] = is_doublet[:num_cells] adata.obs['logit_scores'] = logit_doublet_score[:num_cells] adata.obs['softmax_scores'] = doublet_score[:num_cells] adata.write(os.path.join(args.out_dir, "soloed.h5ad")) if args.plot: import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import seaborn as sns # plot ROC plt.figure() plt.plot(train_fpr, train_tpr, label='Train') plt.plot(test_fpr, test_tpr, label='Test') plt.gca().set_xlabel('False positive rate') plt.gca().set_ylabel('True positive rate') plt.legend() plt.savefig(os.path.join(args.out_dir, 'roc.pdf')) plt.close() # plot accuracy plt.figure() plt.plot(train_t, train_acc, label='Train') plt.plot(test_t, test_acc, label='Test') plt.axvline(0.5, color='black', linestyle='--') plt.gca().set_xlabel('Threshold') plt.gca().set_ylabel('Accuracy') plt.legend() plt.savefig(os.path.join(args.out_dir, 'accuracy.pdf')) plt.close() # plot distributions plt.figure() sns.distplot(test_score[test_y], label='Simulated') sns.distplot(test_score[~test_y], label='Observed') plt.legend() plt.savefig(os.path.join(args.out_dir, 'train_v_test_dist.pdf')) plt.close() plt.figure() sns.distplot(doublet_score[:num_cells], label='Observed') plt.legend() plt.savefig(os.path.join(args.out_dir, 'real_cells_dist.pdf')) plt.close() scvi_umap = umap.UMAP(n_neighbors=16).fit_transform(latent) fig, ax = plt.subplots(1, 1, figsize=(10, 10)) ax.scatter(scvi_umap[:, 0], scvi_umap[:, 1], c=doublet_score[:num_cells], s=8, cmap="GnBu") ax.set_xlabel("UMAP 1") ax.set_ylabel("UMAP 2") ax.set_xticks([], []) ax.set_yticks([], []) fig.savefig(os.path.join(args.out_dir, 'umap_solo_scores.pdf'))
def test_differential_expression(save_path): dataset = CortexDataset(save_path=save_path) n_cells = len(dataset) all_indices = np.arange(n_cells) vae = VAE(dataset.nb_genes, dataset.n_batches) trainer = UnsupervisedTrainer(vae, dataset, train_size=0.5, use_cuda=use_cuda) trainer.train(n_epochs=2) post = trainer.create_posterior(vae, dataset, shuffle=False, indices=all_indices) with tempfile.TemporaryDirectory() as temp_dir: posterior_save_path = os.path.join(temp_dir, "posterior_data") post.save_posterior(posterior_save_path) new_vae = VAE(dataset.nb_genes, dataset.n_batches) new_post = load_posterior(posterior_save_path, model=new_vae, use_cuda=False) assert np.array_equal(new_post.indices, post.indices) assert np.array_equal(new_post.gene_dataset.X, post.gene_dataset.X) # Sample scale example px_scales = post.scale_sampler( n_samples_per_cell=4, n_samples=None, selection=all_indices )["scale"] assert ( px_scales.shape[1] == dataset.nb_genes ), "posterior scales should have shape (n_samples, n_genes)" # Differential expression different models idx_1 = [1, 2, 3] idx_2 = [4, 5, 6, 7] de_dataframe = post.differential_expression_score( idx1=idx_1, idx2=idx_2, n_samples=10, mode="vanilla", use_permutation=True, M_permutation=100, ) de_dataframe = post.differential_expression_score( idx1=idx_1, idx2=idx_2, n_samples=10, mode="change", use_permutation=True, M_permutation=100, cred_interval_lvls=[0.5, 0.95], ) print(de_dataframe.keys()) assert ( de_dataframe["confidence_interval_0.5_min"] <= de_dataframe["confidence_interval_0.5_max"] ).all() assert ( de_dataframe["confidence_interval_0.95_min"] <= de_dataframe["confidence_interval_0.95_max"] ).all() # DE estimation example de_probabilities = de_dataframe.loc[:, "proba_de"] assert ((0.0 <= de_probabilities) & (de_probabilities <= 1.0)).all() # Test totalVI DE sp = os.path.join(save_path, "10X") dataset = Dataset10X(dataset_name="pbmc_10k_protein_v3", save_path=sp) n_cells = len(dataset) all_indices = np.arange(n_cells) vae = TOTALVI( dataset.nb_genes, len(dataset.protein_names), n_batch=dataset.n_batches ) trainer = TotalTrainer( vae, dataset, train_size=0.5, use_cuda=use_cuda, early_stopping_kwargs=None ) trainer.train(n_epochs=2) post = trainer.create_posterior( vae, dataset, shuffle=False, indices=all_indices, type_class=TotalPosterior ) # Differential expression different models idx_1 = [1, 2, 3] idx_2 = [4, 5, 6, 7] de_dataframe = post.differential_expression_score( idx1=idx_1, idx2=idx_2, n_samples=10, mode="vanilla", use_permutation=True, M_permutation=100, ) de_dataframe = post.differential_expression_score( idx1=idx_1, idx2=idx_2, n_samples=10, mode="change", use_permutation=True, M_permutation=100, )
import scanpy as sc import pandas as pd import numpy as np import doubletdetection import sys sys.path.append("./utils/") from utils import seurat_v3_highly_variable_genes # Load data save_path = "data/raw_data/" dataset = Dataset10X( dataset_name="5k_pbmc_protein_v3", save_path=save_path, measurement_names_column=1, dense=True, ) # Filter control proteins non_control_proteins = [] for i, p in enumerate(dataset.protein_names): if not p.startswith("IgG"): non_control_proteins.append(i) else: print(p) dataset.protein_expression = dataset.protein_expression[:, non_control_proteins] dataset.protein_names = dataset.protein_names[non_control_proteins] # Make anndata object
import os import sys import numpy as np import pandas as pd import matplotlib.pyplot as plt from scvi.dataset import Dataset10X from scvi.models import VAE from scvi.inference import UnsupervisedTrainer import torch # load data data_dir, out_file = sys.argv[1], sys.argv[2] gene_dataset = Dataset10X(save_path=os.path.expandusr(data_dir), measurement_names_column=1) # set variables save_path = '.' n_epochs = 400 lr = 1e-3 use_batches = False use_cuda = True # define VAE vae = VAE(gene_dataset.nb_genes, n_batch=gene_dataset.n_batches * use_batches) trainer = UnsupervisedTrainer( vae, gene_dataset, train_size=0.75, use_cuda=use_cuda, frequency=5, )