Blender  V3.3
mathutils_kdtree.c
Go to the documentation of this file.
1 /* SPDX-License-Identifier: GPL-2.0-or-later */
2 
10 #include <Python.h>
11 
12 #include "MEM_guardedalloc.h"
13 
14 #include "BLI_kdtree.h"
15 #include "BLI_utildefines.h"
16 
17 #include "../generic/py_capi_utils.h"
18 #include "../generic/python_utildefines.h"
19 
20 #include "mathutils.h"
21 #include "mathutils_kdtree.h" /* own include */
22 
23 #include "BLI_strict_flags.h"
24 
25 typedef struct {
26  PyObject_HEAD
27  KDTree_3d *obj;
30  uint count_balance; /* size when we last balanced */
31 } PyKDTree;
32 
33 /* -------------------------------------------------------------------- */
34 /* Utility helper functions */
35 
36 static void kdtree_nearest_to_py_tuple(const KDTreeNearest_3d *nearest, PyObject *py_retval)
37 {
38  BLI_assert(nearest->index >= 0);
39  BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);
40 
41  PyTuple_SET_ITEMS(py_retval,
42  Vector_CreatePyObject(nearest->co, 3, NULL),
43  PyLong_FromLong(nearest->index),
44  PyFloat_FromDouble(nearest->dist));
45 }
46 
47 static PyObject *kdtree_nearest_to_py(const KDTreeNearest_3d *nearest)
48 {
49  PyObject *py_retval;
50 
51  py_retval = PyTuple_New(3);
52 
53  kdtree_nearest_to_py_tuple(nearest, py_retval);
54 
55  return py_retval;
56 }
57 
58 static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest_3d *nearest)
59 {
60  PyObject *py_retval;
61 
62  py_retval = PyTuple_New(3);
63 
64  if (nearest->index != -1) {
65  kdtree_nearest_to_py_tuple(nearest, py_retval);
66  }
67  else {
68  PyC_Tuple_Fill(py_retval, Py_None);
69  }
70 
71  return py_retval;
72 }
73 
74 /* -------------------------------------------------------------------- */
75 /* KDTree */
76 
77 /* annoying since arg parsing won't check overflow */
78 #define UINT_IS_NEG(n) ((n) > INT_MAX)
79 
80 static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
81 {
82  uint maxsize;
83  const char *keywords[] = {"size", NULL};
84 
85  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "I:KDTree", (char **)keywords, &maxsize)) {
86  return -1;
87  }
88 
89  if (UINT_IS_NEG(maxsize)) {
90  PyErr_SetString(PyExc_ValueError, "negative 'size' given");
91  return -1;
92  }
93 
94  self->obj = BLI_kdtree_3d_new(maxsize);
95  self->maxsize = maxsize;
96  self->count = 0;
97  self->count_balance = 0;
98 
99  return 0;
100 }
101 
102 static void PyKDTree__tp_dealloc(PyKDTree *self)
103 {
104  BLI_kdtree_3d_free(self->obj);
105  Py_TYPE(self)->tp_free((PyObject *)self);
106 }
107 
108 PyDoc_STRVAR(py_kdtree_insert_doc,
109  ".. method:: insert(co, index)\n"
110  "\n"
111  " Insert a point into the KDTree.\n"
112  "\n"
113  " :arg co: Point 3d position.\n"
114  " :type co: float triplet\n"
115  " :arg index: The index of the point.\n"
116  " :type index: int\n");
117 static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
118 {
119  PyObject *py_co;
120  float co[3];
121  int index;
122  const char *keywords[] = {"co", "index", NULL};
123 
124  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Oi:insert", (char **)keywords, &py_co, &index)) {
125  return NULL;
126  }
127 
128  if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1) {
129  return NULL;
130  }
131 
132  if (index < 0) {
133  PyErr_SetString(PyExc_ValueError, "negative index given");
134  return NULL;
135  }
136 
137  if (self->count >= self->maxsize) {
138  PyErr_SetString(PyExc_RuntimeError, "Trying to insert more items than KDTree has room for");
139  return NULL;
140  }
141 
142  BLI_kdtree_3d_insert(self->obj, index, co);
143  self->count++;
144 
145  Py_RETURN_NONE;
146 }
147 
148 PyDoc_STRVAR(py_kdtree_balance_doc,
149  ".. method:: balance()\n"
150  "\n"
151  " Balance the tree.\n"
152  "\n"
153  ".. note::\n"
154  "\n"
155  " This builds the entire tree, avoid calling after each insertion.\n");
156 static PyObject *py_kdtree_balance(PyKDTree *self)
157 {
158  BLI_kdtree_3d_balance(self->obj);
159  self->count_balance = self->count;
160  Py_RETURN_NONE;
161 }
162 
164  PyObject *py_filter;
165  bool is_error;
166 };
167 
168 static int py_find_nearest_cb(void *user_data, int index, const float co[3], float dist_sq)
169 {
170  UNUSED_VARS(co, dist_sq);
171 
173 
174  PyObject *py_args = PyTuple_New(1);
175  PyTuple_SET_ITEM(py_args, 0, PyLong_FromLong(index));
176  PyObject *result = PyObject_CallObject(data->py_filter, py_args);
177  Py_DECREF(py_args);
178 
179  if (result) {
180  bool use_node;
181  const int ok = PyC_ParseBool(result, &use_node);
182  Py_DECREF(result);
183  if (ok) {
184  return (int)use_node;
185  }
186  }
187 
188  data->is_error = true;
189  return -1;
190 }
191 
192 PyDoc_STRVAR(py_kdtree_find_doc,
193  ".. method:: find(co, filter=None)\n"
194  "\n"
195  " Find nearest point to ``co``.\n"
196  "\n"
197  " :arg co: 3d coordinates.\n"
198  " :type co: float triplet\n"
199  " :arg filter: function which takes an index and returns True for indices to "
200  "include in the search.\n"
201  " :type filter: callable\n"
202  " :return: Returns (:class:`Vector`, index, distance).\n"
203  " :rtype: :class:`tuple`\n");
204 static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
205 {
206  PyObject *py_co, *py_filter = NULL;
207  float co[3];
208  KDTreeNearest_3d nearest;
209  const char *keywords[] = {"co", "filter", NULL};
210 
211  if (!PyArg_ParseTupleAndKeywords(
212  args, kwargs, "O|$O:find", (char **)keywords, &py_co, &py_filter)) {
213  return NULL;
214  }
215 
216  if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1) {
217  return NULL;
218  }
219 
220  if (self->count != self->count_balance) {
221  PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find()");
222  return NULL;
223  }
224 
225  nearest.index = -1;
226 
227  if (py_filter == NULL) {
228  BLI_kdtree_3d_find_nearest(self->obj, co, &nearest);
229  }
230  else {
231  struct PyKDTree_NearestData data = {0};
232 
233  data.py_filter = py_filter;
234  data.is_error = false;
235 
236  BLI_kdtree_3d_find_nearest_cb(self->obj, co, py_find_nearest_cb, &data, &nearest);
237 
238  if (data.is_error) {
239  return NULL;
240  }
241  }
242 
243  return kdtree_nearest_to_py_and_check(&nearest);
244 }
245 
246 PyDoc_STRVAR(py_kdtree_find_n_doc,
247  ".. method:: find_n(co, n)\n"
248  "\n"
249  " Find nearest ``n`` points to ``co``.\n"
250  "\n"
251  " :arg co: 3d coordinates.\n"
252  " :type co: float triplet\n"
253  " :arg n: Number of points to find.\n"
254  " :type n: int\n"
255  " :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
256  " :rtype: :class:`list`\n");
257 static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
258 {
259  PyObject *py_list;
260  PyObject *py_co;
261  float co[3];
262  KDTreeNearest_3d *nearest;
263  uint n;
264  int i, found;
265  const char *keywords[] = {"co", "n", NULL};
266 
267  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OI:find_n", (char **)keywords, &py_co, &n)) {
268  return NULL;
269  }
270 
271  if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1) {
272  return NULL;
273  }
274 
275  if (UINT_IS_NEG(n)) {
276  PyErr_SetString(PyExc_RuntimeError, "negative 'n' given");
277  return NULL;
278  }
279 
280  if (self->count != self->count_balance) {
281  PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_n()");
282  return NULL;
283  }
284 
285  nearest = MEM_mallocN(sizeof(KDTreeNearest_3d) * n, __func__);
286 
287  found = BLI_kdtree_3d_find_nearest_n(self->obj, co, nearest, n);
288 
289  py_list = PyList_New(found);
290 
291  for (i = 0; i < found; i++) {
292  PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
293  }
294 
295  MEM_freeN(nearest);
296 
297  return py_list;
298 }
299 
300 PyDoc_STRVAR(py_kdtree_find_range_doc,
301  ".. method:: find_range(co, radius)\n"
302  "\n"
303  " Find all points within ``radius`` of ``co``.\n"
304  "\n"
305  " :arg co: 3d coordinates.\n"
306  " :type co: float triplet\n"
307  " :arg radius: Distance to search for points.\n"
308  " :type radius: float\n"
309  " :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
310  " :rtype: :class:`list`\n");
311 static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
312 {
313  PyObject *py_list;
314  PyObject *py_co;
315  float co[3];
316  KDTreeNearest_3d *nearest = NULL;
317  float radius;
318  int i, found;
319 
320  const char *keywords[] = {"co", "radius", NULL};
321 
322  if (!PyArg_ParseTupleAndKeywords(
323  args, kwargs, "Of:find_range", (char **)keywords, &py_co, &radius)) {
324  return NULL;
325  }
326 
327  if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1) {
328  return NULL;
329  }
330 
331  if (radius < 0.0f) {
332  PyErr_SetString(PyExc_RuntimeError, "negative radius given");
333  return NULL;
334  }
335 
336  if (self->count != self->count_balance) {
337  PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_range()");
338  return NULL;
339  }
340 
341  found = BLI_kdtree_3d_range_search(self->obj, co, &nearest, radius);
342 
343  py_list = PyList_New(found);
344 
345  for (i = 0; i < found; i++) {
346  PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
347  }
348 
349  if (nearest) {
350  MEM_freeN(nearest);
351  }
352 
353  return py_list;
354 }
355 
356 static PyMethodDef PyKDTree_methods[] = {
357  {"insert", (PyCFunction)py_kdtree_insert, METH_VARARGS | METH_KEYWORDS, py_kdtree_insert_doc},
358  {"balance", (PyCFunction)py_kdtree_balance, METH_NOARGS, py_kdtree_balance_doc},
359  {"find", (PyCFunction)py_kdtree_find, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_doc},
360  {"find_n", (PyCFunction)py_kdtree_find_n, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_n_doc},
361  {"find_range",
362  (PyCFunction)py_kdtree_find_range,
363  METH_VARARGS | METH_KEYWORDS,
364  py_kdtree_find_range_doc},
365  {NULL, NULL, 0, NULL},
366 };
367 
368 PyDoc_STRVAR(py_KDtree_doc,
369  "KdTree(size) -> new kd-tree initialized to hold ``size`` items.\n"
370  "\n"
371  ".. note::\n"
372  "\n"
373  " :class:`KDTree.balance` must have been called before using any of the ``find`` "
374  "methods.\n");
375 PyTypeObject PyKDTree_Type = {
376  PyVarObject_HEAD_INIT(NULL, 0) "KDTree", /* tp_name */
377  sizeof(PyKDTree), /* tp_basicsize */
378  0, /* tp_itemsize */
379  /* methods */
380  (destructor)PyKDTree__tp_dealloc, /* tp_dealloc */
381  (printfunc)NULL, /* tp_print */
382  NULL, /* tp_getattr */
383  NULL, /* tp_setattr */
384  NULL, /* tp_compare */
385  NULL, /* tp_repr */
386  NULL, /* tp_as_number */
387  NULL, /* tp_as_sequence */
388  NULL, /* tp_as_mapping */
389  NULL, /* tp_hash */
390  NULL, /* tp_call */
391  NULL, /* tp_str */
392  NULL, /* tp_getattro */
393  NULL, /* tp_setattro */
394  NULL, /* tp_as_buffer */
395  Py_TPFLAGS_DEFAULT, /* tp_flags */
396  py_KDtree_doc, /* Documentation string */
397  NULL, /* tp_traverse */
398  NULL, /* tp_clear */
399  NULL, /* tp_richcompare */
400  0, /* tp_weaklistoffset */
401  NULL, /* tp_iter */
402  NULL, /* tp_iternext */
403  (struct PyMethodDef *)PyKDTree_methods, /* tp_methods */
404  NULL, /* tp_members */
405  NULL, /* tp_getset */
406  NULL, /* tp_base */
407  NULL, /* tp_dict */
408  NULL, /* tp_descr_get */
409  NULL, /* tp_descr_set */
410  0, /* tp_dictoffset */
411  (initproc)PyKDTree__tp_init, /* tp_init */
412  (allocfunc)PyType_GenericAlloc, /* tp_alloc */
413  (newfunc)PyType_GenericNew, /* tp_new */
414  (freefunc)0, /* tp_free */
415  NULL, /* tp_is_gc */
416  NULL, /* tp_bases */
417  NULL, /* tp_mro */
418  NULL, /* tp_cache */
419  NULL, /* tp_subclasses */
420  NULL, /* tp_weaklist */
421  (destructor)NULL, /* tp_del */
422 };
423 
424 PyDoc_STRVAR(py_kdtree_doc, "Generic 3-dimensional kd-tree to perform spatial searches.");
425 static struct PyModuleDef kdtree_moduledef = {
426  PyModuleDef_HEAD_INIT,
427  "mathutils.kdtree", /* m_name */
428  py_kdtree_doc, /* m_doc */
429  0, /* m_size */
430  NULL, /* m_methods */
431  NULL, /* m_reload */
432  NULL, /* m_traverse */
433  NULL, /* m_clear */
434  NULL, /* m_free */
435 };
436 
437 PyMODINIT_FUNC PyInit_mathutils_kdtree(void)
438 {
439  PyObject *m = PyModule_Create(&kdtree_moduledef);
440 
441  if (m == NULL) {
442  return NULL;
443  }
444 
445  /* Register the 'KDTree' class */
446  if (PyType_Ready(&PyKDTree_Type)) {
447  return NULL;
448  }
449  PyModule_AddType(m, &PyKDTree_Type);
450 
451  return m;
452 }
#define BLI_assert(a)
Definition: BLI_assert.h:46
A KD-tree for nearest neighbor search.
Strict compiler flags for areas of code we want to ensure don't do conversions without us knowing abo...
unsigned int uint
Definition: BLI_sys_types.h:67
#define UNUSED_VARS(...)
Read Guarded memory(de)allocation.
PyObject * self
Definition: bpy_driver.c:165
void * user_data
void(* MEM_freeN)(void *vmemh)
Definition: mallocn.c:27
void *(* MEM_mallocN)(size_t len, const char *str)
Definition: mallocn.c:33
int mathutils_array_parse(float *array, int array_num_min, int array_num_max, PyObject *value, const char *error_prefix)
Definition: mathutils.c:98
PyObject * Vector_CreatePyObject(const float *vec, const int vec_num, PyTypeObject *base_type)
static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
static void kdtree_nearest_to_py_tuple(const KDTreeNearest_3d *nearest, PyObject *py_retval)
#define UINT_IS_NEG(n)
static PyObject * py_kdtree_balance(PyKDTree *self)
PyTypeObject PyKDTree_Type
PyMODINIT_FUNC PyInit_mathutils_kdtree(void)
static PyObject * py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
static PyObject * kdtree_nearest_to_py_and_check(const KDTreeNearest_3d *nearest)
static struct PyModuleDef kdtree_moduledef
static void PyKDTree__tp_dealloc(PyKDTree *self)
static PyMethodDef PyKDTree_methods[]
static int py_find_nearest_cb(void *user_data, int index, const float co[3], float dist_sq)
static PyObject * py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
static PyObject * py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
PyDoc_STRVAR(py_kdtree_insert_doc, ".. method:: insert(co, index)\n" "\n" " Insert a point into the KDTree.\n" "\n" " :arg co: Point 3d position.\n" " :type co: float triplet\n" " :arg index: The index of the point.\n" " :type index: int\n")
static PyObject * kdtree_nearest_to_py(const KDTreeNearest_3d *nearest)
static PyObject * py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
int PyC_ParseBool(PyObject *o, void *p)
void PyC_Tuple_Fill(PyObject *tuple, PyObject *value)
#define PyTuple_SET_ITEMS(op_arg,...)
uint count_balance
PyObject_HEAD KDTree_3d * obj