source/blender/python/mathutils/mathutils_kdtree.c

   1 /*
   2  * ***** BEGIN GPL LICENSE BLOCK *****
   3  *
   4  * This program is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU General Public License
   6  * as published by the Free Software Foundation; either version 2
   7  * of the License, or (at your option) any later version.
   8  *
   9  * This program is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12  * GNU General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU General Public License
  15  * along with this program; if not, write to the Free Software Foundation,
  16  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  17  *
  18  * Contributor(s): Dan Eicher, Campbell Barton
  19  *
  20  * ***** END GPL LICENSE BLOCK *****
  21  */
  22
  23 /** \file blender/python/mathutils/mathutils_kdtree.c
  24  *  \ingroup mathutils
  25  *
  26  * This file defines the 'mathutils.kdtree' module, a general purpose module to access
  27  * blenders kdtree for 3d spatial lookups.
  28  */
  29
  30 #include <Python.h>
  31
  32 #include "MEM_guardedalloc.h"
  33
  34 #include "BLI_utildefines.h"
  35 #include "BLI_kdtree.h"
  36
  37 #include "../generic/py_capi_utils.h"
  38
  39 #include "mathutils.h"
  40 #include "mathutils_kdtree.h"  /* own include */
  41
  42 #include "BLI_strict_flags.h"
  43
  44 typedef struct {
  45         PyObject_HEAD
  46         KDTree *obj;
  47         unsigned int maxsize;
  48         unsigned int count;
  49         unsigned int count_balance;  /* size when we last balanced */
  50 } PyKDTree;
  51
  52
  53 /* -------------------------------------------------------------------- */
  54 /* Utility helper functions */
  55
  56 static void kdtree_nearest_to_py_tuple(const KDTreeNearest *nearest, PyObject *py_retval)
  57 {
  58         BLI_assert(nearest->index >= 0);
  59         BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);
  60
  61         PyTuple_SET_ITEM(py_retval, 0, Vector_CreatePyObject((float *)nearest->co, 3, Py_NEW, NULL));
  62         PyTuple_SET_ITEM(py_retval, 1, PyLong_FromLong(nearest->index));
  63         PyTuple_SET_ITEM(py_retval, 2, PyFloat_FromDouble(nearest->dist));
  64 }
  65
  66 static PyObject *kdtree_nearest_to_py(const KDTreeNearest *nearest)
  67 {
  68         PyObject *py_retval;
  69
  70         py_retval = PyTuple_New(3);
  71
  72         kdtree_nearest_to_py_tuple(nearest, py_retval);
  73
  74         return py_retval;
  75 }
  76
  77 static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest *nearest)
  78 {
  79         PyObject *py_retval;
  80
  81         py_retval = PyTuple_New(3);
  82
  83         if (nearest->index != -1) {
  84                 kdtree_nearest_to_py_tuple(nearest, py_retval);
  85         }
  86         else {
  87                 PyC_Tuple_Fill(py_retval, Py_None);
  88         }
  89
  90         return py_retval;
  91 }
  92
  93
  94 /* -------------------------------------------------------------------- */
  95 /* KDTree */
  96
  97 /* annoying since arg parsing won't check overflow */
  98 #define UINT_IS_NEG(n) ((n) > INT_MAX)
  99
 100 static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
 101 {
 102         unsigned int maxsize;
 103         const char *keywords[] = {"size", NULL};
 104
 105         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *)"I:KDTree", (char **)keywords, &maxsize)) {
 106                 return -1;
 107         }
 108
 109         if (UINT_IS_NEG(maxsize)) {
 110                 PyErr_SetString(PyExc_ValueError, "negative 'size' given");
 111                 return -1;
 112         }
 113
 114         self->obj = BLI_kdtree_new(maxsize);
 115         self->maxsize = maxsize;
 116         self->count = 0;
 117         self->count_balance = 0;
 118
 119         return 0;
 120 }
 121
 122 static void PyKDTree__tp_dealloc(PyKDTree *self)
 123 {
 124         BLI_kdtree_free(self->obj);
 125         Py_TYPE(self)->tp_free((PyObject *)self);
 126 }
 127
 128 PyDoc_STRVAR(py_kdtree_insert_doc,
 129 ".. method:: insert(index, co)\n"
 130 "\n"
 131 "   Insert a point into the KDTree.\n"
 132 "\n"
 133 "   :arg co: Point 3d position.\n"
 134 "   :type co: float triplet\n"
 135 "   :arg index: The index of the point.\n"
 136 "   :type index: int\n"
 137 );
 138 static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
 139 {
 140         PyObject *py_co;
 141         float co[3];
 142         int index;
 143         const char *keywords[] = {"co", "index", NULL};
 144
 145         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Oi:insert", (char **)keywords,
 146                                          &py_co, &index))
 147         {
 148                 return NULL;
 149         }
 150
 151         if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1)
 152                 return NULL;
 153
 154         if (index < 0) {
 155                 PyErr_SetString(PyExc_ValueError, "negative index given");
 156                 return NULL;
 157         }
 158
 159         if (self->count >= self->maxsize) {
 160                 PyErr_SetString(PyExc_RuntimeError, "Trying to insert more items than KDTree has room for");
 161                 return NULL;
 162         }
 163
 164         BLI_kdtree_insert(self->obj, index, co, NULL);
 165         self->count++;
 166
 167         Py_RETURN_NONE;
 168 }
 169
 170 PyDoc_STRVAR(py_kdtree_balance_doc,
 171 ".. method:: balance()\n"
 172 "\n"
 173 "   Balance the tree.\n"
 174 );
 175 static PyObject *py_kdtree_balance(PyKDTree *self)
 176 {
 177         BLI_kdtree_balance(self->obj);
 178         self->count_balance = self->count;
 179         Py_RETURN_NONE;
 180 }
 181
 182 PyDoc_STRVAR(py_kdtree_find_doc,
 183 ".. method:: find(co)\n"
 184 "\n"
 185 "   Find nearest point to ``co``.\n"
 186 "\n"
 187 "   :arg co: 3d coordinates.\n"
 188 "   :type co: float triplet\n"
 189 "   :return: Returns (:class:`Vector`, index, distance).\n"
 190 "   :rtype: :class:`tuple`\n"
 191 );
 192 static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
 193 {
 194         PyObject *py_co;
 195         float co[3];
 196         KDTreeNearest nearest;
 197         const char *keywords[] = {"co", NULL};
 198
 199         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "O:find", (char **)keywords,
 200                                          &py_co))
 201         {
 202                 return NULL;
 203         }
 204
 205         if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1)
 206                 return NULL;
 207
 208         if (self->count != self->count_balance) {
 209                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find()");
 210                 return NULL;
 211         }
 212
 213
 214         nearest.index = -1;
 215
 216         BLI_kdtree_find_nearest(self->obj, co, NULL, &nearest);
 217
 218         return kdtree_nearest_to_py_and_check(&nearest);
 219 }
 220
 221 PyDoc_STRVAR(py_kdtree_find_n_doc,
 222 ".. method:: find_n(co, n)\n"
 223 "\n"
 224 "   Find nearest ``n`` points to ``co``.\n"
 225 "\n"
 226 "   :arg co: 3d coordinates.\n"
 227 "   :type co: float triplet\n"
 228 "   :arg n: Number of points to find.\n"
 229 "   :type n: int\n"
 230 "   :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
 231 "   :rtype: :class:`list`\n"
 232 );
 233 static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
 234 {
 235         PyObject *py_list;
 236         PyObject *py_co;
 237         float co[3];
 238         KDTreeNearest *nearest;
 239         unsigned int n;
 240         int i, found;
 241         const char *keywords[] = {"co", "n", NULL};
 242
 243         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "OI:find_n", (char **)keywords,
 244                                          &py_co, &n))
 245         {
 246                 return NULL;
 247         }
 248
 249         if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1)
 250                 return NULL;
 251
 252         if (UINT_IS_NEG(n)) {
 253                 PyErr_SetString(PyExc_RuntimeError, "negative 'n' given");
 254                 return NULL;
 255         }
 256
 257         if (self->count != self->count_balance) {
 258                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_n()");
 259                 return NULL;
 260         }
 261
 262         nearest = MEM_mallocN(sizeof(KDTreeNearest) * n, __func__);
 263
 264         found = BLI_kdtree_find_nearest_n(self->obj, co, NULL, nearest, n);
 265
 266         py_list = PyList_New(found);
 267
 268         for (i = 0; i < found; i++) {
 269                 PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
 270         }
 271
 272         MEM_freeN(nearest);
 273
 274         return py_list;
 275 }
 276
 277 PyDoc_STRVAR(py_kdtree_find_range_doc,
 278 ".. method:: find_range(co, radius)\n"
 279 "\n"
 280 "   Find all points within ``radius`` of ``co``.\n"
 281 "\n"
 282 "   :arg co: 3d coordinates.\n"
 283 "   :type co: float triplet\n"
 284 "   :arg radius: Distance to search for points.\n"
 285 "   :type radius: float\n"
 286 "   :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
 287 "   :rtype: :class:`list`\n"
 288 );
 289 static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
 290 {
 291         PyObject *py_list;
 292         PyObject *py_co;
 293         float co[3];
 294         KDTreeNearest *nearest = NULL;
 295         float radius;
 296         int i, found;
 297
 298         const char *keywords[] = {"co", "radius", NULL};
 299
 300         if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Of:find_range", (char **)keywords,
 301                                          &py_co, &radius))
 302         {
 303                 return NULL;
 304         }
 305
 306         if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1)
 307                 return NULL;
 308
 309         if (radius < 0.0f) {
 310                 PyErr_SetString(PyExc_RuntimeError, "negative radius given");
 311                 return NULL;
 312         }
 313
 314         if (self->count != self->count_balance) {
 315                 PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_range()");
 316                 return NULL;
 317         }
 318
 319         found = BLI_kdtree_range_search(self->obj, co, NULL, &nearest, radius);
 320
 321         py_list = PyList_New(found);
 322
 323         for (i = 0; i < found; i++) {
 324                 PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
 325         }
 326
 327         if (nearest) {
 328                 MEM_freeN(nearest);
 329         }
 330
 331         return py_list;
 332 }
 333
 334
 335 static PyMethodDef PyKDTree_methods[] = {
 336         {"insert", (PyCFunction)py_kdtree_insert, METH_VARARGS | METH_KEYWORDS, py_kdtree_insert_doc},
 337         {"balance", (PyCFunction)py_kdtree_balance, METH_NOARGS, py_kdtree_balance_doc},
 338         {"find", (PyCFunction)py_kdtree_find, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_doc},
 339         {"find_n", (PyCFunction)py_kdtree_find_n, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_n_doc},
 340         {"find_range", (PyCFunction)py_kdtree_find_range, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_range_doc},
 341         {NULL, NULL, 0, NULL}
 342 };
 343
 344 PyDoc_STRVAR(py_KDtree_doc,
 345 "KdTree(size) -> new kd-tree initialized to hold ``size`` items.\n"
 346 "\n"
 347 ".. note::\n"
 348 "\n"
 349 "   :class:`KDTree.balance` must have been called before using any of the ``find`` methods.\n"
 350 );
 351 PyTypeObject PyKDTree_Type = {
 352         PyVarObject_HEAD_INIT(NULL, 0)
 353         "KDTree",                                    /* tp_name */
 354         sizeof(PyKDTree),                            /* tp_basicsize */
 355         0,                                           /* tp_itemsize */
 356         /* methods */
 357         (destructor)PyKDTree__tp_dealloc,            /* tp_dealloc */
 358         NULL,                                        /* tp_print */
 359         NULL,                                        /* tp_getattr */
 360         NULL,                                        /* tp_setattr */
 361         NULL,                                        /* tp_compare */
 362         NULL,                                        /* tp_repr */
 363         NULL,                                        /* tp_as_number */
 364         NULL,                                        /* tp_as_sequence */
 365         NULL,                                        /* tp_as_mapping */
 366         NULL,                                        /* tp_hash */
 367         NULL,                                        /* tp_call */
 368         NULL,                                        /* tp_str */
 369         NULL,                                        /* tp_getattro */
 370         NULL,                                        /* tp_setattro */
 371         NULL,                                        /* tp_as_buffer */
 372         Py_TPFLAGS_DEFAULT,                          /* tp_flags */
 373         py_KDtree_doc,                               /* Documentation string */
 374         NULL,                                        /* tp_traverse */
 375         NULL,                                        /* tp_clear */
 376         NULL,                                        /* tp_richcompare */
 377         0,                                           /* tp_weaklistoffset */
 378         NULL,                                        /* tp_iter */
 379         NULL,                                        /* tp_iternext */
 380         (struct PyMethodDef *)PyKDTree_methods,      /* tp_methods */
 381         NULL,                                        /* tp_members */
 382         NULL,                                        /* tp_getset */
 383         NULL,                                        /* tp_base */
 384         NULL,                                        /* tp_dict */
 385         NULL,                                        /* tp_descr_get */
 386         NULL,                                        /* tp_descr_set */
 387         0,                                           /* tp_dictoffset */
 388         (initproc)PyKDTree__tp_init,                 /* tp_init */
 389         (allocfunc)PyType_GenericAlloc,              /* tp_alloc */
 390         (newfunc)PyType_GenericNew,                  /* tp_new */
 391         (freefunc)0,                                 /* tp_free */
 392         NULL,                                        /* tp_is_gc */
 393         NULL,                                        /* tp_bases */
 394         NULL,                                        /* tp_mro */
 395         NULL,                                        /* tp_cache */
 396         NULL,                                        /* tp_subclasses */
 397         NULL,                                        /* tp_weaklist */
 398         (destructor) NULL                            /* tp_del */
 399 };
 400
 401 PyDoc_STRVAR(py_kdtree_doc,
 402 "Generic 3-dimentional kd-tree to perform spatial searches."
 403 );
 404 static struct PyModuleDef kdtree_moduledef = {
 405         PyModuleDef_HEAD_INIT,
 406         "mathutils.kdtree",                          /* m_name */
 407         py_kdtree_doc,                               /* m_doc */
 408         0,                                           /* m_size */
 409         NULL,                                        /* m_methods */
 410         NULL,                                        /* m_reload */
 411         NULL,                                        /* m_traverse */
 412         NULL,                                        /* m_clear */
 413         NULL                                         /* m_free */
 414 };
 415
 416 PyMODINIT_FUNC PyInit_mathutils_kdtree(void)
 417 {
 418         PyObject *m = PyModule_Create(&kdtree_moduledef);
 419
 420         if (m == NULL) {
 421                 return NULL;
 422         }
 423
 424         /* Register the 'KDTree' class */
 425         if (PyType_Ready(&PyKDTree_Type)) {
 426                 return NULL;
 427         }
 428         PyModule_AddObject(m, (char *)"KDTree", (PyObject *) &PyKDTree_Type);
 429
 430         return m;
 431 }