def coarse_to_fine_node_map(coarse, fine):
if len(coarse) > 1:
assert len(fine) == len(coarse)
return op2.MixedMap(coarse_to_fine_node_map(c, f) for c, f in zip(coarse, fine))
mesh = coarse.mesh()
assert hasattr(mesh, "_shared_data_cache")
if not (coarse.ufl_element() == fine.ufl_element()):
raise ValueError("Can't transfer between different spaces")
ch, level = get_level(mesh)
fh, fine_level = get_level(fine.mesh())
if ch is not fh:
raise ValueError("Can't map between different hierarchies")
refinements_per_level = ch.refinements_per_level
if refinements_per_level*level + 1 != refinements_per_level*fine_level:
raise ValueError("Can't map between level %s and level %s" % (level, fine_level))
c2f, vperm = ch._cells_vperm[int(level*refinements_per_level)]
key = entity_dofs_key(coarse.fiat_element.entity_dofs()) + (level, )
cache = mesh._shared_data_cache["hierarchy_cell_node_map"]
try:
return cache[key]
except KeyError:
from .impl import create_cell_node_map
map_vals, offset = create_cell_node_map(coarse, fine, c2f, vperm)
return cache.setdefault(key, op2.Map(op2.LocalSet(mesh.cell_set),
fine.node_set,
map_vals.shape[1],
map_vals, offset=offset))
def get_restriction_weights(coarse, fine):
mesh = coarse.mesh()
assert hasattr(mesh, "_shared_data_cache")
cache = mesh._shared_data_cache["hierarchy_restriction_weights"]
key = entity_dofs_key(coarse.fiat_element.entity_dofs())
try:
return cache[key]
except KeyError:
# We hit each fine dof more than once since we loop
# elementwise over the coarse cells. So we need a count of
# how many times we did this to weight the final contribution
# appropriately.
if not (coarse.ufl_element() == fine.ufl_element()):
raise ValueError("Can't transfer between different spaces")
if coarse.fiat_element.entity_dofs() == coarse.fiat_element.entity_closure_dofs():
return cache.setdefault(key, None)
ele = coarse.ufl_element()
if isinstance(ele, ufl.VectorElement):
ele = ele.sub_elements()[0]
weights = firedrake.Function(firedrake.FunctionSpace(fine.mesh(), ele))
else:
weights = firedrake.Function(fine)
c2f_map = coarse_to_fine_node_map(coarse, fine)
kernel = get_count_kernel(c2f_map.arity)
op2.par_loop(kernel, op2.LocalSet(mesh.cell_set),
weights.dat(op2.INC, c2f_map[op2.i[0]]))
weights.assign(1/weights)
return cache.setdefault(key, weights)
def __init__(self, mesh_hierarchy, fses):
"""
Build a hierarchy of function spaces
:arg mesh_hierarchy: a :class:`~.MeshHierarchy` on which to
build the function spaces.
:arg fses: an iterable of :class:`~.FunctionSpace`\s.
"""
self._mesh_hierarchy = mesh_hierarchy
self._hierarchy = tuple(
[set_level(fs, self, lvl) for lvl, fs in enumerate(fses)])
self._map_cache = {}
self._cell_sets = tuple(
op2.LocalSet(m.cell_set) for m in self._mesh_hierarchy)
self._ufl_element = self[0].ufl_element()
self._restriction_weights = None
fiat_element = fses[0].fiat_element
ncelldof = len(
fiat_element.entity_dofs()[fses[0].mesh().cell_dimension()][0])
self._discontinuous = ncelldof == fses[0].cell_node_map().arity
try:
element = self[0].fiat_element
omap = self[1].cell_node_map().values
c2f, vperm = self._mesh_hierarchy._cells_vperm[0]
indices, _ = utils.get_unique_indices(element,
omap[c2f[0, :],
...].reshape(-1),
vperm[0, :],
offset=None)
self._prolong_kernel = utils.get_prolongation_kernel(
element, indices, self.dim)
self._restrict_kernel = utils.get_restriction_kernel(
element, indices, self.dim, no_weights=self._discontinuous)
self._inject_kernel = utils.get_injection_kernel(
element, indices, self.dim)
except:
pass
for V in self:
dm = V._dm
dm.setCoarsen(coarsen)
dm.setRefine(refine)