faiss 0.2.0 → 0.2.1

data/vendor/faiss/faiss/impl/NSG.h

@@ -0,0 +1,199 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#pragma once
+
+#include <memory>
+#include <mutex>
+#include <vector>
+
+#include <omp.h>
+
+#include <faiss/Index.h>
+#include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/random.h>
+
+namespace faiss {
+
+/** Implementation of the Navigating Spreading-out Graph (NSG)
+ * datastructure.
+ *
+ * Fast Approximate Nearest Neighbor Search With The
+ * Navigating Spreading-out Graph
+ *
+ *  Cong Fu, Chao Xiang, Changxu Wang, Deng Cai, VLDB 2019
+ *
+ * This implementation is heavily influenced by the NSG
+ * implementation by ZJULearning Group
+ * (https://github.com/zjulearning/nsg)
+ *
+ * The NSG object stores only the neighbor link structure, see
+ * IndexNSG.h for the full index object.
+ */
+
+struct DistanceComputer; // from AuxIndexStructures
+struct Neighbor;
+struct Node;
+
+namespace nsg {
+
+/***********************************************************
+ * Graph structure to store a graph.
+ *
+ * It is represented by an adjacency matrix `data`, where
+ * data[i, j] is the j-th neighbor of node i.
+ ***********************************************************/
+
+template <class node_t>
+struct Graph {
+    node_t* data;    ///< the flattened adjacency matrix
+    int K;           ///< nb of neighbors per node
+    int N;           ///< total nb of nodes
+    bool own_fields; ///< the underlying data owned by itself or not
+
+    // construct from a known graph
+    Graph(node_t* data, int N, int K)
+            : data(data), K(K), N(N), own_fields(false) {}
+
+    // construct an empty graph
+    // NOTE: the newly allocated data needs to be destroyed at destruction time
+    Graph(int N, int K) : K(K), N(N), own_fields(true) {
+        data = new node_t[N * K];
+    }
+
+    // copy constructor
+    Graph(const Graph& g) : Graph(g.N, g.K) {
+        memcpy(data, g.data, N * K * sizeof(node_t));
+    }
+
+    // release the allocated memory if needed
+    ~Graph() {
+        if (own_fields) {
+            delete[] data;
+        }
+    }
+
+    // access the j-th neighbor of node i
+    inline node_t at(int i, int j) const {
+        return data[i * K + j];
+    }
+
+    // access the j-th neighbor of node i by reference
+    inline node_t& at(int i, int j) {
+        return data[i * K + j];
+    }
+};
+
+DistanceComputer* storage_distance_computer(const Index* storage);
+
+} // namespace nsg
+
+struct NSG {
+    /// internal storage of vectors (32 bits: this is expensive)
+    using storage_idx_t = int;
+
+    /// Faiss results are 64-bit
+    using idx_t = Index::idx_t;
+
+    int ntotal; ///< nb of nodes
+
+    /// construction-time parameters
+    int R; ///< nb of neighbors per node
+    int L; ///< length of the search path at construction time
+    int C; ///< candidate pool size at construction time
+
+    // search-time parameters
+    int search_L; ///< length of the search path
+
+    int enterpoint; ///< enterpoint
+
+    std::shared_ptr<nsg::Graph<int>> final_graph; ///< NSG graph structure
+
+    bool is_built; ///< NSG is built or not
+
+    RandomGenerator rng; ///< random generator
+
+    explicit NSG(int R = 32);
+
+    // build NSG from a KNN graph
+    void build(
+            Index* storage,
+            idx_t n,
+            const nsg::Graph<idx_t>& knn_graph,
+            bool verbose);
+
+    // reset the graph
+    void reset();
+
+    // search interface
+    void search(
+            DistanceComputer& dis,
+            int k,
+            idx_t* I,
+            float* D,
+            VisitedTable& vt) const;
+
+    // Compute the center point
+    void init_graph(Index* storage, const nsg::Graph<idx_t>& knn_graph);
+
+    // Search on a built graph.
+    // If collect_fullset is true, the visited nodes will be
+    // collected in `fullset`.
+    template <bool collect_fullset, class index_t>
+    void search_on_graph(
+            const nsg::Graph<index_t>& graph,
+            DistanceComputer& dis,
+            VisitedTable& vt,
+            int ep,
+            int pool_size,
+            std::vector<Neighbor>& retset,
+            std::vector<Node>& fullset) const;
+
+    // Add reverse links
+    void add_reverse_links(
+            int q,
+            std::vector<std::mutex>& locks,
+            DistanceComputer& dis,
+            nsg::Graph<Node>& graph);
+
+    void sync_prune(
+            int q,
+            std::vector<Node>& pool,
+            DistanceComputer& dis,
+            VisitedTable& vt,
+            const nsg::Graph<idx_t>& knn_graph,
+            nsg::Graph<Node>& graph);
+
+    void link(
+            Index* storage,
+            const nsg::Graph<idx_t>& knn_graph,
+            nsg::Graph<Node>& graph,
+            bool verbose);
+
+    // make NSG be fully connected
+    int tree_grow(Index* storage, std::vector<int>& degrees);
+
+    // count the size of the connected component
+    // using depth first search start by root
+    int dfs(VisitedTable& vt, int root, int cnt) const;
+
+    // attach one unlinked node
+    int attach_unlinked(
+            Index* storage,
+            VisitedTable& vt,
+            VisitedTable& vt2,
+            std::vector<int>& degrees);
+
+    // check the integrity of the NSG built
+    void check_graph() const;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/PolysemousTraining.cpp

@@ -8,18 +8,21 @@
 // -*- c++ -*-
 
 #include <faiss/impl/PolysemousTraining.h>
+#include "faiss/impl/FaissAssert.h"
+
+#include <omp.h>
+#include <stdint.h>
 
-#include <cstdlib>
 #include <cmath>
+#include <cstdlib>
 #include <cstring>
-#include <stdint.h>
 
 #include <algorithm>
 
-#include <faiss/utils/random.h>
-#include <faiss/utils/utils.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/hamming.h>
+#include <faiss/utils/random.h>
+#include <faiss/utils/utils.h>
 
 #include <faiss/impl/FaissAssert.h>
 
@@ -29,16 +32,14 @@
 
 namespace faiss {
 
-
 /****************************************************
  * Optimization code
  ****************************************************/
 
-SimulatedAnnealingParameters::SimulatedAnnealingParameters ()
-{
+SimulatedAnnealingParameters::SimulatedAnnealingParameters() {
     // set some reasonable defaults for the optimization
     init_temperature = 0.7;
-    temperature_decay = pow (0.9, 1/500.);
+    temperature_decay = pow(0.9, 1 / 500.);
     // reduce by a factor 0.9 every 500 it
     n_iter = 500000;
     n_redo = 2;
@@ -50,44 +51,37 @@ SimulatedAnnealingParameters::SimulatedAnnealingParameters ()
 
 // what would the cost update be if iw and jw were swapped?
 // default implementation just computes both and computes the difference
-double PermutationObjective::cost_update (
-        const int *perm, int iw, int jw) const
-{
-    double orig_cost = compute_cost (perm);
+double PermutationObjective::cost_update(const int* perm, int iw, int jw)
+        const {
+    double orig_cost = compute_cost(perm);
 
-    std::vector<int> perm2 (n);
+    std::vector<int> perm2(n);
     for (int i = 0; i < n; i++)
         perm2[i] = perm[i];
     perm2[iw] = perm[jw];
     perm2[jw] = perm[iw];
 
-    double new_cost = compute_cost (perm2.data());
+    double new_cost = compute_cost(perm2.data());
     return new_cost - orig_cost;
 }
 
-
-
-
-SimulatedAnnealingOptimizer::SimulatedAnnealingOptimizer (
-        PermutationObjective *obj,
-        const SimulatedAnnealingParameters &p):
-    SimulatedAnnealingParameters (p),
-    obj (obj),
-    n(obj->n),
-    logfile (nullptr)
-{
-    rnd = new RandomGenerator (p.seed);
-    FAISS_THROW_IF_NOT (n < 100000 && n >=0 );
+SimulatedAnnealingOptimizer::SimulatedAnnealingOptimizer(
+        PermutationObjective* obj,
+        const SimulatedAnnealingParameters& p)
+        : SimulatedAnnealingParameters(p),
+          obj(obj),
+          n(obj->n),
+          logfile(nullptr) {
+    rnd = new RandomGenerator(p.seed);
+    FAISS_THROW_IF_NOT(n < 100000 && n >= 0);
 }
 
-SimulatedAnnealingOptimizer::~SimulatedAnnealingOptimizer ()
-{
+SimulatedAnnealingOptimizer::~SimulatedAnnealingOptimizer() {
     delete rnd;
 }
 
 // run the optimization and return the best result in best_perm
-double SimulatedAnnealingOptimizer::run_optimization (int * best_perm)
-{
+double SimulatedAnnealingOptimizer::run_optimization(int* best_perm) {
     double min_cost = 1e30;
 
     // just do a few runs of the annealing and keep the lowest output cost
@@ -95,84 +89,89 @@ double SimulatedAnnealingOptimizer::run_optimization (int * best_perm)
         std::vector<int> perm(n);
         for (int i = 0; i < n; i++)
             perm[i] = i;
-         if (init_random) {
+        if (init_random) {
             for (int i = 0; i < n; i++) {
-                int j = i + rnd->rand_int (n - i);
-                std::swap (perm[i], perm[j]);
+                int j = i + rnd->rand_int(n - i);
+                std::swap(perm[i], perm[j]);
             }
         }
-         float cost = optimize (perm.data());
-        if (logfile) fprintf (logfile, "\n");
-        if(verbose > 1) {
-            printf ("    optimization run %d: cost=%g %s\n",
-                    it, cost, cost < min_cost ? "keep" : "");
+        float cost = optimize(perm.data());
+        if (logfile)
+            fprintf(logfile, "\n");
+        if (verbose > 1) {
+            printf("    optimization run %d: cost=%g %s\n",
+                   it,
+                   cost,
+                   cost < min_cost ? "keep" : "");
         }
         if (cost < min_cost) {
-            memcpy (best_perm, perm.data(), sizeof(perm[0]) * n);
+            memcpy(best_perm, perm.data(), sizeof(perm[0]) * n);
             min_cost = cost;
         }
     }
-     return min_cost;
+    return min_cost;
 }
 
 // perform the optimization loop, starting from and modifying
 // permutation in-place
-double SimulatedAnnealingOptimizer::optimize (int *perm)
-{
-    double cost = init_cost = obj->compute_cost (perm);
+double SimulatedAnnealingOptimizer::optimize(int* perm) {
+    double cost = init_cost = obj->compute_cost(perm);
     int log2n = 0;
-    while (!(n <= (1 << log2n))) log2n++;
+    while (!(n <= (1 << log2n)))
+        log2n++;
     double temperature = init_temperature;
-     int n_swap = 0, n_hot = 0;
+    int n_swap = 0, n_hot = 0;
     for (int it = 0; it < n_iter; it++) {
         temperature = temperature * temperature_decay;
         int iw, jw;
         if (only_bit_flips) {
-            iw = rnd->rand_int (n);
-            jw = iw ^ (1 << rnd->rand_int (log2n));
+            iw = rnd->rand_int(n);
+            jw = iw ^ (1 << rnd->rand_int(log2n));
         } else {
-            iw = rnd->rand_int (n);
-            jw = rnd->rand_int (n - 1);
-            if (jw == iw) jw++;
+            iw = rnd->rand_int(n);
+            jw = rnd->rand_int(n - 1);
+            if (jw == iw)
+                jw++;
         }
-         double delta_cost = obj->cost_update (perm, iw, jw);
-         if (delta_cost < 0 || rnd->rand_float () < temperature) {
-            std::swap (perm[iw], perm[jw]);
+        double delta_cost = obj->cost_update(perm, iw, jw);
+        if (delta_cost < 0 || rnd->rand_float() < temperature) {
+            std::swap(perm[iw], perm[jw]);
             cost += delta_cost;
             n_swap++;
-            if (delta_cost >= 0) n_hot++;
+            if (delta_cost >= 0)
+                n_hot++;
         }
-         if (verbose > 2 || (verbose > 1 && it % 10000 == 0)) {
-            printf ("      iteration %d cost %g temp %g n_swap %d "
-                    "(%d hot)     \r",
-                    it, cost, temperature, n_swap, n_hot);
+        if (verbose > 2 || (verbose > 1 && it % 10000 == 0)) {
+            printf("      iteration %d cost %g temp %g n_swap %d "
+                   "(%d hot)     \r",
+                   it,
+                   cost,
+                   temperature,
+                   n_swap,
+                   n_hot);
             fflush(stdout);
         }
         if (logfile) {
-            fprintf (logfile, "%d %g %g %d %d\n",
-                    it, cost, temperature, n_swap, n_hot);
+            fprintf(logfile,
+                    "%d %g %g %d %d\n",
+                    it,
+                    cost,
+                    temperature,
+                    n_swap,
+                    n_hot);
         }
-     }
-    if (verbose > 1) printf("\n");
+    }
+    if (verbose > 1)
+        printf("\n");
     return cost;
 }
 
-
-
-
-
 /****************************************************
  * Cost functions: ReproduceDistanceTable
  ****************************************************/
 
-
-
-
-
-
-static inline int hamming_dis (uint64_t a, uint64_t b)
-{
-    return __builtin_popcountl (a ^ b);
+static inline int hamming_dis(uint64_t a, uint64_t b) {
+    return __builtin_popcountl(a ^ b);
 }
 
 namespace {
@@ -182,14 +181,14 @@ struct ReproduceWithHammingObjective : PermutationObjective {
     int nbits;
     double dis_weight_factor;
 
-    static double sqr (double x) { return x * x; }
-
+    static double sqr(double x) {
+        return x * x;
+    }
 
     // weihgting of distances: it is more important to reproduce small
     // distances well
-    double dis_weight (double x) const
-    {
-        return exp (-dis_weight_factor * x);
+    double dis_weight(double x) const {
+        return exp(-dis_weight_factor * x);
     }
 
     std::vector<double> target_dis; // wanted distances (size n^2)
@@ -197,101 +196,105 @@ struct ReproduceWithHammingObjective : PermutationObjective {
 
     // cost = quadratic difference between actual distance and Hamming distance
     double compute_cost(const int* perm) const override {
-      double cost = 0;
-      for (int i = 0; i < n; i++) {
-        for (int j = 0; j < n; j++) {
-          double wanted = target_dis[i * n + j];
-          double w = weights[i * n + j];
-          double actual = hamming_dis(perm[i], perm[j]);
-          cost += w * sqr(wanted - actual);
+        double cost = 0;
+        for (int i = 0; i < n; i++) {
+            for (int j = 0; j < n; j++) {
+                double wanted = target_dis[i * n + j];
+                double w = weights[i * n + j];
+                double actual = hamming_dis(perm[i], perm[j]);
+                cost += w * sqr(wanted - actual);
+            }
         }
-      }
-      return cost;
+        return cost;
     }
 
-
     // what would the cost update be if iw and jw were swapped?
     // computed in O(n) instead of O(n^2) for the full re-computation
     double cost_update(const int* perm, int iw, int jw) const override {
-      double delta_cost = 0;
+        double delta_cost = 0;
 
-      for (int i = 0; i < n; i++) {
-        if (i == iw) {
-          for (int j = 0; j < n; j++) {
-            double wanted = target_dis[i * n + j], w = weights[i * n + j];
-            double actual = hamming_dis(perm[i], perm[j]);
-            delta_cost -= w * sqr(wanted - actual);
-            double new_actual =
-                hamming_dis(perm[jw], perm[j == iw ? jw : j == jw ? iw : j]);
-            delta_cost += w * sqr(wanted - new_actual);
-          }
-        } else if (i == jw) {
-          for (int j = 0; j < n; j++) {
-            double wanted = target_dis[i * n + j], w = weights[i * n + j];
-            double actual = hamming_dis(perm[i], perm[j]);
-            delta_cost -= w * sqr(wanted - actual);
-            double new_actual =
-                hamming_dis(perm[iw], perm[j == iw ? jw : j == jw ? iw : j]);
-            delta_cost += w * sqr(wanted - new_actual);
-          }
-        } else {
-          int j = iw;
-          {
-            double wanted = target_dis[i * n + j], w = weights[i * n + j];
-            double actual = hamming_dis(perm[i], perm[j]);
-            delta_cost -= w * sqr(wanted - actual);
-            double new_actual = hamming_dis(perm[i], perm[jw]);
-            delta_cost += w * sqr(wanted - new_actual);
-          }
-          j = jw;
-          {
-            double wanted = target_dis[i * n + j], w = weights[i * n + j];
-            double actual = hamming_dis(perm[i], perm[j]);
-            delta_cost -= w * sqr(wanted - actual);
-            double new_actual = hamming_dis(perm[i], perm[iw]);
-            delta_cost += w * sqr(wanted - new_actual);
-          }
+        for (int i = 0; i < n; i++) {
+            if (i == iw) {
+                for (int j = 0; j < n; j++) {
+                    double wanted = target_dis[i * n + j],
+                           w = weights[i * n + j];
+                    double actual = hamming_dis(perm[i], perm[j]);
+                    delta_cost -= w * sqr(wanted - actual);
+                    double new_actual = hamming_dis(
+                            perm[jw],
+                            perm[j == iw           ? jw
+                                         : j == jw ? iw
+                                                   : j]);
+                    delta_cost += w * sqr(wanted - new_actual);
+                }
+            } else if (i == jw) {
+                for (int j = 0; j < n; j++) {
+                    double wanted = target_dis[i * n + j],
+                           w = weights[i * n + j];
+                    double actual = hamming_dis(perm[i], perm[j]);
+                    delta_cost -= w * sqr(wanted - actual);
+                    double new_actual = hamming_dis(
+                            perm[iw],
+                            perm[j == iw           ? jw
+                                         : j == jw ? iw
+                                                   : j]);
+                    delta_cost += w * sqr(wanted - new_actual);
+                }
+            } else {
+                int j = iw;
+                {
+                    double wanted = target_dis[i * n + j],
+                           w = weights[i * n + j];
+                    double actual = hamming_dis(perm[i], perm[j]);
+                    delta_cost -= w * sqr(wanted - actual);
+                    double new_actual = hamming_dis(perm[i], perm[jw]);
+                    delta_cost += w * sqr(wanted - new_actual);
+                }
+                j = jw;
+                {
+                    double wanted = target_dis[i * n + j],
+                           w = weights[i * n + j];
+                    double actual = hamming_dis(perm[i], perm[j]);
+                    delta_cost -= w * sqr(wanted - actual);
+                    double new_actual = hamming_dis(perm[i], perm[iw]);
+                    delta_cost += w * sqr(wanted - new_actual);
+                }
+            }
         }
-      }
 
-      return delta_cost;
+        return delta_cost;
     }
 
-
-
-    ReproduceWithHammingObjective (
-           int nbits,
-           const std::vector<double> & dis_table,
-           double dis_weight_factor):
-        nbits (nbits), dis_weight_factor (dis_weight_factor)
-    {
+    ReproduceWithHammingObjective(
+            int nbits,
+            const std::vector<double>& dis_table,
+            double dis_weight_factor)
+            : nbits(nbits), dis_weight_factor(dis_weight_factor) {
         n = 1 << nbits;
-        FAISS_THROW_IF_NOT (dis_table.size() == n * n);
-        set_affine_target_dis (dis_table);
+        FAISS_THROW_IF_NOT(dis_table.size() == n * n);
+        set_affine_target_dis(dis_table);
     }
 
-    void set_affine_target_dis (const std::vector<double> & dis_table)
-    {
+    void set_affine_target_dis(const std::vector<double>& dis_table) {
         double sum = 0, sum2 = 0;
         int n2 = n * n;
         for (int i = 0; i < n2; i++) {
-            sum += dis_table [i];
-            sum2 += dis_table [i] * dis_table [i];
+            sum += dis_table[i];
+            sum2 += dis_table[i] * dis_table[i];
         }
         double mean = sum / n2;
         double stddev = sqrt(sum2 / n2 - (sum / n2) * (sum / n2));
 
-        target_dis.resize (n2);
+        target_dis.resize(n2);
 
         for (int i = 0; i < n2; i++) {
             // the mapping function
-            double td = (dis_table [i] - mean) / stddev * sqrt(nbits / 4) +
-                nbits / 2;
+            double td = (dis_table[i] - mean) / stddev * sqrt(nbits / 4) +
+                    nbits / 2;
             target_dis[i] = td;
             // compute a weight
-            weights.push_back (dis_weight (td));
+            weights.push_back(dis_weight(td));
         }
-
     }
 
     ~ReproduceWithHammingObjective() override {}
@@ -301,27 +304,23 @@ struct ReproduceWithHammingObjective : PermutationObjective {
 
 // weihgting of distances: it is more important to reproduce small
 // distances well
-double ReproduceDistancesObjective::dis_weight (double x) const
-{
-    return exp (-dis_weight_factor * x);
+double ReproduceDistancesObjective::dis_weight(double x) const {
+    return exp(-dis_weight_factor * x);
 }
 
-
-double ReproduceDistancesObjective::get_source_dis (int i, int j) const
-{
-    return source_dis [i * n + j];
+double ReproduceDistancesObjective::get_source_dis(int i, int j) const {
+    return source_dis[i * n + j];
 }
 
 // cost = quadratic difference between actual distance and Hamming distance
-double ReproduceDistancesObjective::compute_cost (const int *perm) const
-{
+double ReproduceDistancesObjective::compute_cost(const int* perm) const {
     double cost = 0;
     for (int i = 0; i < n; i++) {
         for (int j = 0; j < n; j++) {
-            double wanted = target_dis [i * n + j];
-            double w = weights [i * n + j];
-            double actual = get_source_dis (perm[i], perm[j]);
-            cost += w * sqr (wanted - actual);
+            double wanted = target_dis[i * n + j];
+            double w = weights[i * n + j];
+            double actual = get_source_dis(perm[i], perm[j]);
+            cost += w * sqr(wanted - actual);
         }
     }
     return cost;
@@ -329,79 +328,75 @@ double ReproduceDistancesObjective::compute_cost (const int *perm) const
 
 // what would the cost update be if iw and jw were swapped?
 // computed in O(n) instead of O(n^2) for the full re-computation
-double ReproduceDistancesObjective::cost_update(
-        const int *perm, int iw, int jw) const
-{
+double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
+        const {
     double delta_cost = 0;
-     for (int i = 0; i < n; i++) {
+    for (int i = 0; i < n; i++) {
         if (i == iw) {
             for (int j = 0; j < n; j++) {
-                double wanted = target_dis [i * n + j],
-                    w = weights [i * n + j];
-                double actual = get_source_dis (perm[i], perm[j]);
-                delta_cost -= w * sqr (wanted - actual);
-                double new_actual = get_source_dis (
-                       perm[jw],
-                       perm[j == iw ? jw : j == jw ? iw : j]);
-                delta_cost += w * sqr (wanted - new_actual);
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = get_source_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = get_source_dis(
+                        perm[jw],
+                        perm[j == iw           ? jw
+                                     : j == jw ? iw
+                                               : j]);
+                delta_cost += w * sqr(wanted - new_actual);
             }
         } else if (i == jw) {
             for (int j = 0; j < n; j++) {
-                double wanted = target_dis [i * n + j],
-                    w = weights [i * n + j];
-                double actual = get_source_dis (perm[i], perm[j]);
-                delta_cost -= w * sqr (wanted - actual);
-                double new_actual = get_source_dis (
-                       perm[iw],
-                       perm[j == iw ? jw : j == jw ? iw : j]);
-                delta_cost += w * sqr (wanted - new_actual);
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = get_source_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = get_source_dis(
+                        perm[iw],
+                        perm[j == iw           ? jw
+                                     : j == jw ? iw
+                                               : j]);
+                delta_cost += w * sqr(wanted - new_actual);
             }
-        } else  {
+        } else {
             int j = iw;
             {
-                double wanted = target_dis [i * n + j],
-                    w = weights [i * n + j];
-                double actual = get_source_dis (perm[i], perm[j]);
-                delta_cost -= w * sqr (wanted - actual);
-                double new_actual = get_source_dis (perm[i], perm[jw]);
-                delta_cost += w * sqr (wanted - new_actual);
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = get_source_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = get_source_dis(perm[i], perm[jw]);
+                delta_cost += w * sqr(wanted - new_actual);
             }
             j = jw;
             {
-                double wanted = target_dis [i * n + j],
-                    w = weights [i * n + j];
-                double actual = get_source_dis (perm[i], perm[j]);
-                delta_cost -= w * sqr (wanted - actual);
-                double new_actual = get_source_dis (perm[i], perm[iw]);
-                delta_cost += w * sqr (wanted - new_actual);
+                double wanted = target_dis[i * n + j], w = weights[i * n + j];
+                double actual = get_source_dis(perm[i], perm[j]);
+                delta_cost -= w * sqr(wanted - actual);
+                double new_actual = get_source_dis(perm[i], perm[iw]);
+                delta_cost += w * sqr(wanted - new_actual);
             }
         }
     }
-     return delta_cost;
+    return delta_cost;
 }
 
-
-
-ReproduceDistancesObjective::ReproduceDistancesObjective (
-       int n,
-       const double *source_dis_in,
-       const double *target_dis_in,
-       double dis_weight_factor):
-    dis_weight_factor (dis_weight_factor),
-    target_dis (target_dis_in)
-{
+ReproduceDistancesObjective::ReproduceDistancesObjective(
+        int n,
+        const double* source_dis_in,
+        const double* target_dis_in,
+        double dis_weight_factor)
+        : dis_weight_factor(dis_weight_factor), target_dis(target_dis_in) {
     this->n = n;
-    set_affine_target_dis (source_dis_in);
+    set_affine_target_dis(source_dis_in);
 }
 
-void ReproduceDistancesObjective::compute_mean_stdev (
-          const double *tab, size_t n2,
-          double *mean_out, double *stddev_out)
-{
+void ReproduceDistancesObjective::compute_mean_stdev(
+        const double* tab,
+        size_t n2,
+        double* mean_out,
+        double* stddev_out) {
     double sum = 0, sum2 = 0;
     for (int i = 0; i < n2; i++) {
-        sum += tab [i];
-        sum2 += tab [i] * tab [i];
+        sum += tab[i];
+        sum2 += tab[i] * tab[i];
     }
     double mean = sum / n2;
     double stddev = sqrt(sum2 / n2 - (sum / n2) * (sum / n2));
@@ -409,32 +404,34 @@ void ReproduceDistancesObjective::compute_mean_stdev (
     *stddev_out = stddev;
 }
 
-void ReproduceDistancesObjective::set_affine_target_dis (
-          const double *source_dis_in)
-{
+void ReproduceDistancesObjective::set_affine_target_dis(
+        const double* source_dis_in) {
     int n2 = n * n;
 
     double mean_src, stddev_src;
-    compute_mean_stdev (source_dis_in, n2, &mean_src, &stddev_src);
+    compute_mean_stdev(source_dis_in, n2, &mean_src, &stddev_src);
 
     double mean_target, stddev_target;
-    compute_mean_stdev (target_dis, n2, &mean_target, &stddev_target);
+    compute_mean_stdev(target_dis, n2, &mean_target, &stddev_target);
 
-    printf ("map mean %g std %g -> mean %g std %g\n",
-            mean_src, stddev_src, mean_target, stddev_target);
+    printf("map mean %g std %g -> mean %g std %g\n",
+           mean_src,
+           stddev_src,
+           mean_target,
+           stddev_target);
 
-    source_dis.resize (n2);
-    weights.resize (n2);
+    source_dis.resize(n2);
+    weights.resize(n2);
 
     for (int i = 0; i < n2; i++) {
         // the mapping function
-        source_dis[i] = (source_dis_in[i] - mean_src) / stddev_src
-            * stddev_target + mean_target;
+        source_dis[i] =
+                (source_dis_in[i] - mean_src) / stddev_src * stddev_target +
+                mean_target;
 
         // compute a weight
-        weights [i] = dis_weight (target_dis[i]);
+        weights[i] = dis_weight(target_dis[i]);
     }
-
 }
 
 /****************************************************
@@ -444,8 +441,7 @@ void ReproduceDistancesObjective::set_affine_target_dis (
 /// Maintains a 3D table of elementary costs.
 /// Accumulates elements based on Hamming distance comparisons
 template <typename Ttab, typename Taccu>
-struct Score3Computer: PermutationObjective {
-
+struct Score3Computer : PermutationObjective {
     int nc;
 
     // cost matrix of size nc * nc *nc
@@ -453,21 +449,18 @@ struct Score3Computer: PermutationObjective {
     // where x has PQ code i, y- PQ code j and y+ PQ code k
     std::vector<Ttab> n_gt;
 
-
     /// the cost is a triple loop on the nc * nc * nc matrix of entries.
     ///
-    Taccu compute (const int * perm) const
-    {
+    Taccu compute(const int* perm) const {
         Taccu accu = 0;
-        const Ttab *p = n_gt.data();
+        const Ttab* p = n_gt.data();
         for (int i = 0; i < nc; i++) {
-            int ip = perm [i];
+            int ip = perm[i];
             for (int j = 0; j < nc; j++) {
-                int jp = perm [j];
+                int jp = perm[j];
                 for (int k = 0; k < nc; k++) {
-                    int kp = perm [k];
-                    if (hamming_dis (ip, jp) <
-                        hamming_dis (ip, kp)) {
+                    int kp = perm[k];
+                    if (hamming_dis(ip, jp) < hamming_dis(ip, kp)) {
                         accu += *p; // n_gt [ ( i * nc + j) * nc + k];
                     }
                     p++;
@@ -477,7 +470,6 @@ struct Score3Computer: PermutationObjective {
         return accu;
     }
 
-
     /** cost update if entries iw and jw of the permutation would be
      * swapped.
      *
@@ -487,25 +479,23 @@ struct Score3Computer: PermutationObjective {
      * cells. Practical speedup is about 8x, and the code is quite
      * complex :-/
      */
-    Taccu compute_update (const int *perm, int iw, int jw) const
-    {
-        assert (iw != jw);
-        if (iw > jw) std::swap (iw, jw);
+    Taccu compute_update(const int* perm, int iw, int jw) const {
+        assert(iw != jw);
+        if (iw > jw)
+            std::swap(iw, jw);
 
         Taccu accu = 0;
-        const Ttab * n_gt_i = n_gt.data();
+        const Ttab* n_gt_i = n_gt.data();
         for (int i = 0; i < nc; i++) {
-            int ip0 = perm [i];
-            int ip = perm [i == iw ? jw : i == jw ? iw : i];
+            int ip0 = perm[i];
+            int ip = perm[i == iw ? jw : i == jw ? iw : i];
 
-            //accu += update_i (perm, iw, jw, ip0, ip, n_gt_i);
+            // accu += update_i (perm, iw, jw, ip0, ip, n_gt_i);
 
-            accu += update_i_cross (perm, iw, jw,
-                                    ip0, ip, n_gt_i);
+            accu += update_i_cross(perm, iw, jw, ip0, ip, n_gt_i);
 
             if (ip != ip0)
-                accu += update_i_plane (perm, iw, jw,
-                                       ip0, ip, n_gt_i);
+                accu += update_i_plane(perm, iw, jw, ip0, ip, n_gt_i);
 
             n_gt_i += nc * nc;
         }
@@ -513,23 +503,26 @@ struct Score3Computer: PermutationObjective {
         return accu;
     }
 
-
-    Taccu update_i (const int *perm, int iw, int jw,
-                   int ip0, int ip, const Ttab * n_gt_i) const
-    {
+    Taccu update_i(
+            const int* perm,
+            int iw,
+            int jw,
+            int ip0,
+            int ip,
+            const Ttab* n_gt_i) const {
         Taccu accu = 0;
-        const Ttab *n_gt_ij = n_gt_i;
+        const Ttab* n_gt_ij = n_gt_i;
         for (int j = 0; j < nc; j++) {
             int jp0 = perm[j];
-            int jp = perm [j == iw ? jw : j == jw ? iw : j];
+            int jp = perm[j == iw ? jw : j == jw ? iw : j];
             for (int k = 0; k < nc; k++) {
-                int kp0 = perm [k];
-                int kp = perm [k == iw ? jw : k == jw ? iw : k];
-                int ng = n_gt_ij [k];
-                if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
+                int kp0 = perm[k];
+                int kp = perm[k == iw ? jw : k == jw ? iw : k];
+                int ng = n_gt_ij[k];
+                if (hamming_dis(ip, jp) < hamming_dis(ip, kp)) {
                     accu += ng;
                 }
-                if (hamming_dis (ip0, jp0) < hamming_dis (ip0, kp0)) {
+                if (hamming_dis(ip0, jp0) < hamming_dis(ip0, kp0)) {
                     accu -= ng;
                 }
             }
@@ -539,23 +532,27 @@ struct Score3Computer: PermutationObjective {
     }
 
     // 2 inner loops for the case ip0 != ip
-    Taccu update_i_plane (const int *perm, int iw, int jw,
-                         int ip0, int ip, const Ttab * n_gt_i) const
-    {
+    Taccu update_i_plane(
+            const int* perm,
+            int iw,
+            int jw,
+            int ip0,
+            int ip,
+            const Ttab* n_gt_i) const {
         Taccu accu = 0;
-        const Ttab *n_gt_ij = n_gt_i;
+        const Ttab* n_gt_ij = n_gt_i;
 
         for (int j = 0; j < nc; j++) {
             if (j != iw && j != jw) {
                 int jp = perm[j];
                 for (int k = 0; k < nc; k++) {
                     if (k != iw && k != jw) {
-                        int kp = perm [k];
-                        Ttab ng = n_gt_ij [k];
-                        if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
+                        int kp = perm[k];
+                        Ttab ng = n_gt_ij[k];
+                        if (hamming_dis(ip, jp) < hamming_dis(ip, kp)) {
                             accu += ng;
                         }
-                        if (hamming_dis (ip0, jp) < hamming_dis (ip0, kp)) {
+                        if (hamming_dis(ip0, jp) < hamming_dis(ip0, kp)) {
                             accu -= ng;
                         }
                     }
@@ -567,114 +564,128 @@ struct Score3Computer: PermutationObjective {
     }
 
     /// used for the 8 cells were the 3 indices are swapped
-    inline Taccu update_k (const int *perm, int iw, int jw,
-                          int ip0, int ip, int jp0, int jp,
-                          int k,
-                          const Ttab * n_gt_ij) const
-    {
+    inline Taccu update_k(
+            const int* perm,
+            int iw,
+            int jw,
+            int ip0,
+            int ip,
+            int jp0,
+            int jp,
+            int k,
+            const Ttab* n_gt_ij) const {
         Taccu accu = 0;
-        int kp0 = perm [k];
-        int kp = perm [k == iw ? jw : k == jw ? iw : k];
-        Ttab ng = n_gt_ij [k];
-        if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
+        int kp0 = perm[k];
+        int kp = perm[k == iw ? jw : k == jw ? iw : k];
+        Ttab ng = n_gt_ij[k];
+        if (hamming_dis(ip, jp) < hamming_dis(ip, kp)) {
             accu += ng;
         }
-        if (hamming_dis (ip0, jp0) < hamming_dis (ip0, kp0)) {
+        if (hamming_dis(ip0, jp0) < hamming_dis(ip0, kp0)) {
             accu -= ng;
         }
         return accu;
     }
 
     /// compute update on a line of k's, where i and j are swapped
-    Taccu update_j_line (const int *perm, int iw, int jw,
-                        int ip0, int ip, int jp0, int jp,
-                        const Ttab * n_gt_ij) const
-    {
+    Taccu update_j_line(
+            const int* perm,
+            int iw,
+            int jw,
+            int ip0,
+            int ip,
+            int jp0,
+            int jp,
+            const Ttab* n_gt_ij) const {
         Taccu accu = 0;
         for (int k = 0; k < nc; k++) {
-            if (k == iw || k == jw) continue;
-            int kp = perm [k];
-            Ttab ng = n_gt_ij [k];
-            if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
+            if (k == iw || k == jw)
+                continue;
+            int kp = perm[k];
+            Ttab ng = n_gt_ij[k];
+            if (hamming_dis(ip, jp) < hamming_dis(ip, kp)) {
                 accu += ng;
             }
-            if (hamming_dis (ip0, jp0) < hamming_dis (ip0, kp)) {
+            if (hamming_dis(ip0, jp0) < hamming_dis(ip0, kp)) {
                 accu -= ng;
             }
         }
         return accu;
     }
 
-
     /// considers the 2 pairs of crossing lines j=iw or jw and k = iw or kw
-    Taccu update_i_cross (const int *perm, int iw, int jw,
-                        int ip0, int ip, const Ttab * n_gt_i) const
-    {
+    Taccu update_i_cross(
+            const int* perm,
+            int iw,
+            int jw,
+            int ip0,
+            int ip,
+            const Ttab* n_gt_i) const {
         Taccu accu = 0;
-        const Ttab *n_gt_ij = n_gt_i;
+        const Ttab* n_gt_ij = n_gt_i;
 
         for (int j = 0; j < nc; j++) {
             int jp0 = perm[j];
-            int jp = perm [j == iw ? jw : j == jw ? iw : j];
+            int jp = perm[j == iw ? jw : j == jw ? iw : j];
 
-            accu += update_k (perm, iw, jw, ip0, ip, jp0, jp, iw, n_gt_ij);
-            accu += update_k (perm, iw, jw, ip0, ip, jp0, jp, jw, n_gt_ij);
+            accu += update_k(perm, iw, jw, ip0, ip, jp0, jp, iw, n_gt_ij);
+            accu += update_k(perm, iw, jw, ip0, ip, jp0, jp, jw, n_gt_ij);
 
             if (jp != jp0)
-                accu += update_j_line (perm, iw, jw, ip0, ip, jp0, jp, n_gt_ij);
+                accu += update_j_line(perm, iw, jw, ip0, ip, jp0, jp, n_gt_ij);
 
             n_gt_ij += nc;
         }
         return accu;
     }
 
-
     /// PermutationObjective implementeation (just negates the scores
     /// for minimization)
 
     double compute_cost(const int* perm) const override {
-      return -compute(perm);
+        return -compute(perm);
     }
 
     double cost_update(const int* perm, int iw, int jw) const override {
-      double ret = -compute_update(perm, iw, jw);
-      return ret;
+        double ret = -compute_update(perm, iw, jw);
+        return ret;
     }
 
     ~Score3Computer() override {}
 };
 
-
-
-
-
 struct IndirectSort {
-    const float *tab;
-    bool operator () (int a, int b) {return tab[a] < tab[b]; }
+    const float* tab;
+    bool operator()(int a, int b) {
+        return tab[a] < tab[b];
+    }
 };
 
-
-
-struct RankingScore2: Score3Computer<float, double> {
+struct RankingScore2 : Score3Computer<float, double> {
     int nbits;
     int nq, nb;
     const uint32_t *qcodes, *bcodes;
-    const float *gt_distances;
-
-    RankingScore2 (int nbits, int nq, int nb,
-                  const uint32_t *qcodes, const uint32_t *bcodes,
-                  const float *gt_distances):
-        nbits(nbits), nq(nq), nb(nb), qcodes(qcodes),
-        bcodes(bcodes), gt_distances(gt_distances)
-    {
+    const float* gt_distances;
+
+    RankingScore2(
+            int nbits,
+            int nq,
+            int nb,
+            const uint32_t* qcodes,
+            const uint32_t* bcodes,
+            const float* gt_distances)
+            : nbits(nbits),
+              nq(nq),
+              nb(nb),
+              qcodes(qcodes),
+              bcodes(bcodes),
+              gt_distances(gt_distances) {
         n = nc = 1 << nbits;
-        n_gt.resize (nc * nc * nc);
-        init_n_gt ();
+        n_gt.resize(nc * nc * nc);
+        init_n_gt();
     }
 
-
-    double rank_weight (int r)
-    {
+    double rank_weight(int r) {
         return 1.0 / (r + 1);
     }
 
@@ -683,271 +694,290 @@ struct RankingScore2: Score3Computer<float, double> {
     /// they are the ranks of j and k respectively.
     /// specific version for diff-of-rank weighting, cannot optimized
     /// with a cumulative table
-    double accum_gt_weight_diff (const std::vector<int> & a,
-                                 const std::vector<int> & b)
-    {
+    double accum_gt_weight_diff(
+            const std::vector<int>& a,
+            const std::vector<int>& b) {
         int nb = b.size(), na = a.size();
 
         double accu = 0;
         int j = 0;
         for (int i = 0; i < na; i++) {
             int ai = a[i];
-            while (j < nb && ai >= b[j]) j++;
+            while (j < nb && ai >= b[j])
+                j++;
 
             double accu_i = 0;
             for (int k = j; k < b.size(); k++)
-                accu_i += rank_weight (b[k] - ai);
-
-            accu += rank_weight (ai) * accu_i;
+                accu_i += rank_weight(b[k] - ai);
 
+            accu += rank_weight(ai) * accu_i;
         }
         return accu;
     }
 
-    void init_n_gt ()
-    {
+    void init_n_gt() {
         for (int q = 0; q < nq; q++) {
-            const float *gtd = gt_distances + q * nb;
-            const uint32_t *cb = bcodes;// all same codes
-            float * n_gt_q = & n_gt [qcodes[q] * nc * nc];
+            const float* gtd = gt_distances + q * nb;
+            const uint32_t* cb = bcodes; // all same codes
+            float* n_gt_q = &n_gt[qcodes[q] * nc * nc];
 
-            printf("init gt for q=%d/%d    \r", q, nq); fflush(stdout);
+            printf("init gt for q=%d/%d    \r", q, nq);
+            fflush(stdout);
 
-            std::vector<int> rankv (nb);
-            int * ranks = rankv.data();
+            std::vector<int> rankv(nb);
+            int* ranks = rankv.data();
 
             // elements in each code bin, ordered by rank within each bin
-            std::vector<std::vector<int> > tab (nc);
+            std::vector<std::vector<int>> tab(nc);
 
             { // build rank table
                 IndirectSort s = {gtd};
-                for (int j = 0; j < nb; j++) ranks[j] = j;
-                std::sort (ranks, ranks + nb, s);
+                for (int j = 0; j < nb; j++)
+                    ranks[j] = j;
+                std::sort(ranks, ranks + nb, s);
             }
 
             for (int rank = 0; rank < nb; rank++) {
-                int i = ranks [rank];
-                tab [cb[i]].push_back (rank);
+                int i = ranks[rank];
+                tab[cb[i]].push_back(rank);
             }
 
-
             // this is very expensive. Any suggestion for improvement
             // welcome.
             for (int i = 0; i < nc; i++) {
-                std::vector<int> & di = tab[i];
+                std::vector<int>& di = tab[i];
                 for (int j = 0; j < nc; j++) {
-                    std::vector<int> & dj = tab[j];
-                    n_gt_q [i * nc + j] += accum_gt_weight_diff (di, dj);
-
+                    std::vector<int>& dj = tab[j];
+                    n_gt_q[i * nc + j] += accum_gt_weight_diff(di, dj);
                 }
             }
-
         }
-
     }
-
 };
 
-
 /*****************************************
  * PolysemousTraining
  ******************************************/
 
-
-
-PolysemousTraining::PolysemousTraining ()
-{
+PolysemousTraining::PolysemousTraining() {
     optimization_type = OT_ReproduceDistances_affine;
     ntrain_permutation = 0;
     dis_weight_factor = log(2);
+    // max 20 G RAM
+    max_memory = (size_t)(20) * 1024 * 1024 * 1024;
 }
 
-
-
-void PolysemousTraining::optimize_reproduce_distances (
-       ProductQuantizer &pq) const
-{
-
+void PolysemousTraining::optimize_reproduce_distances(
+        ProductQuantizer& pq) const {
     int dsub = pq.dsub;
 
     int n = pq.ksub;
     int nbits = pq.nbits;
 
-#pragma omp parallel for
+    size_t mem1 = memory_usage_per_thread(pq);
+    int nt = std::min(omp_get_max_threads(), int(pq.M));
+    FAISS_THROW_IF_NOT_FMT(
+            mem1 < max_memory,
+            "Polysemous training will use %zd bytes per thread, while the max is set to %zd",
+            mem1,
+            max_memory);
+
+    if (mem1 * nt > max_memory) {
+        nt = max_memory / mem1;
+        fprintf(stderr,
+                "Polysemous training: WARN, reducing number of threads to %d to save memory",
+                nt);
+    }
+
+#pragma omp parallel for num_threads(nt)
     for (int m = 0; m < pq.M; m++) {
         std::vector<double> dis_table;
 
         // printf ("Optimizing quantizer %d\n", m);
 
-        float * centroids = pq.get_centroids (m, 0);
+        float* centroids = pq.get_centroids(m, 0);
 
         for (int i = 0; i < n; i++) {
             for (int j = 0; j < n; j++) {
-                dis_table.push_back (fvec_L2sqr (centroids + i * dsub,
-                                                 centroids + j * dsub,
-                                                 dsub));
+                dis_table.push_back(fvec_L2sqr(
+                        centroids + i * dsub, centroids + j * dsub, dsub));
             }
         }
 
-        std::vector<int> perm (n);
-        ReproduceWithHammingObjective obj (
-               nbits, dis_table,
-               dis_weight_factor);
-
+        std::vector<int> perm(n);
+        ReproduceWithHammingObjective obj(nbits, dis_table, dis_weight_factor);
 
-        SimulatedAnnealingOptimizer optim (&obj, *this);
+        SimulatedAnnealingOptimizer optim(&obj, *this);
 
         if (log_pattern.size()) {
             char fname[256];
-            snprintf (fname, 256, log_pattern.c_str(), m);
-            printf ("opening log file %s\n", fname);
-            optim.logfile = fopen (fname, "w");
-            FAISS_THROW_IF_NOT_MSG (optim.logfile, "could not open logfile");
+            snprintf(fname, 256, log_pattern.c_str(), m);
+            printf("opening log file %s\n", fname);
+            optim.logfile = fopen(fname, "w");
+            FAISS_THROW_IF_NOT_MSG(optim.logfile, "could not open logfile");
         }
-        double final_cost = optim.run_optimization (perm.data());
+        double final_cost = optim.run_optimization(perm.data());
 
         if (verbose > 0) {
-            printf ("SimulatedAnnealingOptimizer for m=%d: %g -> %g\n",
-                    m, optim.init_cost, final_cost);
+            printf("SimulatedAnnealingOptimizer for m=%d: %g -> %g\n",
+                   m,
+                   optim.init_cost,
+                   final_cost);
         }
 
-        if (log_pattern.size()) fclose (optim.logfile);
+        if (log_pattern.size())
+            fclose(optim.logfile);
 
         std::vector<float> centroids_copy;
         for (int i = 0; i < dsub * n; i++)
-            centroids_copy.push_back (centroids[i]);
+            centroids_copy.push_back(centroids[i]);
 
         for (int i = 0; i < n; i++)
-            memcpy (centroids + perm[i] * dsub,
-                    centroids_copy.data() + i * dsub,
-                    dsub * sizeof(centroids[0]));
-
+            memcpy(centroids + perm[i] * dsub,
+                   centroids_copy.data() + i * dsub,
+                   dsub * sizeof(centroids[0]));
     }
-
 }
 
-
-void PolysemousTraining::optimize_ranking (
-      ProductQuantizer &pq, size_t n, const float *x) const
-{
-
+void PolysemousTraining::optimize_ranking(
+        ProductQuantizer& pq,
+        size_t n,
+        const float* x) const {
     int dsub = pq.dsub;
-
     int nbits = pq.nbits;
 
-    std::vector<uint8_t> all_codes (pq.code_size * n);
+    std::vector<uint8_t> all_codes(pq.code_size * n);
 
-    pq.compute_codes (x, all_codes.data(), n);
+    pq.compute_codes(x, all_codes.data(), n);
 
-    FAISS_THROW_IF_NOT (pq.nbits == 8);
+    FAISS_THROW_IF_NOT(pq.nbits == 8);
 
-    if (n == 0)
-        pq.compute_sdc_table ();
+    if (n == 0) {
+        pq.compute_sdc_table();
+    }
 
 #pragma omp parallel for
     for (int m = 0; m < pq.M; m++) {
         size_t nq, nb;
-        std::vector <uint32_t> codes; // query codes, then db codes
-        std::vector <float> gt_distances; // nq * nb matrix of distances
+        std::vector<uint32_t> codes;     // query codes, then db codes
+        std::vector<float> gt_distances; // nq * nb matrix of distances
 
         if (n > 0) {
-            std::vector<float> xtrain (n * dsub);
+            std::vector<float> xtrain(n * dsub);
             for (int i = 0; i < n; i++)
-                memcpy (xtrain.data() + i * dsub,
-                        x + i * pq.d + m * dsub,
-                        sizeof(float) * dsub);
+                memcpy(xtrain.data() + i * dsub,
+                       x + i * pq.d + m * dsub,
+                       sizeof(float) * dsub);
 
-            codes.resize (n);
+            codes.resize(n);
             for (int i = 0; i < n; i++)
-                codes [i] = all_codes [i * pq.code_size + m];
+                codes[i] = all_codes[i * pq.code_size + m];
 
-            nq = n / 4; nb = n - nq;
-            const float *xq = xtrain.data();
-            const float *xb = xq + nq * dsub;
+            nq = n / 4;
+            nb = n - nq;
+            const float* xq = xtrain.data();
+            const float* xb = xq + nq * dsub;
 
-            gt_distances.resize (nq * nb);
+            gt_distances.resize(nq * nb);
 
-            pairwise_L2sqr (dsub,
-                            nq, xq,
-                            nb, xb,
-                            gt_distances.data());
+            pairwise_L2sqr(dsub, nq, xq, nb, xb, gt_distances.data());
         } else {
             nq = nb = pq.ksub;
-            codes.resize (2 * nq);
+            codes.resize(2 * nq);
             for (int i = 0; i < nq; i++)
-                codes[i] = codes [i + nq] = i;
+                codes[i] = codes[i + nq] = i;
 
-            gt_distances.resize (nq * nb);
+            gt_distances.resize(nq * nb);
 
-            memcpy (gt_distances.data (),
-                    pq.sdc_table.data () + m * nq * nb,
-                    sizeof (float) * nq * nb);
+            memcpy(gt_distances.data(),
+                   pq.sdc_table.data() + m * nq * nb,
+                   sizeof(float) * nq * nb);
         }
 
-        double t0 = getmillisecs ();
+        double t0 = getmillisecs();
 
-        PermutationObjective *obj = new RankingScore2 (
-                  nbits, nq, nb,
-                  codes.data(), codes.data() + nq,
-                  gt_distances.data ());
-        ScopeDeleter1<PermutationObjective> del (obj);
+        PermutationObjective* obj = new RankingScore2(
+                nbits,
+                nq,
+                nb,
+                codes.data(),
+                codes.data() + nq,
+                gt_distances.data());
+        ScopeDeleter1<PermutationObjective> del(obj);
 
         if (verbose > 0) {
             printf("   m=%d, nq=%zd, nb=%zd, intialize RankingScore "
                    "in %.3f ms\n",
-                   m, nq, nb, getmillisecs () - t0);
+                   m,
+                   nq,
+                   nb,
+                   getmillisecs() - t0);
         }
 
-        SimulatedAnnealingOptimizer optim (obj, *this);
+        SimulatedAnnealingOptimizer optim(obj, *this);
 
         if (log_pattern.size()) {
             char fname[256];
-            snprintf (fname, 256, log_pattern.c_str(), m);
-            printf ("opening log file %s\n", fname);
-            optim.logfile = fopen (fname, "w");
-            FAISS_THROW_IF_NOT_FMT (optim.logfile,
-                                    "could not open logfile %s", fname);
+            snprintf(fname, 256, log_pattern.c_str(), m);
+            printf("opening log file %s\n", fname);
+            optim.logfile = fopen(fname, "w");
+            FAISS_THROW_IF_NOT_FMT(
+                    optim.logfile, "could not open logfile %s", fname);
         }
 
-        std::vector<int> perm (pq.ksub);
+        std::vector<int> perm(pq.ksub);
 
-        double final_cost = optim.run_optimization (perm.data());
-        printf ("SimulatedAnnealingOptimizer for m=%d: %g -> %g\n",
-                m, optim.init_cost, final_cost);
+        double final_cost = optim.run_optimization(perm.data());
+        printf("SimulatedAnnealingOptimizer for m=%d: %g -> %g\n",
+               m,
+               optim.init_cost,
+               final_cost);
 
-        if (log_pattern.size()) fclose (optim.logfile);
+        if (log_pattern.size())
+            fclose(optim.logfile);
 
-        float * centroids = pq.get_centroids (m, 0);
+        float* centroids = pq.get_centroids(m, 0);
 
         std::vector<float> centroids_copy;
         for (int i = 0; i < dsub * pq.ksub; i++)
-            centroids_copy.push_back (centroids[i]);
+            centroids_copy.push_back(centroids[i]);
 
         for (int i = 0; i < pq.ksub; i++)
-            memcpy (centroids + perm[i] * dsub,
-                    centroids_copy.data() + i * dsub,
-                    dsub * sizeof(centroids[0]));
-
+            memcpy(centroids + perm[i] * dsub,
+                   centroids_copy.data() + i * dsub,
+                   dsub * sizeof(centroids[0]));
     }
-
 }
 
-
-
-void PolysemousTraining::optimize_pq_for_hamming (ProductQuantizer &pq,
-                                                size_t n, const float *x) const
-{
+void PolysemousTraining::optimize_pq_for_hamming(
+        ProductQuantizer& pq,
+        size_t n,
+        const float* x) const {
     if (optimization_type == OT_None) {
-
     } else if (optimization_type == OT_ReproduceDistances_affine) {
-        optimize_reproduce_distances (pq);
+        optimize_reproduce_distances(pq);
     } else {
-        optimize_ranking (pq, n, x);
+        optimize_ranking(pq, n, x);
     }
 
-    pq.compute_sdc_table ();
-
+    pq.compute_sdc_table();
 }
 
+size_t PolysemousTraining::memory_usage_per_thread(
+        const ProductQuantizer& pq) const {
+    size_t n = pq.ksub;
+
+    switch (optimization_type) {
+        case OT_None:
+            return 0;
+        case OT_ReproduceDistances_affine:
+            return n * n * sizeof(double) * 3;
+        case OT_Ranking_weighted_diff:
+            return n * n * n * sizeof(float);
+    }
+
+    FAISS_THROW_MSG("Invalid optmization type");
+    return 0;
+}
 
 } // namespace faiss