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
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