508882fa21
Provide an accelerated neareast neighbor search of points in the dataset using a one layer uniform grid.
238 lines
8.6 KiB
C++
238 lines
8.6 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.
|
|
//
|
|
// Copyright 2017 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
|
|
// Copyright 2014 UT-Battelle, LLC.
|
|
// Copyright 2017 Los Alamos National Security.
|
|
//
|
|
// Under the terms of Contract DE-NA0003525 with NTESS,
|
|
// the U.S. Government retains certain rights in this software.
|
|
//
|
|
// Under the terms of Contract DE-AC52-06NA25396 with Los Alamos National
|
|
// Laboratory (LANL), the U.S. Government retains certain rights in
|
|
// this software.
|
|
//============================================================================
|
|
|
|
#ifndef vtk_m_cont_PointLocatorUniformGrid_h
|
|
#define vtk_m_cont_PointLocatorUniformGrid_h
|
|
|
|
#include <vtkm/cont/ArrayHandleCounting.h>
|
|
#include <vtkm/cont/DeviceAdapter.h>
|
|
#include <vtkm/cont/DeviceAdapterAlgorithm.h>
|
|
#include <vtkm/worklet/DispatcherMapField.h>
|
|
#include <vtkm/worklet/WorkletMapField.h>
|
|
|
|
namespace vtkm
|
|
{
|
|
namespace worklet
|
|
{
|
|
template <typename T>
|
|
class PointLocatorUniformGrid
|
|
{
|
|
public:
|
|
PointLocatorUniformGrid(const vtkm::Vec<T, 3>& _min,
|
|
const vtkm::Vec<T, 3>& _max,
|
|
const vtkm::Vec<vtkm::Id, 3>& _dims)
|
|
: min(_min)
|
|
, max(_max)
|
|
, dims(_dims)
|
|
{
|
|
}
|
|
|
|
class BinPointsWorklet : public vtkm::worklet::WorkletMapField
|
|
{
|
|
public:
|
|
typedef void ControlSignature(FieldIn<> coord, FieldOut<> label);
|
|
|
|
typedef void ExecutionSignature(_1, _2);
|
|
|
|
VTKM_CONT
|
|
BinPointsWorklet(vtkm::Vec<T, 3> _min, vtkm::Vec<T, 3> _max, vtkm::Vec<vtkm::Id, 3> _dims)
|
|
: min(_min)
|
|
, dims(_dims)
|
|
, dxdydz((_max - min) / dims)
|
|
{
|
|
}
|
|
|
|
template <typename CoordVecType, typename IdType>
|
|
VTKM_EXEC void operator()(const CoordVecType& coord, IdType& label) const
|
|
{
|
|
vtkm::Vec<vtkm::Id, 3> ijk = (coord - min) / dxdydz;
|
|
label = ijk[0] + ijk[1] * dims[0] + ijk[2] * dims[0] * dims[1];
|
|
}
|
|
|
|
private:
|
|
vtkm::Vec<T, 3> min;
|
|
vtkm::Vec<vtkm::Id, 3> dims;
|
|
vtkm::Vec<T, 3> dxdydz;
|
|
};
|
|
|
|
class UniformGridSearch : public vtkm::worklet::WorkletMapField
|
|
{
|
|
public:
|
|
typedef void ControlSignature(FieldIn<> query,
|
|
WholeArrayIn<> coordIn,
|
|
WholeArrayIn<IdType> pointId,
|
|
WholeArrayIn<IdType> cellLower,
|
|
WholeArrayIn<IdType> cellUpper,
|
|
FieldOut<IdType> neighborId,
|
|
FieldOut<> distance);
|
|
|
|
typedef void ExecutionSignature(_1, _2, _3, _4, _5, _6, _7);
|
|
|
|
VTKM_CONT
|
|
UniformGridSearch(const vtkm::Vec<T, 3>& _min,
|
|
const vtkm::Vec<T, 3>& _max,
|
|
const vtkm::Vec<vtkm::Id, 3>& _dims)
|
|
: min(_min)
|
|
, dims(_dims)
|
|
, dxdydz((_max - _min) / _dims)
|
|
{
|
|
}
|
|
|
|
|
|
template <typename CoordiVecType,
|
|
typename IdPortalType,
|
|
typename CoordiPortalType,
|
|
typename IdType,
|
|
typename CoordiType>
|
|
VTKM_EXEC void operator()(const CoordiVecType& queryCoord,
|
|
const CoordiPortalType& coordi_Handle,
|
|
const IdPortalType& pointId,
|
|
const IdPortalType& cellLower,
|
|
const IdPortalType& cellUpper,
|
|
IdType& nnId,
|
|
CoordiType& nnDis) const
|
|
{
|
|
auto nlayers = vtkm::Max(vtkm::Max(dims[0], dims[1]), dims[2]);
|
|
|
|
vtkm::Vec<vtkm::Id, 3> xyz = (queryCoord - min) / dxdydz;
|
|
|
|
float min_distance = std::numeric_limits<float>::max();
|
|
vtkm::Id neareast = -1;
|
|
|
|
for (vtkm::Id layer = 0; layer < nlayers; layer++)
|
|
{
|
|
vtkm::Id minx = vtkm::Max(vtkm::Id(), xyz[0] - layer);
|
|
vtkm::Id maxx = vtkm::Min(dims[0] - 1, xyz[0] + layer);
|
|
vtkm::Id miny = vtkm::Max(vtkm::Id(), xyz[1] - layer);
|
|
vtkm::Id maxy = vtkm::Min(dims[1] - 1, xyz[1] + layer);
|
|
vtkm::Id minz = vtkm::Max(vtkm::Id(), xyz[2] - layer);
|
|
vtkm::Id maxz = vtkm::Min(dims[2] - 1, xyz[2] + layer);
|
|
|
|
for (auto i = minx; i <= maxx; i++)
|
|
{
|
|
for (auto j = miny; j <= maxy; j++)
|
|
{
|
|
for (auto k = minz; k <= maxz; k++)
|
|
{
|
|
if (i == (xyz[0] + layer) || i == (xyz[0] - layer) || j == (xyz[1] + layer) ||
|
|
j == (xyz[1] - layer) || k == (xyz[2] + layer) || k == (xyz[2] - layer))
|
|
{
|
|
auto cellid = i + j * dims[0] + k * dims[0] * dims[1];
|
|
auto lower = cellLower.Get(cellid);
|
|
auto upper = cellUpper.Get(cellid);
|
|
for (auto index = lower; index < upper; index++)
|
|
{
|
|
auto pointid = pointId.Get(index);
|
|
auto point = coordi_Handle.Get(pointid);
|
|
auto dx = point[0] - queryCoord[0];
|
|
auto dy = point[1] - queryCoord[1];
|
|
auto dz = point[2] - queryCoord[2];
|
|
|
|
auto distance2 = dx * dx + dy * dy + dz * dz;
|
|
if (distance2 < min_distance)
|
|
{
|
|
neareast = pointid;
|
|
min_distance = distance2;
|
|
nlayers = layer + 2;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
nnId = neareast;
|
|
nnDis = vtkm::Sqrt(min_distance);
|
|
};
|
|
|
|
private:
|
|
vtkm::Vec<T, 3> min;
|
|
vtkm::Vec<vtkm::Id, 3> dims;
|
|
vtkm::Vec<T, 3> dxdydz;
|
|
};
|
|
|
|
/// \brief Construct a 3D uniform grid for nearest neighbor search.
|
|
///
|
|
/// \param coords An ArrayHandle of x, y, z coordinates of input points.
|
|
/// \param device Tag for selecting device adapter
|
|
template <typename DeviceAdapter>
|
|
void Build(const vtkm::cont::ArrayHandle<vtkm::Vec<T, 3>>& coords, DeviceAdapter)
|
|
{
|
|
typedef vtkm::cont::DeviceAdapterAlgorithm<DeviceAdapter> Algorithm;
|
|
// generate unique id for each input point
|
|
vtkm::cont::ArrayHandleCounting<vtkm::Id> pointCounting(0, 1, coords.GetNumberOfValues());
|
|
Algorithm::Copy(pointCounting, pointIds);
|
|
|
|
// bin points into cells and give each of them the cell id.
|
|
BinPointsWorklet cellIdWorklet(min, max, dims);
|
|
vtkm::worklet::DispatcherMapField<BinPointsWorklet> dispatchCellId(cellIdWorklet);
|
|
dispatchCellId.Invoke(coords, cellIds);
|
|
|
|
// Group points of the same cell together by sorting them according to the cell ids
|
|
Algorithm::SortByKey(cellIds, pointIds);
|
|
|
|
// for each cell, find the lower and upper bound of indices to the sorted point ids.
|
|
vtkm::cont::ArrayHandleCounting<vtkm::Id> cell_ids_counting(0, 1, dims[0] * dims[1] * dims[2]);
|
|
Algorithm::UpperBounds(cellIds, cell_ids_counting, cellUpper);
|
|
Algorithm::LowerBounds(cellIds, cell_ids_counting, cellLower);
|
|
};
|
|
|
|
/// \brief Nearest neighbor search using a Uniform Grid
|
|
///
|
|
/// Parallel search of nearesat neighbor for each point in the \c queryPoints in the set of
|
|
/// \c coords. Returns neareast neighbot in \c nearestNeighborIds and distances to nearest
|
|
/// neighbor in \c distances.
|
|
///
|
|
/// \param coords Point coordinates for training dataset.
|
|
/// \param queryPoints Point coordinates to query for nearest neighbors.
|
|
/// \param nearestNeighborIds Neareast neighbor in the training dataset for each points in
|
|
/// the test set
|
|
/// \param distances Distance between query points and their nearest neighbors.
|
|
/// \param device Tag for selecting device adapter.
|
|
template <typename DeviceAdapter>
|
|
void Run(const vtkm::cont::ArrayHandle<vtkm::Vec<T, 3>>& coords,
|
|
const vtkm::cont::ArrayHandle<vtkm::Vec<T, 3>>& queryPoints,
|
|
vtkm::cont::ArrayHandle<vtkm::Id>& nearestNeighborIds,
|
|
vtkm::cont::ArrayHandle<T>& distances,
|
|
DeviceAdapter)
|
|
{
|
|
UniformGridSearch uniformGridSearch(min, max, dims);
|
|
|
|
vtkm::worklet::DispatcherMapField<UniformGridSearch, DeviceAdapter> searchDispatcher(
|
|
uniformGridSearch);
|
|
searchDispatcher.Invoke(
|
|
queryPoints, coords, pointIds, cellLower, cellUpper, nearestNeighborIds, distances);
|
|
};
|
|
|
|
private:
|
|
vtkm::Vec<T, 3> min;
|
|
vtkm::Vec<T, 3> max;
|
|
vtkm::Vec<vtkm::Id, 3> dims;
|
|
|
|
vtkm::cont::ArrayHandle<vtkm::Id> pointIds;
|
|
vtkm::cont::ArrayHandle<vtkm::Id> cellIds;
|
|
vtkm::cont::ArrayHandle<vtkm::Id> cellLower;
|
|
vtkm::cont::ArrayHandle<vtkm::Id> cellUpper;
|
|
};
|
|
}
|
|
}
|
|
#endif //vtk_m_cont_PointLocatorUniformGrid_h
|