Merge pull request #2 from jlvdb/nearest-neighbours

Added fixed-number nearest neighbour search

Author: jlvdb

README.rst

@@ -15,12 +15,9 @@
 
 A fast ball tree implementation for three dimensional (weighted) data with an
 Euclidean distance norm. The base implementation is in `C` and there is a
-wrapper for `Python`.
-
-The tree is optimised towards spatial correlation function calculations since
-the query routines are geared towards range queries, i.e. counting pairs with a
-given (range of) separations. Fixed number nearest neighbour search is currently
-not implemented.
+wrapper for `Python`. The tree is optimised towards spatial correlation function
+calculations since it provides fast counting routinge, e.g. by implementing a
+dualtree query algorithm.
 
 - Code: https://github.com/jlvdb/balltree.git
 - Docs: https://balltree.readthedocs.io/

balltree/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.0.1"
+__version__ = "1.1.0"
 
 from .angulartree import AngularTree
 from .balltree import BallTree, default_leafsize as _df

balltree/angulartree.py

@@ -51,10 +51,10 @@ def from_random(
         Build a new AngularTree instance from randomly generated points.
         
         The (ra, dec) coordinates are generated uniformly in the interval
-        [`ra_min`, `ra_max`) and [`dec_min`, `dec_max`), respectively. `size`
-        controlls the number of points generated. The optional `leafsize`
-        determines when the tree query algorithms switch from traversal to brute
-        force.
+        [``ra_min``, ``ra_max``) and [``dec_min``, ``dec_max``), respectively.
+        ``size`` controlls the number of points generated. The optional
+        ``leafsize`` determines when the tree query algorithms switch from
+        traversal to brute force.
         """
         ((x_min, y_min),) = coord.angular_to_cylinder([ra_min, dec_min])
         ((x_max, y_max),) = coord.angular_to_cylinder([ra_max, dec_max])
@@ -74,11 +74,12 @@ def data(self) -> NDArray:
             np.transpose([data["x"], data["y"], data["z"]])
         )
 
-        dtype = [("ra", "f8"), ("dec", "f8"), ("weight", "f8")]
+        dtype = [("ra", "f8"), ("dec", "f8"), ("weight", "f8"), ("index", "i8")]
         array = np.empty(len(data), dtype=dtype)
         array["ra"] = radec[:, 0]
         array["dec"] = radec[:, 1]
         array["weight"] = data["weight"]
+        array["index"] = data["index"]
         return array
 
     @property
@@ -102,7 +103,7 @@ def radius(self) -> float:
         center = coord.angular_to_euclidean(self.center)[0]
         radec_flat = self.data.view("f8")
         shape = (self.num_points, -1)
-        xyz = coord.angular_to_euclidean(radec_flat.reshape(shape)[:, :-1])
+        xyz = coord.angular_to_euclidean(radec_flat.reshape(shape)[:, :-2])
         # compute the maximum distance from the center project one the sphere
         diff = xyz - center[np.newaxis, :]
         dist = np.sqrt(np.sum(diff**2, axis=1))
@@ -120,15 +121,45 @@ def get_node_data() -> NDArray:
         """
         Collect the meta data of all tree nodes in a numpy array.
 
-        The array fields record `depth` (starting from the root node),
-        `num_points`, `sum_weight`, `x`, `y`, `z` (node center) and node `radius`.
+        The array fields record ``depth`` (starting from the root node),
+        ``num_points``, ``sum_weight``, ``x``, ``y``, ``z`` (node center) and
+        node ``radius``.
 
         .. Note::
             The node coordinates and radius are currently not converted to
             anlges.
         """
         return super().get_node_data()
 
+    def nearest_neighbours(self, radec, k, max_ang=-1.0) -> NDArray:
+        """
+        Query a fixed number of nearest neighbours.
+
+        The query point(s) ``radec`` can be a numpy array of shape (2,) or (N, 2),
+        or an equivalent python object. The number of neighbours ``k`` must be a
+        positive integer and the optional ``max_ang`` parameter puts an upper
+        bound on the angular separation (in radian) to the neighbours.
+
+        Returns an array with fields ``index``, holding the index to the neighbour
+        in the array from which the tree was constructed, and ``angle``, the
+        angular separation in radian. The result is sorted by separation,
+        missing neighbours (e.g. if ``angle > max_ang``) are indicated by an
+        index of -1 and infinite separation.
+        """
+        xyz = coord.angular_to_euclidean(radec)
+        if max_ang > 0:
+            max_dist = coord.angle_to_chorddist(max_ang)
+        else:
+            max_dist = -1.0
+        raw = super().nearest_neighbours(xyz, k, max_dist=max_dist)
+        good = raw["index"] >= 0
+
+        result = np.empty(raw.shape, dtype=[("index", "i8"), ("angle", "f8")])
+        result["index"] = raw["index"]
+        result["angle"][~good] = np.inf
+        result["angle"][good] = coord.chorddist_to_angle(raw["distance"][good])
+        return result
+
     def brute_radius(
         self,
         radec: ArrayLike,

balltree/balltree.c

@@ -4,6 +4,7 @@
 
 #include <stdint.h>
 #include <stdio.h>
+#include <string.h>
 
 #include "point.h"
 #include "balltree.h"
@@ -65,6 +66,7 @@ void inputiterdata_free(InputIterData *);
 static PointBuffer *ptbuf_from_PyObjects(PyObject *xyz_obj, PyObject *weight_obj);
 static PyObject *ptbuf_get_numpy_view(PointBuffer *buffer);
 static PyObject *statvec_get_numpy_array(StatsVector *vec);
+static PyObject *queueitems_get_numpy_array(QueueItem *items, npy_intp size, npy_intp n_items);
 static DistHistogram *disthistogram_from_PyObject(PyObject *edges_obj);
 static PyObject *PyObject_from_disthistogram(DistHistogram *hist);
 static PyObject *PyBallTree_accumulate_radius(PyBallTree *self, count_radius_func accumulator, PyObject *xyz_obj, double radius, PyObject *weight_obj);
@@ -87,6 +89,7 @@ static PyObject *PyBallTree_get_radius(PyBallTree *self, void *closure);
 static PyObject *PyBallTree_str(PyObject *self);
 static PyObject *PyBallTree_to_file(PyBallTree *self, PyObject *args);
 static PyObject *PyBallTree_count_nodes(PyBallTree *self);
+static PyObject *PyBallTree_nearest_neighbours(PyBallTree *self, PyObject *args, PyObject *kwargs);
 static PyObject *PyBallTree_get_node_data(PyBallTree *self);
 static PyObject *PyBallTree_brute_radius(PyBallTree *self, PyObject *args, PyObject *kwargs);
 static PyObject *PyBallTree_count_radius(PyBallTree *self, PyObject *args, PyObject *kwargs);
@@ -230,6 +233,7 @@ PyArrayObject *numpy_array_add_dim(PyArrayObject* array) {
     if (reshaped == NULL) {
         PyErr_SetString(PyExc_MemoryError, "failed to reshape array");
     }
+    Py_DECREF(array);
     return reshaped;
 }
 
@@ -279,7 +283,6 @@ static PyArrayObject *xyz_ensure_2dim_double(PyObject *xyz_obj) {
     npy_int ndim = PyArray_NDIM(xyz_arr);
     if (ndim == 1) {
         xyz_arr_2dim = numpy_array_add_dim(xyz_arr);
-        Py_DECREF(xyz_arr);
         if (xyz_arr_2dim == NULL) {
             return NULL;
         }
@@ -393,7 +396,12 @@ static PointBuffer *ptbuf_from_PyObjects(PyObject *xyz_obj, PyObject *weight_obj
     int64_t idx = 0;
     double x, y, z;
     while (iter_get_next_xyz(data->xyz_iter, &x, &y, &z)) {
-        buffer->points[idx] = (Point){x, y, z, data->weight_buffer[idx]};
+        Point *point = buffer->points + idx;
+        point->x = x;
+        point->y = y;
+        point->z = z;
+        point->weight = data->weight_buffer[idx];
+        // index is already initialised
         ++idx;
     }
     inputiterdata_free(data);
@@ -406,11 +414,12 @@ static PyObject *ptbuf_get_numpy_view(PointBuffer *buffer) {
 
     // construct an appropriate dtype for Point
     PyObject *arr_dtype = Py_BuildValue(
-        "[(ss)(ss)(ss)(ss)]",
+        "[(ss)(ss)(ss)(ss)(ss)]",
         "x", "f8",
         "y", "f8",
         "z", "f8",
-        "weight", "f8"
+        "weight", "f8",
+        "index", "i8"
     );
     if (arr_dtype == NULL) {
         return NULL;
@@ -494,15 +503,48 @@ static PyObject *statvec_get_numpy_array(StatsVector *vec) {
 
     // create an uninitialised array and copy the data into it
     PyObject *array = PyArray_Empty(ndim, shape, arr_descr, 0);
-    Py_DECREF(arr_descr);
     if (array == NULL) {
+        Py_DECREF(arr_descr);
         return NULL;
     }
     void *ptr = PyArray_DATA(array);
     memcpy(ptr, vec->stats, sizeof(NodeStats) * vec->size);
     return array;
 }
 
+static PyObject *queueitems_get_numpy_array(QueueItem *items, npy_intp size, npy_intp n_items) {
+    const npy_intp ndim = 2;
+    npy_intp shape[2] = {size, n_items};
+
+    // construct an appropriate dtype for QueueItem buffer
+    PyObject *arr_dtype = Py_BuildValue(
+        "[(ss)(ss)]",
+        "index", "i8",
+        "distance", "f8"
+    );
+    if (arr_dtype == NULL) {
+        return NULL;
+    }
+
+    // get the numpy API array descriptor
+    PyArray_Descr *arr_descr;
+    int result = PyArray_DescrConverter(arr_dtype, &arr_descr);  // PyArray_Descr **
+    Py_DECREF(arr_dtype);
+    if (result != NPY_SUCCEED) {
+        return NULL;
+    }
+
+    // create an uninitialised array and copy the data into it
+    PyObject *array = PyArray_Empty(ndim, shape, arr_descr, 0);
+    if (array == NULL) {
+        Py_DECREF(arr_descr);
+        return NULL;
+    }
+    void *ptr = PyArray_DATA(array);
+    memcpy(ptr, items, sizeof(QueueItem) * size * n_items);
+    return array;
+}
+
 static PyObject *PyBallTree_accumulate_radius(
     PyBallTree *self,
     count_radius_func accumulator,
@@ -517,7 +559,7 @@ static PyObject *PyBallTree_accumulate_radius(
     // count neighbours for all inputs
     double count = 0.0;
     int64_t idx = 0;
-    Point point;
+    Point point = {0.0, 0.0, 0.0, 0.0, 0};
     while (iter_get_next_xyz(data->xyz_iter, &point.x, &point.y, &point.z)) {
         point.weight = data->weight_buffer[idx];
         count += accumulator(self->balltree, &point, radius);
@@ -545,7 +587,7 @@ static PyObject *PyBallTree_accumulate_range(
     }
     // count neighbours for all inputs
     int64_t idx = 0;
-    Point point;
+    Point point = {0.0, 0.0, 0.0, 0};
     while (iter_get_next_xyz(data->xyz_iter, &point.x, &point.y, &point.z)) {
         point.weight = data->weight_buffer[idx];
         accumulator(self->balltree, &point, hist);
@@ -613,8 +655,8 @@ PyDoc_STRVAR(
     "--\n\n"
     "Build a new BallTree instance from randomly generated points.\n\n"
     "The (x, y, z) coordinates are generated uniformly in the interval\n"
-    "[`low`, `high`), `size` controlls the number of points generated. The\n"
-    "optional `leafsize` determines when the tree query algorithms switch from\n"
+    "[``low``, ``high``), ``size`` controlls the number of points generated. The\n"
+    "optional ``leafsize`` determines when the tree query algorithms switch from\n"
     "traversal to brute force."
 );
 
@@ -794,7 +836,7 @@ static PyObject *PyBallTree_str(PyObject *self) {
     int n_bytes = snprintf(
         buffer,
         sizeof(buffer),
-        "BallTree(num_points=%ld, radius=%lf, center=(%lf, %lf, %lf))",
+        "BallTree(num_points=%lld, radius=%lf, center=(%lf, %lf, %lf))",
         tree->data->size,
         node->ball.radius,
         node->ball.x,
@@ -858,8 +900,9 @@ PyDoc_STRVAR(
     "get_node_data(self) -> NDArray\n"
     "--\n\n"
     "Collect the meta data of all tree nodes in a numpy array.\n\n"
-    "The array fields record `depth` (starting from the root node),\n"
-    "`num_points`, `sum_weight`, `x`, `y`, `z` (node center) and node `radius`."
+    "The array fields record ``depth`` (starting from the root node),\n"
+    "``num_points``, ``sum_weight``, ``x``, ``y``, ``z`` (node center) and node\n"
+    "``radius``."
 );
 
 static PyObject *PyBallTree_get_node_data(PyBallTree *self) {
@@ -873,6 +916,87 @@ static PyObject *PyBallTree_get_node_data(PyBallTree *self) {
     return array;
 }
 
+PyDoc_STRVAR(
+    // .. py:method::
+    nearest_neighbours_doc,
+    "nearest_neighbours(self, xyz: ArrayLike, k: int, max_dist: float = -1.0) -> NDArray\n"
+    "--\n\n"
+    "Query a fixed number of nearest neighbours.\n\n"
+    "The query point(s) ``xyz`` can be a numpy array of shape (3,) or (N, 3),\n"
+    "or an equivalent python object. The number of neighbours ``k`` must be a\n"
+    "positive integer and the optional ``max_dist`` parameter puts an upper\n"
+    "bound on the separation to the neighbours.\n\n"
+    "Returns an array with fields ``index``, holding the index to the neighbour\n"
+    "in the array from which the tree was constructed, and ``distance``, the\n"
+    "separation. The result is sorted by separation, missing neighbours (e.g. if\n"
+    "``distance > max_dist``) are indicated by an index of -1 and infinite\n"
+    "separation.\n"
+);
+
+static PyObject *PyBallTree_nearest_neighbours(
+    PyBallTree *self,
+    PyObject *args,
+    PyObject *kwargs
+) {
+    static char *kwlist[] = {"xyz", "k", "max_dist", NULL};
+    PyObject *xyz_obj;
+    long num_neighbours;  // screw Windows
+    double max_dist = -1.0;
+    if (!PyArg_ParseTupleAndKeywords(
+            args, kwargs, "Ol|d", kwlist,
+            &xyz_obj, &num_neighbours, &max_dist)
+    ) {
+        return NULL;
+    }
+    if (num_neighbours < 1) {
+        PyErr_SetString(PyExc_ValueError, "number of neighbours must be positive");
+        return NULL;
+    }
+
+    InputIterData *data = inputiterdata_new(xyz_obj, Py_None);
+    if (data == NULL) {
+        return NULL;
+    }
+
+    // allocate output buffer
+    size_t n_bytes_queue = num_neighbours * sizeof(QueueItem);
+    QueueItem *result = malloc(data->size * n_bytes_queue);
+    if (result == NULL) {
+        PyErr_SetString(PyExc_MemoryError, "failed to allocate output array");
+        inputiterdata_free(data);
+        return NULL;
+    }
+
+    // find neighbours for all inputs
+    KnnQueue *queue = NULL;
+    PyObject *pyresult = NULL;
+    int64_t idx = 0;
+    Point point = {0.0, 0.0, 0.0, 0.0, 0};
+    while (iter_get_next_xyz(data->xyz_iter, &point.x, &point.y, &point.z)) {
+        queue = balltree_nearest_neighbours(
+            self->balltree,
+            &point,
+            num_neighbours,
+            max_dist
+        );
+        if (queue == NULL) {
+            printf("oops\n");
+            goto error;
+        }
+        // copy result into output buffer
+        memcpy(&result[idx], queue->items, n_bytes_queue);
+        knque_free(queue);
+        idx += num_neighbours;
+    }
+
+    // convert to numpy array
+    pyresult = queueitems_get_numpy_array(result, data->size, num_neighbours);
+error:
+    free(result);
+    inputiterdata_free(data);
+    return pyresult;
+}
+
 PyDoc_STRVAR(
     // .. py:method::
     brute_radius_doc,
@@ -1139,6 +1263,12 @@ static PyMethodDef PyBallTree_methods[] = {
         .ml_flags = METH_NOARGS,
         .ml_doc = get_node_data_doc
     },
+    {
+        .ml_name = "nearest_neighbours",
+        .ml_meth = (PyCFunctionWithKeywords)PyBallTree_nearest_neighbours,
+        .ml_flags = METH_VARARGS | METH_KEYWORDS,
+        .ml_doc = nearest_neighbours_doc
+    },
     {
         .ml_name = "brute_radius",
         .ml_meth = (PyCFunctionWithKeywords)PyBallTree_brute_radius,

include/ballnode.h

@@ -5,6 +5,7 @@
 
 #include "point.h"
 #include "histogram.h"
+#include "queue.h"
 
 #define BALLNODE_IS_LEAF(node) (node)->is_leaf
 
@@ -57,6 +58,7 @@ int bnode_is_leaf(const BallNode *);
 PointSlice bnode_get_ptslc(const BallNode *);
 
 // from ballnode_query.c
+void bnode_nearest_neighbours(const BallNode *, const Point *, KnnQueue *);
 double bnode_count_radius(const BallNode *, const Point *, double);
 void bnode_count_range(const BallNode *, const Point *, DistHistogram *);
 double bnode_dualcount_radius(const BallNode *, const BallNode *, double);