linops: discrete gradient, divergence, and Laplacian

Building on the mesh geometry and adjacency-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 implement gradient (per-vertex scalar field -> per-face vector field) and the cotan-Laplacian (vertex -> vertex). Both cases start with a scalar field \(u_i\) defined per vertex \(i\) of the triangulation. The finite-element gradient is defined for each face \(ijk\), like so: \[ (\nabla u)_{ijk} = \sum_{l\in \{i,j,k\}} u_l \nabla\phi_l \] where \(\phi_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_{ijk}\) is the triangle area, \(\mathbf{v}_i\) are the vertex positions, and \(()^\perp\) 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 \(\alpha_j, \beta_j\) are the two triangle angles “opposite” to the edge \(ij\).

To check for correctness, we can compare with this libigl tutorial, using the test mesh and some random test fields.

# load test data

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = msh.HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)
geommesh = msh.GeomMesh(mesh.vertices, mesh.face_positions)

mesh_3d = TriMesh.read_obj("../test_meshes/disk.obj", dim=3)
geommesh_3d = msh.GeomMesh(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

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

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

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

The cotangent Laplacian is the standard discretization of the Laplace-Beltrami operator on triangle meshes. It is negative semi-definite in our sign convention: for any field \(u\),

\[ \sum_i u_i (\Delta u)_i = -\sum_{ij} w_{ij}(u_i - u_j)^2 \leq 0 \]

where \(w_{ij} = \frac{1}{2}(\cot\alpha_{ij} + \cot\beta_{ij})\) are the cotangent edge weights. On boundary edges only one opposite angle contributes.

The weights depend on the vertex positions only through the edge lengths, so the operator is intrinsic (see notebook 05): the implementation takes edge lengths, and the position-based version simply feeds it geometry.get_he_length. Passing edge lengths measured with a periodic displacement function instead gives the Laplacian on a torus, at no extra implementation cost - see “Periodic boundary conditions” below.

We provide, for each of the two, an operator and a matrix form: - compute_cotan_laplace_intrinsic / compute_cotan_laplace: applies \(L\) to a vertex field via gather/scatter (works with jax.jit and jax.grad). - cotan_laplace_sparse_intrinsic / cotan_laplace_sparse: assembles \(L\) as a sparse BCOO matrix.

---

source

compute_cotan_laplace

def compute_cotan_laplace(
    vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh, vertex_field:Float[Array, 'n_vertices ...'],
    normalize:bool=False
)->Float[Array, 'n_vertices ...']:

Compute cotangent Laplacian of a per-vertex field (natural boundary conditions).

Position-based form of compute_cotan_laplace_intrinsic; see there for the parameters.

---

source

compute_cotan_laplace_intrinsic

def compute_cotan_laplace_intrinsic(
    he_lengths:Float[Array, 'n_hes'], # Edge length per half-edge, e.g. from `geometry.get_he_length` or, under periodic
boundary conditions, `periodic.get_periodic_he_lengths`.
    hemesh:HeMesh, # Half-edge mesh connectivity.
    vertex_field:Float[Array, 'n_vertices ...'], # Per-vertex scalar, vector, or tensor field.
    normalize:bool=False, # If True, return the area-normalized Laplace-Beltrami operator $M^{-1} L u$,
dividing by the robust Voronoi cell area at each vertex.
)->Float[Array, 'n_vertices ...']: # Cotangent Laplacian applied to the field, same shape as ``vertex_field``.

Compute the cotangent Laplacian of a per-vertex field from edge lengths (natural boundary conditions).

---

source

cotan_laplace_sparse

def cotan_laplace_sparse(
    vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh
)->BCOO:

Assemble cotangent Laplacian as a sparse matrix (BCOO).

---

source

cotan_laplace_sparse_intrinsic

def cotan_laplace_sparse_intrinsic(
    he_lengths:Float[Array, 'n_hes'], hemesh:HeMesh
)->BCOO:

Assemble the cotangent Laplacian from edge lengths as a sparse matrix (BCOO).

Same sign convention as compute_cotan_laplace_intrinsic, so cotan_laplace_sparse_intrinsic(l, h) @ u == compute_cotan_laplace_intrinsic(l, h, u).

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

assert rel_err < 1e-10, rel_err
assert rel_err_vec < 1e-10, rel_err_vec
# the sign convention must match igl's (negative semi-definite), not just the magnitude
assert np.allclose(np.asarray(lap_jax), lap_igl, atol=1e-10)
# constants are in the kernel, including at boundary vertices
assert jnp.allclose(compute_cotan_laplace(geommesh.vertices, hemesh, jnp.ones(hemesh.n_vertices)), 0, atol=1e-10)

scalar field rel. error: 1.7329341320735597e-16
vector field rel. error: 1.9265205796945205e-16

apply_and_normalize = compute_cotan_laplace(geommesh.vertices, hemesh, u_vec, normalize=True)

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

assert rel_err_sparse < 1e-10, rel_err_sparse
# the sparse operator must be symmetric with vanishing row sums
assert jnp.allclose(L_sparse.todense(), L_sparse.todense().T, atol=1e-12)
assert jnp.allclose(L_sparse.todense().sum(axis=1), 0, atol=1e-10)

cotan sparse vs apply rel. error: 1.7496118652518946e-16

# scipy <-> BCOO round trip must be exact
assert (scipy_to_bcoo(bcoo_to_scipy(L_sparse)).todense() == L_sparse.todense()).all()
assert bcoo_to_scipy(L_sparse).shape == L_sparse.shape

# diag_jsparse must match np.diag for zero and non-zero offsets, and be jittable
for k in [0, 1, -1, 2]:
    assert jnp.allclose(diag_jsparse(jnp.arange(1., 6.), k).todense(), jnp.diag(jnp.arange(1., 6.), k))
assert jnp.allclose(jax.jit(diag_jsparse)(jnp.arange(1., 6.)).todense(), jnp.diag(jnp.arange(1., 6.)))

Mass matrix (lumped)

The finite-element mass matrix \(M\) appears whenever we discretize a time-dependent PDE. For a lumped (diagonal) mass matrix, \(M_{ii} = A_i\) where \(A_i\) is the Voronoi area associated with vertex \(i\).

---

source

mass_matrix_inv_sparse

def mass_matrix_inv_sparse(
    vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh, area_type:str='voronoi'
)->BCOO:

Assemble the inverse lumped mass matrix as a sparse matrix (BCOO).

Position-based form of mass_matrix_inv_sparse_intrinsic.

---

source

mass_matrix_sparse

def mass_matrix_sparse(
    vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh, area_type:str='voronoi'
)->BCOO:

Assemble lumped (diagonal) mass matrix as a sparse matrix (BCOO).

Position-based form of mass_matrix_sparse_intrinsic; see there for the parameters and the meaning of area_type.

---

source

mass_matrix_inv_sparse_intrinsic

def mass_matrix_inv_sparse_intrinsic(
    he_lengths:Float[Array, 'n_hes'], hemesh:HeMesh, area_type:str='voronoi'
)->BCOO:

Inverse of mass_matrix_sparse_intrinsic. See there for the area conventions.

Note only "voronoi" (mixed) areas are guaranteed positive; the other choices can produce zero or negative areas, and hence a divergent inverse.

---

source

mass_matrix_sparse_intrinsic

def mass_matrix_sparse_intrinsic(
    he_lengths:Float[Array, 'n_hes'], # Edge length per half-edge, e.g. from `geometry.get_he_length` or, under periodic
boundary conditions, `periodic.get_periodic_he_lengths`.
    hemesh:HeMesh, # Half-edge mesh connectivity.
    area_type:str='voronoi', # Choice of dual-cell area definition used on the diagonal.
``"voronoi"`` is the *mixed* Voronoi area of Meyer et al., which is what
``igl.massmatrix(..., MASSMATRIX_TYPE_VORONOI)`` computes and is always
positive; this is the right choice for the cotangent Laplacian.
``"voronoi_exact"`` is the exact (signed) circumcentric area, which can be
NEGATIVE on obtuse triangles and then destroys positive-definiteness of
``M - dt*L``. ``"barycentric"`` is a simpler always-positive approximation.
)->BCOO: # Diagonal sparse mass matrix.

Assemble the lumped (diagonal) mass matrix from edge lengths as a sparse matrix (BCOO).

The lumped mass matrix is diagonal with entries equal to the dual area of each vertex: \(M_{ii} = A_i\).

# 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()
print("M @ M_inv ~ I error:", np.linalg.norm(identity_check - np.eye(hemesh.n_vertices)))

assert rel_err_mass < 1e-12, rel_err_mass
assert np.linalg.norm(identity_check - np.eye(hemesh.n_vertices)) < 1e-10

# The default 'voronoi' area is Meyer's robust Voronoi area, which is what igl computes and is
# always positive. On a mesh with obtuse triangles the exact circumcentric area differs.
mesh_obtuse = TriMesh.read_obj("../test_meshes/sphere_fine_poor.obj", dim=3)
hemesh_obtuse = msh.HeMesh.from_triangles(mesh_obtuse.vertices.shape[0], mesh_obtuse.faces)
M_obtuse = mass_matrix_sparse(mesh_obtuse.vertices, hemesh_obtuse).todense().diagonal()
M_obtuse_igl = igl.massmatrix(np.asarray(mesh_obtuse.vertices), np.asarray(hemesh_obtuse.faces),
                              igl.MASSMATRIX_TYPE_VORONOI).diagonal()
assert np.abs(M_obtuse - M_obtuse_igl).max() < 1e-12
assert (M_obtuse > 0).all()

try:
    mass_matrix_sparse(geommesh.vertices, hemesh, area_type='nonsense')
    raise AssertionError('expected ValueError')
except ValueError:
    pass

mass matrix rel. error: 1.53180380201344e-16
M @ M_inv ~ I error: 4.839349969133127e-16

Warning: readOBJ() ignored non-comment line 3:
  o Icosphere

Periodic boundary conditions

On a periodic domain (a torus, possibly sheared - see notebook 05b) a difference of vertex coordinates across the boundary is meaningless, the position-based operators above cannot be used directly. Their intrinsic forms can: the cotangent Laplacian, the mass matrix and the normal derivative all depend on the geometry only through edge lengths, and edge lengths are well-defined periodically:

displacement_fn = lambda r_1, r_2: per.displacement_periodic(r_1, r_2, L)
he_lengths = per.get_periodic_he_lengths(vertices, hemesh, displacement_fn)

lap = compute_cotan_laplace_intrinsic(he_lengths, hemesh, u)
L_sparse = cotan_laplace_sparse_intrinsic(he_lengths, hemesh)
M = mass_matrix_sparse_intrinsic(he_lengths, hemesh)

The same applies to the intrinsic geometry functions of the geometry module, so an implicit diffusion step, a curvature computation, or a vertex-model energy can be assembled on a periodic mesh with the code used everywhere else. Only genuinely position-dependent operators - the finite-element gradient and divergence below, which return per-face vectors - have no intrinsic form and are not available this way.

# Operators on an actual periodic mesh, via edge lengths measured with a displacement function
mesh_per = TriMesh.read_obj("../test_meshes/torus_2d.obj", dim=2)
hemesh_per = msh.HeMesh.from_triangles(mesh_per.vertices.shape[0], mesh_per.faces)
L_box = jnp.array([1., 1.])
disp = lambda r_1, r_2: per.displacement_periodic(r_1, r_2, L_box)
l_per = per.get_periodic_he_lengths(mesh_per.vertices, hemesh_per, disp)
u_per = jax.random.normal(jax.random.PRNGKey(5), (hemesh_per.n_vertices,))

# constants are in the kernel, and a Fourier mode is an approximate eigenfunction
# with eigenvalue -2*(2*pi)^2 -- the boundary wrap is handled correctly
ones = jnp.ones(hemesh_per.n_vertices)
assert jnp.allclose(compute_cotan_laplace_intrinsic(l_per, hemesh_per, ones), 0, atol=1e-10)

f = jnp.sin(2*jnp.pi*mesh_per.vertices[:, 0]) * jnp.cos(2*jnp.pi*mesh_per.vertices[:, 1])
lap_f = compute_cotan_laplace_intrinsic(l_per, hemesh_per, f, normalize=True)
exact = -2 * (2*jnp.pi)**2 * f
rel_err_per = float(jnp.linalg.norm(lap_f - exact) / jnp.linalg.norm(exact))
print("periodic Laplacian rel. error vs -2k^2 f:", rel_err_per)
assert rel_err_per < 0.05, rel_err_per

# sparse form: agrees with the matrix-free one, symmetric, negative semi-definite
Lp = cotan_laplace_sparse_intrinsic(l_per, hemesh_per)
assert jnp.allclose(Lp @ u_per, compute_cotan_laplace_intrinsic(l_per, hemesh_per, u_per), atol=1e-10)
dense_p = Lp.todense()
assert jnp.allclose(dense_p, dense_p.T, atol=1e-12)
assert jnp.allclose(dense_p.sum(axis=1), 0, atol=1e-10)
assert jnp.linalg.eigvalsh(dense_p).max() < 1e-10

# mass matrix: positive, sums to the total area, and (M - dt*L) is positive definite,
# so an implicit diffusion step can be assembled on a periodic mesh
Mp = mass_matrix_sparse_intrinsic(l_per, hemesh_per)
assert jnp.all(Mp.data > 0)
assert jnp.allclose(Mp.data.sum(), geom.get_area_intrinsic(l_per, hemesh_per), rtol=1e-10)
Mp_inv = mass_matrix_inv_sparse_intrinsic(l_per, hemesh_per)
assert jnp.allclose(Mp.todense() @ Mp_inv.todense(), jnp.eye(hemesh_per.n_vertices), atol=1e-10)
assert jnp.linalg.eigvalsh(Mp.todense() - 0.01 * dense_p).min() > 0

# with a box much larger than the mesh the wrap never triggers, so every periodic operator
# must reduce exactly to its position-based counterpart on an ordinary mesh
disp_big = lambda r_1, r_2: per.displacement_periodic(r_1, r_2, jnp.array([1e6, 1e6]))
l_big = per.get_periodic_he_lengths(geommesh.vertices, hemesh, disp_big)
u_big = jax.random.normal(jax.random.PRNGKey(3), (hemesh.n_vertices,))
assert jnp.allclose(compute_cotan_laplace_intrinsic(l_big, hemesh, u_big),
                    compute_cotan_laplace(geommesh.vertices, hemesh, u_big), atol=1e-8)
assert jnp.allclose(cotan_laplace_sparse_intrinsic(l_big, hemesh).todense(),
                    cotan_laplace_sparse(geommesh.vertices, hemesh).todense(), atol=1e-8)
assert jnp.allclose(mass_matrix_sparse_intrinsic(l_big, hemesh).data,
                    mass_matrix_sparse(geommesh.vertices, hemesh).data, atol=1e-8)
print("periodic Laplacian and mass matrix OK; (M - dt L) is positive definite")

periodic Laplacian rel. error vs -2k^2 f: 0.012032811490690486
periodic Laplacian and mass matrix OK; (M - dt L) is positive definite

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

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

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

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 (3n_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 # (3n_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

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 (2n_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 # (2n_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

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

# 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')

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

assert rel_err_grad < 1e-10, rel_err_grad
# the FE gradient is exact on a linear field
lin = geommesh.vertices @ jnp.array([2., -3.])
assert jnp.allclose(compute_gradient_2d(geommesh.vertices, hemesh, lin), jnp.array([2., -3.]), atol=1e-10)

gradient rel. error: 1.413315746703021e-16

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

assert rel_err_grad_3d < 1e-10, rel_err_grad_3d

gradient rel. error: 1.5657863888820882e-16

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

assert rel_err_g2 < 1e-10, rel_err_g2
assert rel_err_g3 < 1e-10, rel_err_g3
assert rel_err_g2_vec < 1e-10, 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

Divergence

The discrete divergence maps a per-face vector field to a per-vertex scalar field. It is the negative adjoint of the gradient with respect to the area-weighted inner product:

\[ (\mathrm{div}\, V)_i = -\sum_{f \ni i} A_f \, V_f \cdot \nabla\phi_i^{(f)} \]

where \(A_f\) is the face area. With this sign convention, the composed operator satisfies \(L = \mathrm{div} \circ \nabla\) (the negative semi-definite cotan-Laplacian). In matrix form, div = -G^T @ diag(face_areas).

---

source

compute_divergence_3d

def compute_divergence_3d(
    vertices:Float[Array, 'n_vertices 3'], hemesh:HeMesh, face_field:Float[Array, 'n_faces 3 ...']
)->Float[Array, 'n_vertices ...']:

Compute the (integrated) FE divergence of a per-face vector field (3D).

Same as :func:compute_divergence_2d but for surfaces embedded in 3D.

---

source

compute_divergence_2d

def compute_divergence_2d(
    vertices:Float[Array, 'n_vertices 2'], # Vertex positions, shape (n_vertices, 2).
    hemesh:HeMesh, # Half-edge mesh connectivity.
    face_field:Float[Array, 'n_faces 2 ...'], # Per-face vector (or tensor) field, shape (n_faces, 2, ...).
)->Float[Array, 'n_vertices ...']: # Per-vertex field, shape (n_vertices, ...).

Compute the (integrated) FE divergence of a per-face vector field (2D).

Maps a per-face vector field to a per-vertex scalar field. This is the negative adjoint of the FE gradient weighted by face areas. Satisfies compute_cotan_laplace(v, h, u) ≈ compute_divergence_2d(v, h, compute_gradient_2d(v, h, u)).

# Test: div(grad u) == cotan_laplace(u)
key = jax.random.PRNGKey(42)
u_div = jax.random.normal(key, (hemesh.n_vertices,))

lap_direct = compute_cotan_laplace(geommesh.vertices, hemesh, u_div)
grad_u = compute_gradient_2d(geommesh.vertices, hemesh, u_div)
lap_via_div = compute_divergence_2d(geommesh.vertices, hemesh, grad_u)

rel_err_div = jnp.linalg.norm(lap_direct - lap_via_div) / jnp.linalg.norm(lap_direct)
print("div(grad u) vs L u rel. error:", rel_err_div)

# same in 3D
grad_u_3d = compute_gradient_3d(geommesh_3d.vertices, hemesh, u_div)
lap_3d_direct = compute_cotan_laplace(geommesh_3d.vertices, hemesh, u_div)
lap_3d_via_div = compute_divergence_3d(geommesh_3d.vertices, hemesh, grad_u_3d)

rel_err_div_3d = jnp.linalg.norm(lap_3d_direct - lap_3d_via_div) / jnp.linalg.norm(lap_3d_direct)
print("div(grad u) vs L u rel. error (3D):", rel_err_div_3d)

assert rel_err_div < 1e-10, rel_err_div
assert rel_err_div_3d < 1e-10, rel_err_div_3d

# divergence must also accept the vector/tensor face fields that the gradient produces
u_vec_div = geommesh.vertices
grad_vec = compute_gradient_2d(geommesh.vertices, hemesh, u_vec_div)
div_vec = compute_divergence_2d(geommesh.vertices, hemesh, grad_vec)
assert div_vec.shape == u_vec_div.shape
assert jnp.allclose(div_vec, compute_cotan_laplace(geommesh.vertices, hemesh, u_vec_div), atol=1e-10)

# adjointness: <grad u, grad u>_A == -<u, div grad u>
areas = geom.get_triangle_areas(geommesh.vertices, hemesh)
lhs = float(((grad_u * grad_u).sum(-1) * areas).sum())
rhs = -float((u_div * lap_via_div).sum())
assert abs(lhs - rhs) / abs(lhs) < 1e-10, (lhs, rhs)

div(grad u) vs L u rel. error: 2.1665711288568335e-16
div(grad u) vs L u rel. error (3D): 1.8799961803543857e-16

Normal derivative (flux across an edge)

The normal derivative gives the flux of \(\nabla u\) across each half-edge, integrated along the dual edge segment inside the corresponding face. It is the per-half-edge quantity from which the cotangent Laplacian is assembled. Use the normal derivative to impose Neumann boundary conditions or to evaluate a discrete divergence theorem over a sub-region.

Compare igl.normal_derivative.

---

source

normal_derivative_sparse

def normal_derivative_sparse(
    vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh
)->BCOO:

Sparse matrix form of compute_normal_derivative, of shape (n_hes, n_vertices).

normal_derivative_sparse(v, h) @ u == compute_normal_derivative(v, h, u).

---

source

compute_normal_derivative

def compute_normal_derivative(
    vertices:Float[Array, 'n_vertices dim'], hemesh:HeMesh, vertex_field:Float[Array, 'n_vertices ...']
)->Float[Array, 'n_hes ...']:

Integrated normal derivative (flux) of a vertex field across each half-edge.

Position-based form of compute_normal_derivative_intrinsic; see there for details.

---

source

normal_derivative_sparse_intrinsic

def normal_derivative_sparse_intrinsic(
    he_lengths:Float[Array, 'n_hes'], hemesh:HeMesh
)->BCOO:

Sparse matrix form of compute_normal_derivative_intrinsic, of shape (n_hes, n_vertices).

---

source

compute_normal_derivative_intrinsic

def compute_normal_derivative_intrinsic(
    he_lengths:Float[Array, 'n_hes'], # Edge length per half-edge, e.g. from `geometry.get_he_length` or, under periodic
boundary conditions, `periodic.get_periodic_he_lengths`.
    hemesh:HeMesh, # Half-edge mesh.
    vertex_field:Float[Array, 'n_vertices ...'], # Per-vertex scalar, vector, or tensor field.
)->Float[Array, 'n_hes ...']: # Integrated normal derivative per half-edge, 0 on boundary half-edges.

Integrated normal derivative (flux) of a vertex field across each half-edge, from edge lengths.

For half-edge he inside a face, this is the flux of the piecewise-linear gradient through the dual edge segment of that face: D[he] = cot(theta_opposite(he)) / 2 * (u[dest[he]] - u[orig[he]]), where theta_opposite is the corner angle opposite he. It is 0 on boundary half-edges, which have no face.

This is the per-half-edge quantity the cotangent Laplacian is assembled from. The twin half-edge measures the same edge from the adjacent face and with the opposite sign of u[dest] - u[orig], so the total flux across an edge is D - D[twin] and

``compute_cotan_laplace_intrinsic(l, h, u) == -sum_he_to_vertex_incoming(h, D - D[twin])``

holds to machine precision. Note the raw sum of D over all half-edges is not zero: the two half-edges of an edge carry different opposite angles, and it is the combination D - D[twin] that cancels pairwise. Summing that over a set of vertices gives the net flux across the set’s boundary (the discrete divergence theorem); over a closed mesh it is zero.

# the normal derivative is the per-half-edge quantity the Laplacian is assembled from
u_nd = jax.random.normal(jax.random.PRNGKey(11), (hemesh.n_vertices,))
D = compute_normal_derivative(geommesh.vertices, hemesh, u_nd)
assert D.shape == (hemesh.n_hes,)
assert jnp.all(D[hemesh.is_bdry_he] == 0.0)          # no face -> no flux

# combining the two face contributions of each edge recovers the cotan Laplacian.
# D[twin] carries the opposite sign of (u[dest] - u[orig]), hence the minus.
lap_from_flux = -adj.sum_he_to_vertex_incoming(hemesh, D - D[hemesh.twin])
assert jnp.allclose(lap_from_flux, compute_cotan_laplace(geommesh.vertices, hemesh, u_nd), atol=1e-10)

# the sparse form agrees with the matrix-free one, for scalar and vector fields
Dm = normal_derivative_sparse(geommesh.vertices, hemesh)
assert Dm.shape == (hemesh.n_hes, hemesh.n_vertices)
assert jnp.allclose(Dm @ u_nd, D, atol=1e-12)
u_nd_vec = jax.random.normal(jax.random.PRNGKey(12), (hemesh.n_vertices, 3))
assert jnp.allclose(Dm @ u_nd_vec, compute_normal_derivative(geommesh.vertices, hemesh, u_nd_vec), atol=1e-12)

# a constant field has zero flux everywhere
assert jnp.allclose(compute_normal_derivative(geommesh.vertices, hemesh, jnp.ones(hemesh.n_vertices)), 0, atol=1e-12)

# divergence theorem on a CLOSED mesh: the total integrated Laplacian vanishes, because
# each edge's flux enters its two endpoints with opposite signs.
_cm = TriMesh.read_obj("../test_meshes/sphere.obj", dim=3)
_ch = msh.HeMesh.from_triangles(_cm.vertices.shape[0], _cm.faces)
u_c = jax.random.normal(jax.random.PRNGKey(13), (_ch.n_vertices,))
D_c = compute_normal_derivative(_cm.vertices, _ch, u_c)
flux_div = -adj.sum_he_to_vertex_incoming(_ch, D_c - D_c[_ch.twin])
assert jnp.allclose(flux_div.sum(), 0.0, atol=1e-10)
assert jnp.allclose(flux_div, compute_cotan_laplace(_cm.vertices, _ch, u_c), atol=1e-10)
print("normal derivative: reproduces the cotan Laplacian and the sparse form matches")

Warning: readOBJ() ignored non-comment line 3:
  o Icosphere

normal derivative: reproduces the cotan Laplacian and the sparse form matches

Wrapping as linear operators

It’s often useful to think of functions like 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() as a linear operator, which allows us to pass it into iterative linear algebra algorithms.

# "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)

Example: implicit diffusion step

A common use case in simulation is the implicit time step for the heat equation on a mesh. Given per-vertex temperatures \(u^n\) and a time step \(\Delta t\), one solves

\[ (M - \Delta t \, L)\, u^{n+1} = M\, u^n \]

where \(M\) is the mass matrix and \(L\) the cotan-Laplacian. Since \(L\) is negative semi-definite, the system matrix \(M - \Delta t\, L\) is positive definite and can be solved with a conjugate-gradient solver from lineax.

# Implicit diffusion step on the disk mesh

dt = 0.01
M = mass_matrix_sparse(geommesh.vertices, hemesh)
L = cotan_laplace_sparse(geommesh.vertices, hemesh)

# Initial condition: random temperature field
key = jax.random.PRNGKey(7)
u0 = jax.random.normal(key, (hemesh.n_vertices,))

# Right-hand side: M u^n
rhs = M @ u0

# System matrix A = M - dt * L  (positive definite)
# Wrap as a lineax operator so we can use an iterative solver.
def apply_A(x):
    return M @ x - dt * (L @ x)

A_op = lineax.FunctionLinearOperator(apply_A, input_structure=jax.eval_shape(apply_A, u0),
                                     tags=lineax.positive_semidefinite_tag)

# Solve with conjugate gradient
u1 = lineax.linear_solve(A_op, rhs, solver=lineax.CG(rtol=1e-6, atol=1e-10)).value
print("implicit Euler step completed, |u1|_2 =", jnp.linalg.norm(u1))

implicit Euler step completed, |u1|_2 = 5.576995473050827

---

source

linear_op_to_sparse

def linear_op_to_sparse(
    op:Callable, # Linear map taking and returning 1d arrays.
    in_shape:tuple, out_shape:tuple, dtype:Union=None, # Output dtype. Inferred from ``op`` if None.
    chunk_size:int=256, # Number of one-hot probes evaluated per batch.
    tol:float=0.0, # Entries with ``|value| <= tol`` are dropped.
)->BCOO: # Sparse matrix of shape ``(n_out, n_in)``.

Build a sparse matrix for a linear map using batched one-hot probes.

# 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.")
    assert rel_err_sparse < 1e-10, rel_err_sparse

sparse vs lineax rel. error: 0.0

## 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(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

# ~25 s on this 36k-vertex mesh

sparse_laplace_op = linear_op_to_sparse(laplace_op, (hemesh.n_vertices,), (hemesh.n_vertices,))