triangulax

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

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.

Automatic differentiation and just-in-time compilation with JAX

The main feature of triangulax is (forward- and reverse-mode) automatic differentiation. This enables computation of gradients of any mesh-based function. To achieve this, triangulaxuses JAX, a “Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.” Most triangulax tools are compatible with JAX’s just-in-time (JIT) compilation. This delivers high performance in high-level Python (rather than C++) and allows running on your code on GPUs.

Prerequisites: Using triangulax assumes familiarity with triangular meshes (tutorial 0), and basic JAX usage (tutorial 1)

Simulating with automatic differentiation

Consider:

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

These tasks revolve around a mesh-based “energy” (like the Dirichlet energy, the Helfrich 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 with JAX ecosystem libraries like optimistix (optimization) and diffrax (ODE integration).

Inverse problems

Since triangulax is fully JAX-compatible, you can differentiate a simulation with respect to 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 example, consider an elastic energy \(E(\mathbf{v} ; \boldsymbol{\theta})\) that depends on both mesh vertex positions \(\mathbf{v}\) and parameters \(\boldsymbol{\theta}\) (elastic moduli, spring rest lengths, etc). Via automatic differentiation, you can differentiate the minimum-energy configuration \(\mathbf{v}^*[\theta] = \mathrm{argmin} E( \cdot ; \boldsymbol{\theta})\) with respect to the parameters \(\boldsymbol{\theta}\). This means you can use gradient-based optimization to find a \(\boldsymbol{\theta}\) such that the minimum-energy configuration has a desired shape. In the tissue mechanics context, you could 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.

Installation instructions

The triangulax package is hosted on PyPI. Install it as follows:

1. (Recommended) Initialize a virtual environment, for instance with conda:

$ conda env create -n triangulax 
$ conda activate triangulax

Please see the JAX documentation for how to install with GPU support.

2. Install with pip:

$ pip install triangulax 

(Optional) Install optional dependencies for running the tutorials:

$ pip install triangulax[tutorials]

3. Verify installation:

conda run -n triangulax python -c "from triangulax import mesh, geometry"

4. (Optional) Download jupyter notebooks to run tutorials interactively from GitHub

Developer guide

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. Optional: install nbdev if you want to edit the package notebooks

pip install nbdev

4. Install the triangulax package

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

5. Optional: edit the package notebooks and export

# make changes under nbs/ directory
# ...

# export 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
• simulation: Utilities for time-dependent simulations and checkpointing

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