from __future__ import annotations
from . import _scifem # type: ignore
import collections
import dolfinx
import typing
import numpy as np
import numpy.typing as npt
from typing import Protocol
from packaging.version import Version
import basix
import ufl
__all__ = [
"create_entity_markers",
"transfer_meshtags_to_submesh",
"SubmeshData",
"reverse_mark_entities",
"extract_submesh",
"find_interface",
"compute_interface_data",
"compute_subdomain_exterior_facets",
"create_geometry_function_space",
"move",
"copy",
]
if typing.TYPE_CHECKING:
TaggedEntities = (
tuple[int, typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.bool_]]]
| tuple[int, typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.bool_]], bool]
)
class _EntityMap(Protocol):
"""Protocol for EntityMap-like objects."""
sub_topology: dolfinx.mesh.Topology
dim: int
def sub_topology_to_topology(
self, entities: npt.NDArray[np.int32], inverse: bool
) -> npt.NDArray[np.int32]:
"""Map entities between sub-topology and topology."""
...
def get_entity_map(entity_map: _EntityMap | npt.NDArray[np.int32]) -> npt.NDArray[np.int32]:
"""Get an entity map from the sub-topology to the topology.
This function handles both the deprecated construction of an entity map as a numpy array
and the newer `EntityMap` class from `dolfinx.mesh`.
Args:
entity_map: An `EntityMap` object or a numpy array representing the mapping.
Returns:
Mapped indices of entities.
"""
try:
sub_top = entity_map.sub_topology
assert isinstance(sub_top, dolfinx.mesh.Topology)
sub_map = sub_top.index_map(entity_map.dim)
indices = np.arange(sub_map.size_local + sub_map.num_ghosts, dtype=np.int32)
return entity_map.sub_topology_to_topology(indices, inverse=False)
except AttributeError:
return entity_map
[docs]
def create_entity_markers(
domain: dolfinx.mesh.Mesh,
dim: int,
entities_list: list[TaggedEntities],
) -> dolfinx.mesh.MeshTags:
"""Mark entities of specified dimension according to a geometrical marker function.
Args:
domain: A ``dolfinx.mesh.Mesh`` object
dim: Dimension of the entities to mark
entities_list: A list of tuples with the following elements:
- ``index 0``: The tag to assign to the entities
- ``index 1``: A function that takes a point and returns a boolean array
indicating whether the point is inside the entity
- ``index 2``: Optional, if True, the entities will be marked on the boundary
Returns:
A ``dolfinx.mesh.MeshTags`` object with the corresponding entities marked.
If an entity satisfies multiple input marker functions,
it is not deterministic what value the entity gets.
Note:
"""
# Create required connectivities
domain.topology.create_entities(dim)
domain.topology.create_connectivity(dim, 0)
domain.topology.create_connectivity(dim, domain.topology.dim)
# Create marker function
e_map = domain.topology.index_map(dim)
num_entities_local = e_map.size_local + e_map.num_ghosts
markers = np.full(num_entities_local, -1, dtype=np.int32)
locate_entities = lambda on_boundary: (
dolfinx.mesh.locate_entities_boundary if on_boundary else dolfinx.mesh.locate_entities
)
# Concatenate and sort the arrays based on indices
for tagged_entity in entities_list:
on_boundary = False if len(tagged_entity) == 2 else tagged_entity[2]
entities = locate_entities(on_boundary)(domain, dim, tagged_entity[1])
markers[entities] = tagged_entity[0]
facets = np.flatnonzero(markers != -1).astype(np.int32)
return dolfinx.mesh.meshtags(domain, dim, facets, markers[facets])
[docs]
def reverse_mark_entities(
entity_map: dolfinx.common.IndexMap, entities: npt.NDArray[np.int32]
) -> npt.NDArray[np.int32]:
"""Communicate entities marked on a single process to all processes that ghosts or
owns this entity.
Args:
entity_map: Index-map describing entity ownership
entities: Local indices of entities to communicate
Returns:
Local indices marked on any process sharing this entity
"""
comm_vec = dolfinx.la.vector(entity_map, dtype=np.int32)
comm_vec.array[:] = 0
comm_vec.array[entities] = 1
comm_vec.scatter_reverse(dolfinx.la.InsertMode.add)
comm_vec.scatter_forward()
return np.flatnonzero(comm_vec.array).astype(np.int32)
SubmeshData = collections.namedtuple(
"SubmeshData", ["domain", "cell_map", "vertex_map", "node_map", "cell_tag"]
)
[docs]
def find_interface(
cell_tags: dolfinx.mesh.MeshTags,
id_0: tuple[int, ...],
id_1: tuple[int, ...],
) -> npt.NDArray[np.int32]:
"""Given to sets of cells, find the facets that are shared between them.
Args:
cell_tags: MeshTags object marking cells.
id_0: Tags to extract for domain 0
id_1: Tags to extract for domain 1
Returns:
The facets shared between the two domains.
"""
topology = dolfinx.mesh.Topology(cell_tags.topology)
assert topology.dim == cell_tags.dim
tdim = topology.dim
cell_map = topology.index_map(tdim)
# Find all cells on process that has cell with tag(s) id_0
domain_0 = reverse_mark_entities(
cell_map,
cell_tags.indices[
np.isin(cell_tags.values, np.asarray(id_0, dtype=cell_tags.values.dtype))
],
)
# Find all cells on process that has cell with tag(s) id_1
domain_1 = reverse_mark_entities(
cell_map,
cell_tags.indices[
np.isin(cell_tags.values, np.asarray(id_1, dtype=cell_tags.values.dtype))
],
)
# Find all facets connected to each domain
topology.create_connectivity(tdim, tdim - 1)
facet_map = topology.index_map(tdim - 1)
local_facets0 = dolfinx.mesh.compute_incident_entities(topology, domain_0, tdim, tdim - 1)
facets0 = reverse_mark_entities(facet_map, local_facets0)
local_facets1 = dolfinx.mesh.compute_incident_entities(topology, domain_1, tdim, tdim - 1)
facets1 = reverse_mark_entities(facet_map, local_facets1)
# Compute intersecting facets
interface_facets = np.intersect1d(facets0, facets1)
reverse_mark_entities(facet_map, interface_facets)
topology.create_connectivity(tdim - 1, tdim)
f_to_c = topology.connectivity(tdim - 1, tdim)
num_cells_per_facet = f_to_c.offsets[interface_facets + 1] - f_to_c.offsets[interface_facets]
is_interface = interface_facets[num_cells_per_facet == 2]
return is_interface
[docs]
def compute_subdomain_exterior_facets(
mesh: dolfinx.mesh.Mesh, ct: dolfinx.mesh.MeshTags, markers: typing.Sequence[int]
) -> npt.NDArray[np.int32]:
"""Find the the facets that are considered to be on the "exterior" boundary of a subdomain.
The subdomain is defined as the collection of cells in ``ct`` that is marked with any of the
``markers``. The exterior boundary of the subdomain is defined as the collection of facets
that are only connected to a single cell within the subdomain.
Note:
Ghosted facets are included in the resulting array.
Args:
mesh: Mesh to extract subdomains from
ct: MeshTags object marking subdomains
markers: The tags making up the "new" mesh
Returns:
The exterior facets
"""
# Create submesh to find the exterior facet of subdomain
sub_mesh, cell_map, _, _, _ = extract_submesh(
mesh,
ct,
markers,
)
sub_mesh.topology.create_connectivity(sub_mesh.topology.dim - 1, sub_mesh.topology.dim)
sub_facets = dolfinx.mesh.exterior_facet_indices(sub_mesh.topology)
# Map exterior facet to (submesh_cell, local_facet_index) tuples
try:
integration_entities = dolfinx.fem.compute_integration_domains(
dolfinx.fem.IntegralType.exterior_facet, sub_mesh.topology, sub_facets
)
except TypeError:
integration_entities = dolfinx.fem.compute_integration_domains(
dolfinx.fem.IntegralType.exterior_facet,
sub_mesh.topology,
sub_facets,
sub_mesh.topology.dim - 1,
)
integration_entities = integration_entities.reshape(-1, 2)
submap_array = get_entity_map(cell_map)
integration_entities[:, 0] = submap_array[integration_entities[:, 0]]
# Get cell to facet connectivity (parent mesh)
mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim - 1)
num_facets_per_cell = dolfinx.cpp.mesh.cell_num_entities(
mesh.topology.cell_type, mesh.topology.dim - 1
)
c_to_f = mesh.topology.connectivity(mesh.topology.dim, mesh.topology.dim - 1).array.reshape(
-1, num_facets_per_cell
)
# Map (parent_cell, local_facet_index) to facet index (local to process)
parent_facets = c_to_f[integration_entities[:, 0], integration_entities[:, 1]]
facet_map = mesh.topology.index_map(mesh.topology.dim - 1)
# Accumulate ghost facets
return reverse_mark_entities(facet_map, parent_facets)
[docs]
def compute_interface_data(
cell_tags: dolfinx.mesh.MeshTags,
facet_indices: npt.NDArray[np.int32],
include_ghosts: bool = False,
) -> npt.NDArray[np.int32]:
"""
Compute interior facet integrals that are consistently ordered according to the `cell_tags`,
such that the data `(cell0, facet_idx0, cell1, facet_idx1)` is ordered such that
`cell_tags[cell0]`<`cell_tags[cell1]`, i.e the cell with the lowest cell marker
is considered the "+" restriction".
Args:
cell_tags: MeshTags that must contain an integer marker for all cells adjacent
to the `facet_indices`
facet_indices: List of facets (local index) that are on the interface.
include_ghosts: If `True` integration entities will include facets that are ghosts on
the current process. This is for instance useful for interpolation on interior facets.
Returns:
The integration data.
"""
# Future compatibilty check
integration_args: tuple[int] | tuple
if Version("0.10.0") <= Version(dolfinx.__version__):
integration_args = ()
else:
fdim = cell_tags.dim - 1
integration_args = (fdim,)
if include_ghosts:
if len(facet_indices) == 0:
return np.empty((0, 4), dtype=np.int32)
f_to_c = cell_tags.topology.connectivity(cell_tags.topology.dim - 1, cell_tags.topology.dim)
c_to_f = cell_tags.topology.connectivity(cell_tags.topology.dim, cell_tags.topology.dim - 1)
# Extract the cells connected to each facet.
# Assumption is that there can only be two cells per facet, and should always be
# two cells per facet.
num_cells_per_facet = f_to_c.offsets[facet_indices + 1] - f_to_c.offsets[facet_indices]
assert np.all(num_cells_per_facet == 2), "All facets must be interior facets."
facet_pos = np.vstack([f_to_c.offsets[facet_indices], f_to_c.offsets[facet_indices] + 1]).T
cells = f_to_c.array[facet_pos].flatten()
# Extract facets connected to all cells
# Assumption is that all cells have the same number of facets
num_facets_per_cell = c_to_f.offsets[1:] - c_to_f.offsets[:-1]
assert all(
num_facets_per_cell[cells.flatten()] == num_facets_per_cell[cells.flatten()[0]]
), "Cells must have facets."
facets = np.vstack(
[
c_to_f.array[c_to_f.offsets[cells.flatten()] + i]
for i in range(num_facets_per_cell[cells.flatten()[0]])
]
).T
# Repeat facet indices twice to be able to do vectorized match
rep_fi = np.repeat(facet_indices, 2)
indicator = facets == rep_fi[:, None]
_row, local_pos = np.nonzero(indicator)
assert np.unique(_row).shape[0] == len(_row)
idata = np.vstack([cells, local_pos]).T.reshape(-1, 4)
else:
idata = dolfinx.cpp.fem.compute_integration_domains(
dolfinx.fem.IntegralType.interior_facet,
cell_tags.topology,
facet_indices,
*integration_args,
)
ordered_idata = idata.reshape(-1, 4).copy()
switch = cell_tags.values[ordered_idata[:, 0]] > cell_tags.values[ordered_idata[:, 2]]
if True in switch:
ordered_idata[switch, :] = ordered_idata[switch][:, [2, 3, 0, 1]]
return ordered_idata
def create_geometry_function_space(
mesh: dolfinx.mesh.Mesh, N: int | None = None
) -> dolfinx.fem.FunctionSpace:
"""
Reconstruct a vector space with N components using
the geometry dofmap to ensure a 1-1 mapping between mesh nodes and DOFs.
Args:
mesh: The mesh to create the function space on.
N: The number of components. If not provided the geometrical dimension is chosen
"""
geom_imap = mesh.geometry.index_map()
geom_dofmap = mesh.geometry.dofmap
ufl_domain = mesh.ufl_domain()
assert ufl_domain is not None
sub_el = ufl_domain.ufl_coordinate_element().sub_elements[0]
adj_list = dolfinx.cpp.graph.AdjacencyList_int32(geom_dofmap)
value_shape: tuple[int, ...]
if N is None:
ufl_el = basix.ufl.blocked_element(sub_el, shape=(mesh.geometry.dim,))
value_shape = (mesh.geometry.dim,)
N = value_shape[0]
elif N == 1:
ufl_el = sub_el
value_shape = ()
else:
ufl_el = basix.ufl.blocked_element(sub_el, shape=(N,))
value_shape = (N,)
if ufl_el.dtype == np.float32:
_fe_constructor = dolfinx.cpp.fem.FiniteElement_float32
_fem_constructor = dolfinx.cpp.fem.FunctionSpace_float32
elif ufl_el.dtype == np.float64:
_fe_constructor = dolfinx.cpp.fem.FiniteElement_float64
_fem_constructor = dolfinx.cpp.fem.FunctionSpace_float64
else:
raise RuntimeError(f"Unsupported type {ufl_el.dtype}")
try:
cpp_el = _fe_constructor(ufl_el.basix_element._e, block_shape=value_shape, symmetric=False)
except TypeError:
cpp_el = _fe_constructor(ufl_el.basix_element._e, block_size=N, symmetric=False)
dof_layout = dolfinx.cpp.fem.create_element_dof_layout(cpp_el, [])
cpp_dofmap = dolfinx.cpp.fem.DofMap(dof_layout, geom_imap, N, adj_list, N)
# Create function space
try:
cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap)
except TypeError:
cpp_space = _fem_constructor(mesh._cpp_object, cpp_el, cpp_dofmap, value_shape=value_shape)
return dolfinx.fem.FunctionSpace(mesh, ufl_el, cpp_space)
def move(
mesh: dolfinx.mesh.Mesh,
u: dolfinx.fem.Function
| ufl.core.expr.Expr
| typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.inexact]],
):
"""
Move the geometry nodes of a mesh given by the movement of a function u.
Args:
mesh: The mesh to move
u: The displacement as a :py:class:dolfinx.fem.Function`,
:py:class:`ufl.core.expr.Expr` or a lambda function.
"""
V_geom = create_geometry_function_space(mesh)
u_geom = dolfinx.fem.Function(V_geom, dtype=mesh.geometry.x.dtype)
if isinstance(u, dolfinx.fem.Function):
u_geom.interpolate(u)
elif isinstance(u, ufl.core.expr.Expr):
u_compiled = dolfinx.fem.Expression(u, V_geom.element.interpolation_points)
u_geom.interpolate(u_compiled)
else:
u_geom.interpolate(u)
mesh.geometry.x[:, : mesh.geometry.dim] += u_geom.x.array[:].reshape(-1, mesh.geometry.dim)
def copy(mesh: dolfinx.mesh.Mesh) -> dolfinx.mesh.Mesh:
"""
Copy a mesh by copying its geometry.
Note:
The C++ topology object is shared between the two meshes.
Args:
mesh: The mesh.
Returns:
The new mesh
"""
# Initialize C++ mesh
cpp_mesh = type(mesh._cpp_object)(
mesh.comm, mesh.topology._cpp_object, mesh.geometry._cpp_object
)
# Initialize Python wrapper with new symbolic numbering in the `ufl.Mesh`
return dolfinx.mesh.Mesh(cpp_mesh, ufl.Mesh(mesh.ufl_domain().ufl_coordinate_element()))