Function plot_streamlines

python/demo/demo_pyvista.py:312–348

Plot streamlines of a vector field in 3D.

()

def plot_streamlines():
    """Plot streamlines of a vector field in 3D."""
    msh = create_unit_cube(MPI.COMM_WORLD, 4, 4, 4, CellType.hexahedron, dtype=np.float64)
    gdim = msh.geometry.dim
    V = functionspace(msh, ("Discontinuous Lagrange", 2, (gdim,)))
    u = Function(V, dtype=np.float64)
    u.interpolate(lambda x: np.vstack((-(x[1] - 0.5), x[0] - 0.5, np.zeros(x.shape[1]))))

    cells, types, x = plot.vtk_mesh(V)
    num_dofs = x.shape[0]
    values = np.zeros((num_dofs, 3), dtype=np.float64)
    values[:, :gdim] = u.x.array.reshape(num_dofs, V.dofmap.index_map_bs)

    # Create a point cloud of glyphs
    grid = pyvista.UnstructuredGrid(cells, types, x)
    grid["vectors"] = values
    grid.set_active_vectors("vectors")
    glyphs = grid.glyph(orient="vectors", factor=0.1)
    streamlines = grid.streamlines(
        vectors="vectors", return_source=False, source_radius=1, n_points=150
    )

    # Create Create plotter
    plotter = pyvista.Plotter()
    plotter.add_text("Streamlines.", position="upper_edge", font_size=20, color="black")
    plotter.add_mesh(grid, style="wireframe")
    plotter.add_mesh(glyphs)
    plotter.add_mesh(streamlines.tube(radius=0.001))
    plotter.view_xy()
    if pyvista.OFF_SCREEN:
        plotter.screenshot(
            out_folder / f"streamlines_{MPI.COMM_WORLD.rank}.png",
            transparent_background=transparent,
            window_size=[figsize, figsize],
        )
    else:
        plotter.show()


plot_scalar()