def inject_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vf.ufl_domain().coordinates
key = (("inject", )
+ Vf.ufl_element().value_shape()
+ entity_dofs_key(Vc.finat_element.entity_dofs())
+ entity_dofs_key(Vf.finat_element.entity_dofs())
+ entity_dofs_key(Vc.mesh().coordinates.function_space().finat_element.entity_dofs())
+ entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1]
if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs():
return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
coordinates = Vf.ufl_domain().coordinates
evaluate_kernel = compile_element(ufl.Coefficient(Vf))
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
Vf_element = create_element(Vf.ufl_element())
kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_inject(double *R, const double *X, const double *f, const double *Xf)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xfi = Xf + i*%(Xf_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xfi);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) {
abort();
}
const double *fi = f + cell*%(f_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
R[i] = 0;
}
pyop2_kernel_evaluate(R, fi, Xref);
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
"tdim": Vc.ufl_domain().topological_dimension(),
"ncandidate": ncandidate,
"Rdim": numpy.prod(Vf_element.value_shape),
"Xf_cell_inc": coords_element.space_dimension(),
"f_cell_inc": Vf_element.space_dimension()
}
return cache.setdefault(key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def inject_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vf.ufl_domain().coordinates
key = (("inject", ) +
Vf.ufl_element().value_shape() +
entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.mesh().coordinates.function_space().finat_element.entity_dofs()) +
entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1]
if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs():
return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
coordinates = Vf.ufl_domain().coordinates
evaluate_kernel = compile_element(ufl.Coefficient(Vf))
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
Vf_element = create_element(Vf.ufl_element())
kernel = """
%(to_reference)s
%(evaluate)s
void inject_kernel(double *R, const double *X, const double *f, const double *Xf)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xfi = Xf + i*%(Xf_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xfi);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) {
abort();
}
const double *fi = f + cell*%(f_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
R[i] = 0;
}
evaluate_kernel(R, fi, Xref);
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
"tdim": Vc.ufl_domain().topological_dimension(),
"ncandidate": ncandidate,
"Rdim": numpy.prod(Vf_element.value_shape),
"Xf_cell_inc": coords_element.space_dimension(),
"f_cell_inc": Vf_element.space_dimension()
}
return cache.setdefault(key, (op2.Kernel(kernel, name="inject_kernel"), False))
def prolong_kernel(expression):
hierarchy, level = utils.get_level(expression.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
key = (("prolong", )
+ expression.ufl_element().value_shape()
+ entity_dofs_key(expression.function_space().finat_element.entity_dofs())
+ entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
coords_element = create_element(coordinates.ufl_element())
args = eval_args[-1].gencode(not_scope=True)
R, coarse = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_prolong(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) abort();
const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
pyop2_kernel_evaluate(%(R)s, coarsei, Xref);
}
""" % {"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"R": R,
"coarse": coarse,
"ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
"Rdim": numpy.prod(element.value_shape),
"inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
"Xc_cell_inc": coords_element.space_dimension(),
"coarse_cell_inc": element.space_dimension(),
"tdim": mesh.topological_dimension()}
return cache.setdefault(key, op2.Kernel(my_kernel, name="pyop2_kernel_prolong"))
def restrict_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vc.ufl_domain().coordinates
key = (("restrict", ) + Vf.ufl_element().value_shape() +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
element = create_element(Vc.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = eval_args[-1].gencode(not_scope=True)
R, fine = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_restrict(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
const double *Ri = %(R)s + cell*%(coarse_cell_inc)d;
pyop2_kernel_evaluate(Ri, %(fine)s, Xref);
break;
}
}
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
"inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
"Xc_cell_inc": coords_element.space_dimension(),
"coarse_cell_inc": element.space_dimension(),
"args": args,
"R": R,
"fine": fine,
"tdim": mesh.topological_dimension()
}
return cache.setdefault(
key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
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 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.finat_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(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.finat_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.finat_element.entity_dofs(
) == coarse.finat_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, c2f_map.iterset,
weights.dat(op2.INC, c2f_map[op2.i[0]]))
weights.assign(1 / weights)
return cache.setdefault(key, weights)
def get_transfer_kernel(coarse, fine, typ=None):
ch, level = get_level(coarse.mesh())
mesh = ch[0]
assert hasattr(mesh, "_shared_data_cache")
key = entity_dofs_key(coarse.fiat_element.entity_dofs()) + (typ, )
cache = mesh._shared_data_cache["hierarchy_transfer_kernel"]
try:
return cache[key]
except KeyError:
if not (coarse.ufl_element() == fine.ufl_element()):
raise ValueError("Can't transfer between different spaces")
ch, level = get_level(coarse.mesh())
fh, fine_level = get_level(fine.mesh())
refinements_per_level = ch.refinements_per_level
assert ch is fh
assert refinements_per_level*level + 1 == refinements_per_level*fine_level
dim = coarse.dim
element = coarse.fiat_element
omap = fine.cell_node_map().values
c2f, vperm = ch._cells_vperm[int(level*refinements_per_level)]
indices, _ = get_unique_indices(element,
omap[c2f[0, :], ...].reshape(-1),
vperm[0, :],
offset=None)
if typ == "prolong":
kernel = get_prolongation_kernel(element, indices, dim)
elif typ == "inject":
kernel = get_injection_kernel(element, indices, dim)
elif typ == "restrict":
discontinuous = element.entity_dofs() == element.entity_closure_dofs()
kernel = get_restriction_kernel(element, indices, dim, no_weights=discontinuous)
else:
raise ValueError("Unknown transfer type '%s'" % typ)
return cache.setdefault(key, kernel)
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 coarse_node_to_fine_node_map(Vc, Vf):
if len(Vf) > 1:
assert len(Vf) == len(Vc)
return op2.MixedMap(
coarse_node_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.ufl_domain())
hierarchyc, levelc = get_level(Vc.ufl_domain())
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 = (entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf))
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)
return cache.setdefault(
key,
op2.Map(Vc.node_set,
Vf.node_set,
coarse_to_fine_nodes.shape[1],
values=coarse_to_fine_nodes))
def fine_node_to_coarse_node_map(Vf, Vc):
if len(Vf) > 1:
assert len(Vf) == len(Vc)
return op2.MixedMap(
fine_node_to_coarse_node_map(f, c) for f, c in zip(Vf, Vc))
mesh = Vf.mesh()
assert hasattr(mesh, "_shared_data_cache")
hierarchyf, levelf = get_level(Vf.ufl_domain())
hierarchyc, levelc = get_level(Vc.ufl_domain())
if hierarchyc != hierarchyf:
raise ValueError("Can't map across hierarchies")
hierarchy = hierarchyf
if levelc + 1 != levelf:
raise ValueError("Can't map between level %s and level %s" %
(levelc, levelf))
key = (entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf))
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 Vc.mesh().layers != Vf.mesh().layers:
raise ValueError(
"Coarse and fine meshes must have same number of layers")
fine_to_coarse = hierarchy.fine_to_coarse_cells[levelc + 1]
fine_to_coarse_nodes = impl.fine_to_coarse_nodes(
Vf, Vc, fine_to_coarse)
return cache.setdefault(
key,
op2.Map(Vf.node_set,
Vc.node_set,
fine_to_coarse_nodes.shape[1],
values=fine_to_coarse_nodes))
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.ufl_domain())
hierarchyc, levelc = get_level(Vc.ufl_domain())
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 = (entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf))
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 and Vc.mesh().layers != Vf.mesh().layers:
raise ValueError(
"Coarse and fine meshes must have same number of layers")
coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc]
_, ncell = coarse_to_fine.shape
iterset = Vc.mesh().cell_set
arity = Vf.finat_element.space_dimension() * ncell
coarse_to_fine_nodes = numpy.full((iterset.total_size, arity),
-1,
dtype=IntType)
values = Vf.cell_node_map().values[coarse_to_fine, :].reshape(
iterset.size, arity)
coarse_to_fine_nodes[:Vc.mesh().cell_set.size, :] = values
offset = Vf.offset
if offset is not None:
offset = numpy.tile(offset, ncell)
return cache.setdefault(
key,
op2.Map(iterset,
Vf.node_set,
arity=arity,
values=coarse_to_fine_nodes,
offset=offset))
def prolong_kernel(expression):
hierarchy, _ = utils.get_level(expression.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
key = (("prolong", ) +
expression.ufl_element().value_shape() +
entity_dofs_key(expression.function_space().finat_element.entity_dofs()) +
entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = ", ".join(a.gencode(not_scope=True) for a in eval_args)
arg_names = ", ".join(a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
void prolong_kernel(%(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
to_reference_coords_kernel(Xref, X, Xc);
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
evaluate_kernel(%(arg_names)s, Xref);
}
""" % {"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"R": arg_names[0],
"Rdim": numpy.prod(element.value_shape),
"arg_names": arg_names,
"tdim": mesh.topological_dimension()}
return cache.setdefault(key, op2.Kernel(my_kernel, name="prolong_kernel"))
def restrict_kernel(Vf, Vc):
hierarchy, _ = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vc.ufl_domain().coordinates
key = (("restrict", ) +
Vf.ufl_element().value_shape() +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf)
to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = ", ".join(a.gencode(not_scope=True) for a in eval_args)
arg_names = ", ".join(a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
void restrict_kernel(%(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
to_reference_coords_kernel(Xref, X, Xc);
evaluate_kernel(%(arg_names)s, Xref);
}
""" % {"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"arg_names": arg_names,
"tdim": mesh.topological_dimension()}
return cache.setdefault(key, op2.Kernel(my_kernel, name="restrict_kernel"))
def inject_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vf.ufl_domain().coordinates
key = (
("inject", ) + Vf.ufl_element().value_shape() +
entity_dofs_key(Vc.finat_element.entity_dofs()) +
entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(
Vc.mesh().coordinates.function_space().finat_element.entity_dofs())
+ entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1]
if Vc.finat_element.entity_dofs(
) == Vc.finat_element.entity_closure_dofs():
return cache.setdefault(
key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
coordinates = Vf.ufl_domain().coordinates
evaluate_kernel = compile_element(ufl.Coefficient(Vf))
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
Vf_element = create_element(Vf.ufl_element())
kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_inject(double *R, const double *X, const double *f, const double *Xf)
{
double Xref[%(tdim)d];
int cell = -1;
int bestcell = -1;
double bestdist = 1e10;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xfi = Xf + i*%(Xf_cell_inc)d;
double celldist = 2*bestdist;
to_reference_coords_kernel(Xref, X, Xfi);
if (%(inside_cell)s) {
cell = i;
break;
}
%(compute_celldist)s
if (celldist < bestdist) {
bestdist = celldist;
bestcell = i;
}
}
if (cell == -1) {
/* We didn't find a cell that contained this point exactly.
Did we find one that was close enough? */
if (bestdist < 10) {
cell = bestcell;
} else {
fprintf(stderr, "Could not identify cell in transfer operator. Point: ");
for (int coord = 0; coord < %(spacedim)s; coord++) {
fprintf(stderr, "%%.14e ", X[coord]);
}
fprintf(stderr, "\\n");
fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist);
abort();
}
}
const double *fi = f + cell*%(f_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
R[i] = 0;
}
pyop2_kernel_evaluate(R, fi, Xref);
}
""" % {
"to_reference":
str(to_reference_kernel),
"evaluate":
str(evaluate_kernel),
"inside_cell":
inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
"spacedim":
Vc.finat_element.cell.get_spatial_dimension(),
"compute_celldist":
compute_celldist(
Vc.finat_element.cell, X="Xref", celldist="celldist"),
"tdim":
Vc.ufl_domain().topological_dimension(),
"ncandidate":
ncandidate,
"Rdim":
numpy.prod(Vf_element.value_shape),
"Xf_cell_inc":
coords_element.space_dimension(),
"f_cell_inc":
Vf_element.space_dimension()
}
return cache.setdefault(
key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def restrict_kernel(Vf, Vc):
hierarchy, level = utils.get_level(Vc.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = Vc.ufl_domain().coordinates
key = (("restrict", ) + Vf.ufl_element().value_shape() +
entity_dofs_key(Vf.finat_element.entity_dofs()) +
entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
coords_element = create_element(coordinates.ufl_element())
element = create_element(Vc.ufl_element())
eval_args = evaluate_kernel.args[:-1]
args = eval_args[-1].gencode(not_scope=True)
R, fine = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_restrict(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
int bestcell = -1;
double bestdist = 1e10;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
double celldist = 2*bestdist;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
%(compute_celldist)s
/* fprintf(stderr, "cell %%d celldist: %%.14e\\n", i, celldist);
fprintf(stderr, "Xref: %%.14e %%.14e %%.14e\\n", Xref[0], Xref[1], Xref[2]); */
if (celldist < bestdist) {
bestdist = celldist;
bestcell = i;
}
}
if (cell == -1) {
/* We didn't find a cell that contained this point exactly.
Did we find one that was close enough? */
if (bestdist < 10) {
cell = bestcell;
} else {
fprintf(stderr, "Could not identify cell in transfer operator. Point: ");
for (int coord = 0; coord < %(spacedim)s; coord++) {
fprintf(stderr, "%%.14e ", X[coord]);
}
fprintf(stderr, "\\n");
fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist);
abort();
}
}
{
const double *Ri = %(R)s + cell*%(coarse_cell_inc)d;
pyop2_kernel_evaluate(Ri, %(fine)s, Xref);
}
}
""" % {
"to_reference":
str(to_reference_kernel),
"evaluate":
str(evaluate_kernel),
"ncandidate":
hierarchy.fine_to_coarse_cells[levelf].shape[1],
"inside_cell":
inside_check(element.cell, eps=1e-8, X="Xref"),
"compute_celldist":
compute_celldist(element.cell, X="Xref", celldist="celldist"),
"Xc_cell_inc":
coords_element.space_dimension(),
"coarse_cell_inc":
element.space_dimension(),
"args":
args,
"spacedim":
element.cell.get_spatial_dimension(),
"R":
R,
"fine":
fine,
"tdim":
mesh.topological_dimension()
}
return cache.setdefault(
key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
def prolong_kernel(expression):
hierarchy, level = utils.get_level(expression.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
key = (("prolong", ) + expression.ufl_element().value_shape() +
entity_dofs_key(
expression.function_space().finat_element.entity_dofs()) +
entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
coords_element = create_element(coordinates.ufl_element())
args = eval_args[-1].gencode(not_scope=True)
R, coarse = (a.sym.symbol for a in eval_args)
my_kernel = """
%(to_reference)s
%(evaluate)s
void prolong_kernel(double *R, %(args)s, const double *X, const double *Xc)
{
double Xref[%(tdim)d];
int cell = -1;
for (int i = 0; i < %(ncandidate)d; i++) {
const double *Xci = Xc + i*%(Xc_cell_inc)d;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
}
if (cell == -1) abort();
const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
evaluate_kernel(%(R)s, coarsei, Xref);
}
""" % {
"to_reference": str(to_reference_kernel),
"evaluate": str(evaluate_kernel),
"args": args,
"R": R,
"coarse": coarse,
"ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
"Rdim": numpy.prod(element.value_shape),
"inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
"Xc_cell_inc": coords_element.space_dimension(),
"coarse_cell_inc": element.space_dimension(),
"tdim": mesh.topological_dimension()
}
return cache.setdefault(key,
op2.Kernel(my_kernel, name="prolong_kernel"))
def prolong_kernel(expression):
meshc = expression.ufl_domain()
hierarchy, level = utils.get_level(expression.ufl_domain())
levelf = level + Fraction(1 / hierarchy.refinements_per_level)
cache = hierarchy._shared_data_cache["transfer_kernels"]
coordinates = expression.ufl_domain().coordinates
if meshc.cell_set._extruded:
idx = levelf * hierarchy.refinements_per_level
assert idx == int(idx)
level_ratio = (hierarchy._meshes[int(idx)].layers -
1) // (meshc.layers - 1)
else:
level_ratio = 1
key = (("prolong", level_ratio) + expression.ufl_element().value_shape() +
entity_dofs_key(
expression.function_space().finat_element.entity_dofs()) +
entity_dofs_key(
coordinates.function_space().finat_element.entity_dofs()))
try:
return cache[key]
except KeyError:
mesh = coordinates.ufl_domain()
evaluate_kernel = compile_element(expression)
to_reference_kernel = to_reference_coordinates(
coordinates.ufl_element())
element = create_element(expression.ufl_element())
eval_args = evaluate_kernel.args[:-1]
coords_element = create_element(coordinates.ufl_element())
args = eval_args[-1].gencode(not_scope=True)
R, coarse = (a.sym.symbol for a in eval_args)
my_kernel = """#include <petsc.h>
%(to_reference)s
%(evaluate)s
__attribute__((noinline)) /* Clang bug */
static void pyop2_kernel_prolong(PetscScalar *R, %(args)s, const PetscScalar *X, const PetscScalar *Xc)
{
PetscScalar Xref[%(tdim)d];
int cell = -1;
int bestcell = -1;
double bestdist = 1e10;
for (int i = 0; i < %(ncandidate)d; i++) {
const PetscScalar *Xci = Xc + i*%(Xc_cell_inc)d;
double celldist = 2*bestdist;
to_reference_coords_kernel(Xref, X, Xci);
if (%(inside_cell)s) {
cell = i;
break;
}
%(compute_celldist)s
if (celldist < bestdist) {
bestdist = celldist;
bestcell = i;
}
}
if (cell == -1) {
/* We didn't find a cell that contained this point exactly.
Did we find one that was close enough? */
if (bestdist < 10) {
cell = bestcell;
} else {
fprintf(stderr, "Could not identify cell in transfer operator. Point: ");
for (int coord = 0; coord < %(spacedim)s; coord++) {
fprintf(stderr, "%%.14e ", X[coord]);
}
fprintf(stderr, "\\n");
fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist);
abort();
}
}
const PetscScalar *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d;
for ( int i = 0; i < %(Rdim)d; i++ ) {
%(R)s[i] = 0;
}
pyop2_kernel_evaluate(%(R)s, coarsei, Xref);
}
""" % {
"to_reference":
str(to_reference_kernel),
"evaluate":
str(evaluate_kernel),
"args":
args,
"R":
R,
"spacedim":
element.cell.get_spatial_dimension(),
"coarse":
coarse,
"ncandidate":
hierarchy.fine_to_coarse_cells[levelf].shape[1] * level_ratio,
"Rdim":
numpy.prod(element.value_shape),
"inside_cell":
inside_check(element.cell, eps=1e-8, X="Xref"),
"compute_celldist":
compute_celldist(element.cell, X="Xref", celldist="celldist"),
"Xc_cell_inc":
coords_element.space_dimension(),
"coarse_cell_inc":
element.space_dimension(),
"tdim":
mesh.topological_dimension()
}
return cache.setdefault(
key, op2.Kernel(my_kernel, name="pyop2_kernel_prolong"))
