docs: def generate_nb_data(P, R, n_cells, assignments=None):
      """
      Generates negative binomial data
  
      Args:
          P (array): genes x clusters
          R (array): genes x clusters
          n_cells (int): number of cells
          assignments (list): cluster assignment of each cell. Default:
              random uniform
  
      Returns:
          data array with shape genes x cells
          labels - array of cluster labels
      """
      genes, clusters = P.shape
      output = np.zeros((genes, n_cells))
      if assignments is None:
          cluster_probs = np.ones(clusters)/clusters
      labels = []
      for i in range(n_cells):
          if assignments is None:
              c = np.random.choice(range(clusters), p=cluster_probs)
          else:
              c = assignments[i]
          labels.append(c)
          # because numpy's negative binomial, r is the number of successes
          output[:,i] = np.random.negative_binomial(R[:,c], 1.0-P[:,c])
      return output, np.array(labels)