python/nabo.cpp
Go to the documentation of this file.
1 #include <Python.h>
2 #include <boost/python.hpp>
3 #include <numpy/arrayobject.h>
4 #include "../nabo/nabo.h"
5 #include <iostream>
6 #include <cassert>
7 
8 using namespace boost::python;
9 
14 typedef Eigen::Map<NNSNabo::Matrix> MappedEigenDoubleMatrix;
15 typedef Eigen::Map<NNSNabo::IndexMatrix> MappedEigenIndexMatrix;
16 
17 static const double infD = std::numeric_limits<double>::infinity();
18 static const Index maxI = std::numeric_limits<Index>::max();
19 
20 #if PY_MAJOR_VERSION >= 3
21 int
22 init_numpy()
23 {
24  import_array();
25  return 0;
26 }
27 #else
28 void
30 {
31  import_array();
32 }
33 #endif
34 
35 void matrixSizeFromPythonArray(const PyObject* cloudObj, int& rowCount, int& colCount)
36 {
37  assert(PyArray_CHKFLAGS(cloudObj, NPY_C_CONTIGUOUS) || PyArray_CHKFLAGS(cloudObj, NPY_F_CONTIGUOUS));
38  assert(PyArray_NDIM(cloudObj) == 2);
39  const npy_intp *shape = PyArray_DIMS(cloudObj);
40  if (PyArray_CHKFLAGS(cloudObj, NPY_F_CONTIGUOUS))
41  {
42  colCount = shape[1];
43  rowCount = shape[0];
44  }
45  else
46  {
47  colCount = shape[0];
48  rowCount = shape[1];
49  }
50 }
51 
52 void checkPythonArray(const PyObject* cloudObj, const char *paramName)
53 {
54  std::string startMsg("Argument \"");
55  startMsg += paramName;
56  startMsg += "\" ";
57 
58  if (!PyArray_Check(cloudObj))
59  throw std::runtime_error(startMsg + "must be a multi-dimensional array");
60  const int nDim = PyArray_NDIM(cloudObj);
61  if (nDim != 2)
62  throw std::runtime_error(startMsg + "must be a two-dimensional array");
63  if (PyArray_TYPE(cloudObj) != NPY_FLOAT64)
64  throw std::runtime_error(startMsg + "must hold doubles");
65  if (!PyArray_CHKFLAGS(cloudObj, NPY_C_CONTIGUOUS) && !PyArray_CHKFLAGS(cloudObj, NPY_F_CONTIGUOUS))
66  throw std::runtime_error(startMsg + "must be a continuous array");
67 }
68 
69 MappedEigenDoubleMatrix* eigenFromBoostPython(const object cloudIn, const char *paramName)
70 {
71  int dimCount, pointCount;
72  const PyObject *cloudObj(cloudIn.ptr());
73 
74  checkPythonArray(cloudObj, paramName);
75 
76  matrixSizeFromPythonArray(cloudObj, dimCount, pointCount);
77  double* cloudData(reinterpret_cast<double*>(PyArray_DATA(cloudObj)));
78 
79  return new MappedEigenDoubleMatrix(cloudData, dimCount, pointCount);
80 }
81 
82 void eigenFromBoostPython(NNSNabo::Matrix& cloudOut, const object cloudIn, const char *paramName)
83 {
84  int dimCount, pointCount;
85  const PyObject *cloudObj(cloudIn.ptr());
86 
87  checkPythonArray(cloudObj, paramName);
88 
89  matrixSizeFromPythonArray(cloudObj, dimCount, pointCount);
90  cloudOut.resize(dimCount, pointCount);
91 
92  memcpy(cloudOut.data(), PyArray_DATA(cloudObj), pointCount*dimCount*sizeof(double));
93 }
94 
96 {
97 public:
98  NearestNeighbourSearch(const object pycloud, const SearchType searchType = NNSNabo::KDTREE_LINEAR_HEAP, const Index dim = maxI, const dict params = dict())
99  {
100  // build cloud
101  eigenFromBoostPython(cloud, pycloud, "cloud");
102 
103  // build params
104  Nabo::Parameters _params;
105 
106 #if PY_MAJOR_VERSION >=3
107  object it = params.items();
108 #else
109  object it = params.iteritems();
110 #endif
111 
112  for(int i = 0; i < len(params); ++i)
113  {
114  const tuple item(it.attr("next")());
115  const std::string key = extract<std::string>(item[0]);
116  const object val(item[1]);
117  const std::string valType(val.ptr()->ob_type->tp_name);
118  if (valType == "int")
119  {
120  const int iVal = extract<int>(val);
121  if (iVal >= 0)
122  _params[key] = (unsigned)iVal;
123  else
124  _params[key] = iVal;
125  }
126  }
127 
128  // create search
129  nns = NNSNabo ::create(cloud, dim, searchType, 0, _params);
130  }
131 
133  {
134  delete nns;
135  }
136 
137  tuple knn(const object query, const Index k = 1, const double epsilon = 0, const unsigned optionFlags = 0, const double maxRadius = infD)
138  {
139  // map query and create output matrices
140  MappedEigenDoubleMatrix* mappedQuery(eigenFromBoostPython(query, "query"));
141  NNSNabo::IndexMatrix indexMatrix(k, mappedQuery->cols());
142  NNSNabo::Matrix dists2Matrix(k, mappedQuery->cols());
143 
144  // do the search
145  nns->knn(*mappedQuery, indexMatrix, dists2Matrix, k, epsilon, optionFlags, maxRadius);
146 
147  // build resulting python types
148  npy_intp retDims[2] = { mappedQuery->cols(), k };
149  const int dataCount(k * mappedQuery->cols());
150  PyObject* dists2 = PyArray_EMPTY(2, retDims, PyArray_DOUBLE, PyArray_ISFORTRAN(query.ptr()));
151  memcpy(PyArray_DATA(dists2), dists2Matrix.data(), dataCount*sizeof(double));
152  PyObject* indices = PyArray_EMPTY(2, retDims, PyArray_INT, PyArray_ISFORTRAN(query.ptr()));
153  memcpy(PyArray_DATA(indices), indexMatrix.data(), dataCount*sizeof(int));
154  delete mappedQuery;
155 
156  // return results
157  return make_tuple(object(handle<>(indices)), object(handle<>(dists2)));
158  }
159 
160 protected:
163 };
164 
165 BOOST_PYTHON_MEMBER_FUNCTION_OVERLOADS(knn_overloads, knn, 1, 5)
166 
168 {
169  init_numpy();
170 
171  enum_<SearchType>("SearchType", "Type of algorithm used for search.")
172  .value("BRUTE_FORCE", NNSNabo::BRUTE_FORCE)
173  .value("KDTREE_LINEAR_HEAP", NNSNabo::KDTREE_LINEAR_HEAP)
174  .value("KDTREE_TREE_HEAP", NNSNabo::KDTREE_TREE_HEAP)
175  ;
176 
177  enum_<SearchOptionFlags>("SearchOptionFlags", "Flags you can OR when creating search.")
178  .value("ALLOW_SELF_MATCH", NNSNabo::ALLOW_SELF_MATCH)
179  .value("SORT_RESULTS", NNSNabo::SORT_RESULTS)
180  ;
181 
182  class_<NearestNeighbourSearch>(
183  "NearestNeighbourSearch",
184  "Nearest-neighbour search object, containing the data, on which you can do the knn(...) query.\n\n"
185  "The data and query must be continuous numpy arrays.\n"
186  "As numpy proposes both C and Fortran data orders, pynabo\n"
187  "will always consider the contiguous dimension to be coordinates\n"
188  "of points, regardless of order, as this provides the fastest\n"
189  "possible execution. The return values of knn(...) will have\n"
190  "the same order as the query and will have the different results\n"
191  "of each point in the contiguous dimension."
192  ,
193  init<const object, optional<const SearchType, const Index, const dict> >(
194  "Create a nearest-neighbour search.\n\n"
195  "Arguments:\n"
196  " data -- data-point cloud in which to search, must be a numpy array\n"
197  " searchType -- type of search, default: KDTREE_LINEAR_HEAP\n"
198  " dim -- number of dimensions to consider, must be lower or equal to cloud.rows(), default: dim of data\n"
199  " creationOptionFlags -- creation options, a bitwise OR of elements of CreationOptionFlags",
200  args("self", "data", "searchType", "dim", "creationOptionFlags, default: 0")
201  )
202  )
203  .def("knn", &NearestNeighbourSearch::knn,
204  knn_overloads(
205  args("self", "query", "k", "epsilon", "optionFlags", "maxRadius"),
206  "Find the k nearest neighbours of query in data.\n\n"
207  "Arguments:\n"
208  " query -- query points, must be a numpy array\n"
209  " k -- number of nearest neighbour requested, default: 1\n"
210  " epsilon -- maximal ratio of error for approximate search, 0 for exact search; has no effect if the number of neighbour found is smaller than the number requested; default: 0.\n"
211  " optionFlags -- search options, a bitwise OR of elements of SearchOptionFlags, default: 0\n"
212  " maxRadius -- maximum radius in which to search, can be used to prune search, is not affected by epsilon, default: inf\n\n"
213  "Returns:\n"
214  " A tuple of two 2D numpy arrays, the first containing indices to points in data, the other containing squared distances."
215  )
216  )
217  ;
218 }
sort points by distances, when k > 1; do not sort by default
Definition: nabo.h:311
Eigen::Matrix< T, Eigen::Dynamic, Eigen::Dynamic > Matrix
a column-major Eigen matrix in which each column is a point; this matrix has dim rows ...
Definition: nabo.h:263
void checkPythonArray(const PyObject *cloudObj, const char *paramName)
Definition: python/nabo.cpp:52
allows the return of the same point as the query, if this point is in the data cloud; forbidden by de...
Definition: nabo.h:310
static NearestNeighbourSearch * create(const CloudType &cloud, const Index dim=std::numeric_limits< Index >::max(), const SearchType preferedType=KDTREE_LINEAR_HEAP, const unsigned creationOptionFlags=0, const Parameters &additionalParameters=Parameters())
Create a nearest-neighbour search.
Definition: nabo/nabo.cpp:135
kd-tree with tree heap, good for large k (~from 30)
Definition: nabo.h:294
static const double infD
Definition: python/nabo.cpp:17
tuple knn(const object query, const Index k=1, const double epsilon=0, const unsigned optionFlags=0, const double maxRadius=infD)
void matrixSizeFromPythonArray(const PyObject *cloudObj, int &rowCount, int &colCount)
Definition: python/nabo.cpp:35
Nearest neighbour search interface, templatized on scalar type.
Definition: nabo.h:258
SearchType
type of search
Definition: nabo.h:290
nns
Definition: test.py:7
Eigen::Map< NNSNabo::Matrix > MappedEigenDoubleMatrix
Definition: python/nabo.cpp:14
SearchOptionFlags
search option
Definition: nabo.h:308
static const Index maxI
Definition: python/nabo.cpp:18
int Index
an index to a Vector or a Matrix, for refering to data points
Definition: nabo.h:267
NNSNabo::Matrix cloud
Eigen::Matrix< Index, Eigen::Dynamic, Eigen::Dynamic > IndexMatrix
a matrix of indices to data points
Definition: nabo.h:271
void init_numpy()
Definition: python/nabo.cpp:29
NNSNabo::SearchType SearchType
Definition: python/nabo.cpp:12
Parameter vector.
Definition: nabo.h:231
NNSNabo::SearchOptionFlags SearchOptionFlags
Definition: python/nabo.cpp:13
MappedEigenDoubleMatrix * eigenFromBoostPython(const object cloudIn, const char *paramName)
Definition: python/nabo.cpp:69
BOOST_PYTHON_MODULE(pynabo)
NearestNeighbourSearch(const object pycloud, const SearchType searchType=NNSNabo::KDTREE_LINEAR_HEAP, const Index dim=maxI, const dict params=dict())
Definition: python/nabo.cpp:98
NNSNabo::Index Index
Definition: python/nabo.cpp:11
Eigen::Map< NNSNabo::IndexMatrix > MappedEigenIndexMatrix
Definition: python/nabo.cpp:15
Nabo::NNSearchD NNSNabo
Definition: python/nabo.cpp:10
kd-tree with linear heap, good for small k (~up to 30)
Definition: nabo.h:293
brute force, check distance to every point in the data
Definition: nabo.h:292


libnabo
Author(s): Stéphane Magnenat
autogenerated on Mon Feb 28 2022 22:41:38