Find cells from vertices in pyvista PolyData mesh

I have a PolyData mesh, and I am trying to color some cells that contain given vertices. However, I need an efficient way to to this.

For instance, this works, but it is way too slow for a large mesh:

import pyvista as pv
import numpy as np

mesh = pv.Sphere()
# Scalars are zero everywhere
data = [0] * mesh.n_cells

# I want to color the cells which contain these vertices
point_indexes = np.random.randint(0,mesh.n_points,100)

# Loop over all vertices
for point in point_indexes:
    # This next part is inefficient in a large mesh: 
    # I extract the 3D coordinates of the point and look for 
    # a containing cell. I think I should look for a cell containing 
    # the vertex from its index...but I can't find a way to do it!
    point_idx = mesh.points[point]
    cell_idx = mesh.find_containing_cell(point_idx)
    
    # Change the scalar
    data [cell_idx] = 1

# Color the cells touching the vertices
mesh.cell_data['data'] = data 

mesh.plot()

This is what I get, which is fine…I just need to do it much faster!

Answer

If you have the point indices and want to extract every cell that contains any of these points you can use the extract_points() filter whose default adjacent_cells=True, include_cells=True parameters imply that extracting your points (all at once) will also extract every cell they are a part of.

Then all that’s left is using the cell scalars called 'vtkOriginalCellIds' and use those indices to change your scalars in the original mesh. I’ve cleaned up some style things in the code:

import numpy as np
import pyvista as pv

mesh = pv.Sphere()
# start from zero scalars
data = np.zeros(mesh.n_cells)

# point indices to use
rng = np.random.default_rng()
point_indices = rng.integers(0, mesh.n_points, 100)

# extract cells along with the points, assign 1 scalar to cells
extracted = mesh.extract_points(point_indices)
data[extracted.cell_data['vtkOriginalCellIds']] = 1

# assign to cell scalas and plot
mesh.cell_data['data'] = data
plotter = pv.Plotter()
plotter.add_mesh(mesh)
plotter.add_points(mesh.points[point_indices, :], color='forestgreen',
                   point_size=10, render_points_as_spheres=True)
plotter.show()

As you can see, with this approach you indeed get every cell that contains special points, leading to islands of cells around each point. This is what I’d expect based on your description. I suspect the reason why your original output doesn’t reflect this is because finding points by 3d point coordinates is prone to floating-point errors, which makes VTK miss most of the surrounding cells.