nodes: 0, module; 1, function_definition; 2, function_name:max_variance_genes; 3, parameters; 4, block; 5, identifier:data; 6, default_parameter; 7, default_parameter; 8, expression_statement; 9, comment:# TODO: profile, make more efficient for large matrices; 10, comment:# 8000 cells: 0.325 seconds; 11, comment:# top time: sparse.csc_tocsr, csc_matvec, astype, copy, mul_scalar; 12, comment:# 73233 cells: 5.347 seconds, 4.762 s in sparse_var; 13, comment:# csc_tocsr: 1.736 s; 14, comment:# copy: 1.028 s; 15, comment:# astype: 0.999 s; 16, comment:# there is almost certainly something superlinear in this method; 17, comment:# maybe it's to_csr?; 18, expression_statement; 19, if_statement; 20, expression_statement; 21, expression_statement; 22, expression_statement; 23, for_statement; 24, return_statement; 25, identifier:nbins; 26, integer:5; 27, identifier:frac; 28, float:0.2; 29, comment:"""
      This function identifies the genes that have the max variance
      across a number of bins sorted by mean.
  
      Args:
          data (array): genes x cells
          nbins (int): number of bins to sort genes by mean expression level. Default: 10.
          frac (float): fraction of genes to return per bin - between 0 and 1. Default: 0.1
  
      Returns:
          list of gene indices (list of ints)
      """; 30, assignment; 31, call; 32, block; 33, else_clause; 34, assignment; 35, assignment; 36, assignment; 37, identifier:i; 38, call; 39, block; 40, identifier:indices; 41, identifier:indices; 42, list; 43, attribute; 44, argument_list; 45, expression_statement; 46, block; 47, identifier:mean_indices; 48, call; 49, identifier:n_elements; 50, call; 51, identifier:frac_elements; 52, call; 53, identifier:range; 54, argument_list; 55, expression_statement; 56, if_statement; 57, expression_statement; 58, expression_statement; 59, expression_statement; 60, expression_statement; 61, comment:# filter out genes with zero variance; 62, expression_statement; 63, expression_statement; 64, identifier:sparse; 65, identifier:issparse; 66, identifier:data; 67, assignment; 68, expression_statement; 69, expression_statement; 70, attribute; 71, argument_list; 72, identifier:int; 73, argument_list; 74, identifier:int; 75, argument_list; 76, identifier:nbins; 77, assignment; 78, comparison_operator:i==nbins-1; 79, block; 80, assignment; 81, assignment; 82, assignment; 83, assignment; 84, assignment; 85, call; 86, pattern_list; 87, call; 88, assignment; 89, assignment; 90, identifier:means; 91, identifier:argsort; 92, binary_operator:data.shape[0]/nbins; 93, binary_operator:n_elements*frac; 94, identifier:bin_i; 95, subscript; 96, identifier:i; 97, binary_operator:nbins-1; 98, expression_statement; 99, identifier:var_i; 100, subscript; 101, identifier:var_sorted; 102, call; 103, identifier:top_var_indices; 104, subscript; 105, identifier:ind; 106, subscript; 107, identifier:ind; 108, list_comprehension; 109, attribute; 110, argument_list; 111, identifier:means; 112, identifier:var; 113, identifier:sparse_mean_var; 114, argument_list; 115, identifier:means; 116, call; 117, identifier:var; 118, call; 119, subscript; 120, identifier:nbins; 121, identifier:n_elements; 122, identifier:frac; 123, identifier:mean_indices; 124, slice; 125, identifier:nbins; 126, integer:1; 127, assignment; 128, identifier:var; 129, identifier:bin_i; 130, attribute; 131, argument_list; 132, identifier:var_sorted; 133, slice; 134, identifier:bin_i; 135, identifier:top_var_indices; 136, identifier:index; 137, for_in_clause; 138, if_clause; 139, identifier:indices; 140, identifier:extend; 141, identifier:ind; 142, identifier:data; 143, attribute; 144, argument_list; 145, attribute; 146, argument_list; 147, attribute; 148, integer:0; 149, binary_operator:i*n_elements; 150, binary_operator:(i+1)*n_elements; 151, identifier:bin_i; 152, subscript; 153, identifier:var_i; 154, identifier:argsort; 155, binary_operator:len(bin_i) - frac_elements; 156, identifier:index; 157, identifier:ind; 158, comparison_operator:var[index]>0; 159, identifier:data; 160, identifier:mean; 161, integer:1; 162, identifier:data; 163, identifier:var; 164, integer:1; 165, identifier:data; 166, identifier:shape; 167, identifier:i; 168, identifier:n_elements; 169, parenthesized_expression; 170, identifier:n_elements; 171, identifier:mean_indices; 172, slice; 173, call; 174, identifier:frac_elements; 175, subscript; 176, integer:0; 177, binary_operator:i+1; 178, binary_operator:i*n_elements; 179, identifier:len; 180, argument_list; 181, identifier:var; 182, identifier:index; 183, identifier:i; 184, integer:1; 185, identifier:i; 186, identifier:n_elements; 187, identifier:bin_i
code: def max_variance_genes(data, nbins=5, frac=0.2):
      """
      This function identifies the genes that have the max variance
      across a number of bins sorted by mean.
  
      Args:
          data (array): genes x cells
          nbins (int): number of bins to sort genes by mean expression level. Default: 10.
          frac (float): fraction of genes to return per bin - between 0 and 1. Default: 0.1
  
      Returns:
          list of gene indices (list of ints)
      """
      # TODO: profile, make more efficient for large matrices
      # 8000 cells: 0.325 seconds
      # top time: sparse.csc_tocsr, csc_matvec, astype, copy, mul_scalar
      # 73233 cells: 5.347 seconds, 4.762 s in sparse_var
      # csc_tocsr: 1.736 s
      # copy: 1.028 s
      # astype: 0.999 s
      # there is almost certainly something superlinear in this method
      # maybe it's to_csr?
      indices = []
      if sparse.issparse(data):
          means, var = sparse_mean_var(data)
      else:
          means = data.mean(1)
          var = data.var(1)
      mean_indices = means.argsort()
      n_elements = int(data.shape[0]/nbins)
      frac_elements = int(n_elements*frac)
      for i in range(nbins):
          bin_i = mean_indices[i*n_elements : (i+1)*n_elements]
          if i==nbins-1:
              bin_i = mean_indices[i*n_elements :]
          var_i = var[bin_i]
          var_sorted = var_i.argsort()
          top_var_indices = var_sorted[len(bin_i) - frac_elements:]
          ind = bin_i[top_var_indices]
          # filter out genes with zero variance
          ind = [index for index in ind if var[index]>0]
          indices.extend(ind)
      return indices