# the reference areas computed from the metric agree with the triangle areas
metric = get_metric(sphere.vertices, sphere_hemesh)
assert np.allclose(jnp.sqrt(jnp.linalg.det(metric)) / 2,
geom.get_triangle_areas(sphere.vertices, sphere_hemesh))
# diagonal entries are the squared side lengths |u|^2, |v|^2
a, b, c = sphere.vertices[sphere_hemesh.faces.T]
assert np.allclose(metric[:, 0, 0], ((b-a)**2).sum(axis=1))
assert np.allclose(metric[:, 1, 1], ((c-a)**2).sum(axis=1))
assert np.allclose(metric[:, 0, 1], ((b-a)*(c-a)).sum(axis=1))
# the closed-form inverse agrees with jnp.linalg.inv
inverse_metric = get_two_x_two_inverse(metric)
assert np.allclose(inverse_metric @ metric, jnp.eye(2), atol=1e-10)
assert np.allclose(inverse_metric, jnp.linalg.inv(metric))elastic - Discrete elastic energies
This module defines elastic energies for thin shells, bilayers, and membranes for physics-based simulations. The energies are expressed directly using “discrete” quantities (like dihedral angles) which makes their computation (and optimization) on triangular meshes efficient.
All energies share the signature energy(vertices, args) -> scalar, where args is a tuple holding the mesh connectivity and the material parameters. This matches the calling conventions of jax.grad, optimistix, and diffrax, so the energies can be plugged directly into optimizers and ODE solvers. Material parameters can be scalars or per-face/per-edge/per-vertex arrays (inhomogeneous materials).
A note on units: for a shell of thickness \(h\) made from a material with Young’s modulus \(Y\), the in-plane moduli scale as \(\sim Y h\) (units energy/area), while the bending moduli scale as \(\sim Y h^3/12\) (units energy) — thin shells are much easier to bend than to stretch.
Meshes with a boundary are supported (free/natural boundary conditions); meshes with vertices at infinity (see documentation) are not.
For references, please see the following: - “Discrete shells”, E. Grinspun et al., 2006 - “Growth patterns for shape-shifting elastic bilayers”, W. van Rees et al., 2017 - “Physical Simulation of Environmentally Induced Thin Shell Deformation”, H.Y. Chen et al, 2018
In-plane elastic energies
Thin elastic sheets can undergo very large deformations. Geometrically, the reference or rest shape of the sheet can be described by a metric tensor \(g_0\). The in-plane strain of a surface is measured by the (right) Cauchy-Green tensor \(C_f = g_0^{-1} \cdot g_f\), computed from \(g_0\) and the metric \(g_f\) describing the actual configuration of the mesh. A classic choice is the neo-Hookean energy: \[E_{\mathrm{NH}} = \sum_f A^0_f \left[ \frac{\mu}{2}\left(\frac{\mathrm{tr}(C_f)}{\sqrt{\det C_f}} - 2\right) +\frac{K}{2}\left(\sqrt{\det C_f} - 1\right)^2 \right], \qquad C_f = g_0^{-1} g_f\] \(K\) and \(\mu\) are the (2D) bulk and shear moduli (units: energy/area). For \(K=0\), only shear is penalized (the “Least Squares Conformal Mapping” (LSCM) energy). The metric can be easily evaluated for each triangle \(f\) (see below). The total energy sums over all faces, weighted by the reference areas \(A^0_f = \tfrac{1}{2}\sqrt{\det g_0}\): the energy is an integral over the undeformed configuration. Another example is the St.-Venant-Kirchhoff model: \[E_{\mathrm{SK}} = \sum_f A^0_f \left[\alpha \, \mathrm{tr}(C_f - \mathbb{I})^2 + 2\beta \, \mathrm{tr}\!\left((C_f-\mathbb{I})^2\right) \right]\] where \(\alpha,\beta\) are Lamé coefficients and \(C_f - \mathbb{I}\) is a measure of strain (2x the Green-Lagrange strain tensor).
Evaluating the metric for a triangle
The metric for a triangle \(a, b, c\) can be computed from two side vectors \(u=b-a, v=c-a\) as \[g_f = \begin{pmatrix}u^2 & u\cdot v \\ u\cdot v & v^2 \end{pmatrix} \]
Important: The reference metric, a \(2\times 2\) matrix per triangle, also needs to be expressed in the “face basis” \(u, v\).
Since the metrics are \(2\times 2\), we compute \(g_0^{-1}\) with the closed-form adjugate formula (get_inverse_metric) instead of a general linear solve, which would otherwise dominate the energy runtime (see the benchmark below).
get_two_x_two_inverse
def get_two_x_two_inverse(
metric:Float[Array, '*batch 2 2'], # 2x2 matrix/matrices; must be non-singular (e.g. metrics of
non-degenerate triangles).
)->Float[Array, '*batch 2 2']: # Inverse matrices.
Closed-form inverse of 2x2 metric tensor(s), inv(g) = adj(g)/det(g).
Much faster than jnp.linalg.inv/solve for 2x2 matrices.
get_area_from_metric
def get_area_from_metric(
metric:Float[Array, 'n_faces 2 2'], # Per-face metric. Symmetric positive definite for non-degenerate triangles.
)->Float[Array, 'n_faces']: # Per-face area.
Compute triangle areas from metric tensor.
get_metric
def get_metric(
vertices:Float[Array, 'n_vertices dim'], # Vertex positions (dim = 2 or 3).
hemesh:HeMesh, # Half-edge mesh.
)->Float[Array, 'n_faces 2 2']: # Per-face metric. Symmetric positive definite for non-degenerate triangles.
Metric tensor (first fundamental form) per triangle, in the face basis.
For a triangle with vertices a, b, c, the basis is the pair of side vectors u = b - a, v = c - a, so g = [[u.u, u.v], [u.v, v.v]].
get_neo_hookean_energy
def get_neo_hookean_energy(
vertices:Float[Array, 'n_vertices dim'], # Vertex positions (dim = 2 or 3).
args:tuple, # (hemesh, metric_ref, mod_bulk, mod_shear). metric_ref is the reference
metric in the face basis (see [`get_metric`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#get_metric)). The moduli can be scalars
or per-face arrays (inhomogeneous material).
)->Float[Array, '']: # Total elastic energy.
Total neo-Hookean in-plane energy, integrated over the reference configuration: E = sum_f A0_f * W(C_f) with A0_f the reference areas and W the neo-Hookean density (see get_neo_hookean_energy_density).
get_neo_hookean_energy_density
def get_neo_hookean_energy_density(
C:Float[Array, '*batch 2 2'], # Cauchy-Green tensor(s) C = g0^-1 g, positive definite.
mod_bulk:Union, mod_shear:Union
)->Float[Array, '*batch']: # Energy density (units: energy/area).
Neo-Hookean energy density from the Cauchy-Green deformation tensor C: W = mod_shear/2 * (tr(C)/J - 2) + mod_bulk/2 * (J - 1)^2 where J = sqrt(det(C)) is the local area change. See https://en.wikipedia.org/wiki/Neo-Hookean_solid
# at the rest configuration, energy and forces vanish
metric_ref = get_metric(sphere.vertices, sphere_hemesh)
args = (sphere_hemesh, metric_ref, 1.0, 1.0)
assert np.isclose(get_neo_hookean_energy(sphere.vertices, args), 0.0, atol=1e-10)
assert np.allclose(jax.grad(get_neo_hookean_energy)(sphere.vertices, args), 0.0, atol=1e-8)
# uniform scaling x -> s*x is conformal: only the area term contributes,
# E = mod_bulk/2 * (s^2-1)^2 * A0. With mod_bulk = 0 (LSCM), it costs nothing.
s, mod_bulk = 1.3, 2.0
areas = get_area_from_metric(metric_ref)
A0 = areas.sum()
E = get_neo_hookean_energy(s*sphere.vertices, (sphere_hemesh, metric_ref, mod_bulk, 1.0))
assert np.isclose(E, mod_bulk/2 * (s**2-1)**2 * A0)
assert np.isclose(get_neo_hookean_energy(s*sphere.vertices,
(sphere_hemesh, metric_ref, 0.0, 1.0)), 0.0)
# moduli can also be specified per-face (inhomogeneous material)
mods = jnp.linspace(0.5, 2.0, sphere_hemesh.n_faces)
E_inhom = get_neo_hookean_energy(s*sphere.vertices, (sphere_hemesh, metric_ref, mods, mods))
assert np.isclose(E_inhom, (areas * mods/2 * (s**2-1)**2).sum())get_st_venant_kirchhoff_energy
def get_st_venant_kirchhoff_energy(
vertices:Float[Array, 'n_vertices dim'], # Vertex positions (dim = 2 or 3).
args:tuple, # (hemesh, metric_ref, alpha, beta). metric_ref is the reference metric
in the face basis (see [`get_metric`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#get_metric)). The Lame-type coefficients
alpha, beta >= 0 (units: energy/area) can be scalars or per-face arrays.
)->Float[Array, '']: # Total elastic energy.
Total St.-Venant-Kirchhoff in-plane energy: E = sum_f A0_f * (alpha * tr(C - I)^2 + 2 * beta * tr((C - I)^2)) with C = g0^-1 g the Cauchy-Green tensor, integrated over the reference configuration. Quadratic in the (Green-Lagrange) strain.
# uniform scaling x -> s*x gives C = s^2 I, so tr(C-I) = 2(s^2-1) and
# tr((C-I)^2) = 2(s^2-1)^2
s, alpha, beta = 1.2, 1.0, 0.5
metric_ref = get_metric(sphere.vertices, sphere_hemesh)
areas = get_area_from_metric(metric_ref)
A0 = areas.sum()
args = (sphere_hemesh, metric_ref, alpha, beta)
assert np.isclose(get_st_venant_kirchhoff_energy(sphere.vertices, args), 0.0, atol=1e-10)
assert np.isclose(get_st_venant_kirchhoff_energy(s*sphere.vertices, args),
(4*alpha + 4*beta) * (s**2-1)**2 * A0)Bending energies
Bending energies penalize out-of-plane deformation. A simple model uses the dihedral angles \(\theta_{e}\) for each interior edge (the angles between the normal vectors of adjacent triangles): \[E_{\mathrm{B}} = \kappa \sum_e \frac{\ell_e^2}{A_e} (\theta_e -\theta_e^0)^2 \] Where the area \(A_e\) is defined as \(1/3\) of the areas of the adjacent triangles, \(\theta_e^0\) is the target bending angle, and \(\kappa\) is the bending modulus (units: energy). This is the “discrete shells” bending energy; note that it is scale invariant.
get_dihedral_bending_energy
def get_dihedral_bending_energy(
vertices:Float[Array, 'n_vertices 3'], # Vertex positions.
args:tuple, # (hemesh, theta0, mod_bending). The target angles theta0 (radians,
sign convention of `geom.get_dihedral_angles`: positive = convex) and the
bending modulus mod_bending >= 0 (units: energy) can be scalars or
per-half-edge arrays; per-half-edge values must be twin-symmetric.
)->Float[Array, '']: # Total bending energy.
Discrete-shells bending energy from dihedral angles: E = sum_e mod_bending * l_e^2 / A_e * (theta_e - theta0_e)^2 where the sum runs over interior edges, l_e is the edge length and A_e = (A_f + A_f’)/3 the edge area (one third of the adjacent triangle areas). Scale invariant. Boundary edges do not contribute.
# a flat mesh has zero bending energy, and finite forces (boundary edges are masked)
assert np.isclose(get_dihedral_bending_energy(disk.vertices, (disk_hemesh, 0.0, 1.0)), 0.0)
assert np.all(np.isfinite(jax.grad(get_dihedral_bending_energy)(
disk.vertices, (disk_hemesh, 0.0, 1.0))))
# ... and so does any mesh at its rest angles
theta0 = geom.get_dihedral_angles(sphere.vertices, sphere_hemesh)
assert np.isclose(get_dihedral_bending_energy(sphere.vertices, (sphere_hemesh, theta0, 1.0)), 0.0)
# the energy is scale invariant (theta and l^2/A are)
E1 = get_dihedral_bending_energy(sphere.vertices, (sphere_hemesh, 0.0, 1.0))
E2 = get_dihedral_bending_energy(2*sphere.vertices, (sphere_hemesh, 0.0, 1.0))
assert E1 > 0 and np.isclose(E1, E2)Bending energies using the 2nd fundamental form
A second option (van Rees, 2017) is to use the 2nd fundamental form \(b_f\) (the extrinsic curvature tensor) and a variant of the St-Venant-Kirchhoff energy: \[E_{\mathrm{B,SK}} = \sum_f A^0_f \left[\alpha \, \mathrm{tr}\!\left(g_0^{-1}(b_f-b_f^0)\right)^2 + 2\beta \, \mathrm{tr}\!\left((g_0^{-1}(b_f-b_f^0))^2\right) \right]\] Here \(b_f^0\) is the reference curvature. If the sheet prefers to be flat, \(b_f^0 = 0\). Like for the in-plane energies, \(A^0_f\) are the reference areas, and \(\alpha, \beta\) have units of energy.
Evaluating the 2nd fundamental form for a triangle
For a triangle with vertices \(a, b, c\), we defined the edge vectors \(u=b-a, v=c-a\). For an edge \(e\) between two faces \(f, f'\), the per-edge normal is \[n_e = \frac{n_f+n_{f'}}{|n_f+n_{f'}|} \] (for a boundary edge, \(n_e = n_f\)). We label the triangle edges by the opposite vertex (so that \(n_a\) is the normal of the edge opposite to vertex \(a\)). One obtains a finite-difference approximation for \(b\) (Chen, 2018): \[b = \begin{pmatrix}2u\cdot(n_a-n_b) & u\cdot(n_a-n_c) + v\cdot(n_a-n_b) \\ u\cdot(n_a-n_c) + v\cdot(n_a-n_b) & 2v\cdot(n_a-n_c) \end{pmatrix} \] The sign convention is chosen so that a sphere with outward-pointing normals has \(b = g\), i.e. positive mean curvature \(H = \tfrac{1}{2}\mathrm{tr}(g^{-1}b)\), consistent with geom.get_mean_curvature_dihedral (this is the opposite of the textbook convention \(b = n \cdot \partial_i \partial_j x\)).
get_second_fundamental_form
def get_second_fundamental_form(
vertices:Float[Array, 'n_vertices 3'], # Vertex positions.
hemesh:HeMesh, # Half-edge mesh.
)->Float[Array, 'n_faces 2 2']: # Per-face second fundamental form (units: length).
Second fundamental form (extrinsic curvature tensor) per triangle, in the face basis.
Finite-difference approximation from mid-edge normals (Grinspun et al. 2006), expressed in the same basis u, v as get_metric. The shape operator is S = g^-1 b, with mean curvature H = tr(S)/2 and Gaussian curvature K = det(S).
Sign convention: a sphere with outward-pointing normals has b = g (H = +1 for the unit sphere), consistent with geom.get_mean_curvature_dihedral.
# on the unit sphere, b = g: the shape operator g^-1 b is the identity,
# H = tr(g^-1 b)/2 = +1 (outward normals, positive-convex sign convention)
metric = get_metric(sphere.vertices, sphere_hemesh)
b = get_second_fundamental_form(sphere.vertices, sphere_hemesh)
shape_operator = get_two_x_two_inverse(metric) @ b
H_face = jnp.trace(shape_operator, axis1=-2, axis2=-1)/2
assert np.allclose(H_face, 1.0, atol=0.1)
# ... consistent with the vertex-based mean curvature, including the sign
H_vertex = geom.get_mean_curvature_dihedral(sphere.vertices, sphere_hemesh)
assert np.isclose(H_face.mean(), H_vertex.mean(), rtol=0.02)
# a flat disk has b = 0, including at the boundary faces
assert np.allclose(get_second_fundamental_form(disk.vertices, disk_hemesh), 0.0, atol=1e-12)
# the discretization leads to a symmetric form, even though that is not directly manifest
assert np.allclose(b, jnp.transpose(b, axes=(0, 2, 1)), atol=1e-12)get_svk_bending_energy
def get_svk_bending_energy(
vertices:Float[Array, 'n_vertices 3'], # Vertex positions.
args:tuple, # (hemesh, metric_ref, b_ref, alpha, beta). metric_ref and the reference
curvature b_ref are expressed in the face basis (see [`get_metric`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#get_metric) and
[`get_second_fundamental_form`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#get_second_fundamental_form)). The bending moduli alpha, beta >= 0
(units: energy) can be scalars or per-face arrays.
)->Float[Array, '']: # Total bending energy.
St.-Venant-Kirchhoff-type bending energy from the second fundamental form: E = sum_f A0_f * (alpha * tr(S)^2 + 2 * beta * tr(S^2)), S = g0^-1 (b - b0) integrated over the reference configuration. For b0 = 0, the sheet prefers to be flat.
# unit sphere with flat reference curvature (b0 = 0): the shape operator is the
# identity, so the density is 4*alpha + 4*beta and E = 4*pi * (4*alpha + 4*beta)
alpha, beta = 1.0, 0.5
metric_ref = get_metric(sphere.vertices, sphere_hemesh)
b0 = jnp.zeros((sphere_hemesh.n_faces, 2, 2))
E = get_svk_bending_energy(sphere.vertices, (sphere_hemesh, metric_ref, b0, alpha, beta))
assert np.isclose(E, 4*np.pi * (4*alpha + 4*beta), rtol=0.01)
# at the rest configuration (b0 = current curvature), energy and forces vanish
b_rest = get_second_fundamental_form(sphere.vertices, sphere_hemesh)
args_rest = (sphere_hemesh, metric_ref, b_rest, alpha, beta)
assert np.isclose(get_svk_bending_energy(sphere.vertices, args_rest), 0.0, atol=1e-10)
assert np.allclose(jax.grad(get_svk_bending_energy)(sphere.vertices, args_rest), 0.0, atol=1e-8)
# a flat disk with flat reference has zero bending energy, and finite forces
args_disk = (disk_hemesh, get_metric(disk.vertices, disk_hemesh),
jnp.zeros((disk_hemesh.n_faces, 2, 2)), alpha, beta)
assert np.isclose(get_svk_bending_energy(disk.vertices, args_disk), 0.0, atol=1e-12)
assert np.all(np.isfinite(jax.grad(get_svk_bending_energy)(disk.vertices, args_disk)))Membrane energies
Membranes (e.g. lipid bilayers) differ from elastic shells because they are fluid in-plane: material can redistribute within a membrane to relax in-plane stresses through viscosity. However, membranes can still resist out-of-plane bending deformation.
The Helfrich energy is a geometric model of bending energy. It uses the mean and Gaussian curvatures \(H, K\) of the surface \(\mathcal{M}\). The energy reads:
\[E_H =\int dA \left( \frac{\kappa_H}{2}(H-H_0)^2 + \kappa_G K \right) \]
A nonzero spontaneous curvature \(H_0\) means that the membrane “prefers” to be curved. If the surface is closed, the \(\int K\)-term is a topological invariant by the Gauss-Bonnet theorem, \(\int K\, dA = 2\pi\chi\). The discretization preserves this exactly.
get_helfrich_energy
def get_helfrich_energy(
vertices:Float[Array, 'n_vertices 3'], # Vertex positions.
args:tuple, # (hemesh, H0, kappa_H, kappa_K):
the mesh connectivity, the spontaneous curvature H0 (units: 1/length),
the bending modulus kappa_H >= 0 and the Gaussian modulus kappa_K (units: energy).
H0, kappa_H, kappa_K can be scalars or per-vertex arrays.
)->Float[Array, '']: # Total bending energy.
Discrete Helfrich bending energy of a triangulated surface: E = kappa_H/2 * sum_i (H_i - H0)^2 * A_i + kappa_K * sum_i K_i * A_i^bar
The mean-curvature term uses the dihedral-angle mean curvature and mixed Voronoi cell areas A_i. On a closed mesh the Gaussian curvature term is exactly 2pichi*kappa_K (Gauss-Bonnet) and exerts no forces.
# On the unit sphere, up to discretization error, the H-term gives ~2*pi*kappa_H (the exact smooth value),
# and the K-term is exactly 4*pi*kappa_K = 2*pi*chi*kappa_K (discrete Gauss-Bonnet)
E_H = get_helfrich_energy(sphere.vertices, (sphere_hemesh, 0.0, 1.0, 0.0))
assert np.isclose(E_H, 2*np.pi, rtol=0.02)
E_K = get_helfrich_energy(sphere.vertices, (sphere_hemesh, 0.0, 0.0, 1.0))
assert np.isclose(E_K, 4*np.pi)
# for H0 = 0, the energy is scale invariant
args = (sphere_hemesh, 0.0, 1.0, 1.0)
assert np.isclose(get_helfrich_energy(sphere.vertices, args),
get_helfrich_energy(2*sphere.vertices, args))
# a unit sphere with spontaneous curvature H0 = 1 costs (almost) nothing
E_H0 = get_helfrich_energy(sphere.vertices, (sphere_hemesh, 1.0, 1.0, 0.0))
assert E_H0 < 1e-2 * E_HArea and volume constraints
Quadratic penalties on the total surface area and the enclosed volume. These are the usual companions to a bending energy: the Helfrich functional is scale-invariant at \(H_0=0\).
get_volume_constraint_energy
def get_volume_constraint_energy(
vertices:Float[Array, 'n_vertices 3'], # Vertex positions.
args:tuple, # ``(hemesh, volume_0, modulus)``: the target enclosed volume V0 (units: length^3)
and the penalty modulus (units: energy/length^6).
)->Float[Array, '']: # Constraint energy. Zero exactly when the enclosed volume equals V0.
Quadratic penalty on the enclosed volume: E = modulus/2 * (V - V0)^2
Requires a closed, consistently oriented mesh (geometry.get_volume checks this). The volume is signed, so V0 should have the same sign as the initial configuration’s volume – positive for outward-oriented faces.
get_area_constraint_energy
def get_area_constraint_energy(
vertices:Float[Array, 'n_vertices dim'], # Vertex positions.
args:tuple, # ``(hemesh, area_0, modulus)``: the target total area A0 (units: length^2) and the
penalty modulus (units: energy/length^4).
)->Float[Array, '']: # Constraint energy. Zero exactly when the total area equals A0.
Quadratic penalty on the total surface area: E = modulus/2 * (A - A0)^2
Note this is a global constraint on the total area, not a per-face one. Combine it with a bending energy yourself; the energies in this module are deliberately separate so they can be mixed and weighted freely.
from triangulax import trigonometry as _trig# area/volume constraints: zero at the reference, correct scaling, rigid-invariant
A0 = geom.get_area(sphere.vertices, sphere_hemesh)
V0 = geom.get_volume(sphere.vertices, sphere_hemesh)
assert jnp.allclose(get_area_constraint_energy(sphere.vertices, (sphere_hemesh, A0, 1.0)), 0.0, atol=1e-20)
assert jnp.allclose(get_volume_constraint_energy(sphere.vertices, (sphere_hemesh, V0, 1.0)), 0.0, atol=1e-20)
# a uniform dilation by s scales area by s^2 and volume by s^3
s = 1.1
assert jnp.allclose(get_area_constraint_energy(s*sphere.vertices, (sphere_hemesh, A0, 1.0)),
0.5*(s**2 - 1)**2 * A0**2, rtol=1e-10)
assert jnp.allclose(get_volume_constraint_energy(s*sphere.vertices, (sphere_hemesh, V0, 1.0)),
0.5*(s**3 - 1)**2 * V0**2, rtol=1e-10)
# invariant under rigid motion, and the gradient vanishes at the reference
R = _trig.quaternion_to_rot_mat(jnp.array([0.3, 0.5, -0.2, 0.8]))
moved = sphere.vertices @ R.T + jnp.array([1.5, -2.0, 0.7])
for fn, ref in [(get_area_constraint_energy, A0), (get_volume_constraint_energy, V0)]:
assert jnp.allclose(fn(moved, (sphere_hemesh, ref, 1.0)), 0.0, atol=1e-16), fn.__name__
assert jnp.abs(jax.grad(fn)(sphere.vertices, (sphere_hemesh, ref, 1.0))).max() < 1e-12, fn.__name__
# finite and non-zero away from the reference
g_off = jax.grad(fn)(1.05*sphere.vertices, (sphere_hemesh, ref, 1.0))
assert jnp.isfinite(g_off).all() and jnp.abs(g_off).max() > 0, fn.__name__
print(f"area/volume constraints OK (unit sphere A0={A0:.4f} vs 4*pi={4*jnp.pi:.4f}, V0={V0:.4f} vs 4/3*pi={4/3*jnp.pi:.4f})")area/volume constraints OK (unit sphere A0=12.5062 vs 4*pi=12.5664, V0=4.1527 vs 4/3*pi=4.1888)Normal/tangential energies
The helper functions below “project” any energy into a function of only tangential or only normal displacements, relative to the current mesh. This lets any optimizer minimize over normal-only or tangent-only displacements.
The use case is membrane mechanics (tutorial 5): here, tangential displacement driven by a fictitious elastic energy regularizes vertex motion to preserver mesh quality.
vertices_from_tangential
def vertices_from_tangential(
v0:Float[Array, 'n 3'], t:Float[Array, 'n 2'], basis:Float[Array, '2 n 3']
)->Float[Array, 'n 3']:
Reconstruct vertices from base positions + tangential displacement.
basis is expected in the (2, n_vertices, 3) layout returned by geom.get_vertex_tangent_basis(vertices, hemesh).
vertices_from_normal
def vertices_from_normal(
v0:Float[Array, 'n 3'], h:Float[Array, 'n'], normals:Float[Array, 'n 3']
)->Float[Array, 'n 3']:
Reconstruct vertices from base positions + normal displacement.
make_tangential_energy
def make_tangential_energy(
energy_fn:Callable, # Energy with signature ``energy_fn(vertices, args) -> scalar``.
v0:Float[Array, 'n 3'], # Base vertex positions.
basis:Float[Array, '2 n 3'], # Tangent basis, in the ``(2, n_vertices, 3)`` layout returned by
`geom.get_vertex_tangent_basis`.
)->Callable: # Energy with signature ``energy(t, args) -> scalar``.
Wrap energy_fn(vertices, args) to optimize over tangent coords t ∈ ℝⁿˣ².
vertices(t) = v0 + einsum(‘vi,ivd->vd’, t, basis), see vertices_from_tangential.
Any basis (orthonormal or not) which spans the tangent plane is valid, but the same basis must be used in make_tangential_energy and in vertices_from_tangential to recover the 3D vertex positions.
make_normal_energy
def make_normal_energy(
energy_fn:Callable, # Energy with signature ``energy_fn(vertices, args) -> scalar``.
v0:Float[Array, 'n 3'], # Base vertex positions.
normals:Float[Array, 'n 3'], # Per-vertex displacement directions, e.g. `geom.get_vertex_normals`.
)->Callable: # Energy with signature ``energy(h, args) -> scalar``.
Wrap energy_fn(vertices, args) to optimize over normal heights h ∈ ℝⁿ.
vertices(h) = v0 + h[:, None] * normals, see vertices_from_normal.
JAX-compatibility
# all energies are compatible with jax.jit and jax.grad
metric_ref = get_metric(sphere.vertices, sphere_hemesh)
b_ref = get_second_fundamental_form(sphere.vertices, sphere_hemesh)
energies_and_args = [
(get_neo_hookean_energy, (sphere_hemesh, metric_ref, 1.0, 1.0)),
(get_st_venant_kirchhoff_energy, (sphere_hemesh, metric_ref, 1.0, 1.0)),
(get_dihedral_bending_energy, (sphere_hemesh, 0.0, 1.0)),
(get_svk_bending_energy, (sphere_hemesh, metric_ref, b_ref, 1.0, 1.0)),
(get_helfrich_energy, (sphere_hemesh, 0.0, 1.0, 1.0)),]
for energy, args in energies_and_args:
value, grad = jax.jit(jax.value_and_grad(energy))(1.1*sphere.vertices, args)
assert np.isfinite(value) and np.all(np.isfinite(grad))