liblaf.melon.tri ¶

def contains(
    mesh: wp.Mesh, points: Float[Tensor, "*p 3"], max_dist: float = wp.inf
) -> Bool[Tensor, "*p"]:
    shape: Sequence[int] = points.shape[:-1]
    points: Float[Tensor, "p 3"] = points.reshape(-1, 3)
    n_points: int = points.shape[0]
    points_wp: wp.array = wp.from_torch(
        points.contiguous(), wp.vec3f, return_ctype=True
    )
    contains_wp: wp.array = wp.empty((n_points,), wp.bool)
    wp.launch(
        _mesh_query_point_kernel,
        dim=(n_points,),
        inputs=[mesh.id, points_wp, max_dist],
        outputs=[contains_wp],
        stream=warp_stream_from_torch(),
    )
    contains: Bool[Tensor, " p"] = wp.to_torch(contains_wp)
    contains: Bool[Tensor, "*p"] = contains.reshape(shape)
    return contains

def edge_length(mesh: pv.PolyData) -> Float[np.ndarray, " E"]:
    """Compute lengths for all extracted edges of a triangular surface.

    Args:
        mesh: Surface mesh.

    Returns:
        Edge lengths reported by PyVista's cell-size filter.
    """
    edges: pv.PolyData = mesh.extract_all_edges()
    edges: pv.PolyData = edges.compute_cell_sizes(length=True)
    return edges.cell_data["Length"]

def extract_cells(
    mesh: Any, ind: int | VectorLike[int] | VectorLike[bool], *, invert: bool = False
) -> pv.PolyData:
    """Extract selected cells from a surface-like mesh.

    Args:
        mesh: Object convertible to [`pyvista.PolyData`][pyvista.PolyData].
        ind: Cell index, integer indices, or boolean cell mask.
        invert: Extract all cells except the selected cells.

    Returns:
        Extracted surface.
    """
    mesh: pv.PolyData = io.as_polydata(mesh)
    ind: np.ndarray = np.asarray(ind)
    if np.isdtype(ind.dtype, "bool"):
        ind: np.ndarray = np.flatnonzero(ind)
    cells: pv.UnstructuredGrid = mesh.extract_cells(ind, invert=invert)
    surface: pv.PolyData = cells.extract_surface(algorithm=None)
    return surface

def extract_groups(
    mesh: Any, groups: int | str | Iterable[int | str], *, invert: bool = False
) -> pv.PolyData:
    """Extract cells whose `GroupId` matches numeric or named groups.

    Args:
        mesh: Mesh with `GroupId` cell data and optional `GroupName` field data.
        groups: Group id, group name, or iterable of ids and names.
        invert: Extract all cells outside the selected groups.

    Returns:
        Extracted surface.
    """
    return extract_cells(mesh, select_groups(mesh, groups), invert=invert)

def fill_point(
    mesh: pv.PolyData,
    mask: Bool[np.ndarray, " P"],
    names: Iterable[str] | None = None,
    *,
    limit: float = math.inf,
) -> pv.PolyData:
    if names is None:
        names: list[str] = mesh.point_data.keys()
    csgraph: scipy.sparse.coo_array = _make_csgraph(mesh)
    _dist_matrix, _predecessors, sources = scipy.sparse.csgraph.dijkstra(
        csgraph,
        directed=False,
        indices=np.flatnonzero(~mask),
        return_predecessors=True,
        limit=limit * mesh.length,
        min_only=True,
    )
    mask &= sources >= 0
    for name in names:
        data: Float[np.ndarray, "S ..."] = mesh.point_data[name]
        data[mask] = data[sources[mask]]
        mesh.point_data[name] = data
    return mesh

def fix_normals(mesh: Any, *, multibody: bool | None = None) -> tm.Trimesh:
    """Orient triangle normals with Trimesh and return a Trimesh object.

    Args:
        mesh: Object convertible to [`trimesh.Trimesh`][trimesh.Trimesh].
        multibody: Forwarded to [`trimesh.Trimesh.fix_normals`][].

    Returns:
        Mesh with repaired normal orientation.
    """
    mesh: tm.Trimesh = io.as_trimesh(mesh)
    mesh.fix_normals(multibody=multibody)
    return mesh

def geodesic_path(mesh: pv.PolyData, v_start: int, v_end: int) -> pv.PolyData:
    """Find an edge-flip geodesic path between two surface vertices.

    Args:
        mesh: Triangular surface mesh.
        v_start: Start vertex index.
        v_end: End vertex index.

    Returns:
        Polyline following the computed geodesic path.
    """
    solver = pp3d.EdgeFlipGeodesicSolver(mesh.points, mesh.regular_faces)
    points: Float[np.ndarray, "p 3"] = solver.find_geodesic_path(v_start, v_end)
    return pv.lines_from_points(points)

def implicit_distance(
    mesh: wp.Mesh, points: Float[Tensor, "*q 3"], max_dist: float = wp.inf
) -> Float[Tensor, "*q"]:
    shape: Sequence[int] = points.shape[:-1]
    points: Float[Tensor, "q 3"] = points.reshape(-1, 3)
    n_points: int = points.shape[0]
    points_wp: wp.array = wp.from_torch(points, wp.vec3f, return_ctype=True)
    distance_wp: wp.array = wp.empty((n_points,), wp.float32)
    wp.launch(
        _mesh_query_point_kernel,
        (n_points,),
        inputs=[mesh.id, points_wp, max_dist],
        outputs=[distance_wp],
        stream=warp_stream_from_torch(),
    )
    distance: Float[Tensor, " q"] = wp.to_torch(distance_wp)
    distance: Float[Tensor, "*q"] = distance.reshape(shape)
    return distance

def query_ray(
    mesh: wp.Mesh,
    start: Float[Tensor, "*q 3"],
    direction: Float[Tensor, "*q 3"],
    max_t: float = wp.inf,
) -> Float[Tensor, "*q"]:
    shape: Sequence[int] = start.shape[:-1]
    start: Float[Tensor, "q 3"] = start.reshape(-1, 3).contiguous()
    direction: Float[Tensor, "q 3"] = direction.reshape(-1, 3).contiguous()
    start_wp: wp.array = wp.from_torch(start, wp.vec3f, return_ctype=True)
    direction_wp: wp.array = wp.from_torch(direction, wp.vec3f, return_ctype=True)
    distance: Float[Tensor, " q"] = torch.empty((start.shape[0],), dtype=torch.float32)
    distance_wp: wp.array = wp.from_torch(distance, wp.float32, return_ctype=True)
    wp.launch(
        _mesh_query_ray_kernel,
        dim=(start.shape[0],),
        inputs=[mesh.id, start_wp, direction_wp, max_t],
        outputs=[distance_wp],
        stream=warp_stream_from_torch(),
    )
    distance: Float[Tensor, "*q"] = distance.reshape(shape)
    return distance

def select_groups(
    mesh: Any,
    groups: int | str | Iterable[int | str],
    *,
    invert: bool = False,
    preference: pv.FieldAssociation = pv.FieldAssociation.CELL,
) -> Bool[np.ndarray, " cells"]:
    """Build a boolean mask for cells in selected groups.

    Args:
        mesh: Mesh with a group-id array.
        groups: Group id, group name, or iterable of ids and names.
        invert: Invert the selection mask.
        preference: PyVista association used to resolve the `GroupId` array.

    Returns:
        Boolean mask over mesh cells.
    """
    mesh: pv.PolyData = io.as_polydata(mesh)
    group_ids: list[int] = _as_group_ids(mesh, groups)
    mask: Bool[np.ndarray, " cells"] = np.isin(
        mesh.get_array("GroupId", preference),  # ty:ignore[invalid-argument-type]
        group_ids,
        invert=invert,
    )
    return mask