triangulax 0.0.1

JAX-compatible triangular meshes and triangular-mesh-based simulations

triangulax

Overview

This Python package provides data-structures for triangular meshes and geometry processing tools based on JAX, fully compatible with automatic differentiation and just-in-time compilation.

Use cases

Triangular meshes are ubiquitous in computer graphics and in scientific computing, for example in the finite-element method. A major motivation for triangulax are simulations in soft-matter and biophysics. For example, with triangulax, you can simulate:

1. Cell-resolved models of two-dimensional tissue sheets like the self-propelled Voronoi model (tutorial 3).
2. Reaction-diffusion systems on 3D curved surfaces (tutorial 4)
3. Mechanics of membranes and thin elastic shells in 3D (tutorial 5)

triangulax is designed for modularity and flexibility. The library includes a suite of geometry processing tools based on discrete differential geometry (Voronoi duals, curvatures, Laplace operator, …) and represents surfaces as half-edge meshes, which allows implementing custom simulations and geometry operations in any dimension.

Prerequisites: Using triangulax assumes familiarity with triangular meshes, and basic JAX usage. See tutorials 0 and 1.

Automatic differentiation and just-in-time compilation

The main feature of triangulax is compatibility with (forward- and reverse-mode) automatic differentiation. This enables computation of gradients of any mesh-based function. Most tools are also compatible with JAX’s JIT-compilation. This delivers high performance in high-level Python (rather than C++) and allows running on GPUs.

Simulating with automatic differentiation

For example, consider:

1. Flattening or deforming 3D models (computer graphics)
2. Mechanics of thin plates or membranes (mechanics)
3. Cell resolved tissue simulations

These tasks revolve around a mesh-based “energy” (like the Dirichlet variational functional, the Helfrich elastic energy, or the area-perimeter energy, respectively). JAX automatically computes their gradients, making it easy to optimize energies or to simulate forces. For “multiphysics” simulations (for example, reaction-diffusion dynamics on a deforming surface), it suffices to specify the combined energy of the system - all forces and cross-terms are calculated automatically.

To simulate dynamics or minimize energies, triangulax integrates seamlessly with the JAX ecosystem like optimistix (optimization) and diffrax (ODE integration).

Inverse problems and bilevel optimization

Since triangulax is fully JAX-compatible, you can differentiate a simulation w.r.t. its parameters. This means one can apply gradient-based optimization to inverse problems (tutorial 2). Effectively, your simulation becomes a “neural network” which maps initial conditions to simulation results. triangulax can be used with optimistix and diffrax which allows straight-through and adjoint-based differentiation.

For concreteness, consider an elastic energy $E({\mathbf{v}_i}_i ; \boldsymbol{\theta})$ that depends on both mesh vertex positions $\mathbf{v}_i$ and parameters $\boldsymbol{\theta}$ (elastic moduli, spring rest lengths, etc). The minimum-energy configuration $\mathbf{v}i^* = \arg \min E( \cdot ; \boldsymbol{\theta})$ depends on $\boldsymbol{\theta}$, and via automatic differentiation, you can compute $\partial{\boldsymbol{\theta}} \mathbf{v}_i^*$. This means you can use gradient-based optimization to find a $\boldsymbol{\theta}$ such that the minimum-energy configuration has a desired shape. For example, in the tissue mechanics context, you can ask: what do individual cells need to do so that the tissue as a whole takes on a certain shape?

See also

• libigl Geometry processing library with Python bindings. You can use libigl functions on triangulax meshes via the .faces attribute

• VertAX JAX-based simulations of 2D tissues.

Developer guide and installation instructions

Currently, installation from source only - PyPi package coming soon

This package is developed based on Jupyter notebooks, which are converted into python modules using nbdev. Take a look at .github/copilot-instructions.md for details.

Install triangulax in Development mode

1. Clone the GitHub repository

$ git clone https://github.com/nikolas-claussen/triangulax.git

2. Create a conda environment with all Python dependencies

$ conda env create -n triangulax -f environment.yml
$ conda activate triangulax

3. Install the triangulax package

# make sure triangulax package is installed in development mode
$ pip install -e .

4. If necessary, edit the package notebooks and export

# make changes under nbs/ directory
# ...

# compile to have changes apply to triangulax
$ nbdev_export

Documentation

Documentation can be found hosted on this GitHub repository’s pages. Jupyter notebooks tutorials can be found in the nbs/tutorials/ folder.

Usage

triangulax comprises the following modules (see full documentation for details):

• trigonometry: trigonometry and linear algebra utilities
• triangular: input/output for triangular meshes
• mesh: half-edge data structure for triangular meshes, compatible with JAX
• topology: topological modifications (edge flip, collapse, and split)
• adjacency: vertex-vertex, vertex-face, and face-face adjacency operators
• geometry: angles, edge lengths, triangle areas, Voronoi areas, curvatures
• periodic: geometry computations under periodic boundary conditions
• linops: discrete differential operators (Laplacian, mass matrix, gradient)
• algorithms: Delaunay flipping, mesh quality improvement

Minimal example

import igl
import jax
import jax.numpy as jnp
from triangulax import mesh, geometry

# load example mesh and convert to half-edge mesh

vertices, _, _, faces, _, _ = igl.readOBJ("test_meshes/disk.obj")
hemesh = mesh.HeMesh.from_triangles(vertices.shape[0], faces)

# with the half-edge mesh, you can carry out various operations, for example
# compute the coordination number by summing incoming half-edges per vertex

coord_number = jnp.zeros(hemesh.n_vertices)
coord_number = coord_number.at[hemesh.dest].add(jnp.ones(hemesh.n_hes))
print("Mean coordination number:", coord_number.mean())

# Let's define a simple geometric function and compute its gradient with JAX

def mean_voronoi_area(vertices: jax.Array, hemesh: mesh.HeMesh) -> jax.Array:
    """Compute the mean Voronoi area per vertex."""
    voronoi_areas = geometry.get_voronoi_areas(vertices, hemesh)
    return jnp.mean(voronoi_areas)

value, gradient = jax.value_and_grad(mean_voronoi_area)(vertices, hemesh)
print("Mean gradient norm:", jnp.linalg.norm(gradient, axis=1).mean())

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

Mean coordination number: 5.40458
Mean gradient norm: 0.0003638338