liblaf.melon.tet ¶

def fix_winding(mesh: Any, *, check: bool = True) -> pv.UnstructuredGrid:
    """Flip negatively oriented tetrahedra to positive volume.

    Args:
        mesh: Object convertible to [`pyvista.UnstructuredGrid`][pyvista.UnstructuredGrid].
        check: Assert that all tetrahedra have non-negative volume after repair.

    Returns:
        Repaired tetrahedral mesh.
    """
    mesh: pv.UnstructuredGrid = io.as_unstructured_grid(mesh)
    mesh: pv.UnstructuredGrid = mesh.compute_cell_sizes(
        length=False, area=False, volume=True
    )
    flip_mask: Bool[np.ndarray, " C"] = mesh.cell_data["Volume"] < 0
    if np.any(flip_mask):
        mesh: pv.UnstructuredGrid = flip(mesh, flip_mask)
        mesh: pv.UnstructuredGrid = mesh.compute_cell_sizes(
            length=False, area=False, volume=True
        )
        if check:
            assert np.all(mesh.cell_data["Volume"] >= 0)
    return mesh

def flip(mesh: Any, mask: Bool[ArrayLike, " C"]) -> pv.UnstructuredGrid:
    """Flip tetrahedral winding for selected cells.

    Args:
        mesh: Object convertible to [`pyvista.UnstructuredGrid`][pyvista.UnstructuredGrid].
        mask: Boolean mask selecting tetrahedra to flip.

    Returns:
        Mesh with selected tetrahedra rewound.
    """
    mesh: pv.UnstructuredGrid = io.as_unstructured_grid(mesh)
    mask: Bool[np.ndarray, " C"] = np.asarray(mask)
    if not np.any(mask):
        return mesh
    tetras: Integer[np.ndarray, "C 4"] = mesh.cells_dict[pv.CellType.TETRA]  # ty:ignore[invalid-argument-type]
    # ref: <https://felupe.readthedocs.io/en/latest/felupe/mesh.html#felupe.mesh.flip>
    faces: Integer[np.ndarray, " 3"] = np.asarray([0, 1, 2], np.int32)
    tetras[np.ix_(mask, faces)] = tetras[np.ix_(mask, faces[::-1])]
    result: pv.UnstructuredGrid = pv.UnstructuredGrid(
        {pv.CellType.TETRA: tetras}, mesh.points
    )
    result.copy_attributes(mesh)
    return result

def tetra_center_radius(mesh: pv.UnstructuredGrid) -> TetraCenterRadiusResult:
    points: Float[Tensor, "p 3"] = torch.tensor(mesh.points, dtype=torch.float32)
    centers: pv.PolyData = mesh.cell_centers()
    centers: Float[Tensor, "c 3"] = torch.tensor(centers.points, dtype=torch.float32)
    cells: Integer[np.ndarray, "c 4"] = mesh.cells_dict[pv.CellType.TETRA]  # ty:ignore[invalid-argument-type]
    cells: Integer[Tensor, "c 4"] = torch.tensor(cells, dtype=torch.int32)
    radius: Float[Tensor, " c"] = torch.amax(
        (centers[:, torch.newaxis, :] - points[cells]).norm(dim=-1), dim=-1
    )
    return TetraCenterRadiusResult(center=centers, radius=radius)

def tetra_surface_broad_phase(
    centers: pv.PolyData, surface: wp.Mesh
) -> TetraSurfaceBroadPhaseResult:
    radius: Float[np.ndarray, " c"] = np.array(centers.point_data["Radius"])
    points_wp: wp.array = wp.from_numpy(centers.points, wp.vec3f)
    distance_wp: wp.array = wp.zeros((centers.n_points,), wp.float32)
    wp.launch(
        _mesh_query_point_kernel,
        dim=(centers.n_points,),
        inputs=[surface.id, points_wp, wp.inf],
        outputs=[distance_wp],
    )
    distance: Float[np.ndarray, " c"] = distance_wp.numpy()
    return TetraSurfaceBroadPhaseResult(distance=distance, radius=radius)

def tetra_surface_fraction(
    mesh: pv.UnstructuredGrid, surface: wp.Mesh, *, n_samples: int = 1024
) -> Bool[Tensor, "c s"]:
    points: Float[Tensor, "p 3"] = torch.tensor(mesh.points, dtype=torch.float32)
    cells: Integer[Tensor, "c 4"] = torch.tensor(
        mesh.cells_dict[pv.CellType.TETRA],  # ty:ignore[invalid-argument-type]
        dtype=torch.int32,
    )
    barycentric: Float[np.ndarray, "s 4"] = bary.sample(n_samples, 4)
    barycentric: Float[Tensor, "s 4"] = torch.tensor(barycentric, dtype=torch.float32)
    contains: Bool[Tensor, " c s"] = torch.full((mesh.n_cells, n_samples), torch.nan)
    split_size: int = max(1, 100_000_000 // n_samples)
    for chunk in torch.split(torch.arange(mesh.n_cells), split_size):
        chunk_cells: Integer[Tensor, "k 4"] = cells[chunk]
        samples: Float[Tensor, "k s 3"] = einops.einsum(
            points[chunk_cells], barycentric, "k b i, s b -> k s i"
        )
        samples: Float[Tensor, "k*s 3"] = samples.reshape(-1, 3)
        chunk_contains: Bool[Tensor, " k*s"] = _mesh_query_point(surface, samples)
        chunk_contains: Bool[Tensor, "k s"] = chunk_contains.reshape(-1, n_samples)
        contains[chunk] = chunk_contains
    return contains

def volume_fraction(
    mesh: pv.UnstructuredGrid,
    surface: pv.PolyData,
    *,
    n_samples: int = 1024,
    split_size: int = 100_000_000,
) -> Float[np.ndarray, " c"]:
    """Estimate how much of each tetrahedron lies inside a closed surface.

    Tetrahedra whose four vertices are all inside or all outside the surface are
    classified exactly. Boundary tetrahedra are sampled with Sobol
    barycentric coordinates and classified by ray intersections.

    Args:
        mesh: Tetrahedral volume mesh.
        surface: Closed triangular surface used as the inside/outside boundary.
        n_samples: Requested number of samples for each boundary tetrahedron.
            The Sobol sequence rounds this up to a power of two.
        split_size: Approximate maximum number of sampled points to process in
            one chunk.

    Returns:
        One fraction per tetrahedral cell.
    """
    with _torch_device():
        surface_wp: wp.Mesh = io.as_warp_mesh(surface)
        cells: Integer[np.ndarray, " c 4"] = mesh.cells_dict[pv.CellType.TETRA]  # ty:ignore[invalid-argument-type]
        cells: Integer[Tensor, "c 4"] = torch.tensor(cells)
        points: Float[Tensor, "p 3"] = torch.tensor(mesh.points, dtype=torch.float32)
        contains: Bool[Tensor, " p"] = _contains_points(surface_wp, points)

        fraction: Float[Tensor, " c"] = torch.zeros((mesh.n_cells,))
        inside: Bool[Tensor, " c"] = torch.all(contains[cells], dim=-1)
        outside: Bool[Tensor, " c"] = ~torch.any(contains[cells], dim=-1)
        fraction[inside] = 1.0
        fraction[outside] = 0.0

        remainder: Integer[Tensor, " r"] = torch.nonzero(~(inside | outside)).flatten()
        barycentric: Float[Tensor, "s 4"] = _sample_barycentric_coordinates(n_samples)
        n_samples: int = barycentric.shape[0]
        chunk_size: int = max(1, split_size // n_samples)
        for chunk in torch.split(remainder, chunk_size):
            samples: Float[Tensor, "k s 3"] = _barycentric_to_points(
                points, cells[chunk], barycentric
            )
            samples: Float[Tensor, "k*s 3"] = samples.reshape(-1, 3)
            contains: Bool[Tensor, " k*s"] = _contains_points(surface_wp, samples)
            contains: Bool[Tensor, "k s"] = contains.reshape(-1, n_samples)
            fraction[chunk] = torch.mean(contains.float(), dim=-1)

        return fraction.numpy(force=True)