import json
import logging
import shutil
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from mpi4py import MPI
import dolfinx
import io4dolfinx
import ufl
from . import utils
logger = logging.getLogger(__name__)
#: io4dolfinx backends that support checkpointing (mesh, meshtags and functions),
#: mapped to the file suffix used for the geometry file.
IO_BACKEND_SUFFIXES: dict[str, str] = {"adios2": ".bp", "h5py": ".h5"}
IO_BACKENDS: tuple[str, ...] = tuple(IO_BACKEND_SUFFIXES)
DEFAULT_IO_BACKEND = "adios2"
def check_io_backend(backend: str) -> str:
if backend not in IO_BACKEND_SUFFIXES:
raise ValueError(f"Unsupported io backend {backend!r}, expected one of {IO_BACKENDS}")
return backend
[docs]
def geometry_path(folder: str | Path, backend: str = DEFAULT_IO_BACKEND) -> Path:
"""Path of the geometry file in `folder` written with `backend`,
i.e. ``geometry.bp`` for adios2 and ``geometry.h5`` for h5py."""
return Path(folder) / f"geometry{IO_BACKEND_SUFFIXES[check_io_backend(backend)]}"
[docs]
def infer_io_backend(path: str | Path) -> str:
"""Infer the io4dolfinx backend from the suffix of `path`,
defaulting to adios2 for unknown suffixes."""
suffix = Path(path).suffix.lower()
if suffix in (".h5", ".hdf5"):
return "h5py"
return DEFAULT_IO_BACKEND
def _remove_path(path: Path) -> None:
if path.is_dir():
shutil.rmtree(path, ignore_errors=True)
else:
path.unlink(missing_ok=True)
def info_path(path: Path) -> Path:
return path.parent / f"info_{path.stem}.json"
def markers_path(path: Path) -> Path:
return path.parent / f"markers_{path.stem}.json"
def microstructure_path(path: Path) -> Path:
return path.parent / f"microstructure_{path.stem}.json"
def additional_data_path(path: Path) -> Path:
return path.parent / f"additional_data_{path.stem}.json"
def save_additional_data(
additional_data: dict[str, dolfinx.mesh.MeshTags | dolfinx.fem.Function],
path: Path,
mesh: dolfinx.mesh.Mesh,
backend: str | None = None,
) -> None:
backend = backend or infer_io_backend(path)
additional_data_attributes: dict[str, dict[str, Any]] = {}
for name, data in additional_data.items():
if isinstance(data, dolfinx.mesh.MeshTags):
io4dolfinx.write_meshtags(
meshtags=data,
mesh=mesh,
filename=path,
meshtag_name=name,
backend=backend,
)
additional_data_attributes[name] = {"type": "MeshTags"}
elif isinstance(data, dolfinx.fem.Function):
io4dolfinx.write_function(
u=data,
filename=path,
name=name,
backend=backend,
)
additional_data_attributes[name] = {
"type": "Function",
"element": utils.element2array(data.ufl_element()),
}
else:
logger.warning(
f"Additional data '{name}' of type {type(data)} "
"is not supported and will not be saved."
)
if mesh.comm.rank == 0:
additional_data_path(path).write_text(
json.dumps(additional_data_attributes, indent=4, default=utils.json_serial)
)
def load_additional_data(
path: Path,
mesh: dolfinx.mesh.Mesh,
backend: str | None = None,
) -> dict[str, dolfinx.mesh.MeshTags | dolfinx.fem.Function]:
backend = backend or infer_io_backend(path)
additional_data: dict[str, dolfinx.mesh.MeshTags | dolfinx.fem.Function] = {}
if not additional_data_path(path).exists():
return additional_data
if mesh.comm.rank == 0:
additional_data_attributes = json.loads(additional_data_path(path).read_text())
else:
additional_data_attributes = {}
additional_data_attributes = mesh.comm.bcast(additional_data_attributes, root=0)
for name, attr in additional_data_attributes.items():
if attr["type"] == "MeshTags":
mt = io4dolfinx.read_meshtags(
mesh=mesh, meshtag_name=name, filename=path, backend=backend
)
additional_data[name] = mt
elif attr["type"] == "Function":
V = utils.array2functionspace(mesh, tuple(attr["element"]))
f = dolfinx.fem.Function(V, name=name)
io4dolfinx.read_function(u=f, filename=path, name=name, backend=backend)
additional_data[name] = f
else:
logger.warning(
f"Additional data '{name}' of type {attr['type']} "
"is not supported and will not be loaded."
)
return additional_data
[docs]
def save_geometry(
path: Path,
mesh: dolfinx.mesh.Mesh,
markers: dict[str, tuple[int, int]] | None = None,
info: dict[str, Any] | None = None,
cfun: dolfinx.mesh.MeshTags | None = None,
ffun: dolfinx.mesh.MeshTags | None = None,
efun: dolfinx.mesh.MeshTags | None = None,
vfun: dolfinx.mesh.MeshTags | None = None,
f0: dolfinx.fem.Function | None = None,
s0: dolfinx.fem.Function | None = None,
n0: dolfinx.fem.Function | None = None,
additional_data: dict[str, dolfinx.mesh.MeshTags | dolfinx.fem.Function] | None = None,
backend: str | None = None,
) -> None:
"""Save a geometry to `path` using io4dolfinx.
`backend` is the io4dolfinx backend ("adios2" or "h5py"). If None it is
inferred from the suffix of `path` (``.h5``/``.hdf5`` gives h5py, otherwise adios2).
"""
path = Path(path)
backend = check_io_backend(backend or infer_io_backend(path))
if mesh.comm.rank == 0:
_remove_path(path)
# The json files next to the geometry file are shared between the backends,
# so remove geometry files with the same stem written with other backends
for suffix in IO_BACKEND_SUFFIXES.values():
if suffix != path.suffix:
_remove_path(path.with_suffix(suffix))
mesh.comm.barrier()
io4dolfinx.write_mesh(mesh=mesh, filename=path, backend=backend)
if mesh.comm.rank == 0:
if markers is not None:
markers_path(path).write_text(json.dumps(markers, indent=4, default=utils.json_serial))
if info is not None:
info_path(path).write_text(json.dumps(info, indent=4, default=utils.json_serial))
if cfun is not None:
io4dolfinx.write_meshtags(
meshtags=cfun,
mesh=mesh,
filename=path,
meshtag_name="Cell tags",
backend=backend,
)
if ffun is not None:
io4dolfinx.write_meshtags(
meshtags=ffun,
mesh=mesh,
filename=path,
meshtag_name="Facet tags",
backend=backend,
)
if efun is not None:
io4dolfinx.write_meshtags(
meshtags=efun,
mesh=mesh,
filename=path,
meshtag_name="Edge tags",
backend=backend,
)
if vfun is not None:
io4dolfinx.write_meshtags(
meshtags=vfun,
mesh=mesh,
filename=path,
meshtag_name="Vertex tags",
backend=backend,
)
from .fibers.utils import save_microstructure
save_microstructure(
mesh=mesh,
functions=[f for f in (f0, s0, n0) if f is not None],
path=path,
backend=backend,
)
if additional_data is not None:
save_additional_data(additional_data, path, mesh=mesh, backend=backend)
mesh.comm.barrier()
[docs]
@dataclass
class Geometry:
mesh: dolfinx.mesh.Mesh
markers: dict[str, tuple[int, int]] = field(default_factory=dict)
cfun: dolfinx.mesh.MeshTags | None = None
ffun: dolfinx.mesh.MeshTags | None = None
efun: dolfinx.mesh.MeshTags | None = None
vfun: dolfinx.mesh.MeshTags | None = None
f0: dolfinx.fem.Function | None = None
s0: dolfinx.fem.Function | None = None
n0: dolfinx.fem.Function | None = None
info: dict[str, Any] = field(default_factory=dict)
additional_data: dict[str, dolfinx.mesh.MeshTags | dolfinx.fem.Function] = field(
default_factory=dict
)
quadrature_degree: int = 4
[docs]
def save(self, path: str | Path, backend: str | None = None) -> None:
"""Save the geometry to a file using io4dolfinx.
Parameters
----------
path : str | Path
The path to the file where the geometry will be saved.
The file will be created if it does not exist, or overwritten if it does.
backend : str | None, optional
The io4dolfinx backend, "adios2" or "h5py". By default it is inferred
from the suffix of `path` (``.h5`` gives h5py, otherwise adios2).
"""
path = Path(path)
save_geometry(
path=path,
mesh=self.mesh,
markers=self.markers,
cfun=self.cfun,
ffun=self.ffun,
efun=self.efun,
vfun=self.vfun,
f0=self.f0,
s0=self.s0,
n0=self.n0,
additional_data=self.additional_data,
backend=backend,
)
[docs]
def save_folder(self, folder: str | Path, backend: str = DEFAULT_IO_BACKEND) -> None:
"""Save the geometry to a folder containing mesh and markers files.
Parameters
----------
folder : str | Path
The path to the folder where the geometry will be saved.
The folder will be created if it does not exist.
backend : str, optional
The io4dolfinx backend used for the geometry file, "adios2" (written to
``geometry.bp``) or "h5py" (written to ``geometry.h5``), by default "adios2".
"""
comm = self.mesh.comm
folder = Path(folder)
folder.mkdir(parents=True, exist_ok=True)
utils.save_mesh_to_xdmf(
comm=comm,
fname=folder / "mesh.xdmf",
mesh=self.mesh,
ct=self.cfun,
ft=self.ffun,
et=self.efun,
vt=self.vfun,
)
if comm.rank == 0:
(folder / "markers.json").write_text(json.dumps(self.markers, indent=4))
(folder / "info.json").write_text(json.dumps(self.info, indent=4))
save_geometry(
path=geometry_path(folder, backend),
mesh=self.mesh,
markers=self.markers,
cfun=self.cfun,
ffun=self.ffun,
efun=self.efun,
vfun=self.vfun,
f0=self.f0,
s0=self.s0,
n0=self.n0,
additional_data=self.additional_data,
backend=backend,
)
logger.info(f"Geometry saved to {folder}")
@property
def dx(self):
"""Volume measure for the mesh using
the cell function `cfun` if it exists as subdomain data.
"""
return ufl.Measure(
"dx",
domain=self.mesh,
subdomain_data=self.cfun,
metadata={"quadrature_degree": self.quadrature_degree},
)
@property
def ds(self):
"""Surface measure for the mesh using
the facet function `ffun` if it exists as subdomain data.
"""
return ufl.Measure(
"ds",
domain=self.mesh,
subdomain_data=self.ffun,
metadata={"quadrature_degree": self.quadrature_degree},
)
@property
def facet_normal(self) -> ufl.FacetNormal:
"""Facet normal vector for the mesh."""
return ufl.FacetNormal(self.mesh)
@property
def facet_tags(self) -> dolfinx.mesh.MeshTags | None:
return self.ffun
[docs]
def refine(self, n=1, outdir: Path | None = None) -> "Geometry":
"""
Refine the mesh and transfer the meshtags to new geometry.
Also regenerate fibers if `self.info` is found.
If `self.info` is not found, it currently raises a
NotImplementedError, however fiber could be interpolated
from the old mesh to the new mesh but this will result in a
loss of information about the fiber orientation.
Parameters
----------
n : int, optional
Number of times to refine the mesh, by default 1
outdir : Path | None, optional
Output directory to save the refined mesh and meshtags,
by default None in which case the mesh is not saved.
Returns
-------
Geometry
A new Geometry object with the refined mesh and updated meshtags.
Raises
------
NotImplementedError
If `self.info` is not found, indicating that fiber
interpolation after refinement is not implemented yet.
"""
mesh = self.mesh
cfun = self.cfun
ffun = self.ffun
for _ in range(n):
# Newer versions of dolfinx require the edges to be created before refining
mesh.topology.create_entities(1)
new_mesh, parent_cell, parent_facet = dolfinx.mesh.refine(
mesh,
partitioner=None,
option=dolfinx.mesh.RefinementOption.parent_cell_and_facet,
)
assert parent_cell is not None
new_mesh.name = mesh.name
mesh = new_mesh
new_mesh.topology.create_entities(1)
new_mesh.topology.create_connectivity(2, 3)
if cfun is not None:
new_cfun = dolfinx.mesh.transfer_meshtag(cfun, new_mesh, parent_cell, parent_facet)
new_cfun.name = cfun.name
cfun = new_cfun
else:
new_cfun = None
if ffun is not None:
new_ffun = dolfinx.mesh.transfer_meshtag(ffun, new_mesh, parent_cell, parent_facet)
new_ffun.name = ffun.name
ffun = new_ffun
else:
new_ffun = None
if outdir is not None:
outdir = Path(outdir)
outdir.mkdir(parents=True, exist_ok=True)
with dolfinx.io.XDMFFile(new_mesh.comm, outdir / "mesh.xdmf", "w") as xdmf:
xdmf.write_mesh(new_mesh)
if cfun is not None:
xdmf.write_meshtags(
cfun,
new_mesh.geometry,
geometry_xpath=f"/Xdmf/Domain/Grid[@Name='{new_mesh.name}']/Geometry",
)
if ffun is not None:
mesh.topology.create_connectivity(2, 3)
xdmf.write_meshtags(
ffun,
new_mesh.geometry,
geometry_xpath=f"/Xdmf/Domain/Grid[@Name='{new_mesh.name}']/Geometry",
)
if self.info is not None:
info = self.info.copy()
info["refinement"] = n
info.pop("outdir", None)
if outdir is not None:
info["outdir"] = str(outdir)
from .fibers import generate_fibers
f0, s0, n0 = generate_fibers(
mesh=new_mesh, ffun=new_ffun, cfun=new_cfun, markers=self.markers, **info
)
else:
info = None
# Interpolate fibers
raise NotImplementedError(
"Interpolating fibers after refinement is not implemented yet."
)
return Geometry(
mesh=new_mesh,
markers=self.markers,
cfun=new_cfun,
ffun=new_ffun,
f0=f0,
s0=s0,
n0=n0,
)
[docs]
@classmethod
def from_file(
cls,
comm: MPI.Comm,
path: str | Path,
ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.none,
backend: str | None = None,
) -> "Geometry":
"""Read geometry from a file using io4dolfinx.
Parameters
----------
comm : MPI.Comm
The MPI communicator to use for reading the mesh.
path : str | Path
The path to the file containing the geometry data.
ghost_mode : dolfinx.mesh.GhostMode, optional
The ghost mode to use when reading the mesh, by default
dolfinx.mesh.GhostMode.none.
Note that if you need to refine the mesh later, you should
use dolfinx.mesh.GhostMode.none.
backend : str | None, optional
The io4dolfinx backend, "adios2" or "h5py". By default it is inferred
from the suffix of `path` (``.h5`` gives h5py, otherwise adios2).
Returns
-------
Geometry
An instance of the Geometry class containing the mesh, markers, and functions.
"""
path = Path(path)
if not path.exists():
raise ValueError(f"File {path} does not exist")
backend = check_io_backend(backend or infer_io_backend(path))
if markers_path(path).exists():
if comm.rank == 0:
markers = json.loads(markers_path(path).read_text())
else:
markers = {}
markers = comm.bcast(markers, root=0)
else:
markers = {}
if info_path(path).exists():
if comm.rank == 0:
info = json.loads(info_path(path).read_text())
else:
info = {}
info = comm.bcast(info, root=0)
else:
info = {}
mesh = io4dolfinx.read_mesh(
comm=comm,
filename=path,
ghost_mode=ghost_mode,
backend=backend,
)
# markers = io4dolfinx.read_attributes(comm=comm, filename=path, name="markers")
tags: dict[str, dolfinx.mesh.MeshTags | None] = {}
for name, meshtag_name in (
("cfun", "Cell tags"),
("ffun", "Facet tags"),
("efun", "Edge tags"),
("vfun", "Vertex tags"),
):
try:
tags[name] = io4dolfinx.read_meshtags(
mesh=mesh, meshtag_name=meshtag_name, filename=path, backend=backend
)
except (KeyError, IndexError):
logger.debug(f"{name} not found in {path}")
tags[name] = None
functions: dict[str, dolfinx.fem.Function] = {}
# if function_space_data is None:
# function_space_data = io4dolfinx.read_attributes(
# comm=comm, filename=path, name="function_space"
# )
if microstructure_path(path).exists():
if comm.rank == 0:
function_space_data = json.loads(microstructure_path(path).read_text())
else:
function_space_data = {}
function_space_data = comm.bcast(function_space_data, root=0)
else:
function_space_data = {}
assert isinstance(function_space_data, dict), "function_space_data must be a dictionary"
for name, el in function_space_data.items():
V = utils.array2functionspace(mesh, tuple(el))
f = dolfinx.fem.Function(V, name=name)
try:
io4dolfinx.read_function(u=f, filename=path, name=name, backend=backend)
except KeyError:
continue
else:
functions[name] = f
additional_data = load_additional_data(path=path, mesh=mesh, backend=backend)
return cls(
mesh=mesh,
markers=markers,
info=info,
additional_data=additional_data,
cfun=tags.get("cfun"),
ffun=tags.get("ffun"),
efun=tags.get("efun"),
vfun=tags.get("vfun"),
f0=functions.get("f0"),
s0=functions.get("s0"),
n0=functions.get("n0"),
)
[docs]
@classmethod
def from_folder(
cls,
comm: MPI.Comm,
folder: str | Path,
ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.none,
) -> "Geometry":
"""Read geometry from a folder containing mesh and markers files.
Parameters
----------
comm : MPI.Comm
The MPI communicator to use for reading the mesh and markers.
folder : str | Path
The path to the folder containing the geometry data.
ghost_mode : dolfinx.mesh.GhostMode, optional
The ghost mode to use when reading the mesh, by default
dolfinx.mesh.GhostMode.none.
Note that if you need to refine the mesh later, you should
use dolfinx.mesh.GhostMode.none.
Returns
-------
Geometry
An instance of the Geometry class containing the mesh, markers, and functions.
Raises
------
ValueError
If the required mesh file is not found in the specified folder.
"""
folder = Path(folder)
logger.info(f"Reading geometry from {folder}")
if not folder.exists():
raise ValueError(f"Folder {folder} does not exist")
for backend in IO_BACKENDS:
path = geometry_path(folder, backend)
if path.exists():
logger.debug(f"Reading geometry from {path.name}")
return cls.from_file(comm=comm, path=path, ghost_mode=ghost_mode, backend=backend)
logger.warning(
"geometry.bp/geometry.h5 not found, reading mesh and microstructure separately. "
"This may lead to inconsistent dof numbering in the mesh and microstructure. "
"Try deleting the folder and regenerating the mesh."
)
# Read mesh
if (folder / "mesh.xdmf").exists():
logger.debug("Reading mesh")
mesh, tags = utils.read_mesh(
comm=comm, filename=folder / "mesh.xdmf", ghost_mode=ghost_mode
)
else:
raise ValueError("No mesh file found")
# Read markers
if (folder / "markers.json").exists():
logger.debug("Reading markers")
if comm.rank == 0:
markers = json.loads((folder / "markers.json").read_text())
else:
markers = {}
markers = comm.bcast(markers, root=0)
else:
markers = {}
if (folder / "microstructure.json").exists():
if comm.rank == 0:
microstructure = json.loads((folder / "microstructure.json").read_text())
else:
microstructure = {}
microstructure = comm.bcast(microstructure, root=0)
else:
microstructure = {}
functions: dict[str, dolfinx.fem.Function] = {}
microstructure_path = folder / "microstructure.bp"
if microstructure_path.exists():
logger.debug("Reading microstructure")
# function_space = io4dolfinx.read_attributes(
# comm=MPI.COMM_WORLD, filename=microstructure_path, name="function_space"
# )
for name, el in microstructure.items():
logger.debug(f"Reading {name}")
V = utils.array2functionspace(mesh, tuple(el))
f = dolfinx.fem.Function(V, name=name)
try:
io4dolfinx.read_function(u=f, filename=microstructure_path, name=name)
except KeyError:
continue
else:
functions[name] = f
if (folder / "info.json").exists():
logger.debug("Reading info")
if comm.rank == 0:
info = json.loads((folder / "info.json").read_text())
else:
info = {}
info = comm.bcast(info, root=0)
else:
info = {}
return cls(
mesh=mesh,
markers=markers,
info=info,
cfun=tags.get("cfun"),
ffun=tags.get("ffun"),
efun=tags.get("efun"),
vfun=tags.get("vfun"),
f0=functions.get("f0"),
s0=functions.get("s0"),
n0=functions.get("n0"),
)
[docs]
def rotate(self, target_normal, base_marker):
"""Rotate the geometry so that the base normal aligns with the target normal.
Parameters
----------
target_normal : np.ndarray
The target normal vector to align the base normal with.
base_marker : str
The marker for the base of the geometry.
Returns
-------
Geometry
The rotated Geometry object. Only returned if inplace is False.
"""
from . import utils
msg = (
f"Base marker '{base_marker}' not found in markers. "
f"Available markers: {list(self.markers.keys())}"
)
assert base_marker in self.markers, msg
fields = [f for f in (self.f0, self.s0, self.n0) if f is not None]
mesh_rotated, R_matrix, fields_rotated = utils.rotate_geometry_and_fields(
mesh=self.mesh,
ffun=self.ffun,
base_marker=self.markers[base_marker][0],
target_normal=target_normal,
fields=fields,
)
if self.f0 is not None:
f0, s0, n0 = fields_rotated
else:
f0 = None
s0 = None
n0 = None
self.info["rotation_matrix"] = R_matrix.tolist()
return Geometry(
mesh=mesh_rotated,
markers=self.markers,
cfun=self.cfun,
ffun=self.ffun,
efun=self.efun,
vfun=self.vfun,
f0=f0,
s0=s0,
n0=n0,
info=self.info,
)
