nabo.cpp
Go to the documentation of this file.
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                 // build cloud
00086                 eigenFromBoostPython(cloud, pycloud, "cloud");
00087                 
00088                 // build params
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                 // create search
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                 // map query and create output matrices
00119                 MappedEigenDoubleMatrix* mappedQuery(eigenFromBoostPython(query, "query"));
00120                 NNSNabo::IndexMatrix indexMatrix(k, mappedQuery->cols());
00121                 NNSNabo::Matrix dists2Matrix(k, mappedQuery->cols());
00122                 
00123                 // do the search
00124                 nns->knn(*mappedQuery, indexMatrix, dists2Matrix, k, epsilon, optionFlags, maxRadius);
00125                 
00126                 // build resulting python types
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                 // return results
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 }


libnabo
Author(s): Stéphane Magnenat
autogenerated on Thu Sep 10 2015 10:54:55