Initial commit

.gitignore

+/makefile.inc
+/tests/demo_ivfpq_indexing
+*.swp
+*.swo
+*.o
+*.a
+*.dSYM
+*.so
+*.pyc
+*~

AutoTune.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#ifndef FAISS_AUTO_TUNE_H
+#define FAISS_AUTO_TUNE_H
+
+#include <vector>
+
+#include "Index.h"
+
+namespace faiss {
+
+
+/**
+ * Evaluation criterion. Returns a performance measure in [0,1],
+ * higher is better.
+ */
+struct AutoTuneCriterion {
+    typedef Index::idx_t idx_t;
+    idx_t nq;  ///< nb of queries this criterion is evaluated on
+    idx_t nnn; ///< nb of NNs that the query should request
+    idx_t gt_nnn; ///< nb of GT NNs required to evaluate crterion
+
+    std::vector<float> gt_D;  ///< Ground-truth distances (size nq * gt_nnn)
+    std::vector<idx_t> gt_I;  ///< Ground-truth indexes (size nq * gt_nnn)
+
+    AutoTuneCriterion (idx_t nq, idx_t nnn);
+
+    /** Intitializes the gt_D and gt_I vectors. Must be called before evaluating
+     *
+     * @param gt_D_in  size nq * gt_nnn
+     * @param gt_I_in  size nq * gt_nnn
+     */
+    void set_groundtruth (int gt_nnn, const float *gt_D_in,
+                          const idx_t *gt_I_in);
+
+    /** Evaluate the criterion.
+     *
+     * @param D  size nq * nnn
+     * @param I  size nq * nnn
+     * @return the criterion, between 0 and 1. Larger is better.
+     */
+    virtual double evaluate (const float *D, const idx_t *I) const = 0;
+
+    virtual ~AutoTuneCriterion () {}
+
+};
+
+struct OneRecallAtRCriterion: AutoTuneCriterion {
+
+    idx_t R;
+
+    OneRecallAtRCriterion (idx_t nq, idx_t R);
+
+    virtual double evaluate (const float *D, const idx_t *I) const override;
+
+    virtual ~OneRecallAtRCriterion () {}
+};
+
+
+struct IntersectionCriterion: AutoTuneCriterion {
+
+    idx_t R;
+
+    IntersectionCriterion (idx_t nq, idx_t R);
+
+    virtual double evaluate (const float *D, const idx_t *I) const override;
+
+    virtual ~IntersectionCriterion () {}
+};
+
+/**
+ * Maintains a list of experimental results. Each operating point is a
+ * (perf, t, key) triplet, where higher perf and lower t is
+ * better. The key field is an arbitrary identifier for the operating point
+ */
+
+struct OperatingPoint {
+    double perf;     ///< performance measure (output of a Criterion)
+    double t;        ///< corresponding execution time (ms)
+    std::string key; ///< key that identifies this op pt
+    long cno;        ///< integer identifer
+};
+
+struct OperatingPoints {
+    /// all operating points
+    std::vector<OperatingPoint> all_pts;
+
+    /// optimal operating points, sorted by perf
+    std::vector<OperatingPoint> optimal_pts;
+
+    // begins with a single operating point: t=0, perf=0
+    OperatingPoints ();
+
+    /// add operating points from other to this, with a prefix to the keys
+    int merge_with (const OperatingPoints &other,
+                    const std::string & prefix = "");
+
+    void clear ();
+
+    /// add a performance measure. Return whether it is an optimal point
+    bool add (double perf, double t, const std::string & key, size_t cno = 0);
+
+    /// get time required to obtain a given performance measure
+    double t_for_perf (double perf) const;
+
+    /// easy-to-read output
+    void display (bool only_optimal = true) const;
+
+    /// output to a format easy to digest by gnuplot
+    void all_to_gnuplot (const char *fname) const;
+    void optimal_to_gnuplot (const char *fname) const;
+
+};
+
+/// possible values of a parameter, sorted from least to most expensive/accurate
+struct ParameterRange {
+    std::string name;
+    std::vector<double> values;
+};
+
+/** Uses a-priori knowledge on the Faiss indexes to extract tunable parameters.
+ */
+struct ParameterSpace {
+    /// all tunable parameters
+    std::vector<ParameterRange> parameter_ranges;
+
+    // exploration parameters
+
+    /// verbosity during exploration
+    int verbose;
+
+    /// nb of experiments during optimization (0 = try all combinations)
+    int n_experiments;
+
+    /// maximum number of queries to submit at a time.
+    size_t batchsize;
+
+    /// use multithreading over batches (useful to benchmark
+    /// independent single-searches)
+    bool thread_over_batches;
+
+    ParameterSpace ();
+
+    /// nb of combinations, = product of values sizes
+    size_t n_combinations () const;
+
+    /// returns whether combinations c1 >= c2 in the tuple sense
+    bool combination_ge (size_t c1, size_t c2) const;
+
+    /// get string representation of the combination
+    std::string combination_name (size_t cno) const;
+
+    /// print a description on stdout
+    void display () const;
+
+    /// add a new parameter
+    ParameterRange &add_range(const char * name);
+
+    /// initialize with reasonable parameters for the index
+    virtual void initialize (const Index * index);
+
+    /// set a combination of parameters on an index
+    void set_index_parameters (Index *index, size_t cno) const;
+
+    /// set a combination of parameters described by a string
+    void set_index_parameters (Index *index, const char *param_string) const;
+
+    /// set one of the parameters
+    virtual void set_index_parameter (
+        Index * index, const std::string & name, double val) const;
+
+    /** find an upper bound on the performance and a lower bound on t
+     * for configuration cno given another operating point op */
+    void update_bounds (size_t cno, const OperatingPoint & op,
+                        double *upper_bound_perf,
+                        double *lower_bound_t) const;
+
+    /** explore operating points
+     * @param index   index to run on
+     * @param xq      query vectors (size nq * index.d)
+     * @param crit    selection criterion
+     * @param ops     resutling operating points
+     */
+    void explore (Index *index,
+                  size_t nq, const float *xq,
+                  const AutoTuneCriterion & crit,
+                  OperatingPoints * ops)  const;
+
+    virtual ~ParameterSpace () {}
+};
+
+/** Build and index with the sequence of processing steps described in
+ *  the string. */
+Index *index_factory (int d, const char *description,
+                      MetricType metric = METRIC_L2);
+
+
+
+} // namespace faiss
+
+
+
+#endif

AuxIndexStructures.cpp

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+// -*- c++ -*-
+
+#include "AuxIndexStructures.h"
+
+#include <cstring>
+
+namespace faiss {
+
+
+/***********************************************************************
+ * RangeSearchResult
+ ***********************************************************************/
+
+RangeSearchResult::RangeSearchResult (size_t nq): nq (nq) {
+    lims = new size_t [nq + 1];
+    memset (lims, 0, sizeof(*lims) * (nq + 1));
+    labels = nullptr;
+    distances = nullptr;
+    buffer_size = 1024 * 256;
+}
+
+/// called when lims contains the nb of elements result entries
+/// for each query
+void RangeSearchResult::do_allocation () {
+    size_t ofs = 0;
+    for (int i = 0; i < nq; i++) {
+        size_t n = lims[i];
+        lims [i] = ofs;
+        ofs += n;
+    }
+    lims [nq] = ofs;
+    labels = new idx_t [ofs];
+    distances = new float [ofs];
+}
+
+RangeSearchResult::~RangeSearchResult () {
+    delete [] labels;
+    delete [] distances;
+    delete [] lims;
+}
+
+/***********************************************************************
+ * BufferList
+ ***********************************************************************/
+
+
+BufferList::BufferList (size_t buffer_size):
+    buffer_size (buffer_size)
+{
+    wp = buffer_size;
+}
+
+BufferList::~BufferList ()
+{
+    for (int i = 0; i < buffers.size(); i++) {
+        delete [] buffers[i].ids;
+        delete [] buffers[i].dis;
+    }
+}
+
+
+
+void BufferList::append_buffer ()
+{
+        Buffer buf = {new idx_t [buffer_size], new float [buffer_size]};
+        buffers.push_back (buf);
+        wp = 0;
+}
+
+/// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
+/// tables dest_ids, dest_dis
+void BufferList::copy_range (size_t ofs, size_t n,
+                             idx_t * dest_ids, float *dest_dis)
+{
+    size_t bno = ofs / buffer_size;
+    ofs -= bno * buffer_size;
+    while (n > 0) {
+        size_t ncopy = ofs + n < buffer_size ? n : buffer_size - ofs;
+        Buffer buf = buffers [bno];
+        memcpy (dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
+        memcpy (dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
+        dest_ids += ncopy;
+        dest_dis += ncopy;
+        ofs = 0;
+        bno ++;
+            n -= ncopy;
+    }
+}
+
+
+/***********************************************************************
+ * RangeSearchPartialResult
+ ***********************************************************************/
+
+
+RangeSearchPartialResult::RangeSearchPartialResult (RangeSearchResult * res_in):
+    BufferList(res_in->buffer_size),
+    res(res_in)
+{}
+
+
+/// begin a new result
+RangeSearchPartialResult::QueryResult &
+    RangeSearchPartialResult::new_result (idx_t qno)
+{
+    QueryResult qres = {qno, 0, this};
+    queries.push_back (qres);
+    return queries.back();
+}
+
+
+void RangeSearchPartialResult::finalize ()
+{
+    set_lims ();
+#pragma omp barrier
+
+#pragma omp single
+    res->do_allocation ();
+
+#pragma omp barrier
+    set_result ();
+}
+
+
+/// called by range_search before do_allocation
+void RangeSearchPartialResult::set_lims ()
+{
+    for (int i = 0; i < queries.size(); i++) {
+        QueryResult & qres = queries[i];
+        res->lims[qres.qno] = qres.nres;
+    }
+}
+
+/// called by range_search after do_allocation
+void RangeSearchPartialResult::set_result (bool incremental)
+{
+    size_t ofs = 0;
+    for (int i = 0; i < queries.size(); i++) {
+        QueryResult & qres = queries[i];
+
+        copy_range (ofs, qres.nres,
+                    res->labels + res->lims[qres.qno],
+                    res->distances + res->lims[qres.qno]);
+        if (incremental) {
+            res->lims[qres.qno] += qres.nres;
+        }
+        ofs += qres.nres;
+    }
+}
+
+
+IDSelectorRange::IDSelectorRange (idx_t imin, idx_t imax):
+    imin (imin), imax (imax)
+{
+}
+
+bool IDSelectorRange::is_member (idx_t id) const
+{
+    return id >= imin && id < imax;
+}
+
+
+
+IDSelectorBatch::IDSelectorBatch (long n, const idx_t *indices)
+{
+    nbits = 0;
+    while (n > (1L << nbits)) nbits++;
+    nbits += 5;
+    // for n = 1M, nbits = 25 is optimal, see P56659518
+
+    mask = (1L << nbits) - 1;
+    bloom.resize (1UL << (nbits - 3), 0);
+    for (long i = 0; i < n; i++) {
+        long id = indices[i];
+        set.insert(id);
+        id &= mask;
+        bloom[id >> 3] |= 1 << (id & 7);
+    }
+}
+
+bool IDSelectorBatch::is_member (idx_t i) const
+{
+    long im = i & mask;
+    if(!(bloom[im>>3] & (1 << (im & 7)))) {
+        return 0;
+    }
+    return set.count(i);
+}
+
+
+
+
+
+
+
+}

AuxIndexStructures.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+// -*- c++ -*-
+// Auxiliary index structures, that are used in indexes but that can
+// be forward-declared
+
+#ifndef FAISS_AUX_INDEX_STRUCTURES_H
+#define FAISS_AUX_INDEX_STRUCTURES_H
+
+#include <vector>
+
+#if __cplusplus >= 201103L
+#include <unordered_set>
+#endif
+
+#include <set>
+
+
+#include "Index.h"
+
+namespace faiss {
+
+/** The objective is to have a simple result structure while
+ *  minimizing the number of mem copies in the result. The method
+ *  do_allocation can be overloaded to allocate the result tables in
+ *  the matrix type of a srcipting language like Lua or Python. */
+struct RangeSearchResult {
+    size_t nq;      ///< nb of queries
+    size_t *lims;   ///< size (nq + 1)
+
+    typedef Index::idx_t idx_t;
+
+    idx_t *labels;     ///< result for query i is labels[lims[i]:lims[i+1]]
+    float *distances;  ///< corresponding distances (not sorted)
+
+    size_t buffer_size; ///< size of the result buffers used
+
+    /// lims must be allocated on input to range_search.
+    explicit RangeSearchResult (size_t nq);
+
+    /// called when lims contains the nb of elements result entries
+    /// for each query
+    virtual void do_allocation ();
+
+    virtual ~RangeSearchResult ();
+};
+
+
+/** Encapsulates a set of ids to remove. */
+struct IDSelector {
+    typedef Index::idx_t idx_t;
+    virtual bool is_member (idx_t id) const = 0;
+    virtual ~IDSelector() {}
+};
+
+
+
+/** remove ids between [imni, imax) */
+struct IDSelectorRange: IDSelector {
+    idx_t imin, imax;
+
+    IDSelectorRange (idx_t imin, idx_t imax);
+    virtual bool is_member (idx_t id) const override;
+    virtual ~IDSelectorRange () {}
+
+};
+
+
+/** Remove ids from a set. Repetitions of ids in the indices set
+ * passed to the constructor does not hurt performance. The hash
+ * function used for the bloom filter and GCC's implementation of
+ * unordered_set are just the least significant bits of the id. This
+ * works fine for random ids or ids in sequences but will produce many
+ * hash collisions if lsb's are always the same */
+struct IDSelectorBatch: IDSelector {
+
+#if __cplusplus >= 201103L
+    std::unordered_set<idx_t> set;
+#else
+    std::set<idx_t> set;
+#endif
+
+    typedef unsigned char uint8_t;
+    std::vector<uint8_t> bloom; // assumes low bits of id are a good hash value
+    int nbits;
+    idx_t mask;
+
+    IDSelectorBatch (long n, const idx_t *indices);
+    virtual bool is_member (idx_t id) const override;
+    virtual ~IDSelectorBatch() {}
+
+};
+
+
+// Below are structures used only by Index implementations
+
+
+
+/** List of temporary buffers used to store results before they are
+ *  copied to the RangeSearchResult object. */
+struct BufferList {
+    typedef Index::idx_t idx_t;
+
+    // buffer sizes in # entries
+    size_t buffer_size;
+
+    struct Buffer {
+        idx_t *ids;
+        float *dis;
+    };
+
+    std::vector<Buffer> buffers;
+    size_t wp; ///< write pointer in the last buffer.
+
+    explicit BufferList (size_t buffer_size);
+
+    ~BufferList ();
+
+    // create a new buffer
+    void append_buffer ();
+
+    inline void add (idx_t id, float dis)
+    {
+        if (wp == buffer_size) { // need new buffer
+            append_buffer();
+        }
+        Buffer & buf = buffers.back();
+        buf.ids [wp] = id;
+        buf.dis [wp] = dis;
+        wp++;
+    }
+
+    /// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
+    /// tables dest_ids, dest_dis
+    void copy_range (size_t ofs, size_t n,
+                     idx_t * dest_ids, float *dest_dis);
+
+};
+
+
+
+/// the entries in the buffers are split per query
+struct RangeSearchPartialResult: BufferList {
+    RangeSearchResult * res;
+
+    explicit RangeSearchPartialResult (RangeSearchResult * res_in);
+
+    /// result structure for a single query
+    struct QueryResult {
+        idx_t qno;
+        size_t nres;
+        RangeSearchPartialResult * pres;
+        inline void add (float dis, idx_t id) {
+            nres++;
+            pres->add (id, dis);
+        }
+    };
+
+    std::vector<QueryResult> queries;
+
+    /// begin a new result
+    QueryResult & new_result (idx_t qno);
+
+
+    void finalize ();
+
+
+    /// called by range_search before do_allocation
+    void set_lims ();
+
+    /// called by range_search after do_allocation
+    void set_result (bool incremental = false);
+
+};
+
+
+}; // namespace faiss
+
+
+
+#endif

CONTRIBUTING.md

+# Contributing to Faiss
+
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Our Development Process
+
+We mainly develop Faiss within Facebook. Sometimes, we will sync the 
+github version of Faiss with the internal state.
+
+## Pull Requests
+
+We welcome pull requests that add significant value to Faiss. 
+
+1. Fork the repo and create your branch from `master`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+There is a Facebook internal test suite for Faiss, and we need to run 
+all changes to Faiss through it.
+
+## Contributor License Agreement ("CLA")
+
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## Coding Style  
+* 4 spaces for indentation rather than tabs
+* 80 character line length
+* C++03 for the main Faiss
+
+
+## License
+By contributing to Faiss, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.
+
+
+

Clustering.cpp

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+/* Copyright 2004-present Facebook. All Rights Reserved.
+   kmeans clustering routines
+*/
+
+#include "Clustering.h"
+
+
+
+#include <cmath>
+#include <cstdio>
+#include <cstring>
+
+#include "utils.h"
+#include "FaissAssert.h"
+#include "IndexFlat.h"
+
+namespace faiss {
+
+ClusteringParameters::ClusteringParameters ():
+    niter(25),
+    nredo(1),
+    verbose(false), spherical(false),
+    update_index(false),
+    min_points_per_centroid(39),
+    max_points_per_centroid(256),
+    seed(1234)
+{}
+// 39 corresponds to 10000 / 256 -> to avoid warnings on PQ tests with randu10k
+
+
+Clustering::Clustering (int d, int k):
+    d(d), k(k) {}
+
+Clustering::Clustering (int d, int k, const ClusteringParameters &cp):
+    ClusteringParameters (cp), d(d), k(k) {}
+
+
+
+static double imbalance_factor (int n, int k, long *assign) {
+    std::vector<int> hist(k, 0);
+    for (int i = 0; i < n; i++)
+        hist[assign[i]]++;
+
+    double tot = 0, uf = 0;
+
+    for (int i = 0 ; i < k ; i++) {
+        tot += hist[i];
+        uf += hist[i] * (double) hist[i];
+    }
+    uf = uf * k / (tot * tot);
+
+    return uf;
+}
+
+
+
+
+void Clustering::train (idx_t nx, const float *x_in, Index & index) {
+    FAISS_ASSERT (nx >= k ||
+                  !"need at least as many training points as clusters");
+
+    double t0 = getmillisecs();
+
+    // yes it is the user's responsibility, but it may spare us some
+    // hard-to-debug reports.
+    for (size_t i = 0; i < nx * d; i++) {
+        FAISS_ASSERT (finite (x_in[i]) ||
+                      !"input contains NaN's or Inf's");
+    }
+
+    const float *x = x_in;
+
+    if (nx > k * max_points_per_centroid) {
+        if (verbose)
+            printf("Sampling a subset of %ld / %ld for training\n",
+                   k * max_points_per_centroid, nx);
+        int *perm = new int[nx];
+        rand_perm (perm, nx, seed);
+        nx = k * max_points_per_centroid;
+        float * x_new = new float [nx * d];
+        for (idx_t i = 0; i < nx; i++)
+            memcpy (x_new + i * d, x + perm[i] * d, sizeof(x_new[0]) * d);
+        delete [] perm;
+        x = x_new;
+    } else if (nx < k * min_points_per_centroid) {
+        fprintf (stderr,
+                 "WARNING clustering %ld points to %ld centroids: "
+                 "please provide at least %ld training points\n",
+                 nx, k, idx_t(k) * min_points_per_centroid);
+    }
+
+    if (verbose)
+        printf("Clustering %d points in %ldD to %ld clusters, "
+               "redo %d times, %d iterations\n",
+               int(nx), d, k, nredo, niter);
+
+
+
+    idx_t * assign = new idx_t[nx];
+    float * dis = new float[nx];
+
+    float best_err = 1e50;
+    double t_search_tot = 0;
+    if (verbose) {
+        printf("  Preprocessing in %5g s\n",
+               (getmillisecs() - t0)/1000.);
+    }
+    t0 = getmillisecs();
+
+    for (int redo = 0; redo < nredo; redo++) {
+
+        std::vector<float> buf_centroids;
+
+        std::vector<float> &cur_centroids =
+            nredo == 1 ? centroids : buf_centroids;
+
+        if (verbose && nredo > 1) {
+            printf("Outer iteration %d / %d\n", redo, nredo);
+        }
+
+        if (cur_centroids.size() == 0) {
+            // initialize centroids with random points from the dataset
+            cur_centroids.resize (d * k);
+            int *perm = new int[nx];
+            rand_perm (perm, nx, seed + 1 + redo * 15486557L);
+#pragma omp parallel for
+            for (int i = 0; i < k ; i++)
+                memcpy (&cur_centroids[i * d], x + perm[i] * d,
+                        d * sizeof (float));
+            delete [] perm;
+        } else { // assume user provides some meaningful initialization
+            FAISS_ASSERT (cur_centroids.size() == d * k);
+            FAISS_ASSERT (nredo == 1 ||
+                          !"will redo with same initialization");
+        }
+
+        if (spherical)
+            fvec_renorm_L2 (d, k, cur_centroids.data());
+
+        if (!index.is_trained)
+            index.train (k, cur_centroids.data());
+
+        FAISS_ASSERT (index.ntotal == 0 );
+        index.add (k, cur_centroids.data());
+        float err = 0;
+        for (int i = 0; i < niter; i++) {
+            double t0s = getmillisecs();
+            index.search (nx, x, 1, dis, assign);
+            t_search_tot += getmillisecs() - t0s;
+
+            err = 0;
+            for (int j = 0; j < nx; j++)
+                err += dis[j];
+            obj.push_back (err);
+
+            int nsplit = km_update_centroids (x, cur_centroids.data(),
+                                              assign, d, k, nx);
+
+            if (verbose) {
+                printf ("  Iteration %d (%5g s, search %5g s): "
+                        "objective=%g imbalance=%g nsplit=%d       \r",
+                        i, (getmillisecs() - t0) / 1000.0,
+                        t_search_tot / 1000,
+                        err, imbalance_factor (nx, k, assign),
+                        nsplit);
+                fflush (stdout);
+            }
+
+            if (spherical)
+                fvec_renorm_L2 (d, k, cur_centroids.data());
+
+            index.reset ();
+            if (update_index)
+                index.train (k, cur_centroids.data());
+
+            assert (index.ntotal == 0);
+            index.add (k, centroids.data());
+        }
+        if (verbose) printf("\n");
+        if (nredo > 1) {
+            if (err < best_err) {
+                if (verbose)
+                    printf ("Keep new clusters\n");
+                centroids = cur_centroids;
+                best_err = err;
+            }
+        }
+    }
+
+    delete [] assign;
+    delete [] dis;
+    if (x_in != x) delete [] x;
+}
+
+float kmeans_clustering (size_t d, size_t n, size_t k,
+                         const float *x,
+                         float *centroids)
+{
+    Clustering clus (d, k);
+    clus.verbose = d * n * k > (1L << 30);
+    // display logs if > 1Gflop per iteration
+    IndexFlatL2 index (d);
+    clus.train (n, x, index);
+    memcpy(centroids, clus.centroids.data(), sizeof(*centroids) * d * k);
+    return clus.obj.back();
+}
+
+} // namespace faiss

Clustering.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+// -*- c++ -*-
+
+#ifndef FAISS_CLUSTERING_H
+#define FAISS_CLUSTERING_H
+#include "Index.h"
+
+#include <vector>
+
+namespace faiss {
+
+
+/** Class for the clustering parameters. Can be passed to the
+ * constructor of the Clustering object.
+ */
+struct ClusteringParameters {
+    int niter;          ///< clustering iterations
+    int nredo;          ///< redo clustering this many times and keep best
+
+
+    bool verbose;
+    bool spherical;     ///< do we want normalized centroids?
+    bool update_index;  ///< update index after each iteration?
+
+    int min_points_per_centroid; ///< otherwise you get a warning
+    int max_points_per_centroid;  ///< to limit size of dataset
+
+    int seed; ///< seed for the random number generator
+
+    /// sets reasonable defaults
+    ClusteringParameters ();
+};
+
+
+/** clustering based on assignment - centroid update iterations
+ *
+ * The clustering is based on an Index object that assigns training
+ * points to the centroids. Therefore, at each iteration the centroids
+ * are added to the index.
+ *
+ * On output, the centoids table is set to the latest version
+ * of the centroids and they are also added to the index. If the
+ * centroids table it is not empty on input, it is also used for
+ * initialization.
+ *
+ * To do several clusterings, just call train() several times on
+ * different training sets, clearing the centroid table in between.
+ */
+struct Clustering: ClusteringParameters {
+    typedef Index::idx_t idx_t;
+    size_t d;              ///< dimension of the vectors
+    size_t k;              ///< nb of centroids
+
+    /// centroids (k * d)
+    std::vector<float> centroids;
+
+    /// objective values (sum of distances reported by index) over
+    /// iterations
+    std::vector<float> obj;
+
+    /// the only mandatory parameters are k and d
+    Clustering (int d, int k);
+    Clustering (int d, int k, const ClusteringParameters &cp);
+
+    /// Index is used during the assignment stage
+    virtual void train (idx_t n, const float * x, faiss::Index & index);
+
+    virtual ~Clustering() {}
+};
+
+
+/** simplified interface
+ *
+ * @param d dimension of the data
+ * @param n nb of training vectors
+ * @param k nb of output centroids
+ * @param x training set (size n * d)
+ * @param centroids output centroids (size k * d)
+ * @return final quantization error
+ */
+float kmeans_clustering (size_t d, size_t n, size_t k,
+                         const float *x,
+                         float *centroids);
+
+
+
+}
+
+
+#endif

FaissAssert.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+
+#ifndef FAISS_ASSERT_INCLUDED
+#define FAISS_ASSERT_INCLUDED
+
+#include <cstdlib>
+#include <cstdio>
+
+/// Asserts that risk to be triggered by user input
+#define FAISS_ASSERT(X) ({if (! (X)) {                             \
+    fprintf (stderr, "Faiss assertion %s failed in %s at %s:%d", \
+             #X, __PRETTY_FUNCTION__, __FILE__, __LINE__);       \
+    abort(); }})
+
+
+
+#endif

Heap.cpp

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+/* Copyright 2004-present Facebook. All Rights Reserved. */
+/* Function for soft heap */
+
+#include "Heap.h"
+
+
+namespace faiss {
+
+
+
+
+
+
+
+
+
+template <typename C>
+void HeapArray<C>::heapify ()
+{
+#pragma omp parallel for
+    for (size_t j = 0; j < nh; j++)
+        heap_heapify<C> (k, val + j * k, ids + j * k);
+}
+
+template <typename C>
+void HeapArray<C>::reorder ()
+{
+#pragma omp parallel for
+    for (size_t j = 0; j < nh; j++)
+        heap_reorder<C> (k, val + j * k, ids + j * k);
+}
+
+template <typename C>
+void HeapArray<C>::addn (size_t nj, const T *vin, TI j0,
+                         size_t i0, long ni)
+{
+    if (ni == -1) ni = nh;
+    assert (i0 >= 0 && i0 + ni <= nh);
+#pragma omp parallel for
+    for (size_t i = i0; i < i0 + ni; i++) {
+        T * __restrict simi = get_val(i);
+        TI * __restrict idxi = get_ids (i);
+        const T *ip_line = vin + (i - i0) * nj;
+
+        for (size_t j = 0; j < nj; j++) {
+            T ip = ip_line [j];
+            if (C::cmp(simi[0], ip)) {
+                heap_pop<C> (k, simi, idxi);
+                heap_push<C> (k, simi, idxi, ip, j + j0);
+            }
+        }
+    }
+}
+
+template <typename C>
+void HeapArray<C>::addn_with_ids (
+     size_t nj, const T *vin, const TI *id_in,
+     long id_stride, size_t i0, long ni)
+{
+    if (id_in == nullptr) {
+        addn (nj, vin, 0, i0, ni);
+        return;
+    }
+    if (ni == -1) ni = nh;
+    assert (i0 >= 0 && i0 + ni <= nh);
+#pragma omp parallel for
+    for (size_t i = i0; i < i0 + ni; i++) {
+        T * __restrict simi = get_val(i);
+        TI * __restrict idxi = get_ids (i);
+        const T *ip_line = vin + (i - i0) * nj;
+        const TI *id_line = id_in + (i - i0) * id_stride;
+
+        for (size_t j = 0; j < nj; j++) {
+            T ip = ip_line [j];
+            if (C::cmp(simi[0], ip)) {
+                heap_pop<C> (k, simi, idxi);
+                heap_push<C> (k, simi, idxi, ip, id_line [j]);
+            }
+        }
+    }
+}
+
+template <typename C>
+void HeapArray<C>::per_line_extrema (
+                   T * out_val,
+                   TI * out_ids) const
+{
+#pragma omp parallel for
+    for (size_t j = 0; j < nh; j++) {
+        long imin = -1;
+        typename C::T xval = C::Crev::neutral ();
+        const typename C::T * x_ = val + j * k;
+        for (size_t i = 0; i < k; i++)
+            if (C::cmp (x_[i], xval)) {
+                xval = x_[i];
+                imin = i;
+            }
+        if (out_val)
+            out_val[j] = xval;
+
+        if (out_ids) {
+            if (ids && imin != -1)
+                out_ids[j] = ids [j * k + imin];
+            else
+                out_ids[j] = imin;
+        }
+    }
+}
+
+
+
+
+// explicit instanciations
+
+template class HeapArray<CMin <float, long> >;
+template class HeapArray<CMax <float, long> >;
+template class HeapArray<CMin <int, long> >;
+template class HeapArray<CMax <int, long> >;
+
+
+
+
+
+}  // END namespace fasis

INSTALL

+
+**********************************************************
+** INSTALL file for Faiss (Fair AI Similarity Search).  **
+**********************************************************
+
+
+The Faiss installation works in 3 steps, from easiest to most
+involved:
+
+1. compile the C++ core and examples
+
+2. compile the Python interface
+
+3. compile GPU part
+
+Steps 2 and 3 depend on 1, but they are otherwise independent.
+
+General compilation instructions
+================================
+
+Faiss has been tested only on x86_64 machines on Linux and Mac OS.
+
+Faiss is compiled via a Makefile. The system-dependent configuration
+of the Makefile is in an include file, makefile.inc. The variables in
+makefile.inc must be set by hand.
+
+Faiss requires a C++ compiler that understands:
+- the Intel intrinsics for SSE instructions
+- the GCC intrinsic for the popcount instruction
+- basic OpenMP
+
+There are a few models for makefile.inc in the example_makefiles/
+subdirectory. Copy the relevant one for your system and adjust to your
+needs.
+
+Faiss comes as a .a archive, that can be linked with executables or
+dynamic libraries (useful for the Python wrapper).
+
+Step 1: Compiling the C++ Faiss
+===============================
+
+The CPU version of Faiss is written in C++03, so it should compile
+even with relatively old C++ compilers.
+
+BLAS/Lapack
+-----------
+
+The only variables that need to be configured for the C++ Faiss are
+the BLAS/Lapack flags (a linear aglebra software package). It needs a
+flag telling whether BLAS/Lapack uses 32 or 64 bit integers and the
+linking flags. Faiss uses the Fortran 77 interface of BLAS/Lapack and
+thus does not need an include path.
+
+There are several BLAS implementations, depending on the OS and
+machine. To have reasonable performance, the BLAS library should be
+multithreaded. See the example makefile.inc's for hints and examples
+on how to set the flags.
+
+To check that the link flags are correct, and verify whether the
+implementation uses 32 or 64 bit integers, you can
+
+  make tests/test_blas
+
+and run
+
+  ./tests/test_blas
+
+Testing Faiss
+-------------
+
+Once the proper BLAS flags are set, the library should compile
+smoothly by running
+
+  make
+
+A basic usage example is in
+
+  tests/demo_ivfpq_indexing
+
+it makes a small index, stores it and performs some searches. A normal
+runtime is around 20s. With a fast machine and Intel MKL's BLAS it
+runs in 2.5s.
+
+A real-life benchmark
+---------------------
+
+A bit longer example runs and evaluates Faiss on the SIFT1M
+dataset. To run it, please download the ANN_SIFT1M dataset from
+
+http://corpus-texmex.irisa.fr/
+
+and unzip it to the sudirectory sift1M.
+
+Then compile and run
+
+  make tests/demo_sift1M
+  tests/demo_sift1M
+
+This is a demonstration of the high-level auto-tuning API. You can try
+setting a different index_key to find the indexing structure that
+gives the best performance.
+
+
+Step 2: Compiling the Python interface
+======================================
+
+The Python interface is provided via SWIG (Simple Wrapper and
+Interface Generator) and an additional level of manual wrappers (in faiss.py).
+
+SWIG generates two wrapper files: a Python file (swigfaiss.py) and a
+C++ file that must be compiled to a dynamic library (_swigfaiss.so).
+
+The C++ compilation to the dynamic library requires to set:
+
+- SHAREDFLAGS: system-specific flags to generate a dynamic library
+
+- PYTHONCFLAGS: include flags for Python
+
+See the example makefile.inc's on how to set the flags.
+
+
+Testing the Python wrapper
+--------------------------
+
+Often, a successful compile does not mean that the library works,
+because missing symbols are detected only at runtime. You should be
+able to load the Faiss dynamic library:
+
+  python -c "import faiss"
+
+In case of failure, it reports the first missing symbol. To see all
+missing symbols (on Linux), use
+
+  ldd -r _swigfaiss.so
+
+Sometimes, problems (eg with BLAS libraries) appear only when actually
+calling a BLAS function. A simple way to check this
+
+python -c "import faiss, numpy
+faiss.Kmeans(10, 20).train(numpy.random.rand(1000, 10).astype('float32'))"
+
+
+Real-life test
+--------------
+
+The following script extends the demo_sift1M test to several types of
+indexes:
+
+  python python/demo_auto_tune.py
+
+It will cycle through a few types of indexes and find optimal
+operating points. You can play around with the types of indexes.
+
+
+Step 3: Compiling the GPU implementation
+========================================
+
+There is a GPU-specific Makefile in the gpu/ directory. It depends on
+the same ../makefile.inc for system-specific variables. You need
+libfaiss.a from Step 1 for this to work.
+
+The GPU version is a superset of the CPU version. In addition it
+requires:
+
+- a C++11 compliant compiler (and flags)
+
+- the cuda compiler and related libraries (Cublas)
+
+See the example makefile on how to set the flags.
+
+The nvcc-specific flags to pass to the compiler:
+
+   -gencode arch=compute_35,code="compute_35"
+   -gencode arch=compute_52,code="compute_52"
+   -gencode arch=compute_60,code="compute_60"
+
+Most other flags are related to the C++11 compiler used by nvcc to
+complile the actual C++ code. They are normally just transmitted by
+nvcc, except some of them that are not recognized and that should be
+escaped by prefixing them with -Xcompiler. Also link flags that are
+prefixed with -Wl, should be passed with -Xlinker.
+
+Then compile with
+
+  cd gpu; make
+
+You may want to add '-j 10' to use 10 threads during compile.
+
+Testing the GPU implementation
+------------------------------
+
+Compile the example with
+
+  cd gpu; make test/demo_ivfpq_indexing_gpu
+
+This produce the GPU code equivalent to the CPU
+demo_ivfpq_indexing. It also shows how to translate indexed from/to
+the GPU.
+
+Compiling the Python interface with GPU support
+-----------------------------------------------
+
+Given step 2, adding support of the GPU from Python is quite
+straightforward. Run
+
+  cd gpu; make py
+
+The import is the same for the GPU version and the CPU-only
+version.
+
+  python -c "import faiss"
+
+Faiss tries to load the GPU version first, and in case of failure,
+loads the CPU-only version. To investigate more closely the cause of
+a failure, you can run:
+
+  python -c "import _swigfaiss_gpu"
+
+Python example with GPU support
+-------------------------------
+
+The auto-tuning example above also runs on the GPU. Edit
+
+  python/demo_auto_tune.py
+
+to enable and run it.
+
+Hot to use Faiss in your own projects
+=====================================
+
+C++ 
+---
+
+The makefile generates a static and a dynamic library
+
+  libfaiss.a
+  libfaiss.so (or libfaiss.dylib)
+
+the executable should be linked to one of these. If you use 
+the static version (.a), add the LDFLAGS found in the Makefile.
+
+Python
+------
+
+To import Faiss in your own Python project, you need the files 
+
+  faiss.py 
+  swigfaiss.py  / swigfaiss_gpu.py
+  _swigfaiss.so / _swigfaiss_gpu.so 
+
+to be visible in the PYTHONPATH or in the current directory. 
+Then Faiss can be used in python with
+
+  import faiss
+

Index.cpp

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+
+#include "IndexFlat.h"
+#include "FaissAssert.h"
+
+namespace faiss {
+
+
+void Index::range_search (idx_t , const float *, float,
+                          RangeSearchResult *) const
+{
+    FAISS_ASSERT (!"range search not implemented");
+}
+
+void Index::assign (idx_t n, const float * x, idx_t * labels, idx_t k)
+{
+    float * distances = new float[n * k];
+    search (n, x, k, distances, labels);
+    delete [] distances;
+}
+
+
+void Index::add_with_ids (idx_t n, const float * x, const long *xids)
+{
+    FAISS_ASSERT (!"add_with_ids not implemented for this type of index");
+}
+
+
+long Index::remove_ids (const IDSelector & sel)
+{
+    FAISS_ASSERT (!"remove_ids not implemented for this type of index");
+    return -1;
+}
+
+
+void Index::reconstruct (idx_t, float * ) const {
+    FAISS_ASSERT (! "Can not compute reconstruct without knowing howto\n");
+}
+
+
+void Index::reconstruct_n (idx_t i0, idx_t ni, float *recons) const {
+    for (idx_t i = 0; i < ni; i++) {
+        reconstruct (i0 + i, recons + i * d);
+    }
+}
+
+
+
+void Index::compute_residual (const float * x,
+                              float * residual, idx_t key) const {
+    reconstruct (key, residual);
+    for (size_t i = 0; i < d; i++)
+        residual[i] = x[i] - residual[i];
+}
+
+
+void Index::display () const {
+    printf ("Index: %s  -> %ld elements\n", typeid (*this).name(), ntotal);
+}
+
+
+
+
+}

Index.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+// -*- c++ -*-
+
+#ifndef FAISS_INDEX_H
+#define FAISS_INDEX_H
+
+
+#include <cstdio>
+#include <typeinfo>
+#include <string>
+#include <sstream>
+
+
+/**
+ * @namespace faiss
+ *
+ * Throughout the library, vectors are provided as float * pointers.
+ * Most algorithms can be optimized when several vectors are processed
+ * (added/searched) together in a batch. In this case, they are passed
+ * in as a matrix. When n vectors of size d are provided as float * x,
+ * component j of vector i is
+ *
+ *   x[ i * d + j ]
+ *
+ * where 0 <= i < n and 0 <= j < d. In other words, matrices are
+ * always compact. When specifying the size of the matrix, we call it
+ * an n*d matrix, which implies a row-major storage.
+ */
+
+
+namespace faiss {
+
+
+/// Some algorithms support both an inner product vetsion and a L2 search version.
+enum MetricType {
+    METRIC_INNER_PRODUCT = 0,
+    METRIC_L2 = 1,
+};
+
+
+/// Forward declarations see AuxIndexStructures.h
+struct IDSelector;
+struct RangeSearchResult;
+
+/** Abstract structure for an index
+ *
+ * Supports adding vertices and searching them.
+ *
+ * Currently only asymmetric queries are supported:
+ * database-to-database queries are not implemented.
+ */
+struct Index {
+    std::string index_typename;
+
+    typedef long idx_t;    ///< all indices are this type
+
+    int d;                 ///< vector dimension
+    idx_t ntotal;          ///< total nb of indexed vectors
+    bool verbose;          ///< verbosity level
+
+    /// set if the Index does not require training, or if training is done already
+    bool is_trained;
+
+    /// type of metric this index uses for search
+    MetricType metric_type;
+
+    explicit Index (idx_t d = 0, MetricType metric = METRIC_INNER_PRODUCT):
+                    index_typename ("Undefined Index typename"),
+                    d(d),
+                    ntotal(0),
+                    verbose(false),
+                    is_trained(true),
+                    metric_type (metric) {}
+
+    virtual ~Index () {  }
+
+
+    /** Perform training on a representative set of vectors
+     *
+     * @param n      nb of training vectors
+     * @param x      training vecors, size n * d
+     */
+    virtual void train (idx_t n, const float *x) {
+        // does nothing by default
+    }
+
+    /** Add n vectors of dimension d to the index.
+     *
+     * Vectors are implicitly assigned labels ntotal .. ntotal + n - 1
+     * This function slices the input vectors in chuncks smaller than
+     * blocksize_add and calls add_core.
+     * @param x      input matrix, size n * d
+     */
+    virtual void add (idx_t n, const float *x) = 0;
+
+    /** Same as add, but stores xids instead of sequential ids.
+     *
+     * The default implementation fails with an assertion, as it is
+     * not supported by all indexes.
+     *
+     * @param xids if non-null, ids to store for the vectors (size n)
+     */
+    virtual void add_with_ids (idx_t n, const float * x, const long *xids);
+
+    /** query n vectors of dimension d to the index.
+     *
+     * return at most k vectors. If there are not enough results for a
+     * query, the result array is padded with -1s.
+     *
+     * @param x           input vectors to search, size n * d
+     * @param labels      output labels of the NNs, size n*k
+     * @param distances   output pairwise distances, size n*k
+     */
+    virtual void search (idx_t n, const float *x, idx_t k,
+                         float *distances, idx_t *labels) const = 0;
+
+    /** query n vectors of dimension d to the index.
+     *
+     * return all vectors with distance < radius. Note that many
+     * indexes do not implement the range_search (only the k-NN search
+     * is mandatory).
+     *
+     * @param x           input vectors to search, size n * d
+     * @param radius      search radius
+     * @param result      result table
+     */
+    virtual void range_search (idx_t n, const float *x, float radius,
+                               RangeSearchResult *result) const;
+
+    /** return the indexes of the k vectors closest to the query x.
+     *
+     * This function is identical as search but only return labels of neighbors.
+     * @param x           input vectors to search, size n * d
+     * @param labels      output labels of the NNs, size n*k
+     */
+    void assign (idx_t n, const float * x, idx_t * labels, idx_t k = 1);
+
+    /// removes all elements from the database.
+    virtual void reset() = 0;
+
+    /** removes IDs from the index. Not supported by all indexes
+     */
+    virtual long remove_ids (const IDSelector & sel);
+
+    /** Reconstruct a stored vector (or an approximation if lossy coding)
+     *
+     * this function may not be defined for some indexes
+     * @param key         id of the vector to reconstruct
+     * @param recons      reconstucted vector (size d)
+     */
+    virtual void reconstruct (idx_t key, float * recons) const;
+
+
+    /** Reconstruct vectors i0 to i0 + ni - 1
+     *
+     * this function may not be defined for some indexes
+     * @param recons      reconstucted vector (size ni * d)
+     */
+    virtual void reconstruct_n (idx_t i0, idx_t ni, float *recons) const;
+
+
+    /** Computes a residual vector after indexing encoding.
+     *
+     * The residual vector is the difference between a vector and the
+     * reconstruction that can be decoded from its representation in
+     * the index. The residual can be used for multiple-stage indexing
+     * methods, like IndexIVF's methods.
+     *
+     * @param x           input vector, size d
+     * @param residual    output residual vector, size d
+     * @param key         encoded index, as returned by search and assign
+     */
+    void compute_residual (const float * x, float * residual, idx_t key) const;
+
+    /** Display the actual class name and some more info */
+    void display () const;
+
+    /** Return the typeName of the index (which includes main parameters */
+    virtual std::string get_typename () const {
+        return index_typename; }
+
+    virtual void set_typename () = 0 ;
+
+
+};
+
+}
+
+
+#endif

IndexFlat.cpp

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+
+#include "IndexFlat.h"
+
+#include <cstring>
+#include "utils.h"
+#include "Heap.h"
+
+#include "FaissAssert.h"
+
+namespace faiss {
+
+IndexFlat::IndexFlat (idx_t d, MetricType metric):
+            Index(d, metric)
+{
+    set_typename();
+}
+
+
+void IndexFlat::set_typename()
+{
+    std::stringstream s;
+    if (metric_type == METRIC_INNER_PRODUCT)
+        s << "IP";
+    else if (metric_type == METRIC_L2)
+        s << "L2";
+    else s << "??";
+    index_typename = s.str();
+}
+
+
+void IndexFlat::add (idx_t n, const float *x) {
+    for (idx_t i = 0; i < n * d; i++)
+        xb.push_back (x[i]);
+    ntotal += n;
+}
+
+
+void IndexFlat::reset() {
+    xb.clear();
+    ntotal = 0;
+}
+
+
+void IndexFlat::search (idx_t n, const float *x, idx_t k,
+                               float *distances, idx_t *labels) const
+{
+    // we see the distances and labels as heaps
+
+    if (metric_type == METRIC_INNER_PRODUCT) {
+        float_minheap_array_t res = {
+            size_t(n), size_t(k), labels, distances};
+        knn_inner_product (x, xb.data(), d, n, ntotal, &res);
+    } else if (metric_type == METRIC_L2) {
+        float_maxheap_array_t res = {
+            size_t(n), size_t(k), labels, distances};
+        knn_L2sqr (x, xb.data(), d, n, ntotal, &res);
+    }
+}
+
+void IndexFlat::range_search (idx_t n, const float *x, float radius,
+                              RangeSearchResult *result) const
+{
+    switch (metric_type) {
+        case METRIC_INNER_PRODUCT:
+            range_search_inner_product (x, xb.data(), d, n, ntotal,
+                                        radius, result);
+            break;
+        case METRIC_L2:
+            range_search_L2sqr (x, xb.data(), d, n, ntotal, radius, result);
+            break;
+    }
+}
+
+
+void IndexFlat::compute_distance_subset (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            const idx_t *labels) const
+{
+    switch (metric_type) {
+        case METRIC_INNER_PRODUCT:
+            fvec_inner_products_by_idx (
+                 distances,
+                 x, xb.data(), labels, d, n, k);
+            break;
+        case METRIC_L2:
+            fvec_L2sqr_by_idx (
+                 distances,
+                 x, xb.data(), labels, d, n, k);
+            break;
+    }
+
+}
+
+
+
+void IndexFlat::reconstruct (idx_t key, float * recons) const
+{
+    memcpy (recons, &(xb[key * d]), sizeof(*recons) * d);
+}
+
+/***************************************************
+ * IndexFlatL2BaseShift
+ ***************************************************/
+
+IndexFlatL2BaseShift::IndexFlatL2BaseShift (idx_t d, size_t nshift, const float *shift):
+    IndexFlatL2 (d), shift (nshift)
+{
+    memcpy (this->shift.data(), shift, sizeof(float) * nshift);
+}
+
+void IndexFlatL2BaseShift::search (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            idx_t *labels) const
+{
+    FAISS_ASSERT(shift.size() == ntotal);
+
+    float_maxheap_array_t res = {
+        size_t(n), size_t(k), labels, distances};
+    knn_L2sqr_base_shift (x, xb.data(), d, n, ntotal, &res, shift.data());
+}
+
+
+
+/***************************************************
+ * IndexRefineFlat
+ ***************************************************/
+
+IndexRefineFlat::IndexRefineFlat (Index *base_index):
+    Index (base_index->d, base_index->metric_type),
+    refine_index (base_index->d, base_index->metric_type),
+    base_index (base_index), own_fields (false),
+    k_factor (1)
+{
+    is_trained = base_index->is_trained;
+    assert (base_index->ntotal == 0 ||
+            !"base_index should be empty in the beginning");
+    set_typename ();
+}
+
+IndexRefineFlat::IndexRefineFlat () {
+    base_index = nullptr;
+    own_fields = false;
+    k_factor = 1;
+}
+
+void IndexRefineFlat::set_typename ()
+{
+    std::stringstream s;
+    s << "Refine" << '[' << base_index->get_typename()
+      << ',' << refine_index.get_typename() << ']';
+    index_typename = s.str();
+}
+
+
+void IndexRefineFlat::train (idx_t n, const float *x)
+{
+    base_index->train (n, x);
+    is_trained = true;
+}
+
+void IndexRefineFlat::add (idx_t n, const float *x) {
+    FAISS_ASSERT (is_trained);
+    base_index->add (n, x);
+    refine_index.add (n, x);
+    ntotal = refine_index.ntotal;
+}
+
+void IndexRefineFlat::reset ()
+{
+    base_index->reset ();
+    refine_index.reset ();
+    ntotal = 0;
+}
+
+namespace {
+typedef faiss::Index::idx_t idx_t;
+
+template<class C>
+static void reorder_2_heaps (
+      idx_t n,
+      idx_t k, idx_t *labels, float *distances,
+      idx_t k_base, const idx_t *base_labels, const float *base_distances)
+{
+#pragma omp parallel for
+    for (idx_t i = 0; i < n; i++) {
+        idx_t *idxo = labels + i * k;
+        float *diso = distances + i * k;
+        const idx_t *idxi = base_labels + i * k_base;
+        const float *disi = base_distances + i * k_base;
+
+        heap_heapify<C> (k, diso, idxo, disi, idxi, k);
+        if (k_base != k) { // add remaining elements
+            heap_addn<C> (k, diso, idxo, disi + k, idxi + k, k_base - k);
+        }
+        heap_reorder<C> (k, diso, idxo);
+    }
+}
+
+
+}
+
+
+void IndexRefineFlat::search (
+              idx_t n, const float *x, idx_t k,
+              float *distances, idx_t *labels) const
+{
+    FAISS_ASSERT (is_trained);
+    idx_t k_base = idx_t (k * k_factor);
+    idx_t * base_labels = labels;
+    float * base_distances = distances;
+
+    if (k != k_base) {
+        base_labels = new idx_t [n * k_base];
+        base_distances = new float [n * k_base];
+    }
+
+    base_index->search (n, x, k_base, base_distances, base_labels);
+
+    for (int i = 0; i < n * k_base; i++)
+        assert (base_labels[i] >= -1 &&
+                base_labels[i] < ntotal);
+
+    // compute refined distances
+    refine_index.compute_distance_subset (
+        n, x, k_base, base_distances, base_labels);
+
+    // sort and store result
+    if (metric_type == METRIC_L2) {
+        typedef CMax <float, idx_t> C;
+        reorder_2_heaps<C> (
+            n, k, labels, distances,
+            k_base, base_labels, base_distances);
+
+    } else if (metric_type == METRIC_INNER_PRODUCT) {
+        typedef CMin <float, idx_t> C;
+        reorder_2_heaps<C> (
+            n, k, labels, distances,
+            k_base, base_labels, base_distances);
+    }
+
+    if (k != k_base) {
+        delete [] base_labels;
+        delete [] base_distances;
+    }
+}
+
+
+
+IndexRefineFlat::~IndexRefineFlat ()
+{
+    if (own_fields) delete base_index;
+}
+
+/***************************************************
+ * IndexFlat1D
+ ***************************************************/
+
+
+IndexFlat1D::IndexFlat1D (bool continuous_update):
+    IndexFlatL2 (1),
+    continuous_update (continuous_update)
+{
+}
+
+/// if not continuous_update, call this between the last add and
+/// the first search
+void IndexFlat1D::update_permutation ()
+{
+    perm.resize (ntotal);
+    if (ntotal < 1000000) {
+        fvec_argsort (ntotal, xb.data(), (size_t*)perm.data());
+    } else {
+        fvec_argsort_parallel (ntotal, xb.data(), (size_t*)perm.data());
+    }
+}
+
+void IndexFlat1D::add (idx_t n, const float *x)
+{
+    IndexFlatL2::add (n, x);
+    if (continuous_update)
+        update_permutation();
+}
+
+void IndexFlat1D::reset()
+{
+    IndexFlatL2::reset();
+    perm.clear();
+}
+
+void IndexFlat1D::search (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            idx_t *labels) const
+{
+    FAISS_ASSERT (perm.size() == ntotal ||
+                  !"Call update_permutation before search");
+
+#pragma omp parallel for
+    for (idx_t i = 0; i < n; i++) {
+
+        float q = x[i]; // query
+        float *D = distances + i * k;
+        idx_t *I = labels + i * k;
+
+        // binary search
+        idx_t i0 = 0, i1 = ntotal;
+        idx_t wp = 0;
+
+        if (xb[perm[i0]] > q) {
+            i1 = 0;
+            goto finish_right;
+        }
+
+        if (xb[perm[i1 - 1]] <= q) {
+            i0 = i1 - 1;
+            goto finish_left;
+        }
+
+        while (i0 + 1 < i1) {
+            idx_t imed = (i0 + i1) / 2;
+            if (xb[perm[imed]] <= q) i0 = imed;
+            else                    i1 = imed;
+        }
+
+        // query is between xb[perm[i0]] and xb[perm[i1]]
+        // expand to nearest neighs
+
+        while (wp < k) {
+            float xleft = xb[perm[i0]];
+            float xright = xb[perm[i1]];
+
+            if (q - xleft < xright - q) {
+                D[wp] = q - xleft;
+                I[wp] = perm[i0];
+                i0--; wp++;
+                if (i0 < 0) { goto finish_right; }
+            } else {
+                D[wp] = xright - q;
+                I[wp] = perm[i1];
+                i1++; wp++;
+                if (i1 >= ntotal) { goto finish_left; }
+            }
+        }
+        goto done;
+
+    finish_right:
+        // grow to the right from i1
+        while (wp < k) {
+            if (i1 < ntotal) {
+                D[wp] = xb[perm[i1]] - q;
+                I[wp] = perm[i1];
+                i1++;
+            } else {
+                D[wp] = 1.0 / 0.0;
+                I[wp] = -1;
+            }
+            wp++;
+        }
+        goto done;
+
+    finish_left:
+        // grow to the left from i0
+        while (wp < k) {
+            if (i0 >= 0) {
+                D[wp] = q - xb[perm[i0]];
+                I[wp] = perm[i0];
+                i0--;
+            } else {
+                D[wp] = 1.0 / 0.0;
+                I[wp] = -1;
+            }
+            wp++;
+        }
+    done:  ;
+    }
+
+}
+
+
+
+} // namespace faiss

IndexFlat.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+// -*- c++ -*-
+
+#ifndef INDEX_FLAT_H
+#define INDEX_FLAT_H
+
+#include <vector>
+
+#include "Index.h"
+
+
+namespace faiss {
+
+/** Index that stores the full vectors and performs exhaustive search */
+struct IndexFlat: Index {
+    /// database vectors, size ntotal * d
+    std::vector<float> xb;
+
+    explicit IndexFlat (idx_t d, MetricType metric = METRIC_INNER_PRODUCT);
+
+    virtual void set_typename() override;
+
+    virtual void add (idx_t n, const float *x) override;
+
+    virtual void reset() override;
+
+    virtual void search (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            idx_t *labels) const override;
+
+
+
+    virtual void range_search (
+            idx_t n,
+            const float *x,
+            float radius,
+            RangeSearchResult *result) const override;
+
+    virtual void reconstruct (idx_t key, float * recons)
+        const override;
+
+
+    /** compute distance with a subset of vectors
+     *
+     * @param x       query vectors, size n * d
+     * @param labels  indices of the vectors that should be compared
+     *                for each query vector, size n * k
+     * @param distances
+     *                corresponding output distances, size n * k
+     */
+    void compute_distance_subset (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            const idx_t *labels) const;
+
+    IndexFlat () {}
+};
+
+
+
+struct IndexFlatIP:IndexFlat {
+    explicit IndexFlatIP (idx_t d): IndexFlat (d, METRIC_INNER_PRODUCT) {}
+    IndexFlatIP () {}
+};
+
+
+struct IndexFlatL2:IndexFlat {
+    explicit IndexFlatL2 (idx_t d): IndexFlat (d, METRIC_L2) {}
+    IndexFlatL2 () {}
+};
+
+
+// same as an IndexFlatL2 but a value is subtracted from each distance
+struct IndexFlatL2BaseShift: IndexFlatL2 {
+    std::vector<float> shift;
+
+    IndexFlatL2BaseShift (idx_t d, size_t nshift, const float *shift);
+
+    virtual void search (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            idx_t *labels) const override;
+};
+
+
+/** Index that queries in a base_index (a fast one) and refines the
+ *  results with an exact search, hopefully improving the results.
+ */
+struct IndexRefineFlat: Index {
+
+    /// storage for full vectors
+    IndexFlat refine_index;
+
+    /// faster index to pre-select the vectors that should be filtered
+    Index *base_index;
+    bool own_fields;  ///< should the base index be deallocated?
+
+    /// factor between k requested in search and the k requested from
+    /// the base_index (should be >= 1)
+    float k_factor;
+
+    explicit IndexRefineFlat (Index *base_index);
+
+    IndexRefineFlat ();
+
+    virtual void train (idx_t n, const float *x) override;
+
+    virtual void add (idx_t n, const float *x) override;
+
+    virtual void reset() override;
+
+    virtual void search (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            idx_t *labels) const override;
+
+    virtual void set_typename () override;
+
+    virtual ~IndexRefineFlat ();
+};
+
+
+/// optimized version for 1D "vectors"
+struct IndexFlat1D:IndexFlatL2 {
+    bool continuous_update; ///< is the permutation updated continuously?
+
+    std::vector<idx_t> perm; ///< sorted database indices
+
+    explicit IndexFlat1D (bool continuous_update=true);
+
+    /// if not continuous_update, call this between the last add and
+    /// the first search
+    void update_permutation ();
+
+    virtual void add (idx_t n, const float *x) override;
+
+    virtual void reset() override;
+
+    /// Warn: the distances returned are L1 not L2
+    virtual void search (
+            idx_t n,
+            const float *x,
+            idx_t k,
+            float *distances,
+            idx_t *labels) const override;
+
+
+};
+
+
+}
+
+#endif

IndexIVF.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#ifndef FAISS_INDEX_IVF_H
+#define FAISS_INDEX_IVF_H
+
+
+#include <vector>
+
+
+#include "Index.h"
+#include "Clustering.h"
+#include "Heap.h"
+
+
+namespace faiss {
+
+
+
+/** Index based on a inverted file (IVF)
+ *
+ * In the inverted file, the quantizer (an Index instance) provides a
+ * quantization index for each vector to be added. The quantization
+ * index maps to a list (aka inverted list or posting list), where the
+ * id of the vector is then stored.
+ *
+ * At search time, the vector to be searched is also quantized, and
+ * only the list corresponding to the quantization index is
+ * searched. This speeds up the search by making it
+ * non-exhaustive. This can be relaxed using multi-probe search: a few
+ * (nprobe) quantization indices are selected and several inverted
+ * lists are visited.
+ *
+ * Sub-classes implement a post-filtering of the index that refines
+ * the distance estimation from the query to databse vectors.
+ */
+struct IndexIVF: Index {
+    size_t nlist;             ///< number of possible key values
+    size_t nprobe;            ///< number of probes at query time
+
+    Index * quantizer;        ///< quantizer that maps vectors to inverted lists
+    bool quantizer_trains_alone;   ///< just pass over the trainset to quantizer
+    bool own_fields;          ///< whether object owns the quantizer
+
+    ClusteringParameters cp; ///< to override default clustering params
+
+    std::vector < std::vector<long> > ids;  ///< Inverted lists for indexes
+
+    /// map for direct access to the elements. Enables reconstruct().
+    bool maintain_direct_map;
+    std::vector <long> direct_map;
+
+    /** The Inverted file takes a quantizer (an Index) on input,
+     * which implements the function mapping a vector to a list
+     * identifier. The pointer is borrowed: the quantizer should not
+     * be deleted while the IndexIVF is in use.
+     */
+    IndexIVF (Index * quantizer, size_t d, size_t nlist,
+              MetricType metric = METRIC_INNER_PRODUCT);
+
+    virtual void reset () override;
+
+    /// Trains the quantizer and calls train_residual to train sub-quantizers
+    virtual void train (idx_t n, const float *x) override;
+
+    /// Quantizes x and calls add_with_key
+    virtual void add (idx_t n, const float *x) override;
+
+
+
+    /// Sub-classes that encode the residuals can train their encoders here
+    /// does nothing by default
+    virtual void train_residual (idx_t n, const float *x);
+
+    /** moves the entries from another dataset to self. On output,
+     * other is empty. add_id is added to all moved ids (for
+     * sequential ids, this would be this->ntotal */
+    virtual void merge_from (IndexIVF &other, idx_t add_id);
+
+    /** implemented by sub-classes */
+    virtual void merge_from_residuals (IndexIVF &other) = 0;
+
+    virtual ~IndexIVF();
+
+    /// intialize a direct map
+    void make_direct_map ();
+
+    /// 1= perfectly balanced, >1: imbalanced
+    double imbalance_factor () const;
+
+    /// display some stats about the inverted lists
+    void print_stats () const;
+
+    IndexIVF ();
+};
+
+
+
+
+/** Inverted file with stored vectors. Here the inverted file
+ * pre-selects the vectors to be searched, but they are not otherwise
+ * encoded.
+ */
+struct IndexIVFFlat: IndexIVF {
+    /** Inverted list of original vectors. Each list is a nl * d
+     * matrix, where nl is the nb of vectors stored in the list. */
+    std::vector < std::vector<float> > vecs;
+
+    IndexIVFFlat (
+            Index * quantizer, size_t d, size_t nlist_,
+            MetricType = METRIC_INNER_PRODUCT);
+
+    /** Return the typeName of the index (which includes main parameters */
+    virtual void set_typename () override;
+
+
+    /// same as add_with_ids, with precomputed coarse quantizer
+    void add_core (idx_t n, const float * x, const long *xids,
+                   const long *precomputed_idx);
+
+
+    /// implemented for all IndexIVF* classes
+    virtual void add_with_ids (idx_t n, const float * x, const long *xids)
+        override;
+
+    virtual void search (
+            idx_t n, const float *x, idx_t k,
+            float *distances, idx_t *labels) const override;
+
+    virtual void range_search (
+            idx_t n, const float *x, float radius,
+            RangeSearchResult *result) const override;
+
+    /** copy a subset of the entries index to the other index
+     *
+     * if subset_type == 0: copies ids in [a1, a2)
+     * if subset_type == 1: copies ids if id % a1 == a2
+     */
+    void copy_subset_to (IndexIVFFlat & other, int subset_type,
+                         long a1, long a2) const;
+
+
+    virtual void reset() override;
+
+    virtual long remove_ids (const IDSelector & sel) override;
+
+    /// Implementation of the search for the inner product metric
+    void search_knn_inner_product (
+            size_t nx, const float * x,
+            const long * keys,
+            float_minheap_array_t * res) const;
+
+    /// Implementation of the search for the L2 metric
+    void search_knn_L2sqr (
+            size_t nx, const float * x,
+            const long * keys,
+            float_maxheap_array_t * res) const;
+
+    virtual void reconstruct (idx_t key, float * recons) const override;
+
+    virtual void merge_from_residuals (IndexIVF &other) override;
+
+
+    IndexIVFFlat () {}
+};
+
+
+
+
+
+
+} // namespace faiss
+
+
+
+
+
+#endif

IndexIVFPQ.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#ifndef FAISS_INDEX_IVFPQ_H
+#define FAISS_INDEX_IVFPQ_H
+
+
+#include <vector>
+
+#include "IndexIVF.h"
+#include "IndexPQ.h"
+
+
+namespace faiss {
+
+
+
+/** Inverted file with Product Quantizer encoding. Each residual
+ * vector is encoded as a product quantizer code.
+ */
+struct IndexIVFPQ: IndexIVF {
+    bool by_residual;              ///< Encode residual or plain vector?
+    int use_precomputed_table;     ///< if by_residual, build precompute tables
+    size_t code_size;              ///< code size per vector in bytes
+    ProductQuantizer pq;           ///< produces the codes
+
+    bool do_polysemous_training;   ///< reorder PQ centroids after training?
+    PolysemousTraining *polysemous_training; ///< if NULL, use default
+
+    // search-time parameters
+    size_t scan_table_threshold;   ///< use table computation or on-the-fly?
+    size_t max_codes;              ///< max nb of codes to visit to do a query
+    int polysemous_ht;             ///< Hamming thresh for polysemous filtering
+
+    std::vector < std::vector<uint8_t> > codes; // binary codes, size nlist
+
+    /// if use_precompute_table
+    /// size nlist * pq.M * pq.ksub
+    std::vector <float> precomputed_table;
+
+    IndexIVFPQ (
+            Index * quantizer, size_t d, size_t nlist,
+            size_t M, size_t nbits_per_idx);
+
+    virtual void set_typename () override;
+
+    virtual void add_with_ids (
+            idx_t n, const float *x,
+            const long *xids = nullptr) override;
+
+    /// same as add_core, also:
+    /// - output 2nd level residuals if residuals_2 != NULL
+    /// - use precomputed list numbers if precomputed_idx != NULL
+    void add_core_o (idx_t n, const float *x,
+                     const long *xids, float *residuals_2,
+                     const long *precomputed_idx = nullptr);
+
+    virtual void search (
+            idx_t n, const float *x, idx_t k,
+            float *distances, idx_t *labels) const override;
+
+    virtual void reset () override;
+
+    virtual long remove_ids (const IDSelector & sel) override;
+
+    /// trains the product quantizer
+    virtual void train_residual(idx_t n, const float *x) override;
+
+    /// same as train_residual, also output 2nd level residuals
+    void train_residual_o (idx_t n, const float *x, float *residuals_2);
+
+
+    /** Reconstruct a subset of the indexed vectors
+     *
+     * @param i0     first vector to reconstruct
+     * @param ni     nb of vectors to reconstruct
+     * @param recons output array of reconstructed vectors, size ni * d
+     */
+    virtual void reconstruct_n (idx_t i0, idx_t ni, float *recons)
+        const override;
+
+    virtual void reconstruct (idx_t key, float * recons)
+        const override;
+
+    /** Find exact duplicates in the dataset.
+     *
+     * the duplicates are returned in pre-allocated arrays (see the
+     * max sizes).
+     *
+     * @params lims   limits between groups of duplicates
+     *                (max size ntotal / 2 + 1)
+     * @params ids    ids[lims[i]] : ids[lims[i+1]-1] is a group of
+     *                duplicates (max size ntotal)
+     * @return n      number of groups found
+     */
+    size_t find_duplicates (idx_t *ids, size_t *lims) const;
+
+    // map a vector to a binary code knowning the index
+    void encode (long key, const float * x, uint8_t * code) const;
+
+    /// same as encode, for multiple points at once
+    void encode_multiple (size_t n, const long *keys,
+                          const float * x, uint8_t * codes) const;
+
+    /** search a set of vectors, that are pre-quantized by the IVF
+     *  quantizer. Fill in the corresponding heaps with the query
+     *  results.
+     *
+     * @param nx     nb of vectors to query
+     * @param qx     query vectors, size nx * d
+     * @param keys   coarse quantization indices, size nx * nprobe
+     * @param coarse_dis
+     *               distances to coarse centroids, size nx * nprobe
+     * @param res    heaps for all the results, gives the nprobe
+     * @param store_pairs store inv list index + inv list offset
+     *                     instead in upper/lower 32 bit of result,
+     *                     instead of ids (used for reranking).
+     */
+    virtual void search_knn_with_key (
+            size_t nx,
+            const float * qx,
+            const long * keys,
+            const float * coarse_dis,
+            float_maxheap_array_t* res,
+            bool store_pairs = false) const;
+
+    /// build precomputed table
+    void precompute_table ();
+
+    /// used to implement merging
+    virtual void merge_from_residuals (IndexIVF &other) override;
+
+
+    /** copy a subset of the entries index to the other index
+     *
+     * if subset_type == 0: copies ids in [a1, a2)
+     * if subset_type == 1: copies ids if id % a1 == a2
+     */
+    void copy_subset_to (IndexIVFPQ & other, int subset_type,
+                         long a1, long a2) const;
+
+    IndexIVFPQ ();
+
+};
+
+
+/// statistics are robust to internal threading, but not if
+/// IndexIVFPQ::search is called by multiple threads
+struct IndexIVFPQStats {
+    size_t nq;       // nb of queries run
+    size_t nlist;    // nb of inverted lists scanned
+    size_t ncode;    // nb of codes visited
+    size_t nrefine;  // nb of refines (IVFPQR)
+
+    size_t n_hamming_pass;
+    // nb of passed Hamming distance tests (for polysemous)
+
+    // timings measured with the CPU RTC
+    // on all threads
+    size_t assign_cycles;
+    size_t search_cycles;
+    size_t refine_cycles; // only for IVFPQR
+
+    // single thread (double-counted with search_cycles)
+    size_t init_query_cycles;
+    size_t init_list_cycles;
+    size_t scan_cycles;
+    size_t heap_cycles;
+
+    IndexIVFPQStats () {reset (); }
+    void reset ();
+};
+
+// global var that collects them all
+extern IndexIVFPQStats indexIVFPQ_stats;
+
+
+
+/** Index with an additional level of PQ refinement */
+struct IndexIVFPQR: IndexIVFPQ {
+    ProductQuantizer refine_pq;           ///< 3rd level quantizer
+    std::vector <uint8_t> refine_codes;   ///< corresponding codes
+
+    /// factor between k requested in search and the k requested from the IVFPQ
+    float k_factor;
+
+    IndexIVFPQR (
+            Index * quantizer, size_t d, size_t nlist,
+            size_t M, size_t nbits_per_idx,
+            size_t M_refine, size_t nbits_per_idx_refine);
+
+    virtual void set_typename () override;
+
+    virtual void reset() override;
+
+    virtual long remove_ids (const IDSelector & sel) override;
+
+    /// trains the two product quantizers
+    virtual void train_residual (idx_t n, const float *x) override;
+
+    virtual void add_with_ids (idx_t n, const float *x, const long *xids)
+        override;
+
+    /// same as add_with_ids, but optionally use the precomputed list ids
+    void add_core (idx_t n, const float *x, const long *xids,
+                     const long *precomputed_idx = nullptr);
+
+
+    virtual void reconstruct_n (idx_t i0, idx_t ni, float *recons)
+        const override;
+
+    virtual void search (
+            idx_t n, const float *x, idx_t k,
+            float *distances, idx_t *labels) const override;
+
+    virtual void merge_from_residuals (IndexIVF &other) override;
+
+    IndexIVFPQR();
+};
+
+
+/** Index with 32-bit ids and flat tables. Must be constructed from an
+ *  exisiting IndexIVFPQ. Cannot be copy-constructed/assigned. The
+ *  actual data is stored in the compact_* tables, the ids and codes
+ *  tables are not used.  */
+struct IndexIVFPQCompact: IndexIVFPQ {
+
+    explicit IndexIVFPQCompact (const IndexIVFPQ &other);
+
+    /// how were the compact tables allocated?
+    enum Alloc_type_t {
+        Alloc_type_none,     ///< alloc from outside
+        Alloc_type_new,      ///< was allocated with new
+        Alloc_type_mmap      ///< was mmapped
+    };
+
+    Alloc_type_t alloc_type;
+
+    uint32_t *limits;        ///< size nlist + 1
+    uint32_t *compact_ids;   ///< size ntotal
+    uint8_t *compact_codes;  ///< size ntotal * code_size
+
+    // file and buffer this was mmapped (will be unmapped when object
+    // is deleted)
+    char * mmap_buffer;
+    long mmap_length;
+
+    virtual void search_knn_with_key (
+            size_t nx,
+            const float * qx,
+            const long * keys,
+            const float * coarse_dis,
+            float_maxheap_array_t * res,
+            bool store_pairs = false) const override;
+
+    /// the three following functions will fail at runtime
+    virtual void add (idx_t, const float *) override;
+    virtual void reset () override;
+    virtual void train (idx_t, const float *) override;
+
+    virtual ~IndexIVFPQCompact ();
+
+    IndexIVFPQCompact ();
+
+};
+
+
+
+} // namespace faiss
+
+
+
+
+
+#endif

IndexLSH.cpp

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+
+#include "IndexLSH.h"
+
+#include <cstdio>
+#include <cstring>
+
+#include <algorithm>
+
+#include "utils.h"
+#include "hamming.h"
+#include "FaissAssert.h"
+
+namespace faiss {
+
+/***************************************************************
+ * IndexLSH
+ ***************************************************************/
+
+
+IndexLSH::IndexLSH (idx_t d, int nbits, bool rotate_data, bool train_thresholds):
+    Index(d), nbits(nbits), rotate_data(rotate_data),
+    train_thresholds (train_thresholds), rrot(d, nbits)
+{
+    is_trained = !train_thresholds;
+
+    bytes_per_vec = (nbits + 7) / 8;
+
+    if (rotate_data) {
+        rrot.init(5);
+    } else {
+        FAISS_ASSERT(d >= nbits);
+    }
+    set_typename();
+}
+
+IndexLSH::IndexLSH ():
+    nbits (0), bytes_per_vec(0), rotate_data (false), train_thresholds (false)
+{
+}
+
+
+void IndexLSH::set_typename()
+{
+    std::stringstream s;
+    s << "LSH_" << nbits << (rotate_data ? "r" : "");
+    index_typename = s.str();
+}
+
+const float * IndexLSH::apply_preprocess (idx_t n, const float *x) const
+{
+
+    float *xt = nullptr;
+    if (rotate_data) {
+        // also applies bias if exists
+        xt = rrot.apply (n, x);
+    } else if (d != nbits) {
+        xt = new float [nbits * n];
+        float *xp = xt;
+        for (idx_t i = 0; i < n; i++) {
+            const float *xl = x + i * d;
+            for (int j = 0; j < nbits; j++)
+                *xp++ = xl [j];
+        }
+    }
+
+    if (train_thresholds) {
+
+        if (xt == NULL) {
+            xt = new float [nbits * n];
+            memcpy (xt, x, sizeof(*x) * n * nbits);
+        }
+
+        float *xp = xt;
+        for (idx_t i = 0; i < n; i++)
+            for (int j = 0; j < nbits; j++)
+                *xp++ -= thresholds [j];
+    }
+
+    return xt ? xt : x;
+}
+
+
+
+void IndexLSH::train (idx_t n, const float *x)
+{
+    if (train_thresholds) {
+        thresholds.resize (nbits);
+        train_thresholds = false;
+        const float *xt = apply_preprocess (n, x);
+        train_thresholds = true;
+
+        float * transposed_x = new float [n * nbits];
+
+        for (idx_t i = 0; i < n; i++)
+            for (idx_t j = 0; j < nbits; j++)
+                transposed_x [j * n + i] = xt [i * nbits + j];
+        if (xt != x) delete [] xt;
+
+        for (idx_t i = 0; i < nbits; i++) {
+            float *xi = transposed_x + i * n;
+            // std::nth_element
+            std::sort (xi, xi + n);
+            if (n % 2 == 1)
+                thresholds [i] = xi [n / 2];
+            else
+                thresholds [i] = (xi [n / 2 - 1] + xi [n / 2]) / 2;
+
+        }
+    }
+    is_trained = true;
+}
+
+
+void IndexLSH::add (idx_t n, const float *x)
+{
+    FAISS_ASSERT (is_trained);
+    const float *xt = apply_preprocess (n, x);
+
+    codes.resize ((ntotal + n) * bytes_per_vec);
+    fvecs2bitvecs (xt, &codes[ntotal * bytes_per_vec], nbits, n);
+    if (x != xt)
+        delete [] xt;
+    ntotal += n;
+}
+
+
+void IndexLSH::search (
+        idx_t n,
+        const float *x,
+        idx_t k,
+        float *distances,
+        idx_t *labels) const
+{
+    FAISS_ASSERT (is_trained);
+    const float *xt = apply_preprocess (n, x);
+
+    uint8_t * qcodes = new uint8_t [n * bytes_per_vec];
+    fvecs2bitvecs (xt, qcodes, nbits, n);
+
+    if (x != xt)
+        delete [] xt;
+
+    int * idistances = new int [n * k];
+    int_maxheap_array_t res = { size_t(n), size_t(k), labels, idistances};
+
+    hammings_knn (&res, qcodes, codes.data(),
+                  ntotal, bytes_per_vec, true);
+
+    delete [] qcodes;
+
+    // convert distances to floats
+    for (int i = 0; i < k * n; i++)
+        distances[i] = idistances[i];
+    delete [] idistances;
+
+}
+
+
+void IndexLSH::transfer_thresholds (LinearTransform *vt) {
+    if (!train_thresholds) return;
+    FAISS_ASSERT (nbits == vt->d_out);
+    if (!vt->have_bias) {
+        vt->b.resize (nbits, 0);
+        vt->have_bias = true;
+    }
+    for (int i = 0; i < nbits; i++)
+        vt->b[i] -= thresholds[i];
+    train_thresholds = false;
+    thresholds.clear();
+}
+
+void IndexLSH::reset() {
+    codes.clear();
+    ntotal = 0;
+}
+
+
+
+} // namespace faiss

IndexLSH.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#ifndef INDEX_LSH_H
+#define INDEX_LSH_H
+
+#include <vector>
+
+#include "Index.h"
+#include "VectorTransform.h"
+
+namespace faiss {
+
+
+/** The sign of each vector component is put in a binary signature */
+struct IndexLSH:Index {
+    typedef unsigned char uint8_t;
+
+    int nbits;              ///< nb of bits per vector
+    int bytes_per_vec;      ///< nb of 8-bits per encoded vector
+    bool rotate_data;       ///< whether to apply a random rotation to input
+    bool train_thresholds;  ///< whether we train thresholds or use 0
+
+    RandomRotationMatrix rrot; ///< optional random rotation
+
+    std::vector <float> thresholds; ///< thresholds to compare with
+
+    /// encoded dataset
+    std::vector<uint8_t> codes;
+
+    IndexLSH (
+            idx_t d, int nbits,
+            bool rotate_data = true,
+            bool train_thresholds = false);
+
+    /** Preprocesses and resizes the input to the size required to
+     * binarize the data
+     *
+     * @param x input vectors, size n * d
+     * @return output vectors, size n * bits. May be the same pointer
+     *         as x, otherwise it should be deleted by the caller
+     */
+    const float *apply_preprocess (idx_t n, const float *x) const;
+
+    virtual void set_typename () override;
+
+    virtual void train (idx_t n, const float *x) override;
+
+    virtual void add (idx_t n, const float *x) override;
+
+    virtual void search (
+            idx_t n,
+            const float *x, idx_t k,
+            float *distances,
+            idx_t *labels) const override;
+
+    virtual void reset() override;
+
+    /// transfer the thresholds to a pre-processing stage (and unset
+    /// train_thresholds)
+    void transfer_thresholds (LinearTransform * vt);
+
+    virtual ~IndexLSH () {}
+
+    IndexLSH ();
+};
+
+
+
+}
+
+
+
+
+
+
+#endif

IndexPQ.h

+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved.
+// -*- c++ -*-
+
+#ifndef FAISS_INDEX_PQ_H
+#define FAISS_INDEX_PQ_H
+
+#include <stdint.h>
+
+#include <vector>
+
+#include "Index.h"
+#include "ProductQuantizer.h"
+#include "PolysemousTraining.h"
+
+namespace faiss {
+
+
+/** Index based on a product quantizer. Stored vectors are
+ * approximated by PQ codes. */
+struct IndexPQ: Index {
+
+    /// The product quantizer used to encode the vectors
+    ProductQuantizer pq;
+
+    /// Codes. Size ntotal * pq.code_size
+    std::vector<uint8_t> codes;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    IndexPQ (int d,                    ///< dimensionality of the input vectors
+             size_t M,                 ///< number of subquantizers
+             size_t nbits,             ///< number of bit per subvector index
+             MetricType metric = METRIC_L2);
+
+    IndexPQ ();
+
+    virtual void set_typename () override;
+
+    virtual void train (idx_t n, const float *x) override;
+
+    virtual void add (idx_t n, const float *x) override;
+
+    virtual void search (
+            idx_t n, const float *x, idx_t k,
+            float *distances, idx_t *labels) const override;
+
+    virtual void reset() override;
+
+    virtual void reconstruct_n (idx_t i0, idx_t ni, float *recons)
+        const override;
+
+    virtual void reconstruct (idx_t key, float * recons)
+        const override;
+
+    /******************************************************
+     * Polysemous codes implementation
+     ******************************************************/
+    bool do_polysemous_training; ///< false = standard PQ
+
+    /// parameters used for the polysemous training
+    PolysemousTraining polysemous_training;
+
+    /// how to perform the search in search_core
+    enum Search_type_t {
+        ST_PQ,             ///< asymmetric product quantizer (default)
+        ST_HE,             ///< Hamming distance on codes
+        ST_generalized_HE, ///< nb of same codes
+        ST_SDC,            ///< symmetric product quantizer (SDC)
+        ST_polysemous,     ///< HE filter (using ht) + PQ combination
+        ST_polysemous_generalize,  ///< Filter on generalized Hamming
+    };
+
+    Search_type_t search_type;
+
+    // just encode the sign of the components, instead of using the PQ encoder
+    // used only for the queries
+    bool encode_signs;
+
+    /// Hamming threshold used for polysemy
+    int polysemous_ht;
+
+    // actual polysemous search
+    void search_core_polysemous (idx_t n, const float *x, idx_t k,
+                               float *distances, idx_t *labels) const;
+
+    /// prepare query for a polysemous search, but instead of
+    /// computing the result, just get the histogram of Hamming
+    /// distances. May be computed on a provided dataset if xb != NULL
+    /// @param dist_histogram (M * nbits + 1)
+    void hamming_distance_histogram (idx_t n, const float *x,
+                                     idx_t nb, const float *xb,
+                                     long *dist_histogram);
+
+    /** compute pairwise distances between queries and database
+     *
+     * @param n    nb of query vectors
+     * @param x    query vector, size n * d
+     * @param dis  output distances, size n * ntotal
+     */
+    void hamming_distance_table (idx_t n, const float *x,
+                                 int32_t *dis) const;
+
+};
+
+
+/// statistics are robust to internal threading, but not if
+/// IndexPQ::search is called by multiple threads
+struct IndexPQStats {
+    size_t nq;       // nb of queries run
+    size_t ncode;    // nb of codes visited
+
+    size_t n_hamming_pass; // nb of passed Hamming distance tests (for polysemy)
+
+    IndexPQStats () {reset (); }
+    void reset ();
+};
+
+extern IndexPQStats indexPQ_stats;
+
+
+
+/** Quantizer where centroids are virtual: they are the Cartesian
+ *  product of sub-centroids. */
+struct MultiIndexQuantizer: Index  {
+    ProductQuantizer pq;
+
+    MultiIndexQuantizer (int d,         ///< dimension of the input vectors
+                         size_t M,      ///< number of subquantizers
+                         size_t nbits); ///< number of bit per subvector index
+
+    virtual void set_typename ();
+
+    virtual void train (idx_t n, const float *x);
+
+
+    virtual void search (idx_t n, const float *x, idx_t k,
+                              float *distances, idx_t *labels) const;
+
+    /// add and reset will crash at runtime
+    virtual void add (idx_t n, const float *x);
+    virtual void reset ();
+
+    MultiIndexQuantizer () {}
+
+    virtual void reconstruct (idx_t key, float * recons) const;
+};
+
+
+} // namespace faiss
+
+
+
+#endif

README.md

+# Faiss 
+
+Faiss is a library for efficient similarity search and clustering of dense vectors. It contains algorithms that search in sets of vectors of any size, up to ones that possibly do not fit in RAM. It also contains supporting code for evaluation and parameter tuning. Faiss is written in C++ with complete wrappers for Python/numpy. Some of the most useful algorithms are implemented on the GPU. It is developed by [Facebook AI Research](https://research.fb.com/category/facebook-ai-research-fair/).
+
+## Introduction
+
+Faiss contains several methods for similarity search. It assumes that the instances are represented as vectors and are identified by an integer, and that the vectors can be compared with L2 distances or dot products. Vectors that are similar to a query vector are those that have the lowest L2 distance or the highest dot product with the query vector. It also supports cosine similarity, since this is a dot product on normalized vectors.
+
+Most of the methods, like those based on binary vectors and compact quantization codes, solely use a compressed representation of the vectors and do not require to keep the original vectors. This generally comes at the cost of a less precise search but these methods can scale to billions of vectors in main memory on a single server. 
+
+The GPU implementation can accept input from either CPU or GPU memory. On a server with GPUs, the GPU indexes can be used a drop-in replacement for the CPU indexes (e.g., replace `IndexFlatL2` with `GpuIndexFlatL2`) and copies to/from GPU memory are handled automatically.
+
+## Building 
+
+The library is mostly implemented in C++, with optional GPU support provided via CUDA, and an optional Python interface. It compiles with a Makefile. See [INSTALL](INSTALL) for details.
+
+## How Faiss works
+
+Faiss is built around an index type that stores a set of vectors, and provides a function to search in them with L2 and/or dot product vector comparison. Some index types are simple baselines, such as exact search. Most of the available indexing structures correspond to various trade-offs with respect to
+
+- search time
+- search quality
+- memory used per index vector 
+- training time
+- need for external data for unsupervised training
+
+## Full documentation of Faiss
+
+The full documentation, including a tutorial can be found in the [wiki page](http://github.com/facebookresearch/faiss/wiki).
+See the [doxygen documentation](http://rawgithub.com/facebookresearch/faiss/blob/master/docs/html/index.html) for per-class information.
+
+## Join the Faiss community
+
+We monitor the [issues page](http://github.com/facebookresearch/faiss/issues) of the repository. You can report bugs, ask questions, etc.
+
+## License
+
+Faiss is licenced under CC-by-NC, see the LICENCE file for details. This licence may be relaxed to BSD in the future.
+