docs: def nb_cluster(data, k, P_init=None, R_init=None, assignments=None, means=None, max_iters=10):
      """
      Performs negative binomial clustering on the given data. If some genes have mean > variance, then these genes are fitted to a Poisson distribution.
  
      Args:
          data (array): genes x cells
          k (int): number of clusters
          P_init (array): NB success prob param - genes x k. Default: random
          R_init (array): NB stopping param - genes x k. Default: random
          assignments (array): cells x 1 array of integers 0...k-1. Default: kmeans-pp (poisson)
          means (array): initial cluster means (for use with kmeans-pp to create initial assignments). Default: None
          max_iters (int): default: 100
  
      Returns:
          assignments (array): 1d array of length cells, containing integers 0...k-1
          P (array): genes x k - value is 0 for genes with mean > var
          R (array): genes x k - value is inf for genes with mean > var
      """
      genes, cells = data.shape
      if P_init is None:
          P_init = np.random.random((genes, k))
      if R_init is None:
          R_init = np.random.randint(1, data.max(), (genes, k))
          R_init = R_init.astype(float)
      if assignments is None:
          _, assignments = kmeans_pp(data, k, means)
      means = np.zeros((genes, k))
          #assignments = np.array([np.random.randint(0,k) for i in range(cells)])
      old_assignments = np.copy(assignments)
      # If mean > variance, then fall back to Poisson, since NB
      # distribution can't handle that case.
      for i in range(max_iters):
          # estimate params from assigned cells
          nb_gene_indices = fit_cluster(data, assignments, k, P_init, R_init, means)
          # re-calculate assignments
          lls = nb_ll(data[nb_gene_indices, :], P_init[nb_gene_indices,:], R_init[nb_gene_indices,:])
          lls += pois_ll.poisson_ll(data[~nb_gene_indices,:], means[~nb_gene_indices,:])
          # set NB params to failure values
          P_init[~nb_gene_indices,:] = 0
          R_init[~nb_gene_indices,:] = np.inf
          for c in range(cells):
              assignments[c] = np.argmax(lls[c,:])
          if np.equal(assignments,old_assignments).all():
              break
          old_assignments = np.copy(assignments)
      return assignments, P_init, R_init