Patch D133: Python wrapper for BLI_kdtree (adds mathutils.kdtree)
authorCampbell Barton <ideasman42@gmail.com>
Mon, 6 Jan 2014 09:32:34 +0000 (20:32 +1100)
committerCampbell Barton <ideasman42@gmail.com>
Mon, 6 Jan 2014 09:32:34 +0000 (20:32 +1100)
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

doc/python_api/examples/mathutils.kdtree.py [new file with mode: 0644]
doc/python_api/sphinx_doc_gen.py
source/blender/python/intern/bpy_interface.c
source/blender/python/mathutils/CMakeLists.txt
source/blender/python/mathutils/mathutils.c
source/blender/python/mathutils/mathutils.h
source/blender/python/mathutils/mathutils_kdtree.c [new file with mode: 0644]
source/blender/python/mathutils/mathutils_kdtree.h [new file with mode: 0644]
source/tests/bl_pyapi_mathutils.py

diff --git a/doc/python_api/examples/mathutils.kdtree.py b/doc/python_api/examples/mathutils.kdtree.py
new file mode 100644 (file)
index 0000000..d367582
--- /dev/null
@@ -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)
+
index a81286335699e2a9fc56ad2c921101ffdf32d826..4cea81d5cf01085372b1f6259971d9efe6b27c2e 100644 (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",
     }
index 68816b728a76b46c53949386d63f6d389e5f54b1..aafb5e3d6422681f78f6abdc570efbb580805732 100644 (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},
index dae1c665ebc440e66eaf9ee45650f580043e923b..133b8d3895c61e17039b09866e1d6a4263289ae2 100644 (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
 )
 
index 8f94fc8b467603cde5e06d19fd0c9823811f993e..dd3e5dec8a4a7e5f575420170258151d7d6c48fa 100644 (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);
index c465c27cbb5de8033cec3141350a8591a5551ec8..df1d5704190e599fb2f1cb5b24a4278f489025f4 100644 (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 *);
diff --git a/source/blender/python/mathutils/mathutils_kdtree.c b/source/blender/python/mathutils/mathutils_kdtree.c
new file mode 100644 (file)
index 0000000..aa9c7ee
--- /dev/null
@@ -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 (file)
index 0000000..8421661
--- /dev/null
@@ -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__ */
index 1754644e813db03a7279b8f4039fdebd15dbdb4c..bbb4d7f355fe77a3864479bf1eea5e987e62fd31 100644 (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()