christiangeorgelucas/spatial-tools

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import BoundingBoxInput, BoundingBoxResult
from nodes._common import NodeInputError, array_to_vector, check_matrix_shape, err, matrix_to_array

def bounding_box(ax: AxiomContext, input: BoundingBoxInput) -> BoundingBoxResult:
    """Compute the axis-aligned bounding box of a point set: the
    per-dimension minimum and maximum across all rows of `points`, via
    numpy.min/numpy.max.
    """
    try:
        check_matrix_shape(input.points, name="points")
        points = matrix_to_array(input.points)

        return BoundingBoxResult(
            min=array_to_vector(points.min(axis=0)), max=array_to_vector(points.max(axis=0))
        )
    except NodeInputError as e:
        return BoundingBoxResult(error=e.error)
    except Exception as e:
        return BoundingBoxResult(error=err("BOUNDING_BOX_FAILED", str(e)))

from scipy.spatial import ConvexHull as ScipyConvexHull
from scipy.spatial import QhullError

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import ConvexHullInput, ConvexHullResult
from nodes._common import (
    NodeInputError,
    array_to_int_matrix,
    array_to_int_vector,
    check_matrix_shape,
    err,
    matrix_to_array,
)

MIN_DIM = 2


def convex_hull(ax: AxiomContext, input: ConvexHullInput) -> ConvexHullResult:
    """Compute the convex hull of a general N-dimensional point set (D >= 2)
    via scipy.spatial.ConvexHull (Qhull). Returns the hull's
    extreme-point (vertex) indices, its facet simplices, its D-dimensional
    volume (area, for D=2), and its (D-1)-dimensional surface measure
    (perimeter, for D=2). Fewer than D+1 points returns
    INSUFFICIENT_POINTS; coplanar/collinear (rank-deficient) points return
    DEGENERATE_INPUT.
    """
    try:
        check_matrix_shape(input.points)
        points = matrix_to_array(input.points)
        n, d = points.shape
        if d < MIN_DIM:
            raise NodeInputError("INVALID_ARGUMENT", f"points must have D >= {MIN_DIM}, got D={d}")
        if n < d + 1:
            raise NodeInputError(
                "INSUFFICIENT_POINTS", f"need at least {d + 1} points for a {d}D hull, got {n}"
            )

        try:
            hull = ScipyConvexHull(points)
        except QhullError as e:
            raise NodeInputError(
                "DEGENERATE_INPUT", f"points are degenerate (coplanar/collinear, rank < {d}): {e}"
            )

        return ConvexHullResult(
            vertices=array_to_int_vector(hull.vertices),
            simplices=array_to_int_matrix(hull.simplices),
            volume=float(hull.volume),
            area=float(hull.area),
        )
    except NodeInputError as e:
        return ConvexHullResult(error=e.error)
    except Exception as e:
        return ConvexHullResult(error=err("HULL_FAILED", str(e)))

from scipy.spatial.distance import cdist

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import CdistInput, CdistResult
from nodes._common import (
    NodeInputError,
    array_to_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
    resolve_full_metric_kwargs,
)


def cross_distance_matrix(ax: AxiomContext, input: CdistInput) -> CdistResult:
    """Compute the full cross distance matrix between two point sets
    `points_a` (N x D) and `points_b` (M x D), via
    scipy.spatial.distance.cdist. Same metric/p/inv_cov options as
    PairwiseDistanceMatrix. Returns an N x M matrix,
    distances[i][j] = distance(points_a[i], points_b[j]).
    """
    try:
        check_matrix_shape(input.points_a, name="points_a")
        check_matrix_shape(input.points_b, name="points_b")
        a = matrix_to_array(input.points_a, name="points_a")
        b = matrix_to_array(input.points_b, name="points_b")
        if a.shape[1] != b.shape[1]:
            raise NodeInputError(
                "DIMENSION_MISMATCH", f"points_a has {a.shape[1]} columns, points_b has {b.shape[1]}"
            )

        metric, kwargs = resolve_full_metric_kwargs(
            input.metric, input.p, input.inv_cov if input.HasField("inv_cov") else None,
            require_inv_cov_for_mahalanobis=False,
        )

        distances = cdist(a, b, metric=metric, **kwargs)

        return CdistResult(distances=array_to_matrix(distances))
    except NodeInputError as e:
        return CdistResult(error=e.error)
    except Exception as e:
        return CdistResult(error=err("DISTANCE_FAILED", str(e)))

from scipy.spatial import Delaunay
from scipy.spatial import QhullError

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import DelaunayInput, DelaunayResult
from nodes._common import (
    NodeInputError,
    array_to_int_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
)

ALLOWED_DIMS = (2, 3)


def delaunay_triangulate(ax: AxiomContext, input: DelaunayInput) -> DelaunayResult:
    """Compute the Delaunay triangulation of a 2D or 3D point set via
    scipy.spatial.Delaunay (Qhull). Returns each simplex as D+1 point
    indices (a triangle for 2D, a tetrahedron for 3D) plus the per-simplex
    neighbor structure (-1 where a facet has no neighbor, i.e. lies on the
    convex hull boundary). Points must be 2D or 3D; fewer than D+1 points
    returns INSUFFICIENT_POINTS; collinear/coplanar points that admit no
    triangulation return DEGENERATE_INPUT rather than crashing.
    """
    try:
        check_matrix_shape(input.points)
        points = matrix_to_array(input.points)
        n, d = points.shape
        if d not in ALLOWED_DIMS:
            raise NodeInputError("INVALID_ARGUMENT", f"points must be 2D or 3D, got D={d}")
        if n < d + 1:
            raise NodeInputError(
                "INSUFFICIENT_POINTS", f"need at least {d + 1} points for a {d}D triangulation, got {n}"
            )

        try:
            tri = Delaunay(points)
        except QhullError as e:
            raise NodeInputError(
                "DEGENERATE_INPUT",
                f"points are degenerate (collinear/coplanar) for a {d}D triangulation: {e}",
            )

        return DelaunayResult(
            simplices=array_to_int_matrix(tri.simplices),
            neighbors=array_to_int_matrix(tri.neighbors),
        )
    except NodeInputError as e:
        return DelaunayResult(error=e.error)
    except Exception as e:
        return DelaunayResult(error=err("TRIANGULATION_FAILED", str(e)))

import numpy as np
from scipy.spatial import Delaunay, QhullError

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import FindSimplexInput, FindSimplexResult
from nodes._common import (
    NodeInputError,
    array_to_int_vector,
    array_to_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
)

ALLOWED_DIMS = (2, 3)


def find_simplex(ax: AxiomContext, input: FindSimplexInput) -> FindSimplexResult:
    """Locate which Delaunay simplex (of `points`' 2D/3D triangulation, via
    scipy.spatial.Delaunay.find_simplex) each row of `queries` falls into,
    and its barycentric weights within that simplex — set-up for
    barycentric interpolation. Returns -1 for a query point outside the
    triangulation's convex hull (with an all-zero barycentric row).
    """
    try:
        check_matrix_shape(input.points, name="points")
        check_matrix_shape(input.queries, name="queries")
        points = matrix_to_array(input.points)
        queries = matrix_to_array(input.queries)
        n, d = points.shape
        if d not in ALLOWED_DIMS:
            raise NodeInputError("INVALID_ARGUMENT", f"points must be 2D or 3D, got D={d}")
        if points.shape[1] != queries.shape[1]:
            raise NodeInputError(
                "DIMENSION_MISMATCH",
                f"points has {points.shape[1]} columns, queries has {queries.shape[1]}",
            )
        if n < d + 1:
            raise NodeInputError(
                "INSUFFICIENT_POINTS", f"need at least {d + 1} points for a {d}D triangulation, got {n}"
            )

        try:
            tri = Delaunay(points)
        except QhullError as e:
            raise NodeInputError(
                "DEGENERATE_INPUT", f"points are degenerate (collinear/coplanar) for a {d}D triangulation: {e}"
            )

        simplex_idx = tri.find_simplex(queries)
        q = queries.shape[0]
        bary = np.zeros((q, d + 1), dtype=np.float64)

        valid = simplex_idx >= 0
        if np.any(valid):
            valid_simplex = simplex_idx[valid]
            X = tri.transform[valid_simplex, :d, :d]
            Y = queries[valid] - tri.transform[valid_simplex, d, :]
            b = np.einsum("ijk,ik->ij", X, Y)
            last = 1.0 - b.sum(axis=1, keepdims=True)
            bary[valid] = np.hstack([b, last])

        return FindSimplexResult(
            simplex_indices=array_to_int_vector(simplex_idx),
            barycentric_coords=array_to_matrix(bary),
        )
    except NodeInputError as e:
        return FindSimplexResult(error=e.error)
    except Exception as e:
        return FindSimplexResult(error=err("FIND_SIMPLEX_FAILED", str(e)))

from scipy.spatial.distance import directed_hausdorff

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import HausdorffInput, HausdorffResult
from nodes._common import (
    NodeInputError,
    array_to_int_vector,
    check_matrix_shape,
    err,
    matrix_to_array,
)

def hausdorff_distance(ax: AxiomContext, input: HausdorffInput) -> HausdorffResult:
    """Compute the symmetric Hausdorff distance between two point sets
    `points_a` and `points_b`, via
    max(scipy.spatial.distance.directed_hausdorff(A,B), directed_hausdorff(B,A)).
    Returns the distance, which direction achieved it (from_a), and the
    achieving index pair.
    """
    try:
        check_matrix_shape(input.points_a, name="points_a")
        check_matrix_shape(input.points_b, name="points_b")
        a = matrix_to_array(input.points_a, name="points_a")
        b = matrix_to_array(input.points_b, name="points_b")
        if a.shape[1] != b.shape[1]:
            raise NodeInputError(
                "DIMENSION_MISMATCH", f"points_a has {a.shape[1]} columns, points_b has {b.shape[1]}"
            )

        d_ab, i_a, j_b = directed_hausdorff(a, b)
        d_ba, i_b, j_a = directed_hausdorff(b, a)

        if d_ab >= d_ba:
            return HausdorffResult(
                distance=float(d_ab), from_a=True, index_pair=array_to_int_vector([i_a, j_b])
            )
        return HausdorffResult(
            distance=float(d_ba), from_a=False, index_pair=array_to_int_vector([i_b, j_a])
        )
    except NodeInputError as e:
        return HausdorffResult(error=e.error)
    except Exception as e:
        return HausdorffResult(error=err("DISTANCE_FAILED", str(e)))

import numpy as np
from scipy.spatial import cKDTree

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import KNNInput, KNNResult
from nodes._common import (
    NodeInputError,
    array_to_int_matrix,
    array_to_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
    resolve_kdtree_metric,
)


def k_nearest_neighbors(ax: AxiomContext, input: KNNInput) -> KNNResult:
    """Build a KD-tree over `points` (scipy.spatial.cKDTree) and query the
    k nearest neighbors of each row in `queries`. metric is "euclidean"
    (default), "cityblock", or "chebyshev". Returns Q x k index and
    distance matrices in ascending-distance order per row. k is rejected
    with INVALID_ARGUMENT if it exceeds len(points) — every returned row is
    always fully populated with k real neighbors, never padded.
    """
    try:
        check_matrix_shape(input.points, name="points")
        check_matrix_shape(input.queries, name="queries")
        points = matrix_to_array(input.points)
        queries = matrix_to_array(input.queries)
        if points.shape[1] != queries.shape[1]:
            raise NodeInputError(
                "DIMENSION_MISMATCH",
                f"points has {points.shape[1]} columns, queries has {queries.shape[1]}",
            )

        k = input.k
        if k < 1 or k > points.shape[0]:
            raise NodeInputError(
                "INVALID_ARGUMENT",
                f"k must satisfy 1 <= k <= len(points) = {points.shape[0]} (got {k})",
            )

        metric, p = resolve_kdtree_metric(input.metric)

        tree = cKDTree(points)
        dist, idx = tree.query(queries, k=k, p=p)
        dist = np.atleast_2d(dist.reshape(queries.shape[0], k))
        idx = np.atleast_2d(idx.reshape(queries.shape[0], k)).astype(np.int64)
        idx[np.isinf(dist)] = -1

        return KNNResult(indices=array_to_int_matrix(idx), distances=array_to_matrix(dist))
    except NodeInputError as e:
        return KNNResult(error=e.error)
    except Exception as e:
        return KNNResult(error=err("QUERY_FAILED", str(e)))

import numpy as np
from scipy.spatial import cKDTree

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import KNNGraphInput, KNNGraphResult
from nodes._common import (
    NodeInputError,
    array_to_int_matrix,
    array_to_vector,
    check_matrix_shape,
    err,
    matrix_to_array,
)


def knn_graph(ax: AxiomContext, input: KNNGraphInput) -> KNNGraphResult:
    """Build a spatial k-nearest-neighbor GRAPH over `points` (euclidean):
    for every point, find its k nearest OTHER points (self excluded) via
    scipy.spatial.cKDTree, and return the resulting directed edge list
    (i -> j meaning "j is one of i's k nearest neighbors") with per-edge
    distances. Not symmetrized — union/intersect the edges yourself for a
    mutual/undirected graph.
    """
    try:
        check_matrix_shape(input.points, name="points")
        points = matrix_to_array(input.points)
        n = points.shape[0]

        k = input.k
        max_k = n - 1
        if k < 1 or k > max_k:
            raise NodeInputError(
                "INVALID_ARGUMENT", f"k must satisfy 1 <= k <= len(points)-1 = {max_k} (got {k})"
            )

        tree = cKDTree(points)
        # Query k+1 to include self, then drop the self-match per row.
        dist, idx = tree.query(points, k=k + 1, p=2)
        dist = np.atleast_2d(dist)
        idx = np.atleast_2d(idx)

        edges = []
        weights = []
        for i in range(n):
            count = 0
            for j, dj in zip(idx[i], dist[i]):
                if int(j) == i:
                    continue
                edges.append((i, int(j)))
                weights.append(float(dj))
                count += 1
                if count == k:
                    break

        return KNNGraphResult(
            edges=array_to_int_matrix(np.array(edges, dtype=np.int64) if edges else np.empty((0, 2))),
            weights=array_to_vector(weights),
        )
    except NodeInputError as e:
        return KNNGraphResult(error=e.error)
    except Exception as e:
        return KNNGraphResult(error=err("GRAPH_FAILED", str(e)))

from scipy.spatial import cKDTree

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import NearestSinglePointInput, NearestSinglePointResult
from nodes._common import (
    NodeInputError,
    check_matrix_shape,
    err,
    matrix_to_array,
    resolve_kdtree_metric,
    vector_to_array,
)

def nearest_single_point(ax: AxiomContext, input: NearestSinglePointInput) -> NearestSinglePointResult:
    """Find the single nearest point in `points` to one `query` point, via
    scipy.spatial.cKDTree — a convenience node for the common single-query
    case (see KNearestNeighbors for the batched-queries, k>=1 general case).
    metric is "euclidean" (default), "cityblock", or "chebyshev".
    """
    try:
        check_matrix_shape(input.points, name="points")
        points = matrix_to_array(input.points)
        query = vector_to_array(input.query, name="query")
        if len(query) == 0:
            raise NodeInputError("EMPTY_INPUT", "query must have at least one dimension")
        if points.shape[1] != len(query):
            raise NodeInputError(
                "DIMENSION_MISMATCH",
                f"points has {points.shape[1]} columns, query has {len(query)} dimensions",
            )

        metric, p = resolve_kdtree_metric(input.metric)

        tree = cKDTree(points)
        dist, idx = tree.query(query, k=1, p=p)

        return NearestSinglePointResult(index=int(idx), distance=float(dist))
    except NodeInputError as e:
        return NearestSinglePointResult(error=e.error)
    except Exception as e:
        return NearestSinglePointResult(error=err("QUERY_FAILED", str(e)))

from scipy.spatial.distance import pdist

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import PdistInput, PdistResult
from nodes._common import (
    NodeInputError,
    array_to_vector,
    check_matrix_shape,
    err,
    matrix_to_array,
    resolve_full_metric_kwargs,
)

def pairwise_distance_matrix(ax: AxiomContext, input: PdistInput) -> PdistResult:
    """Compute the condensed (upper-triangular) pairwise distance vector
    among all rows of `points`, via scipy.spatial.distance.pdist. metric is
    "euclidean" (default), "cityblock", "cosine", "chebyshev", "minkowski"
    (uses `p`), "hamming", "jaccard", or "mahalanobis" (uses `inv_cov` if
    supplied, else scipy's own internally-derived sample
    inverse-covariance). Returns the condensed vector (length n*(n-1)/2)
    plus `n` so the caller can index or squareform-expand it.
    """
    try:
        check_matrix_shape(input.points, min_rows=2, name="points")
        points = matrix_to_array(input.points)
        n = points.shape[0]

        metric, kwargs = resolve_full_metric_kwargs(
            input.metric, input.p, input.inv_cov if input.HasField("inv_cov") else None,
            require_inv_cov_for_mahalanobis=False,
        )

        condensed = pdist(points, metric=metric, **kwargs)

        return PdistResult(condensed=array_to_vector(condensed), n=n)
    except NodeInputError as e:
        return PdistResult(error=e.error)
    except Exception as e:
        return PdistResult(error=err("DISTANCE_FAILED", str(e)))

import numpy as np
from scipy.spatial.distance import pdist

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import PointDistanceInput, PointDistanceResult
from nodes._common import (
    NodeInputError,
    err,
    resolve_full_metric_kwargs,
    vector_to_array,
)

def point_distance(ax: AxiomContext, input: PointDistanceInput) -> PointDistanceResult:
    """Compute the distance between exactly two D-dimensional points under
    a chosen metric, via scipy.spatial.distance. metric is "euclidean"
    (default), "cityblock", "cosine", "chebyshev", "minkowski" (uses `p`),
    "hamming", "jaccard", or "mahalanobis" (requires `inv_cov` to be
    supplied explicitly — with only two points there is no sample to
    derive a covariance from). `a` and `b` must have equal length.
    """
    try:
        a = vector_to_array(input.a, name="a")
        b = vector_to_array(input.b, name="b")
        if len(a) == 0 or len(b) == 0:
            raise NodeInputError("EMPTY_INPUT", "a and b must each have at least one dimension")
        if len(a) != len(b):
            raise NodeInputError("DIMENSION_MISMATCH", f"a has {len(a)} dims, b has {len(b)} dims")

        metric, kwargs = resolve_full_metric_kwargs(
            input.metric, input.p, input.inv_cov if input.HasField("inv_cov") else None,
            require_inv_cov_for_mahalanobis=True,
        )

        # Route through pdist on the 2-row stack so this node's numeric
        # behavior is exactly consistent with Pdist/Cdist for the same
        # metric (rather than a separately hand-rolled formula).
        distance = float(pdist(np.vstack([a, b]), metric=metric, **kwargs)[0])

        return PointDistanceResult(distance=distance)
    except NodeInputError as e:
        return PointDistanceResult(error=e.error)
    except Exception as e:
        return PointDistanceResult(error=err("DISTANCE_FAILED", str(e)))

import numpy as np
from scipy.spatial import ConvexHull, QhullError

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import PointInHullInput, PointInHullResult
from nodes._common import (
    NodeInputError,
    array_to_bool_vector,
    check_matrix_shape,
    err,
    matrix_to_array,
)

MIN_DIM = 2
DEFAULT_TOL = 1e-9


def point_in_convex_hull(ax: AxiomContext, input: PointInHullInput) -> PointInHullResult:
    """Test whether each point in `queries` lies inside (or on the boundary
    of, within `tol`) the convex hull of `points`, via
    scipy.spatial.ConvexHull's facet equations (hull.equations) — the
    standard scipy-documented point-in-hull technique. `tol` defaults to
    1e-9 when <= 0.
    """
    try:
        check_matrix_shape(input.points, name="points")
        check_matrix_shape(input.queries, name="queries")
        points = matrix_to_array(input.points)
        queries = matrix_to_array(input.queries)
        n, d = points.shape
        if d < MIN_DIM:
            raise NodeInputError("INVALID_ARGUMENT", f"points must have D >= {MIN_DIM}, got D={d}")
        if points.shape[1] != queries.shape[1]:
            raise NodeInputError(
                "DIMENSION_MISMATCH",
                f"points has {points.shape[1]} columns, queries has {queries.shape[1]}",
            )
        if n < d + 1:
            raise NodeInputError(
                "INSUFFICIENT_POINTS", f"need at least {d + 1} points for a {d}D hull, got {n}"
            )

        try:
            hull = ConvexHull(points)
        except QhullError as e:
            raise NodeInputError(
                "DEGENERATE_INPUT", f"points are degenerate (coplanar/collinear, rank < {d}): {e}"
            )

        tol = input.tol if input.tol > 0 else DEFAULT_TOL
        # hull.equations: (n_facets, D+1) rows [normal..., offset]; a point
        # is inside/on-boundary iff normal . p + offset <= tol for every facet.
        eq = hull.equations
        inside = np.all(eq[:, :-1] @ queries.T + eq[:, -1:] <= tol, axis=0)

        return PointInHullResult(inside=array_to_bool_vector(inside))
    except NodeInputError as e:
        return PointInHullResult(error=e.error)
    except Exception as e:
        return PointInHullResult(error=err("HULL_TEST_FAILED", str(e)))

from scipy.spatial import procrustes

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import ProcrustesInput, ProcrustesResult
from nodes._common import (
    NodeInputError,
    array_to_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
)

def procrustes_disparity(ax: AxiomContext, input: ProcrustesInput) -> ProcrustesResult:
    """Compute the Procrustes disparity between two equally-shaped point
    sets (shape alignment), via scipy.spatial.procrustes — standardizes
    both point sets and finds the optimal rotation of `points_b` onto
    `points_a`, returning the sum-of-squared-errors disparity (0 =
    identical shape) plus both standardized/aligned point sets.
    """
    try:
        check_matrix_shape(input.points_a, min_rows=2, name="points_a")
        check_matrix_shape(input.points_b, min_rows=2, name="points_b")
        a = matrix_to_array(input.points_a, name="points_a")
        b = matrix_to_array(input.points_b, name="points_b")

        if a.shape != b.shape:
            raise NodeInputError(
                "DIMENSION_MISMATCH", f"points_a has shape {a.shape}, points_b has shape {b.shape}"
            )

        try:
            aligned_a, aligned_b, disparity = procrustes(a, b)
        except ValueError as e:
            raise NodeInputError("DEGENERATE_INPUT", str(e))

        return ProcrustesResult(
            disparity=float(disparity),
            aligned_a=array_to_matrix(aligned_a),
            aligned_b=array_to_matrix(aligned_b),
        )
    except NodeInputError as e:
        return ProcrustesResult(error=e.error)
    except Exception as e:
        return ProcrustesResult(error=err("PROCRUSTES_FAILED", str(e)))

from scipy.spatial import cKDTree

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import QueryPairsInput, QueryPairsResult
from nodes._common import (
    NodeInputError,
    array_to_int_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
    resolve_kdtree_metric,
)


def query_pairs(ax: AxiomContext, input: QueryPairsInput) -> QueryPairsResult:
    """Find every pair of points within `distance` of each other, inside a
    single point set, via scipy.spatial.cKDTree.query_pairs. metric is
    "euclidean" (default), "cityblock", or "chebyshev". Returns each
    qualifying pair (i, j) with i < j as row indices into `points`.
    """
    try:
        check_matrix_shape(input.points, name="points")
        points = matrix_to_array(input.points)
        if input.distance <= 0:
            raise NodeInputError("INVALID_ARGUMENT", f"distance must be > 0 (got {input.distance})")

        metric, p = resolve_kdtree_metric(input.metric)

        tree = cKDTree(points)
        pairs = tree.query_pairs(r=input.distance, p=p, output_type="ndarray")

        return QueryPairsResult(pairs=array_to_int_matrix(pairs))
    except NodeInputError as e:
        return QueryPairsResult(error=e.error)
    except Exception as e:
        return QueryPairsResult(error=err("QUERY_FAILED", str(e)))

from scipy.spatial import cKDTree

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import RadiusNeighborsInput, RadiusNeighborsResult
from nodes._common import (
    NodeInputError,
    check_matrix_shape,
    err,
    matrix_to_array,
    ragged_to_int_matrix,
    resolve_kdtree_metric,
)


def radius_neighbors(ax: AxiomContext, input: RadiusNeighborsInput) -> RadiusNeighborsResult:
    """Build a KD-tree over `points` and, for each row in `queries`, return
    every point within `radius` (a "ball query"), via
    scipy.spatial.cKDTree.query_ball_point. metric is "euclidean" (default),
    "cityblock", or "chebyshev". Result rows are ragged (variable neighbor
    count per query) and not guaranteed sorted by distance.
    """
    try:
        check_matrix_shape(input.points, name="points")
        check_matrix_shape(input.queries, name="queries")
        points = matrix_to_array(input.points)
        queries = matrix_to_array(input.queries)
        if points.shape[1] != queries.shape[1]:
            raise NodeInputError(
                "DIMENSION_MISMATCH",
                f"points has {points.shape[1]} columns, queries has {queries.shape[1]}",
            )
        if input.radius <= 0:
            raise NodeInputError("INVALID_ARGUMENT", f"radius must be > 0 (got {input.radius})")

        metric, p = resolve_kdtree_metric(input.metric)

        tree = cKDTree(points)
        neighbor_lists = tree.query_ball_point(queries, r=input.radius, p=p)

        return RadiusNeighborsResult(indices=ragged_to_int_matrix(neighbor_lists))
    except NodeInputError as e:
        return RadiusNeighborsResult(error=e.error)
    except Exception as e:
        return RadiusNeighborsResult(error=err("QUERY_FAILED", str(e)))

from scipy.spatial import QhullError, Voronoi

from gen.axiom_context import AxiomContext
from gen.messages_pb2 import VoronoiInput, VoronoiResult
from nodes._common import (
    NodeInputError,
    array_to_int_vector,
    array_to_matrix,
    check_matrix_shape,
    err,
    matrix_to_array,
    ragged_to_int_matrix,
)

ALLOWED_DIMS = (2, 3)


def voronoi_diagram(ax: AxiomContext, input: VoronoiInput) -> VoronoiResult:
    """Compute the Voronoi diagram of a 2D or 3D point set via
    scipy.spatial.Voronoi (Qhull) — the geometric dual of Delaunay
    triangulation. Returns the finite Voronoi vertices, each input point's
    (possibly unbounded, ragged) region as a list of vertex indices with -1
    marking an unbounded direction, the point_region mapping from input
    point to its region, and the ridge structure.
    """
    try:
        check_matrix_shape(input.points)
        points = matrix_to_array(input.points)
        n, d = points.shape
        if d not in ALLOWED_DIMS:
            raise NodeInputError("INVALID_ARGUMENT", f"points must be 2D or 3D, got D={d}")
        if n < d + 1:
            raise NodeInputError(
                "INSUFFICIENT_POINTS", f"need at least {d + 1} points for a {d}D Voronoi diagram, got {n}"
            )

        try:
            vor = Voronoi(points)
        except QhullError as e:
            raise NodeInputError(
                "DEGENERATE_INPUT",
                f"points are degenerate (collinear/coplanar) for a {d}D Voronoi diagram: {e}",
            )

        return VoronoiResult(
            vertices=array_to_matrix(vor.vertices),
            regions=ragged_to_int_matrix(vor.regions),
            point_region=array_to_int_vector(vor.point_region),
            ridge_points=ragged_to_int_matrix(vor.ridge_points),
            ridge_vertices=ragged_to_int_matrix(vor.ridge_vertices),
        )
    except NodeInputError as e:
        return VoronoiResult(error=e.error)
    except Exception as e:
        return VoronoiResult(error=err("VORONOI_FAILED", str(e)))