from __future__ import annotations
import numpy
from fractions import Fraction
from mpi4py import MPI
from pyop2 import op2
import pyop2
from firedrake.utils import IntType
from firedrake.functionspacedata import entity_dofs_key
import finat.ufl
import firedrake
from firedrake.cython import mgimpl as impl
[docs]
def identity_node_map(V):
cache = V.mesh()._shared_data_cache["hierarchy_identity_node_map"]
key = (V.ufl_element(), V.boundary_set)
try:
return cache[key]
except KeyError:
values = numpy.arange(V.node_set.total_size, dtype=IntType).reshape(-1, 1)
return cache.setdefault(key, op2.Map(V.node_set, V.node_set, 1, values=values))
[docs]
def fine_node_to_coarse_node_map(Vf, Vc):
if len(Vf) > 1:
assert len(Vf) == len(Vc)
return op2.MixedMap(map(fine_node_to_coarse_node_map, Vf, Vc))
mesh = Vf.mesh()
assert hasattr(mesh, "_shared_data_cache")
hierarchyf, levelf = get_level(Vf.mesh())
hierarchyc, levelc = get_level(Vc.mesh())
if hierarchyc != hierarchyf:
raise ValueError("Can't map across hierarchies")
hierarchy = hierarchyf
increment = Fraction(1, hierarchyf.refinements_per_level)
if levelc + increment != levelf:
raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
key = _cache_key(Vc, Vf)
cache = mesh._shared_data_cache["hierarchy_fine_node_to_coarse_node_map"]
try:
return cache[key]
except KeyError:
assert Vc.extruded == Vf.extruded
if Vc.mesh().variable_layers or Vf.mesh().variable_layers:
raise NotImplementedError("Not implemented for variable layers, sorry")
if Vc.extruded and not ((Vf.mesh().layers - 1)/(Vc.mesh().layers - 1)).is_integer():
raise ValueError("Coarse and fine meshes must have an integer ratio of layers")
fine_to_coarse = hierarchy.fine_to_coarse_cells[levelf]
fine_to_coarse_nodes = impl.fine_to_coarse_nodes(Vf, Vc, fine_to_coarse)
# Detect fine nodes that only touch orphaned fine cells with no coarse parent
orphaned = (fine_to_coarse_nodes[:Vf.node_set.size] < 0).all(axis=1).any()
with pyop2.mpi.temp_internal_comm(Vf.comm) as icomm:
orphaned = icomm.allreduce(bool(orphaned), op=MPI.LOR)
if orphaned:
raise NotImplementedError("Transfer is not implemented for fine nodes that only touch "
"orphaned fine cells with no coarse parent, as in a "
"SubmeshHierarchy of interior facets")
return cache.setdefault(key, op2.Map(Vf.node_set, Vc.node_set,
fine_to_coarse_nodes.shape[1],
values=fine_to_coarse_nodes))
[docs]
def coarse_node_to_fine_node_map(Vc, Vf):
if len(Vf) > 1:
assert len(Vf) == len(Vc)
return op2.MixedMap(map(coarse_node_to_fine_node_map, Vf, Vc))
mesh = Vc.mesh()
assert hasattr(mesh, "_shared_data_cache")
hierarchyf, levelf = get_level(Vf.mesh())
hierarchyc, levelc = get_level(Vc.mesh())
if hierarchyc != hierarchyf:
raise ValueError("Can't map across hierarchies")
hierarchy = hierarchyf
increment = Fraction(1, hierarchyf.refinements_per_level)
if levelc + increment != levelf:
raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
key = _cache_key(Vc, Vf)
cache = mesh._shared_data_cache["hierarchy_coarse_node_to_fine_node_map"]
try:
return cache[key]
except KeyError:
assert Vc.extruded == Vf.extruded
if Vc.mesh().variable_layers or Vf.mesh().variable_layers:
raise NotImplementedError("Not implemented for variable layers, sorry")
if Vc.extruded and not ((Vf.mesh().layers - 1)/(Vc.mesh().layers - 1)).is_integer():
raise ValueError("Coarse and fine meshes must have an integer ratio of layers")
coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc]
coarse_to_fine_nodes = impl.coarse_to_fine_nodes(Vc, Vf, coarse_to_fine)
# op2.Map cannot hold the -1 that pads a short row, so fill each
# padded slot with a real entry from its own row. The injection
# kernel picks the candidate that matches the coarse node's physical
# location, so a repeated entry changes nothing.
valid = coarse_to_fine_nodes >= 0
nonempty = valid.any(axis=1)
if not Vc.comm.allreduce(bool(nonempty[:Vc.node_set.size].all()), op=MPI.LAND):
raise RuntimeError("Adaptive coarse-to-fine map has empty node candidates")
replacement = numpy.zeros(coarse_to_fine_nodes.shape[0],
dtype=coarse_to_fine_nodes.dtype)
rows = numpy.nonzero(nonempty)[0]
replacement[rows] = coarse_to_fine_nodes[rows, valid[rows].argmax(axis=1)]
coarse_to_fine_nodes = numpy.where(valid, coarse_to_fine_nodes,
replacement[:, None])
return cache.setdefault(key, op2.Map(Vc.node_set, Vf.node_set,
coarse_to_fine_nodes.shape[1],
values=coarse_to_fine_nodes))
[docs]
def coarse_cell_to_fine_node_map(Vc, Vf):
if len(Vf) > 1:
assert len(Vf) == len(Vc)
return op2.MixedMap(coarse_cell_to_fine_node_map(f, c) for f, c in zip(Vf, Vc))
mesh = Vc.mesh()
assert hasattr(mesh, "_shared_data_cache")
hierarchyf, levelf = get_level(Vf.mesh())
hierarchyc, levelc = get_level(Vc.mesh())
if hierarchyc != hierarchyf:
raise ValueError("Can't map across hierarchies")
hierarchy = hierarchyf
increment = Fraction(1, hierarchyf.refinements_per_level)
if levelc + increment != levelf:
raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
key = _cache_key(Vc, Vf, needs_coarse_entity_dofs=False)
cache = mesh._shared_data_cache["hierarchy_coarse_cell_to_fine_node_map"]
try:
return cache[key]
except KeyError:
assert Vc.extruded == Vf.extruded
if Vc.mesh().variable_layers or Vf.mesh().variable_layers:
raise NotImplementedError("Not implemented for variable layers, sorry")
if Vc.extruded:
level_ratio = (Vf.mesh().layers - 1) // (Vc.mesh().layers - 1)
else:
level_ratio = 1
coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc]
_, ncell = coarse_to_fine.shape
iterset = Vc.mesh().cell_set
fine_per_cell = Vf.finat_element.space_dimension()
arity = fine_per_cell * ncell
coarse_to_fine_nodes = numpy.full((iterset.total_size, arity*level_ratio), -1, dtype=IntType)
# The DG injection kernel skips the padded slots of a row, but PyOP2
# still reads through them. Fill each one with the row's first child.
children = coarse_to_fine[:iterset.size, :]
children = numpy.where(children >= 0, children, children[:, :1])
values = Vf.cell_node_map().values[children]
if Vc.extruded:
# Keep the layers of each child together, so that the children of
# a coarse cell come before its padded slots.
values = values[:, :, None, :] + numpy.arange(level_ratio)[:, None] * Vf.offset
coarse_to_fine_nodes[:iterset.size, :] = values.reshape(iterset.size, -1)
offset = Vf.offset
if offset is not None:
offset = numpy.tile(offset*level_ratio, ncell*level_ratio)
return cache.setdefault(key, op2.Map(iterset, Vf.node_set,
arity=arity*level_ratio, values=coarse_to_fine_nodes,
offset=offset))
[docs]
def coarse_cell_child_count(
Vc: firedrake.functionspaceimpl.WithGeometry,
Vf: firedrake.functionspaceimpl.WithGeometry,
) -> op2.Dat:
"""Count the fine cells that each coarse cell was refined into.
A row of `HierarchyBase.coarse_to_fine_cells` is as wide as the busiest
coarse cell's count, so its width overstates how many children most cells
have. The DG injection kernel reads this count to stop at a coarse cell's
own children, and so leaves the padding alone.
Parameters
----------
Vc : firedrake.functionspaceimpl.WithGeometry
The coarse function space.
Vf : firedrake.functionspaceimpl.WithGeometry
The fine function space, on the next level of the same hierarchy.
Returns
-------
pyop2.types.dat.Dat
One count per cell of ``Vc``'s mesh, over that mesh's cell set. Halo
cells are left at zero: a par_loop visits the core and owned parts
only, so the kernel never reads them.
"""
mesh = Vc.mesh()
assert hasattr(mesh, "_shared_data_cache")
hierarchyf, levelf = get_level(Vf.mesh())
hierarchyc, levelc = get_level(Vc.mesh())
if hierarchyc != hierarchyf:
raise ValueError("Can't map across hierarchies")
hierarchy = hierarchyf
increment = Fraction(1, hierarchyf.refinements_per_level)
if levelc + increment != levelf:
raise ValueError(f"Can't map between level {levelc} and level {levelf}")
key = (levelc, Vc.extruded and (Vf.mesh().layers, Vc.mesh().layers))
cache = mesh._shared_data_cache["hierarchy_coarse_cell_child_count"]
try:
return cache[key]
except KeyError:
if Vc.extruded:
level_ratio = (Vf.mesh().layers - 1) // (Vc.mesh().layers - 1)
else:
level_ratio = 1
coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc]
iterset = mesh.cell_set
counts = numpy.zeros(iterset.total_size, dtype=IntType)
# Each child of a coarse cell becomes level_ratio cells once extruded.
counts[:iterset.size] = (coarse_to_fine[:iterset.size] >= 0).sum(axis=1) * level_ratio
# A count belongs to a base cell, and every layer of that cell shares
# it. An ExtrudedSet holds no data of its own, so hang the counts off
# the base set that it was built on.
dset = op2.DataSet(iterset.parent if Vc.extruded else iterset, 1)
return cache.setdefault(key, op2.Dat(dset, counts, dtype=IntType))
[docs]
def physical_node_locations(V):
element = V.ufl_element()
if V.value_shape:
assert isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement))
element = element.sub_elements[0]
mesh = V.mesh()
# This is a defaultdict, so the first time we access the key we
# get a fresh dict for the cache.
cache = mesh.geometric_shared_data_cache["hierarchy_physical_node_locations"]
key = (element, V.boundary_set)
try:
return cache[key]
except KeyError:
Vc = V.collapse().reconstruct(element=finat.ufl.VectorElement(element, dim=mesh.geometric_dimension))
# FIXME: This is unsafe for DG coordinates and CG target spaces.
locations = firedrake.assemble(firedrake.interpolate(firedrake.SpatialCoordinate(mesh), Vc))
return cache.setdefault(key, locations)
[docs]
def set_level(obj, hierarchy, level):
"""Attach hierarchy and level info to an object."""
setattr(obj.topological, "__level_info__", (hierarchy, level))
return obj
[docs]
def get_level(obj):
"""Try and obtain hierarchy and level info from an object.
If no level info is available, return ``None, None``."""
try:
return getattr(obj.topological, "__level_info__")
except AttributeError:
return None, None
[docs]
def has_level(obj):
"""Does the provided object have level info?"""
return hasattr(obj.topological, "__level_info__")
def _cache_key(Vc, Vf, needs_coarse_entity_dofs=True):
"""Construct a cache key for node maps"""
_, levelf = get_level(Vf.mesh())
_, levelc = get_level(Vc.mesh())
if needs_coarse_entity_dofs:
key = entity_dofs_key(Vc.finat_element.entity_dofs())
else:
key = ()
key += entity_dofs_key(Vf.finat_element.entity_dofs())
key += (levelc, levelf)
key += (Vc.boundary_set, Vf.boundary_set)
return key
