import tempfile
import typing
import xml.etree.ElementTree as ET
from enum import IntEnum
from functools import lru_cache
from pathlib import Path
from typing import Iterable, NamedTuple
from mpi4py import MPI
import basix
import dolfinx
import fenicsx_compat
import numpy as np
import numpy.typing as npt
from structlog import get_logger
logger = get_logger()
quads = ("Quadrature", "Q", "Quad", "quadrature", "q", "quad")
QUADNR = 42
[docs]
def parse_element(
space_string: str, mesh: dolfinx.mesh.Mesh, dim: int, discontinuous: bool = False
) -> basix.ufl._ElementBase:
"""
Parse a string representation of a basix element family
"""
family_str, degree_str = space_string.split("_")
kwargs: dict[str, typing.Any] = {"degree": int(degree_str), "cell": mesh.basix_cell()}
if dim > 1:
if family_str in quads:
kwargs["value_shape"] = (dim,)
else:
kwargs["shape"] = (dim,)
el: basix.ufl._ElementBase
if family_str in ["Lagrange", "P", "CG"]:
el = basix.ufl.element(family=basix.ElementFamily.P, discontinuous=discontinuous, **kwargs)
elif family_str in ["Discontinuous Lagrange", "DG", "dP"]:
el = basix.ufl.element(family=basix.ElementFamily.P, discontinuous=True, **kwargs)
elif family_str in quads:
el = basix.ufl.quadrature_element(scheme="default", **kwargs)
else:
families = list(basix.ElementFamily.__members__.keys())
msg = f"Unknown element family: {family_str!r}, available families: {families}"
raise ValueError(msg)
return el
[docs]
def space_from_string(
space_string: str, mesh: dolfinx.mesh.Mesh, dim: int, discontinuous: bool = False
) -> dolfinx.fem.FunctionSpace:
"""
Constructed a finite elements space from a string
representation of the space
Arguments
---------
space_string : str
A string on the form {family}_{degree} which
determines the space. Example 'Lagrange_1'.
mesh : df.Mesh
The mesh
dim : int
1 for scalar space, 3 for vector space.
discontinuous: bool
If true force element to be discontinuous, by default False
"""
el = parse_element(space_string, mesh, dim, discontinuous=discontinuous)
return dolfinx.fem.functionspace(mesh, el)
def element2array(el: basix.ufl._BlockedElement) -> np.ndarray:
if el.family_name in quads:
return np.array(
[QUADNR, int(el.cell_type), int(el.degree), 0],
dtype=np.uint8,
)
else:
return np.array(
[
int(el.basix_element.family),
int(el.cell_type),
int(el.degree),
int(el.basix_element.discontinuous),
],
dtype=np.uint8,
)
_E = typing.TypeVar("_E", bound=IntEnum)
def number2Enum(num: int, enum: Iterable[_E]) -> _E:
for e in enum:
if int(e) == num:
return e
raise ValueError(f"Invalid value {num} for enum {enum}")
def array2element(arr: np.ndarray) -> basix.ufl._ElementBase:
cell_type = number2Enum(arr[1], basix.CellType)
degree = int(arr[2])
discontinuous = bool(arr[3])
if arr[0] == QUADNR:
return basix.ufl.quadrature_element(
scheme="default", degree=degree, value_shape=(3,), cell=cell_type
)
else:
family = number2Enum(arr[0], basix.ElementFamily)
# TODO: Shape is hardcoded to (3,) for now, but this should also be stored
return basix.ufl.element(
family=family,
cell=cell_type,
degree=degree,
discontinuous=discontinuous,
shape=(3,),
lagrange_variant=basix.LagrangeVariant.unset,
)
@lru_cache
def array2functionspace(mesh: dolfinx.mesh.Mesh, arr: np.ndarray) -> dolfinx.fem.FunctionSpace:
el = array2element(arr)
return dolfinx.fem.functionspace(mesh, el)
def handle_mesh_name(mesh_name: str = "") -> Path:
if mesh_name == "":
fd, mesh_name = tempfile.mkstemp(suffix=".msh")
return Path(mesh_name).with_suffix(".msh")
def json_serial(obj):
if isinstance(obj, (np.ndarray)):
return obj.tolist()
else:
try:
return str(obj)
except Exception:
raise TypeError("Type %s not serializable" % type(obj))
[docs]
class GMshGeometry(NamedTuple):
mesh: dolfinx.mesh.Mesh
cfun: dolfinx.mesh.MeshTags | None
ffun: dolfinx.mesh.MeshTags | None
efun: dolfinx.mesh.MeshTags | None
vfun: dolfinx.mesh.MeshTags | None
markers: dict[str, tuple[int, int]]
[docs]
def tag_dimensions(tdim: int) -> list[tuple[str, str, int]]:
"""The tags of a mesh of topological dimension `tdim` and the dimension of
their entities. A mesh of dimension 2 has no edge tags: its edges are its
facets."""
tags = [("cfun", "Cell tags", tdim), ("ffun", "Facet tags", tdim - 1)]
if tdim == 3:
tags.append(("efun", "Edge tags", 1))
tags.append(("vfun", "Vertex tags", 0))
return tags
def read_mesh(
comm,
filename: str | Path,
ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.none,
) -> tuple[dolfinx.mesh.Mesh, dict[str, dolfinx.mesh.MeshTags]]:
tags = {}
with dolfinx.io.XDMFFile(comm, filename, "r") as xdmf:
mesh = xdmf.read_mesh(name="Mesh", ghost_mode=ghost_mode)
for var, name, dim in tag_dimensions(mesh.topology.dim):
mesh.topology.create_connectivity(dim, mesh.topology.dim)
try:
tags[var] = xdmf.read_meshtags(mesh, name=name)
except (RuntimeError, IndexError):
# A missing grid raises RuntimeError in older versions of dolfinx
# and IndexError in newer ones
continue
return mesh, tags
def gmsh2dolfin(
comm: MPI.Comm,
msh_file,
rank: int = 0,
ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.none,
gdim: int = 3,
) -> GMshGeometry:
logger.debug(f"Convert file {msh_file} to dolfin")
outdir = Path(msh_file).parent
outdir.mkdir(parents=True, exist_ok=True)
import gmsh
if comm.rank == rank:
gmsh.initialize()
gmsh.model.add("Mesh from file")
gmsh.merge(str(msh_file))
mesh_data = fenicsx_compat.model_to_mesh(
gmsh.model, comm, rank, gdim=gdim, ghost_mode=ghost_mode
)
if comm.rank == rank:
gmsh.finalize()
mesh = mesh_data.mesh
markers_ = mesh_data.physical_groups
ct = mesh_data.cell_tags
tdim = mesh.topology.dim
if ct is None:
ct = dolfinx.mesh.meshtags(
mesh, tdim, np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
)
ft = mesh_data.facet_tags
if ft is None:
ft = dolfinx.mesh.meshtags(
mesh, tdim - 1, np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
)
if hasattr(mesh_data, "edge_tags"):
ridge_tags = mesh_data.edge_tags # type: ignore[attr-defined]
else:
ridge_tags = mesh_data.ridge_tags
if hasattr(mesh_data, "vertex_tags"):
peak_tags = mesh_data.vertex_tags # type: ignore[attr-defined]
else:
peak_tags = mesh_data.peak_tags
# The ridges are the edges of a 3D mesh but the vertices of a 2D one
et: dolfinx.mesh.MeshTags | None
if tdim == 3:
et, vt = ridge_tags, peak_tags
if et is None:
et = dolfinx.mesh.meshtags(
mesh, 1, np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
)
else:
et, vt = None, ridge_tags
if vt is None:
vt = dolfinx.mesh.meshtags(
mesh, 0, np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32)
)
markers = {k: (v[1], v[0]) for k, v in markers_.items()}
mesh.name = "Mesh"
ct.name = "Cell tags"
ft.name = "Facet tags"
if et is not None:
et.name = "Edge tags"
vt.name = "Vertex tags"
save_mesh_to_xdmf(comm, outdir / "mesh.xdmf", mesh, ct, ft, et, vt)
return GMshGeometry(mesh, ct, ft, et, vt, markers)
[docs]
def transfer_to_parent(
u_sub: dolfinx.fem.Function,
V: dolfinx.fem.FunctionSpace,
parent_cells: npt.NDArray[np.int32],
) -> dolfinx.fem.Function:
"""Copy a function on a submesh into `V` on the parent mesh, cell by cell.
`V` must use the same element as the space of `u_sub`, and `parent_cells` is
the parent cell of each cell of the submesh (see `extract_submesh`). The
degrees of freedom of the parent outside the submesh are zero. The local
dofs are copied as is, which is exact for elements whose dofs do not need
permuting between the meshes, such as Lagrange of degree at most 2 and
quadrature elements.
"""
u_sub.x.scatter_forward()
u = dolfinx.fem.Function(V, name=u_sub.name)
# Count the number of writes to each dof, as a dof shared between processes
# may only be written to on some of them
count = dolfinx.fem.Function(V)
bs = V.dofmap.bs
sub_dofs = u_sub.function_space.dofmap.list[: len(parent_cells)]
parent_dofs = V.dofmap.list[parent_cells]
for b in range(bs):
u.x.array[parent_dofs * bs + b] = u_sub.x.array[sub_dofs * bs + b]
count.x.array[parent_dofs * bs + b] = 1.0
u.x.scatter_reverse(dolfinx.la.InsertMode.add)
count.x.scatter_reverse(dolfinx.la.InsertMode.add)
num_owned = V.dofmap.index_map.size_local * bs
written = count.x.array[:num_owned] > 0
u.x.array[:num_owned][written] /= count.x.array[:num_owned][written]
u.x.scatter_forward()
return u
def save_mesh_to_xdmf(
comm: MPI.Comm,
fname: Path,
mesh: dolfinx.mesh.Mesh,
ct: dolfinx.mesh.MeshTags | None,
ft: dolfinx.mesh.MeshTags | None,
et: dolfinx.mesh.MeshTags | None,
vt: dolfinx.mesh.MeshTags | None,
):
domain = mesh.ufl_domain()
assert domain is not None
if domain.ufl_coordinate_element().degree > 2:
# XDMF cannot represent cells of higher order, but they are saved to geometry.bp
logger.info(f"Skip saving {fname}: XDMF only supports meshes of order at most 2")
return
tdim = mesh.topology.dim
# Save tags to xdmf
with dolfinx.io.XDMFFile(comm, fname, "w") as xdmf:
xdmf.write_mesh(mesh)
if ct is not None:
xdmf.write_meshtags(
ct,
mesh.geometry,
geometry_xpath=f"/Xdmf/Domain/Grid[@Name='{mesh.name}']/Geometry",
)
if ft is not None:
mesh.topology.create_connectivity(ft.dim, tdim)
xdmf.write_meshtags(
ft,
mesh.geometry,
geometry_xpath=f"/Xdmf/Domain/Grid[@Name='{mesh.name}']/Geometry",
)
if et is not None:
mesh.topology.create_connectivity(et.dim, tdim)
xdmf.write_meshtags(
et,
mesh.geometry,
geometry_xpath=f"/Xdmf/Domain/Grid[@Name='{mesh.name}']/Geometry",
)
if vt is not None:
mesh.topology.create_connectivity(vt.dim, tdim)
xdmf.write_meshtags(
vt,
mesh.geometry,
geometry_xpath=f"/Xdmf/Domain/Grid[@Name='{mesh.name}']/Geometry",
)
def create_xdmf_pointcloud(filename: Path, us: typing.Sequence[dolfinx.fem.Function]) -> None:
# Adopted from https://gist.github.com/jorgensd/8bae61ad7a0c211570dff0116a68a356
if len(us) == 0:
return
import adios2
u = us[0]
points = u.function_space.tabulate_dof_coordinates()
h5name = filename.with_suffix(".h5")
bs = u.function_space.dofmap.index_map_bs
comm = u.function_space.mesh.comm
num_dofs_global = u.function_space.dofmap.index_map.size_global
num_dofs_local = u.function_space.dofmap.index_map.size_local
local_range = np.array(u.function_space.dofmap.index_map.local_range, dtype=np.int64)
# Write XDMF on rank 0
if comm.rank == 0:
xdmf = ET.Element("XDMF")
xdmf.attrib["Version"] = "3.0"
xdmf.attrib["xmlns:xi"] = "http://www.w3.org/2001/XInclude"
domain = ET.SubElement(xdmf, "Domain")
grid = ET.SubElement(domain, "Grid")
grid.attrib["GridType"] = "Uniform"
grid.attrib["Name"] = "Point Cloud"
topology = ET.SubElement(grid, "Topology")
topology.attrib["NumberOfElements"] = str(num_dofs_global)
topology.attrib["TopologyType"] = "PolyVertex"
topology.attrib["NodesPerElement"] = "1"
geometry = ET.SubElement(grid, "Geometry")
geometry.attrib["GeometryType"] = "XY" if points.shape[1] == 2 else "XYZ"
for u in us:
it0 = ET.SubElement(geometry, "DataItem")
it0.attrib["Dimensions"] = f"{num_dofs_global} {points.shape[1]}"
it0.attrib["Format"] = "HDF"
it0.text = f"{h5name.name}:/Step0/Points"
attrib = ET.SubElement(grid, "Attribute")
attrib.attrib["Name"] = u.name
if bs == 1:
attrib.attrib["AttributeType"] = "Scalar"
else:
attrib.attrib["AttributeType"] = "Vector"
attrib.attrib["Center"] = "Node"
it1 = ET.SubElement(attrib, "DataItem")
it1.attrib["Dimensions"] = f"{num_dofs_global} {bs}"
it1.attrib["Format"] = "HDF"
it1.text = f"{h5name.name}:/Step0/Values_{u.name}"
text = [
'<?xml version="1.0"?>\n<!DOCTYPE Xdmf SYSTEM "Xdmf.dtd" []>\n',
ET.tostring(xdmf, encoding="unicode"),
]
filename.write_text("".join(text))
# Create ADIOS2 reader
# start = time.perf_counter()
adios = adios2.ADIOS(comm)
io = adios.DeclareIO("Point cloud writer")
io.SetEngine("HDF5")
outfile = io.Open(h5name.as_posix(), adios2.Mode.Write)
points_out = points[:num_dofs_local, :]
pointvar = io.DefineVariable(
"Points",
points_out,
shape=[num_dofs_global, points.shape[1]],
start=[local_range[0], 0],
count=[num_dofs_local, points.shape[1]],
)
outfile.Put(pointvar, points_out)
for u in us:
data = u.x.array[: num_dofs_local * bs].reshape(-1, bs)
valuevar = io.DefineVariable(
f"Values_{u.name}",
data,
shape=[num_dofs_global, bs],
start=[local_range[0], 0],
count=[num_dofs_local, bs],
)
outfile.Put(valuevar, data)
outfile.PerformPuts()
outfile.Close()
assert adios.RemoveIO("Point cloud writer")
[docs]
class BaseData(NamedTuple):
centroid: npt.NDArray[np.float64]
vector: npt.NDArray[np.float64]
normal: npt.NDArray[np.float64]
[docs]
def compute_base_data(
mesh: dolfinx.mesh.Mesh,
facet_tags: dolfinx.mesh.MeshTags,
marker,
) -> BaseData:
"""Compute the centroid, vector and normal of the base
Parameters
----------
mesh : dolfinx.mesh.Mesh
The mesh
facet_tags : dolfinx.mesh.MeshTags
The facet tags
marker : _type_
Marker for the base
Returns
-------
BaseData
NamedTuple containing the centroid, vector and normal of the base
"""
base_facets = facet_tags.find(marker)
base_midpoints = mesh.comm.gather(
dolfinx.mesh.compute_midpoints(mesh, 2, base_facets),
root=0,
)
base_vector = np.zeros(3)
base_centroid = np.zeros(3)
base_normal = np.zeros(3)
if mesh.comm.rank == 0:
assert base_midpoints is not None
bm = np.concatenate(base_midpoints)
base_centroid = bm.mean(axis=0)
# print("Base centroid", len(base_midpoints))
base_points_centered = bm - base_centroid
u, s, vh = np.linalg.svd(base_points_centered)
base_normal = vh[-1, :]
# Initialize vector to be used for cross product
vector_init = np.array([0, 1, 0])
# If the normal is parallel to the initial vector, change the initial vector
if np.allclose(np.abs(base_normal), np.abs(vector_init)):
vector_init = np.array([0, 0, 1])
# Find two vectors in the plane, orthogonal to the normal
vector = np.cross(base_normal, vector_init)
base_vector = np.cross(base_normal, vector)
base_centroid = mesh.comm.bcast(base_centroid, root=0)
base_vector = mesh.comm.bcast(base_vector, root=0)
base_normal = mesh.comm.bcast(base_normal, root=0)
return BaseData(centroid=base_centroid, vector=base_vector, normal=base_normal)
[docs]
def rotate_geometry_and_fields(
mesh: dolfinx.mesh.Mesh,
ffun: dolfinx.mesh.MeshTags,
base_marker: int,
target_normal: typing.Sequence[float] = tuple((1.0, 0.0, 0.0)),
fields: typing.Sequence[dolfinx.fem.Function] | None = None,
) -> tuple[dolfinx.mesh.Mesh, np.ndarray, list[dolfinx.fem.Function]]:
"""
Rotates the mesh and vector fields so that the base normal points in the +x direction.
"""
comm = mesh.comm
if comm.size > 1:
raise NotImplementedError("rotate_geometry_and_fields only works in serial.")
if fields is None:
fs: list[dolfinx.fem.Function] = []
else:
fs = list(fields)
# Compute current base normal
# We use pulse's built-in helper to find the normal of the surface marked as base
base_data = compute_base_data(mesh, ffun, base_marker)
current_normal = base_data.normal
# Create rotation matrix
R_matrix = np.eye(3)
target_normal_normalized: np.ndarray = np.array(target_normal) / np.linalg.norm(target_normal)
# Compute Rotation Matrix
# We find the rotation that aligns current_normal to target_normal
if np.allclose(current_normal, target_normal_normalized):
logger.info("Geometry already aligned.")
return mesh, R_matrix, fs
# Simple check to avoid singular rotation if vectors are opposite
if np.allclose(current_normal, -target_normal_normalized):
R_matrix = np.diag([-1.0, -1.0, 1.0]) # 180 flip
else:
# Calculate cross product and angle
v = np.cross(current_normal, target_normal_normalized)
s = np.linalg.norm(v)
c = np.dot(current_normal, target_normal_normalized)
if s > 1e-8:
# Skew-symmetric cross-product matrix of v
vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
R_matrix = np.eye(3) + vx + vx @ vx * ((1 - c) / (s**2))
new_fields = []
for f in fs:
# Ensure it is a vector function
bs = f.function_space.dofmap.bs
if bs != 3:
continue
# Get underlying array
values = f.x.array
# Reshape to (Num_DOFs, 3)
num_dofs = len(values) // 3
vectors = values.reshape((num_dofs, 3))
# Rotate vectors: v_new = R * v_old
vectors_rotated = vectors @ R_matrix.T
# Assign back
f.x.array[:] = vectors_rotated.flatten()
new_fields.append(f)
# Rotate Mesh Coordinates
# geometry.x is N x 3. We apply x_new = R * x_old^T -> x_new^T = x_old * R^T
coords = mesh.geometry.x[:, :3]
mesh.geometry.x[:, :3] = coords @ R_matrix.T
logger.info(f"Rotated geometry. Base normal {current_normal} aligned to {target_normal}")
logger.debug(f"Rotation matrix:\n{R_matrix}")
return mesh, R_matrix, new_fields
