NearestNeighbours.h

Go to the documentation of this file.

// Mantid Repository : https://github.com/mantidproject/mantid
//
// Copyright © 2016 ISIS Rutherford Appleton Laboratory UKRI,
//   NScD Oak Ridge National Laboratory, European Spallation Source,
//   Institut Laue - Langevin & CSNS, Institute of High Energy Physics, CAS
// SPDX - License - Identifier: GPL - 3.0 +
#pragma once
 
#include "MantidKernel/ANN/ANN.h"
#include "MantidKernel/DllConfig.h"
 
#include <Eigen/Core>
#include <memory>
#include <vector>
 
namespace Mantid {
namespace Kernel {
 
//------------------------------------------------------------------------------
// Helper classes
//------------------------------------------------------------------------------
 
class NNDataPoints {
public:
  NNDataPoints(const int nPts, const int nElems) : m_nPts(nPts) { m_data = annAllocPts(m_nPts, nElems); }
 
  ~NNDataPoints() { annDeallocPts(m_data); }
 
  ANNpointArray rawData() { return m_data; }
 
  ANNcoord *mutablePoint(const int i) {
    if (i < m_nPts)
      return m_data[i];
    else
      return nullptr;
  }
 
private:
  const int m_nPts;
  ANNpointArray m_data;
};
 
//------------------------------------------------------------------------------
// NearestNeighbours implementation
//------------------------------------------------------------------------------
 
template <int N = 3> class DLLExport NearestNeighbours {
 
public:
  // typedefs for code brevity
  using VectorType = Eigen::Matrix<double, N, 1>;
  using NearestNeighbourResults = std::vector<std::tuple<VectorType, size_t, double>>;
 
  NearestNeighbours(const std::vector<VectorType> &points) {
    const auto numPoints = static_cast<int>(points.size());
    if (numPoints == 0)
      std::runtime_error("Need at least one point to initialise NearestNeighbours.");
 
    m_dataPoints = std::make_unique<NNDataPoints>(numPoints, static_cast<int>(N));
 
    for (size_t i = 0; i < points.size(); ++i) {
      Eigen::Map<VectorType>(m_dataPoints->mutablePoint(static_cast<int>(i)), N, 1) = points[i];
    }
    m_kdTree = std::make_unique<ANNkd_tree>(m_dataPoints->rawData(), numPoints, static_cast<int>(N));
  }
 
  ~NearestNeighbours() { annClose(); }
 
  NearestNeighbours(const NearestNeighbours &) = delete;
 
  NearestNeighbourResults findNearest(const VectorType &pos, const size_t k = 1, const double error = 0.0) {
    const auto numNeighbours = static_cast<int>(k);
    // create arrays to store the indices & distances of nearest neighbours
    auto nnIndexList = std::unique_ptr<ANNidx[]>(new ANNidx[numNeighbours]);
    auto nnDistList = std::unique_ptr<ANNdist[]>(new ANNdist[numNeighbours]);
 
    // create ANNpoint from Eigen array
    auto point = std::unique_ptr<ANNcoord[]>(annAllocPt(N));
    Eigen::Map<VectorType>(point.get(), N, 1) = pos;
 
    // find the k nearest neighbours
    m_kdTree->annkSearch(point.get(), numNeighbours, nnIndexList.get(), nnDistList.get(), error);
 
    return makeResults(k, std::move(nnIndexList), std::move(nnDistList));
  }
 
private:
  NearestNeighbourResults makeResults(const size_t k, const std::unique_ptr<ANNidx[]> nnIndexList,
                                      const std::unique_ptr<ANNdist[]> nnDistList) {
    NearestNeighbourResults results;
    results.reserve(k);
 
    for (size_t i = 0; i < k; ++i) {
      // create Eigen array from ANNpoint
      auto pos = m_dataPoints->mutablePoint(nnIndexList[i]);
      VectorType point = Eigen::Map<VectorType>(pos, N, 1);
      results.emplace_back(point, nnIndexList[i], nnDistList[i]);
    }
 
    return results;
  }
 
  std::unique_ptr<NNDataPoints> m_dataPoints;
  std::unique_ptr<ANNkd_tree> m_kdTree;
};
} // namespace Kernel
} // namespace Mantid