from time import time
import math, os
from sklearn import metrics
from sklearn.cluster import KMeans
from sklearn.preprocessing import OneHotEncoder
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.nn import Parameter
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from scMDC_batch import scMultiClusterBatch
import numpy as np
import collections
import h5py
import scanpy as sc
from preprocess import read_dataset, normalize
from utils import *
if __name__ == "__main__":
# setting the hyper parameters
import argparse
parser = argparse.ArgumentParser(description='train',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--n_clusters', default=27, type=int)
parser.add_argument('--cutoff', default=0.5, type=float, help='Start to train combined layer after what ratio of epoch')
parser.add_argument('--batch_size', default=256, type=int)
parser.add_argument('--data_file', default='Normalized_filtered_BMNC_GSE128639_Seurat.h5')
parser.add_argument('--maxiter', default=5000, type=int)
parser.add_argument('--pretrain_epochs', default=400, type=int)
parser.add_argument('--gamma', default=.1, type=float,
help='coefficient of clustering loss')
parser.add_argument('--tau', default=1., type=float,
help='weight of clustering loss')
parser.add_argument('--phi1', default=0.001, type=float,
help='coefficient of KL loss in pretraining stage')
parser.add_argument('--phi2', default=0.001, type=float,
help='coefficient of KL loss in clustering stage')
parser.add_argument('--update_interval', default=1, type=int)
parser.add_argument('--tol', default=0.001, type=float)
parser.add_argument('--lr', default=1., type=float)
parser.add_argument('--ae_weights', default=None)
parser.add_argument('--save_dir', default='results/')
parser.add_argument('--ae_weight_file', default='AE_weights_1.pth.tar')
parser.add_argument('--resolution', default=0.2, type=float)
parser.add_argument('--n_neighbors', default=30, type=int)
parser.add_argument('--embedding_file', action='store_true', default=False)
parser.add_argument('--prediction_file', action='store_true', default=False)
parser.add_argument('-el','--encodeLayer', nargs='+', default=[256,64,32,16])
parser.add_argument('-dl1','--decodeLayer1', nargs='+', default=[16,64,256])
parser.add_argument('-dl2','--decodeLayer2', nargs='+', default=[16,20])
parser.add_argument('--sigma1', default=2.5, type=float)
parser.add_argument('--sigma2', default=1.5, type=float)
parser.add_argument('--f1', default=1000, type=float, help='Number of mRNA after feature selection')
parser.add_argument('--f2', default=2000, type=float, help='Number of ADT/ATAC after feature selection')
parser.add_argument('--filter1', action='store_true', default=False, help='Do mRNA selection')
parser.add_argument('--filter2', action='store_true', default=False, help='Do ADT/ATAC selection')
parser.add_argument('--nbatch', default=2, type=int)
parser.add_argument('--run', default=1, type=int)
parser.add_argument('--device', default='cuda')
parser.add_argument('--no_labels', action='store_true', default=False)
args = parser.parse_args()
print(args)
data_mat = h5py.File(args.data_file)
x1 = np.array(data_mat['X1'])
x2 = np.array(data_mat['X2'])
if not args.no_labels:
y = np.array(data_mat['Y'])
b = np.array(data_mat['Batch'])
enc = OneHotEncoder()
enc.fit(b.reshape(-1, 1))
B = enc.transform(b.reshape(-1, 1)).toarray()
data_mat.close()
#Gene filter
if args.filter1:
importantGenes = geneSelection(x1, n=args.f1, plot=False)
x1 = x1[:, importantGenes]
if args.filter2:
importantGenes = geneSelection(x2, n=args.f2, plot=False)
x2 = x2[:, importantGenes]
# preprocessing scRNA-seq read counts matrix
adata1 = sc.AnnData(x1)
#adata1.obs['Group'] = y
adata1 = read_dataset(adata1,
transpose=False,
test_split=False,
copy=True)
adata1 = normalize(adata1,
size_factors=True,
normalize_input=True,
logtrans_input=True)
adata2 = sc.AnnData(x2)
#adata2.obs['Group'] = y
adata2 = read_dataset(adata2,
transpose=False,
test_split=False,
copy=True)
adata2 = normalize(adata2,
size_factors=True,
normalize_input=True,
logtrans_input=True)
input_size1 = adata1.n_vars
input_size2 = adata2.n_vars
print(args)
encodeLayer = list(map(int, args.encodeLayer))
decodeLayer1 = list(map(int, args.decodeLayer1))
decodeLayer2 = list(map(int, args.decodeLayer2))
model = scMultiClusterBatch(input_dim1=input_size1, input_dim2=input_size2, n_batch = args.nbatch, tau=args.tau,
encodeLayer=encodeLayer, decodeLayer1=decodeLayer1, decodeLayer2=decodeLayer2,
activation='elu', sigma1=args.sigma1, sigma2=args.sigma2, gamma=args.gamma,
cutoff = args.cutoff, phi1=args.phi1, phi2=args.phi2, device=args.device).to(args.device)
print(str(model))
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
t0 = time()
if args.ae_weights is None:
model.pretrain_autoencoder(X1=adata1.X, X_raw1=adata1.raw.X, sf1=adata1.obs.size_factors,
X2=adata2.X, X_raw2=adata2.raw.X, sf2=adata2.obs.size_factors, B = B, batch_size=args.batch_size,
epochs=args.pretrain_epochs, ae_weights=args.ae_weight_file)
else:
if os.path.isfile(args.ae_weights):
print("==> loading checkpoint '{}'".format(args.ae_weights))
checkpoint = torch.load(args.ae_weights)
model.load_state_dict(checkpoint['ae_state_dict'])
else:
print("==> no checkpoint found at '{}'".format(args.ae_weights))
raise ValueError
print('Pretraining time: %d seconds.' % int(time() - t0))
#get k
latent = model.encodeBatch(torch.tensor(adata1.X).to(args.device), torch.tensor(adata2.X).to(args.device), torch.tensor(B).to(args.device), batch_size=args.batch_size)
latent = latent.cpu().numpy()
if args.n_clusters == -1:
n_clusters = GetCluster(latent, res=args.resolution, n=args.n_neighbors)
else:
print("n_cluster is defined as " + str(args.n_clusters))
n_clusters = args.n_clusters
if not args.no_labels:
y_pred,_ = model.fit(X1=adata1.X, X_raw1=adata1.raw.X, sf1=adata1.obs.size_factors,
X2=adata2.X, X_raw2=adata2.raw.X, sf2=adata2.obs.size_factors, B=B, y=y,
n_clusters=n_clusters, batch_size=args.batch_size, num_epochs=args.maxiter,
update_interval=args.update_interval, tol=args.tol, lr=args.lr, save_dir=args.save_dir)
else:
y_pred,_ = model.fit(X1=adata1.X, X_raw1=adata1.raw.X, sf1=adata1.obs.size_factors,
X2=adata2.X, X_raw2=adata2.raw.X, sf2=adata2.obs.size_factors, B=B, y=None,
n_clusters=n_clusters, batch_size=args.batch_size, num_epochs=args.maxiter,
update_interval=args.update_interval, tol=args.tol, lr=args.lr, save_dir=args.save_dir)
print('Total time: %d seconds.' % int(time() - t0))
if args.prediction_file:
np.savetxt(args.save_dir + "/" + str(args.run) + "_pred.csv", y_pred, delimiter=",")
if args.embedding_file:
final_latent = model.encodeBatch(torch.tensor(adata1.X).to(args.device), torch.tensor(adata2.X).to(args.device), torch.tensor(B).to(args.device), batch_size=args.batch_size)
final_latent = final_latent.cpu().numpy()
np.savetxt(args.save_dir + "/" + str(args.run) + "_embedding.csv", final_latent, delimiter=",")
if not args.no_labels:
y_pred_ = best_map(y, y_pred)
ami = np.round(metrics.adjusted_mutual_info_score(y, y_pred), 5)
nmi = np.round(metrics.normalized_mutual_info_score(y, y_pred), 5)
ari = np.round(metrics.adjusted_rand_score(y, y_pred), 5)
print('Final: AMI= %.4f, NMI= %.4f, ARI= %.4f' % (ami, nmi, ari))
else:
print("No labels for evaluation!")
Datasets
Models
