--- a
+++ b/src/scotv2.py
@@ -0,0 +1,329 @@
+"""
+Author: Pinar Demetci
+Principal Investigator: Ritambhara Singh, Ph.D. from Brown University
+08 August 2021
+Updated: 23 February 2023
+SCOTv2 algorithm: Single Cell alignment using Optimal Transport version 2
+Correspondence: pinar_demetci@brown.edu, ritambhara@brown.edu
+"""
+
+### Import python packages we depend on:
+import numpy as np
+import torch
+import ot
+import scipy
+# For computing graph distances:
+from scipy.sparse.csgraph import dijkstra
+from scipy.sparse import csr_matrix
+from sklearn.neighbors import kneighbors_graph
+
+# For pre-processing, normalization
+from sklearn.preprocessing import StandardScaler, normalize
+
+
+class SCOTv2(object):
+	"""
+	SCOT algorithm for unsupervised alignment of single-cell multi-omic data.
+	https://www.biorxiv.org/content/10.1101/2020.04.28.066787v2 (original preprint)
+	https://www.liebertpub.com/doi/full/10.1089/cmb.2021.0446 (Journal of Computational Biology publication through RECOMB 2021 conference)
+
+	Input: domain1, domain2 in form of numpy arrays/matrices, where the rows correspond to samples and columns correspond to features.
+	Returns: aligned domain 1, aligned domain 2 in form of numpy arrays/matrices projected on domain 1
+
+	Example use:
+	# Given two numpy matrices, domain1 and domain2, where the rows are cells and columns are different genomic features:
+	scot= SCOT(domain1, domain2)
+	aligned_domain1, aligned_domain2 = scot.align(k=20, e=1e-3)
+
+	#If you can't pick the parameters k and e, you can try out our unsupervised self-tuning heuristic by running:
+	scot= SCOT(domain1, domain2)
+	aligned_domain1, aligned_domain2 = scot.align(selfTune=True)
+
+	Required parameters:
+	- k: Number of neighbors to be used when constructing kNN graphs. Default= min(min(n_1, n_2), 50), where n_i, for i=1,2 corresponds to the number of samples in the i^th domain.
+	- e: Regularization constant for the entropic regularization term in entropic Gromov-Wasserstein optimal transport formulation. Default= 1e-3 
+   
+	Optional parameters:
+
+	- normalize= Determines whether to normalize input data ahead of alignment. True or False (boolean parameter). Default = True.
+	- norm= Determines what sort of normalization to run, "l2", "l1", "max", "zscore". Default="l2" 
+	- mode: "connectivity" or "distance". Determines whether to use a connectivity graph (adjacency matrix of 1s/0s based on whether nodes are connected) or a distance graph (adjacency matrix entries weighted by distances between nodes). Default="connectivity"  
+	- metric: Sets the metric to use while constructing nearest neighbor graphs. some possible choices are "correlation", "minkowski".  "correlation" is Pearson's correlation and "minkowski" is equivalent to Euclidean distance in its default form (). Default= "correlation". 
+	- verbose: Prints loss while optimizing the optimal transport formulation. Default=True
+	- XontoY: Determines the direction of barycentric projection. True or False (boolean parameter). If True, projects domain1 onto domain2. If False, projects domain2 onto domain1. Default=True.
+
+	Note: If you want to specify the marginal distributions of the input domains and not use uniform distribution, please set the attributes p and q to the distributions of your choice (for domain 1, and 2, respectively) 
+			after initializing a SCOT class instance and before running alignment and set init_marginals=False in .align() parameters
+	"""
+
+	def __init__(self, data):
+
+		assert type(data)==list and len(data)>=2, "As input, SCOTv2 requires a list, containing at least two numpy arrays to be aligned.  \
+				Each numpy array/matrix corresponds to a dataset, with samples (cells) in rows and features (latent representations or genomic features) in columns. \
+				We recommend using latent representations (e.g. principal components for RNA-seq and topics - via cisTopic- for ATAC-seq/Methyl-seq)."
+		self.data=data
+		self.marginals=[] # Holds the empirical probability distributions over samples in each dataset
+		self.graphs=[] # Holds graphs per dataset
+		self.graphDists=[] # Holds intra-domain graph distances for each input dataset
+		self.couplings=[] # Holds coupling matrices
+		self.gwdists=[] # Gromov-Wasserstein distances between domains after alignment
+		self.flags = [] # Holds alignment convergence flags (booleans: True/False)
+
+		self.aligned_data=[]
+
+	def _init_marginals(self):
+		# Without any prior information, we set the probabilities to what we observe empirically: uniform over all observed sample
+		for i in range(len(self.data)):
+			num_cells=self.data[i].shape[0]
+			marginalDist=torch.ones(num_cells)/num_cells
+			self.marginals.append(marginalDist)
+		return self.marginals
+
+	def _normalize(self, norm="l2", bySample=True):
+		assert (norm in ["l1","l2","max", "zscore"]), "Norm argument has to be either one of 'max', 'l1', 'l2' or 'zscore'.\
+		 If you would like to perform another type of normalization, please give SCOT the normalized data and set the argument 'normalize=False' when running the algorithm. \
+		 We have found l2 normalization to empirically perform better with single-cell sequencing datasets, including when using latent representations. "
+
+		for i in range(len(self.data)):
+			if norm=="zscore":
+				scaler=StandardScaler()
+				self.data[i]=scaler.fit_transform(self.data[i])
+			else:
+				if (bySample==True or bySample==None):
+					axis=1
+				else:
+					axis=0
+				self.data[i] =normalize(self.data[i], norm=norm, axis=axis)
+		return self.data # Normalized data
+
+	def construct_graph(self, k=20, mode= "connectivity", metric="correlation"):
+		assert (mode in ["connectivity", "distance"]), "Norm argument has to be either one of 'connectivity', or 'distance'. "
+		if mode=="connectivity":
+			include_self=True
+		else:
+			include_self=False
+
+		for i in range(len(self.data)):
+			self.graphs.append(kneighbors_graph(self.data[i], n_neighbors=k, mode=mode, metric=metric, include_self=include_self))
+
+		return self.graphs
+
+	def init_graph_distances(self):
+		for i in range(len(self.data)):
+			# Compute shortest distances
+			shortestPath=dijkstra(csgraph= csr_matrix(self.graphs[i]), directed=False, return_predecessors=False)
+			# Deal with unconnected stuff (infinities):
+			Max_dist=np.nanmax(shortestPath[shortestPath != np.inf])
+			shortestPath[shortestPath > Max_dist] = Max_dist
+			# Finnally, normalize the distance matrix:
+			self.graphDists.append(shortestPath/shortestPath.max())
+
+		return self.graphDists
+
+	def _exp_sinkhorn_solver(self, ecost, u, v,a,b, mass, eps, rho, rho2, nits_sinkhorn, tol_sinkhorn):
+			"""
+			Parameters
+			----------
+			- ecost: torch.Tensor of size [size_X, size_Y]
+					 Exponential kernel generated from the local cost based on the current coupling.  
+			- u: torch.Tensor of size [size_X[0]].
+				 First dual potential defined on X.
+			- v: torch.Tensor of size [size_Y[0]].
+				 Second dual potential defined on Y. 
+			- mass: torch.Tensor of size [1]. 
+					Mass of the current coupling.
+			- nits_sinkhorn: int. 
+							 Maximum number of iterations to update Sinkhorn potentials in inner loop.
+			- tol_sinkhorn: float
+							Tolerance on convergence of Sinkhorn potentials.
+
+			Returns
+			----------
+			u: torch.Tensor of size [size_X[0]]
+			   First dual potential of Sinkhorn algorithm
+			v: torch.Tensor of size [size_Y[0]]
+			   Second dual potential of Sinkhorn algorithm
+			logpi: torch.Tensor of size [size_X, size_Y]
+				   Optimal transport plan in log-space.
+			"""
+			# Initialize potentials by finding best translation
+			if u is None or v is None:
+				u, v = torch.ones_like(a), torch.ones_like(b)
+			k = (a * u ** (-eps / rho)).sum()+ (b * v ** (-eps / rho)).sum()
+			k = k / (2 * (u[:, None] * v[None, :] * ecost *a[:, None] * b[None, :]).sum())
+			z = (0.5 * mass * eps) / (2.0 + 0.5 * (eps / rho) + 0.5 * (eps / rho2))
+			k = k ** z
+			u,v= u * k, v * k
+
+			# perform Sinkhorn updates in LSE form
+			for j in range(nits_sinkhorn):
+				u_prev = u.clone()
+				v = torch.einsum("ij,i->j", ecost, a * u) ** (-1.0 / (1.0 + eps / rho))
+				u = torch.einsum("ij,j->i", ecost, b * v) ** (-1.0 / (1.0 + eps / rho2))
+				if (u.log() - u_prev.log()).abs().max().item() * eps < tol_sinkhorn:
+					break
+			pi = u[:, None] * v[None, :] * ecost * a[:, None] * b[None, :]
+			return u, v, pi
+
+	def exp_unbalanced_gw(self,a, dx, b, dy, eps=0.01, rho=1.0, rho2=None, nits_plan=3000, tol_plan=1e-6, nits_sinkhorn=3000, tol_sinkhorn=1e-6):
+		if rho2 is None:
+			rho2 = rho #KL divergence coefficient doesn't have to be the same for both couplings. 
+					   #But, to keep #hyperparameters low, we default to using the same coefficient. 
+					   #Someone else playing with our code could assign a rho2 different than rho, though.
+
+		# Initialize the coupling and local costs
+		pi= a[:, None]* b[None, :] / (a.sum() * b.sum()).sqrt()
+		pi_prev = torch.zeros_like(pi)
+		up, vp = None, None
+
+		for i in range(nits_plan):
+			pi_prev = pi.clone()
+			mp = pi.sum()
+
+			#Compute the current local cost:
+			distxy = torch.einsum("ij,kj->ik", dx, torch.einsum("kl,jl->kj", dy, pi))
+			kl_pi = torch.sum(pi * (pi / (a[:, None] * b[None, :]) + 1e-10).log())
+			mu, nu = torch.sum(pi, dim=1), torch.sum(pi, dim=0)
+			distxx = torch.einsum("ij,j->i", dx ** 2, mu)
+			distyy = torch.einsum("kl,l->k", dy ** 2, nu)
+			lcost = (distxx[:, None] + distyy[None, :] - 2 * distxy) + eps * kl_pi
+			if rho < float("Inf"):
+				lcost = (lcost+ rho* torch.sum(mu * (mu / a + 1e-10).log()))
+			if rho2 < float("Inf"):
+				lcost = (lcost+ rho2* torch.sum(nu * (nu / b + 1e-10).log()))
+			ecost = (-lcost /(mp * eps)).exp()
+
+			if (i%10)==0:
+				print("Unbalanced GW step:", i)
+			#compute the coupling via sinkhorn
+			up, vp, pi = self._exp_sinkhorn_solver(ecost, up, vp, a, b, mp, eps, rho, rho2,nits_sinkhorn, tol_sinkhorn)
+			
+			flag=True
+			if torch.any(torch.isnan(pi)):
+				flag=False
+
+			pi = (mp / pi.sum()).sqrt() * pi
+			if (pi - pi_prev).abs().max().item() < tol_plan:
+				break
+		return pi, flag
+
+	def find_correspondences(self, normalize=True, norm="l2", bySample=True, k=20, mode= "connectivity", metric="correlation",  eps=0.01, rho=1.0, rho2=None):
+		# Normalize 
+		if normalize:
+			self._normalize(norm=norm, bySample=bySample)
+		# Initialize inputs for (unbalanced) Gromov-Wasserstein optimal transport:
+		self._init_marginals()
+		print("computing intra-domain graph distances")
+		self.construct_graph(k=k, mode=mode, metric=metric)
+		self.init_graph_distances()
+		# Run pairwise dataset alignments:
+		for i in range(len(self.data)-1):
+			print("running pairwise dataset alignments")
+			a,b =torch.Tensor(self.marginals[0]), torch.Tensor(self.marginals[i+1])
+			dx, dy= torch.Tensor(self.graphDists[0]), torch.Tensor(self.graphDists[i+1])
+			coupling, flag=self.exp_unbalanced_gw(a, dx, b, dy, eps=eps, rho=rho, rho2=rho2, nits_plan=3000, tol_plan=1e-6, nits_sinkhorn=3000, tol_sinkhorn=1e-6)
+			self.couplings.append(coupling)
+			self.flags.append(flag)
+			if flag==False:
+					raise Exception(
+					f"Solver got NaN plan with params (eps, rho, rho2) "
+					f" = {eps, rho, rho2}. Try increasing argument eps")
+		return self.couplings
+
+	def barycentric_projection(self):
+		aligned_datasets=[self.data[0]]
+		for i in range(0,len(self.couplings)):
+			coupling=np.transpose(self.couplings[i].numpy())
+			weights=np.sum(coupling, axis = 1)
+			projected_data=np.matmul((coupling/ weights[:, None]), self.data[0])
+			aligned_datasets.append(projected_data)
+		return aligned_datasets
+
+	def coembed_datasets(self, Lambda=1.0, out_dim=10):
+		"""
+		Co-embeds datasets in a shared space.
+		Implementation is based on Cao et al 2022 (Pamona)
+		"""
+		n_datasets = len(self.data)
+		H0 = []
+		L = []
+		for i in range(n_datasets-1):
+			self.couplings[i] = self.couplings[i]*np.shape(self.data[i])[0]
+
+		for i in range(n_datasets):    
+			graph_data = self.graphs[i] + self.graphs[i].T.multiply(self.graphs[i].T > self.graphs[i]) - \
+				self.graphs[i].multiply(self.graphs[i].T > self.graphs[i])
+			W = np.array(graph_data.todense())
+			index_pos = np.where(W>0)
+			W[index_pos] = 1/W[index_pos] 
+			D = np.diag(np.dot(W, np.ones(np.shape(W)[1])))
+			L.append(D - W)
+
+		Sigma_x = []
+		Sigma_y = []
+		for i in range(n_datasets-1):
+			Sigma_y.append(np.diag(np.dot(np.transpose(np.ones(np.shape(self.couplings[i])[0])), self.couplings[i])))
+			Sigma_x.append(np.diag(np.dot(self.couplings[i], np.ones(np.shape(self.couplings[i])[1]))))
+
+		S_xy = self.couplings[0]
+		S_xx = L[0] + Lambda*Sigma_x[0]
+		S_yy = L[-1] +Lambda*Sigma_y[0]
+		for i in range(1, n_datasets-1):
+			S_xy = np.vstack((S_xy, self.couplings[i]))
+			S_xx = block_diag(S_xx, L[i] + Lambda*Sigma_x[i])
+			S_yy = S_yy + Lambda*Sigma_y[i]
+
+		v, Q = np.linalg.eig(S_xx)
+		v = v + 1e-12   
+		V = np.diag(v**(-0.5))
+		H_x = np.dot(Q, np.dot(V, np.transpose(Q)))
+
+		v, Q = np.linalg.eig(S_yy)
+		v = v + 1e-12      
+		V = np.diag(v**(-0.5))
+		H_y = np.dot(Q, np.dot(V, np.transpose(Q)))
+
+		H = np.dot(H_x, np.dot(S_xy, H_y))
+		U, sigma, V = np.linalg.svd(H)
+
+		num = [0]
+		for i in range(n_datasets-1):
+			num.append(num[i]+len(self.data[i]))
+
+		U, V = U[:,:out_dim], np.transpose(V)[:,:out_dim]
+
+		fx = np.dot(H_x, U)
+		fy = np.dot(H_y, V)
+
+		integrated_data = []
+		for i in range(n_datasets-1):
+			integrated_data.append(fx[num[i]:num[i+1]])
+
+		integrated_data.append(fy)
+
+		return integrated_data
+
+	def align(self,normalize=True, norm="l2", bySample=True, k=20, mode= "connectivity", metric="correlation",  eps=0.01, rho=1.0, rho2=None, projMethod="embedding", Lambda=1.0, out_dim=10):
+		assert projMethod in ["embedding", "barycentric"], "The input to the parameter 'projMethod' needs to be one of \
+								'embedding' (if co-embedding them in a new shared space) or 'barycentric' (if using barycentric projection)"
+		self.find_correspondences(normalize=normalize, norm=norm, bySample=bySample, k=k, mode=mode, metric=metric,  eps=eps, rho=rho, rho2=rho2)
+		print("FLAGS", self.flags)
+		if projMethod=="embedding":
+			integrated_data=self.coembed_datasets(Lambda=Lambda, out_dim=out_dim)
+		else:
+			integrated_data=self.barycentric_projection()
+		self.integrated_data=integrated_data
+		return integrated_data
+	
+# X=np.load("../data/SNARE/SNAREseq_atac_feat.npy")[0:1000,:]
+# Y=np.load("../data/SNARE/SNAREseq_rna_feat.npy")
+# print(X.shape, Y.shape)
+# SCOT=SCOTv2([Y,X])
+# aligned_datasets=SCOT.align(normalize=True, k=50, eps=0.005, rho=0.1, projMethod="barycentric")
+# print(len(aligned_datasets))
+# print(aligned_datasets[0].shape)
+# print(aligned_datasets[1].shape)
+# # np.save("aligned_atac.npy", aligned_datasets[1])
+# np.save("aligned_rna.npy", aligned_datasets[0])
+
+