kdtree/python_wrapper.cc at master · hpenedones/kdtree · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <array>
#include <list>

#include "kdtree.h"
#include "point.h"

namespace py = pybind11;

// Bind Point class for specific dimensions
template <int N>
void bind_point(py::module& m, const std::string& name) {
    auto point_class =
        py::class_<Point<N>>(m, name.c_str())
            .def(py::init<int, const std::array<float, N>&>(), py::arg("id"), py::arg("coords"))
            .def("id", &Point<N>::id)
            .def("__getitem__",
                 [](const Point<N>& p, int axis) {
                     int idx = axis;
                     if (idx < 0) {
                         idx += N;  // Support Python-style negative indices
                     }
                     if (idx < 0 || idx >= N) {
                         throw py::index_error("Point index out of range");
                     }
                     return p[idx];
                 })
            .def_static("dimensions", &Point<N>::dimensions)
            .def("__repr__", [](const Point<N>& p) {
                std::string repr =
                    "Point" + std::to_string(N) + "D(id=" + std::to_string(p.id()) + ", coords=[";
                for (int i = 0; i < N; i++) {
                    if (i > 0)
                        repr += ", ";
                    repr += std::to_string(p[i]);
                }
                repr += "])";
                return repr;
            });

    // Bind dimension-specific accessors only when valid
    if (N >= 1) {
        point_class.def("x", &Point<N>::x);
    }
    if (N >= 2) {
        point_class.def("y", &Point<N>::y);
    }
    if (N >= 3) {
        point_class.def("z", &Point<N>::z);
    }
}

// Bind Kdtree class for specific dimensions
template <int N>
void bind_kdtree(py::module& m, const std::string& name) {
    py::class_<Kdtree<N>>(m, name.c_str())
        .def(py::init<const Point<N>&, int>(), py::arg("point"), py::arg("split_axis") = 0)
        .def("insert", &Kdtree<N>::insert, py::arg("point"))
        .def("get_nearby_points", &Kdtree<N>::get_nearby_points, py::arg("point"),
             py::arg("radius"));
}

PYBIND11_MODULE(pykdtree, m) {
    m.doc() = "Python bindings for N-dimensional k-d tree library";

    // Bind 2D types
    bind_point<2>(m, "Point2D");
    bind_kdtree<2>(m, "Kdtree2D");

    // Bind 3D types
    bind_point<3>(m, "Point3D");
    bind_kdtree<3>(m, "Kdtree3D");

    // Bind 4D types
    bind_point<4>(m, "Point4D");
    bind_kdtree<4>(m, "Kdtree4D");
}