faiss/contrib/clustering.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
This contrib module contains a few routines useful to do clustering variants.
"""

import numpy as np
import faiss
import time
from multiprocessing.pool import ThreadPool


try:
    import scipy.sparse
except ImportError:
    print("scipy not accessible, Python k-means will not work")


def print_nop(*arg, **kwargs):
    pass


def two_level_clustering(
    xt, nc1, nc2, rebalance=True, clustering_niter=25, **args
):
    """
    perform 2-level clustering on a training set xt
    nc1 and nc2 are the number of clusters at each level, the final number of
    clusters is nc2. Additional arguments are passed to the Kmeans object.

    Rebalance allocates the number of sub-clusters depending on the number of
    first-level assignment.
    """
    d = xt.shape[1]

    verbose = args.get("verbose", False)

    log = print if verbose else print_nop

    log(
        f"2-level clustering of {xt.shape} nb 1st level clusters = {nc1} "
        f"total {nc2}"
    )
    log("perform coarse training")

    km = faiss.Kmeans(
        d, nc1, niter=clustering_niter, max_points_per_centroid=2000, **args
    )
    km.train(xt)

    iteration_stats = [km.iteration_stats]
    log()

    log("assigning the training set")
    t0 = time.time()
    _, assign1 = km.assign(xt)
    bc = np.bincount(assign1, minlength=nc1)
    log(
        f"done in {time.time() - t0:.2f} s. "
        f"Sizes of clusters {min(bc)}-{max(bc)}"
    )
    o = assign1.argsort()
    del km

    if not rebalance:
        # make sure the sub-clusters sum up to exactly nc2
        cc = np.arange(nc1 + 1) * nc2 // nc1
        all_nc2 = cc[1:] - cc[:-1]
    else:
        bc_sum = np.cumsum(bc)
        all_nc2 = bc_sum * nc2 // bc_sum[-1]
        all_nc2[1:] -= all_nc2[:-1]
        assert sum(all_nc2) == nc2
        log(f"nb 2nd-level centroids {min(all_nc2)}-{max(all_nc2)}")

    # train sub-clusters
    i0 = 0
    c2 = []
    t0 = time.time()
    for c1 in range(nc1):
        nc2 = int(all_nc2[c1])
        log(
            f"[{time.time() - t0:.2f} s] training sub-cluster "
            f"{c1}/{nc1} nc2={nc2}\r",
            end="",
            flush=True,
        )
        i1 = i0 + bc[c1]
        subset = o[i0:i1]
        assert np.all(assign1[subset] == c1)
        km = faiss.Kmeans(d, nc2, **args)
        xtsub = xt[subset]
        km.train(xtsub)
        iteration_stats.append(km.iteration_stats)
        c2.append(km.centroids)
        del km
        i0 = i1
    log(f"done in {time.time() - t0:.2f} s")
    return np.vstack(c2), iteration_stats


def train_ivf_index_with_2level(index, xt, **args):
    """
    Applies 2-level clustering to an index_ivf embedded in an index.
    """
    # handle PreTransforms
    index = faiss.downcast_index(index)
    if isinstance(index, faiss.IndexPreTransform):
        for i in range(index.chain.size()):
            vt = index.chain.at(i)
            vt.train(xt)
            xt = vt.apply(xt)
        train_ivf_index_with_2level(index.index, xt, **args)
        index.is_trained = True
        return
    assert isinstance(index, faiss.IndexIVF)
    assert index.metric_type == faiss.METRIC_L2
    # now do 2-level clustering
    nc1 = int(np.sqrt(index.nlist))
    print("REBALANCE=", args)

    centroids, _ = two_level_clustering(xt, nc1, index.nlist, **args)
    index.quantizer.train(centroids)
    index.quantizer.add(centroids)
    # finish training
    index.train(xt)


def balanced_assignment_with_penalties(
    x, centroids, alpha=0.03, num_iter=20, maxk=100
):
    """
    Assign vectors x to centroids with a balance constraint.

    Iteratively adjusts per-cluster penalties so that oversized clusters
    become less attractive. At each iteration the penalized distance for
    cluster c is ``d(x, c)^2 + penalty_c^2`` and the penalty is updated as
    ``penalty_c *= (binsize_c / n_opt) ** alpha`` where ``n_opt = n / nc``.

    A single kNN call (with *maxk* neighbors) is done upfront; subsequent
    iterations only re-weight among those candidates, making the routine
    fast even for large datasets.

    Reference: "Balancing clusters to reduce response time variability in
    large scale image search", Tavenard et al., CBMI 2011.
    https://inria.hal.science/inria-00576886/document
    See also notebook N10159950.

    Args:
        x:          (n, d) float32 array of vectors to assign.
        centroids:  (nc, d) float32 array of cluster centroids.
        alpha:      exponent that controls how aggressively penalties grow.
                    Higher values yield more balanced clusters at the cost
                    of higher MSE.  Typical range: 0.01 – 0.1.
        num_iter:   number of penalty-update iterations.
        maxk:       number of nearest centroids to consider per vector.
                    Must be <= nc.

    Returns:
        assign:  (n,) int64 array of centroid indices.
        stats:   dict with keys

                 - *imf*: imbalance factor (1.0 = perfectly balanced)
                 - *mse*: mean squared error of the assignment
                 - *binsize_min*, *binsize_max*: smallest / largest cluster
                 - *penalty_min*, *penalty_max*: penalty value range
                 - *alpha*: the alpha value used
    """

    nc = len(centroids)
    n = len(x)
    nopt = n / nc  # targed bin sizes

    # we assign to the top-maxk clusters. The final assignment will pick
    # among these clusters.
    full_d2, full_assign = faiss.knn(x, centroids, maxk)

    # scalar penalty for each cluster
    penalties = np.ones(nc, dtype=np.float32)

    for it in range(num_iter):
        # compute penalized assignment
        penalties2 = penalties**2
        full_d2_penalized = full_d2 + penalties2[full_assign]
        a0 = full_d2_penalized.argmin(axis=1)
        assign = np.take_along_axis(full_assign, a0[:, None], axis=1).ravel()
        binsizes = np.bincount(assign, minlength=nc)
        # print(imbalance_factor(nc, assign), mse, int(binsizes.min()),
        #       int(binsizes.max()))
        penalties *= (binsizes / nopt) ** alpha

    stats = dict(
        alpha=alpha,
        imf=imbalance_factor(nc, assign),
        mse=((x - centroids[assign]) ** 2).sum(1).mean(),  # recompute MSE
        binsize_min=int(binsizes.min()),
        binsize_max=int(binsizes.max()),
        penalty_min=penalties.min(),
        penalty_max=penalties.max(),
    )

    return assign, stats


###############################################################################
# K-means implementation in Python
#
# It relies on DatasetAssign, an abstraction of the training vectors that offers
# the minimal set of operations to perform k-means clustering.
###############################################################################


class DatasetAssign:
    """Wrapper for a matrix that offers a function to assign the vectors
    to centroids. All other implementations offer the same interface"""

    def __init__(self, x):
        self.x = np.ascontiguousarray(x, dtype="float32")

    def count(self):
        return self.x.shape[0]

    def dim(self):
        return self.x.shape[1]

    def get_subset(self, indices):
        return self.x[indices]

    def perform_search(self, centroids):
        return faiss.knn(self.x, centroids, 1)

    def assign_to(self, centroids, weights=None):
        D, I = self.perform_search(centroids)

        I = I.ravel()
        D = D.ravel()
        nc, d = centroids.shape
        sum_per_centroid = np.zeros((nc, d), dtype="float32")
        if weights is None:
            np.add.at(sum_per_centroid, I, self.x)
        else:
            np.add.at(sum_per_centroid, I, weights[:, np.newaxis] * self.x)

        return I, D, sum_per_centroid


class DatasetAssignGPU(DatasetAssign):
    """GPU version of the previous"""

    def __init__(self, x, gpu_id, verbose=False):
        DatasetAssign.__init__(self, x)
        index = faiss.IndexFlatL2(x.shape[1])
        if gpu_id >= 0:
            self.index = faiss.index_cpu_to_gpu(
                faiss.StandardGpuResources(), gpu_id, index
            )
        else:
            # -1 -> assign to all GPUs
            self.index = faiss.index_cpu_to_all_gpus(index)

    def perform_search(self, centroids):
        self.index.reset()
        self.index.add(centroids)
        return self.index.search(self.x, 1)


def sparse_assign_to_dense(xq, xb, xq_norms=None, xb_norms=None):
    """assignment function for xq is sparse, xb is dense
    uses a matrix multiplication. The squared norms can be provided if
    available.
    """
    nq = xq.shape[0]
    nb = xb.shape[0]
    if xb_norms is None:
        xb_norms = (xb**2).sum(1)
    if xq_norms is None:
        xq_norms = np.array(xq.power(2).sum(1))
    d2 = xb_norms - 2 * xq @ xb.T
    I = d2.argmin(axis=1)
    D = d2.ravel()[I + np.arange(nq) * nb] + xq_norms.ravel()
    return D, I


def sparse_assign_to_dense_blocks(
    xq, xb, xq_norms=None, xb_norms=None, qbs=16384, bbs=16384, nt=None
):
    """
    decomposes the sparse_assign_to_dense function into blocks to avoid a
    possible memory blow up. Can be run in multithreaded mode, because scipy's
    sparse-dense matrix multiplication is single-threaded.
    """
    nq = xq.shape[0]
    nb = xb.shape[0]
    D = np.empty(nq, dtype="float32")
    D.fill(np.inf)
    I = -np.ones(nq, dtype=int)

    if xb_norms is None:
        xb_norms = (xb**2).sum(1)

    def handle_query_block(i):
        xq_block = xq[i : i + qbs]
        Iblock = I[i : i + qbs]
        Dblock = D[i : i + qbs]
        if xq_norms is None:
            xq_norms_block = np.array(xq_block.power(2).sum(1))
        else:
            xq_norms_block = xq_norms[i : i + qbs]
        for j in range(0, nb, bbs):
            Di, Ii = sparse_assign_to_dense(
                xq_block,
                xb[j : j + bbs],
                xq_norms=xq_norms_block,
                xb_norms=xb_norms[j : j + bbs],
            )
            if j == 0:
                Iblock[:] = Ii
                Dblock[:] = Di
            else:
                mask = Di < Dblock
                Iblock[mask] = Ii[mask] + j
                Dblock[mask] = Di[mask]

    if nt == 0 or nt == 1 or nq <= qbs:
        list(map(handle_query_block, range(0, nq, qbs)))
    else:
        pool = ThreadPool(nt)
        pool.map(handle_query_block, range(0, nq, qbs))

    return D, I


class DatasetAssignSparse(DatasetAssign):
    """Wrapper for a matrix that offers a function to assign the vectors
    to centroids. All other implementations offer the same interface"""

    def __init__(self, x):
        assert x.__class__ == scipy.sparse.csr_matrix
        self.x = x
        self.squared_norms = np.array(x.power(2).sum(1))

    def get_subset(self, indices):
        return np.array(self.x[indices].todense())

    def perform_search(self, centroids):
        return sparse_assign_to_dense_blocks(
            self.x, centroids, xq_norms=self.squared_norms
        )

    def assign_to(self, centroids, weights=None):
        D, I = self.perform_search(centroids)

        I = I.ravel()
        D = D.ravel()
        n = self.x.shape[0]
        if weights is None:
            weights = np.ones(n, dtype="float32")
        nc = len(centroids)

        m = scipy.sparse.csc_matrix(
            (weights, I, np.arange(n + 1)), shape=(nc, n)
        )
        sum_per_centroid = np.array((m * self.x).todense())

        return I, D, sum_per_centroid


def imbalance_factor(k, assign):
    assign = np.ascontiguousarray(assign, dtype="int64")
    return faiss.imbalance_factor(len(assign), k, faiss.swig_ptr(assign))


def check_if_torch(x):
    if x.__class__ == np.ndarray:
        return False
    import torch

    if isinstance(x, torch.Tensor):
        return True
    raise NotImplementedError(f"Unknown tensor type {type(x)}")


def reassign_centroids(hassign, centroids, rs=None):
    """reassign centroids when some of them collapse"""
    if rs is None:
        rs = np.random
    k, d = centroids.shape
    nsplit = 0
    is_torch = check_if_torch(centroids)

    empty_cents = np.where(hassign == 0)[0]

    if len(empty_cents) == 0:
        return 0

    if is_torch:
        import torch

        fac = torch.ones_like(centroids[0])
    else:
        fac = np.ones_like(centroids[0])
    fac[::2] += 1 / 1024.0
    fac[1::2] -= 1 / 1024.0

    # this is a single pass unless there are more than k/2
    # empty centroids
    while len(empty_cents) > 0:
        # choose which centroids to split (numpy)
        probas = hassign.astype("float") - 1
        probas[probas < 0] = 0
        probas /= probas.sum()
        nnz = (probas > 0).sum()

        nreplace = min(nnz, empty_cents.size)
        cjs = rs.choice(k, size=nreplace, p=probas)

        for ci, cj in zip(empty_cents[:nreplace], cjs):

            c = centroids[cj]
            centroids[ci] = c * fac
            centroids[cj] = c / fac

            hassign[ci] = hassign[cj] // 2
            hassign[cj] -= hassign[ci]
            nsplit += 1

        empty_cents = empty_cents[nreplace:]

    return nsplit


def kmeans(
    k,
    data,
    niter=25,
    seed=1234,
    checkpoint=None,
    verbose=True,
    return_stats=False,
):
    """Pure python kmeans implementation. Follows the Faiss C++ version
    quite closely, but takes a DatasetAssign instead of a training data
    matrix. Also redo is not implemented.

    For the torch implementation, the centroids are tensors (possibly on GPU),
    but the indices remain numpy on CPU.
    """
    n, d = data.count(), data.dim()
    log = print if verbose else print_nop

    log(
        (
            "Clustering %d points in %dD to %d clusters, "
            + "%d iterations seed %d"
        )
        % (n, d, k, niter, seed)
    )

    rs = np.random.RandomState(seed)
    print("preproc...")
    t0 = time.time()
    # initialization
    perm = rs.choice(n, size=k, replace=False)
    centroids = data.get_subset(perm)
    is_torch = check_if_torch(centroids)

    iteration_stats = []

    log("  done")
    t_search_tot = 0
    obj = []
    for i in range(niter):
        t0s = time.time()

        log("assigning", end="\r", flush=True)
        assign, D, sums = data.assign_to(centroids)

        log("compute centroids", end="\r", flush=True)

        t_search_tot += time.time() - t0s

        err = D.sum()
        if is_torch:
            err = err.item()
        obj.append(err)

        hassign = np.bincount(assign, minlength=k)

        fac = hassign.reshape(-1, 1).astype("float32")
        fac[fac == 0] = 1  # quiet warning
        if is_torch:
            import torch

            fac = torch.from_numpy(fac).to(sums.device)

        centroids = sums / fac

        nsplit = reassign_centroids(hassign, centroids, rs)

        s = {
            "obj": err,
            "time": (time.time() - t0),
            "time_search": t_search_tot,
            "imbalance_factor": imbalance_factor(k, assign),
            "nsplit": nsplit,
        }

        log(
            (
                "  Iteration %d (%.2f s, search %.2f s): "
                "objective=%g imbalance=%.3f nsplit=%d"
            )
            % (
                i,
                s["time"],
                s["time_search"],
                err,
                s["imbalance_factor"],
                nsplit,
            )
        )
        iteration_stats.append(s)

        if checkpoint is not None:
            log("storing centroids in", checkpoint)
            if is_torch:
                import torch

                torch.save(centroids, checkpoint)
            else:
                np.save(checkpoint, centroids)

    if return_stats:
        return centroids, iteration_stats
    else:
        return centroids