## `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 ∼ *Y**h* (units
energy/area), while the bending moduli scale as ∼ *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](https://nikolas-claussen.github.io/triangulax/src/02_halfedge_datastructure.html#boundary-and-the-vertex-at-infinity))
are not.
For references, please see the following: - [“Discrete shells”, E.
Grinspun et al., 2006](https://dl.acm.org/doi/10.5555/846276.846284) -
[“Growth patterns for shape-shifting elastic bilayers”, W. van Rees et
al., 2017](http://www.pnas.org/cgi/doi/10.1073/pnas.1709025114) -
[“Physical Simulation of Environmentally Induced Thin Shell
Deformation”, H.Y. Chen et al,
2018](https://doi.org/10.1145/3197517.3201395)
### 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 ⋅ *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 *μ* 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*SK = ∑*f**A**f*0\[*α* tr(*C**f* − 𝕀)2 + 2*β* tr ((*C**f* − 𝕀)2)\]
where *α*, *β* are Lamé coefficients and *C**f* − 𝕀 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 × 2 matrix per triangle, also
needs to be expressed in the “face basis” *u*, *v*.
Since the metrics are 2 × 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).
------------------------------------------------------------------------
source
### get_two_x_two_inverse
``` python
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.
------------------------------------------------------------------------
source
### get_area_from_metric
``` python
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.*
------------------------------------------------------------------------
source
### get_metric
``` python
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\]\].
``` python
# 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))
```
------------------------------------------------------------------------
source
### get_neo_hookean_energy
``` python
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`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#get_neo_hookean_energy_density)).
------------------------------------------------------------------------
source
### get_neo_hookean_energy_density
``` python
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
``` python
# 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())
```
------------------------------------------------------------------------
source
### get_st_venant_kirchhoff_energy
``` python
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.
``` python
# 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* *θ**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, *θ**e*0 is the target bending
angle, and *κ* is the bending modulus (units: energy). This is the
“discrete shells” bending energy; note that it is scale invariant.
------------------------------------------------------------------------
source
### get_dihedral_bending_energy
``` python
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.
``` python
# 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*B, SK = ∑*f**A**f*0\[*α* tr (*g*0−1(*b**f* − *b**f*0))2 + 2*β* tr ((*g*0−1(*b**f* − *b**f*0))2)\]
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**f*0 are the reference
areas, and *α*, *β* 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* ⋅ ∂*i*∂*j**x*).
------------------------------------------------------------------------
source
### get_second_fundamental_form
``` python
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`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#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`.
``` python
# 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)
```
------------------------------------------------------------------------
source
### get_svk_bending_energy
``` python
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.
``` python
# 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 ℳ. 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 ∫*K*-term
is a topological invariant by the Gauss-Bonnet theorem,
∫*K* *d**A* = 2*π**χ*. The discretization preserves this exactly.
------------------------------------------------------------------------
source
### get_helfrich_energy
``` python
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 2*pi*chi\*kappa_K (Gauss-Bonnet) and exerts no forces.
``` python
# 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_H
```
### Area 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.
------------------------------------------------------------------------
source
### get_volume_constraint_energy
``` python
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.
------------------------------------------------------------------------
source
### get_area_constraint_energy
``` python
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.
``` python
from triangulax import trigonometry as _trig
```
``` python
# 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.
------------------------------------------------------------------------
source
### vertices_from_tangential
``` python
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)`.
------------------------------------------------------------------------
source
### vertices_from_normal
``` python
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.*
------------------------------------------------------------------------
source
### make_tangential_energy
``` python
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`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#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.
------------------------------------------------------------------------
source
### make_normal_energy
``` python
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`](https://nikolas-claussen.github.io/triangulax/src/elasticity.html#vertices_from_normal).
### JAX-compatibility
``` python
# 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))
```