spatial-graph 0.0.1__py3-none-any.whl

spatial_graph/rtree/line_rtree.py

@@ -0,0 +1,156 @@
+from .rtree import RTree
+import numpy as np
+
+
+class LineRTree(RTree):
+    pyx_item_t_declaration = """
+    cdef struct item_t:
+        item_base_t u
+        item_base_t v
+        bool corner_mask[DIMS]
+"""
+
+    c_item_t_declaration = """
+typedef struct item_t {
+    item_base_t u;
+    item_base_t v;
+    bool corner_mask[DIMS];
+} item_t;
+"""
+
+    c_converter_functions = """
+inline item_t convert_pyx_to_c_item(pyx_item_t *pyx_item, coord_t *start, coord_t *end) {
+    item_t item;
+    coord_t tmp;
+    item.u = (*pyx_item)[0];
+    item.v = (*pyx_item)[1];
+    for (int d = 0; d < DIMS; d++) {
+        item.corner_mask[d] = (start[d] < end[d]);
+        if (!item.corner_mask[d]) {
+            // swap coordinates to create bounding box
+            tmp = start[d];
+            start[d] = end[d];
+            end[d] = tmp;
+        }
+    }
+    return item;
+}
+inline void copy_c_to_pyx_item(const item_t c_item, pyx_item_t *pyx_item) {
+    (*pyx_item)[0] = c_item.u;
+    (*pyx_item)[1] = c_item.v;
+}
+"""
+
+    c_equal_function = """
+inline int equal(const item_t a, const item_t b) {
+    return (a.u == b.u && a.v == b.v);
+}
+"""
+
+    c_distance_function = """
+inline coord_t length2(const coord_t x[]) {
+    coord_t length2 = 0;
+    for (int d = 0; d < DIMS; d++) {
+        length2 += pow(x[d], 2);
+    }
+    return length2;
+}
+
+inline coord_t point_segment_dist2(const coord_t point[], const coord_t start[], const coord_t end[]) {
+
+    coord_t a[DIMS];
+    coord_t b[DIMS];
+    coord_t alpha = 0;
+
+    for (int d = 0; d < DIMS; d++) {
+
+        // subtract "start" from "end" and "point" to get "a" and "b"
+        a[d] = end[d] - start[d];
+        b[d] = point[d] - start[d];
+
+        // compute dot product "alpha" of "a" and "b"
+        alpha += a[d] * b[d];
+    }
+
+    // normalize dot product
+    alpha /= length2(a);
+
+    // clip at 0 and 1 (beginning and end of line segment)
+    alpha = min0(1, max0(0, alpha));
+
+    for (int d = 0; d < DIMS; d++) {
+
+        // multiply "a" by "alpha" to obtain closest segment point to "b"
+        a[d] *= alpha;
+
+        // subtract "b" from "a" to get offset
+        a[d] -= b[d];
+    }
+
+    // compute squared length of offset
+    return length2(a);
+}
+
+inline coord_t distance(const coord_t point[], const struct rect *rect, const struct item_t item) {
+    coord_t start[DIMS];
+    coord_t end[DIMS];
+    for (int d = 0; d < DIMS; d++) {
+        if (item.corner_mask[d]) {
+            start[d] = rect->min[d];
+            end[d] = rect->max[d];
+        } else {
+            start[d] = rect->max[d];
+            end[d] = rect->min[d];
+        }
+    }
+    return point_segment_dist2(point, start, end);
+}
+"""
+
+    def insert_line(self, line, start, end):
+        """Convenience function to insert a single line. To insert multiple
+        lines in bulk, please use the faster `insert_lines`.
+
+        Args:
+
+            line (`item_dtype`):
+
+                The line identifier (as passed as the `item_dtype` to the
+                constructor).
+
+            start (`ndarray`, shape `(d,)`):
+
+                The coordinates of the start of the line.
+
+            end (`ndarray`, shape `(d,)`):
+
+                The coordinates of the end of the line.
+        """
+        lines = np.array([line], dtype=self.item_dtype.base)
+        starts = start[np.newaxis]
+        ends = end[np.newaxis]
+        return self.insert_bb_items(lines, starts, ends)
+
+    def insert_lines(self, lines, starts, ends):
+        """Insert a list of lines.
+
+        Args:
+
+            lines (`ndarray`, shape `(n, [m])`):
+
+                Array containing the line identifiers (as passed as the
+                `item_dtype` to the constructor). If the identifiers are an
+                array of size `m`, the expected shape is `(n, m)` where `n` is
+                the number of lines, otherwise the shape is just `(n,)`.
+
+            starts (`ndarray`, shape `(n, d)`):
+
+                The coordinates of the start of each line.
+
+            ends (`ndarray`, shape `(n, d)`):
+
+                The coordinates of the end of each line.
+        """
+        # we just forward to bb insert, "start" and "end" will be used to compute
+        # the bounding box in our custom converter above
+        return self.insert_bb_items(lines, starts, ends)

spatial_graph/rtree/point_rtree.py

@@ -0,0 +1,6 @@
+from .rtree import RTree
+
+
+# no modifications to plain RTree needed
+class PointRTree(RTree):
+    pass

spatial_graph/rtree/rtree.py

@@ -0,0 +1,125 @@
+from typing import ClassVar
+import sys
+import witty
+import numpy as np
+from Cheetah.Template import Template
+from pathlib import Path
+from ..dtypes import DType
+
+DEFINE_MACROS = [("RTREE_NOATOMICS", "1")] if sys.platform == "win32" else []
+
+
+class RTree:
+    """A generic RTree implementation, compiled on-the-fly during
+    instantiation.
+
+    Args:
+
+        item_dtype (``string``):
+
+            The C type of the items to hold. Can be a scalar (e.g. ``uint64``)
+            or an array of scalars (e.g., "uint64[3]").
+
+        coord_dtype (``string``):
+
+            The scalar C type to use for coordinates (e.g., ``float``).
+
+        dims (``int``):
+
+            The dimension of the r-tree.
+
+    Subclassing:
+
+        This generic implementation can be subclassed and modified in the
+        following ways:
+
+        The class members ``pyx_item_t_declaration`` and
+        ``c_item_t_declaration`` can be overwritten to use custom ``item_t``
+        structures. This will also require overwriting the
+        ``c_converter_functions`` to translate between the PYX interface (where
+        items are scalars or C arrays of scalars) and the C interface (the
+        custom ``item_t`` type.
+
+        The following constants and typedefs are available to use in the
+        provided code:
+
+            DIMS:
+
+                A constant set to the value of ``dims``.
+
+            item_base_t:
+
+                The scalar type of the item (e.g., ``uint64``), regardless of
+                whether this is a scalar or array item.
+    """
+
+    # overwrite in subclasses for custom item_t structures
+    pyx_item_t_declaration: ClassVar[str] = ""
+    c_item_t_declaration: ClassVar[str] = ""
+
+    # overwrite in subclasses for custom converters
+    c_converter_functions: ClassVar[str] = ""
+
+    # overwrite in subclasses for custom item comparison code
+    c_equal_function: ClassVar[str] = ""
+
+    # overwrite in subclasses for custom distance computation
+    c_distance_function: ClassVar[str] = ""
+
+    def __new__(
+        cls,
+        item_dtype,
+        coord_dtype,
+        dims,
+    ):
+        item_dtype = DType(item_dtype)
+        coord_dtype = DType(coord_dtype)
+
+        ############################################
+        # create wrapper from template and compile #
+        ############################################
+
+        src_dir = Path(__file__).parent
+        wrapper_template = Template(
+            file=str(src_dir / "wrapper_template.pyx"),
+            compilerSettings={"directiveStartToken": "%"},
+        )
+        wrapper_template.item_dtype = item_dtype
+        wrapper_template.coord_dtype = coord_dtype
+        wrapper_template.dims = dims
+        wrapper_template.c_distance_function = cls.c_distance_function
+        wrapper_template.pyx_item_t_declaration = cls.pyx_item_t_declaration
+        wrapper_template.c_item_t_declaration = cls.c_item_t_declaration
+        wrapper_template.c_converter_functions = cls.c_converter_functions
+        wrapper_template.c_equal_function = cls.c_equal_function
+
+        wrapper = witty.compile_module(
+            str(wrapper_template),
+            source_files=[
+                src_dir / "src" / "rtree.h",
+                src_dir / "src" / "rtree.c",
+                src_dir / "src" / "config.h",
+            ],
+            extra_compile_args=["/O2" if sys.platform == "win32" else "-O3"],
+            include_dirs=[str(src_dir)],
+            language="c",
+            quiet=True,
+            define_macros=DEFINE_MACROS,
+        )
+        RTreeType = type(cls.__name__, (cls, wrapper.RTree), {})
+        return wrapper.RTree.__new__(RTreeType)
+
+    def __init__(self, item_dtype, coord_dtype, dims):
+        super().__init__()
+        self.item_dtype = DType(item_dtype)
+
+    def insert_point_item(self, item, position):
+        items = np.array([item], dtype=self.item_dtype.base)
+        positions = position[np.newaxis]
+        return self.insert_point_items(items, positions)
+
+    def delete_item(self, item, bb_min, bb_max=None):
+        items = np.array([item], dtype=self.item_dtype.base)
+        bb_mins = bb_min[np.newaxis, :]
+        bb_maxs = None if bb_max is None else bb_max[np.newaxis, :]
+        return self.delete_items(items, bb_mins, bb_maxs)

spatial_graph/rtree/src/ARCHITECTURE.md

@@ -0,0 +1,162 @@
+History
+=======
+
+[Implementation from Joshua J Baker](https://github.com/tidwall/rtree.c), under
+an MIT-style license (see LICENSE).
+
+* modified to work with [witty](https://github.com/funkelab/witty)
+* added kNN search
+
+Data Structures
+===============
+
+R-Tree Itself
+-------------
+
+`rect`
+`item`
+    * `data` = node ID (type `item_data_t` defined in wrapper)
+`node`
+    * `rects`: list of `rect`s
+    * `kind`:
+        * `LEAF`: contains a list of `item`s, one for each `rect`
+        * `BRANCH`: contains a list of `node`s, one for each `rect`
+
+Priority Queue for kNN
+----------------------
+
+`element`
+    * `distance`
+    * `kind` (same enum as rtree kind):
+        * `LEAF`: contains a `node*`
+        * `BRANCH`: contains a `node*`
+        * `ITEM`: contains an `item`
+`priority_queue`
+    * dynamically allocated array of `element`s
+
+Implementation for Line Segments
+================================
+
+`item`:
+    * `u`, `v`: IDs of endpoint nodes
+    * `corner_mask`: (n-1)D bitmask for line segment corner in bounding box
+
+`distance(coord_t point[], rect *bb, item *segment)`
+    * compute distance of `point` to line `segment` bound by `bb`
+`distance_bb(coord_t point[], rect *bb)`
+    * previous implementation, computes distance to bounding box
+
+`element`
+    * `kind`
+        * `LEAF`, `BRANCH` as before
+        * `ITEM_BY_BB`: item, distance computed using bounding box
+        * `ITEM`: item, distance computed exactly
+
+
+Generalize for Arbitrary Items
+==============================
+
+User-provided "templates":
+
+Level 0 (C)
+-----------
+
+`DIMS`
+    * a `#define`, the number of dimensions
+
+`KNN_USE_EXACT_DISTANCE`
+    * an optional `#define`
+    * if defined, the distance function below needs to be provided
+    * if not defined, only distances to bounding boxes are used for the kNN search
+
+`coord_t`
+    * the scalar type of coordinates
+
+`item_t`
+    * a type (can be a `struct`) representing an item
+    * keep small (e.g., only the ID of a node)
+
+`distance(coord_t point[], rect *bb, item_t *item)`
+    * a function that computes the exact distance of a point to the item
+    * `point`: the query point
+    * `bb`: the bounding box of the item (as given during insert)
+    * `item`: the item to compute the distance for
+    * if not given, should not attempt to compute exact distance and return
+      item based on `distance_bb` right away
+
+Level 1 (PYX)
+-------------
+
+This level holds glue code to bridge between C and Python.
+
+
+### Definitions
+
+`NUM_DIMS`
+    * the number of dimensions
+    * sets definition of `DIMS` used in Level 0
+
+`KNN_USE_EXACT_DISTANCE`
+    * as described above
+
+### Data Types
+
+`coord_t`
+    * needs to be provided as:
+        * C code: `typedef`
+        * PYX code: `ctypedef`
+
+`item_t`
+    * needs to be provided as
+        * C code: this is the `item_t` used in Level 0 (`typedef` or `struct`)
+        * PYX code: this is just to wrap `item_t` to be used in PYX code (not user facing)
+          ⇒ TODO: is that really needed?
+
+`pyx_item_t`
+    * semantically, the Python-wrapped version of `item_t` (user facing for
+      insert and query)
+    * has to be something that can be placed in a numpy array / memory view
+        * this is part of the interface for efficient insert and query
+        * similar to attributes, nodes, and edge lists in the graph, the
+          interface uses numpy arrays
+    * this means it is either a numpy type or an array of that, i.e.:
+        * dtype    -> dtype[::1] for efficient insert/query
+        * dtype[n] -> dtype[:, ::1] for efficient insert/query (where n is the size of the item)
+    * needs to be provided as:
+        * C code: `typedef`
+        * PYX code: `ctypedef`
+
+`NP_ITEM_DTYPE`
+    * the `dtype` of `pyx_item_t` (e.g., `uint64`), used to create numpy arrays
+
+`API_ITEMS_MEMVIEW_TYPE`
+    * the type of a list of `pyx_item_t`s, either
+        * dtype[::1] or
+        * dtype[:, ::1]
+
+### Strings
+
+`ITEM_LENGTH`
+    * `d` if an array type, empty string otherwise
+
+`ITEMS_EXTRA_DIMS_0`
+    * extra indices to index an element: `item[i, 0]` would be ", 0"
+
+### Functions
+
+`item_t pyx_to_c_item(pyx_item_t pyx_item, *)`
+    * needs to be provided as:
+        * C code: function implementation
+        * PYX code: `cdef` function declaration
+    ⇒ TODO: can be PYX function only?
+
+`void copy_c_to_pyx_item(const item_t item, pyx_item_t *pyx_item)`
+    * copy and convert C-item `item` to address of `pyx_item`
+    * needs to be provided as:
+        * C code: function implementation
+        * PYX code: `cdef` function declaration
+    ⇒ TODO: can be PYX function only?
+
+`distance(coord_t point[], rect *bb, item_t *item)`
+    * only C implementation needed
+    * skip if `KNN_USE_EXACT_DISTANCE` is not set

spatial_graph/rtree/src/LICENSE

@@ -0,0 +1,20 @@
+The MIT License (MIT)
+
+Copyright (c) 2020 Joshua J Baker
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software is furnished to do so,
+subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

spatial_graph/rtree/src/config.h

@@ -0,0 +1,13 @@
+// rtree configuration header
+
+#ifndef __RTREE_CONFIG_H
+#define __RTREE_CONFIG_H
+
+// the maximal number of items per node in the rtree
+#define MAXITEMS 64
+
+// the initial number of items in the priority queue to find nearest neighbors
+// (will keep doubling as needed to accommodate more)
+#define INITIAL_QUEUE_SIZE 256
+
+#endif // __RTREE_CONFIG_H