Well, did it with wrapping. I haven't tested correctness, but this seems to work.
(The array annotation you can remove or copy from here)
def get_delaunay_pose_neighbours(poses: Array['N',3], meters_per_turn: float) -> Array['P,2', int]:
""" Do delaunay triangulation in pose-space. The trickiness here is in doing it correctly for the wrapped yaw dimension
"""
yaws = poses[:, [2]] # poses are (x, y, yaw)
wrapped_poses = np.hstack([poses[:, :2], yaws + np.where(yaws < np.pi, 2 * np.pi, -2 * np.pi)])
vertices = np.vstack([poses, wrapped_poses]) * (1, 1, meters_per_turn)
pairs = get_delaunay_neighbour_indices(vertices).T
pairs = pairs % len(poses) # Collapse the wrapping
pairs = pairs[pairs[:, 0] != pairs[:, 1]] # Remove self-pairs
pairs = get_unique_rows(pairs) # Remove duplicate pairs (from the wrapping)
return pairs
def get_delaunay_neighbour_indices(vertices: Array['N,D', int]) -> Array['2,P', int]:
"""
Find each pair of neighbouring vertices in the Delaunay triangulation.
# TODO: See if we can get better performance. This takes about 30ms on a BPearl Scan.
:param vertices: The vertices of the points to perform Delaunay triangulation on
:return: The pairs of indices of vertices
"""
tri = Delaunay(vertices)
spacing_indices, neighbours = tri.vertex_neighbor_vertices
ixs = np.zeros((2, len(neighbours)), dtype=int)
ixs[0, spacing_indices[1:np.argmax(spacing_indices)]] = 1 # The argmax is unfortuantely needed when multiple final elements the same
ixs[0, :] = np.cumsum(ixs[0, :])
ixs[1, :] = neighbours
return ixs
def get_unique_rows(arr: Array['N,D', float]) -> Array['M,D', float]:
""" Get unique rows, given the arrays"""
row_view = collapse_rows_into_values(arr)
values = np.unique(row_view)
values_recast = values.view(arr.dtype).reshape(len(values), -1)
return values_recast
def collapse_rows_into_values(arr: Array['N,D', Any]) -> Array['N', Any]:
""" Collapse rows into single values. Useful for finding unique rows with np.unique """
return np.ascontiguousarray(arr).view(np.dtype((np.void, arr.dtype.itemsize * np.prod(arr.shape[1:])))).ravel()