00001 #include <Python.h>
00002 #include <boost/python.hpp>
00003 #include <numpy/arrayobject.h>
00004 #include "../nabo/nabo.h"
00005 #include <iostream>
00006 #include <cassert>
00007
00008 using namespace boost::python;
00009
00010 typedef Nabo::NNSearchD NNSNabo;
00011 typedef NNSNabo::Index Index;
00012 typedef NNSNabo::SearchType SearchType;
00013 typedef NNSNabo::SearchOptionFlags SearchOptionFlags;
00014 typedef Eigen::Map<NNSNabo::Matrix> MappedEigenDoubleMatrix;
00015 typedef Eigen::Map<NNSNabo::IndexMatrix> MappedEigenIndexMatrix;
00016
00017 static const double infD = std::numeric_limits<double>::infinity();
00018 static const Index maxI = std::numeric_limits<Index>::max();
00019
00020 void matrixSizeFromPythonArray(const PyObject* cloudObj, int& rowCount, int& colCount)
00021 {
00022 assert(PyArray_CHKFLAGS(cloudObj, NPY_C_CONTIGUOUS) || PyArray_CHKFLAGS(cloudObj, NPY_F_CONTIGUOUS));
00023 assert(PyArray_NDIM(cloudObj) == 2);
00024 const npy_intp *shape = PyArray_DIMS(cloudObj);
00025 if (PyArray_CHKFLAGS(cloudObj, NPY_F_CONTIGUOUS))
00026 {
00027 colCount = shape[1];
00028 rowCount = shape[0];
00029 }
00030 else
00031 {
00032 colCount = shape[0];
00033 rowCount = shape[1];
00034 }
00035 }
00036
00037 void checkPythonArray(const PyObject* cloudObj, const char *paramName)
00038 {
00039 std::string startMsg("Argument \"");
00040 startMsg += paramName;
00041 startMsg += "\" ";
00042
00043 if (!PyArray_Check(cloudObj))
00044 throw std::runtime_error(startMsg + "must be a multi-dimensional array");
00045 const int nDim = PyArray_NDIM(cloudObj);
00046 if (nDim != 2)
00047 throw std::runtime_error(startMsg + "must be a two-dimensional array");
00048 if (PyArray_TYPE(cloudObj) != NPY_FLOAT64)
00049 throw std::runtime_error(startMsg + "must hold doubles");
00050 if (!PyArray_CHKFLAGS(cloudObj, NPY_C_CONTIGUOUS) && !PyArray_CHKFLAGS(cloudObj, NPY_F_CONTIGUOUS))
00051 throw std::runtime_error(startMsg + "must be a continuous array");
00052 }
00053
00054 MappedEigenDoubleMatrix* eigenFromBoostPython(const object cloudIn, const char *paramName)
00055 {
00056 int dimCount, pointCount;
00057 const PyObject *cloudObj(cloudIn.ptr());
00058
00059 checkPythonArray(cloudObj, paramName);
00060
00061 matrixSizeFromPythonArray(cloudObj, dimCount, pointCount);
00062 double* cloudData(reinterpret_cast<double*>(PyArray_DATA(cloudObj)));
00063
00064 return new MappedEigenDoubleMatrix(cloudData, dimCount, pointCount);
00065 }
00066
00067 void eigenFromBoostPython(NNSNabo::Matrix& cloudOut, const object cloudIn, const char *paramName)
00068 {
00069 int dimCount, pointCount;
00070 const PyObject *cloudObj(cloudIn.ptr());
00071
00072 checkPythonArray(cloudObj, paramName);
00073
00074 matrixSizeFromPythonArray(cloudObj, dimCount, pointCount);
00075 cloudOut.resize(dimCount, pointCount);
00076
00077 memcpy(cloudOut.data(), PyArray_DATA(cloudObj), pointCount*dimCount*sizeof(double));
00078 }
00079
00080 class NearestNeighbourSearch
00081 {
00082 public:
00083 NearestNeighbourSearch(const object pycloud, const SearchType searchType = NNSNabo::KDTREE_LINEAR_HEAP, const Index dim = maxI, const dict params = dict())
00084 {
00085
00086 eigenFromBoostPython(cloud, pycloud, "cloud");
00087
00088
00089 Nabo::Parameters _params;
00090 object it = params.iteritems();
00091 for(int i = 0; i < len(params); ++i)
00092 {
00093 const tuple item(it.attr("next")());
00094 const std::string key = extract<std::string>(item[0]);
00095 const object val(item[1]);
00096 const std::string valType(val.ptr()->ob_type->tp_name);
00097 if (valType == "int")
00098 {
00099 const int iVal = extract<int>(val);
00100 if (iVal >= 0)
00101 _params[key] = (unsigned)iVal;
00102 else
00103 _params[key] = iVal;
00104 }
00105 }
00106
00107
00108 nns = NNSNabo ::create(cloud, dim, searchType, 0, _params);
00109 }
00110
00111 ~NearestNeighbourSearch()
00112 {
00113 delete nns;
00114 }
00115
00116 tuple knn(const object query, const Index k = 1, const double epsilon = 0, const unsigned optionFlags = 0, const double maxRadius = infD)
00117 {
00118
00119 MappedEigenDoubleMatrix* mappedQuery(eigenFromBoostPython(query, "query"));
00120 NNSNabo::IndexMatrix indexMatrix(k, mappedQuery->cols());
00121 NNSNabo::Matrix dists2Matrix(k, mappedQuery->cols());
00122
00123
00124 nns->knn(*mappedQuery, indexMatrix, dists2Matrix, k, epsilon, optionFlags, maxRadius);
00125
00126
00127 npy_intp retDims[2] = { mappedQuery->cols(), k };
00128 const int dataCount(k * mappedQuery->cols());
00129 PyObject* dists2 = PyArray_EMPTY(2, retDims, PyArray_DOUBLE, PyArray_ISFORTRAN(query.ptr()));
00130 memcpy(PyArray_DATA(dists2), dists2Matrix.data(), dataCount*sizeof(double));
00131 PyObject* indices = PyArray_EMPTY(2, retDims, PyArray_INT, PyArray_ISFORTRAN(query.ptr()));
00132 memcpy(PyArray_DATA(indices), indexMatrix.data(), dataCount*sizeof(int));
00133 delete mappedQuery;
00134
00135
00136 return make_tuple(object(handle<>(indices)), object(handle<>(dists2)));
00137 }
00138
00139 protected:
00140 NNSNabo *nns;
00141 NNSNabo::Matrix cloud;
00142 };
00143
00144 BOOST_PYTHON_MEMBER_FUNCTION_OVERLOADS(knn_overloads, knn, 1, 5)
00145
00146 BOOST_PYTHON_MODULE(pynabo)
00147 {
00148 import_array();
00149
00150 enum_<SearchType>("SearchType", "Type of algorithm used for search.")
00151 .value("BRUTE_FORCE", NNSNabo::BRUTE_FORCE)
00152 .value("KDTREE_LINEAR_HEAP", NNSNabo::KDTREE_LINEAR_HEAP)
00153 .value("KDTREE_TREE_HEAP", NNSNabo::KDTREE_TREE_HEAP)
00154 ;
00155
00156 enum_<SearchOptionFlags>("SearchOptionFlags", "Flags you can OR when creating search.")
00157 .value("ALLOW_SELF_MATCH", NNSNabo::ALLOW_SELF_MATCH)
00158 .value("SORT_RESULTS", NNSNabo::SORT_RESULTS)
00159 ;
00160
00161 class_<NearestNeighbourSearch>(
00162 "NearestNeighbourSearch",
00163 "Nearest-neighbour search object, containing the data, on which you can do the knn(...) query.\n\n"
00164 "The data and query must be continuous numpy arrays.\n"
00165 "As numpy proposes both C and Fortran data orders, pynabo\n"
00166 "will always consider the contiguous dimension to be coordinates\n"
00167 "of points, regardless of order, as this provides the fastest\n"
00168 "possible execution. The return values of knn(...) will have\n"
00169 "the same order as the query and will have the different results\n"
00170 "of each point in the contiguous dimension."
00171 ,
00172 init<const object, optional<const SearchType, const Index, const dict> >(
00173 "Create a nearest-neighbour search.\n\n"
00174 "Arguments:\n"
00175 " data -- data-point cloud in which to search, must be a numpy array\n"
00176 " searchType -- type of search, default: KDTREE_LINEAR_HEAP\n"
00177 " dim -- number of dimensions to consider, must be lower or equal to cloud.rows(), default: dim of data\n"
00178 " creationOptionFlags -- creation options, a bitwise OR of elements of CreationOptionFlags",
00179 args("self", "data", "searchType", "dim", "creationOptionFlags, default: 0")
00180 )
00181 )
00182 .def("knn", &NearestNeighbourSearch::knn,
00183 knn_overloads(
00184 args("self", "query", "k", "epsilon", "optionFlags", "maxRadius"),
00185 "Find the k nearest neighbours of query in data.\n\n"
00186 "Arguments:\n"
00187 " query -- query points, must be a numpy array\n"
00188 " k -- number of nearest neighbour requested, default: 1\n"
00189 " 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"
00190 " optionFlags -- search options, a bitwise OR of elements of SearchOptionFlags, default: 0\n"
00191 " maxRadius -- maximum radius in which to search, can be used to prune search, is not affected by epsilon, default: inf\n\n"
00192 "Returns:\n"
00193 " A tuple of two 2D numpy arrays, the first containing indices to points in data, the other containing squared distances."
00194 )
00195 )
00196 ;
00197 }