trigonometry: Trigonometry

This module provides basic trigonometric, vector, and linear algebra functions for triangular meshes. These are scalar or single-triangle operations designed to be vectorized across meshes using jax.vmap.

Throughout, we provide type signatures for all functions using jaxtyping. We use jnp (= jax.numpy) instead of numpy, and follow JAX’s functional programming paradigm (see JAX — the sharp bits).

Safe division

Degenerate configurations (zero-length edges, collapsed triangles) can make divisors vanish. safe_divide and safe_normalize guard against this so that BOTH the value and its gradient stay finite, and - unlike clipping the divisor to a fixed value - they are scale independent, so a small mesh is not mistaken for a degenerate one.

source

safe_normalize


def safe_normalize(
    x:Float[Array, '... dim'], axis:int=-1
)->Float[Array, '... dim']:

Return x / |x|, or the zero vector where |x| == 0.

NaN-gradient-safe and scale independent (see safe_divide). The sum of squares is guarded directly rather than going through jnp.linalg.norm, whose gradient is itself NaN at the origin.

source

safe_norm


def safe_norm(
    x:Float[Array, '... dim'], axis:int=-1
)->Float[Array, '...']:

Return |x|, with a finite (zero) gradient where |x| == 0.

jnp.linalg.norm has a NaN gradient at the origin, which propagates into every quantity computed from edge lengths as soon as a single edge collapses. The sum of squares is guarded before the square root instead.

source

safe_divide


def safe_divide(
    numerator:Array, denominator:Array, fill:float=0.0
)->Array:

Divide, returning fill where the denominator is exactly zero.

Both the value and its gradient stay finite at denominator == 0: the division uses a denominator of 1 wherever the true denominator vanishes, so reverse-mode never differentiates x / 0 (a bare jnp.where(d != 0, x / d, fill) would still produce a NaN gradient from the untaken branch).

This is scale independent: a small-but-valid denominator is divided by faithfully rather than being treated as zero, and a genuine zero yields fill instead of a huge finite number. numerator broadcasts against denominator as usual.

Triangle areas and circumcenters

source

get_polygon_area


def get_polygon_area(
    pts:Float[Array, 'n_vertices 2'], # Ordered polygon vertices.
)->Float[Array, '']: # Signed area.

Signed area of a 2D simple polygon (shoelace formula).

Positive for counter-clockwise vertex ordering. This matches the sign convention of get_oriented_triangle_area.

source

get_triangle_area


def get_triangle_area(
    a:Float[Array, 'dim'], b:Float[Array, 'dim'], c:Float[Array, 'dim']
)->Float[Array, '']: # Triangle area.

Unsigned area of triangle with vertices a, b, c.

source

get_triangle_area_from_sides


def get_triangle_area_from_sides(
    u:Float[Array, 'dim'], v:Float[Array, 'dim']
)->Float[Array, '']: # Triangle area.

Unsigned area of triangle edge vectors u, v

source

get_oriented_triangle_area


def get_oriented_triangle_area(
    a:Float[Array, 'dim'], b:Float[Array, 'dim'], c:Float[Array, 'dim']
)->Float[Array, '*']: # Scalar (dim=2) or area-weighted normal vector (dim=3).
In 2d, positive for counter-clockwise (a, b, c).

Signed area of triangle with vertices a, b, c.

source

get_circumcenter


def get_circumcenter(
    a:Float[Array, 'dim'], b:Float[Array, 'dim'], c:Float[Array, 'dim']
)->Float[Array, 'dim']: # Circumcenter coordinates.

Circumcenter of triangle with vertices a, b, c via barycentric coordinates.

Vector operations

source

get_cot_between_vectors


def get_cot_between_vectors(
    a:Float[Array, 'dim'], b:Float[Array, 'dim']
)->Float[Array, '']: # Cotangent value.

Cotangent of the angle between vectors a and b.

source

get_angle_between_vectors


def get_angle_between_vectors(
    a:Float[Array, 'dim'], b:Float[Array, 'dim']
)->Float[Array, '']: # Angle in radians, in $[0, \pi]$.

Unsigned angle between vectors a and b.

source

get_signed_angle_between_vectors


def get_signed_angle_between_vectors(
    a:Float[Array, '2'], b:Float[Array, '2']
)->Float[Array, '']: # Signed angle in radians, in $(-\pi, \pi]$.

Signed angle from 2D vector a to b (CCW positive).

source

get_projector


def get_projector(
    normal:Float[Array, 'dim'], # Normal vector (need not be unit length).
)->Float[Array, 'dim dim']: # Projection matrix onto the plane orthogonal to normal.

Tangent-plane projector \(P = I - \hat{n} \otimes \hat{n}\).

source

project_out_vector


def project_out_vector(
    a:Float[Array, 'dim'], # Input vector.
    b:Float[Array, 'dim'], # Normal direction to project out.
)->Float[Array, 'dim']: # Component of a orthogonal to b.

Project vector a onto the plane orthogonal to b.

source

project_on_vector


def project_on_vector(
    a:Float[Array, 'dim'], # Vector to project.
    b:Float[Array, 'dim'], # Target direction.
)->Float[Array, 'dim']: # Projection of a onto b.

Project vector a onto vector b.

source

get_tetrahedron_volume


def get_tetrahedron_volume(
    a:Float[Array, '3'], b:Float[Array, '3'], c:Float[Array, '3']
)->Float[Array, '']: # Signed volume (positive if a, b, c form a right-handed frame).

Signed volume of tetrahedron with edge vectors a, b, c from a common vertex.

# get_polygon_area is positive for counter-clockwise ordering, and agrees in
# sign and magnitude with get_oriented_triangle_area on triangles.
ccw_tri = jnp.array([[0., 0.], [1., 0.], [0., 1.]])
cw_tri = ccw_tri[::-1]

assert jnp.allclose(get_polygon_area(ccw_tri), 0.5)
assert jnp.allclose(get_polygon_area(cw_tri), -0.5)
assert jnp.allclose(get_polygon_area(ccw_tri), get_oriented_triangle_area(*ccw_tri))
assert jnp.allclose(get_polygon_area(cw_tri), get_oriented_triangle_area(*cw_tri))
assert jnp.allclose(jnp.abs(get_polygon_area(ccw_tri)), get_triangle_area(*ccw_tri))

# counter-clockwise unit square
assert jnp.allclose(get_polygon_area(jnp.array([[0., 0.], [1., 0.], [1., 1.], [0., 1.]])), 1.0)

# translation invariance
assert jnp.allclose(get_polygon_area(ccw_tri + jnp.array([3., -7.])), 0.5)

Intrinsic geometry

Many mesh quantities (angles, triangle areas, cotangent weights, circumcenters, …) can be computed purely intrinsically from edge lengths \(\ell_{ij}\), without explicit reference to vertex coordinates.

This matters beyond elegance: whenever vertex positions are unavailable or ambiguous - most importantly under periodic or Lees-Edwards boundary conditions, where a difference of coordinates is meaningless across the boundary - edge lengths are still well defined. The functions below are the single-triangle building blocks; their mesh-level counterparts are in the geometry module (“Intrinsic geometry from edge lengths”) and the linops module.

source

get_circumcenter_from_lengths


def get_circumcenter_from_lengths(
    la:Float[Array, ''], # Length of side opposite vertex a (i.e. edge bc).
    lb:Float[Array, ''], # Length of side opposite vertex b (i.e. edge ca).
    lc:Float[Array, ''], # Length of side opposite vertex c (i.e. edge ab).
)->Float[Array, '3']: # Normalized barycentric coordinates [lambda_a, lambda_b, lambda_c].

Circumcenter in barycentric coordinates from edge lengths.

To recover Cartesian coordinates: \(u = \lambda_a \, a + \lambda_b \, b + \lambda_c \, c\).

source

get_cotangents_from_lengths


def get_cotangents_from_lengths(
    la:Float[Array, ''], # Length of side opposite vertex a (i.e. edge bc).
    lb:Float[Array, ''], # Length of side opposite vertex b (i.e. edge ca).
    lc:Float[Array, ''], # Length of side opposite vertex c (i.e. edge ab).
)->Float[Array, '3']: # Cotangents [cot_alpha, cot_beta, cot_gamma] at vertices a, b, c.

Cotangents of interior angles from side lengths.

Uses \(\cot(\alpha) = (b^2 + c^2 - a^2) / (4 \cdot \text{area})\).

source

get_angles_from_lengths


def get_angles_from_lengths(
    la:Float[Array, ''], # Length of side opposite vertex a (i.e. edge bc).
    lb:Float[Array, ''], # Length of side opposite vertex b (i.e. edge ca).
    lc:Float[Array, ''], # Length of side opposite vertex c (i.e. edge ab).
)->Float[Array, '3']: # Angles [alpha, beta, gamma] at vertices a, b, c respectively.

Interior angles from side lengths.

Uses \(\theta_a = \mathrm{atan2}(4A, b^2 + c^2 - a^2)\) with \(A\) the triangle area, the intrinsic counterpart of get_angle_between_vectors. This is accurate for angles near 0 and \(\pi\), where \(\arccos\) of the law of cosines loses about half the available digits.

source

get_triangle_area_from_lengths


def get_triangle_area_from_lengths(
    la:Float[Array, ''], # Length of side opposite vertex a (i.e. edge bc).
    lb:Float[Array, ''], # Length of side opposite vertex b (i.e. edge ca).
    lc:Float[Array, ''], # Length of side opposite vertex c (i.e. edge ab).
)->Float[Array, '']: # Area of the triangle.

Triangle area from side lengths using Heron’s formula (Kahan’s stable form).

The sides are sorted as \(a \geq b \geq c\) before evaluating \(A = \frac{1}{4}\sqrt{(a+(b+c))(c-(a-b))(c+(a-b))(a+(b-c))}\). The textbook form \(\sqrt{s(s-a)(s-b)(s-c)}\) suffers catastrophic cancellation for thin triangles: for a triangle of unit side and height \(10^{-8}\) it returns exactly 0 (float64), which silently deletes the triangle from cotangent weights and areas.

Returns 0 (with zero gradient, not NaN) when the side lengths violate the triangle inequality or describe a degenerate triangle.

# Test intrinsic functions against extrinsic (vertex-based) implementations
triangles = [
    jnp.array([[0., 0.], [1., 0.], [0., 1.]]),
    jnp.array([[0., 0.], [3., 0.], [1.5, 2.]]),
    jnp.array([[0., 0., 0.], [1., 0., 0.], [0., 1., 1.]]),
    jnp.array([[1., 2.], [4., 6.], [7., 1.]]),]

for tri in triangles:
    a, b, c = tri
    la, lb, lc = jnp.linalg.norm(b - c), jnp.linalg.norm(c - a), jnp.linalg.norm(a - b)

    # Area
    assert jnp.allclose(get_triangle_area(a, b, c),
                         get_triangle_area_from_lengths(la, lb, lc), atol=1e-10)
    # Angles
    angles_ext = jnp.stack([get_angle_between_vectors(b - a, c - a),
                            get_angle_between_vectors(a - b, c - b),
                            get_angle_between_vectors(a - c, b - c)])
    assert jnp.allclose(angles_ext, get_angles_from_lengths(la, lb, lc), atol=1e-10)

    # Cotangents
    cots_ext = jnp.stack([get_cot_between_vectors(b - a, c - a),
                          get_cot_between_vectors(a - b, c - b),
                          get_cot_between_vectors(a - c, b - c)])
    assert jnp.allclose(cots_ext, get_cotangents_from_lengths(la, lb, lc), atol=1e-10)

    # Circumcenter
    cc_ext = get_circumcenter(a, b, c)
    bary = get_circumcenter_from_lengths(la, lb, lc)
    cc_int = bary[0] * a + bary[1] * b + bary[2] * c
    assert jnp.allclose(cc_ext, cc_int, atol=1e-10)

# Sliver triangles: base 1, height h, so the exact area is h/2. The textbook Heron
# formula returns exactly 0 below h ~ 1e-8 (float64); the stable form does not.
for h in [1e-4, 1e-6, 1e-8]:
    la, lb, lc = jnp.sqrt(0.25 + h**2), jnp.sqrt(0.25 + h**2), jnp.array(1.0)
    assert jnp.isclose(get_triangle_area_from_lengths(la, lb, lc), h / 2, rtol=1e-4), h
    # the two small angles are resolved too (arccos of the law of cosines would not)
    alpha, beta, gamma = get_angles_from_lengths(la, lb, lc)
    assert jnp.isclose(alpha, 2 * h, rtol=1e-4) and jnp.isclose(gamma, jnp.pi - 4 * h, rtol=1e-6), h

# degenerate side lengths: zero area with a finite (zero) gradient, not NaN
assert get_triangle_area_from_lengths(jnp.array(1.), jnp.array(.5), jnp.array(.5)) == 0.0
assert jnp.isfinite(jnp.array(jax.grad(get_triangle_area_from_lengths, argnums=(0, 1, 2))(
    jnp.array(1.), jnp.array(.5), jnp.array(.5)))).all()
# violating the triangle inequality gives 0, not NaN
assert get_triangle_area_from_lengths(jnp.array(5.), jnp.array(1.), jnp.array(1.)) == 0.0

Rotation matrices and normals

source

quaternion_to_rot_mat


def quaternion_to_rot_mat(
    q:Float[Array, '4'], # Quaternion $[w, x, y, z]$ (normalized internally).
)->Float[Array, '3 3']: # Rotation matrix.

Convert unit quaternion to 3D rotation matrix.

See Quaternion rotation <https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation>_.

source

get_triangle_normal


def get_triangle_normal(
    a:Float[Array, '3'], b:Float[Array, '3'], c:Float[Array, '3']
)->Float[Array, '3']: # Unit normal vector.

Unit normal of triangle abc (right-hand rule on a -> b -> c).

source

get_perp_2d


def get_perp_2d(
    x:Float[Array, '... 2'], # Input vector(s).
)->Float[Array, '... 2']: # Perpendicular vector(s).

90-degree clockwise rotation of a 2D vector: \((x, y) \to (y, -x)\).

source

get_rot_mat


def get_rot_mat(
    theta:float, # Rotation angle in radians.
)->Float[Array, '2 2']: # Rotation matrix.

2D counter-clockwise rotation matrix for angle theta.

Barycentric coordinates

source

get_barycentric_coordinates


def get_barycentric_coordinates(
    point:Float[Array, 'dim'], # Query point.
    a:Float[Array, 'dim'], b:Float[Array, 'dim'], c:Float[Array, 'dim']
)->Float[Array, '3']: # Barycentric coordinates [lambda_a, lambda_b, lambda_c] summing to 1.

Barycentric coordinates of point w.r.t. triangle abc.

Uses the area-ratio method: each coordinate is the ratio of the sub-triangle area to the full triangle area. In 3D the point is projected onto the triangle plane (least-squares sense).

vertices = jnp.array([[0., 0.], [0., 1.], [1., 0.]])
point1 = jnp.array([0.5, 0.2])

get_barycentric_coordinates(point1, *vertices)
Array([0.3, 0.2, 0.5], dtype=float64)
vertices2 = jnp.array([[0., 0., 0.], [0., 1., 0.], [1., 0., 0.]])
point2 = jnp.array([0.5, 0.2, 0.])

get_barycentric_coordinates(point2, *vertices2)
Array([0.3, 0.2, 0.5], dtype=float64)
point3 = jnp.array([0.5, 0.2, 1.])

get_barycentric_coordinates(point3, *vertices2)
Array([0.3, 0.2, 0.5], dtype=float64)

Rodrigues rotation

source

rotate_around_axis


def rotate_around_axis(
    v:Float[Array, '3'], # Vector to rotate.
    axis:Float[Array, '3'], # Unit rotation axis.
    angle:Float[Array, ''], # Rotation angle in radians.
)->Float[Array, '3']: # Rotated vector.

Rotate 3D vector v by angle (radians) around unit axis using Rodrigues’ formula.

# test: rotating x-axis by 90° around z-axis should give y-axis
x = jnp.array([1., 0., 0.])
z = jnp.array([0., 0., 1.])
result = rotate_around_axis(x, z, jnp.array(jnp.pi/2))
assert jnp.allclose(result, jnp.array([0., 1., 0.]), atol=1e-10)

# test: rotating by 0 should return the same vector
assert jnp.allclose(rotate_around_axis(x, z, jnp.array(0.)), x, atol=1e-10)

# test: rotating around the vector itself should return it
v = jnp.array([1., 2., 3.])
axis = v / jnp.linalg.norm(v)
assert jnp.allclose(rotate_around_axis(v, axis, jnp.array(1.23)), v, atol=1e-10)