Skip to content

Commit

Permalink
update nanoflann
Browse files Browse the repository at this point in the history
  • Loading branch information
j042 committed Jan 3, 2024
1 parent 7fcba64 commit 941705c
Show file tree
Hide file tree
Showing 2 changed files with 136 additions and 24 deletions.
158 changes: 135 additions & 23 deletions third_party/nanoflann.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
*
* Copyright 2008-2009 Marius Muja ([email protected]). All rights reserved.
* Copyright 2008-2009 David G. Lowe ([email protected]). All rights reserved.
* Copyright 2011-2022 Jose Luis Blanco ([email protected]).
* Copyright 2011-2023 Jose Luis Blanco ([email protected]).
* All rights reserved.
*
* THE BSD LICENSE
Expand Down Expand Up @@ -50,6 +50,7 @@
#include <cassert>
#include <cmath> // for abs()
#include <cstdlib> // for abs()
#include <cstdint>
#include <functional> // std::reference_wrapper
#include <future>
#include <istream>
Expand All @@ -60,7 +61,7 @@
#include <vector>

/** Library version: 0xMmP (M=Major,m=minor,P=patch) */
#define NANOFLANN_VERSION 0x150
#define NANOFLANN_VERSION 0x151

// Avoid conflicting declaration of min/max macros in Windows headers
#if !defined(NOMINMAX) && \
Expand Down Expand Up @@ -158,6 +159,8 @@ inline typename std::enable_if<!has_assign<Container>::value, void>::type

/** @addtogroup result_sets_grp Result set classes
* @{ */

/** Result set for KNN searches (N-closest neighbors) */
template <
typename _DistanceType, typename _IndexType = size_t,
typename _CountType = size_t>
Expand Down Expand Up @@ -190,8 +193,93 @@ class KNNResultSet
}

CountType size() const { return count; }
bool empty() const { return count == 0; }
bool full() const { return count == capacity; }

/**
* Called during search to add an element matching the criteria.
* @return true if the search should be continued, false if the results are
* sufficient
*/
bool addPoint(DistanceType dist, IndexType index)
{
CountType i;
for (i = count; i > 0; --i)
{
/** If defined and two points have the same distance, the one with
* the lowest-index will be returned first. */
#ifdef NANOFLANN_FIRST_MATCH
if ((dists[i - 1] > dist) ||
((dist == dists[i - 1]) && (indices[i - 1] > index)))
{
#else
if (dists[i - 1] > dist)
{
#endif
if (i < capacity)
{
dists[i] = dists[i - 1];
indices[i] = indices[i - 1];
}
}
else
break;
}
if (i < capacity)
{
dists[i] = dist;
indices[i] = index;
}
if (count < capacity) count++;

// tell caller that the search shall continue
return true;
}

DistanceType worstDist() const { return dists[capacity - 1]; }
};

/** Result set for RKNN searches (N-closest neighbors with a maximum radius) */
template <
typename _DistanceType, typename _IndexType = size_t,
typename _CountType = size_t>
class RKNNResultSet
{
public:
using DistanceType = _DistanceType;
using IndexType = _IndexType;
using CountType = _CountType;

bool full() const { return count == capacity; }
private:
IndexType* indices;
DistanceType* dists;
CountType capacity;
CountType count;
DistanceType maximumSearchDistanceSquared;

public:
explicit RKNNResultSet(
CountType capacity_, DistanceType maximumSearchDistanceSquared_)
: indices(nullptr),
dists(nullptr),
capacity(capacity_),
count(0),
maximumSearchDistanceSquared(maximumSearchDistanceSquared_)
{
}

void init(IndexType* indices_, DistanceType* dists_)
{
indices = indices_;
dists = dists_;
count = 0;
if (capacity)
dists[capacity - 1] = maximumSearchDistanceSquared;
}

CountType size() const { return count; }
bool empty() const { return count == 0; }
bool full() const { return count == capacity; }

/**
* Called during search to add an element matching the criteria.
Expand Down Expand Up @@ -744,7 +832,7 @@ struct SearchParameters
*/
class PooledAllocator
{
static constexpr size_t WORDSIZE = 16;
static constexpr size_t WORDSIZE = 16; //WORDSIZE must >= 8
static constexpr size_t BLOCKSIZE = 8192;

/* We maintain memory alignment to word boundaries by requiring that all
Expand All @@ -753,9 +841,7 @@ class PooledAllocator
/* Minimum number of bytes requested at a time from the system. Must be
* multiple of WORDSIZE. */

using Offset = uint32_t;
using Size = uint32_t;
using Dimension = int32_t;
using Size = size_t;

Size remaining_ = 0; //!< Number of bytes left in current block of storage
void* base_ = nullptr; //!< Pointer to base of current block of storage
Expand Down Expand Up @@ -817,9 +903,9 @@ class PooledAllocator

/* Allocate new storage. */
const Size blocksize =
(size + sizeof(void*) + (WORDSIZE - 1) > BLOCKSIZE)
? size + sizeof(void*) + (WORDSIZE - 1)
: BLOCKSIZE;
size > BLOCKSIZE
? size + WORDSIZE
: BLOCKSIZE + WORDSIZE;

// use the standard C malloc to allocate memory
void* m = ::malloc(blocksize);
Expand All @@ -833,12 +919,8 @@ class PooledAllocator
static_cast<void**>(m)[0] = base_;
base_ = m;

Size shift = 0;
// int size_t = (WORDSIZE - ( (((size_t)m) + sizeof(void*)) &
// (WORDSIZE-1))) & (WORDSIZE-1);

remaining_ = blocksize - sizeof(void*) - shift;
loc_ = (static_cast<char*>(m) + sizeof(void*) + shift);
remaining_ = blocksize - WORDSIZE;
loc_ = static_cast<char*>(m) + WORDSIZE;
}
void* rloc = loc_;
loc_ = static_cast<char*>(loc_) + size;
Expand Down Expand Up @@ -1220,25 +1302,26 @@ class KDTreeBaseClass
}
ElementType max_spread = -1;
cutfeat = 0;
ElementType min_elem = 0, max_elem = 0;
for (Dimension i = 0; i < dims; ++i)
{
ElementType span = bbox[i].high - bbox[i].low;
if (span > (1 - EPS) * max_span)
{
ElementType min_elem, max_elem;
computeMinMax(obj, ind, count, i, min_elem, max_elem);
ElementType spread = max_elem - min_elem;
ElementType min_elem_, max_elem_;
computeMinMax(obj, ind, count, i, min_elem_, max_elem_);
ElementType spread = max_elem_ - min_elem_;
if (spread > max_spread)
{
cutfeat = i;
max_spread = spread;
min_elem = min_elem_;
max_elem = max_elem_;
}
}
}
// split in the middle
DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2;
ElementType min_elem, max_elem;
computeMinMax(obj, ind, count, cutfeat, min_elem, max_elem);

if (split_val < min_elem)
cutval = min_elem;
Expand Down Expand Up @@ -1680,6 +1763,34 @@ class KDTreeSingleIndexAdaptor
return resultSet.size();
}

/**
* Find the first N neighbors to \a query_point[0:dim-1] within a maximum
* radius. The output is given as a vector of pairs, of which the first
* element is a point index and the second the corresponding distance.
* Previous contents of \a IndicesDists are cleared.
*
* \sa radiusSearch, findNeighbors
* \return Number `N` of valid points in the result set.
*
* \note If L2 norms are used, all returned distances are actually squared
* distances.
*
* \note Only the first `N` entries in `out_indices` and `out_distances`
* will be valid. Return is less than `num_closest` only if the
* number of elements in the tree is less than `num_closest`.
*/
Size rknnSearch(
const ElementType* query_point, const Size num_closest,
IndexType* out_indices, DistanceType* out_distances,
const DistanceType& radius) const
{
nanoflann::RKNNResultSet<DistanceType, IndexType> resultSet(
num_closest, radius);
resultSet.init(out_indices, out_distances);
findNeighbors(resultSet, query_point);
return resultSet.size();
}

/** @} */

public:
Expand Down Expand Up @@ -2060,11 +2171,12 @@ class KDTreeSingleIndexDynamicAdaptor_
*/
Size knnSearch(
const ElementType* query_point, const Size num_closest,
IndexType* out_indices, DistanceType* out_distances) const
IndexType* out_indices, DistanceType* out_distances,
const SearchParameters& searchParams = {}) const
{
nanoflann::KNNResultSet<DistanceType, IndexType> resultSet(num_closest);
resultSet.init(out_indices, out_distances);
findNeighbors(resultSet, query_point);
findNeighbors(resultSet, query_point, searchParams);
return resultSet.size();
}

Expand Down

0 comments on commit 941705c

Please sign in to comment.