Make refine_surface not distort the point cloud to handle horizontal extensions.

discretize/tree_mesh.py

@@ -101,6 +101,7 @@
 import scipy.sparse as sp
 from discretize.utils.code_utils import deprecate_property
 from scipy.spatial import Delaunay
+from discretize.utils.mesh_utils import extend_surface_boundary
 
 
 class TreeMesh(
@@ -644,6 +645,7 @@ def refine_surface(
         pad_down=True,
         finalize=True,
         diagonal_balance=None,
+        extrapolate_surface=None,
     ):
         """Refine along a surface triangulated from xyz to the prescribed level.
 
@@ -652,8 +654,14 @@ def refine_surface(
         level.
 
         It also optionally pads the surface at each level based on the number of
-        padding cells at each dimension. It does so by stretching the bounding box of
-        the surface in each dimension.
+        padding cells at each dimension. Vertically, the surface is shifted down
+        (and/or up) by the padding amount. Horizontally, the padding is done in one
+        of two ways, set by ``extrapolate_surface``. If ``False`` (the current
+        default), the bounding box of the surface is stretched in each dimension. If
+        ``True``, points are added outward from every node along the boundary of the
+        triangulation, along the outward bisector of the adjacent boundary edges
+        (equivalent to a nearest neighbor extrapolation of the surface), see
+        :func:`~discretize.utils.extend_surface_boundary`.
 
         Parameters
         ----------
@@ -680,10 +688,26 @@ def refine_surface(
         diagonal_balance : None or bool, optional
             Whether to balance cells diagonally in the refinement, `None` implies using
             the same setting used to instantiate the TreeMesh`.
+        extrapolate_surface : None or bool, optional
+            Whether to pad horizontally by extrapolating the surface outward from the
+            nodes on its boundary, instead of stretching its bounding box. `None`
+            implies `False`, but issues a warning if horizontal padding is requested,
+            as the default will change to `True` in discretize v1.0. This has no
+            effect if there is no horizontal padding.
 
         See Also
         --------
-        refine_triangle, refine_vertical_trianglular_prism
+        refine_triangle, refine_vertical_trianglular_prism, discretize.utils.extend_surface_boundary
+
+        Notes
+        -----
+        When ``extrapolate_surface=True`` and the surface has a non-convex boundary,
+        the outward bisectors of different boundary nodes can intersect if the
+        horizontal padding is large compared to
+        the size of local features of the boundary (e.g. notches or concavities). In
+        that case the padded region can fold over itself and the result may be
+        unreliable. Keep the padding small relative to such features, or use a convex
+        boundary.
 
         Examples
         --------
@@ -743,46 +767,77 @@ def refine_surface(
         ):
             raise ValueError("incorrect dimension for padding_cells_by_level.")
 
+        has_h_pad = (
+            np.any(padding_cells_by_level[:, :-1])
+            if padding_cells_by_level.ndim == 2
+            else np.any(padding_cells_by_level)
+        )
         if self.dim == 2:
             xyz = np.asarray(xyz)
             sorter = np.argsort(xyz[:, 0])
             xyz = xyz[sorter]
-            n_ps = len(xyz)
-            inds = np.arange(n_ps)
-            simps1 = np.c_[inds[:-1], inds[1:], inds[:-1]] + [0, 0, n_ps]
-            simps2 = np.c_[inds[:-1], inds[1:], inds[1:]] + [n_ps, n_ps, 0]
-            simps = np.r_[simps1, simps2]
+            simps = np.c_[np.arange(len(xyz) - 1), np.arange(1, len(xyz))]
+        elif isinstance(xyz, tuple):
+            xyz, simps = xyz
+            xyz = np.asarray(xyz)
+            simps = np.asarray(simps)
         else:
-            if isinstance(xyz, tuple):
-                xyz, simps = xyz
-                xyz = np.asarray(xyz)
-                simps = np.asarray(simps)
-            else:
-                xyz = np.asarray(xyz)
-                triang = Delaunay(xyz[:, :2])
-                simps = triang.simplices
-            n_ps = len(xyz)
-
-        # calculate bounding box for padding
-        bb_min = np.min(xyz, axis=0)[:-1]
-        bb_max = np.max(xyz, axis=0)[:-1]
-        half_width = (bb_max - bb_min) / 2
-        center = (bb_max + bb_min) / 2
-        points = np.empty((n_ps, self.dim))
-        points[:, -1] = xyz[:, -1]
+            xyz = np.asarray(xyz)
+            simps = Delaunay(xyz[:, :2]).simplices
+        n_ps = len(xyz)
+
+        if extrapolate_surface is None:
+            extrapolate_surface = False
+            if has_h_pad:
+                warnings.warn(
+                    "In discretize v1.0 refine_surface will change the default "
+                    "behavior of horizontal padding to extrapolate the surface "
+                    "outwards from the nodes on its boundary, instead of stretching "
+                    "its bounding box. To use the new behavior now, set "
+                    "extrapolate_surface=True. To keep the current behavior, "
+                    "explicitly set extrapolate_surface=False.",
+                    FutureWarning,
+                    stacklevel=2,
+                )
+
+        # Add points outward from the boundary of the triangulation, so that padding
+        # is a scaling of their offsets, without needing to re-triangulate.
+        if has_h_pad and extrapolate_surface:
+            xyz, simps, pad_dirs = extend_surface_boundary((xyz, simps))
+        else:
+            pad_dirs = np.zeros((n_ps, self.dim - 1))
+        stretch = has_h_pad and not extrapolate_surface
+        if stretch:
+            # bounding box of the surface, to stretch for padding
+            bb_min = np.min(xyz, axis=0)[:-1]
+            bb_max = np.max(xyz, axis=0)[:-1]
+            half_width = (bb_max - bb_min) / 2
+            center = (bb_max + bb_min) / 2
+        n_pts = len(xyz)
+
+        if self.dim == 2:
+            simps1 = np.c_[simps, simps[:, 0]] + [0, 0, n_pts]
+            simps2 = np.c_[simps, simps[:, 1]] + [n_pts, n_pts, 0]
+            simps = np.r_[simps1, simps2]
+
+        points = np.empty((n_pts, self.dim))
 
         h_min = np.r_[[h.min() for h in self.h]]
         pad = 0.0
         for lv, n_pad in zip(np.arange(level, 1, -1), padding_cells_by_level):
             pad += n_pad * h_min * 2 ** (self.max_level - lv)
             h = 0
+            if stretch:
+                horizontal_expansion = (half_width + pad[:-1]) / half_width
+                points[:, :-1] = horizontal_expansion * (xyz[:, :-1] - center) + center
+            else:
+                points[:, :-1] = xyz[:, :-1] + pad[:-1] * pad_dirs
+            points[:, -1] = xyz[:, -1]
             if pad_up:
                 h += pad[-1]
             if pad_down:
                 h += pad[-1]
                 points[:, -1] = xyz[:, -1] - pad[-1]
-            horizontal_expansion = (half_width + pad[:-1]) / half_width
-            points[:, :-1] = horizontal_expansion * (xyz[:, :-1] - center) + center
             if self.dim == 2:
                 triangles = np.r_[points, points + [0, h]][simps]
                 self.refine_triangle(

discretize/utils/__init__.py

@@ -92,6 +92,7 @@
   active_from_xyz
   mesh_builder_xyz
   slice_mesh
+  extend_surface_boundary
 
 Utilities for Curvilinear Meshes
 --------------------------------
@@ -139,6 +140,7 @@
     mesh_builder_xyz,
     example_simplex_mesh,
     slice_mesh,
+    extend_surface_boundary,
 )
 from discretize.utils.curvilinear_utils import (
     example_curvilinear_grid,

discretize/utils/mesh_utils.py

@@ -1362,3 +1362,176 @@ def example_simplex_mesh(rect_shape):
 closest_points = deprecate_function(
     closest_points_index, "closest_points", removal_version="1.0.0", error=True
 )
+
+
+def _outward_boundary_directions(xy, simps, max_miter=5.0):
+    """Outward offset directions for every boundary node of a 2D triangulation.
+
+    Parameters
+    ----------
+    xy : (n_points, 2) numpy.ndarray
+        The horizontal locations of the points.
+    simps : (n_simplices, 3) numpy.ndarray
+        The triangles of the triangulation.
+    max_miter : float
+        Maximum allowed length of the offset direction, (for very sharp corners).
+
+    Returns
+    -------
+    nodes : (n_boundary,) numpy.ndarray
+        Indices of the nodes on the boundary.
+    directions : (n_boundary, 2) numpy.ndarray
+        Outward directions such that offsetting a node by ``d * directions`` moves
+        the adjacent boundary edges out by a perpendicular distance ``d`` (a miter).
+    edges : (n_edges, 2) numpy.ndarray
+        The boundary edges (as node indices).
+    """
+    edges = np.r_[simps[:, [0, 1]], simps[:, [1, 2]], simps[:, [2, 0]]]
+    third = np.r_[simps[:, 2], simps[:, 0], simps[:, 1]]
+    _, inverse, counts = np.unique(
+        np.sort(edges, axis=1), axis=0, return_inverse=True, return_counts=True
+    )
+    on_boundary = counts[inverse.reshape(-1)] == 1
+    edges = edges[on_boundary]
+    third = third[on_boundary]
+
+    tangent = xy[edges[:, 1]] - xy[edges[:, 0]]
+    normals = np.c_[tangent[:, 1], -tangent[:, 0]]
+    length = np.linalg.norm(normals, axis=1)
+    length[length == 0] = 1.0
+    normals /= length[:, None]
+    # point away from the triangle the edge belongs to
+    inward = np.sum(normals * (xy[third] - xy[edges[:, 0]]), axis=1) > 0
+    normals[inward] *= -1
+
+    nodes = np.unique(edges)
+    local = np.searchsorted(nodes, edges)
+    n_sum = np.zeros((len(nodes), 2))
+    n_count = np.zeros(len(nodes))
+    for k in range(2):
+        np.add.at(n_sum, local[:, k], normals)
+        np.add.at(n_count, local[:, k], 1)
+
+    bisector = n_sum / n_count[:, None]
+    b_len = np.linalg.norm(bisector, axis=1)
+    b_len[b_len == 0] = 1.0
+    bisector /= b_len[:, None]
+
+    # mean of cos(half angle) between the bisector and adjacent edge normals
+    cos_half = np.zeros(len(nodes))
+    for k in range(2):
+        np.add.at(cos_half, local[:, k], np.sum(bisector[local[:, k]] * normals, 1))
+    cos_half /= n_count
+    scale = 1.0 / np.clip(cos_half, 1.0 / max_miter, 1.0)
+    return nodes, bisector * scale[:, None], edges
+
+
+def extend_surface_boundary(xyz, max_miter=5.0):
+    """Extend a triangulated surface outward from every node on its boundary.
+
+    A new point is added for every node along the boundary of the triangulation
+    (not just the convex hull). Each new point starts at the same location as the
+    node it comes from, and is offset horizontally along the outward bisector of the
+    boundary edges adjacent to that node. The new points are triangulated with the
+    original ones once, so extending the surface by a distance ``d`` (a nearest
+    neighbor extrapolation of the surface) is just a scaling of the returned
+    directions.
+
+    Parameters
+    ----------
+    xyz : (n_points, dim) array_like or tuple
+        Surface points, where ``dim`` is 2 or 3. The last column is the vertical
+        coordinate. A point cloud is triangulated (Delaunay) along the horizontal
+        dimensions (in 2D it is sorted along x and connected as a line). Alternatively
+        supply your own connectivity as a tuple of ``(xyz, simplices)``, where
+        ``simplices`` are line segments (2D) or triangles (3D) indexing ``xyz``.
+    max_miter : float, optional
+        Maximum length of the offset direction, which limits how far the extension
+        of very sharp boundary corners can reach.
+
+    Returns
+    -------
+    xyz_ext : (n_points + n_new, dim) numpy.ndarray
+        The original (sorted, in 2D) points followed by the new points, which have
+        the vertical coordinate of the node they extend.
+    simplices_ext : (n_simplices + n_new_simplices, dim) numpy.ndarray
+        The simplices connecting the original and new points.
+    directions : (n_points + n_new, dim - 1) numpy.ndarray
+        Horizontal offset directions, which are zero for the original points. The
+        extended surface for a distance ``d`` (a scalar, or one value per horizontal
+        dimension) is ``xyz_ext[:, :-1] + d * directions``. The directions are
+        scaled such that boundary edges move out by a perpendicular distance ``d``.
+
+    Raises
+    ------
+    ValueError
+        If ``xyz`` has 2 columns and all of its points have the same x location.
+
+    Notes
+    -----
+    For surfaces with a non-convex boundary, the outward bisectors of different
+    boundary nodes can intersect if the distance is large compared to the size of
+    local features of the boundary (e.g. notches or concavities). The extended
+    surface can then fold over itself and be unreliable.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from discretize.utils import extend_surface_boundary
+    >>> x, y = np.mgrid[0:1:3j, 0:1:3j]
+    >>> xyz = np.stack([x, y, x + y], axis=-1).reshape(-1, 3)
+    >>> xyz_ext, simps, dirs = extend_surface_boundary(xyz)
+    >>> padded = xyz_ext[:, :-1] + 0.1 * dirs
+    >>> padded.shape
+    (17, 2)
+    """
+    if isinstance(xyz, tuple):
+        xyz, simplices = xyz
+        simplices = np.asarray(simplices)
+    else:
+        simplices = None
+    xyz = np.asarray(xyz, dtype=float)
+    dim = xyz.shape[1]
+    if dim not in (2, 3):
+        raise ValueError("xyz must have 2 or 3 columns.")
+    n_ps = len(xyz)
+    if dim == 2:
+        if np.ptp(xyz[:, 0]) == 0:
+            raise ValueError(
+                "The surface has no horizontal extent, so it has no outward direction."
+            )
+        if simplices is None:
+            xyz = xyz[np.argsort(xyz[:, 0])]
+            simplices = np.c_[np.arange(n_ps - 1), np.arange(1, n_ps)]
+        # boundary nodes are those on a single segment
+        ends, counts = np.unique(simplices, return_counts=True)
+        nodes = ends[counts == 1]
+        # neighbor of each end point, to point away from
+        neighbor = np.empty(n_ps, dtype=int)
+        neighbor[simplices[:, 0]] = simplices[:, 1]
+        neighbor[simplices[:, 1]] = simplices[:, 0]
+        dirs = np.sign(xyz[nodes, :1] - xyz[neighbor[nodes], :1])
+        # vertical segment at an end: fall back to which side of the middle it is on
+        tie = dirs == 0
+        dirs[tie] = np.sign(xyz[nodes, :1] - xyz[:, :1].mean())[tie]
+        dirs[dirs == 0] = 1.0
+        new_inds = n_ps + np.arange(len(nodes))
+        simplices = np.r_[simplices, np.c_[nodes, new_inds]]
+    else:
+        if simplices is None:
+            simplices = Delaunay(xyz[:, :2]).simplices
+        nodes, dirs, edges = _outward_boundary_directions(
+            xyz[:, :2], simplices, max_miter
+        )
+        new_of_node = np.full(n_ps, -1)
+        new_of_node[nodes] = n_ps + np.arange(len(nodes))
+        i, j = edges[:, 0], edges[:, 1]
+        simplices = np.r_[
+            simplices,
+            np.c_[i, j, new_of_node[j]],
+            np.c_[i, new_of_node[j], new_of_node[i]],
+        ]
+    xyz_ext = np.r_[xyz, xyz[nodes]]
+    directions = np.zeros((len(xyz_ext), dim - 1))
+    directions[n_ps:] = dirs
+    return xyz_ext, simplices, directions