mirror of
https://github.com/blender/blender
synced 2026-09-29 04:37:17 +03:00
PyAPI: support 2D KDTree's for mathutils.kdtree.KDTree
Support 2D KDTree's as well as 3D, allowing KDTree's to be built from 2D data, using 2D coordinates for lookups/searching. - Add dimension keyword argument for new KDTree's. - Add KDTree.dimension read-only accessor. - Add `Doxygen` documentation. - Update tests. Ref !159342
This commit is contained in:
Oxicid
Campbell Barton
parent
26744b6e2e
commit
e81d8ae45a
2 changed files with 344 additions and 136 deletions
|
|
@ -6,7 +6,7 @@
|
|||
* \ingroup mathutils
|
||||
*
|
||||
* This file defines the 'mathutils.kdtree' module, a general purpose module to access
|
||||
* blenders kdtree for 3d spatial lookups.
|
||||
* blenders kdtree for 2D/3D spatial lookups.
|
||||
*/
|
||||
|
||||
#include <Python.h>
|
||||
|
|
@ -28,45 +28,57 @@ namespace blender {
|
|||
|
||||
struct PyKDTree {
|
||||
PyObject_HEAD
|
||||
KDTree<float3> *obj;
|
||||
/* Used for 2D/3D KDTrees. */
|
||||
void *obj;
|
||||
uint maxsize;
|
||||
uint count;
|
||||
uint count_balance; /* size when we last balanced */
|
||||
int dimensions;
|
||||
};
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/* Utility helper functions */
|
||||
/** \name Utility helper functions
|
||||
* \{ */
|
||||
|
||||
static void kdtree_nearest_to_py_tuple(const KDTreeNearest<float3> *nearest, PyObject *py_retval)
|
||||
/**
|
||||
* Access the tree, callers must pass a `CoordT` matching #PyKDTree::dimensions.
|
||||
*/
|
||||
template<typename CoordT> static KDTree<CoordT> *pykdtree_tree_get(PyKDTree *self)
|
||||
{
|
||||
/* Null when `__init__` never ran or raised an error, `dimensions` is unset in that case. */
|
||||
BLI_assert(self->obj == nullptr || CoordT::type_length == self->dimensions);
|
||||
return reinterpret_cast<KDTree<CoordT> *>(self->obj);
|
||||
}
|
||||
|
||||
template<typename CoordT>
|
||||
static void kdtree_nearest_to_py_tuple(const KDTreeNearest<CoordT> *nearest, PyObject *py_retval)
|
||||
{
|
||||
BLI_assert(nearest->index >= 0);
|
||||
BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);
|
||||
|
||||
PyTuple_SET_ITEMS(py_retval,
|
||||
Vector_CreatePyObject(nearest->co, 3, nullptr),
|
||||
Vector_CreatePyObject(nearest->co, CoordT::type_length, nullptr),
|
||||
PyLong_FromLong(nearest->index),
|
||||
PyFloat_FromDouble(nearest->dist));
|
||||
}
|
||||
|
||||
static PyObject *kdtree_nearest_to_py(const KDTreeNearest<float3> *nearest)
|
||||
template<typename CoordT>
|
||||
static PyObject *kdtree_nearest_to_py(const KDTreeNearest<CoordT> *nearest)
|
||||
{
|
||||
PyObject *py_retval;
|
||||
PyObject *py_retval = PyTuple_New(3);
|
||||
|
||||
py_retval = PyTuple_New(3);
|
||||
|
||||
kdtree_nearest_to_py_tuple(nearest, py_retval);
|
||||
kdtree_nearest_to_py_tuple<CoordT>(nearest, py_retval);
|
||||
|
||||
return py_retval;
|
||||
}
|
||||
|
||||
static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest<float3> *nearest)
|
||||
template<typename CoordT>
|
||||
static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest<CoordT> *nearest)
|
||||
{
|
||||
PyObject *py_retval;
|
||||
|
||||
py_retval = PyTuple_New(3);
|
||||
PyObject *py_retval = PyTuple_New(3);
|
||||
|
||||
if (nearest->index != -1) {
|
||||
kdtree_nearest_to_py_tuple(nearest, py_retval);
|
||||
kdtree_nearest_to_py_tuple<CoordT>(nearest, py_retval);
|
||||
}
|
||||
else {
|
||||
PyC_Tuple_Fill(py_retval, Py_None);
|
||||
|
|
@ -75,8 +87,11 @@ static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest<float3> *nea
|
|||
return py_retval;
|
||||
}
|
||||
|
||||
/** \} */
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/* KDTree */
|
||||
/** \name KDTree
|
||||
* \{ */
|
||||
|
||||
/* annoying since arg parsing won't check overflow */
|
||||
#define UINT_IS_NEG(n) ((n) > INT_MAX)
|
||||
|
|
@ -84,14 +99,22 @@ static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest<float3> *nea
|
|||
static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
|
||||
{
|
||||
uint maxsize;
|
||||
const char *keywords[] = {"size", nullptr};
|
||||
int dimensions = 3;
|
||||
const char *keywords[] = {
|
||||
"size",
|
||||
"dimensions",
|
||||
nullptr,
|
||||
};
|
||||
|
||||
if (!PyArg_ParseTupleAndKeywords(args,
|
||||
kwargs,
|
||||
"I" /* `size` */
|
||||
"I" /* `size` */
|
||||
"|$" /* Optional, keyword only arguments. */
|
||||
"i" /* `dimensions` */
|
||||
":KDTree",
|
||||
const_cast<char **>(keywords),
|
||||
&maxsize))
|
||||
&maxsize,
|
||||
&dimensions))
|
||||
{
|
||||
return -1;
|
||||
}
|
||||
|
|
@ -101,21 +124,44 @@ static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
|
|||
return -1;
|
||||
}
|
||||
|
||||
self->obj = kdtree_new<float3>(maxsize);
|
||||
if (dimensions == 2) {
|
||||
self->obj = kdtree_new<float2>(maxsize);
|
||||
}
|
||||
else if (dimensions == 3) {
|
||||
self->obj = kdtree_new<float3>(maxsize);
|
||||
}
|
||||
else {
|
||||
PyErr_SetString(PyExc_ValueError, "dimensions must be 2 or 3");
|
||||
return -1;
|
||||
}
|
||||
|
||||
self->maxsize = maxsize;
|
||||
self->count = 0;
|
||||
/* Initialize `uint-max` to avoid crashes on unbalanced trees. */
|
||||
self->count_balance = uint(-1);
|
||||
self->dimensions = dimensions;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
static void PyKDTree__tp_dealloc(PyKDTree *self)
|
||||
{
|
||||
kdtree_free<float3>(self->obj);
|
||||
if (self->dimensions == 2) {
|
||||
kdtree_free<float2>(pykdtree_tree_get<float2>(self));
|
||||
}
|
||||
else {
|
||||
kdtree_free<float3>(pykdtree_tree_get<float3>(self));
|
||||
}
|
||||
|
||||
Py_TYPE(self)->tp_free(reinterpret_cast<PyObject *>(self));
|
||||
}
|
||||
|
||||
/** \} */
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Methods: Insert
|
||||
* \{ */
|
||||
|
||||
PyDoc_STRVAR(
|
||||
/* Wrap. */
|
||||
py_kdtree_insert_doc,
|
||||
|
|
@ -123,7 +169,7 @@ PyDoc_STRVAR(
|
|||
"\n"
|
||||
" Insert a point into the KDTree.\n"
|
||||
"\n"
|
||||
" :param co: Point 3d position.\n"
|
||||
" :param co: Point position. Can be 2D or 3D, based on the KDTree dimensions.\n"
|
||||
" :type co: Sequence[float]\n"
|
||||
" :param index: The index of the point (must be non-negative).\n"
|
||||
" :type index: int\n");
|
||||
|
|
@ -146,7 +192,9 @@ static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwar
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1) {
|
||||
if (mathutils_array_parse(
|
||||
co, self->dimensions, self->dimensions, py_co, "insert: invalid 'co' arg") == -1)
|
||||
{
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
|
|
@ -160,12 +208,23 @@ static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwar
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
kdtree_insert<float3>(self->obj, index, co);
|
||||
if (self->dimensions == 2) {
|
||||
kdtree_insert<float2>(pykdtree_tree_get<float2>(self), index, co);
|
||||
}
|
||||
else {
|
||||
kdtree_insert<float3>(pykdtree_tree_get<float3>(self), index, co);
|
||||
}
|
||||
self->count++;
|
||||
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
/** \} */
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Methods: Balance
|
||||
* \{ */
|
||||
|
||||
PyDoc_STRVAR(
|
||||
/* Wrap. */
|
||||
py_kdtree_balance_doc,
|
||||
|
|
@ -178,20 +237,30 @@ PyDoc_STRVAR(
|
|||
" This builds the entire tree, avoid calling after each insertion.\n");
|
||||
static PyObject *py_kdtree_balance(PyKDTree *self)
|
||||
{
|
||||
kdtree_balance<float3>(self->obj);
|
||||
if (self->dimensions == 2) {
|
||||
kdtree_balance<float2>(pykdtree_tree_get<float2>(self));
|
||||
}
|
||||
else {
|
||||
kdtree_balance<float3>(pykdtree_tree_get<float3>(self));
|
||||
}
|
||||
|
||||
self->count_balance = self->count;
|
||||
Py_RETURN_NONE;
|
||||
}
|
||||
|
||||
/** \} */
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Methods: Find
|
||||
* \{ */
|
||||
|
||||
struct PyKDTree_NearestData {
|
||||
PyObject *py_filter;
|
||||
bool is_error;
|
||||
};
|
||||
|
||||
static int py_find_nearest_cb(void *user_data, int index, const float3 &co, float dist_sq)
|
||||
static int py_find_nearest_cb(void *user_data, int index)
|
||||
{
|
||||
UNUSED_VARS(co, dist_sq);
|
||||
|
||||
PyKDTree_NearestData *data = static_cast<PyKDTree_NearestData *>(user_data);
|
||||
|
||||
PyObject *py_args = PyTuple_New(1);
|
||||
|
|
@ -212,6 +281,32 @@ static int py_find_nearest_cb(void *user_data, int index, const float3 &co, floa
|
|||
return -1;
|
||||
}
|
||||
|
||||
template<typename CoordT>
|
||||
static PyObject *py_kdtree_find_impl(PyKDTree *self, const float *co, PyObject *py_filter)
|
||||
{
|
||||
KDTree<CoordT> *tree = pykdtree_tree_get<CoordT>(self);
|
||||
KDTreeNearest<CoordT> nearest;
|
||||
nearest.index = -1;
|
||||
|
||||
if (py_filter == Py_None) {
|
||||
kdtree_find_nearest<CoordT>(tree, co, &nearest);
|
||||
}
|
||||
else {
|
||||
PyKDTree_NearestData data{py_filter, false};
|
||||
|
||||
kdtree_find_nearest_cb<CoordT>(
|
||||
tree, co, &nearest, [&](int index, const CoordT & /*co_nearest*/, float /*dist_sq*/) {
|
||||
return py_find_nearest_cb(&data, index);
|
||||
});
|
||||
|
||||
if (data.is_error) {
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
return kdtree_nearest_to_py_and_check<CoordT>(&nearest);
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(
|
||||
/* Wrap. */
|
||||
py_kdtree_find_doc,
|
||||
|
|
@ -219,7 +314,7 @@ PyDoc_STRVAR(
|
|||
"\n"
|
||||
" Find nearest point to ``co``.\n"
|
||||
"\n"
|
||||
" :param co: 3D coordinate.\n"
|
||||
" :param co: Point position. Can be 2D or 3D, based on the KDTree dimensions.\n"
|
||||
" :type co: Sequence[float]\n"
|
||||
" :param filter: function which takes an index and returns True for indices to "
|
||||
"include in the search.\n"
|
||||
|
|
@ -231,7 +326,6 @@ static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs
|
|||
{
|
||||
PyObject *py_co, *py_filter = Py_None;
|
||||
float co[3];
|
||||
KDTreeNearest<float3> nearest;
|
||||
const char *keywords[] = {"co", "filter", nullptr};
|
||||
|
||||
if (!PyArg_ParseTupleAndKeywords(args,
|
||||
|
|
@ -247,7 +341,9 @@ static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1) {
|
||||
if (mathutils_array_parse(
|
||||
co, self->dimensions, self->dimensions, py_co, "find: invalid 'co' arg") == -1)
|
||||
{
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
|
|
@ -256,28 +352,34 @@ static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
nearest.index = -1;
|
||||
|
||||
if (py_filter == Py_None) {
|
||||
kdtree_find_nearest<float3>(self->obj, co, &nearest);
|
||||
if (self->dimensions == 2) {
|
||||
return py_kdtree_find_impl<float2>(self, co, py_filter);
|
||||
}
|
||||
else {
|
||||
PyKDTree_NearestData data = {nullptr};
|
||||
return py_kdtree_find_impl<float3>(self, co, py_filter);
|
||||
}
|
||||
|
||||
data.py_filter = py_filter;
|
||||
data.is_error = false;
|
||||
/** \} */
|
||||
|
||||
kdtree_find_nearest_cb<float3>(
|
||||
self->obj, co, &nearest, [&](int index, const float3 &co_nearest, float dist_sq) {
|
||||
return py_find_nearest_cb(&data, index, co_nearest, dist_sq);
|
||||
});
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Methods: Find N
|
||||
* \{ */
|
||||
|
||||
if (data.is_error) {
|
||||
return nullptr;
|
||||
}
|
||||
template<typename CoordT>
|
||||
static PyObject *py_kdtree_find_n_impl(PyKDTree *self, const float *co, uint n)
|
||||
{
|
||||
KDTreeNearest<CoordT> *nearest = MEM_new_array_uninitialized<KDTreeNearest<CoordT>>(n, __func__);
|
||||
|
||||
const int found = kdtree_find_nearest_n<CoordT>(pykdtree_tree_get<CoordT>(self), co, nearest, n);
|
||||
|
||||
PyObject *py_list = PyList_New(found);
|
||||
|
||||
for (int i = 0; i < found; i++) {
|
||||
PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py<CoordT>(&nearest[i]));
|
||||
}
|
||||
|
||||
return kdtree_nearest_to_py_and_check(&nearest);
|
||||
MEM_delete(nearest);
|
||||
|
||||
return py_list;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(
|
||||
|
|
@ -287,7 +389,7 @@ PyDoc_STRVAR(
|
|||
"\n"
|
||||
" Find nearest ``n`` points to ``co``.\n"
|
||||
"\n"
|
||||
" :param co: 3D coordinate.\n"
|
||||
" :param co: Point position. Can be 2D or 3D, based on the KDTree dimensions.\n"
|
||||
" :type co: Sequence[float]\n"
|
||||
" :param n: Number of points to find.\n"
|
||||
" :type n: int\n"
|
||||
|
|
@ -295,12 +397,9 @@ PyDoc_STRVAR(
|
|||
" :rtype: list[tuple[:class:`Vector`, int, float]]\n");
|
||||
static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
|
||||
{
|
||||
PyObject *py_list;
|
||||
PyObject *py_co;
|
||||
float co[3];
|
||||
KDTreeNearest<float3> *nearest;
|
||||
uint n;
|
||||
int i, found;
|
||||
const char *keywords[] = {"co", "n", nullptr};
|
||||
|
||||
if (!PyArg_ParseTupleAndKeywords(args,
|
||||
|
|
@ -315,7 +414,9 @@ static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwar
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1) {
|
||||
if (mathutils_array_parse(
|
||||
co, self->dimensions, self->dimensions, py_co, "find_n: invalid 'co' arg") == -1)
|
||||
{
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
|
|
@ -329,17 +430,34 @@ static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwar
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
nearest = MEM_new_array_uninitialized<KDTreeNearest<float3>>(n, __func__);
|
||||
if (self->dimensions == 2) {
|
||||
return py_kdtree_find_n_impl<float2>(self, co, n);
|
||||
}
|
||||
return py_kdtree_find_n_impl<float3>(self, co, n);
|
||||
}
|
||||
|
||||
found = kdtree_find_nearest_n<float3>(self->obj, co, nearest, n);
|
||||
/** \} */
|
||||
|
||||
py_list = PyList_New(found);
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Methods: Find Range
|
||||
* \{ */
|
||||
|
||||
for (i = 0; i < found; i++) {
|
||||
PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
|
||||
template<typename CoordT>
|
||||
static PyObject *py_kdtree_find_range_impl(PyKDTree *self, const float *co, float radius)
|
||||
{
|
||||
KDTreeNearest<CoordT> *nearest = nullptr;
|
||||
const int found = kdtree_range_search<CoordT>(
|
||||
pykdtree_tree_get<CoordT>(self), co, &nearest, radius);
|
||||
|
||||
PyObject *py_list = PyList_New(found);
|
||||
|
||||
for (int i = 0; i < found; i++) {
|
||||
PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py<CoordT>(&nearest[i]));
|
||||
}
|
||||
|
||||
MEM_delete(nearest);
|
||||
if (nearest) {
|
||||
MEM_delete(nearest);
|
||||
}
|
||||
|
||||
return py_list;
|
||||
}
|
||||
|
|
@ -351,7 +469,7 @@ PyDoc_STRVAR(
|
|||
"\n"
|
||||
" Find all points within ``radius`` of ``co``.\n"
|
||||
"\n"
|
||||
" :param co: 3D coordinate.\n"
|
||||
" :param co: Point position. Can be 2D or 3D, based on the KDTree dimensions.\n"
|
||||
" :type co: Sequence[float]\n"
|
||||
" :param radius: Maximum distance to search for points.\n"
|
||||
" :type radius: float\n"
|
||||
|
|
@ -359,12 +477,9 @@ PyDoc_STRVAR(
|
|||
" :rtype: list[tuple[:class:`Vector`, int, float]]\n");
|
||||
static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
|
||||
{
|
||||
PyObject *py_list;
|
||||
PyObject *py_co;
|
||||
float co[3];
|
||||
KDTreeNearest<float3> *nearest = nullptr;
|
||||
float radius;
|
||||
int i, found;
|
||||
|
||||
const char *keywords[] = {"co", "radius", nullptr};
|
||||
|
||||
|
|
@ -380,7 +495,9 @@ static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1) {
|
||||
if (mathutils_array_parse(
|
||||
co, self->dimensions, self->dimensions, py_co, "find_range: invalid 'co' arg") == -1)
|
||||
{
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
|
|
@ -394,21 +511,66 @@ static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *
|
|||
return nullptr;
|
||||
}
|
||||
|
||||
found = kdtree_range_search<float3>(self->obj, co, &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 (self->dimensions == 2) {
|
||||
return py_kdtree_find_range_impl<float2>(self, co, radius);
|
||||
}
|
||||
|
||||
if (nearest) {
|
||||
MEM_delete(nearest);
|
||||
}
|
||||
|
||||
return py_list;
|
||||
return py_kdtree_find_range_impl<float3>(self, co, radius);
|
||||
}
|
||||
|
||||
/** \} */
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Type: Get/Set Item Implementation
|
||||
* \{ */
|
||||
|
||||
/* `KDTree.dimensions`. */
|
||||
PyDoc_STRVAR(
|
||||
/* Wrap. */
|
||||
py_kdtree_dimensions_doc,
|
||||
"KDTree dimensions.\n"
|
||||
"\n"
|
||||
":type: int\n");
|
||||
static PyObject *py_kdtree_dimensions_get(PyKDTree *self, void * /*closure*/)
|
||||
{
|
||||
return PyLong_FromLong(self->dimensions);
|
||||
}
|
||||
|
||||
/** \} */
|
||||
|
||||
/* -------------------------------------------------------------------- */
|
||||
/** \name KDTree Type: Get/Set Item Definitions
|
||||
* \{ */
|
||||
|
||||
#ifdef __GNUC__
|
||||
# ifdef __clang__
|
||||
# pragma clang diagnostic push
|
||||
# pragma clang diagnostic ignored "-Wcast-function-type"
|
||||
# else
|
||||
# pragma GCC diagnostic push
|
||||
# pragma GCC diagnostic ignored "-Wcast-function-type"
|
||||
# endif
|
||||
#endif
|
||||
|
||||
static PyGetSetDef PyKDTree_getseters[] = {
|
||||
{"dimensions",
|
||||
reinterpret_cast<getter>(py_kdtree_dimensions_get),
|
||||
static_cast<setter>(nullptr),
|
||||
py_kdtree_dimensions_doc,
|
||||
nullptr},
|
||||
|
||||
{nullptr, nullptr, nullptr, nullptr, nullptr} /* Sentinel */
|
||||
};
|
||||
|
||||
#ifdef __GNUC__
|
||||
# ifdef __clang__
|
||||
# pragma clang diagnostic pop
|
||||
# else
|
||||
# pragma GCC diagnostic pop
|
||||
# endif
|
||||
#endif
|
||||
|
||||
/** \} */
|
||||
|
||||
#ifdef __GNUC__
|
||||
# ifdef __clang__
|
||||
# pragma clang diagnostic push
|
||||
|
|
@ -454,12 +616,14 @@ static PyMethodDef PyKDTree_methods[] = {
|
|||
PyDoc_STRVAR(
|
||||
/* Wrap. */
|
||||
py_KDtree_doc,
|
||||
".. class:: KDTree(size)\n"
|
||||
".. class:: KDTree(size, *, dimensions=3)\n"
|
||||
"\n"
|
||||
" KDTree(size) -> new kd-tree initialized to hold up to ``size`` items.\n"
|
||||
" KDTree(size, *, dimensions=3) -> new kd-tree initialized to hold up to ``size`` items.\n"
|
||||
"\n"
|
||||
" :param size: Maximum number of items.\n"
|
||||
" :type size: int\n"
|
||||
" :param dimensions: The dimensions of the tree (2 or 3).\n"
|
||||
" :type dimensions: int\n"
|
||||
"\n"
|
||||
" .. note::\n"
|
||||
"\n"
|
||||
|
|
@ -495,7 +659,7 @@ PyTypeObject PyKDTree_Type = {
|
|||
/*tp_iternext*/ nullptr,
|
||||
/*tp_methods*/ static_cast<PyMethodDef *>(PyKDTree_methods),
|
||||
/*tp_members*/ nullptr,
|
||||
/*tp_getset*/ nullptr,
|
||||
/*tp_getset*/ PyKDTree_getseters,
|
||||
/*tp_base*/ nullptr,
|
||||
/*tp_dict*/ nullptr,
|
||||
/*tp_descr_get*/ nullptr,
|
||||
|
|
@ -520,7 +684,7 @@ PyTypeObject PyKDTree_Type = {
|
|||
PyDoc_STRVAR(
|
||||
/* Wrap. */
|
||||
py_kdtree_doc,
|
||||
"Generic 3-dimensional kd-tree to perform spatial searches.");
|
||||
"Generic 2D/3D kd-tree to perform spatial searches.");
|
||||
static PyModuleDef kdtree_moduledef = {
|
||||
/*m_base*/ PyModuleDef_HEAD_INIT,
|
||||
/*m_name*/ "mathutils.kdtree",
|
||||
|
|
|
|||
|
|
@ -1005,19 +1005,16 @@ class TypeTesting(unittest.TestCase):
|
|||
|
||||
class KDTreeTesting(unittest.TestCase):
|
||||
@staticmethod
|
||||
def kdtree_create_grid_3d_data(tot):
|
||||
index = 0
|
||||
def kdtree_create_grid_data(tot, dimensions):
|
||||
import itertools
|
||||
mul = 1.0 / (tot - 1)
|
||||
for x in range(tot):
|
||||
for y in range(tot):
|
||||
for z in range(tot):
|
||||
yield (x * mul, y * mul, z * mul), index
|
||||
index += 1
|
||||
for index, co in enumerate(itertools.product(range(tot), repeat=dimensions)):
|
||||
yield tuple(axis * mul for axis in co), index
|
||||
|
||||
@staticmethod
|
||||
def kdtree_create_grid_3d(tot, *, filter_fn=None):
|
||||
k = kdtree.KDTree(tot * tot * tot)
|
||||
for co, index in KDTreeTesting.kdtree_create_grid_3d_data(tot):
|
||||
def kdtree_create_grid(tot, dimensions, *, filter_fn=None):
|
||||
k = kdtree.KDTree(tot ** dimensions, dimensions=dimensions)
|
||||
for co, index in KDTreeTesting.kdtree_create_grid_data(tot, dimensions):
|
||||
if (filter_fn is not None) and (not filter_fn(co, index)):
|
||||
continue
|
||||
k.insert(co, index)
|
||||
|
|
@ -1029,13 +1026,15 @@ class KDTreeTesting(unittest.TestCase):
|
|||
self.assertAlmostEqual(first[1], second[1], places=places, msg=msg, delta=delta)
|
||||
self.assertAlmostEqual(first[2], second[2], places=places, msg=msg, delta=delta)
|
||||
|
||||
def test_kdtree_single(self):
|
||||
co = (0,) * 3
|
||||
def _test_kdtree_single_test_impl(self, dimensions):
|
||||
# Use a different value for each axis to detect axis mix-ups.
|
||||
co = tuple(range(5, 5 + dimensions))
|
||||
index = 2
|
||||
|
||||
k = kdtree.KDTree(1)
|
||||
k = kdtree.KDTree(1, dimensions=dimensions)
|
||||
k.insert(co, index)
|
||||
k.balance()
|
||||
self.assertEqual(k.dimensions, dimensions)
|
||||
|
||||
co_found, index_found, dist_found = k.find(co)
|
||||
|
||||
|
|
@ -1043,10 +1042,16 @@ class KDTreeTesting(unittest.TestCase):
|
|||
self.assertEqual(index_found, index)
|
||||
self.assertEqual(dist_found, 0.0)
|
||||
|
||||
def test_kdtree_empty(self):
|
||||
co = (0,) * 3
|
||||
def test_kdtree_single_2d(self):
|
||||
self._test_kdtree_single_test_impl(dimensions=2)
|
||||
|
||||
k = kdtree.KDTree(0)
|
||||
def test_kdtree_single_3d(self):
|
||||
self._test_kdtree_single_test_impl(dimensions=3)
|
||||
|
||||
def _test_kdtree_empty_test_impl(self, dimensions):
|
||||
co = (0,) * dimensions
|
||||
|
||||
k = kdtree.KDTree(0, dimensions=dimensions)
|
||||
k.balance()
|
||||
|
||||
co_found, index_found, dist_found = k.find(co)
|
||||
|
|
@ -1055,95 +1060,134 @@ class KDTreeTesting(unittest.TestCase):
|
|||
self.assertIsNone(index_found)
|
||||
self.assertIsNone(dist_found)
|
||||
|
||||
def test_kdtree_line(self):
|
||||
def test_kdtree_empty_2d(self):
|
||||
self._test_kdtree_empty_test_impl(dimensions=2)
|
||||
|
||||
def test_kdtree_empty_3d(self):
|
||||
self._test_kdtree_empty_test_impl(dimensions=3)
|
||||
|
||||
def _test_kdtree_line_test_impl(self, dimensions):
|
||||
tot = 10
|
||||
|
||||
k = kdtree.KDTree(tot)
|
||||
k = kdtree.KDTree(tot, dimensions=dimensions)
|
||||
|
||||
for i in range(tot):
|
||||
k.insert((i,) * 3, i)
|
||||
k.insert((i,) * dimensions, i)
|
||||
|
||||
k.balance()
|
||||
|
||||
co_found, index_found, dist_found = k.find((-1,) * 3)
|
||||
self.assertEqual(tuple(co_found), (0,) * 3)
|
||||
# The nearest point is one unit away on every axis.
|
||||
dist_expect = math.sqrt(dimensions)
|
||||
|
||||
co_found, index_found, dist_found = k.find((tot,) * 3)
|
||||
self.assertEqual(tuple(co_found), (tot - 1,) * 3)
|
||||
co_found, index_found, dist_found = k.find((-1,) * dimensions)
|
||||
self.assertEqual(tuple(co_found), (0,) * dimensions)
|
||||
self.assertEqual(index_found, 0)
|
||||
self.assertAlmostEqual(dist_found, dist_expect)
|
||||
|
||||
def test_kdtree_grid(self):
|
||||
co_found, index_found, dist_found = k.find((tot,) * dimensions)
|
||||
self.assertEqual(tuple(co_found), (tot - 1,) * dimensions)
|
||||
self.assertEqual(index_found, tot - 1)
|
||||
self.assertAlmostEqual(dist_found, dist_expect)
|
||||
|
||||
def test_kdtree_line_2d(self):
|
||||
self._test_kdtree_line_test_impl(dimensions=2)
|
||||
|
||||
def test_kdtree_line_3d(self):
|
||||
self._test_kdtree_line_test_impl(dimensions=3)
|
||||
|
||||
def _test_kdtree_grid_test_impl(self, dimensions):
|
||||
size = 10
|
||||
k = self.kdtree_create_grid_3d(size)
|
||||
k = self.kdtree_create_grid(size, dimensions)
|
||||
self.assertEqual(k.dimensions, dimensions)
|
||||
|
||||
# find_range
|
||||
ret = k.find_range((0.5,) * 3, 2.0)
|
||||
self.assertEqual(len(ret), size * size * size)
|
||||
ret = k.find_range((0.5,) * dimensions, 2.0)
|
||||
self.assertEqual(len(ret), size ** dimensions)
|
||||
self.assertEqual(len(ret[0][0]), dimensions)
|
||||
|
||||
ret = k.find_range((1.0,) * 3, 1.0 / size)
|
||||
ret = k.find_range((1.0,) * dimensions, 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((1.0,) * dimensions, 2.0 / size)
|
||||
self.assertEqual(len(ret), 2 ** dimensions)
|
||||
|
||||
ret = k.find_range((10,) * 3, 0.5)
|
||||
ret = k.find_range((10,) * dimensions, 0.5)
|
||||
self.assertEqual(len(ret), 0)
|
||||
|
||||
# find_n
|
||||
tot = 0
|
||||
ret = k.find_n((1.0,) * 3, tot)
|
||||
ret = k.find_n((1.0,) * dimensions, tot)
|
||||
self.assertEqual(len(ret), tot)
|
||||
|
||||
tot = 10
|
||||
ret = k.find_n((1.0,) * 3, tot)
|
||||
ret = k.find_n((1.0,) * dimensions, tot)
|
||||
self.assertEqual(len(ret), tot)
|
||||
self.assertEqual(len(ret[0][0]), dimensions)
|
||||
self.assertEqual(ret[0][2], 0.0)
|
||||
|
||||
tot = size * size * size
|
||||
ret = k.find_n((1.0,) * 3, tot)
|
||||
tot = size ** dimensions
|
||||
ret = k.find_n((1.0,) * dimensions, tot)
|
||||
self.assertEqual(len(ret), tot)
|
||||
|
||||
def test_kdtree_grid_filter_simple(self):
|
||||
def test_kdtree_grid_2d(self):
|
||||
self._test_kdtree_grid_test_impl(dimensions=2)
|
||||
|
||||
def test_kdtree_grid_3d(self):
|
||||
self._test_kdtree_grid_test_impl(dimensions=3)
|
||||
|
||||
def _test_kdtree_grid_filter_simple_test_impl(self, dimensions):
|
||||
size = 10
|
||||
k = self.kdtree_create_grid_3d(size)
|
||||
k = self.kdtree_create_grid(size, dimensions)
|
||||
|
||||
# filter exact index
|
||||
ret_regular = k.find((1.0,) * 3)
|
||||
ret_filter = k.find((1.0,) * 3, filter=lambda i: i == ret_regular[1])
|
||||
ret_regular = k.find((1.0,) * dimensions)
|
||||
ret_filter = k.find((1.0,) * dimensions, filter=lambda i: i == ret_regular[1])
|
||||
self.assertEqual(ret_regular, ret_filter)
|
||||
ret_filter = k.find((-1.0,) * 3, filter=lambda i: i == ret_regular[1])
|
||||
ret_filter = k.find((-1.0,) * dimensions, filter=lambda i: i == ret_regular[1])
|
||||
self.assertEqual(ret_regular[:2], ret_filter[:2]) # ignore distance
|
||||
|
||||
def test_kdtree_grid_filter_pairs(self):
|
||||
def test_kdtree_grid_filter_simple_2d(self):
|
||||
self._test_kdtree_grid_filter_simple_test_impl(dimensions=2)
|
||||
|
||||
def test_kdtree_grid_filter_simple_3d(self):
|
||||
self._test_kdtree_grid_filter_simple_test_impl(dimensions=3)
|
||||
|
||||
def _test_kdtree_grid_filter_pairs_test_impl(self, dimensions):
|
||||
import itertools
|
||||
size = 10
|
||||
k_all = self.kdtree_create_grid_3d(size)
|
||||
k_odd = self.kdtree_create_grid_3d(size, filter_fn=lambda co, i: (i % 2) == 1)
|
||||
k_evn = self.kdtree_create_grid_3d(size, filter_fn=lambda co, i: (i % 2) == 0)
|
||||
k_all = self.kdtree_create_grid(size, dimensions)
|
||||
k_odd = self.kdtree_create_grid(size, dimensions, filter_fn=lambda co, i: (i % 2) == 1)
|
||||
k_evn = self.kdtree_create_grid(size, dimensions, filter_fn=lambda co, i: (i % 2) == 0)
|
||||
|
||||
samples = 5
|
||||
mul = 1 / (samples - 1)
|
||||
for x in range(samples):
|
||||
for y in range(samples):
|
||||
for z in range(samples):
|
||||
co = (x * mul, y * mul, z * mul)
|
||||
for co_grid in itertools.product(range(samples), repeat=dimensions):
|
||||
co = tuple(axis * mul for axis in co_grid)
|
||||
|
||||
ret_regular = k_odd.find(co)
|
||||
self.assertEqual(ret_regular[1] % 2, 1)
|
||||
ret_filter = k_all.find(co, filter=lambda i: (i % 2) == 1)
|
||||
self.assertAlmostEqualVector(ret_regular, ret_filter)
|
||||
ret_regular = k_odd.find(co)
|
||||
self.assertEqual(ret_regular[1] % 2, 1)
|
||||
ret_filter = k_all.find(co, filter=lambda i: (i % 2) == 1)
|
||||
self.assertAlmostEqualVector(ret_regular, ret_filter)
|
||||
|
||||
ret_regular = k_evn.find(co)
|
||||
self.assertEqual(ret_regular[1] % 2, 0)
|
||||
ret_filter = k_all.find(co, filter=lambda i: (i % 2) == 0)
|
||||
self.assertAlmostEqualVector(ret_regular, ret_filter)
|
||||
ret_regular = k_evn.find(co)
|
||||
self.assertEqual(ret_regular[1] % 2, 0)
|
||||
ret_filter = k_all.find(co, filter=lambda i: (i % 2) == 0)
|
||||
self.assertAlmostEqualVector(ret_regular, ret_filter)
|
||||
|
||||
# filter out all values (search odd tree for even values and the reverse)
|
||||
co = (0,) * 3
|
||||
co = (0,) * dimensions
|
||||
ret_filter = k_odd.find(co, filter=lambda i: (i % 2) == 0)
|
||||
self.assertEqual(ret_filter[1], None)
|
||||
|
||||
ret_filter = k_evn.find(co, filter=lambda i: (i % 2) == 1)
|
||||
self.assertEqual(ret_filter[1], None)
|
||||
|
||||
def test_kdtree_grid_filter_pairs_2d(self):
|
||||
self._test_kdtree_grid_filter_pairs_test_impl(dimensions=2)
|
||||
|
||||
def test_kdtree_grid_filter_pairs_3d(self):
|
||||
self._test_kdtree_grid_filter_pairs_test_impl(dimensions=3)
|
||||
|
||||
def test_kdtree_invalid_size(self):
|
||||
with self.assertRaises(ValueError):
|
||||
kdtree.KDTree(-1)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue