How do i efficiently find the distance from a set of lines/rays to a point cloud

How can i efficiently find the euclidean distance from N unbounded rays (parametrized by a point and a direction) and M points in 3D, using Python/Numpy/PyTorch? The goal is to end up with N distances, from each ray to its nearest point.

The naive solution is to compute the distance between each ray and each point, but this has complexity O(NM). Does there exist any algorithm that may speed up this query? Perhaps one based on r-trees?

Solution 1:

Answering my own question:

I ended up using a ray-marching approximation, based on BallTree

from sklearn.neighbors import BallTree
import numpy as np

def rays_pc_distance(
        ray_origins : np.ndarray, # (N, 3)
        ray_dirs    : np.ndarray, # (N, 3)
        points      : np.ndarray, # (M, 3)
        n_steps     : int = 10,
        ) -> np.ndarray:

    index = BallTree(points)
    min_d = index.query(ray_origins + np.zeros_like(ray_dirs), k=1, return_distance=True)[0]
    acc_d = min_d.copy()
    for _ in range(n_steps):
        current_d = index.query(ray_origins + acc_d * ray_dirs, k=1, return_distance=True)[0]
        np.minimum(current_d, min_d, out=min_d)
        acc_d += current_d #* 0.8    

    return min_d # (N, 1)

Please let me know if there is something seriously wrong with this approach.