mesh: Half-edge meshes

In the TriMesh class, we represent a mesh a list of triangles. However, many common operations are difficult with this data structure. For example, how do you get all the neighbors of a given vertex, or compute the area of a dual cell?

For simulation and geometry processing, we need a different representation of the adjacency information. Typically, this is achieved by a half-edge mesh (HE) data structure. We represent the HE data structure by 3 sets of integer index arrays:

1. Vertices: 1 \((N_V,)\) array, whose entry for vertex \(i\) is an arbitrary HE incident on \(i\)
2. Edges: 6 \((2N_E,)\) arrays, [origin, dest, nxt, prv, twin, face] for each half-edge (face is undefined for boundary half-edges).
3. Faces, 1 \((N_F, 1)\) array, whose entry for face \(i\) is an arbitrary HE in \(i\). (Not to be confused with the \((N_F, 3)\) array of vertex IDs used previously).

Additionally, there are two float arrays for vertex and face positions, as previously. However, we split combinatorial and geometric information - a HeMesh class for the combinatorics, and a couple of regular arrays for the vertex positions, face positions, and vertex/half-edge/face attributes. The latter are packaged into a GeomMesh class. Together, the pair (GeomMesh, HeMesh) describes a mesh (like vertices/faces pair). A named tuple Mesh combines the two.

The first task is to create a helper function to plot mesh connectivity, and to create the half-edge connectivity matrices from the more conventional list-of-triangles format. The latter is somewhat involved.

For JAX compatibility, the mesh module uses jax.numpy instead of standard numpy for all numerical arrays, follows a functional programming style (no in-place mutation), and registers the HeMesh and GeomMesh dataclasses as JAX pytrees (this enables automatic differentiation and JIT-compilation with custom datastructures).

---

source

label_plot

def label_plot(
    vertices:Float[Array, 'n_vertices 2'], faces:Int[Array, 'n_faces 3'], hemesh:Union=None, vertex_labels:bool=True,
    face_labels:bool=True, ax:Union=None, fontsize:Union=None
)->None:

For debugging purposes. Plot triangular mesh with face/vertex labels in black/blue. If hemesh is not None, the connectivity info from it is used to plot the half-edge labels.

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)

plt.triplot(*mesh.vertices.T, mesh.faces)
label_plot(mesh.vertices, mesh.faces, fontsize=10)
plt.axis("equal")

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

(np.float64(-1.10003475),
 np.float64(1.09628575),
 np.float64(-1.09934025),
 np.float64(1.09050125))

---

source

get_half_edge_arrays_vectorized

def get_half_edge_arrays_vectorized(
    n_vertices:int, faces:Int[Array, 'n_faces 3']
)->list:

Get half-edge data structure arrays from faces (vectorized). Returned arrays are dtype int32. For internal use to construct HeMesh objects.

Returns: incident, orig, dest, twin, nxt, prv, heface, face_incident

mesh_high_res = TriMesh.read_obj("../test_meshes/torus_high_resolution.obj")
mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)

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

59.5 ms ± 968 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# test vectorized vs reference implementation for two meshes

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
ref = get_half_edge_arrays(mesh.vertices.shape[0], mesh.faces)
fast = get_half_edge_arrays_vectorized(mesh.vertices.shape[0], mesh.faces)

print("Equal?", all([jnp.array_equal(a, b) for a, b in zip(ref, fast)]))

mesh = TriMesh.read_obj("../test_meshes/sphere.obj")
ref = get_half_edge_arrays(mesh.vertices.shape[0], mesh.faces)
fast = get_half_edge_arrays_vectorized(mesh.vertices.shape[0], mesh.faces)

print("Equal?", all([jnp.array_equal(a, b) for a, b in zip(ref, fast)]))

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

Equal? True
Equal? True

1.08 ms ± 29.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

---

source

HeMesh

def HeMesh(
    incident:Int[Array, '*batch n_vertices'], orig:Int[Array, '*batch n_hes'], dest:Int[Array, '*batch n_hes'],
    twin:Int[Array, '*batch n_hes'], nxt:Int[Array, '*batch n_hes'], prv:Int[Array, '*batch n_hes'],
    heface:Int[Array, '*batch n_hes'], face_incident:Int[Array, '*batch n_faces'], inf_vertices:Union=()
)->None:

Half-edge mesh data structure for triangular meshes.

A half-edge mesh is described by a set of half-edges and several arrays that specify their connectivity (see full explanation in mesh module docs). This class serves as a container for multiple arrays. For compatibility with JAX, after initialization, do not modify these arrays in-place; always return a new HeMesh object. The mesh vertices may live in whatever dimension (or in periodic BC) - this does not affect the connectivity bookkeeping.

Half-edge meshes are initialized from a list of triangles and a number of vertices, and can return the original triangles (e.g., to save as a .obj).

All information and methods are purely “combinatorial”. The HeMesh class does not contain the vertex or face positions. These are saved in the GeomHeMesh class that combines a HeMesh (combinatorics) with a couple of other arrays (geometry).

Comparing two HeMeshes checks for equality of all arrays they contain, not for graph isomorphism (equivalence up to vertex renaming).

—Conventions—

For vertices, the incident half-edge points away from the vertex.

To describe the mesh boundary, there are two options: 1. Initialize from a triangulation with a boundary. Half-edges without a face (boundary) are assigned heface=-1. 2. Initialize from a triangulation without boundary, where certain vertices are “at infinity”. They should have coordinates [np.inf, np.inf]. Each infinity vertex corresponds to one boundary. For a single boundary, the vertex at infinity is, by convention, the final one. Mixing the two conventions will lead to errors.

Starting from a set of triangles, the half-edges are initialized as follows: The 1st N_edges half-edges are (origin_vertex, destination_vertex), in lexicographic order, with origin_vertex < destination_vertex. The 2nd N_edges are their twins, in the same order.

Attributes

incident : Int[jax.Array, “n_vertices”]

orig : Int[jax.Array, “n_hes”]

dest : Int[jax.Array, “n_hes”]

nxt : Int[jax.Array, “n_hes”]

prv : Int[jax.Array, “n_hes”]

twin : Int[jax.Array, “n_hes”]

heface : Int[jax.Array, “n_hes”]

face_incident : Int[jax.Array, “n_faces”]

inf_vertices : tuple[Int]

Property methods (use like attributes)

n_vertices : int

n_hes : int

n_faces : int

n_items : tuple[int, int, int]

faces : Int[jax.Array, “n_faces 3”]

has_inf_vertex : bool

is_inf_face : Bool[jax.Array, “n_faces”]

is_unique : Bool[jax.Array, “n_hes”]

is_inf_he : Bool[jax.Array, “n_hes”]

is_bdry_he : Bool[jax.Array, “n_hes”]

is_bdry_edge : Bool[jax.Array, “n_hes”]

is_bdry : Bool[jax.Array, “n_vertices”]

Static methods

from_triangles : tuple[int, Int[jax.Array, “n_faces 3”], Int[jax.Array, “n_boundaries”] -> HeMesh

The arguments are the number of vertices, the n_faces * 3 array of faces, and, optionally, a list of “infinity” vertices.

Class methods

iterate_around_vertex : int -> Int[jax.Array, “n_neighbors”]

save : str -> None:

Static methods

load : str -> HeMesh

---

source

test_mesh_validity

def test_mesh_validity(
    h:HeMesh, # Half-edge mesh to validate.
    verbose:bool=False, # Return diagnostic message.
)->bool: # True if the mesh is valid. Also returns a string with a message if verbose = True

Test if a mesh is valid. Returns True if valid, optionally, returns message.

Checks both for consistency of the HeMesh datastructure AND whether the mesh is edge- and vertex manifold (defines a 2D surface).

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)

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

test_mesh_validity(hemesh)

True

# hemeshes can be compared for equality and are registered as pytrees

leafs, ts = jax.tree_util.tree_flatten(hemesh)
assert len(leafs) == 8                                  # the 8 connectivity arrays
assert jax.tree_util.tree_unflatten(ts, leafs) == hemesh

assert hemesh == hemesh
assert not (hemesh == "not a mesh")

# equality must be EXACT: these are integer index arrays, so a tolerance-based
# comparison would silently accept an off-by-one index on a large mesh.
corrupted = HeMesh(hemesh.incident, hemesh.orig, hemesh.dest,
                   hemesh.twin.at[0].set(hemesh.twin[0] + 1),
                   hemesh.nxt, hemesh.prv, hemesh.heface, hemesh.face_incident, hemesh.inf_vertices)
assert not (hemesh == corrupted)

hemesh

HeMesh(N_V=131, N_HE=708, N_F=224)

# test iteration around vertex
hemesh.dest[hemesh.iterate_around_vertex(69)], hemesh.orig[hemesh.iterate_around_vertex(56)]

(Array([80, 68, 56, 46], dtype=int32),
 Array([56, 56, 56, 56, 56, 56, 56], dtype=int32))

# boundary in cc-wise order
(hemesh.orig[187], hemesh.dest[187]), hemesh.heface[187], hemesh.is_bdry_he[187],

((Array(58, dtype=int32), Array(70, dtype=int32)),
 Array(-1, dtype=int32),
 Array(True, dtype=bool))

hemesh.is_bdry_he[187], hemesh.is_bdry_he[541], hemesh.heface[541]

(Array(True, dtype=bool), Array(False, dtype=bool), Array(145, dtype=int32))

# to model mesh boundaries, we can add an "infinity" vertex. Not done here, see below
hemesh.has_inf_vertex, hemesh.inf_vertices

(False, ())

fig = plt.figure(figsize=(14,14))

plt.triplot(*mesh.vertices.T, hemesh.faces)
label_plot(mesh.vertices, hemesh.faces, fontsize=10, hemesh=hemesh, face_labels=False)
plt.axis("equal")

(np.float64(-1.10003475),
 np.float64(1.09628575),
 np.float64(-1.09934025),
 np.float64(1.09050125))

# here is how you would do traversal of vertex neighbors with jax.lax. In JAX, the output size needs to be fixed
# ahead of time, so this requires padding and setting a cap on vertex valence (inefficient and error-prone).

self = hemesh
max_valence = 10
v = 10

initial = jnp.hstack([jnp.array([self.incident[v]]), -1*jnp.ones(max_valence-1, dtype=int)])
jax.lax.fori_loop(1, max_valence, lambda i, x: x.at[i].set(self.twin[x[i-1]]), initial)

Array([ 47, 401,  47, 401,  47, 401,  47, 401,  47, 401], dtype=int64)

Computing with half-edge meshes

By using the arrays of a half-edge meshes to index vertex- or face-positions (in increasingly complex ways), we can compute all sorts of quantities of interests associated with a mesh, for example the edge lengths.

edges = mesh.vertices[hemesh.orig]-mesh.vertices[hemesh.dest]
lengths = jnp.linalg.norm(edges, axis=-1)

Boundary and the vertex at infinity

So far, our mesh representations TriMesh and HeMesh work for triangular meshes with and without boundary. In the HeMesh class, boundary half-edges are assigned to a fictitious -1 face. This convention has a downside. It is not possible to modify the boundary loop of the mesh by edge flips - doing so would result in an invalid state. In a simulation, this artificially limits the mesh’s ability to deform. Instead, we can add “vertices at infinity” and connect all edges in a given boundary to \(\infty\). This turns the mesh into a topological sphere. Now, one can flip boundary edges without the overall number of half-edges changing (so the array shape stays the same). Multiple boundaries are also supported. Each boundary corresponds to a distinct \(\infty\)-vertex (for example, 2 for a cylinder).

The coordinates of the fictitious vertices are set to [np.inf, np.inf] by convention. The boundary is found by iterating around \(\infty\). By convention, \(\infty\)-vertices, if they exist, are the final vertices of the mesh (don’t rely on this - implementation detail).

We generally assume that the mesh has only a single connected component.

The HeMesh class can deal with both the -1-face and the \(\infty\)-vertices conventions. The latter are listed in the inf_vertices attribute of a HeMesh.

---

source

connect_boundary_to_infinity

def connect_boundary_to_infinity(
    vertices:Float[Array, 'n_vertices 2'], # Vertex positions.
    faces:Int[Array, 'n_faces 3'], # Faces (triangles) as list of vertex indices.
)->tuple: # Vertex positions with infinity vertices appended.
One infinity vertex per boundary loop.

Connect boundary loop(s) to infinity.

New vertices are appended to the end of vertex array and have coordinates [np.inf, np.inf].

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)

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

new_vertices, new_faces, infinity_vertices = connect_boundary_to_infinity(mesh.vertices, mesh.faces)
mesh_infty = TriMesh(vertices=new_vertices, faces=new_faces)
hemesh_infty = HeMesh.from_triangles(mesh_infty.vertices.shape[0], mesh_infty.faces,
                                     inf_vertices=infinity_vertices)

# to get back the original faces/vertices, do this:

_ = hemesh_infty.faces[~hemesh_infty.is_inf_face]

# if you want to re-index the triangles so they only refer to non-infinity vertices:
_ = igl.remove_unreferenced(new_vertices, np.asarray(hemesh_infty.faces[~hemesh_infty.is_inf_face]) )

# the "vertex at infinity" convention must agree with the heface == -1 convention
assert test_mesh_validity(hemesh_infty)
assert (hemesh.is_bdry == (hemesh_infty.is_bdry[:-1] > 0)).all()
assert hemesh_infty.has_inf_vertex and not hemesh.has_inf_vertex

# boundary loops must have the same vertices AND the same orientation under both
# conventions (the inf-vertex traversal runs the other way round and is reversed)
loop_a, loop_b = np.asarray(hemesh.bdry_loops[0]), np.asarray(hemesh_infty.bdry_loops[0])
assert set(loop_a) == set(loop_b)
shift = np.where(loop_b == loop_a[0])[0][0]
assert (np.roll(loop_b, -shift) == loop_a).all(), "boundary loop orientation differs between conventions"

Topological summary

These are combinatorial invariants of the connectivity: they take a HeMesh and return a plain Python int/bool. They are host-side helpers (not jittable), which is fine since connectivity is static. Fictitious faces and vertices from the “vertex at infinity” boundary convention are excluded, so both boundary conventions give the same answer.

---

source

get_genus

def get_genus(
    hemesh:HeMesh
)->int:

Genus of the surface, from chi = 2n_components - 2genus - n_boundary_loops.

0 for a sphere or disk, 1 for a torus. Assumes an orientable mesh (which the half-edge construction guarantees). Raises if the result is not an integer, which indicates an invalid mesh.

---

source

get_n_connected_components

def get_n_connected_components(
    hemesh:HeMesh
)->int:

Number of edge-connected components of the mesh.

---

source

get_euler_characteristic

def get_euler_characteristic(
    hemesh:HeMesh
)->int:

Euler characteristic chi = V - E + F.

2 for a sphere, 1 for a disk, 0 for a torus.

---

source

get_n_boundary_loops

def get_n_boundary_loops(
    hemesh:HeMesh
)->int:

Number of boundary loops (0 for a closed mesh, 1 for a disk, 2 for a cylinder).

---

source

is_closed

def is_closed(
    hemesh:HeMesh
)->bool:

True if the mesh has no boundary.

---

source

get_n_vertices_edges_faces

def get_n_vertices_edges_faces(
    hemesh:HeMesh
)->tuple:

Counts of real (non-fictitious) vertices, edges, and faces.

---

source

get_real_faces

def get_real_faces(
    hemesh:HeMesh
)->Int[ndarray, 'n_real_faces 3']:

Faces as an int64 numpy array, with fictitious (infinity) faces removed.

Useful to hand a mesh to igl, which requires int64 and has no notion of the “vertex at infinity” boundary convention.

# sphere: chi = 2, genus 0, closed;  disk: chi = 1, one boundary loop;  torus: chi = 0, genus 1
_s = TriMesh.read_obj("../test_meshes/sphere.obj", dim=3)
_t = TriMesh.read_obj("../test_meshes/torus.obj", dim=3)
sphere_h = HeMesh.from_triangles(_s.vertices.shape[0], _s.faces)
torus_h = HeMesh.from_triangles(_t.vertices.shape[0], _t.faces)

assert (get_euler_characteristic(sphere_h), get_genus(sphere_h), is_closed(sphere_h)) == (2, 0, True)
assert (get_euler_characteristic(torus_h), get_genus(torus_h), is_closed(torus_h)) == (0, 1, True)
assert (get_euler_characteristic(hemesh), get_genus(hemesh), is_closed(hemesh)) == (1, 0, False)

assert get_n_boundary_loops(sphere_h) == 0 and get_n_boundary_loops(hemesh) == 1
assert get_n_connected_components(sphere_h) == 1

# V - E + F, with E consistent with the half-edge count (no infinity vertices here)
for h in [sphere_h, torus_h, hemesh]:
    n_v, n_e, n_f = get_n_vertices_edges_faces(h)
    assert (n_v, n_e, n_f) == (h.n_vertices, h.n_hes // 2, h.n_faces)

# the two boundary conventions must agree on every invariant
assert get_euler_characteristic(hemesh_infty) == get_euler_characteristic(hemesh)
assert get_n_boundary_loops(hemesh_infty) == get_n_boundary_loops(hemesh)
assert get_genus(hemesh_infty) == get_genus(hemesh)
assert not is_closed(hemesh_infty)
print("topology:  sphere chi=2 g=0 | disk chi=1 b=1 | torus chi=0 g=1")

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

topology:  sphere chi=2 g=0 | disk chi=1 b=1 | torus chi=0 g=1

Mesh geometry and per-mesh variables

Mesh geometry (vertex and face positions) and per-mesh-item (per-face, per-half-edge, per-vertex) variables are combined into a second data class, the GeomMesh.

---

source

GeomMesh

def GeomMesh(
    vertices:Float[Array, '*batch n_vertices dim'], face_positions:Float[Array, '*batch n_faces dim']=<factory>,
    vertex_attribs:dict=<factory>, he_attribs:dict=<factory>, face_attribs:dict=<factory>
)->None:

Data class for holding mesh geometry and mesh-associated variables. To be combined with a HeMesh to specify the connectivity.

One array (for vertex positions) must always be present. A second, optional, standard entry is a set of positions for each face. The mesh coordinates can live in any dimension

Optionally, vertices, half-edges, and faces can have attributes (stored as dictionaries). The keys of the dictionary should be taken from a suitable ‘enum’. The values are ndarrays, whose 0th axis is (vertices/edges/faces). These attribute dicts are initialized empty and can be set afterwards.

Unlike HeMesh, this class is intentionally not frozen. Vertex positions and per-mesh attributes may be updated directly (e.g. during a simulation step), whereas mesh connectivity (HeMesh) should never be edited by hand.

This class stores no element counts of its own: the number of vertices, half-edges, and faces belongs to the HeMesh. Use check_compatibility(hemesh) to confirm a geometry and a connectivity match.

Attributes

vertices : Float[jax.Array, “n_vertices dim”]

face_positions : Float[jax.Array, “n_faces dim”]

vertex_attribs : dict[IntEnum, Float[jax.Array, “n_vertices …”]]

he_attribs : dict[IntEnum, Float[jax.Array, “n_hes …”]]

face_attribs : dict[IntEnum, Float[jax.Array, “n_faces …”]]

Property methods (use like attributes)

dim : int

Class methods

validate_dimensions : None

check_compatibility : HeMesh -> bool

Static methods

load : str -> GeomMesh

---

source

Mesh

def Mesh(
    args:VAR_POSITIONAL, kwargs:VAR_KEYWORD
):

Combine geometric and connectivity info into a single object.

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)
geommesh = GeomMesh(mesh.vertices, mesh.face_positions)
combined_mesh = Mesh(geommesh, hemesh)
combined_mesh

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

Mesh(geommesh=GeomMesh(D=2, N_V=131, N_HE=N/A, N_F=224), hemesh=HeMesh(N_V=131, N_HE=708, N_F=224))

leafs, ts = jax.tree_util.tree_flatten(geommesh) # also a pytree
ts

PyTreeDef(CustomNode(GeomMesh[()], [*, *, {}, {}, {}]))

geommesh, hemesh.n_vertices, geommesh.vertices.shape, geommesh.check_compatibility(hemesh), geommesh == geommesh

(GeomMesh(D=2, N_V=131, N_HE=N/A, N_F=224), 131, (131, 2), True, True)

---

source

cellplot

def cellplot(
    hemesh:HeMesh, face_positions:Float[Array, 'n_faces 2'], cell_colors:Union=None, mpl_polygon_kwargs:Union=None
)->PatchCollection:

Plot a cell tesselation.

cell_colors can be either a single color (for all cells) or a vector of rgba values. Only interior cells are plotted.

plt.triplot(*geommesh.vertices.T, hemesh.faces)
polygons = cellplot(hemesh, geommesh.face_positions,
                    cell_colors=np.array([0,0,1,0.5]), mpl_polygon_kwargs={"lw": 1, "ec": "k"})
ax = plt.gca()
ax.add_collection(polygons)

plt.axis("equal")

(np.float64(-1.10003475),
 np.float64(1.09628575),
 np.float64(-1.09934025),
 np.float64(1.09050125))

Vertex, half-edge, and face properties

In simulations, we will often want to attach extra information to a mesh’s vertices/edges/faces. In the GeomMesh class, these are saved in three dictionaries, vertex_attribs, he_attribs, face_attribs. Each key/value pair represents one property (for example, the cell target area). All values are arrays, and the first axis corresponds to the number of vertices/half-edges/faces, respectively. To keep track of the possible attributes, we use IntEnum’s as keys (this also ensures keys are hashable, as required by JAX).

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)
geommmesh = GeomMesh(mesh.vertices, mesh.face_positions)

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

# this is how you set up an enum. It is important to use IntEnum, so we can _order_ the enums.
# The precise Enum you will use depends on your application.

class VertexAttribs(IntEnum):
    TARGET_AREA = 1
    TARGET_PERIMETER = 2

class HeAttribs(IntEnum):
    EDGE_TENSION = 1

class FaceAttribs(IntEnum):
    FACE_AREA = 1

# you can iterate over enums, and they are hashable. The latter is essential for JAX!
print([a for a in VertexAttribs])
# there are multiple ways to access enum entries:
hash(VertexAttribs.TARGET_PERIMETER), HeAttribs.EDGE_TENSION, HeAttribs['EDGE_TENSION'], HeAttribs.EDGE_TENSION.name

[<VertexAttribs.TARGET_AREA: 1>, <VertexAttribs.TARGET_PERIMETER: 2>]

(2, <HeAttribs.EDGE_TENSION: 1>, <HeAttribs.EDGE_TENSION: 1>, 'EDGE_TENSION')

# at initialization, a HeMesh's attribute dictionaries are empty
geommmesh.vertex_attribs

{}

# set some attributes

key1 = jax.random.key(0)
_, key2 = jax.random.split(key1)
_, key3 = jax.random.split(key2)

geommmesh = dataclasses.replace(geommmesh, vertex_attribs={VertexAttribs.TARGET_AREA: jax.random.normal(key=key1,shape=hemesh.n_vertices),
                                                           VertexAttribs.TARGET_PERIMETER: jax.random.normal(key=key2, shape=hemesh.n_vertices)})
geommmesh = dataclasses.replace(geommmesh, he_attribs={HeAttribs.EDGE_TENSION: jax.random.normal(key=key3, shape=hemesh.n_hes)})

geommmesh.he_attribs.keys()

dict_keys([<HeAttribs.EDGE_TENSION: 1>])

Batching

In our simulations, we may want to “batch” over several initial conditions/random seeds/etc. (analogous to batching over training data in normal ML). JAX can efficiently and concisely vectorize operations over such “batch axes” with jax.vmap.

To batch over our custom data structures, we need to convert a list of HeMesh/GeomMeshe instances into a single mesh with a batch axis for the various arrays. Luckily, this can be done using JAX’s pytree tools. The resulting meshes have an extra “batch” axis in all their arrays.

---

source

tree_unstack

def tree_unstack(
    xb:PyTree, axis:int=0
)->list:

Unstack a batched pytree along axis into a list of pytrees.

---

source

tree_stack

def tree_stack(
    xs:list, axis:int=0
)->PyTree:

Stack a sequence of identical-structure pytrees along a new axis.

## Let us create a bunch of meshes with different initial positions and see if we can batch over them using vmap

key = jax.random.key(0)
sigma = 0.02

batch_geom = []
batch_he = []
for i in range(3):
    key, subkey = jax.random.split(key)
    random_noise = jax.random.normal(subkey, shape=geommmesh.vertices.shape)
    batch_geom.append(dataclasses.replace(geommmesh, vertices=geommmesh.vertices+sigma*random_noise))
    batch_he.append(copy.copy(hemesh))

# define a test function to appy over the batch

def test_function(geommesh: GeomMesh, hemesh: HeMesh) -> Float[jax.Array, " n_vertices"]:
    """Dummy test function."""
    return jnp.ones(hemesh.n_vertices)

# naive batching does not work. JAX needs a "struct-of-arrays", but a list of HeMeshes is an "array-of-structs"
# see https://stackoverflow.com/questions/79123001/storing-and-jax-vmap-over-pytrees

try:
    jax.vmap(test_function)(batch_geom, batch_he)
except ValueError as e:
    print("Expected error:", e)

Expected error: vmap got inconsistent sizes for array axes to be mapped:
  * most axes (21 of them) had size 708, e.g. axis 0 of argument geommesh[0].he_attribs[<HeAttribs.EDGE_TENSION: 1>] of type float64[708];
  * some axes (12 of them) had size 131, e.g. axis 0 of argument geommesh[0].vertices of type float64[131,2];
  * some axes (6 of them) had size 224, e.g. axis 0 of argument geommesh[0].face_positions of type float64[224,2]

# instead, we use a jax.tree.map to "push" the list axis into the underlying arrays.
# the resulting meshes have an extra batch dimension in all of their arrays.

batch_he_array = tree_stack(batch_he)
batch_geom_array = tree_stack(batch_geom)
batch_he_array, batch_geom_array, batch_geom_array.vertices.shape

(HeMesh(N_V=131, N_HE=708, N_F=224),
 GeomMesh(D=2, N_V=3, N_HE=3, N_F=3),
 (3, 131, 2))

# now it works! The result is a single object with batch axis

batch_out =jax.vmap(test_function)(batch_geom_array, batch_he_array)
batch_out.shape

(3, 131)

# we can unpack things again into a list of meshes

isinstance(tree_unstack(batch_out), list)

True

Saving to disk

We save and load TriMesh meshes as standard .obj files (with the hack of using vn lines for the face positions). The HeMesh class is basically a collection of arrays, which we can save to disk using numpy.

from tempfile import TemporaryFile

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)

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

outfile = TemporaryFile()

hemesh.save(outfile)
_ = outfile.seek(0) # simulates closing & reopening file
npzfile = np.load(outfile)
npzfile.files

['incident',
 'orig',
 'dest',
 'twin',
 'nxt',
 'prv',
 'heface',
 'face_incident',
 'inf_vertices']

outfile = TemporaryFile()

hemesh.save(outfile)
_ = outfile.seek(0) # simulates closing & reopening file
reloaded = HeMesh.load(outfile)

assert np.allclose(reloaded.faces, hemesh.faces)
assert reloaded == hemesh                       # every connectivity array round-trips
assert reloaded.inf_vertices == hemesh.inf_vertices

# test GeomMesh save/load round-trip with IntEnum keys

class _TestVA(IntEnum):
    A = 1
    B = 2
class _TestHA(IntEnum):
    C = 1

mesh = TriMesh.read_obj("../test_meshes/disk.obj", dim=2)
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)
gm = GeomMesh(mesh.vertices, mesh.face_positions,
              vertex_attribs={_TestVA.A: jnp.ones(hemesh.n_vertices),
                              _TestVA.B: jnp.zeros(hemesh.n_vertices)},
              he_attribs={_TestHA.C: jnp.ones(hemesh.n_hes)})

outfile = TemporaryFile()
gm.save(outfile)
_ = outfile.seek(0)

# with enum classes: keys round-trip as IntEnum members
gm_loaded = GeomMesh.load(outfile, vertex_attribs_enum=_TestVA, he_attribs_enum=_TestHA)
assert all(isinstance(k, _TestVA) for k in gm_loaded.vertex_attribs), "vertex keys not IntEnum"
assert all(isinstance(k, _TestHA) for k in gm_loaded.he_attribs), "he keys not IntEnum"
assert gm == gm_loaded, "loaded GeomMesh not equal to original"

# without enum classes: keys are strings (backward compat)
_ = outfile.seek(0)
gm_str = GeomMesh.load(outfile)
assert all(isinstance(k, str) for k in gm_str.vertex_attribs), "expected string keys"

print("GeomMesh save/load round-trip OK")

GeomMesh save/load round-trip OK

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