Cr8-xyz/blender
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Patch D133: Python wrapper for BLI_kdtree (adds mathutils.kdtree)

Browse Source 
Originally by Dan Eicher, with my own fixes and adjustments (see patch page for details).

For details there are unit tests and api example usage.

	doc/python_api/sphinx-in-tmp/menu_id.png
This commit is contained in:
Campbell Barton
2014-01-06 20:32:34 +11:00
parent 90efa345c2
commit 4848f9029a
9 changed files with 618 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
doc/python_api/examples/mathutils.kdtree.py Normal file
View File

@@ -0,0 +1,37 @@
import mathutils
# create a kd-tree from a mesh
from bpy import context
obj = context.object
# 3d cursor relative to the object data
co_find = context.scene.cursor_location * obj.matrix_world.inverted()
mesh = obj.data
size = len(mesh.vertices)
kd = mathutils.kdtree.KDTree(size)
for i, v in enumerate(mesh.vertices):
kd.insert(v.co, i)
kd.balance()
# Find the closest point to the center
co_find = (0.0, 0.0, 0.0)
co, index, dist = kd.find(co_find)
print("Close to center:", co, index, dist)
# Find the closest 10 points to the 3d cursor
print("Close 10 points")
for (co, index, dist) in kd.find_n(co_find, 10):
print(" ", co, index, dist)
# Find points within a radius of the 3d cursor
print("Close points within 0.5 distance")
co_find = context.scene.cursor_location
for (co, index, dist) in kd.find_range(co_find, 0.5):
print(" ", co, index, dist)

4
doc/python_api/sphinx_doc_gen.py
View File

@@ -271,6 +271,7 @@ else:
"gpu",
"mathutils",
"mathutils.geometry",
"mathutils.kdtree",
"mathutils.noise",
"freestyle",
]
@@ -1625,7 +1626,7 @@ def write_rst_contents(basepath):
standalone_modules = (
# mathutils
"mathutils", "mathutils.geometry", "mathutils.noise",
"mathutils", "mathutils.geometry", "mathutils.kdtree", "mathutils.noise",
# misc
"freestyle", "bgl", "blf", "gpu", "aud", "bpy_extras",
# bmesh, submodules are in own page
@@ -1776,6 +1777,7 @@ def write_rst_importable_modules(basepath):
"bpy.props" : "Property Definitions",
"mathutils" : "Math Types & Utilities",
"mathutils.geometry" : "Geometry Utilities",
"mathutils.kdtree" : "KDTree Utilities",
"mathutils.noise" : "Noise Utilities",
"freestyle" : "Freestyle Data Types & Operators",
}

1
source/blender/python/intern/bpy_interface.c
View File

@@ -213,6 +213,7 @@ static struct _inittab bpy_internal_modules[] = {
{(char *)"mathutils", PyInit_mathutils},
// {(char *)"mathutils.geometry", PyInit_mathutils_geometry},
// {(char *)"mathutils.noise", PyInit_mathutils_noise},
// {(char *)"mathutils.kdtree", PyInit_mathutils_kdtree},
{(char *)"_bpy_path", BPyInit__bpy_path},
{(char *)"bgl", BPyInit_bgl},
{(char *)"blf", BPyInit_blf},

2
source/blender/python/mathutils/CMakeLists.txt
View File

@@ -38,6 +38,7 @@ set(SRC
mathutils_Quaternion.c
mathutils_Vector.c
mathutils_geometry.c
mathutils_kdtree.c
mathutils_noise.c
mathutils.h
@@ -47,6 +48,7 @@ set(SRC
mathutils_Quaternion.h
mathutils_Vector.h
mathutils_geometry.h
mathutils_kdtree.h
mathutils_noise.h
)

5
source/blender/python/mathutils/mathutils.c
View File

@@ -515,6 +515,11 @@ PyMODINIT_FUNC PyInit_mathutils(void)
PyModule_AddObject(mod, "noise", (submodule = PyInit_mathutils_noise()));
PyDict_SetItemString(sys_modules, PyModule_GetName(submodule), submodule);
Py_INCREF(submodule);
/* KDTree submodule */
PyModule_AddObject(mod, "kdtree", (submodule = PyInit_mathutils_kdtree()));
PyDict_SetItemString(sys_modules, PyModule_GetName(submodule), submodule);
Py_INCREF(submodule);
#endif
mathutils_matrix_row_cb_index = Mathutils_RegisterCallback(&mathutils_matrix_row_cb);

1
source/blender/python/mathutils/mathutils.h
View File

@@ -58,6 +58,7 @@ typedef struct {
/* utility submodules */
#include "mathutils_geometry.h"
#include "mathutils_noise.h"
#include "mathutils_kdtree.h"
PyObject *BaseMathObject_owner_get(BaseMathObject *self, void *);
PyObject *BaseMathObject_is_wrapped_get(BaseMathObject *self, void *);

429
source/blender/python/mathutils/mathutils_kdtree.c Normal file
View File

@@ -0,0 +1,429 @@
/*
* ***** BEGIN GPL LICENSE BLOCK *****
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software Foundation,
* Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
*
* Contributor(s): Dan Eicher, Campbell Barton
*
* ***** END GPL LICENSE BLOCK *****
*/
/** \file blender/python/mathutils/mathutils_kdtree.c
* \ingroup mathutils
*
* This file defines the 'mathutils.kdtree' module, a general purpose module to access
* blenders kdtree for 3d spatial lookups.
*/
#include <Python.h>
#include "MEM_guardedalloc.h"
#include "BLI_utildefines.h"
#include "BLI_kdtree.h"
#include "../generic/py_capi_utils.h"
#include "mathutils.h"
#include "BLI_strict_flags.h"
typedef struct {
PyObject_HEAD
KDTree *obj;
unsigned int maxsize;
unsigned int count;
unsigned int count_balance; /* size when we last balanced */
} PyKDTree;
/* -------------------------------------------------------------------- */
/* Utility helper functions */
static void kdtree_nearest_to_py_tuple(const KDTreeNearest *nearest, PyObject *py_retval)
{
BLI_assert(nearest->index >= 0);
BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);
PyTuple_SET_ITEM(py_retval, 0, Vector_CreatePyObject((float *)nearest->co, 3, Py_NEW, NULL));
PyTuple_SET_ITEM(py_retval, 1, PyLong_FromLong(nearest->index));
PyTuple_SET_ITEM(py_retval, 2, PyFloat_FromDouble(nearest->dist));
}
static PyObject *kdtree_nearest_to_py(const KDTreeNearest *nearest)
{
PyObject *py_retval;
py_retval = PyTuple_New(3);
kdtree_nearest_to_py_tuple(nearest, py_retval);
return py_retval;
}
static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest *nearest)
{
PyObject *py_retval;
py_retval = PyTuple_New(3);
if (nearest->index != -1) {
kdtree_nearest_to_py_tuple(nearest, py_retval);
}
else {
PyC_Tuple_Fill(py_retval, Py_None);
}
return py_retval;
}
/* -------------------------------------------------------------------- */
/* KDTree */
/* annoying since arg parsing won't check overflow */
#define UINT_IS_NEG(n) ((n) > INT_MAX)
static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
unsigned int maxsize;
const char *keywords[] = {"size", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *)"I:KDTree", (char **)keywords, &maxsize)) {
return -1;
}
if (UINT_IS_NEG(maxsize)) {
PyErr_SetString(PyExc_ValueError, "negative 'size' given");
return -1;
}
self->obj = BLI_kdtree_new(maxsize);
self->maxsize = maxsize;
self->count = 0;
self->count_balance = 0;
return 0;
}
static void PyKDTree__tp_dealloc(PyKDTree *self)
{
BLI_kdtree_free(self->obj);
Py_TYPE(self)->tp_free((PyObject *)self);
}
PyDoc_STRVAR(py_kdtree_insert_doc,
".. method:: insert(index, co)\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 *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
PyObject *py_co;
float co[3];
int index;
const char *keywords[] = {"co", "index", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Oi:insert", (char **)keywords,
&py_co, &index))
{
return NULL;
}
if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1)
return NULL;
if (index < 0) {
PyErr_SetString(PyExc_ValueError, "negative index given");
return NULL;
}
if (self->count >= self->maxsize) {
PyErr_SetString(PyExc_RuntimeError, "Trying to insert more items than KDTree has room for");
return NULL;
}
BLI_kdtree_insert(self->obj, index, co, NULL);
self->count++;
Py_RETURN_NONE;
}
PyDoc_STRVAR(py_kdtree_balance_doc,
".. method:: balance()\n"
"\n"
" Balance the tree.\n"
);
static PyObject *py_kdtree_balance(PyKDTree *self)
{
BLI_kdtree_balance(self->obj);
self->count_balance = self->count;
Py_RETURN_NONE;
}
PyDoc_STRVAR(py_kdtree_find_doc,
".. method:: find(co)\n"
"\n"
" Find nearest point to ``co``.\n"
"\n"
" :arg co: 3d coordinates.\n"
" :type co: float triplet\n"
" :return: Returns (:class:`Vector`, index, distance).\n"
" :rtype: :class:`tuple`\n"
);
static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
PyObject *py_co;
float co[3];
KDTreeNearest nearest;
const char *keywords[] = {"co", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "O:find", (char **)keywords,
&py_co))
{
return NULL;
}
if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1)
return NULL;
if (self->count != self->count_balance) {
PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find()");
return NULL;
}
nearest.index = -1;
BLI_kdtree_find_nearest(self->obj, co, NULL, &nearest);
return kdtree_nearest_to_py_and_check(&nearest);
}
PyDoc_STRVAR(py_kdtree_find_n_doc,
".. method:: find_n(co, n)\n"
"\n"
" Find nearest ``n`` points to ``co``.\n"
"\n"
" :arg co: 3d coordinates.\n"
" :type co: float triplet\n"
" :arg n: Number of points to find.\n"
" :type n: int\n"
" :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
" :rtype: :class:`list`\n"
);
static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
PyObject *py_list;
PyObject *py_co;
float co[3];
KDTreeNearest *nearest;
unsigned int n;
int i, found;
const char *keywords[] = {"co", "n", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "OI:find_n", (char **)keywords,
&py_co, &n))
{
return NULL;
}
if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1)
return NULL;
if (UINT_IS_NEG(n)) {
PyErr_SetString(PyExc_RuntimeError, "negative 'n' given");
return NULL;
}
if (self->count != self->count_balance) {
PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_n()");
return NULL;
}
nearest = MEM_mallocN(sizeof(KDTreeNearest) * n, __func__);
found = BLI_kdtree_find_nearest_n(self->obj, co, NULL, nearest, n);
py_list = PyList_New(found);
for (i = 0; i < found; i++) {
PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
}
MEM_freeN(nearest);
return py_list;
}
PyDoc_STRVAR(py_kdtree_find_range_doc,
".. method:: find_range(co, radius)\n"
"\n"
" Find all points within ``radius`` of ``co``.\n"
"\n"
" :arg co: 3d coordinates.\n"
" :type co: float triplet\n"
" :arg radius: Distance to search for points.\n"
" :type radius: float\n"
" :return: Returns a list of tuples (:class:`Vector`, index, distance).\n"
" :rtype: :class:`list`\n"
);
static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
PyObject *py_list;
PyObject *py_co;
float co[3];
KDTreeNearest *nearest = NULL;
float radius;
int i, found;
const char *keywords[] = {"co", "radius", NULL};
if (!PyArg_ParseTupleAndKeywords(args, kwargs, (char *) "Of:find_range", (char **)keywords,
&py_co, &radius))
{
return NULL;
}
if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1)
return NULL;
if (radius < 0.0f) {
PyErr_SetString(PyExc_RuntimeError, "negative radius given");
return NULL;
}
if (self->count != self->count_balance) {
PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_range()");
return NULL;
}
found = BLI_kdtree_range_search(self->obj, co, NULL, &nearest, radius);
py_list = PyList_New(found);
for (i = 0; i < found; i++) {
PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
}
if (nearest) {
MEM_freeN(nearest);
}
return py_list;
}
static PyMethodDef PyKDTree_methods[] = {
{"insert", (PyCFunction)py_kdtree_insert, METH_VARARGS | METH_KEYWORDS, py_kdtree_insert_doc},
{"balance", (PyCFunction)py_kdtree_balance, METH_NOARGS, py_kdtree_balance_doc},
{"find", (PyCFunction)py_kdtree_find, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_doc},
{"find_n", (PyCFunction)py_kdtree_find_n, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_n_doc},
{"find_range", (PyCFunction)py_kdtree_find_range, METH_VARARGS | METH_KEYWORDS, py_kdtree_find_range_doc},
{NULL, NULL, 0, NULL}
};
PyDoc_STRVAR(py_KDtree_doc,
"KdTree(size) -> new kd-tree initialized to hold ``size`` items.\n"
"\n"
".. note::\n"
"\n"
" :class:`KDTree.balance` must have been called before using any of the ``find`` methods.\n"
);
PyTypeObject PyKDTree_Type = {
PyVarObject_HEAD_INIT(NULL, 0)
"KDTree", /* tp_name */
sizeof(PyKDTree), /* tp_basicsize */
0, /* tp_itemsize */
/* methods */
(destructor)PyKDTree__tp_dealloc, /* tp_dealloc */
NULL, /* tp_print */
NULL, /* tp_getattr */
NULL, /* tp_setattr */
NULL, /* tp_compare */
NULL, /* tp_repr */
NULL, /* tp_as_number */
NULL, /* tp_as_sequence */
NULL, /* tp_as_mapping */
NULL, /* tp_hash */
NULL, /* tp_call */
NULL, /* tp_str */
NULL, /* tp_getattro */
NULL, /* tp_setattro */
NULL, /* tp_as_buffer */
Py_TPFLAGS_DEFAULT, /* tp_flags */
py_KDtree_doc, /* Documentation string */
NULL, /* tp_traverse */
NULL, /* tp_clear */
NULL, /* tp_richcompare */
0, /* tp_weaklistoffset */
NULL, /* tp_iter */
NULL, /* tp_iternext */
(struct PyMethodDef *)PyKDTree_methods, /* tp_methods */
NULL, /* tp_members */
NULL, /* tp_getset */
NULL, /* tp_base */
NULL, /* tp_dict */
NULL, /* tp_descr_get */
NULL, /* tp_descr_set */
0, /* tp_dictoffset */
(initproc)PyKDTree__tp_init, /* tp_init */
(allocfunc)PyType_GenericAlloc, /* tp_alloc */
(newfunc)PyType_GenericNew, /* tp_new */
(freefunc)0, /* tp_free */
NULL, /* tp_is_gc */
NULL, /* tp_bases */
NULL, /* tp_mro */
NULL, /* tp_cache */
NULL, /* tp_subclasses */
NULL, /* tp_weaklist */
(destructor) NULL /* tp_del */
};
PyDoc_STRVAR(py_kdtree_doc,
"Generic 3-dimentional kd-tree to perform spatial searches."
);
static struct PyModuleDef kdtree_moduledef = {
PyModuleDef_HEAD_INIT,
"mathutils.kdtree", /* m_name */
py_kdtree_doc, /* m_doc */
0, /* m_size */
NULL, /* m_methods */
NULL, /* m_reload */
NULL, /* m_traverse */
NULL, /* m_clear */
NULL /* m_free */
};
PyMODINIT_FUNC PyInit_mathutils_kdtree(void)
{
PyObject *m = PyModule_Create(&kdtree_moduledef);
if (m == NULL) {
return NULL;
}
/* Register the 'KDTree' class */
if (PyType_Ready(&PyKDTree_Type)) {
return NULL;
}
PyModule_AddObject(m, (char *)"KDTree", (PyObject *) &PyKDTree_Type);
return m;
}

33
source/blender/python/mathutils/mathutils_kdtree.h Normal file
View File

@@ -0,0 +1,33 @@
/*
* ***** BEGIN GPL LICENSE BLOCK *****
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License
* as published by the Free Software Foundation; either version 2
* of the License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, write to the Free Software Foundation,
* Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
*
* ***** END GPL LICENSE BLOCK *****
*/
/** \file blender/python/mathutils/mathutils_kdtree.h
* \ingroup mathutils
*/
#ifndef __MATHUTILS_KDTREE_H__
#define __MATHUTILS_KDTREE_H__
PyObject *PyInit_mathutils_kdtree(void);
extern PyTypeObject PyKDTree_Type;
#endif /* __MATHUTILS_KDTREE_H__ */

108
source/tests/bl_pyapi_mathutils.py
View File

@@ -1,7 +1,8 @@
# ./blender.bin --background -noaudio --python source/tests/bl_pyapi_mathutils.py
import unittest
from test import support
from mathutils import Matrix, Vector
from mathutils import kdtree
class MatrixTesting(unittest.TestCase):
def test_matrix_column_access(self):
@@ -148,9 +149,114 @@ class MatrixTesting(unittest.TestCase):
self.assertEqual(mat * mat, prod_mat)
class KDTreeTesting(unittest.TestCase):
@staticmethod
def kdtree_create_grid_3d(tot):
k = kdtree.KDTree(tot * tot * tot)
index = 0
mul = 1.0 / (tot - 1)
for x in range(tot):
for y in range(tot):
for z in range(tot):
k.insert((x * mul, y * mul, z * mul), index)
index += 1
k.balance()
return k
def test_kdtree_single(self):
co = (0,) * 3
index = 2
k = kdtree.KDTree(1)
k.insert(co, index)
k.balance()
co_found, index_found, dist_found = k.find(co)
self.assertEqual(tuple(co_found), co)
self.assertEqual(index_found, index)
self.assertEqual(dist_found, 0.0)
def test_kdtree_empty(self):
co = (0,) * 3
k = kdtree.KDTree(0)
k.balance()
co_found, index_found, dist_found = k.find(co)
self.assertIsNone(co_found)
self.assertIsNone(index_found)
self.assertIsNone(dist_found)
def test_kdtree_line(self):
tot = 10
k = kdtree.KDTree(tot)
for i in range(tot):
k.insert((i,) * 3, i)
k.balance()
co_found, index_found, dist_found = k.find((-1,) * 3)
self.assertEqual(tuple(co_found), (0,) * 3)
co_found, index_found, dist_found = k.find((tot,) * 3)
self.assertEqual(tuple(co_found), (tot - 1,) * 3)
def test_kdtree_grid(self):
size = 10
k = self.kdtree_create_grid_3d(size)
# find_range
ret = k.find_range((0.5,) * 3, 2.0)
self.assertEqual(len(ret), size * size * size)
ret = k.find_range((1.0,) * 3, 1.0 / size)
self.assertEqual(len(ret), 1)
ret = k.find_range((1.0,) * 3, 2.0 / size)
self.assertEqual(len(ret), 8)
ret = k.find_range((10,) * 3, 0.5)
self.assertEqual(len(ret), 0)
# find_n
tot = 0
ret = k.find_n((1.0,) * 3, tot)
self.assertEqual(len(ret), tot)
tot = 10
ret = k.find_n((1.0,) * 3, tot)
self.assertEqual(len(ret), tot)
self.assertEqual(ret[0][2], 0.0)
tot = size * size * size
ret = k.find_n((1.0,) * 3, tot)
self.assertEqual(len(ret), tot)
def test_kdtree_invalid_size(self):
with self.assertRaises(ValueError):
kdtree.KDTree(-1)
def test_kdtree_invalid_balance(self):
co = (0,) * 3
index = 2
k = kdtree.KDTree(2)
k.insert(co, index)
k.balance()
k.insert(co, index)
with self.assertRaises(RuntimeError):
k.find(co)
def test_main():
try:
support.run_unittest(MatrixTesting)
support.run_unittest(KDTreeTesting)
except:
import traceback
traceback.print_exc()