densemaps.torch.nn_utils

Functions

compute_sqdistmat(X, Y[, normalized])

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

nn_query(X, Y[, use_keops])

Computes the nearest neighbor query between two sets of points X and Y.

nn_query_dense(X, Y)

Computes the nearest neighbor query between two sets of points X and Y.

nn_query_keops(X, Y)

Computes the nearest neighbor query between two sets of points X and Y.
densemaps.torch.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).

Returns:

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

Return type:

torch.Tensor

densemaps.torch.nn_utils.nn_query(X, Y, use_keops=None)

Computes the nearest neighbor query between two sets of points X and Y. Use KeOps for efficient computation if available and more than 25k points.

If not on GPU, uses numpy nn

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).

  • use_keops (bool) – Whether to use KeOps for efficient computation. If None, use KeOps if available.

Returns:

The indices of the nearest neighbors of each point of Y in X, of shape (M,) or (B, M).

Return type:

torch.Tensor

densemaps.torch.nn_utils.nn_query_dense(X, Y)

Computes the nearest neighbor query between two sets of points X and Y. Use dense computation with PyTorch.

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).

Returns:

The indices of the nearest neighbors of each point of Y in X, of shape (M,) or (B, M).

Return type:

torch.Tensor

densemaps.torch.nn_utils.nn_query_keops(X, Y)

Computes the nearest neighbor query between two sets of points X and Y. Use KeOps for efficient computation.

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).

Returns:

The indices of the nearest neighbors of each point of Y in X, of shape (M,) or (B, M).

Return type:

torch.Tensor