Source code for dolfinx.plot

# Copyright (C) 2021 Jørgen S. Dokken and Garth N. Wells
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Support functions for plotting"""

import functools
import warnings
import numpy as np

from dolfinx import cpp, fem

# Permutation for DOLFINx DG layout to VTK
# Note that third order tetrahedrons has a special ordering:
# https://gitlab.kitware.com/vtk/vtk/-/issues/17746
_perm_dg = {cpp.mesh.CellType.triangle: {1: [0, 1, 2], 2: [0, 2, 5, 1, 4, 3], 3: [0, 3, 9, 1, 2, 6, 8, 7, 4, 5],
                                         4: [0, 4, 14, 1, 2, 3, 8, 11, 13, 12, 9, 5, 6, 7, 10]},
            cpp.mesh.CellType.tetrahedron: {1: [0, 1, 2, 3], 2: [0, 2, 5, 9, 1, 4, 5, 6, 7, 8],
                                            3: [0, 3, 9, 19, 1, 2, 6, 8, 7, 4, 10, 16, 12, 17, 15, 18, 11, 14, 13, 5]}}
_perm_dq = {cpp.mesh.CellType.quadrilateral: {1: [0, 1, 3, 2], 2: [0, 2, 8, 6, 1, 5, 7, 3, 4],
                                              3: [0, 3, 15, 12, 1, 2, 7, 11, 13, 14, 4, 8, 5, 6, 9, 10]},
            cpp.mesh.CellType.hexahedron: {1: [0, 1, 3, 2, 4, 5, 7, 6],
                                           2: [0, 2, 8, 6, 18, 20, 26, 24, 1, 5, 7, 3, 19,
                                               23, 25, 21, 9, 11, 17, 15, 12, 14, 10, 16, 4, 22, 14]}}

# NOTE: Edge visualization of higher order elements are sketchy, see:
# https://github.com/pyvista/pyvista/issues/947


# NOTE: These dictionaries and following function should be replaced by
# cpp.io.get_vtk_cell_type when plotting module has better support for
# arbitrary lagrangian elements
_first_order_vtk = {cpp.mesh.CellType.interval: 3,
                    cpp.mesh.CellType.triangle: 5,
                    cpp.mesh.CellType.quadrilateral: 9,
                    cpp.mesh.CellType.tetrahedron: 10,
                    cpp.mesh.CellType.hexahedron: 12}
_cell_degree_triangle = {3: 1, 6: 2, 10: 3, 15: 4, 21: 5, 28: 6, 36: 7, 45: 8, 55: 9}
_cell_degree_tetrahedron = {4: 1, 10: 2, 20: 3}
_cell_degree_hexahedron = {8: 1, 27: 2}


def _element_degree(cell_type: cpp.mesh.CellType, num_nodes: int):
    """Determine the degree of a cell by the number of nodes"""
    if cell_type == cpp.mesh.CellType.triangle:
        return _cell_degree_triangle[num_nodes]
    elif cell_type == cpp.mesh.CellType.point:
        return 1
    elif cell_type == cpp.mesh.CellType.interval:
        return num_nodes - 1
    elif cell_type == cpp.mesh.CellType.tetrahedron:
        return _cell_degree_tetrahedron[num_nodes]
    elif cell_type == cpp.mesh.CellType.quadrilateral:
        return int(np.sqrt(num_nodes) - 1)
    elif cell_type == cpp.mesh.CellType.hexahedron:
        return _cell_degree_hexahedron[num_nodes]


[docs]@functools.singledispatch def create_vtk_topology(mesh: cpp.mesh.Mesh, dim: int, entities=None): """Create vtk mesh topology data for mesh entities of a given dimension. The vertex indices in the returned topology array are the indices for the associated entry in the mesh geometry. """ if entities is None: num_cells = mesh.topology.index_map(dim).size_local entities = np.arange(num_cells, dtype=np.int32) else: num_cells = len(entities) # Get the indices in the geometry array that correspong to the # topology vertices geometry_entities = cpp.mesh.entities_to_geometry(mesh, dim, entities, False) # Array holding the cell type (shape) for each cell if mesh.topology.cell_type == cpp.mesh.CellType.prism: raise RuntimeError("Plotting of prism meshes not supported") e_type = cpp.mesh.cell_entity_type(mesh.topology.cell_type, dim, 0) degree = _element_degree(e_type, geometry_entities.shape[1]) if degree == 1: cell_types = np.full(num_cells, _first_order_vtk[e_type]) else: warnings.warn("Plotting of higher order mesh topologies is experimental.") cell_types = np.full(num_cells, cpp.io.get_vtk_cell_type(mesh, dim)) # Get cell data and the DOLFINx -> VTK permutation array num_vertices_per_cell = geometry_entities.shape[1] map_vtk = np.argsort(cpp.io.perm_vtk(e_type, num_vertices_per_cell)) # Create mesh topology topology = np.zeros((num_cells, num_vertices_per_cell + 1), dtype=np.int32) topology[:, 0] = num_vertices_per_cell topology[:, 1:] = geometry_entities[:, map_vtk] return topology.reshape(1, -1)[0], cell_types
@create_vtk_topology.register(fem.FunctionSpace) def _(V: fem.FunctionSpace, entities=None): """Creates a vtk mesh topology (topology array and array of cell types) that is based on degree of freedom coordinate. Note that this function supports Lagrange elements (continuous and discontinuous) only. """ family = V.ufl_element().family() if not (family in ['Discontinuous Lagrange', "Lagrange", "DQ", "Q"]): raise RuntimeError("Can only create meshes from Continuous or Discontinuous function-spaces") if V.ufl_element().degree() == 0: raise RuntimeError("Cannot create topology from cellwise constants.") mesh = V.mesh if entities is None: num_cells = mesh.topology.index_map(mesh.topology.dim).size_local entities = np.arange(num_cells, dtype=np.int32) else: num_cells = entities.size dofmap = V.dofmap num_dofs_per_cell = V.dofmap.dof_layout.num_dofs degree = V.ufl_element().degree() cell_type = mesh.topology.cell_type if family == "Discontinuous Lagrange": perm = np.array(_perm_dg[cell_type][degree], dtype=np.int32) elif family == "DQ": perm = np.array(_perm_dq[cell_type][degree], dtype=np.int32) else: perm = np.argsort(cpp.io.perm_vtk(cell_type, num_dofs_per_cell)) if degree == 1: cell_types = np.full(num_cells, _first_order_vtk[mesh.topology.cell_type]) else: warnings.warn("Plotting of higher order functions is experimental.") cell_types = np.full(num_cells, cpp.io.get_vtk_cell_type(mesh, mesh.topology.dim)) topology = np.zeros((num_cells, num_dofs_per_cell + 1), dtype=np.int32) topology[:, 0] = num_dofs_per_cell dofmap_ = dofmap.list.array.reshape(dofmap.list.num_nodes, num_dofs_per_cell) topology[:, 1:] = dofmap_[:num_cells, perm] return topology.reshape(1, -1)[0], cell_types