From 7ce5ed26dbc729121b47c5920e4dbdf13b2b29c1 Mon Sep 17 00:00:00 2001 From: Tim Heap Date: Thu, 31 Oct 2024 14:43:01 +1100 Subject: [PATCH] Speed up triangulation --- src/emsarray/operations/triangulate.py | 254 ++++++++++++------ .../triangulate/test_triangulate_dataset.py | 55 ++-- 2 files changed, 202 insertions(+), 107 deletions(-) diff --git a/src/emsarray/operations/triangulate.py b/src/emsarray/operations/triangulate.py index 335d002..8981814 100644 --- a/src/emsarray/operations/triangulate.py +++ b/src/emsarray/operations/triangulate.py @@ -1,19 +1,20 @@ """ Operations for making a triangular mesh out of the polygons of a dataset. """ -from typing import cast - import numpy +import pandas +import shapely import xarray from shapely.geometry import LineString, MultiPoint, Polygon Vertex = tuple[float, float] -Triangle = tuple[int, int, int] +VertexTriangle = tuple[Vertex, Vertex, Vertex] +IndexTriangle = tuple[int, int, int] def triangulate_dataset( dataset: xarray.Dataset, -) -> tuple[list[Vertex], list[Triangle], list[int]]: +) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: """ Triangulate the polygon cells of a dataset @@ -28,17 +29,22 @@ def triangulate_dataset( Returns ------- - tuple of vertices, triangles, and cell indexes. - A tuple of three lists is returned, + tuple of vertices, triangles, and `cell_indices` + A tuple of three numpy arrays is returned, containing vertices, triangles, and cell indexes respectively. - Each vertex is a tuple of (x, y) or (lon, lat) coordinates. + `vertices` is a numpy array of shape (V, 2) + where V is the number of unique vertices in the dataset. + The vertex coordinates are in (x, y) or (lon, lat) order. + + `triangles` is a numpy array of shape (T, 3) + where T is the number of triangles in the dataset. + Each triangle is a set of three vertex indices. + + `cell_indices` is a numpy list of length T. + Each entry indicates which polygon from the dataset a triangle is a part of. - Each triangle is a tuple of three integers, - indicating which vertices make up the triangle. - The cell indexes tie the triangles to the original cell polygon, - allowing you to plot data on the triangle mesh. Examples -------- @@ -87,79 +93,160 @@ def triangulate_dataset( """ polygons = dataset.ems.polygons - # Getting all the vertices is easy - extract them from the polygons. - # By going through a set, this will deduplicate the vertices. - # Back to a list and we have a stable order - vertices: list[Vertex] = list({ - vertex - for polygon in polygons - if polygon is not None - for vertex in polygon.exterior.coords - }) - - # This maps between a vertex tuple and its index. - # Vertex positions are (probably) floats. For grid datasets, where cells - # are implicitly defined by their centres, be careful to compute cell - # vertices in consistent ways. Float equality is tricky! - vertex_indexes = {vertex: index for index, vertex in enumerate(vertices)} - - # Each cell polygon needs to be triangulated, - # while also recording the convention native index of the cell, - # so that we can later correlate cell data with the triangles. - polygons_with_index = [ - (polygon, index) - for index, polygon in enumerate(polygons) - if polygon is not None] - triangles_with_index = list( - (tuple(vertex_indexes[vertex] for vertex in triangle_coords), dataset_index) - for polygon, dataset_index in polygons_with_index - for triangle_coords in _triangulate_polygon(polygon) - ) - triangles: list[Triangle] = [tri for tri, index in triangles_with_index] # type: ignore - indexes = [index for tri, index in triangles_with_index] - - return (vertices, triangles, indexes) - - -def _triangulate_polygon(polygon: Polygon) -> list[tuple[Vertex, Vertex, Vertex]]: + # Find all the unique coordinates and assign them each a unique index + all_coords = shapely.get_coordinates(polygons) + vertex_index = pandas.MultiIndex.from_arrays(all_coords.T).drop_duplicates() + vertex_series = pandas.Series(numpy.arange(len(vertex_index)), index=vertex_index) + vertex_coords = numpy.array(vertex_index.to_list()) + + polygon_length = shapely.get_num_coordinates(polygons) + + # Count the total number of triangles. + # A polygon with n sides can be decomposed in to n-2 triangles, + # however polygon_length counts an extra vertex because of the closed rings. + total_triangles = numpy.sum(polygon_length[numpy.nonzero(polygon_length)] - 3) + + # Pre-allocate some numpy arrays that will store the triangle data. + # This is much faster than building up these data structures iteratively. + face_indices = numpy.empty(total_triangles, dtype=int) + triangle_coords = numpy.empty((total_triangles, 3, 2), dtype=float) + # This is the index of which face we will populate next in the above arrays + current_face = 0 + + def _add_triangles(face_index: int, vertex_triangles: numpy.ndarray) -> None: + """ + Append some triangles to the face_indices and triangle_coords arrays. + + Parameters + ---------- + face_index : int + The face index of all the triangles + vertex_triangles : numpy.ndarray + The triangles for this face as an (n, 3, 2) array, + where n is the number of triangles in the face. + """ + nonlocal current_face + current_length = len(vertex_triangles) + face_indices[current_face:current_face + current_length] = face_index + triangle_coords[current_face:current_face + current_length] = vertex_triangles + current_face += current_length + + # Find all concave polygons by comparing each polygon to its convex hull. + # A convex polygon is its own convex hull, + # while the convex hull of a concave polygon + # will always have fewer vertices than the original polygon. + # Comparing the number of vertices is a shortcut. + convex_hulls = shapely.convex_hull(polygons) + convex_hull_length = shapely.get_num_coordinates(convex_hulls) + polygon_is_concave = numpy.flatnonzero(convex_hull_length != polygon_length) + + # Categorize each polygon by length, skipping concave polygons. + # We will handle them separately. + polygon_length[polygon_is_concave] = 0 + unique_lengths = numpy.unique(polygon_length) + + # Triangulate polygons in batches of identical sizes. + # This allows the coordinates to be manipulated efficiently. + for unique_length in unique_lengths: + if unique_length == 0: + # Any `None` polygons will have a length of 0, + # and any concave polygons have been set to 0. + continue + + same_length_face_indices = numpy.flatnonzero(polygon_length == unique_length) + same_length_polygons = polygons[same_length_face_indices] + vertex_triangles = _triangulate_polygons_by_length(same_length_polygons) + + for face_index, triangles in zip(same_length_face_indices, vertex_triangles): + _add_triangles(face_index, triangles) + + # Triangulate each concave polygon using a slower manual method. + # Anecdotally concave polygons are very rare, + # so using a slower method isn't an issue. + for face_index in polygon_is_concave: + polygon = polygons[face_index] + triangles = _triangulate_concave_polygon(polygon) + _add_triangles(face_index, triangles) + + # Check that we have handled each triangle we expected. + assert current_face == total_triangles + + # Make a DataFrame. By manually constructing Series the data in the + # underlying numpy arrays will be used in place. + face_triangle_df = pandas.DataFrame({ + 'face_indices': pandas.Series(face_indices), + 'x0': pandas.Series(triangle_coords[:, 0, 0]), + 'y0': pandas.Series(triangle_coords[:, 0, 1]), + 'x1': pandas.Series(triangle_coords[:, 1, 0]), + 'y1': pandas.Series(triangle_coords[:, 1, 1]), + 'x2': pandas.Series(triangle_coords[:, 2, 0]), + 'y2': pandas.Series(triangle_coords[:, 2, 1]), + }, copy=False) + + joined_df = face_triangle_df\ + .join(vertex_series.rename('v0'), on=['x0', 'y0'])\ + .join(vertex_series.rename('v1'), on=['x1', 'y1'])\ + .join(vertex_series.rename('v2'), on=['x2', 'y2']) + + faces = joined_df['face_indices'].to_numpy() + triangles = joined_df[['v0', 'v1', 'v2']].to_numpy() + + return vertex_coords, triangles, faces + + +def _triangulate_polygons_by_length(polygons: numpy.ndarray) -> numpy.ndarray: """ - Triangulate a polygon. + Triangulate a list of convex polygons of equal length. - .. note:: + Parameters + ---------- + polygons : numpy.ndarray of shapely.Polygon + The polygons to triangulate. + These must all have the same number of vertices + and must all be convex. - This currently only supports simple polygons - polygons that do not - intersect themselves and do not have holes. + Returns + ------- + numpy.ndarray + A numpy array of shape (# polygons, # triangles, 3, 2), + where `# polygons` is the length of `polygons` + and `# triangles` is the number of triangles each polygon is decomposed in to. + """ + vertex_count = len(polygons[0].exterior.coords) - 1 + + # An array of shape (len(polygons), vertex_count, 2) + coordinates = shapely.get_coordinates(shapely.get_exterior_ring(polygons)) + coordinates = coordinates.reshape((len(polygons), vertex_count + 1, 2)) + coordinates = coordinates[:, :-1, :] + + # Arrays of shape (len(polygons), vertex_count - 2, 2) + v0 = numpy.repeat( + coordinates[:, 0, :].reshape((-1, 1, 2)), + repeats=vertex_count - 2, + axis=1) + v1 = coordinates[:, 1:-1] + v2 = coordinates[:, 2:] + + # An array of shape (len(polygons), vertex_count - 2, 3, 2) + triangles: numpy.ndarray = numpy.stack([v0, v1, v2], axis=2) + return triangles - Examples - -------- - .. code-block:: python +def _triangulate_concave_polygon(polygon: Polygon) -> numpy.ndarray: + """ + Triangulate a single convex polygon. - >>> polygon = Polygon([(0, 0), (2, 0), (2, 2), (1, 3), (0, 2), (0, 0)]) - >>> for triangle in triangulate_polygon(polygon): - ... print(triangle.wkt) - POLYGON ((0 0, 2 0, 2 2, 0 0)) - POLYGON ((0 0, 2 2, 1 3, 0 0)) - POLYGON ((0 0, 1 3, 0 2, 0 0)) + Parameters + ---------- + polygon : shapely.Polygon + The polygon to triangulate. - See Also - -------- - :func:`triangulate_dataset`, - `Polygon triangulation `_ + Returns + ------- + numpy.ndarray + A numpy array of shape (# triangles, 3, 2), + where `# triangles` is the number of triangles the polygon is decomposed in to. """ - # The 'ear clipping' method used below is correct for all polygons, but not - # performant. If the polygon is convex we can use a shortcut method. - if polygon.equals(polygon.convex_hull): - # Make a fan triangulation. For a polygon with n vertices the triangles - # will have vertices: - # (1, 2, 3), (1, 3, 4), (1, 4, 5), ... (1, n-1, n) - exterior_vertices = numpy.array(polygon.exterior.coords)[:-1] - num_triangles = len(exterior_vertices) - 2 - v0 = numpy.broadcast_to(exterior_vertices[0], (num_triangles, 2)) - v1 = exterior_vertices[1:-1] - v2 = exterior_vertices[2:] - return list(zip(map(tuple, v0), map(tuple, v1), map(tuple, v2))) - # This is the 'ear clipping' method of polygon triangulation. # In any simple polygon, there is guaranteed to be at least two 'ears' # - three neighbouring vertices whos diagonal is inside the polygon. @@ -171,7 +258,12 @@ def _triangulate_polygon(polygon: Polygon) -> list[tuple[Vertex, Vertex, Vertex] # Most polygons will be either squares, convex quadrilaterals, or convex # polygons. - triangles: list[tuple[Vertex, Vertex, Vertex]] = [] + # A triangle with n vertices will have n - 2 triangles. + # Because the exterior is a closed loop, we need to subtract 3. + triangle_count = len(polygon.exterior.coords) - 3 + triangles = numpy.empty((triangle_count, 3, 2)) + triangle_index = 0 + # Note that shapely polygons with n vertices will be closed, and thus have # n+1 coordinates. We trim that superfluous coordinate off in the next line while len(polygon.exterior.coords) > 4: @@ -192,7 +284,8 @@ def _triangulate_polygon(polygon: Polygon) -> list[tuple[Vertex, Vertex, Vertex] # that ear will result in two disconnected polygons. and exterior.intersection(diagonal).equals(multipoint) ): - triangles.append((coords[i], coords[i + 1], coords[i + 2])) + triangles[triangle_index] = coords[i:i + 3] + triangle_index += 1 polygon = Polygon(coords[:i + 1] + coords[i + 2:]) break else: @@ -201,8 +294,7 @@ def _triangulate_polygon(polygon: Polygon) -> list[tuple[Vertex, Vertex, Vertex] f"Could not find interior diagonal for polygon! {polygon.wkt}") # The trimmed polygon is now a triangle. Add it to the list and we are done! + triangles[triangle_index] = polygon.exterior.coords[:-1] + assert (triangle_index + 1) == triangle_count - triangles.append(cast( - tuple[Vertex, Vertex, Vertex], - tuple(map(tuple, polygon.exterior.coords[:-1])))) return triangles diff --git a/tests/operations/triangulate/test_triangulate_dataset.py b/tests/operations/triangulate/test_triangulate_dataset.py index 50cbdf9..5bb9c3d 100644 --- a/tests/operations/triangulate/test_triangulate_dataset.py +++ b/tests/operations/triangulate/test_triangulate_dataset.py @@ -1,15 +1,14 @@ from collections import defaultdict -from functools import reduce import numpy -import pytest import shapely import xarray from shapely.geometry import Polygon import emsarray from emsarray.operations.triangulate import ( - _triangulate_polygon, triangulate_dataset + _triangulate_concave_polygon, _triangulate_polygons_by_length, + triangulate_dataset ) @@ -114,17 +113,35 @@ def check_triangulation( triangles = cell_triangles[index] # Turn them in to polygons... - polygons = [ - Polygon([vertices[i] for i in triangle]) - for triangle in triangles - ] - # Union the those together in to one large polygon - union = reduce(lambda a, b: a.union(b), polygons) + reconstructed_polygon = shapely.unary_union(shapely.polygons(vertices[triangles])) # Check it matches - assert polygon.equals(union) + assert polygon.equals(reconstructed_polygon) -def test_triangulate_polygon(): +def test_triangulate_polygons_by_length(): + hexagon_template = numpy.array([[-2, 0], [-1, -1], [1, -1], [2, 0], [1, 1], [-1, 1], [-2, 0]]) + polygons = shapely.polygons([ + hexagon_template + (0, 0), + hexagon_template + (3, 1), + hexagon_template + (6, 0), + hexagon_template + (0, 2), + hexagon_template + (6, 2), + ]) + polygon_triangles = _triangulate_polygons_by_length(polygons) + + # Shape is (# polygons, # triangles, # vertices, # coords). + # Each hexagon can be decomposed in to 4 triangles, + # each triangle has three vertices, + # and each vertex has two coordinates. + assert polygon_triangles.shape == (len(polygons), 4, 3, 2) + + # Reconstruct each polygon and check it equals itself. + for polygon, triangles in zip(polygons, polygon_triangles): + reconstructed_polygon = shapely.unary_union(shapely.polygons(triangles)) + assert polygon.equals(reconstructed_polygon) + + +def test_triangulate_convex_polygon(): # These coordinates are carefully chosen to produce a polygon that: # * is not convex # * has three non-sequential vertices in a row @@ -135,22 +152,8 @@ def test_triangulate_polygon(): assert polygon.is_valid assert polygon.is_simple - triangles = _triangulate_polygon(polygon) + triangles = _triangulate_concave_polygon(polygon) assert len(triangles) == len(coords) - 2 union = shapely.union_all(shapely.polygons(numpy.array(triangles))) assert union.equals(polygon) - - -@pytest.mark.parametrize("poly", [ - # Polygon with a hole - Polygon( - [(0, 0), (3, 0), (3, 3), (0, 3), (0, 0)], - [[(1, 1), (1, 2), (2, 2), (2, 1), (1, 1)]], - ), - # Polygon that intersects itself - Polygon([(0, 0), (1, 0), (0, 1), (1, 1), (0, 0)]), -]) -def test_triangulate_polygon_non_simple(poly): - with pytest.raises(ValueError): - _triangulate_polygon(poly)