faiss/tests/test_standalone_codec.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.

"""test byte codecs"""

from __future__ import print_function
import numpy as np
import unittest
import faiss
import platform

from common_faiss_tests import get_dataset_2
from faiss.contrib.datasets import SyntheticDataset
from faiss.contrib.inspect_tools import get_additive_quantizer_codebooks


class TestEncodeDecode(unittest.TestCase):

    def do_encode_twice(self, factory_key):
        d = 96
        nb = 1000
        nq = 0
        nt = 2000

        xt, x, _ = get_dataset_2(d, nt, nb, nq)

        assert x.size > 0

        codec = faiss.index_factory(d, factory_key)

        codec.train(xt)

        codes = codec.sa_encode(x)
        x2 = codec.sa_decode(codes)

        codes2 = codec.sa_encode(x2)

        if "IVF" in factory_key or "RQ" in factory_key:
            # some rows are not reconstructed exactly because they
            # flip into another quantization cell
            nrowdiff = (codes != codes2).any(axis=1).sum()
            self.assertTrue(nrowdiff < 10)
        else:
            self.assertTrue(np.all(codes == codes2))

        x3 = codec.sa_decode(codes2)

        if "IVF" in factory_key or "RQ" in factory_key:
            diffs = np.abs(x2 - x3).sum(axis=1)
            avg = np.abs(x2).sum(axis=1).mean()
            diffs.sort()
            assert diffs[-10] < avg * 1e-5
        else:
            self.assertTrue(np.allclose(x2, x3))

    def test_SQ8(self):
        self.do_encode_twice("SQ8")

    def test_SQtqmse1(self):
        self.do_encode_twice("SQtqmse1")

    def test_SQtqmse2(self):
        self.do_encode_twice("SQtqmse2")

    def test_SQtqmse3(self):
        self.do_encode_twice("SQtqmse3")

    def test_SQtqmse4(self):
        self.do_encode_twice("SQtqmse4")

    def test_SQtqmse8(self):
        self.do_encode_twice("SQtqmse8")

    def test_IVFSQ8(self):
        self.do_encode_twice("IVF256,SQ8")

    def test_PCAIVFSQ8(self):
        self.do_encode_twice("PCAR32,IVF256,SQ8")

    def test_PQ6x8(self):
        self.do_encode_twice("PQ6np")

    def test_PQ6x6(self):
        self.do_encode_twice("PQ6x6np")

    def test_IVFPQ6x8np(self):
        self.do_encode_twice("IVF512,PQ6np")

    def test_LSH(self):
        self.do_encode_twice("LSHrt")

    def test_RQ6x8(self):
        self.do_encode_twice("RQ6x8")


class TestIndexEquiv(unittest.TestCase):

    def do_test(self, key1, key2):
        d = 96
        nb = 1000
        nq = 0
        nt = 2000

        xt, x, _ = get_dataset_2(d, nt, nb, nq)

        codec_ref = faiss.index_factory(d, key1)
        codec_ref.train(xt)

        code_ref = codec_ref.sa_encode(x)
        x_recons_ref = codec_ref.sa_decode(code_ref)

        codec_new = faiss.index_factory(d, key2)
        codec_new.pq = codec_ref.pq

        # replace quantizer, avoiding mem leak
        oldq = codec_new.q1.quantizer
        oldq.this.own()
        codec_new.q1.own_fields = False
        codec_new.q1.quantizer = codec_ref.quantizer
        codec_new.is_trained = True

        code_new = codec_new.sa_encode(x)
        x_recons_new = codec_new.sa_decode(code_new)

        self.assertTrue(np.all(code_new == code_ref))
        self.assertTrue(np.all(x_recons_new == x_recons_ref))

        codec_new_2 = faiss.deserialize_index(faiss.serialize_index(codec_new))

        code_new = codec_new_2.sa_encode(x)
        x_recons_new = codec_new_2.sa_decode(code_new)

        self.assertTrue(np.all(code_new == code_ref))
        self.assertTrue(np.all(x_recons_new == x_recons_ref))

        # Verify deserialized index is serializable again
        codec_new_3 = faiss.deserialize_index(
            faiss.serialize_index(codec_new_2)
        )
        code_new_3 = codec_new_3.sa_encode(x)
        self.assertTrue(np.all(code_new_3 == code_ref))

    def test_IVFPQ(self):
        self.do_test("IVF512,PQ6np", "Residual512,PQ6")

    def test_IMI(self):
        self.do_test("IMI2x5,PQ6np", "Residual2x5,PQ6")


class TestAccuracy(unittest.TestCase):
    """comparative accuracy of a few types of indexes"""

    def compare_accuracy(self, lowac, highac, max_errs=(1e10, 1e10)):
        d = 96
        nb = 1000
        nq = 0
        nt = 2000

        xt, x, _ = get_dataset_2(d, nt, nb, nq)

        errs = []

        for factory_string in lowac, highac:

            codec = faiss.index_factory(d, factory_string)
            print("sa codec: code size %d" % codec.sa_code_size())
            codec.train(xt)

            codes = codec.sa_encode(x)
            x2 = codec.sa_decode(codes)

            err = ((x - x2) ** 2).sum()
            errs.append(err)

        self.assertGreater(errs[0], errs[1])

        self.assertGreater(max_errs[0], errs[0])
        self.assertGreater(max_errs[1], errs[1])

        # just a small IndexLattice I/O test
        if "Lattice" in highac:
            codec2 = faiss.deserialize_index(faiss.serialize_index(codec))
            codes = codec2.sa_encode(x)
            x3 = codec2.sa_decode(codes)
            self.assertTrue(np.all(x2 == x3))

            # Verify deserialized index is serializable again
            codec3 = faiss.deserialize_index(faiss.serialize_index(codec2))
            codes3 = codec3.sa_encode(x)
            x4 = codec3.sa_decode(codes3)
            self.assertTrue(np.all(x2 == x4))

    def test_SQ(self):
        self.compare_accuracy("SQ4", "SQ8")

    def test_SQ2(self):
        self.compare_accuracy("SQ6", "SQ8")

    def test_SQ3(self):
        self.compare_accuracy("SQ8", "SQfp16")

    def test_SQ4(self):
        self.compare_accuracy("SQ8", "SQbf16")

    def test_SQtqmse(self):
        self.compare_accuracy("SQtqmse4", "SQtqmse8")

    def test_SQtqmse_ordering(self):
        d = 96
        nb = 1000
        nq = 0
        nt = 2000

        xt, x, _ = get_dataset_2(d, nt, nb, nq)

        errs = []
        for factory_string in (
            "SQtqmse1",
            "SQtqmse2",
            "SQtqmse3",
            "SQtqmse4",
            "SQtqmse8",
        ):
            codec = faiss.index_factory(d, factory_string)
            codec.train(xt)
            codes = codec.sa_encode(x)
            x2 = codec.sa_decode(codes)
            errs.append(((x - x2) ** 2).sum())

        for i in range(len(errs) - 1):
            self.assertGreaterEqual(errs[i], errs[i + 1])

    def test_PQ(self):
        self.compare_accuracy("PQ6x8np", "PQ8x8np")

    def test_PQ2(self):
        self.compare_accuracy("PQ8x6np", "PQ8x8np")

    def test_IVFvsPQ(self):
        self.compare_accuracy("PQ8np", "IVF256,PQ8np")

    def test_Lattice(self):
        # measured low/high: 20946.244, 5277.483
        self.compare_accuracy(
            "ZnLattice3x10_4", "ZnLattice3x20_4", (22000, 5400)
        )

    def test_Lattice2(self):
        # here the difference is actually tiny
        # measured errs: [16403.072, 15967.735]
        self.compare_accuracy(
            "ZnLattice3x12_1", "ZnLattice3x12_7", (18000, 16000)
        )


swig_ptr = faiss.swig_ptr


class LatticeTest(unittest.TestCase):
    """Low-level lattice tests"""

    def test_repeats(self):
        rs = np.random.RandomState(123)
        dim = 32
        for _i in range(1000):
            vec = np.floor((rs.rand(dim) ** 7) * 3).astype("float32")
            vecs = vec.copy()
            vecs.sort()
            repeats = faiss.Repeats(dim, swig_ptr(vecs))
            code = repeats.encode(swig_ptr(vec))
            vec2 = np.zeros(dim, dtype="float32")
            repeats.decode(code, swig_ptr(vec2))
            assert np.all(vec == vec2)

    def test_ZnSphereCodec_encode_centroid(self):
        dim = 8
        r2 = 5
        ref_codec = faiss.ZnSphereCodec(dim, r2)
        codec = faiss.ZnSphereCodecRec(dim, r2)
        assert ref_codec.nv == codec.nv
        s = set()
        for i in range(ref_codec.nv):
            c = np.zeros(dim, dtype="float32")
            ref_codec.decode(i, swig_ptr(c))
            code = codec.encode_centroid(swig_ptr(c))
            assert 0 <= code < codec.nv
            s.add(code)
        assert len(s) == codec.nv

    def test_ZnSphereCodecRec(self):
        dim = 16
        r2 = 6
        codec = faiss.ZnSphereCodecRec(dim, r2)
        for i in range(codec.nv):
            c = np.zeros(dim, dtype="float32")
            codec.decode(i, swig_ptr(c))
            code = codec.encode_centroid(swig_ptr(c))
            assert code == i

    def run_ZnSphereCodecAlt(self, dim, r2):
        # dim = 32
        # r2 = 14
        codec = faiss.ZnSphereCodecAlt(dim, r2)
        rs = np.random.RandomState(123)
        n = 100
        codes = rs.randint(codec.nv, size=n, dtype="uint64")
        x = np.empty((n, dim), dtype="float32")
        codec.decode_multi(n, swig_ptr(codes), swig_ptr(x))
        codes2 = np.empty(n, dtype="uint64")
        codec.encode_multi(n, swig_ptr(x), swig_ptr(codes2))

        assert np.all(codes == codes2)

    def test_ZnSphereCodecAlt32(self):
        self.run_ZnSphereCodecAlt(32, 14)

    def test_ZnSphereCodecAlt24(self):
        self.run_ZnSphereCodecAlt(24, 14)

    def test_lattice_index(self):
        index = faiss.index_factory(96, "ZnLattice3x10_4")
        rs = np.random.RandomState(123)
        xq = rs.randn(10, 96).astype("float32")
        xb = rs.randn(20, 96).astype("float32")
        index.train(xb)
        index.add(xb)
        D, I = index.search(xq, 5)
        for i in range(10):
            recons = index.reconstruct_batch(I[i, :])
            ref_dis = ((recons - xq[i]) ** 2).sum(1)
            np.testing.assert_allclose(D[i, :], ref_dis, atol=1e-4)


class TestBitstring(unittest.TestCase):

    def test_rw(self):
        """Low-level bit string tests"""
        rs = np.random.RandomState(1234)
        nbyte = 1000
        sz = 0

        bs = np.ones(nbyte, dtype="uint8")
        bw = faiss.BitstringWriter(swig_ptr(bs), nbyte)

        if False:
            ctrl = [(7, 0x35), (13, 0x1D74)]
            for nbit, x in ctrl:
                bw.write(x, nbit)
        else:
            ctrl = []
            while True:
                nbit = int(1 + 62 * rs.rand() ** 4)
                if sz + nbit > nbyte * 8:
                    break
                x = int(rs.randint(1 << nbit, dtype="int64"))
                bw.write(x, nbit)
                ctrl.append((nbit, x))
                sz += nbit

        bignum = 0
        sz = 0
        for nbit, x in ctrl:
            bignum |= x << sz
            sz += nbit

        for i in range(nbyte):
            self.assertTrue(((bignum >> (i * 8)) & 255) == bs[i])

        br = faiss.BitstringReader(swig_ptr(bs), nbyte)

        for nbit, xref in ctrl:
            xnew = br.read(nbit)
            self.assertTrue(xnew == xref)

    def test_arrays(self):
        nbit = 5
        M = 10
        n = 20
        rs = np.random.RandomState(123)
        a = rs.randint(1 << nbit, size=(n, M), dtype="int32")
        b = faiss.pack_bitstrings(a, nbit)
        c = faiss.unpack_bitstrings(b, M, nbit)
        np.testing.assert_array_equal(a, c)

    def test_arrays_variable_size(self):
        nbits = [10, 5, 3, 12, 6, 7, 4]
        n = 20
        rs = np.random.RandomState(123)
        a = rs.randint(1 << 16, size=(n, len(nbits)), dtype="int32")
        a &= (1 << np.array(nbits)) - 1
        b = faiss.pack_bitstrings(a, nbits)
        c = faiss.unpack_bitstrings(b, nbits)
        np.testing.assert_array_equal(a, c)


class TestIVFTransfer(unittest.TestCase):

    def test_transfer(self):

        ds = SyntheticDataset(32, 2000, 200, 100)
        index = faiss.index_factory(ds.d, "IVF20,SQ8")
        index.train(ds.get_train())
        index.add(ds.get_database())
        Dref, Iref = index.search(ds.get_queries(), 10)
        index.reset()

        codes = index.sa_encode(ds.get_database())
        index.add_sa_codes(codes)

        Dnew, Inew = index.search(ds.get_queries(), 10)

        np.testing.assert_array_equal(Iref, Inew)
        np.testing.assert_array_equal(Dref, Dnew)


class TestIDMap(unittest.TestCase):
    def test_idmap(self):
        ds = SyntheticDataset(32, 2000, 200, 100)
        ids = np.random.randint(10000, size=ds.nb, dtype="int64")
        index = faiss.index_factory(ds.d, "IDMap2,PQ8x2")
        index.train(ds.get_train())
        index.add_with_ids(ds.get_database(), ids)
        Dref, Iref = index.search(ds.get_queries(), 10)

        index.reset()

        index.train(ds.get_train())
        codes = index.index.sa_encode(ds.get_database())
        index.add_sa_codes(codes, ids)
        Dnew, Inew = index.search(ds.get_queries(), 10)

        np.testing.assert_array_equal(Iref, Inew)
        np.testing.assert_array_equal(Dref, Dnew)

    def _check_idmap_ivf(self, factory, with_reconstruct):
        # IVF stores ids in its inverted lists, so add_sa_codes must not
        # forward them to the sub-index (issue #5333). Use two independent
        # indexes so the reference (add_with_ids) and fix-under-test
        # (add_sa_codes) paths don't share state.
        ds = SyntheticDataset(32, 2000, 200, 100)
        ids = (1000 + np.arange(ds.nb)).astype('int64')

        ref = faiss.index_factory(ds.d, factory)
        ref.train(ds.get_train())
        faiss.downcast_index(ref.index).nprobe = 4
        ref.add_with_ids(ds.get_database(), ids)
        Dref, Iref = ref.search(ds.get_queries(), 10)

        test = faiss.index_factory(ds.d, factory)
        test.train(ds.get_train())
        faiss.downcast_index(test.index).nprobe = 4
        codes = test.index.sa_encode(ds.get_database())
        test.add_sa_codes(codes, ids)
        Dnew, Inew = test.search(ds.get_queries(), 10)

        np.testing.assert_array_equal(Iref, Inew)
        np.testing.assert_array_equal(Dref, Dnew)

        if with_reconstruct:
            # reconstruct should work too (rev_map is updated) — IDMap2 only
            faiss.extract_index_ivf(test).make_direct_map()
            for key in ids[:5]:
                recons = test.reconstruct(int(key))
                np.testing.assert_array_almost_equal(
                    recons, ds.get_database()[int(key) - 1000])

    def test_idmap_ivf(self):
        self._check_idmap_ivf("IDMap2,IVF20,Flat", with_reconstruct=True)

    def test_idmap_ivf_base(self):
        # guards the base IndexIDMap fix (nullptr instead of xids) on its own,
        # independently of the IDMap2 rev_map maintenance
        self._check_idmap_ivf("IDMap,IVF20,Flat", with_reconstruct=False)


class TestRefine(unittest.TestCase):

    def test_refine(self):
        """Make sure that IndexRefine can function as a standalone codec"""

        ds = SyntheticDataset(32, 500, 100, 0)
        index = faiss.index_factory(ds.d, "RQ2x5,Refine(ITQ,LSHt)")
        index.train(ds.get_train())

        index1 = index.base_index
        index2 = index.refine_index

        codes12 = index.sa_encode(ds.get_database())
        codes1 = index1.sa_encode(ds.get_database())
        codes2 = index2.sa_encode(ds.get_database())

        np.testing.assert_array_equal(codes12, np.hstack((codes1, codes2)))

    def test_refine_sa_decode(self):
        """Make sure IndexRefine.sa_decode correctly extracts refine codes"""

        ds = SyntheticDataset(32, 500, 100, 0)
        index = faiss.index_factory(ds.d, "PQ4,Refine(SQ8)")
        index.train(ds.get_train())

        xb = ds.get_database()
        codes = index.sa_encode(xb)
        x_decoded = index.sa_decode(codes)

        refine_codes = index.refine_index.sa_encode(xb)
        x_decoded_ref = index.refine_index.sa_decode(refine_codes)

        np.testing.assert_allclose(x_decoded, x_decoded_ref)

    def test_equiv_rcq_rq(self):
        """make sure that the codes generated by the standalone codec are
        the same
        between an
           IndexRefine with ResidualQuantizer
        and
           IVF with ResidualCoarseQuantizer
        both are the centroid id concatenated with the code.
        """
        ds = SyntheticDataset(16, 400, 100, 0)
        index1 = faiss.index_factory(ds.d, "RQ2x3,Refine(Flat)")
        index1.train(ds.get_train())
        irq = faiss.downcast_index(index1.base_index)
        # because the default beam factor for RCQ is 4
        irq.rq.max_beam_size = 4

        index2 = faiss.index_factory(ds.d, "IVF64(RCQ2x3),Flat")
        index2.train(ds.get_train())
        quantizer = faiss.downcast_index(index2.quantizer)
        quantizer.rq = irq.rq
        index2.is_trained = True

        codes1 = index1.sa_encode(ds.get_database())
        codes2 = index2.sa_encode(ds.get_database())

        np.testing.assert_array_equal(codes1, codes2)

    @unittest.skipIf(
        platform.system() == "Windows",
        "Does not work on Windows after numpy 2 upgrade.",
    )
    def test_equiv_sh(self):
        """make sure that the IVFSpectralHash sa_encode function gives the same
        result as the concatenated RQ + LSH index sa_encode"""
        ds = SyntheticDataset(32, 500, 100, 0)
        index1 = faiss.index_factory(ds.d, "RQ1x4,Refine(ITQ16,LSH)")
        index1.train(ds.get_train())

        # reproduce this in an IndexIVFSpectralHash
        coarse_quantizer = faiss.IndexFlat(ds.d)
        rq = faiss.downcast_index(index1.base_index).rq
        centroids = get_additive_quantizer_codebooks(rq)[0]
        coarse_quantizer.add(centroids)

        encoder = faiss.downcast_index(index1.refine_index)

        # larger than the magnitude of the vectors
        # negative because otherwise the bits are flipped
        period = -100000.0

        index2 = faiss.IndexIVFSpectralHash(
            coarse_quantizer,
            ds.d,
            coarse_quantizer.ntotal,
            encoder.sa_code_size() * 8,
            period,
        )

        # replace with the vt of the encoder. Binarization is performed by
        # the IndexIVFSpectralHash itself
        index2.replace_vt(encoder)

        codes1 = index1.sa_encode(ds.get_database())
        codes2 = index2.sa_encode(ds.get_database())

        np.testing.assert_array_equal(codes1, codes2)