densemaps.numpy.nn_utils

Functions

compute_sqdistmat(X, Y[, normalized])

Computes the pairwise squared Euclidean distance matrix between two sets of points X and Y.

knn_query(X, Y[, k, return_distance, n_jobs])

Query nearest neighbors.
densemaps.numpy.nn_utils.knn_query(X, Y, k=1, return_distance=False, n_jobs=1)

Query nearest neighbors.

Parameters:

  • X (np.ndarray) – (N1,p) or (B, N1, p) first collection

  • Y (np.ndarray) – (N2,p) or (B, N2, p) second collection

  • k (int) – number of neighbors to look for

  • return_distance (bool) – hether to return the nearest neighbor distance

  • n_jobs (int) – number of parallel jobs. Set to -1 to use all processes

Returns:

  • dists (np.ndarray, optional) – (n2,k) or (n2,) if k=1 (with optional first batch dimension)- ONLY if return_distance is False. Nearest neighbor distance.

  • matches (np.ndarray) – (n2,k) or (n2,) if k=1 (with optional first batch dimension)- nearest neighbor

densemaps.numpy.nn_utils.compute_sqdistmat(X, Y, normalized=False)

Computes the pairwise squared Euclidean distance matrix between two sets of points X and Y.

Parameters:

  • X (torch.Tensor) – The first set of points, of shape (N, D) or (B, N, D).

  • Y (torch.Tensor) – The second set of points, of shape (M, D) or (B, M, D).

  • normalized (bool) – Whether the points are normalized to have unit norm.

Returns:

distmat – The pairwise squared Euclidean distance matrix between X and Y, of shape (N, M) or (B, N, M).

Return type:

torch.Tensor