gemlib.spatial.sparse_pdist

Contents

gemlib.spatial.sparse_pdist#

gemlib.spatial.sparse_pdist(coords, include_fn=<function include_all>, chunk_size=None)#

Compute a sparse distance matrix

Compute a sparse Euclidean distance matrix between all pairs of coords such that the distance is less than max_dist.

Parameters:

  • coords (ArrayLike) – a [N, D] array of coordinates

  • include_fn (Callable[[TypeAliasForwardRef('ArrayLike')], ndarray[tuple[int, ...], dtype[bool]]]) – a callable that takes a float representing the distance between two points, and returns True if the distance should be included as a “non-zero” element of the returned sparse matrix.

  • batch_size – If memory is limited, compute the distances in batches of [batch_size, N] stripes.

Returns:

A sparse tensor of Euclidean distances satisfying include_fn.

Return type:

Array

Example

>>> import numpy as np
>>> from gemlib.spatial import sparse_pdist
>>> coords = np.random.uniform(size=(1000, 2))
>>> d_sparse = sparse_pdist(coords, max_dist=0.01, batch_size=200)
>>> d_sparse
SparseTensor(indices=tf.Tensor(
[[  0   0]
 [  1   1]
 [  2   2]
 ...
 [997 997]
 [998 998]
 [999 999]], shape=(1316, 2), dtype=int64), values=tf.Tensor(
[0.00000000e+00 2.22044605e-16 0.00000000e+00 ... 0.00000000e+00
 0.00000000e+00 0.00000000e+00], shape=(1316,), dtype=float64),
dense_shape=tf.Tensor([1000 1000], shape=(2,), dtype=int64))
Contents