text: import numpy as np
  import scipy
  import scipy.stats
  import scipy.sparse 
  from sklearn.decomposition import PCA,TruncatedSVD
  from sklearn.neighbors import NearestNeighbors
  import time
  
  ########## PREPROCESSING PIPELINE
  
  def print_optional(string, verbose=True):
      if verbose:
          print(string)
      return
  
  def pipeline_normalize(self, postnorm_total=None):
      ''' Total counts normalization '''
      if postnorm_total is None:
          postnorm_total = self._total_counts_obs.mean()
  
      self._E_obs_norm = tot_counts_norm(self._E_obs, target_total=postnorm_total, total_counts=self._total_counts_obs)
  
      if self._E_sim is not None:
          self._E_sim_norm = tot_counts_norm(self._E_sim, target_total=postnorm_total, total_counts=self._total_counts_sim)
      return
  
  def pipeline_get_gene_filter(self, min_counts=3, min_cells=3, min_gene_variability_pctl=85):
      ''' Identify highly variable genes expressed above a minimum level '''
      self._gene_filter = filter_genes(self._E_obs_norm,
                                          min_counts=min_counts,
                                          min_cells=min_cells,
                                          min_vscore_pctl=min_gene_variability_pctl)
      return
  
  def pipeline_apply_gene_filter(self):
      if self._E_obs is not None:
          self._E_obs = self._E_obs[:,self._gene_filter]
      if self._E_obs_norm is not None:
          self._E_obs_norm = self._E_obs_norm[:,self._gene_filter]
      if self._E_sim is not None:
          self._E_sim = self._E_sim[:,self._gene_filter]
      if self._E_sim_norm is not None:
          self._E_sim_norm = self._E_sim_norm[:,self._gene_filter]
      return
  
  def pipeline_mean_center(self):
      gene_means = self._E_obs_norm.mean(0)
      self._E_obs_norm = self._E_obs_norm - gene_means
      if self._E_sim_norm is not None:
          self._E_sim_norm = self._E_sim_norm - gene_means
      return 
  
  def pipeline_normalize_variance(self):
      gene_stdevs = np.sqrt(sparse_var(self._E_obs_norm))
      self._E_obs_norm = sparse_multiply(self._E_obs_norm.T, 1/gene_stdevs).T
      if self._E_sim_norm is not None:
          self._E_sim_norm = sparse_multiply(self._E_sim_norm.T, 1/gene_stdevs).T
      return 
  
  def pipeline_zscore(self):
      gene_means = self._E_obs_norm.mean(0)
      gene_stdevs = np.sqrt(sparse_var(self._E_obs_norm))
      self._E_obs_norm = np.array(sparse_zscore(self._E_obs_norm, gene_means, gene_stdevs))
      if self._E_sim_norm is not None:
          self._E_sim_norm = np.array(sparse_zscore(self._E_sim_norm, gene_means, gene_stdevs))
      return
  
  def pipeline_log_transform(self, pseudocount=1):
      self._E_obs_norm = log_normalize(self._E_obs_norm, pseudocount)
      if self._E_sim_norm is not None:
          self._E_sim_norm = log_normalize(self._E_sim_norm, pseudocount)
      return
  
  def pipeline_truncated_svd(self, n_prin_comps=30, random_state=0, algorithm='arpack'):
      svd = TruncatedSVD(n_components=n_prin_comps, random_state=random_state, algorithm=algorithm).fit(self._E_obs_norm)
      self.set_manifold(svd.transform(self._E_obs_norm), svd.transform(self._E_sim_norm)) 
      return
      
  def pipeline_pca(self, n_prin_comps=50, random_state=0, svd_solver='arpack'):
      if scipy.sparse.issparse(self._E_obs_norm):
          X_obs = self._E_obs_norm.toarray()
      else:
          X_obs = self._E_obs_norm
      if scipy.sparse.issparse(self._E_sim_norm):
          X_sim = self._E_sim_norm.toarray()
      else:
          X_sim = self._E_sim_norm
  
      pca = PCA(n_components=n_prin_comps, random_state=random_state, svd_solver=svd_solver).fit(X_obs)
      self.set_manifold(pca.transform(X_obs), pca.transform(X_sim)) 
      return
  
  def matrix_multiply(X, Y):
      if not type(X) == np.ndarray:
          if scipy.sparse.issparse(X):
              X = X.toarray()
          else:
              X = np.array(X)
      if not type(Y) == np.ndarray:
          if scipy.sparse.issparse(Y):
              Y = Y.toarray()
          else:
              Y = np.array(Y)
      return np.dot(X,Y)
  
  def log_normalize(X,pseudocount=1):
      X.data = np.log10(X.data + pseudocount)
      return X
  
  ########## LOADING DATA
  def load_genes(filename, delimiter='\t', column=0, skip_rows=0):
      gene_list = []
      gene_dict = {}
  
      with open(filename) as f:
          for iL in range(skip_rows):
              f.readline()
          for l in f:
              gene = l.strip('\n').split(delimiter)[column]
              if gene in gene_dict:
                  gene_dict[gene] += 1
                  gene_list.append(gene + '__' + str(gene_dict[gene]))
                  if gene_dict[gene] == 2:
                      i = gene_list.index(gene)
                      gene_list[i] = gene + '__1'
              else: 
                  gene_dict[gene] = 1
                  gene_list.append(gene)
      return gene_list
  
  
  def make_genes_unique(orig_gene_list):
      gene_list = []
      gene_dict = {}
  
      for gene in orig_gene_list:
          if gene in gene_dict:
              gene_dict[gene] += 1
              gene_list.append(gene + '__' + str(gene_dict[gene]))
              if gene_dict[gene] == 2:
                  i = gene_list.index(gene)
                  gene_list[i] = gene + '__1'
          else:
             gene_dict[gene] = 1
             gene_list.append(gene)
      return gene_list
  
  ########## USEFUL SPARSE FUNCTIONS
  
  def sparse_var(E, axis=0):
      ''' variance across the specified axis '''
  
      mean_gene = E.mean(axis=axis).A.squeeze()
      tmp = E.copy()
      tmp.data **= 2
      return tmp.mean(axis=axis).A.squeeze() - mean_gene ** 2
  
  def sparse_multiply(E, a):
      ''' multiply each row of E by a scalar '''
  
      nrow = E.shape[0]
      w = scipy.sparse.lil_matrix((nrow, nrow))
      w.setdiag(a)
      return w * E
  
  def sparse_zscore(E, gene_mean=None, gene_stdev=None):
      ''' z-score normalize each column of E '''
  
      if gene_mean is None:
          gene_mean = E.mean(0)
      if gene_stdev is None:
          gene_stdev = np.sqrt(sparse_var(E))
      return sparse_multiply((E - gene_mean).T, 1/gene_stdev).T
  
  def subsample_counts(E, rate, original_totals, random_seed=0):
      if rate < 1:
          np.random.seed(random_seed)
          E.data = np.random.binomial(np.round(E.data).astype(int), rate)
          current_totals = E.sum(1).A.squeeze()
          unsampled_orig_totals = original_totals - current_totals
          unsampled_downsamp_totals = np.random.binomial(np.round(unsampled_orig_totals).astype(int), rate)
          final_downsamp_totals = current_totals + unsampled_downsamp_totals
      else:
          final_downsamp_totals = original_totals
      return E, final_downsamp_totals
  
  
  ########## GENE FILTERING
  
  def runningquantile(x, y, p, nBins):
  
      ind = np.argsort(x)
      x = x[ind]
      y = y[ind]
  
      dx = (x[-1] - x[0]) / nBins
      xOut = np.linspace(x[0]+dx/2, x[-1]-dx/2, nBins)
  
      yOut = np.zeros(xOut.shape)
  
      for i in range(len(xOut)):
          ind = np.nonzero((x >= xOut[i]-dx/2) & (x < xOut[i]+dx/2))[0]
          if len(ind) > 0:
              yOut[i] = np.percentile(y[ind], p)
          else:
              if i > 0:
                  yOut[i] = yOut[i-1]
              else:
                  yOut[i] = np.nan
  
      return xOut, yOut
  
  
  def get_vscores(E, min_mean=0, nBins=50, fit_percentile=0.1, error_wt=1):
      '''
      Calculate v-score (above-Poisson noise statistic) for genes in the input counts matrix
      Return v-scores and other stats
      '''
  
      ncell = E.shape[0]
  
      mu_gene = E.mean(axis=0).A.squeeze()
      gene_ix = np.nonzero(mu_gene > min_mean)[0]
      mu_gene = mu_gene[gene_ix]
  
      tmp = E[:,gene_ix]
      tmp.data **= 2
      var_gene = tmp.mean(axis=0).A.squeeze() - mu_gene ** 2
      del tmp
      FF_gene = var_gene / mu_gene
  
      data_x = np.log(mu_gene)
      data_y = np.log(FF_gene / mu_gene)
  
      x, y = runningquantile(data_x, data_y, fit_percentile, nBins)
      x = x[~np.isnan(y)]
      y = y[~np.isnan(y)]
  
      gLog = lambda input: np.log(input[1] * np.exp(-input[0]) + input[2])
      h,b = np.histogram(np.log(FF_gene[mu_gene>0]), bins=200)
      b = b[:-1] + np.diff(b)/2
      max_ix = np.argmax(h)
      c = np.max((np.exp(b[max_ix]), 1))
      errFun = lambda b2: np.sum(abs(gLog([x,c,b2])-y) ** error_wt)
      b0 = 0.1
      b = scipy.optimize.fmin(func = errFun, x0=[b0], disp=False)
      a = c / (1+b) - 1
  
  
      v_scores = FF_gene / ((1+a)*(1+b) + b * mu_gene);
      CV_eff = np.sqrt((1+a)*(1+b) - 1);
      CV_input = np.sqrt(b);
  
      return v_scores, CV_eff, CV_input, gene_ix, mu_gene, FF_gene, a, b
  
  def filter_genes(E, base_ix = [], min_vscore_pctl = 85, min_counts = 3, min_cells = 3, show_vscore_plot = False, sample_name = ''):
      ''' 
      Filter genes by expression level and variability
      Return list of filtered gene indices
      '''
  
      if len(base_ix) == 0:
          base_ix = np.arange(E.shape[0])
  
      Vscores, CV_eff, CV_input, gene_ix, mu_gene, FF_gene, a, b = get_vscores(E[base_ix, :])
      ix2 = Vscores>0
      Vscores = Vscores[ix2]
      gene_ix = gene_ix[ix2]
      mu_gene = mu_gene[ix2]
      FF_gene = FF_gene[ix2]
      min_vscore = np.percentile(Vscores, min_vscore_pctl)
      ix = (((E[:,gene_ix] >= min_counts).sum(0).A.squeeze() >= min_cells) & (Vscores >= min_vscore))
      
      if show_vscore_plot:
          import matplotlib.pyplot as plt
          x_min = 0.5*np.min(mu_gene)
          x_max = 2*np.max(mu_gene)
          xTh = x_min * np.exp(np.log(x_max/x_min)*np.linspace(0,1,100))
          yTh = (1 + a)*(1+b) + b * xTh
          plt.figure(figsize=(8, 6));
          plt.scatter(np.log10(mu_gene), np.log10(FF_gene), c = [.8,.8,.8], alpha = 0.3, edgecolors=None);
          plt.scatter(np.log10(mu_gene)[ix], np.log10(FF_gene)[ix], c = [0,0,0], alpha = 0.3, edgecolors=None);
          plt.plot(np.log10(xTh),np.log10(yTh));
          plt.title(sample_name)
          plt.xlabel('log10(mean)');
          plt.ylabel('log10(Fano factor)');
          plt.show()
  
      return gene_ix[ix]
  
  ########## CELL NORMALIZATION
  
  def tot_counts_norm(E, total_counts = None, exclude_dominant_frac = 1, included = [], target_total = None):
      ''' 
      Cell-level total counts normalization of input counts matrix, excluding overly abundant genes if desired.
      Return normalized counts, average total counts, and (if exclude_dominant_frac < 1) list of genes used to calculate total counts 
      '''
  
      E = E.tocsc()
      ncell = E.shape[0]
      if total_counts is None:
          if len(included) == 0:
              if exclude_dominant_frac == 1:
                  tots_use = E.sum(axis=1)
              else:
                  tots = E.sum(axis=1)
                  wtmp = scipy.sparse.lil_matrix((ncell, ncell))
                  wtmp.setdiag(1. / tots)
                  included = np.asarray(~(((wtmp * E) > exclude_dominant_frac).sum(axis=0) > 0))[0,:]
                  tots_use = E[:,included].sum(axis = 1)
                  print('Excluded %i genes from normalization' %(np.sum(~included)))
          else:
              tots_use = E[:,included].sum(axis = 1)
      else:
          tots_use = total_counts.copy()
  
      if target_total is None:
          target_total = np.mean(tots_use)
  
      w = scipy.sparse.lil_matrix((ncell, ncell))
      w.setdiag(float(target_total) / tots_use)
      Enorm = w * E
  
      return Enorm.tocsc()
  
  ########## DIMENSIONALITY REDUCTION
  
  def get_pca(E, base_ix=[], numpc=50, keep_sparse=False, normalize=True, random_state=0, svd_solver='arpack'):
      '''
      Run PCA on the counts matrix E, gene-level normalizing if desired
      Return PCA coordinates
      '''
      # If keep_sparse is True, gene-level normalization maintains sparsity
      #     (no centering) and TruncatedSVD is used instead of normal PCA.
  
      if len(base_ix) == 0:
          base_ix = np.arange(E.shape[0])
  
      if keep_sparse:
          if normalize:
              zstd = np.sqrt(sparse_var(E[base_ix, :]))
              Z = sparse_multiply(E.T, 1 / zstd).T
          else:
              Z = E
          pca = TruncatedSVD(n_components=numpc, random_state=random_state, algorithm=svd_solver)
  
      else:
          if normalize:
              zmean = E[base_ix,:].mean(0)
              zstd = np.sqrt(sparse_var(E[base_ix, :]))
              Z = sparse_multiply((E - zmean).T, 1 / zstd).T
          else:
              Z = E
          pca = PCA(n_components=numpc, random_state=random_state, svd_solver=svd_solver)
  
      pca.fit(Z[base_ix,:])
      return pca.transform(Z)
  
  
  def preprocess_and_pca(E, total_counts_normalize=True, norm_exclude_abundant_gene_frac=1, min_counts=3, min_cells=5, min_vscore_pctl=85, gene_filter=None, num_pc=50, sparse_pca=False, zscore_normalize=True, show_vscore_plot=False):
      '''
      Total counts normalize, filter genes, run PCA
      Return PCA coordinates and filtered gene indices
      '''
  
      if total_counts_normalize:
          print('Total count normalizing')
          E = tot_counts_norm(E, exclude_dominant_frac=norm_exclude_abundant_gene_frac)[0]
  
      if gene_filter is None:
          print('Finding highly variable genes')
          gene_filter = filter_genes(E, min_vscore_pctl=min_vscore_pctl, min_counts=min_counts, min_cells=min_cells, show_vscore_plot=show_vscore_plot)
  
      print('Using %i genes for PCA' %len(gene_filter))
      PCdat = get_pca(E[:,gene_filter], numpc=num_pc, keep_sparse=sparse_pca, normalize=zscore_normalize)
  
      return PCdat, gene_filter
  
  ########## GRAPH CONSTRUCTION
  
  def get_knn_graph(X, k=5, dist_metric='euclidean', approx=False, return_edges=True, random_seed=0):
      '''
      Build k-nearest-neighbor graph
      Return edge list and nearest neighbor matrix
      '''
  
      t0 = time.time()
      if approx:
          try:
              from annoy import AnnoyIndex
          except:
              approx = False
              print('Could not find library "annoy" for approx. nearest neighbor search')
      if approx:
          #print('Using approximate nearest neighbor search')
  
          if dist_metric == 'cosine':
              dist_metric = 'angular'
          npc = X.shape[1]
          ncell = X.shape[0]
          annoy_index = AnnoyIndex(npc, metric=dist_metric)
          annoy_index.set_seed(random_seed)
  
          for i in range(ncell):
              annoy_index.add_item(i, list(X[i,:]))
          annoy_index.build(10) # 10 trees
  
          knn = []
          for iCell in range(ncell):
              knn.append(annoy_index.get_nns_by_item(iCell, k + 1)[1:])
          knn = np.array(knn, dtype=int)
  
      else:
          #print('Using sklearn NearestNeighbors')
  
          if dist_metric == 'cosine':
              nbrs = NearestNeighbors(n_neighbors=k, metric=dist_metric, algorithm='brute').fit(X)
          else:
              nbrs = NearestNeighbors(n_neighbors=k, metric=dist_metric).fit(X)
          knn = nbrs.kneighbors(return_distance=False)
  
      if return_edges:
          links = set([])
          for i in range(knn.shape[0]):
              for j in knn[i,:]:
                  links.add(tuple(sorted((i,j))))
  
          t_elapse = time.time() - t0
          #print('kNN graph built in %.3f sec' %(t_elapse))
  
          return links, knn
      return knn
  
  def build_adj_mat(edges, n_nodes):
      A = scipy.sparse.lil_matrix((n_nodes, n_nodes))
      for e in edges:
          i, j = e
          A[i,j] = 1
          A[j,i] = 1
      return A.tocsc()
  
  ########## 2-D EMBEDDINGS
  
  def get_umap(X, n_neighbors=10, min_dist=0.1, metric='euclidean', random_state=0):
      import umap
      return umap.UMAP(n_neighbors=n_neighbors, min_dist=min_dist, metric=metric, random_state=random_state).fit_transform(X) 
  
  def get_tsne(X, angle=0.5, perplexity=30, random_state=0, verbose=False):
      from sklearn.manifold import TSNE
      return TSNE(angle=angle, perplexity=perplexity, random_state=random_state, verbose=verbose).fit_transform(X)
  
  def get_force_layout(X, n_neighbors=5, approx_neighbors=False, n_iter=300, verbose=False):
      edges = get_knn_graph(X, k=n_neighbors, approx=approx_neighbors, return_edges=True)[0]
      return run_force_layout(edges, X.shape[0], verbose=verbose)
  
  def run_force_layout(links, n_cells, n_iter=100, edgeWeightInfluence=1, barnesHutTheta=2, scalingRatio=1, gravity=0.05, jitterTolerance=1, verbose=False):
      from fa2 import ForceAtlas2
      import networkx as nx
  
      G = nx.Graph()
      G.add_nodes_from(range(n_cells))
      G.add_edges_from(list(links))
  
      forceatlas2 = ForceAtlas2(
                    # Behavior alternatives
                    outboundAttractionDistribution=False,  # Dissuade hubs
                    linLogMode=False,  # NOT IMPLEMENTED
                    adjustSizes=False,  # Prevent overlap (NOT IMPLEMENTED)
                    edgeWeightInfluence=edgeWeightInfluence,
  
                    # Performance
                    jitterTolerance=jitterTolerance,  # Tolerance
                    barnesHutOptimize=True,
                    barnesHutTheta=barnesHutTheta,
                    multiThreaded=False,  # NOT IMPLEMENTED
  
                    # Tuning
                    scalingRatio=scalingRatio,
                    strongGravityMode=False,
                    gravity=gravity,
                    # Log
                    verbose=verbose)
  
      positions = forceatlas2.forceatlas2_networkx_layout(G, pos=None, iterations=n_iter)
      positions = np.array([positions[i] for i in sorted(positions.keys())])
      return positions
  
  ########## CLUSTERING
  
  def get_spectral_clusters(A, k):
      from sklearn.cluster import SpectralClustering
      spec = SpectralClustering(n_clusters=k, random_state = 0, affinity = 'precomputed', assign_labels = 'discretize')
      return spec.fit_predict(A)
  
  
  def get_louvain_clusters(nodes, edges):
      import networkx as nx
      import community
      
      G = nx.Graph()
      G.add_nodes_from(nodes)
      G.add_edges_from(edges)
      
      return np.array(list(community.best_partition(G).values()))
  
  ########## GENE ENRICHMENT
  
  def rank_enriched_genes(E, gene_list, cell_mask, min_counts=3, min_cells=3, verbose=False):
      gix = (E[cell_mask,:]>=min_counts).sum(0).A.squeeze() >= min_cells
      print_optional('%i cells in group' %(sum(cell_mask)), verbose)
      print_optional('Considering %i genes' %(sum(gix)), verbose)
      
      gene_list = gene_list[gix]
      
      z = sparse_zscore(E[:,gix])
      scores = z[cell_mask,:].mean(0).A.squeeze()
      o = np.argsort(-scores)
      
      return gene_list[o], scores[o]
      
  
  ########## PLOTTING STUFF
  
  def darken_cmap(cmap, scale_factor):
      cdat = np.zeros((cmap.N, 4))
      for ii in range(cdat.shape[0]):
          curcol = cmap(ii)
          cdat[ii,0] = curcol[0] * scale_factor
          cdat[ii,1] = curcol[1] * scale_factor
          cdat[ii,2] = curcol[2] * scale_factor
          cdat[ii,3] = 1
      cmap = cmap.from_list(cmap.N, cdat)
      return cmap
  
  def custom_cmap(rgb_list):
      import matplotlib.pyplot as plt
      rgb_list = np.array(rgb_list)
      cmap = plt.cm.Reds
      cmap = cmap.from_list(rgb_list.shape[0],rgb_list)
      return cmap
  
  def plot_groups(x, y, groups, lim_buffer = 50, saving = False, fig_dir = './', fig_name = 'fig', res = 300, close_after = False, title_size = 12, point_size = 3, ncol = 5):
      import matplotlib.pyplot as plt
  
      n_col = int(ncol)
      ngroup = len(np.unique(groups))
      nrow = int(np.ceil(ngroup / float(ncol)))
      fig = plt.figure(figsize = (14, 3 * nrow))
      for ii, c in enumerate(np.unique(groups)):
          ax = plt.subplot(nrow, ncol, ii+1)
          ix = groups == c
  
          ax.scatter(x[~ix], y[~ix], s = point_size, c = [.8,.8,.8], edgecolors=None)
          ax.scatter(x[ix], y[ix], s = point_size, c = [0,0,0], edgecolors=None)
          ax.set_xticks([])
          ax.set_yticks([])
          ax.set_xlim([min(x) - lim_buffer, max(x) + lim_buffer])
          ax.set_ylim([min(y) - lim_buffer, max(y) + lim_buffer])
  
          ax.set_title(str(c), fontsize = title_size)
  
      fig.tight_layout()
  
      if saving:
          if not os.path.exists(fig_dir):
              os.makedirs(fig_dir)
          plt.savefig(fig_dir + '/' + fig_name + '.png', dpi=res)
  
      if close_after:
          plt.close()
  