All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/matthijs/github_faiss/faiss/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
faiss::knn_L2sqr_base_shift
void knn_L2sqr_base_shift(const float *x, const float *y, size_t d, size_t nx, size_t ny, float_maxheap_array_t *res, const float *base_shift)
Definition: utils.cpp:870
faiss::IndexFlatL2BaseShift::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexFlat.cpp:124
faiss::IndexRefineFlat::reset
virtual void reset() override
removes all elements from the database.
Definition: IndexFlat.cpp:184
faiss::CMax
Definition: Heap.h:49
faiss::Index::reset
virtual void reset()=0
removes all elements from the database.
faiss::IndexFlat1D::continuous_update
bool continuous_update
is the permutation updated continuously?
Definition: IndexFlat.h:143
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::IndexFlat::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexFlat.cpp:54
faiss::IndexFlat::reset
virtual void reset() override
removes all elements from the database.
Definition: IndexFlat.cpp:48
faiss::IndexFlat1D::update_permutation
void update_permutation()
Definition: IndexFlat.cpp:283
faiss::IndexFlat::reconstruct
virtual void reconstruct(idx_t key, float *recons) const override
Definition: IndexFlat.cpp:109
faiss::IndexFlat1D::add
virtual void add(idx_t n, const float *x) override
Definition: IndexFlat.cpp:293
faiss::IndexRefineFlat::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexFlat.cpp:219
faiss::IndexRefineFlat::base_index
Index * base_index
faster index to pre-select the vectors that should be filtered
Definition: IndexFlat.h:111
faiss::IndexRefineFlat::refine_index
IndexFlat refine_index
storage for full vectors
Definition: IndexFlat.h:108
faiss::IndexRefineFlat::own_fields
bool own_fields
should the base index be deallocated?
Definition: IndexFlat.h:112
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::IndexFlat::range_search
virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result) const override
Definition: IndexFlat.cpp:70
faiss::IndexRefineFlat::train
virtual void train(idx_t n, const float *x) override
Definition: IndexFlat.cpp:171
faiss::Index::add
virtual void add(idx_t n, const float *x)=0
faiss::HeapArray
Definition: Heap.h:351
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::range_search_inner_product
void range_search_inner_product(const float *x, const float *y, size_t d, size_t nx, size_t ny, float radius, RangeSearchResult *res)
same as range_search_L2sqr for the inner product similarity
Definition: utils.cpp:1166
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::Index::get_typename
virtual std::string get_typename() const
Definition: Index.h:188
faiss::knn_inner_product
void knn_inner_product(const float *x, const float *y, size_t d, size_t nx, size_t ny, float_minheap_array_t *res)
Definition: utils.cpp:830
faiss::IndexFlat::add
virtual void add(idx_t n, const float *x) override
Definition: IndexFlat.cpp:41
faiss::IndexFlat1D::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Warn: the distances returned are L1 not L2.
Definition: IndexFlat.cpp:306
faiss::Index::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
faiss::range_search_L2sqr
void range_search_L2sqr(const float *x, const float *y, size_t d, size_t nx, size_t ny, float radius, RangeSearchResult *res)
Definition: utils.cpp:1151
faiss::RangeSearchResult
Definition: AuxIndexStructures.h:35
faiss::IndexFlat::compute_distance_subset
void compute_distance_subset(idx_t n, const float *x, idx_t k, float *distances, const idx_t *labels) const
Definition: IndexFlat.cpp:85
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::Index
Definition: Index.h:61
faiss::knn_L2sqr
void knn_L2sqr(const float *x, const float *y, size_t d, size_t nx, size_t ny, float_maxheap_array_t *res)
Definition: utils.cpp:850
faiss::Index::is_trained
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:71
faiss::Index::train
virtual void train(idx_t n, const float *x)
Definition: Index.h:92
faiss::IndexFlat::xb
std::vector< float > xb
database vectors, size ntotal * d
Definition: IndexFlat.h:26
faiss::IndexRefineFlat::k_factor
float k_factor
Definition: IndexFlat.h:116
faiss::IndexFlat1D::reset
virtual void reset() override
removes all elements from the database.
Definition: IndexFlat.cpp:300
faiss::IndexFlat1D::perm
std::vector< idx_t > perm
sorted database indices
Definition: IndexFlat.h:145
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44
faiss::CMin
Definition: Heap.h:54
faiss::IndexRefineFlat::add
virtual void add(idx_t n, const float *x) override
Definition: IndexFlat.cpp:177