## Finite-element gradient and cotan-Laplacian
Building on the mesh geometry and adjancency-based operators, we can now
define two important linear operators that depend both on mesh
connectivity and on mesh geometry. They are the (discrete,
triangulation-based) equivalent of the gradient and Laplace-Beltrami
operator. The latter is known as the *cotan Laplacian*.
We now implement gradient (per-vertex scalar field -\> per-face vector
field) and the cotan-Laplacian (vertex -\> vertex) using gather/scatter
ops. In both cases, we start with a scalar field *u**i*
defined per vertex *i* of the triangulation. The finite-element gradient
is defined for each face *i**j**k*, like so:
(∇*u*)*i**j**k* = ∑*l* ∈ {*i*, *j*, *k*}*u**l*∇*ϕ**l*
where *ϕ**i* is a linear finite element test function (linear
Lagrange element) and has gradient
$$
\nabla\phi_i = \frac{1}{2a\_{ijk}} (\mathbf{v}\_k-\mathbf{v}\_j)^\perp
$$
plus cyclic permutations. Here, *a**i**j**k* is the triangle
area, **v***i* are the vertex positions, and ()⟂
denotes rotation by 90 degrees (in 3D, you rotate about the triangle
normal).
The cot-Laplacian computes the following per-vertex field:
$$
(\Delta u)\_i = \frac{1}{2} \sum\_{j} (\cot\alpha_j +\cot\beta_j) (u_j-u_i)
$$
The sum is over adjacent vertices, and
*α**j*, *β**j* are the two triangle angles
“opposite” to the edge *i**j*.
To check for correctness, we can compare with [this `libgigl`
tutorial](https://libigl.github.io/libigl-python-bindings/tut-chapter1/),
using the test mesh and some random test fields.
``` python
from triangulax.triangular import TriMesh
```
``` python
# load test data
mesh = TriMesh.read_obj("../test_meshes/disk.obj")
hemesh = msh.HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)
geommesh = msh.GeomMesh(*hemesh.n_items, mesh.vertices, mesh.face_positions)
mesh_3d = TriMesh.read_obj("../test_meshes/disk.obj", dim=3)
geommesh_3d = msh.GeomMesh(*hemesh.n_items, mesh_3d.vertices, mesh_3d.face_positions)
```
Warning: readOBJ() ignored non-comment line 3:
o flat_tri_ecmc
Warning: readOBJ() ignored non-comment line 3:
o flat_tri_ecmc
### Sparse matrix utility functions
------------------------------------------------------------------------
source
### diag_jsparse
``` python
def diag_jsparse(
v:Float[Array, 'N'], k:int=0
)->BCOO:
```
*Construct a diagonal jax.sparse array. Plugin replacement for np.diag*
------------------------------------------------------------------------
source
### bcoo_to_scipy
``` python
def bcoo_to_scipy(
A:BCOO, # Input JAX sparse matrix
)->csr_matrix: # Equivalent SciPy sparse matrix
```
*Convert a JAX BCOO sparse matrix to a SciPy CSR sparse matrix.*
------------------------------------------------------------------------
source
### scipy_to_bcoo
``` python
def scipy_to_bcoo(
A, # Input sparse matrix (CSR or CSC recommended)
)->BCOO: # Equivalent JAX sparse matrix
```
*Convert a SciPy sparse matrix (CSC or CSR) to a JAX BCOO sparse matrix*
without converting to dense.
### Cotan-Laplacian
------------------------------------------------------------------------
source
### compute_cotan_laplace
``` python
def compute_cotan_laplace(
vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh, vertex_field:Float[Array, 'n_vertices ...']
)->Float[Array, 'n_vertices ...']:
```
*Compute cotangent laplacian of a per-vertex field (natural boundary
conditions).*
------------------------------------------------------------------------
source
### cotan_laplace_sparse
``` python
def cotan_laplace_sparse(
vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh
)->BCOO:
```
*Assemble cotangent Laplacian as a sparse matrix (BCOO).*
``` python
# Test against libigl cotmatrix (natural boundary conditions)
key = jax.random.PRNGKey(0)
u = jax.random.normal(key, (hemesh.n_vertices,))
u_vec = jax.random.normal(key, (hemesh.n_vertices, 3))
L = igl.cotmatrix(np.asarray(geommesh.vertices), np.asarray(hemesh.faces))
lap_jax = compute_cotan_laplace(geommesh.vertices, hemesh, u)
lap_igl = L @ np.asarray(u)
rel_err = np.linalg.norm(np.asarray(lap_jax) - lap_igl) / np.linalg.norm(lap_igl)
print("scalar field rel. error:", rel_err)
lap_jax_vec = compute_cotan_laplace(geommesh.vertices, hemesh, u_vec)
lap_igl_vec = L @ np.asarray(u_vec)
rel_err_vec = np.linalg.norm(np.asarray(lap_jax_vec) - lap_igl_vec) / np.linalg.norm(lap_igl_vec)
print("vector field rel. error:", rel_err_vec)
```
scalar field rel. error: 1.7692000627878292e-16
vector field rel. error: 2.011929541056845e-16
``` python
# test sparse cotan Laplacian vs apply function
key = jax.random.PRNGKey(0)
u_test = jax.random.normal(key, (hemesh.n_vertices,))
L_sparse = cotan_laplace_sparse(geommesh.vertices, hemesh)
lap_sparse = L_sparse @ u_test
lap_apply = compute_cotan_laplace(geommesh.vertices, hemesh, u_test)
rel_err_sparse = jnp.linalg.norm(lap_sparse - lap_apply) / jnp.linalg.norm(lap_apply)
print("cotan sparse vs apply rel. error:", rel_err_sparse)
```
cotan sparse vs apply rel. error: 1.302894564211555e-16
``` python
bcoo_to_scipy(L_sparse), (scipy_to_bcoo(bcoo_to_scipy(L_sparse)).todense() == L_sparse.todense()).all()
```
(,
Array(True, dtype=bool))
### Mass matrix (lumped)
The finite-element mass matrix *M* appears whenever we discretize a
time-dependent PDE. For a lumped (diagonal) mass matrix,
*M**i**i* = *A**i* where *A**i* is the
Voronoi area associated with vertex *i*.
------------------------------------------------------------------------
source
### mass_matrix_inv_sparse
``` python
def mass_matrix_inv_sparse(
vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh
)->BCOO: # Diagonal sparse inverse mass matrix.
```
*Assemble inverse lumped mass matrix as a sparse matrix (BCOO).*
------------------------------------------------------------------------
source
### mass_matrix_sparse
``` python
def mass_matrix_sparse(
vertices:Float[Array, 'n_vertices dim'], # Vertex positions.
hemesh:HeMesh, # Half-edge mesh connectivity.
)->BCOO: # Diagonal sparse mass matrix.
```
*Assemble lumped (diagonal) mass matrix as a sparse matrix (BCOO).*
The lumped mass matrix is diagonal with entries equal to the Voronoi
dual area of each vertex: *M**i**i* = *A**i*.
``` python
# Test mass matrix against igl.massmatrix (Voronoi type)
M_jax = mass_matrix_sparse(geommesh.vertices, hemesh)
M_igl = igl.massmatrix(np.asarray(geommesh.vertices), np.asarray(hemesh.faces), igl.MASSMATRIX_TYPE_VORONOI)
rel_err_mass = np.linalg.norm(M_jax.todense() - M_igl.todense()) / np.linalg.norm(M_igl.todense())
print("mass matrix rel. error:", rel_err_mass)
# Test inverse
M_inv_jax = mass_matrix_inv_sparse(geommesh.vertices, hemesh)
identity_check = (M_jax @ M_inv_jax.todense()).todense()
print("M @ M_inv ~ I error:", np.linalg.norm(identity_check - np.eye(hemesh.n_vertices)))
```
### Finite-element gradient
Not to be confused with the discrete-exterior-calculus operators, which
only depend on mesh connectivity, not geometry.
------------------------------------------------------------------------
source
### compute_gradient_3d
``` python
def compute_gradient_3d(
vertices:Float[Array, 'n_vertices 3'], hemesh:HeMesh, vertex_field:Float[Array, 'n_vertices ...']
)->Float[Array, 'n_faces 3 ...']:
```
*Compute the linear finite-element gradient (constant per face).*
------------------------------------------------------------------------
source
### compute_gradient_2d
``` python
def compute_gradient_2d(
vertices:Float[Array, 'n_vertices 2'], hemesh:HeMesh, vertex_field:Float[Array, 'n_vertices ...']
)->Float[Array, 'n_faces 2 ...']:
```
*Compute the linear finite-element gradient (constant per face).*
------------------------------------------------------------------------
source
### gradient_sparse_3d
``` python
def gradient_sparse_3d(
vertices:Float[Array, 'n_vertices 3'], hemesh:HeMesh
)->BCOO:
```
*Assemble FE gradient in 3D as a sparse matrix (BCOO).*
Returns a matrix G with shape (3*n_faces, n_vertices) such that for a
scalar per-vertex field u (n_vertices,), the per-face gradients are
obtained via: g_flat = G @ u \# (3*n_faces,) g = g_flat.reshape((3,
n_faces)).T \# (n_faces, 3)
This row layout matches libigl’s `grad` operator convention (component
blocks).
------------------------------------------------------------------------
source
### gradient_sparse_2d
``` python
def gradient_sparse_2d(
vertices:Float[Array, 'n_vertices 2'], hemesh:HeMesh
)->BCOO:
```
*Assemble FE gradient in 2D as a sparse matrix (BCOO).*
Returns a matrix G with shape (2*n_faces, n_vertices) such that for a
scalar per-vertex field u (n_vertices,), the per-face gradients are
obtained via: g_flat = G @ u \# (2*n_faces,) g = g_flat.reshape((2,
n_faces)).T \# (n_faces, 2)
This row layout matches libigl’s `grad` operator convention (component
blocks).
------------------------------------------------------------------------
source
### reshape_face_gradient
``` python
def reshape_face_gradient(
grad_flat:Float[Array, 'dim_n_faces ...'], # Output of `G @ u`, with shape `(dim*n_faces, ...)`.
n_faces:int, # Number of mesh faces.
dim:int, # Spatial dimension (2 or 3).
)->Float[Array, 'n_faces dim ...']: # Reshaped gradient with shape `(n_faces, dim, ...)`, matching the output convention
of `compute_gradient_2d/3d`.
```
*Reshape a flattened FE gradient into per-face vectors.*
This is meant to be used with `gradient_sparse_2d/3d` (and any similar
operator that stacks components in blocks), where applying the sparse
matrix yields an array of shape `(dim*n_faces, ...)` (for
scalar/vector/tensor per-vertex fields).
``` python
# here's how to compute the gradient in libigl
grad_matrix = igl.grad(np.asarray(geommesh.vertices), np.asarray(hemesh.faces))
# calculate the gradient of field by matrix multiplication
grad_igl = grad_matrix @ np.asarray(u)
# order='F' copied from igl tutorial
grad_igl = grad_igl.reshape((hemesh.n_faces, geommesh.dim), order='F')
```
``` python
# test jax and libigl implementations
grad_jax = compute_gradient_2d(geommesh.vertices, hemesh, u)
rel_err_grad = np.linalg.norm(np.asarray(grad_jax) - grad_igl) / np.linalg.norm(grad_igl)
print("gradient rel. error:", rel_err_grad)
```
gradient rel. error: 1.413315746703021e-16
``` python
# same test, in 3d
grad_matrix_3d = igl.grad(np.asarray(geommesh_3d.vertices), np.asarray(hemesh.faces))
grad_igl_3d = grad_matrix_3d @ np.asarray(u)
grad_igl_3d = grad_igl_3d.reshape((hemesh.n_faces, geommesh_3d.dim), order='F')
grad_jax_3d = compute_gradient_3d(geommesh_3d.vertices, hemesh, u)
rel_err_grad_3d = np.linalg.norm(np.asarray(grad_jax_3d) - grad_igl_3d) / np.linalg.norm(grad_igl_3d)
print("gradient rel. error:", rel_err_grad_3d)
```
gradient rel. error: 1.5657863888820882e-16
``` python
# Test sparse gradient operators vs apply functions
key = jax.random.PRNGKey(123)
u_test = jax.random.normal(key, (hemesh.n_vertices,))
G2 = gradient_sparse_2d(geommesh.vertices, hemesh)
g2 = reshape_face_gradient(G2 @ u_test, hemesh.n_faces, dim=2)
g2_apply = compute_gradient_2d(geommesh.vertices, hemesh, u_test)
rel_err_g2 = jnp.linalg.norm(g2 - g2_apply) / jnp.linalg.norm(g2_apply)
print("2D grad sparse vs apply rel. error:", rel_err_g2)
G3 = gradient_sparse_3d(geommesh_3d.vertices, hemesh)
g3 = reshape_face_gradient(G3 @ u_test, hemesh.n_faces, dim=3)
g3_apply = compute_gradient_3d(geommesh_3d.vertices, hemesh, u_test)
rel_err_g3 = jnp.linalg.norm(g3 - g3_apply) / jnp.linalg.norm(g3_apply)
print("3D grad sparse vs apply rel. error:", rel_err_g3)
# quick sanity check for vector/tensor fields: u has extra axes
u_vec = jax.random.normal(key, (hemesh.n_vertices, 3))
g2_vec = reshape_face_gradient(G2 @ u_vec, hemesh.n_faces, dim=2)
g2_vec_apply = compute_gradient_2d(geommesh.vertices, hemesh, u_vec)
rel_err_g2_vec = jnp.linalg.norm(g2_vec - g2_vec_apply) / jnp.linalg.norm(g2_vec_apply)
print("2D grad (vector field) sparse vs apply rel. error:", rel_err_g2_vec)
```
2D grad sparse vs apply rel. error: 8.71017994729607e-17
3D grad sparse vs apply rel. error: 9.602668379845331e-17
2D grad (vector field) sparse vs apply rel. error: 8.285943150518157e-17
### Wrapping as linear operators
It’s often useful to think of functions like
[`compute_cotan_laplace()`](https://nikolas-claussen.github.io/triangulax/src/linear_operators.html#compute_cotan_laplace)
as a linear operator on fields on meshes. For example, imagine you want
to solve the Laplace equation on a mesh with fixed vertex positions and
connectivity. You will want to use a linear solver. Luckily, most such
solvers only need to be able to compute the action of a linear operator
on an input vector, and don’t need an explicit matrix representation.
In the JAX ecosystem, the `lineax` library defines linear solvers. We
can *wrap*
[`compute_cotan_laplace()`](https://nikolas-claussen.github.io/triangulax/src/linear_operators.html#compute_cotan_laplace)
as a linear operator, which allows us to pass it into iterative linear
algebra algorithms.
``` python
# "bake in" the connectivity and vertex positions
laplace_op = functools.partial(compute_cotan_laplace, geommesh.vertices, hemesh)
_ = laplace_op(u) # you can apply this to vertex-fields
# define the linear operator
laplace_op_lx = lineax.FunctionLinearOperator(laplace_op, input_structure=jax.eval_shape(laplace_op, u))
# now you can use the linear operator to compute matrix representations, solve linear systems, etc.
mat = laplace_op_lx.as_matrix()
mat.shape
```
(131, 131)
------------------------------------------------------------------------
source
### linear_op_to_sparse
``` python
def linear_op_to_sparse(
op:callable, in_shape:tuple, out_shape:tuple, dtype:Union=None, chunk_size:int=256, tol:float=0.0
)->BCOO:
```
*Build a sparse matrix for a linear map using batched one-hot probes.*
Note: this function is general, but not necessarily very efficient for
large matrix sizes.
``` python
# compare sparse construction to lineax dense matrix (small meshes only)
if hemesh.n_vertices <= 2000:
laplace_op_local = functools.partial(compute_cotan_laplace, geommesh.vertices, hemesh)
laplace_op_lx_local = lineax.FunctionLinearOperator(laplace_op_local,
input_structure=jax.eval_shape(laplace_op_local, u))
sp_mat = linear_op_to_sparse(laplace_op_local, (hemesh.n_vertices,), (hemesh.n_vertices,))
mat_dense = laplace_op_lx_local.as_matrix()
rel_err_sparse = jnp.linalg.norm(sp_mat.todense() - mat_dense) / jnp.linalg.norm(mat_dense)
print("sparse vs lineax rel. error:", rel_err_sparse)
else:
print("Skipping dense comparison for large mesh.")
```
sparse vs lineax rel. error: 0.0
``` python
## now let's try with a large mesh
mesh = TriMesh.read_obj("test_meshes/torus_high_resolution.obj")
hemesh = msh.HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)
geommesh = msh.GeomMesh(*hemesh.n_items, mesh.vertices, mesh.face_positions)
laplace_op = jax.jit(functools.partial(compute_cotan_laplace, geommesh.vertices, hemesh))
```
Warning: readOBJ() ignored non-comment line 3:
o Torus
``` python
sparse_laplace_op = linear_op_to_sparse(laplace_op, (hemesh.n_vertices,), (hemesh.n_vertices,))
```