diff options
author | Campbell Barton <ideasman42@gmail.com> | 2014-01-06 13:32:34 +0400 |
---|---|---|
committer | Campbell Barton <ideasman42@gmail.com> | 2014-01-06 13:32:34 +0400 |
commit | 4848f9029a7cf70324888f18ccfcdea3fa50392b (patch) | |
tree | 4a7735e561efaa96240756c9c585e891315bc24f | |
parent | 90efa345c2b512c7dbc615b6210127070ff8b03f (diff) |
Patch D133: Python wrapper for BLI_kdtree (adds mathutils.kdtree)
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
-rw-r--r-- | doc/python_api/examples/mathutils.kdtree.py | 37 | ||||
-rw-r--r-- | doc/python_api/sphinx_doc_gen.py | 4 | ||||
-rw-r--r-- | source/blender/python/intern/bpy_interface.c | 1 | ||||
-rw-r--r-- | source/blender/python/mathutils/CMakeLists.txt | 2 | ||||
-rw-r--r-- | source/blender/python/mathutils/mathutils.c | 5 | ||||
-rw-r--r-- | source/blender/python/mathutils/mathutils.h | 1 | ||||
-rw-r--r-- | source/blender/python/mathutils/mathutils_kdtree.c | 429 | ||||
-rw-r--r-- | source/blender/python/mathutils/mathutils_kdtree.h | 33 | ||||
-rw-r--r-- | source/tests/bl_pyapi_mathutils.py | 108 |
9 files changed, 618 insertions, 2 deletions
diff --git a/doc/python_api/examples/mathutils.kdtree.py b/doc/python_api/examples/mathutils.kdtree.py new file mode 100644 index 00000000000..d367582e5be --- /dev/null +++ b/doc/python_api/examples/mathutils.kdtree.py @@ -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) + diff --git a/doc/python_api/sphinx_doc_gen.py b/doc/python_api/sphinx_doc_gen.py index a8128633569..4cea81d5cf0 100644 --- a/doc/python_api/sphinx_doc_gen.py +++ b/doc/python_api/sphinx_doc_gen.py @@ -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", } diff --git a/source/blender/python/intern/bpy_interface.c b/source/blender/python/intern/bpy_interface.c index 68816b728a7..aafb5e3d642 100644 --- a/source/blender/python/intern/bpy_interface.c +++ b/source/blender/python/intern/bpy_interface.c @@ -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}, diff --git a/source/blender/python/mathutils/CMakeLists.txt b/source/blender/python/mathutils/CMakeLists.txt index dae1c665ebc..133b8d3895c 100644 --- a/source/blender/python/mathutils/CMakeLists.txt +++ b/source/blender/python/mathutils/CMakeLists.txt @@ -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 ) diff --git a/source/blender/python/mathutils/mathutils.c b/source/blender/python/mathutils/mathutils.c index 8f94fc8b467..dd3e5dec8a4 100644 --- a/source/blender/python/mathutils/mathutils.c +++ b/source/blender/python/mathutils/mathutils.c @@ -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); diff --git a/source/blender/python/mathutils/mathutils.h b/source/blender/python/mathutils/mathutils.h index c465c27cbb5..df1d5704190 100644 --- a/source/blender/python/mathutils/mathutils.h +++ b/source/blender/python/mathutils/mathutils.h @@ -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 *); diff --git a/source/blender/python/mathutils/mathutils_kdtree.c b/source/blender/python/mathutils/mathutils_kdtree.c new file mode 100644 index 00000000000..aa9c7eecc6b --- /dev/null +++ b/source/blender/python/mathutils/mathutils_kdtree.c @@ -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; +} diff --git a/source/blender/python/mathutils/mathutils_kdtree.h b/source/blender/python/mathutils/mathutils_kdtree.h new file mode 100644 index 00000000000..84216617712 --- /dev/null +++ b/source/blender/python/mathutils/mathutils_kdtree.h @@ -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__ */ diff --git a/source/tests/bl_pyapi_mathutils.py b/source/tests/bl_pyapi_mathutils.py index 1754644e813..bbb4d7f355f 100644 --- a/source/tests/bl_pyapi_mathutils.py +++ b/source/tests/bl_pyapi_mathutils.py @@ -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() |