faiss 0.2.0 → 0.2.4

data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp

@@ -7,43 +7,48 @@
 
 // -*- c++ -*-
 
-
 #include <faiss/IndexIVFSpectralHash.h>
 
-#include <memory>
-#include <algorithm>
 #include <stdint.h>
+#include <algorithm>
+#include <memory>
 
+#include <faiss/IndexLSH.h>
+#include <faiss/IndexPreTransform.h>
+#include <faiss/VectorTransform.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
 #include <faiss/utils/hamming.h>
 #include <faiss/utils/utils.h>
-#include <faiss/impl/FaissAssert.h>
-#include <faiss/impl/AuxIndexStructures.h>
-#include <faiss/VectorTransform.h>
 
 namespace faiss {
 
-
-IndexIVFSpectralHash::IndexIVFSpectralHash (
-        Index * quantizer, size_t d, size_t nlist,
-        int nbit, float period):
-    IndexIVF (quantizer, d, nlist, (nbit + 7) / 8, METRIC_L2),
-    nbit (nbit), period (period), threshold_type (Thresh_global)
-{
-    FAISS_THROW_IF_NOT (code_size % 4 == 0);
-    RandomRotationMatrix *rr = new RandomRotationMatrix (d, nbit);
-    rr->init (1234);
+IndexIVFSpectralHash::IndexIVFSpectralHash(
+        Index* quantizer,
+        size_t d,
+        size_t nlist,
+        int nbit,
+        float period)
+        : IndexIVF(quantizer, d, nlist, (nbit + 7) / 8, METRIC_L2),
+          nbit(nbit),
+          period(period),
+          threshold_type(Thresh_global) {
+    RandomRotationMatrix* rr = new RandomRotationMatrix(d, nbit);
+    rr->init(1234);
     vt = rr;
     own_fields = true;
     is_trained = false;
 }
 
-IndexIVFSpectralHash::IndexIVFSpectralHash():
-    IndexIVF(), vt(nullptr), own_fields(false),
-    nbit(0), period(0), threshold_type(Thresh_global)
-{}
+IndexIVFSpectralHash::IndexIVFSpectralHash()
+        : IndexIVF(),
+          vt(nullptr),
+          own_fields(false),
+          nbit(0),
+          period(0),
+          threshold_type(Thresh_global) {}
 
-IndexIVFSpectralHash::~IndexIVFSpectralHash ()
-{
+IndexIVFSpectralHash::~IndexIVFSpectralHash() {
     if (own_fields) {
         delete vt;
     }
@@ -51,35 +56,33 @@ IndexIVFSpectralHash::~IndexIVFSpectralHash ()
 
 namespace {
 
-
-float median (size_t n, float *x) {
+float median(size_t n, float* x) {
     std::sort(x, x + n);
     if (n % 2 == 1) {
-        return x [n / 2];
+        return x[n / 2];
     } else {
-        return (x [n / 2 - 1] + x [n / 2]) / 2;
+        return (x[n / 2 - 1] + x[n / 2]) / 2;
     }
 }
 
-}
-
+} // namespace
 
-void IndexIVFSpectralHash::train_residual (idx_t n, const float *x)
-{
+void IndexIVFSpectralHash::train_residual(idx_t n, const float* x) {
     if (!vt->is_trained) {
-        vt->train (n, x);
+        vt->train(n, x);
     }
 
     if (threshold_type == Thresh_global) {
         // nothing to do
         return;
-    } else if (threshold_type == Thresh_centroid ||
-        threshold_type == Thresh_centroid_half) {
+    } else if (
+            threshold_type == Thresh_centroid ||
+            threshold_type == Thresh_centroid_half) {
         // convert all centroids with vt
-        std::vector<float> centroids (nlist * d);
-        quantizer->reconstruct_n (0, nlist, centroids.data());
+        std::vector<float> centroids(nlist * d);
+        quantizer->reconstruct_n(0, nlist, centroids.data());
         trained.resize(nlist * nbit);
-        vt->apply_noalloc (nlist, centroids.data(), trained.data());
+        vt->apply_noalloc(nlist, centroids.data(), trained.data());
         if (threshold_type == Thresh_centroid_half) {
             for (size_t i = 0; i < nlist * nbit; i++) {
                 trained[i] -= 0.25 * period;
@@ -90,12 +93,12 @@ void IndexIVFSpectralHash::train_residual (idx_t n, const float *x)
     // otherwise train medians
 
     // assign
-    std::unique_ptr<idx_t []> idx (new idx_t [n]);
-    quantizer->assign (n, x, idx.get());
+    std::unique_ptr<idx_t[]> idx(new idx_t[n]);
+    quantizer->assign(n, x, idx.get());
 
     std::vector<size_t> sizes(nlist + 1);
     for (size_t i = 0; i < n; i++) {
-        FAISS_THROW_IF_NOT (idx[i] >= 0);
+        FAISS_THROW_IF_NOT(idx[i] >= 0);
         sizes[idx[i]]++;
     }
 
@@ -107,10 +110,10 @@ void IndexIVFSpectralHash::train_residual (idx_t n, const float *x)
     }
 
     // transform
-    std::unique_ptr<float []> xt (vt->apply (n, x));
+    std::unique_ptr<float[]> xt(vt->apply(n, x));
 
     // transpose + reorder
-    std::unique_ptr<float []> xo (new float[n * nbit]);
+    std::unique_ptr<float[]> xo(new float[n * nbit]);
 
     for (size_t i = 0; i < n; i++) {
         size_t idest = sizes[idx[i]]++;
@@ -119,14 +122,14 @@ void IndexIVFSpectralHash::train_residual (idx_t n, const float *x)
         }
     }
 
-    trained.resize (n * nbit);
+    trained.resize(n * nbit);
     // compute medians
 #pragma omp for
     for (int i = 0; i < nlist; i++) {
         size_t i0 = i == 0 ? 0 : sizes[i - 1];
         size_t i1 = sizes[i];
         for (int j = 0; j < nbit; j++) {
-            float *xoi = xo.get() + i0 + n * j;
+            float* xoi = xo.get() + i0 + n * j;
             if (i0 == i1) { // nothing to train
                 trained[i * nbit + j] = 0.0;
             } else if (i1 == i0 + 1) {
@@ -138,75 +141,71 @@ void IndexIVFSpectralHash::train_residual (idx_t n, const float *x)
     }
 }
 
-
 namespace {
 
-void binarize_with_freq(size_t nbit, float freq,
-                        const float *x, const float *c,
-                        uint8_t *codes)
-{
-    memset (codes, 0, (nbit + 7) / 8);
+void binarize_with_freq(
+        size_t nbit,
+        float freq,
+        const float* x,
+        const float* c,
+        uint8_t* codes) {
+    memset(codes, 0, (nbit + 7) / 8);
     for (size_t i = 0; i < nbit; i++) {
         float xf = (x[i] - c[i]);
-        int xi = int(floor(xf * freq));
-        int bit = xi & 1;
+        int64_t xi = int64_t(floor(xf * freq));
+        int64_t bit = xi & 1;
         codes[i >> 3] |= bit << (i & 7);
     }
 }
 
+}; // namespace
 
-};
-
-
-
-void IndexIVFSpectralHash::encode_vectors(idx_t n, const float* x_in,
-                                          const idx_t *list_nos,
-                                          uint8_t * codes,
-                                          bool include_listnos) const
-{
-    FAISS_THROW_IF_NOT (is_trained);
+void IndexIVFSpectralHash::encode_vectors(
+        idx_t n,
+        const float* x_in,
+        const idx_t* list_nos,
+        uint8_t* codes,
+        bool include_listnos) const {
+    FAISS_THROW_IF_NOT(is_trained);
     float freq = 2.0 / period;
-
-    FAISS_THROW_IF_NOT_MSG (!include_listnos, "listnos encoding not supported");
+    size_t coarse_size = include_listnos ? coarse_code_size() : 0;
 
     // transform with vt
-    std::unique_ptr<float []> x (vt->apply (n, x_in));
+    std::unique_ptr<float[]> x(vt->apply(n, x_in));
 
-#pragma omp parallel
-    {
-        std::vector<float> zero (nbit);
+    std::vector<float> zero(nbit);
 
-        // each thread takes care of a subset of lists
 #pragma omp for
-        for (idx_t i = 0; i < n; i++) {
-            int64_t list_no = list_nos [i];
-
-            if (list_no >= 0) {
-                const float *c;
-                if (threshold_type == Thresh_global) {
-                    c = zero.data();
-                } else {
-                    c = trained.data() + list_no * nbit;
-                }
-                binarize_with_freq (nbit, freq,
-                                    x.get() + i * nbit, c,
-                                    codes + i * code_size) ;
+    for (idx_t i = 0; i < n; i++) {
+        int64_t list_no = list_nos[i];
+        uint8_t* code = codes + i * (code_size + coarse_size);
+
+        if (list_no >= 0) {
+            if (coarse_size) {
+                encode_listno(list_no, code);
+            }
+            const float* c;
+
+            if (threshold_type == Thresh_global) {
+                c = zero.data();
+            } else {
+                c = trained.data() + list_no * nbit;
             }
+            binarize_with_freq(
+                    nbit, freq, x.get() + i * nbit, c, code + coarse_size);
+        } else {
+            memset(code, 0, code_size + coarse_size);
         }
     }
 }
 
 namespace {
 
-
-template<class HammingComputer>
-struct IVFScanner: InvertedListScanner {
-
+template <class HammingComputer>
+struct IVFScanner : InvertedListScanner {
     // copied from index structure
-    const IndexIVFSpectralHash *index;
-    size_t code_size;
+    const IndexIVFSpectralHash* index;
     size_t nbit;
-    bool store_pairs;
 
     float period, freq;
     std::vector<float> q;
@@ -216,61 +215,57 @@ struct IVFScanner: InvertedListScanner {
 
     using idx_t = Index::idx_t;
 
-    IVFScanner (const IndexIVFSpectralHash * index,
-                bool store_pairs):
-        index (index),
-        code_size(index->code_size),
-        nbit(index->nbit),
-        store_pairs(store_pairs),
-        period(index->period), freq(2.0 / index->period),
-        q(nbit), zero(nbit), qcode(code_size),
-        hc(qcode.data(), code_size)
-    {
+    IVFScanner(const IndexIVFSpectralHash* index, bool store_pairs)
+            : index(index),
+              nbit(index->nbit),
+              period(index->period),
+              freq(2.0 / index->period),
+              q(nbit),
+              zero(nbit),
+              qcode(index->code_size),
+              hc(qcode.data(), index->code_size) {
+        this->store_pairs = store_pairs;
+        this->code_size = index->code_size;
     }
 
-
-    void set_query (const float *query) override {
+    void set_query(const float* query) override {
         FAISS_THROW_IF_NOT(query);
         FAISS_THROW_IF_NOT(q.size() == nbit);
-        index->vt->apply_noalloc (1, query, q.data());
+        index->vt->apply_noalloc(1, query, q.data());
 
-        if (index->threshold_type ==
-            IndexIVFSpectralHash::Thresh_global) {
-            binarize_with_freq
-                (nbit, freq, q.data(), zero.data(), qcode.data());
-            hc.set (qcode.data(), code_size);
+        if (index->threshold_type == IndexIVFSpectralHash::Thresh_global) {
+            binarize_with_freq(nbit, freq, q.data(), zero.data(), qcode.data());
+            hc.set(qcode.data(), code_size);
         }
     }
 
-    idx_t list_no;
-
-    void set_list (idx_t list_no, float /*coarse_dis*/) override {
+    void set_list(idx_t list_no, float /*coarse_dis*/) override {
         this->list_no = list_no;
         if (index->threshold_type != IndexIVFSpectralHash::Thresh_global) {
-            const float *c = index->trained.data() + list_no * nbit;
-            binarize_with_freq (nbit, freq, q.data(), c, qcode.data());
-            hc.set (qcode.data(), code_size);
+            const float* c = index->trained.data() + list_no * nbit;
+            binarize_with_freq(nbit, freq, q.data(), c, qcode.data());
+            hc.set(qcode.data(), code_size);
         }
     }
 
-    float distance_to_code (const uint8_t *code) const final {
-        return hc.hamming (code);
+    float distance_to_code(const uint8_t* code) const final {
+        return hc.hamming(code);
     }
 
-    size_t scan_codes (size_t list_size,
-                       const uint8_t *codes,
-                       const idx_t *ids,
-                       float *simi, idx_t *idxi,
-                       size_t k) const override
-    {
+    size_t scan_codes(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float* simi,
+            idx_t* idxi,
+            size_t k) const override {
         size_t nup = 0;
         for (size_t j = 0; j < list_size; j++) {
+            float dis = hc.hamming(codes);
 
-            float dis = hc.hamming (codes);
-
-            if (dis < simi [0]) {
-                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
-                maxheap_replace_top (k, simi, idxi, dis, id);
+            if (dis < simi[0]) {
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                maxheap_replace_top(k, simi, idxi, dis, id);
                 nup++;
             }
             codes += code_size;
@@ -278,34 +273,31 @@ struct IVFScanner: InvertedListScanner {
         return nup;
     }
 
-    void scan_codes_range (size_t list_size,
-                           const uint8_t *codes,
-                           const idx_t *ids,
-                           float radius,
-                           RangeQueryResult & res) const override
-    {
+    void scan_codes_range(
+            size_t list_size,
+            const uint8_t* codes,
+            const idx_t* ids,
+            float radius,
+            RangeQueryResult& res) const override {
         for (size_t j = 0; j < list_size; j++) {
-            float dis = hc.hamming (codes);
+            float dis = hc.hamming(codes);
             if (dis < radius) {
-                int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
-                res.add (dis, id);
+                int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
+                res.add(dis, id);
             }
             codes += code_size;
         }
     }
-
-
 };
 
 } // anonymous namespace
 
-InvertedListScanner* IndexIVFSpectralHash::get_InvertedListScanner
-    (bool store_pairs) const
-{
+InvertedListScanner* IndexIVFSpectralHash::get_InvertedListScanner(
+        bool store_pairs) const {
     switch (code_size) {
 #define HANDLE_CODE_SIZE(cs) \
-    case cs: \
-        return new IVFScanner<HammingComputer ## cs> (this, store_pairs)
+    case cs:                 \
+        return new IVFScanner<HammingComputer##cs>(this, store_pairs)
         HANDLE_CODE_SIZE(4);
         HANDLE_CODE_SIZE(8);
         HANDLE_CODE_SIZE(16);
@@ -314,17 +306,38 @@ InvertedListScanner* IndexIVFSpectralHash::get_InvertedListScanner
         HANDLE_CODE_SIZE(64);
 #undef HANDLE_CODE_SIZE
         default:
-            if (code_size % 8 == 0) {
-                return new IVFScanner<HammingComputerM8>(this, store_pairs);
-            } else if (code_size % 4 == 0) {
-                return new IVFScanner<HammingComputerM4>(this, store_pairs);
-            } else {
-                FAISS_THROW_MSG("not supported");
-            }
+            return new IVFScanner<HammingComputerDefault>(this, store_pairs);
     }
-
 }
 
+void IndexIVFSpectralHash::replace_vt(VectorTransform* vt_in, bool own) {
+    FAISS_THROW_IF_NOT(vt_in->d_out == nbit);
+    FAISS_THROW_IF_NOT(vt_in->d_in == d);
+    if (own_fields) {
+        delete vt;
+    }
+    vt = vt_in;
+    threshold_type = Thresh_global;
+    is_trained = quantizer->is_trained && quantizer->ntotal == nlist &&
+            vt->is_trained;
+    own_fields = own;
+}
 
+/*
+    Check that the encoder is a single vector transform followed by a LSH
+    that just does thresholding.
+    If this is not the case, the linear transform + threhsolds of the IndexLSH
+    should be merged into the VectorTransform (which is feasible).
+*/
+
+void IndexIVFSpectralHash::replace_vt(IndexPreTransform* encoder, bool own) {
+    FAISS_THROW_IF_NOT(encoder->chain.size() == 1);
+    auto sub_index = dynamic_cast<IndexLSH*>(encoder->index);
+    FAISS_THROW_IF_NOT_MSG(sub_index, "final index should be LSH");
+    FAISS_THROW_IF_NOT(sub_index->nbits == nbit);
+    FAISS_THROW_IF_NOT(!sub_index->rotate_data);
+    FAISS_THROW_IF_NOT(!sub_index->train_thresholds);
+    replace_vt(encoder->chain[0], own);
+}
 
-}  // namespace faiss
+} // namespace faiss

data/vendor/faiss/faiss/IndexIVFSpectralHash.h

@@ -10,15 +10,14 @@
 #ifndef FAISS_INDEX_IVFSH_H
 #define FAISS_INDEX_IVFSH_H
 
-
 #include <vector>
 
 #include <faiss/IndexIVF.h>
 
-
 namespace faiss {
 
 struct VectorTransform;
+struct IndexPreTransform;
 
 /** Inverted list that stores binary codes of size nbit. Before the
  * binary conversion, the dimension of the vectors is transformed from
@@ -27,49 +26,63 @@ struct VectorTransform;
  * Each coordinate is subtracted from a value determined by
  * threshold_type, and split into intervals of size period. Half of
  * the interval is a 0 bit, the other half a 1.
+ *
  */
-struct IndexIVFSpectralHash: IndexIVF {
-
-    VectorTransform *vt; // transformation from d to nbit dim
+struct IndexIVFSpectralHash : IndexIVF {
+    /// transformation from d to nbit dim
+    VectorTransform* vt;
+    /// own the vt
     bool own_fields;
 
+    /// nb of bits of the binary signature
     int nbit;
+    /// interval size for 0s and 1s
     float period;
 
     enum ThresholdType {
-        Thresh_global,
-        Thresh_centroid,
-        Thresh_centroid_half,
-        Thresh_median
+        Thresh_global,        ///< global threshold at 0
+        Thresh_centroid,      ///< compare to centroid
+        Thresh_centroid_half, ///< central interval around centroid
+        Thresh_median         ///< median of training set
     };
     ThresholdType threshold_type;
 
-    // size nlist * nbit or 0 if Thresh_global
+    /// Trained threshold.
+    /// size nlist * nbit or 0 if Thresh_global
     std::vector<float> trained;
 
-    IndexIVFSpectralHash (Index * quantizer, size_t d, size_t nlist,
-                          int nbit, float period);
+    IndexIVFSpectralHash(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            int nbit,
+            float period);
 
-    IndexIVFSpectralHash ();
+    IndexIVFSpectralHash();
 
     void train_residual(idx_t n, const float* x) override;
 
-    void encode_vectors(idx_t n, const float* x,
-                        const idx_t *list_nos,
-                        uint8_t * codes,
-                        bool include_listnos = false) const override;
+    void encode_vectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            uint8_t* codes,
+            bool include_listnos = false) const override;
 
-    InvertedListScanner *get_InvertedListScanner (bool store_pairs)
-        const override;
+    InvertedListScanner* get_InvertedListScanner(
+            bool store_pairs) const override;
 
-    ~IndexIVFSpectralHash () override;
+    /** replace the vector transform for an empty (and possibly untrained) index
+     */
+    void replace_vt(VectorTransform* vt, bool own = false);
 
-};
+    /** convenience function to get the VT from an index constucted by an
+     * index_factory (should end in "LSH") */
+    void replace_vt(IndexPreTransform* index, bool own = false);
 
+    ~IndexIVFSpectralHash() override;
+};
 
-
-
-}; // namespace faiss
-
+} // namespace faiss
 
 #endif