# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-unsafe
from typing import Tuple
import torch
# ------------------------ Laplacian Matrices ------------------------ #
# This file contains implementations of differentiable laplacian matrices.
# These include
# 1) Standard Laplacian matrix
# 2) Cotangent Laplacian matrix
# 3) Norm Laplacian matrix
# -------------------------------------------------------------------- #
[docs]
def laplacian(verts: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
"""
Computes the laplacian matrix.
The definition of the laplacian is
L[i, j] = -1 , if i == j
L[i, j] = 1 / deg(i) , if (i, j) is an edge
L[i, j] = 0 , otherwise
where deg(i) is the degree of the i-th vertex in the graph.
Args:
verts: tensor of shape (V, 3) containing the vertices of the graph
edges: tensor of shape (E, 2) containing the vertex indices of each edge
Returns:
L: Sparse FloatTensor of shape (V, V)
"""
V = verts.shape[0]
e0, e1 = edges.unbind(1)
idx01 = torch.stack([e0, e1], dim=1) # (E, 2)
idx10 = torch.stack([e1, e0], dim=1) # (E, 2)
idx = torch.cat([idx01, idx10], dim=0).t() # (2, 2*E)
# First, we construct the adjacency matrix,
# i.e. A[i, j] = 1 if (i,j) is an edge, or
# A[e0, e1] = 1 & A[e1, e0] = 1
ones = torch.ones(idx.shape[1], dtype=torch.float32, device=verts.device)
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
A = torch.sparse.FloatTensor(idx, ones, (V, V))
# the sum of i-th row of A gives the degree of the i-th vertex
deg = torch.sparse.sum(A, dim=1).to_dense()
# We construct the Laplacian matrix by adding the non diagonal values
# i.e. L[i, j] = 1 ./ deg(i) if (i, j) is an edge
deg0 = deg[e0]
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
deg0 = torch.where(deg0 > 0.0, 1.0 / deg0, deg0)
deg1 = deg[e1]
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
deg1 = torch.where(deg1 > 0.0, 1.0 / deg1, deg1)
val = torch.cat([deg0, deg1])
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
L = torch.sparse.FloatTensor(idx, val, (V, V))
# Then we add the diagonal values L[i, i] = -1.
idx = torch.arange(V, device=verts.device)
idx = torch.stack([idx, idx], dim=0)
ones = torch.ones(idx.shape[1], dtype=torch.float32, device=verts.device)
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
L -= torch.sparse.FloatTensor(idx, ones, (V, V))
return L
[docs]
def cot_laplacian(
verts: torch.Tensor, faces: torch.Tensor, eps: float = 1e-12
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Returns the Laplacian matrix with cotangent weights and the inverse of the
face areas.
Args:
verts: tensor of shape (V, 3) containing the vertices of the graph
faces: tensor of shape (F, 3) containing the vertex indices of each face
Returns:
2-element tuple containing
- **L**: Sparse FloatTensor of shape (V,V) for the Laplacian matrix.
Here, L[i, j] = cot a_ij + cot b_ij iff (i, j) is an edge in meshes.
See the description above for more clarity.
- **inv_areas**: FloatTensor of shape (V,) containing the inverse of sum of
face areas containing each vertex
"""
V, F = verts.shape[0], faces.shape[0]
face_verts = verts[faces]
v0, v1, v2 = face_verts[:, 0], face_verts[:, 1], face_verts[:, 2]
# Side lengths of each triangle, of shape (sum(F_n),)
# A is the side opposite v1, B is opposite v2, and C is opposite v3
A = (v1 - v2).norm(dim=1)
B = (v0 - v2).norm(dim=1)
C = (v0 - v1).norm(dim=1)
# Area of each triangle (with Heron's formula); shape is (sum(F_n),)
s = 0.5 * (A + B + C)
# note that the area can be negative (close to 0) causing nans after sqrt()
# we clip it to a small positive value
# pyre-fixme[16]: `float` has no attribute `clamp_`.
area = (s * (s - A) * (s - B) * (s - C)).clamp_(min=eps).sqrt()
# Compute cotangents of angles, of shape (sum(F_n), 3)
A2, B2, C2 = A * A, B * B, C * C
cota = (B2 + C2 - A2) / area
cotb = (A2 + C2 - B2) / area
cotc = (A2 + B2 - C2) / area
cot = torch.stack([cota, cotb, cotc], dim=1)
cot /= 4.0
# Construct a sparse matrix by basically doing:
# L[v1, v2] = cota
# L[v2, v0] = cotb
# L[v0, v1] = cotc
ii = faces[:, [1, 2, 0]]
jj = faces[:, [2, 0, 1]]
idx = torch.stack([ii, jj], dim=0).view(2, F * 3)
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
L = torch.sparse.FloatTensor(idx, cot.view(-1), (V, V))
# Make it symmetric; this means we are also setting
# L[v2, v1] = cota
# L[v0, v2] = cotb
# L[v1, v0] = cotc
L += L.t()
# For each vertex, compute the sum of areas for triangles containing it.
idx = faces.view(-1)
inv_areas = torch.zeros(V, dtype=torch.float32, device=verts.device)
val = torch.stack([area] * 3, dim=1).view(-1)
inv_areas.scatter_add_(0, idx, val)
idx = inv_areas > 0
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
inv_areas[idx] = 1.0 / inv_areas[idx]
inv_areas = inv_areas.view(-1, 1)
return L, inv_areas
[docs]
def norm_laplacian(
verts: torch.Tensor, edges: torch.Tensor, eps: float = 1e-12
) -> torch.Tensor:
"""
Norm laplacian computes a variant of the laplacian matrix which weights each
affinity with the normalized distance of the neighboring nodes.
More concretely,
L[i, j] = 1. / wij where wij = ||vi - vj|| if (vi, vj) are neighboring nodes
Args:
verts: tensor of shape (V, 3) containing the vertices of the graph
edges: tensor of shape (E, 2) containing the vertex indices of each edge
Returns:
L: Sparse FloatTensor of shape (V, V)
"""
edge_verts = verts[edges] # (E, 2, 3)
v0, v1 = edge_verts[:, 0], edge_verts[:, 1]
# Side lengths of each edge, of shape (E,)
w01 = 1.0 / ((v0 - v1).norm(dim=1) + eps)
# Construct a sparse matrix by basically doing:
# L[v0, v1] = w01
# L[v1, v0] = w01
e01 = edges.t() # (2, E)
V = verts.shape[0]
# pyre-fixme[16]: Module `sparse` has no attribute `FloatTensor`.
L = torch.sparse.FloatTensor(e01, w01, (V, V))
L = L + L.t()
return L