densemaps.torch¶

PyTorch backend for densemaps (CUDA / keops capable).

Importing this subpackage requires torch. The memory-scalable KernelDistMap additionally requires pykeops (install the [keops] extra); it raises a clear ImportError at construction time if keops is missing.

class densemaps.torch.PointWiseMap(tensor_names=None)¶

Root class representing a pointwise map. Not supposed to be instanciated in itself.

A pointwise, denoted \(P : S_2 \to S_1\), is a map that associates to each point x in \(S_2\) a point \(P(x)\) in \(S_1\). It can usually be represented as a \(n_2 \times n_1\) matrix, where \(n_2\) is the number of points in \(S_2\) and \(n_1\) the number of points in \(S_1\).

Given a pointwise map \(P\), the pullback of a function \(f : S_1 \to R\) is a function \(f_{pb} : S_2 \to R\) defined by \(f_{pb}(x) = f(P(x))\). In practice it can easily be computed by matrix multiplication: \(f_{pb} = P f\).

In practice, we usually don’t need to use the exact values inside \(P\), but rather only care about multiplying with some functions, extracting maximal values per-row or per-column, or summing on rows or columns.

In torch, all maps can be represented in a batched way, with an additional dimension at the beginning of the tensor.

tensor_names¶

Names of the tensors stored in the object. Used to transfer to GPU.

Type:: list of str

is_index_map¶

Whether the map is intrinsically stored as a one-nonzero-per-row index array (i.e. a vertex-to-vertex map). False for all matrix-valued representations.

Type:: bool

is_index_map = False¶: Intrinsically an index (vertex-to-vertex) representation. Overridden by P2PMap.

property shape¶

Shape of the map.

Returns:: shape – Depends on the representation
Return type:: tuple

property ndim¶

Number of dimensions of the map.

Returns:: ndim – Number of dimensions
Return type:: int

pull_back(f)¶

Pull back a function \(f\). Three possibilities:

f is a function on S1, of shape (N1,), then the output is a function on S2, of shape (N2,)
f represents multiple function on S1, of shape (N1, p), then the output is a function on S2, of shape (N2, p)
f represents a batch multiple function on S1, of shape (B, N1, p), then the output is a function on S2, of shape (B, N2, p)

Note tht the case where f is a batch of a single function (B, N1) is not supported, and one should then use f[…, None]

Parameters:: f (torch.Tensor) – (N1,), (N1, p) or (B, N1, p)
Returns:: f_pb – (N2,), (N2, p) or (B, N2, p)
Return type:: torch.Tensor

get_nn()¶

Outputs the nearest neighbor map. The nearest neighbor map is the map that associates to each point of S2 the index of the closest point in S1.

Returns:: nn_map – (N2,) or (B, N2) depending on the representation
Return type:: torch.Tensor

property mT¶

Returns the transpose of the map.

Returns the transpose of the matrix representation of the map.

Returns:: map_t – Transpose map
Return type:: PointWiseMap

class densemaps.torch.SparseMap(map)¶

Map represented by a sparse matrix.

The sparse matrix is of size (N2, N1).

Can actually represent a densemap too, but it’s not recommended.

Parameters:: map (scipy.sparse.csr_matrix) – (N2, N1) or (B, N2, N1)

property shape¶

Shape of the map.

Returns:: shape – returns (N2, N1) or (B, N2, N1)
Return type:: tuple

property mT¶

Returns the transpose of the map.

Returns the transpose of the matrix representation of the map.

Returns:: map_t – Transpose map
Return type:: SparseMap

pull_back(f)¶

Pull back a function \(f\). Three possibilities:

f is a function on S1, of shape (N1,), then the output is a function on S2, of shape (N2,)
f represents multiple function on S1, of shape (N1, p), then the output is a function on S2, of shape (N2, p)
f represents a batch multiple function on S1, of shape (B, N1, p), then the output is a function on S2, of shape (B, N2, p)

Note tht the case where f is a batch of a single function (B, N1) is not supported, and one should then use f[…, None]

Parameters:: f (torch.Tensor) – (N1,), (N1, p) or (B, N1, p)
Returns:: f_pb – (N2,), (N2, p) or (B, N2, p)
Return type:: torch.Tensor

get_nn()¶

Outputs the nearest neighbor map. The nearest neighbor map associates to each point of S2 the index of the highest-weight point in S1 (argmax along the last dimension).

Returns:: nn_map – (N2,) or (B, N2) depending on the representation
Return type:: torch.Tensor

to_dense()¶: Returns the dense (N2, N1) (or (B, N2, N1)) tensor representation.

class densemaps.torch.P2PMap(p2p_21, n1=None)¶

Point to point map from a set S2 to a set S1. Defined by a map \(P_{21} : S2 \to S1\) or a tensor p2p_21 of size (n_2), where p2p_21[i] is the index of the point in S1 closest to the point i in S2. Batched versions are accepted

Parameters:

p2p_21 (th.Tensor) – (n2,) or (B, n2)
n1 (int or None) – Number of points in S1. If None, n1 = p2p.max()+1

is_index_map = True¶: P2PMap is stored as a one-index-per-row tensor (see shape).

property n1¶

Number of vertices on the first shape. Estimated if not provided as input.

Returns:: n1
Return type:: int

property shape¶

Shape of the map, as a matrix.

Even though the map is stored as an index array (see is_index_map), its shape is reported as a matrix for consistency with the other representations. The underlying index tensor is available as self.p2p_21 (shape (N2,)/(B, N2)).

Returns:: shape – (N2, N1), or (B, N2, N1) for a batched map
Return type:: torch.Size

pull_back(f)¶

Pull back a function \(f\). Four possibilities:

f is a function on S1, of shape (N1,), then the output is a function on S2, of shape (N2,) (or (B, N2,) if the map is batched)
f represents multiple function on S1, of shape (N1, p), then the output is a function on S2, of shape (N2, p) (or (B, N2, p) if the map is batched)
f represents a batch multiple function on S1, of shape (B, N1, p), then the output is a function on S2, of shape (B, N2, p)

Note tht the case where f is a batch of a single function (B, N1) is not supported, and one should then use f[…, None]

Parameters:: f (torch.Tensor) – (N1,), (N1, p) or (B, N1, p)
Returns:: f_pb – (N2,), (N2, p) or (B, N2, p)
Return type:: torch.Tensor

get_nn()¶

Outputs the nearest neighbor map. The nearest neighbor map is the same as the input.

Returns:: nn_map – (N2,) or (B, N2) depending on the representation
Return type:: torch.Tensor

property mT¶

Returns the transpose of the map.

Returns the transpose of the matrix representation of the map.

Returns:: map_t – Transpose map
Return type:: PointWiseMap

to_dense()¶: Returns the dense (N2, N1) (or (B, N2, N1)) tensor representation.

class densemaps.torch.PreciseMap(v2face_21, bary_coords, faces1, n1=None)¶

Point to barycentric map from a set S2 to a surface S1. Batched Version is not supported yet.

Is represented as a sparse matrix of size (N2, N1), where there are at most 3 non-zero entries per row, which sum to 1.

Note

Batched version is not supported yet.

Parameters:

v2face_21 (torch.Tensor) – (n2,) Indices of the faces of S1 closest to each point of S2.
bary_coords (torch.Tensor) – (n2, 3) Barycentric coordinates of the points of S2 in the faces of S1.
faces1 (torch.Tensor) – (N1, 3) All the Faces of S1.
n1 (int, optional) – Number of vertices of S1. If None, inferred as faces1.max()+1, which is only correct if every vertex of S1 belongs to at least one face.

property shape¶

Shape of the map.

Returns:: shape – Depends on the representation
Return type:: tuple

pull_back(f)¶

Pull back f using the map T.

Parameters:¶

f : (N1,), (N1, p) or (B, N, p)

returns:: pull_back
rtype:: (N2, p) or (B, N2, p)

get_nn()¶

Outputs the nearest neighbor map. The nearest neighbor map is the map that associates to each point of S2 the index of the closest point in S1.

Returns:: nn_map – (N2,) or (B, N2) depending on the representation
Return type:: torch.Tensor

property mT¶

Returns the transpose of the map.

Returns the transpose of the matrix representation of the map.

Returns:: map_t – Transpose map
Return type:: PointWiseMap

to_dense()¶: Returns the dense (N2, N1) tensor representation.

class densemaps.torch.EmbP2PMap(emb1, emb2)¶

Point to point map, computed from embeddings.

Simple wrapper around P2PMap

Parameters:

emb1 (torch.Tensor) – (N1, p) or (B, N1, p)
emb2 (torch.Tensor) – (N2, p) or (B, N2, p)

class densemaps.torch.EmbPreciseMap(emb1, emb2, faces1, clear_cache=True, use_keops=None)¶

Point to barycentric map, computed from embeddings.

Simple wrapper around PreciseMap

Note

Batched version is not supported yet.

Parameters:

emb1 (torch.Tensor) – (N1, K)
emb2 (torch.Tensor) – (N2, K)
faces1 (torch.Tensor) – (N1, 3)
clear_cache (bool) – Whether to free GPU memory after the projection (torch.cuda.empty_cache()). Only useful on CUDA; a bit slower.
use_keops (bool, optional) – Passed to the nearest-vertex query. Set False to skip KeOps and use plain torch.

class densemaps.torch.KernelDenseDistMap(log_matrix, lse_row=None, lse_col=None)¶

Map represented by a row-normalized dense matrix obtained from an element-wise exponential.

The matrix is of size (N2, N1), and has values \(P_{ij} = \frac{1}{\sum_j \exp(D_{ij})} \exp(D_{ij})\) where \(D\) is some matrix

Only D has to be provided

Parameters:

log_matrix (torch.Tensor) – (N2, N1) or (B, N2, N1), the matrix D
lse_row (torch.Tensor, optional) – (N2,) or (B, N2). The logsumexp on rows
lse_col (torch.Tensor) – (N1,) or (B, N1). The logsumexp on columns

to_dense()¶: Public alias for _to_dense(). Returns the (N2, N1) dense matrix.

pull_back(f)¶

Pull back a function \(f\). Four possibilities:

f is a function on S1, of shape (N1,), then the output is a function on S2, of shape (N2,)
f represents multiple function on S1, of shape (N1, p), then the output is a function on S2, of shape (N2, p)
f represents a batch multiple function on S1, of shape (B, N1, p), then the output is a function on S2, of shape (B, N2, p)

Note tht the case where f is a batch of a single function (B, N1) is not supported, and one should then use f[…, None]

Parameters:: f (torch.Tensor) – (N1,), (N1, p) or (B, N1, p)
Returns:: f_pb – (N2,), (N2, p) or (B, N2, p)
Return type:: torch.Tensor

get_nn()¶

Outputs the nearest neighbor map. The nearest neighbor map is the map that associates to each point of S2 the index of the closest point in S1. Simple argmax along the last dimension.

Returns:: nn_map – (N2,) on the representation
Return type:: torch.Tensor

reverse()¶

Row-normalized reverse map, from \(S_1\) to \(S_2\).

This is not the literal matrix transpose (see mT): the reverse map is re-normalized so that its rows sum to 1, i.e. it is the row-normalized version of \(\Pi^\top\). Use this to get a valid (row-stochastic) map in the reverse direction.

Returns:: map_rev – Row-normalized reverse map
Return type:: KernelDenseDistMap

property mT¶

Literal matrix transpose \(\Pi^\top\) of the (dense, row-normalized) map.

Note

The transpose of a row-stochastic matrix is not row-stochastic. For a valid reverse map (rows summing to 1), use reverse() instead.

Returns:: map_t – Transpose map, wrapping the dense (N1, N2) matrix
Return type:: SparseMap

property shape¶

Shape of the map.

Returns:: shape – Depends on the representation
Return type:: tuple

class densemaps.torch.EmbKernelDenseDistMap(emb1, emb2, blur=None, normalize=False, normalize_emb=False, dist_type='sqdist')¶

Kernel Map, computed from embeddings.

Simple wrapper around KernelDenseDistMap.

Kernel has the form \(\exp\big(-\frac{s(x,y)}{2\sigma^2}\big)\), where \(s(x,y)\) is either: - the negative squared distance between the embeddings of x and y. - the (positive) inner product between the embeddings of x and y (potentially normalized).

Parameters:

emb1 (torch.Tensor) – (N1, p) or (B, N1, p) embedding on first shape
emb2 (torch.Tensor) – (N2, p) or (B, N2, p) embedding on second shape
blur (float) – Standard deviation of the Gaussian kernel.
normalize (bool) – Normalize the blur by the maximum distance between embedding points
normalize_emb (bool) – Normalize the embeddings.
dist_type (string) – {“sqdist”, “inner”} Type of score to use.

class densemaps.torch.KernelDistMap(emb1, emb2, normalize=False, blur=None, normalize_emb=False, dist_type='sqdist')¶

Memory-Scalable Version of EmbKernelDenseDistMap

Row normalized version fo the kernel map of the form \(\exp\big(-\frac{s(x,y)}{2\sigma^2}\big)\), where \(s(x,y)\) is either: - the negative squared distance between the embeddings of x and y. - the (positive) inner product between the embeddings of x and y (potentially normalized).

Parameters:

emb1 (torch.Tensor) – (N1, p) or (B, N1, p) embedding on first shape
emb2 (torch.Tensor) – (N2, p) or (B, N2, p) embedding on second shape
blur (float) – Standard deviation of the Gaussian kernel.
normalize (bool) – Normalize the blur by the maximum distance between embedding points
normalize_emb (bool) – Normalize the embeddings.
dist_type (string) – {“sqdist”, “inner”} Type of score to use.

property shape¶

Shape of the map.

Returns:: shape – Depends on the representation
Return type:: tuple

get_pull_back_formula(dim)¶: B, N1, 1 -> B, N2, 1

to_dense()¶: Materializes the full (N2, N1) dense matrix.

Note

This defeats the memory-scalable design of KernelDistMap. It is only intended for inspection / small problems.

pull_back(f)¶

Pull back a function \(f\). Four possibilities:

f is a function on S1, of shape (N1,), then the output is a function on S2, of shape (N2,)
f represents multiple function on S1, of shape (N1, p), then the output is a function on S2, of shape (N2, p)
f represents a batch multiple function on S1, of shape (B, N1, p), then the output is a function on S2, of shape (B, N2, p)

Note tht the case where f is a batch of a single function (B, N1) is not supported, and one should then use f[…, None]

Parameters:: f (np.ndarray) – (N1,), (N1, p) or (B, N1, p)
Returns:: f_pb – (N2,), (N2, p) or (B, N2, p)
Return type:: np.ndarray

get_nn()¶

Highest-weight point of S1 for each point of S2.

For dist_type="sqdist" this is the nearest neighbour (min squared distance). For dist_type="inner" the kernel is monotonically increasing in the inner product, so it is the point of maximum inner product, not the nearest neighbour.

reverse()¶

Row-normalized reverse map, from \(S_1\) to \(S_2\) (memory-scalable).

This is not the literal matrix transpose (see mT): it is a valid row-stochastic KernelDistMap in the reverse direction, obtained by swapping the two embeddings.

property mT¶

Literal matrix transpose \(\Pi^\top\).

Note

The transpose of a row-stochastic matrix is not row-stochastic, and it cannot be expressed as a memory-scalable KernelDistMap. This materializes the full dense (N1, N2) matrix, defeating memory scalability. For a scalable reverse map use reverse() instead.

Returns:: map_t – Transpose map, wrapping the dense (N1, N2) matrix
Return type:: SparseMap

densemaps.torch.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_query_dist(X, Y, use_keops=None)¶

Computes, for each point of Y, the (Euclidean) distance to its nearest neighbor in X.

Same backend-selection logic as nn_query(), but returns distances rather than indices.

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 distance from each point of Y to its nearest neighbor in X, of shape (M,) or (B, M).

Return type:

torch.Tensor

densemaps.torch.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_query_precise_torch(vert_emb, faces, points_emb, return_dist=False, batch_size=None, clear_cache=True, use_keops=None)¶

Project a pointcloud on a p-dimensional triangle mesh

Parameters:

vert_emb – (n1, p) coordinates of the mesh vertices
faces – (m1, 3) faces of the mesh defined as indices of vertices
points_emb – (n2, p) coordinates of the pointcloud
return_dist (bool) – Whether to return the distance to the nearest vertex
batch_size (int, optional) – If precompute_dmin is False, projects batches of points on the surface
clear_cache (bool) – Whether to clear cache after computation
use_keops (bool, optional) – Passed through for the Delta_min query (see project_pc_to_triangles()). Set False to skip KeOps and use plain torch.

Returns:

face_match (torch.Tensor) – (n2,) - indices of the face assigned to each point.
bary_coord (torch.Tensor) – (n2,3) - barycentric coordinates of each point within the face.

densemaps.torch.project_pc_to_triangles(vert_emb, faces, points_emb, precompute_dmin=True, batch_size=None, verbose=False, use_keops=None)¶

Project a pointcloud on a set of triangles in p-dimension. Projection is defined as barycentric coordinates on one of the triangle.

Line i for the output has 3 non-zero values at indices j,k and l of the vertices of the triangle point i is projected on.

Note

This picks the closest face per point (an argmin), so it is not differentiable. It always runs under torch.no_grad(), so it never keeps a gradient graph on its large temporary (n_points, n_faces) tensors, whatever the caller’s grad setting.

Parameters:

vert_emb (torch.Tensor) – (n1, p) coordinates of the mesh vertices
faces (torch.Tensor) – (m1, 3) faces of the mesh defined as indices of vertices
points_emb (torch.Tensor) – (n2, p) coordinates of the pointcloud
precompute_dmin (bool) – Whether to precompute all the values of delta_min. Faster but heavier in memory.
batch_size (int, optional) – If precompute_dmin is False, projects batches of points on the surface
use_keops (bool, optional) – Passed to nn_query_dist() for the Delta_min query. If None, KeOps is used when available and the problem is large enough. Set False to skip KeOps and use plain torch (handy to check whether KeOps is the one holding onto GPU memory).

Returns:

face_match (torch.Tensor) – (n2,) - indices of the face assigned to each point.
bary_coord (torch.Tensor) – (n2,3) - barycentric coordinates of each point within the face.

Modules

`maps`
`nn_utils`
`point_to_triangle`