vtk-m/vtkm/worklet/KdTree3D.h
2020-06-01 16:58:32 -04:00

77 lines
2.9 KiB
C++

//============================================================================
// Copyright (c) Kitware, Inc.
// All rights reserved.
// See LICENSE.txt for details.
//
// This software is distributed WITHOUT ANY WARRANTY; without even
// the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the above copyright notice for more information.
//============================================================================
#ifndef vtkm_m_worklet_KdTree3D_h
#define vtkm_m_worklet_KdTree3D_h
#include <vtkm/worklet/spatialstructure/KdTree3DConstruction.h>
#include <vtkm/worklet/spatialstructure/KdTree3DNNSearch.h>
namespace vtkm
{
namespace worklet
{
class KdTree3D
{
public:
KdTree3D() = default;
/// \brief Construct a 3D KD-tree for 3D point positions.
///
/// \tparam CoordType type of the x, y, z component of the point coordinates.
/// \tparam CoordStorageTag
/// \param coords An ArrayHandle of x, y, z coordinates of input points.
///
template <typename CoordType, typename CoordStorageTag>
void Build(const vtkm::cont::ArrayHandle<vtkm::Vec<CoordType, 3>, CoordStorageTag>& coords)
{
vtkm::worklet::spatialstructure::KdTree3DConstruction().Run(
coords, this->PointIds, this->SplitIds);
}
/// \brief Nearest neighbor search using KD-Tree
///
/// Parallel search of nearest neighbor for each point in the \c queryPoints in the the set of
/// \c coords. Returns nearest neighbor in \c nearestNeighborId and distance to nearest neighbor
/// in \c distances.
///
/// \tparam CoordType
/// \tparam CoordStorageTag1
/// \tparam CoordStorageTag2
/// \tparam DeviceAdapter
/// \param coords Point coordinates for training data set (haystack)
/// \param queryPoints Point coordinates to query for nearest neighbor (needles).
/// \param nearestNeighborIds Nearest neighbor in the traning data set for each points in the
/// testing set
/// \param distances Distances between query points and their nearest neighbors.
/// \param deviceId Tag for selecting device adapter.
template <typename CoordType,
typename CoordStorageTag1,
typename CoordStorageTag2,
typename DeviceAdapter>
void Run(const vtkm::cont::ArrayHandle<vtkm::Vec<CoordType, 3>, CoordStorageTag1>& coords,
const vtkm::cont::ArrayHandle<vtkm::Vec<CoordType, 3>, CoordStorageTag2>& queryPoints,
vtkm::cont::ArrayHandle<vtkm::Id>& nearestNeighborIds,
vtkm::cont::ArrayHandle<CoordType>& distances,
DeviceAdapter deviceId)
{
vtkm::worklet::spatialstructure::KdTree3DNNSearch().Run(
coords, this->PointIds, this->SplitIds, queryPoints, nearestNeighborIds, distances, deviceId);
}
private:
vtkm::cont::ArrayHandle<vtkm::Id> PointIds;
vtkm::cont::ArrayHandle<vtkm::Id> SplitIds;
};
}
} // namespace vtkm::worklet
#endif // vtkm_m_worklet_Kdtree3D_h