liblaf.melon.xfer ¶

def tri_cell_to_tri_point(
    source: pv.PolyData,
    target: pv.PolyData,
    names: Iterable[str],
    *,
    tolerance: float | None = None,
    snap_to_closest_point: bool = True,
) -> pv.PolyData:
    """Sample triangle cell arrays onto target surface points.

    Args:
        source: Surface carrying the source cell-data arrays.
        target: Surface whose points receive sampled point-data arrays.
        names: Cell-data array names to transfer.
        tolerance: Optional PyVista sampling tolerance.
        snap_to_closest_point: Forwarded to [`pyvista.DataSetFilters.sample`][].

    Returns:
        The mutated `target` mesh.
    """
    src: pv.PolyData = pv.PolyData()
    src.copy_structure(source)
    src.cell_data.update({name: source.cell_data[name] for name in names})
    result: pv.PolyData = target.sample(
        src, tolerance=tolerance, snap_to_closest_point=snap_to_closest_point
    )
    result.point_data.pop("vtkGhostType", None)
    target.point_data.update(result.point_data)
    return target

def tri_point_to_tet_point(
    source: pv.PolyData, target: pv.UnstructuredGrid, fill_values: Mapping[str, Any]
) -> pv.UnstructuredGrid:
    """Transfer point arrays from a surface to nearest tetrahedral points.

    Args:
        source: Surface with source point-data arrays.
        target: Tetrahedral mesh whose point data should be filled.
        fill_values: Mapping from point-data array name to default value for
            target points that do not receive a nearest source value.

    Returns:
        The mutated `target` mesh.
    """
    kdtree: scipy.spatial.KDTree = scipy.spatial.KDTree(target.points)
    _d, indices = kdtree.query(source.points)
    for name, fill_value in fill_values.items():
        source_data: np.ndarray = source.point_data[name]
        target_data: np.ndarray = np.full(
            (target.n_points, *source_data.shape[1:]), fill_value, source_data.dtype
        )
        target_data[indices] = source_data
        target.point_data[name] = target_data
    return target

def tri_point_to_tri_point(
    source: pv.PolyData,
    target: pv.PolyData,
    names: Iterable[str] | None = None,
    *,
    max_dist: float = 0.01,
    normal_threshold: float = 0.8,
) -> pv.PolyData:
    if names is None:
        names: list[str] = source.point_data.keys()
    source_wp: wp.Mesh = io.as_warp_mesh(source)
    points: wp.array = wp.from_numpy(target.points, wp.vec3f)
    normals: wp.array = wp.from_numpy(target.point_normals, wp.vec3f)
    result_wp: wp.array = wp.zeros((target.n_points), wp.bool)
    face_wp: wp.array = wp.zeros((target.n_points), wp.int32)
    barycentric_wp: wp.array = wp.zeros((target.n_points,), wp.vec3f)
    wp.launch(
        _mesh_query_point_kernel,
        dim=(target.n_points,),
        inputs=[
            source_wp.id,
            points,
            normals,
            max_dist * source.length,
            normal_threshold,
        ],
        outputs=[result_wp, barycentric_wp, face_wp],
    )
    result: Bool[np.ndarray, " T"] = result_wp.numpy()
    face: Integer[np.ndarray, " T"] = face_wp.numpy()
    barycentric: Float[np.ndarray, "T 3"] = barycentric_wp.numpy()
    indices: Integer[np.ndarray, "T 3"] = source.regular_faces[face]
    for name in names:
        source_data: Float[np.ndarray, "S ..."] = source.point_data[name]
        target_data: Float[np.ndarray, "T ..."] = einops.einsum(
            barycentric, source_data[indices], "t i, t i ... -> t ..."
        )
        mask: Bool[np.ndarray, "T ..."] = np.expand_dims(
            result, axis=tuple(range(1, target_data.ndim))
        )
        target.point_data[name] = np.where(mask, target_data, np.nan)
    return target