--- a
+++ b/util.py
@@ -0,0 +1,307 @@
+from __future__ import division
+import scipy.sparse as sp 
+import scipy.io
+import numpy as np
+import copy
+from scipy import pi
+import sys
+import pandas as pd
+from os.path import join
+import gzip
+from scipy.io import mmwrite,mmread
+from sklearn.decomposition import PCA,TruncatedSVD
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.metrics import pairwise_distances
+from sklearn.cluster import KMeans
+from sklearn.metrics.cluster import normalized_mutual_info_score, adjusted_rand_score
+from sklearn.metrics.cluster import homogeneity_score, adjusted_mutual_info_score
+from sklearn.preprocessing import MinMaxScaler,MaxAbsScaler
+import metric
+import matplotlib
+matplotlib.use('agg')
+import matplotlib.pyplot as plt
+import seaborn as sns
+sns.set_style("whitegrid", {'axes.grid' : False})
+matplotlib.rc('xtick', labelsize=20) 
+matplotlib.rc('ytick', labelsize=20) 
+matplotlib.rcParams.update({'font.size': 22})
+
+def compute_inertia(a, X, norm=True):
+    if norm:
+        W = [np.sum(pairwise_distances(X[a == c, :]))/(2.*sum(a == c)) for c in np.unique(a)]
+        return np.sum(W)
+    else:
+        W = [np.mean(pairwise_distances(X[a == c, :])) for c in np.unique(a)]
+        return np.mean(W)
+
+#gap statistic 
+def compute_gap(clustering, data, k_max=10, n_references=100):
+    if len(data.shape) == 1:
+        data = data.reshape(-1, 1)
+    reference_inertia = []
+    for k in range(2, k_max+1):
+        reference = np.random.rand(*data.shape)
+        mins = np.min(data,axis=0)
+        maxs = np.max(data,axis=0)
+        reference = reference*(maxs-mins)+mins
+        local_inertia = []
+        for _ in range(n_references):
+            clustering.n_clusters = k
+            assignments = clustering.fit_predict(reference)
+            local_inertia.append(compute_inertia(assignments, reference))
+        reference_inertia.append(np.mean(local_inertia))
+    
+    ondata_inertia = []
+    for k in range(2, k_max+1):
+        clustering.n_clusters = k
+        assignments = clustering.fit_predict(data)
+        ondata_inertia.append(compute_inertia(assignments, data))
+        
+    gap = np.log(reference_inertia)-np.log(ondata_inertia)
+    return gap, np.log(reference_inertia), np.log(ondata_inertia)
+
+
+
+#scATAC data
+class scATAC_Sampler(object):
+    def __init__(self,name,dim=20,low=0.03,has_label=True):
+        self.name = name
+        self.dim = dim
+        self.has_label = has_label
+        X = pd.read_csv('datasets/%s/sc_mat.txt'%name,sep='\t',header=0,index_col=[0]).values
+        if has_label:
+            labels = [item.strip() for item in open('datasets/%s/label.txt'%name).readlines()]
+            uniq_labels = list(np.unique(labels))
+            Y = np.array([uniq_labels.index(item) for item in labels])
+            #X,Y = self.filter_cells(X,Y,min_peaks=10)
+            self.Y = Y
+        X = self.filter_peaks(X,low)
+        #TF-IDF transformation
+        nfreqs = 1.0 * X / np.tile(np.sum(X,axis=0), (X.shape[0],1))
+        X  = nfreqs * np.tile(np.log(1 + 1.0 * X.shape[1] / np.sum(X,axis=1)).reshape(-1,1), (1,X.shape[1]))
+        X = X.T #(cells, peaks)
+        X = MinMaxScaler().fit_transform(X)
+        #PCA transformation
+        pca = PCA(n_components=dim, random_state=3456).fit(X)
+        X = pca.transform(X)
+        self.X = X
+        self.total_size = self.X.shape[0]
+
+
+    def filter_peaks(self,X,ratio):
+        ind = np.sum(X>0,axis=1) > X.shape[1]*ratio
+        return X[ind,:]
+
+    def filter_cells(self,X,Y,min_peaks):
+        ind = np.sum(X>0,axis=0) > min_peaks
+        return X[:,ind], Y[ind]
+
+    def correlation(self,X,Y,heatmap=False):
+        nb_classes = len(set(Y))
+        print(nb_classes)
+        km = KMeans(n_clusters=nb_classes,random_state=0).fit(X)
+        label_kmeans = km.labels_
+        purity = metric.compute_purity(label_kmeans, Y)
+        nmi = normalized_mutual_info_score(Y, label_kmeans)
+        ari = adjusted_rand_score(Y, label_kmeans)
+        homogeneity = homogeneity_score(Y, label_kmeans)
+        ami = adjusted_mutual_info_score(Y, label_kmeans)
+        print('NMI = {}, ARI = {}, Purity = {},AMI = {}, Homogeneity = {}'.format(nmi,ari,purity,ami,homogeneity))
+ 
+    def train(self, batch_size):
+        indx = np.random.randint(low = 0, high = self.total_size, size = batch_size)
+
+        if self.has_label:
+            return self.X[indx, :], self.Y[indx]
+        else:
+            return self.X[indx, :]
+
+    def load_all(self):
+        if self.has_label:
+            return self.X, self.Y
+        else:
+            return self.X 
+
+
+#load data from 10x Genomic paired ARC technology
+class ARC_Sampler(object):
+    def __init__(self,name='PBMC10k',n_components=50,scale=10000,filter_feat=True,filter_cell=False,random_seed=1234,mode=1, \
+        min_rna_c=0,max_rna_c=None,min_atac_c=0,max_atac_c=None):
+        #c:cell, g:gene, l:locus
+        self.name = name
+        self.mode = mode
+        self.min_rna_c = min_rna_c
+        self.max_rna_c = max_rna_c
+        self.min_atac_c = min_atac_c
+        self.max_atac_c = max_atac_c
+
+        self.rna_mat, self.atac_mat, self.genes, self.peaks = self.load_data(filter_feat,filter_cell)
+
+        self.rna_mat = self.rna_mat*scale/self.rna_mat.sum(1)
+        self.rna_mat = np.log10(self.rna_mat+1)
+        self.atac_mat = self.atac_mat*scale/self.atac_mat.sum(1)
+        self.atac_mat = np.log10(self.atac_mat+1)
+
+        self.rna_reducer = PCA(n_components=n_components, random_state=random_seed)
+        self.rna_reducer.fit(self.rna_mat)
+        self.pca_rna_mat = self.rna_reducer.transform(self.rna_mat)
+
+        self.atac_reducer = PCA(n_components=n_components, random_state=random_seed)
+        self.atac_reducer.fit(self.atac_mat)
+        self.pca_atac_mat = self.atac_reducer.transform(self.atac_mat)
+
+    def load_data(self,filter_feat,filter_cell):
+        mtx_file = 'datasets/%s/matrix.mtx'%self.name
+        feat_file = 'datasets/%s/features.tsv'%self.name
+        barcode_file = 'datasets/%s/barcodes.tsv'%self.name
+        combined_mat = mmread(mtx_file).T.tocsr() #(cells, genes+peaks)
+        cells = [item.strip() for item in open(barcode_file).readlines()] 
+        genes = [item.split('\t')[1] for item in open(feat_file).readlines() if item.split('\t')[2]=="Gene Expression"]
+        peaks = [item.split('\t')[1] for item in open(feat_file).readlines() if item.split('\t')[2]=="Peaks"]
+        assert len(genes)+len(peaks) == combined_mat.shape[1]
+        rna_mat = combined_mat[:,:len(genes)]
+        atac_mat = combined_mat[:,len(genes):]
+        print('scRNA-seq: ', rna_mat.shape, 'scATAC-seq: ', atac_mat.shape)
+
+        if filter_feat:
+            rna_mat, atac_mat, genes, peaks = self.filter_feats(rna_mat, atac_mat, genes, peaks)
+            print('scRNA-seq filtered: ', rna_mat.shape, 'scATAC-seq filtered: ', atac_mat.shape)
+        return rna_mat, atac_mat, genes, peaks
+
+    def filter_feats(self, rna_mat_sp, atac_mat_sp, genes, peaks):
+        #filter genes
+        gene_select = np.array((rna_mat_sp>0).sum(axis=0)).squeeze() > self.min_rna_c
+        if self.max_rna_c is not None:
+            gene_select *= np.array((rna_mat_sp>0).sum(axis=0)).squeeze() < self.max_rna_c
+        rna_mat_sp = rna_mat_sp[:,gene_select]
+        genes = np.array(genes)[gene_select]
+        #filter peaks
+        locus_select = np.array((atac_mat_sp>0).sum(axis=0)).squeeze() > self.min_atac_c
+        if self.max_atac_c is not None:
+            locus_select *= np.array((atac_mat_sp>0).sum(axis=0)).squeeze() < self.max_atac_c
+        atac_mat_sp = atac_mat_sp[:,locus_select]
+        peaks = np.array(peaks)[locus_select]
+        return rna_mat_sp, atac_mat_sp, genes, peaks
+
+    def get_batch(self,batch_size):
+        idx = np.random.randint(low = 0, high = self.atac_mat.shape[0], size = batch_size)
+        if self.mode == 1:
+            return self.pca_rna_mat[idx,:]
+        elif self.mode == 2:
+            return self.pca_atac_mat[idx,:]
+        elif self.mode == 3:
+            return np.hstack((self.pca_rna_mat, self.pca_atac_mat))[idx,:]
+        else:
+            print('Wrong mode!')
+            sys.exit()
+    
+    def load_all(self):
+        #mode: 1 only scRNA-seq, 2 only scATAC-seq, 3 both
+        if self.mode == 1:
+            return self.pca_rna_mat
+        elif self.mode == 2:
+            return self.pca_atac_mat
+        elif self.mode == 3:
+            return np.hstack((self.pca_rna_mat, self.pca_atac_mat))
+        else:
+            print('Wrong mode!')
+            sys.exit()
+
+
+#sample continuous (Gaussian) and discrete (Catagory) latent variables together
+class Mixture_sampler(object):
+    def __init__(self, nb_classes, N, dim, sd, scale=1):
+        self.nb_classes = nb_classes
+        self.total_size = N
+        self.dim = dim
+        self.sd = sd 
+        self.scale = scale
+        np.random.seed(1024)
+        self.X_c = self.scale*np.random.normal(0, self.sd**2, (self.total_size,self.dim))
+        #self.X_c = self.scale*np.random.uniform(-1, 1, (self.total_size,self.dim))
+        self.label_idx = np.random.randint(low = 0 , high = self.nb_classes, size = self.total_size)
+        self.X_d = np.eye(self.nb_classes)[self.label_idx]
+        self.X = np.hstack((self.X_c,self.X_d))
+    
+    def train(self,batch_size,weights=None):
+        X_batch_c = self.scale*np.random.normal(0, 1, (batch_size,self.dim))
+        #X_batch_c = self.scale*np.random.uniform(-1, 1, (batch_size,self.dim))
+        if weights is None:
+            weights = np.ones(self.nb_classes, dtype=np.float64) / float(self.nb_classes)
+        label_batch_idx =  np.random.choice(self.nb_classes, size=batch_size, replace=True, p=weights)
+        X_batch_d = np.eye(self.nb_classes)[label_batch_idx]
+        return X_batch_c, X_batch_d
+
+    def load_all(self):
+        return self.X_c, self.X_d
+
+#sample continuous (Gaussian Mixture) and discrete (Catagory) latent variables together
+class Mixture_sampler_v2(object):
+    def __init__(self, nb_classes, N, dim, weights=None,sd=0.5):
+        self.nb_classes = nb_classes
+        self.total_size = N
+        self.dim = dim
+        np.random.seed(1024)
+        if nb_classes<=dim:
+            self.mean = np.random.uniform(-5,5,size =(nb_classes, dim))
+            #self.mean = np.zeros((nb_classes,dim))
+            #self.mean[:,:nb_classes] = np.eye(nb_classes)
+        else:
+            if dim==2:
+                self.mean = np.array([(np.cos(2*np.pi*idx/float(self.nb_classes)),np.sin(2*np.pi*idx/float(self.nb_classes))) for idx in range(self.nb_classes)])
+            else:
+                self.mean = np.zeros((nb_classes,dim))
+                self.mean[:,:2] = np.array([(np.cos(2*np.pi*idx/float(self.nb_classes)),np.sin(2*np.pi*idx/float(self.nb_classes))) for idx in range(self.nb_classes)])
+        self.cov = [sd**2*np.eye(dim) for item in range(nb_classes)]
+        if weights is None:
+            weights = np.ones(self.nb_classes, dtype=np.float64) / float(self.nb_classes)
+        self.Y = np.random.choice(self.nb_classes, size=N, replace=True, p=weights)
+        self.X_c = np.array([np.random.multivariate_normal(mean=self.mean[i],cov=self.cov[i]) for i in self.Y],dtype='float64')
+        self.X_d = np.eye(self.nb_classes)[self.Y]
+        self.X = np.hstack((self.X_c,self.X_d))
+
+    def train(self, batch_size, label = False):
+        indx = np.random.randint(low = 0, high = self.total_size, size = batch_size)
+        if label:
+            return self.X_c[indx, :], self.X_d[indx, :], self.Y[indx, :]
+        else:
+            return self.X_c[indx, :], self.X_d[indx, :]
+
+    def get_batch(self,batch_size,weights=None):
+        if weights is None:
+            weights = np.ones(self.nb_classes, dtype=np.float64) / float(self.nb_classes)
+        label_batch_idx =  np.random.choice(self.nb_classes, size=batch_size, replace=True, p=weights)
+        return self.X_c[label_batch_idx, :], self.X_d[label_batch_idx, :]
+    def predict_onepoint(self,array):#return component index with max likelyhood
+        from scipy.stats import multivariate_normal
+        assert len(array) == self.dim
+        return np.argmax([multivariate_normal.pdf(array,self.mean[idx],self.cov[idx]) for idx in range(self.nb_classes)])
+
+    def predict_multipoints(self,arrays):
+        assert arrays.shape[-1] == self.dim
+        return map(self.predict_onepoint,arrays)
+    def load_all(self):
+        return self.X_c, self.X_d, self.label_idx
+
+#get a batch of data from previous 50 batches, add stochastic
+class DataPool(object):
+    def __init__(self, maxsize=50):
+        self.maxsize = maxsize
+        self.nb_batch = 0
+        self.pool = []
+
+    def __call__(self, data):
+        if self.nb_batch < self.maxsize:
+            self.pool.append(data)
+            self.nb_batch += 1
+            return data
+        if np.random.rand() > 0.5:
+            results=[]
+            for i in range(len(data)):
+                idx = int(np.random.rand()*self.maxsize)
+                results.append(copy.copy(self.pool[idx])[i])
+                self.pool[idx][i] = data[i]
+            return results
+        else:
+            return data
+