def _interpolator(V, dat, expr, subset):
to_element = create_element(V.ufl_element(), vector_is_mixed=False)
to_pts = []
if V.ufl_element().mapping() != "identity":
raise NotImplementedError(
"Can only interpolate onto elements "
"with affine mapping. Try projecting instead")
for dual in to_element.dual_basis():
if not isinstance(dual, FIAT.functional.PointEvaluation):
raise NotImplementedError(
"Can only interpolate onto point "
"evaluation operators. Try projecting instead")
pts, = dual.pt_dict.keys()
to_pts.append(pts)
if len(expr.ufl_shape) != len(V.ufl_element().value_shape()):
raise RuntimeError(
'Rank mismatch: Expression rank %d, FunctionSpace rank %d' %
(len(expr.ufl_shape), len(V.ufl_element().value_shape())))
if expr.ufl_shape != V.ufl_element().value_shape():
raise RuntimeError(
'Shape mismatch: Expression shape %r, FunctionSpace shape %r' %
(expr.ufl_shape, V.ufl_element().value_shape()))
mesh = V.ufl_domain()
coords = mesh.coordinates
if not isinstance(expr, (firedrake.Expression, SubExpression)):
if expr.ufl_domain() and expr.ufl_domain() != V.mesh():
raise NotImplementedError(
"Interpolation onto another mesh not supported.")
if expr.ufl_shape != V.shape:
raise ValueError(
"UFL expression has incorrect shape for interpolation.")
ast, oriented, coefficients = compile_ufl_kernel(expr, to_pts, coords)
kernel = op2.Kernel(ast, ast.name)
indexed = True
elif hasattr(expr, "eval"):
kernel, oriented, coefficients = compile_python_kernel(
expr, to_pts, to_element, V, coords)
indexed = False
elif expr.code is not None:
kernel, oriented, coefficients = compile_c_kernel(
expr, to_pts, to_element, V, coords)
indexed = True
else:
raise RuntimeError(
"Attempting to evaluate an Expression which has no value.")
cell_set = coords.cell_set
if subset is not None:
assert subset.superset == cell_set
cell_set = subset
args = [kernel, cell_set]
copy_back = False
if dat in set((c.dat for c in coefficients)):
output = dat
dat = op2.Dat(dat.dataset)
copy_back = True
if indexed:
args.append(dat(op2.WRITE, V.cell_node_map()[op2.i[0]]))
else:
args.append(dat(op2.WRITE, V.cell_node_map()))
if oriented:
co = mesh.cell_orientations()
args.append(co.dat(op2.READ, co.cell_node_map()[op2.i[0]]))
for coefficient in coefficients:
m_ = coefficient.cell_node_map()
if indexed:
args.append(coefficient.dat(op2.READ, m_ and m_[op2.i[0]]))
else:
args.append(coefficient.dat(op2.READ, m_))
for o in coefficients:
domain = o.ufl_domain()
if domain is not None and domain.topology != mesh.topology:
raise NotImplementedError(
"Interpolation onto another mesh not supported.")
if copy_back:
return partial(op2.par_loop, *args), partial(dat.copy, output)
else:
return (partial(op2.par_loop, *args), )
def get_restriction_kernel(fiat_element,
unique_indices,
dim=1,
no_weights=False):
weights = get_restriction_weights(fiat_element)[unique_indices].T
ncdof = weights.shape[0]
nfdof = weights.shape[1]
arglist = [
ast.Decl("double", ast.Symbol("coarse", (ncdof * dim, ))),
ast.Decl("double",
ast.Symbol("*restrict *restrict fine", ()),
qualifiers=["const"])
]
if not no_weights:
arglist.append(
ast.Decl("double",
ast.Symbol("*restrict *restrict count_weights", ()),
qualifiers=["const"]))
all_ones = np.allclose(weights, 1.0)
if all_ones:
w = []
else:
w_sym = ast.Symbol("weights", (ncdof, nfdof))
init = ast.ArrayInit(format_array_literal(weights))
w = [ast.Decl("double", w_sym, init, qualifiers=["const"])]
i = ast.Symbol("i", ())
j = ast.Symbol("j", ())
k = ast.Symbol("k", ())
fine = ast.Symbol("fine", (j, k))
if no_weights:
if all_ones:
assign = fine
else:
assign = ast.Prod(fine, ast.Symbol("weights", (i, j)))
else:
if all_ones:
assign = ast.Prod(fine, ast.Symbol("count_weights", (j, 0)))
else:
assign = ast.Prod(
fine,
ast.Prod(ast.Symbol("weights", (i, j)),
ast.Symbol("count_weights", (j, 0))))
assignment = ast.Incr(
ast.Symbol("coarse", (ast.Sum(k, ast.Prod(i, ast.c_sym(dim))), )),
assign)
k_loop = ast.For(ast.Decl("int", k, ast.c_sym(0)),
ast.Less(k, ast.c_sym(dim)), ast.Incr(k, ast.c_sym(1)),
ast.Block([assignment], open_scope=True))
j_loop = ast.For(ast.Decl("int", j, ast.c_sym(0)),
ast.Less(j, ast.c_sym(nfdof)), ast.Incr(j, ast.c_sym(1)),
ast.Block([k_loop], open_scope=True))
i_loop = ast.For(ast.Decl("int", i, ast.c_sym(0)),
ast.Less(i, ast.c_sym(ncdof)), ast.Incr(i, ast.c_sym(1)),
ast.Block([j_loop], open_scope=True))
k = ast.FunDecl("void",
"restriction",
arglist,
ast.Block(w + [i_loop]),
pred=["static", "inline"])
return op2.Kernel(k, "restriction", opts=parameters["coffee"])
def compile_ufl_kernel(expression, to_pts, to_element, fs):
import collections
from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks
from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering
from ufl.algorithms.apply_derivatives import apply_derivatives
from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering
from ufl.algorithms import extract_arguments, extract_coefficients
from gem import gem, impero_utils
from tsfc import fem, ufl_utils
from tsfc.coffee import generate as generate_coffee
from tsfc.kernel_interface import (KernelBuilderBase,
needs_cell_orientations,
cell_orientations_coffee_arg)
# Imitate the compute_form_data processing pipeline
#
# Unfortunately, we cannot call compute_form_data here, since
# we only have an expression, not a form
expression = apply_algebra_lowering(expression)
expression = apply_derivatives(expression)
expression = apply_function_pullbacks(expression)
expression = apply_geometry_lowering(expression)
expression = apply_derivatives(expression)
expression = apply_geometry_lowering(expression)
expression = apply_derivatives(expression)
# Replace coordinates (if any)
if expression.ufl_domain():
assert fs.mesh() == expression.ufl_domain()
expression = ufl_utils.replace_coordinates(expression, fs.mesh().coordinates)
if extract_arguments(expression):
return ValueError("Cannot interpolate UFL expression with Arguments!")
builder = KernelBuilderBase()
args = []
coefficients = extract_coefficients(expression)
for i, coefficient in enumerate(coefficients):
args.append(builder.coefficient(coefficient, "w_%d" % i))
point_index = gem.Index(name='p')
ir = fem.process('cell', fs.mesh().ufl_cell(), to_pts, None,
point_index, (), expression,
builder.coefficient_mapper,
collections.defaultdict(gem.Index))
assert len(ir) == 1
# Deal with non-scalar expressions
tensor_indices = ()
if fs.shape:
tensor_indices = tuple(gem.Index() for s in fs.shape)
ir = [gem.Indexed(ir[0], tensor_indices)]
# Build kernel body
return_var = gem.Variable('A', (len(to_pts),) + fs.shape)
return_expr = gem.Indexed(return_var, (point_index,) + tensor_indices)
impero_c = impero_utils.compile_gem([return_expr], ir, [point_index])
body = generate_coffee(impero_c, index_names={point_index: 'p'})
oriented = needs_cell_orientations(ir)
if oriented:
args.insert(0, cell_orientations_coffee_arg)
# Build kernel
args.insert(0, ast.Decl("double", ast.Symbol('A', rank=(len(to_pts),) + fs.shape)))
kernel_code = builder.construct_kernel("expression_kernel", args, body)
return op2.Kernel(kernel_code, kernel_code.name), oriented, coefficients
def test_global_inc_init_not_zero(self, backend, elems, g):
"""Increment a global initialized with a non-zero value."""
k = """void k(unsigned int* inc) { (*inc) += 1; }"""
g.data[0] = 10
op2.par_loop(op2.Kernel(k, 'k'), elems, g(op2.INC))
assert g.data[0] == elems.size + 10
def test_mismatch_set_raises_error(self, backend, elems, x):
"""The iterset of the parloop should match the dataset of the direct dat."""
kernel_wo = """void kernel_wo(unsigned int* x) { *x = 42; }"""
with pytest.raises(MapValueError):
op2.par_loop(op2.Kernel(kernel_wo, "kernel_wo"),
op2.Set(elems.size), x(op2.WRITE))
def test_direct_loop_inc(self, backend, xtr_nodes):
dat = op2.Dat(xtr_nodes)
k = 'void k(double *x) { *x += 1.0; }'
dat.data[:] = 0
op2.par_loop(op2.Kernel(k, 'k'), dat.dataset.set, dat(op2.INC))
assert numpy.allclose(dat.data[:], 1.0)
def rhs():
kernel_code = FlatBlock("""
void rhs(double** localTensor, double* c0[2], double* c1[1])
{
double CG1[3][6] = { { 0.09157621, 0.09157621, 0.81684757,
0.44594849, 0.44594849, 0.10810302 },
{ 0.09157621, 0.81684757, 0.09157621,
0.44594849, 0.10810302, 0.44594849 },
{ 0.81684757, 0.09157621, 0.09157621,
0.10810302, 0.44594849, 0.44594849 } };
double d_CG1[3][6][2] = { { { 1., 0. },
{ 1., 0. },
{ 1., 0. },
{ 1., 0. },
{ 1., 0. },
{ 1., 0. } },
{ { 0., 1. },
{ 0., 1. },
{ 0., 1. },
{ 0., 1. },
{ 0., 1. },
{ 0., 1. } },
{ { -1.,-1. },
{ -1.,-1. },
{ -1.,-1. },
{ -1.,-1. },
{ -1.,-1. },
{ -1.,-1. } } };
double w[6] = { 0.05497587, 0.05497587, 0.05497587, 0.11169079,
0.11169079, 0.11169079 };
double c_q1[6];
double c_q0[6][2][2];
for(int i_g = 0; i_g < 6; i_g++)
{
c_q1[i_g] = 0.0;
for(int q_r_0 = 0; q_r_0 < 3; q_r_0++)
{
c_q1[i_g] += c1[q_r_0][0] * CG1[q_r_0][i_g];
};
for(int i_d_0 = 0; i_d_0 < 2; i_d_0++)
{
for(int i_d_1 = 0; i_d_1 < 2; i_d_1++)
{
c_q0[i_g][i_d_0][i_d_1] = 0.0;
for(int q_r_0 = 0; q_r_0 < 3; q_r_0++)
{
c_q0[i_g][i_d_0][i_d_1] += c0[q_r_0][i_d_0] * d_CG1[q_r_0][i_g][i_d_1];
};
};
};
};
for(int i_r_0 = 0; i_r_0 < 3; i_r_0++)
{
for(int i_g = 0; i_g < 6; i_g++)
{
double ST1 = 0.0;
ST1 += CG1[i_r_0][i_g] * c_q1[i_g] * (c_q0[i_g][0][0] * c_q0[i_g][1][1] + -1 * c_q0[i_g][0][1] * c_q0[i_g][1][0]);
localTensor[i_r_0][0] += ST1 * w[i_g];
};
};
}""")
return op2.Kernel(kernel_code, "rhs")
def k1_write_to_dat(cls):
k = """
void k(unsigned int *x, unsigned int *g) { *x = *g; }
"""
return op2.Kernel(k, "k")
def k1_inc_to_global(cls):
k = """
void k(unsigned int *x, unsigned int *g) { *g += *x; }
"""
return op2.Kernel(k, "k")
def k2_write_to_dat(cls, request):
k = """
void k(unsigned int *x, unsigned int *g) { *x = g[0] + g[1]; }
"""
return op2.Kernel(k, "k")
def k2_inc_to_global(cls):
k = """
void k(unsigned int *x, unsigned int *g) { g[0] += x[0]; g[1] += x[1]; }
"""
return op2.Kernel(k, "k")
def init_cell_orientations(self, expr):
"""Compute and initialise :attr:`cell_orientations` relative to a specified orientation.
:arg expr: an :class:`.Expression` evaluated to produce a
reference normal direction.
"""
import firedrake.function as function
import firedrake.functionspace as functionspace
if expr.value_shape()[0] != 3:
raise NotImplementedError('Only implemented for 3-vectors')
if self.ufl_cell() not in (ufl.Cell('triangle', 3), ufl.Cell("quadrilateral", 3), ufl.OuterProductCell(ufl.Cell('interval'), ufl.Cell('interval'), gdim=3)):
raise NotImplementedError('Only implemented for triangles and quadrilaterals embedded in 3d')
if hasattr(self.topology, '_cell_orientations'):
raise RuntimeError("init_cell_orientations already called, did you mean to do so again?")
v0 = lambda x: ast.Symbol("v0", (x,))
v1 = lambda x: ast.Symbol("v1", (x,))
n = lambda x: ast.Symbol("n", (x,))
x = lambda x: ast.Symbol("x", (x,))
coords = lambda x, y: ast.Symbol("coords", (x, y))
body = []
body += [ast.Decl("double", v(3)) for v in [v0, v1, n, x]]
body.append(ast.Decl("double", "dot"))
body.append(ast.Assign("dot", 0.0))
body.append(ast.Decl("int", "i"))
# if triangle, use v0 = x1 - x0, v1 = x2 - x0
# otherwise, for the various quads, use v0 = x2 - x0, v1 = x1 - x0
# recall reference element ordering:
# triangle: 2 quad: 1 3
# 0 1 0 2
if self.ufl_cell() == ufl.Cell('triangle', 3):
body.append(ast.For(ast.Assign("i", 0), ast.Less("i", 3), ast.Incr("i", 1),
[ast.Assign(v0("i"), ast.Sub(coords(1, "i"), coords(0, "i"))),
ast.Assign(v1("i"), ast.Sub(coords(2, "i"), coords(0, "i"))),
ast.Assign(x("i"), 0.0)]))
else:
body.append(ast.For(ast.Assign("i", 0), ast.Less("i", 3), ast.Incr("i", 1),
[ast.Assign(v0("i"), ast.Sub(coords(2, "i"), coords(0, "i"))),
ast.Assign(v1("i"), ast.Sub(coords(1, "i"), coords(0, "i"))),
ast.Assign(x("i"), 0.0)]))
# n = v0 x v1
body.append(ast.Assign(n(0), ast.Sub(ast.Prod(v0(1), v1(2)), ast.Prod(v0(2), v1(1)))))
body.append(ast.Assign(n(1), ast.Sub(ast.Prod(v0(2), v1(0)), ast.Prod(v0(0), v1(2)))))
body.append(ast.Assign(n(2), ast.Sub(ast.Prod(v0(0), v1(1)), ast.Prod(v0(1), v1(0)))))
body.append(ast.For(ast.Assign("i", 0), ast.Less("i", 3), ast.Incr("i", 1),
[ast.Incr(x(j), coords("i", j)) for j in range(3)]))
body.extend([ast.FlatBlock("dot += (%(x)s) * n[%(i)d];\n" % {"x": x_, "i": i})
for i, x_ in enumerate(expr.code)])
body.append(ast.Assign("orientation[0][0]", ast.Ternary(ast.Less("dot", 0), 1, 0)))
kernel = op2.Kernel(ast.FunDecl("void", "cell_orientations",
[ast.Decl("int**", "orientation"),
ast.Decl("double**", "coords")],
ast.Block(body)),
"cell_orientations")
# Build the cell orientations as a DG0 field (so that we can
# pass it in for facet integrals and the like)
fs = functionspace.FunctionSpace(self, 'DG', 0)
cell_orientations = function.Function(fs, name="cell_orientations", dtype=np.int32)
op2.par_loop(kernel, self.cell_set,
cell_orientations.dat(op2.WRITE, cell_orientations.cell_node_map()),
self.coordinates.dat(op2.READ, self.coordinates.cell_node_map()))
self.topology._cell_orientations = cell_orientations
def exterior_facet_boundary_node_map(self, V, method):
"""Return the :class:`pyop2.Map` from exterior facets to nodes
on the boundary.
:arg V: The function space.
:arg method: The method for determining boundary nodes. See
:class:`~.DirichletBC` for details.
"""
try:
return self.map_caches["boundary_node"][method]
except KeyError:
pass
el = V.finat_element
dim = self.mesh.facet_dimension()
if method == "topological":
boundary_dofs = el.entity_closure_dofs()[dim]
elif method == "geometric":
# This function is only called on extruded meshes when
# asking for the nodes that live on the "vertical"
# exterior facets.
boundary_dofs = entity_support_dofs(el, dim)
nodes_per_facet = \
len(boundary_dofs[0])
# HACK ALERT
# The facet set does not have a halo associated with it, since
# we only construct halos for DoF sets. Fortunately, this
# loop is direct and we already have all the correct
# information available locally. So We fake a set of the
# correct size and carry out a direct loop
facet_set = op2.Set(self.mesh.exterior_facets.set.total_size,
comm=self.mesh.comm)
fs_dat = op2.Dat(
facet_set**el.space_dimension(),
data=V.exterior_facet_node_map().values_with_halo.view())
facet_dat = op2.Dat(facet_set**nodes_per_facet, dtype=numpy.int32)
# Ensure these come out in sorted order.
local_facet_nodes = numpy.array(
[boundary_dofs[e] for e in sorted(boundary_dofs.keys())])
# Helper function to turn the inner index of an array into c
# array literals.
c_array = lambda xs: "{" + ", ".join(map(str, xs)) + "}"
# AST for: l_nodes[facet[0]][n]
rank_ast = ast.Symbol("l_nodes",
rank=(ast.Symbol("facet", rank=(0, )), "n"))
body = ast.Block([
ast.Decl("int",
ast.Symbol("l_nodes",
(len(el.cell.topology[dim]), nodes_per_facet)),
init=ast.ArrayInit(
c_array(map(c_array, local_facet_nodes))),
qualifiers=["const"]),
ast.For(
ast.Decl("int", "n", 0), ast.Less("n", nodes_per_facet),
ast.Incr("n", 1),
ast.Assign(ast.Symbol("facet_nodes", ("n", )),
ast.Symbol("cell_nodes", (rank_ast, ))))
])
kernel = op2.Kernel(
ast.FunDecl("void", "create_bc_node_map", [
ast.Decl("int*", "cell_nodes"),
ast.Decl("int*", "facet_nodes"),
ast.Decl("unsigned int*", "facet")
], body), "create_bc_node_map")
local_facet_dat = op2.Dat(
facet_set**self.mesh.exterior_facets._rank,
self.mesh.exterior_facets.local_facet_dat.data_ro_with_halos,
dtype=numpy.uintc)
op2.par_loop(kernel, facet_set, fs_dat(op2.READ), facet_dat(op2.WRITE),
local_facet_dat(op2.READ))
if self.extruded:
offset = self.offset[boundary_dofs[0]]
else:
offset = None
val = op2.Map(facet_set,
self.node_set,
nodes_per_facet,
facet_dat.data_ro_with_halos,
name="exterior_facet_boundary_node",
offset=offset)
self.map_caches["boundary_node"][method] = val
return val
def compile_c_kernel(expression, to_pts, to_element, fs, coords):
"""Produce a :class:`PyOP2.Kernel` from the c expression provided."""
coords_space = coords.function_space()
coords_element = create_element(coords_space.ufl_element(),
vector_is_mixed=False)
names = {v[0] for v in expression._user_args}
X = list(coords_element.tabulate(0, to_pts).values())[0]
# Produce C array notation of X.
X_str = "{{" + "},\n{".join([",".join(map(str, x)) for x in X.T]) + "}}"
A = utils.unique_name("A", names)
X = utils.unique_name("X", names)
x_ = utils.unique_name("x_", names)
k = utils.unique_name("k", names)
d = utils.unique_name("d", names)
i_ = utils.unique_name("i", names)
# x is a reserved name.
x = "x"
if "x" in names:
raise ValueError(
"cannot use 'x' as a user-defined Expression variable")
ass_exp = [
ast.Assign(ast.Symbol(A, (k, ), ((len(expression.code), i), )),
ast.FlatBlock("%s" % code))
for i, code in enumerate(expression.code)
]
dim = coords_space.value_size
ndof = to_element.space_dimension()
xndof = coords_element.space_dimension()
nfdof = to_element.space_dimension() * numpy.prod(fs.value_size, dtype=int)
init_X = ast.Decl(typ="double",
sym=ast.Symbol(X, rank=(ndof, xndof)),
qualifiers=["const"],
init=X_str)
init_x = ast.Decl(typ="double",
sym=ast.Symbol(x, rank=(coords_space.value_size, )))
init_pi = ast.Decl(typ="double",
sym="pi",
qualifiers=["const"],
init="3.141592653589793")
init = ast.Block([init_X, init_x, init_pi])
incr_x = ast.Incr(
ast.Symbol(x, rank=(d, )),
ast.Prod(ast.Symbol(X, rank=(k, i_)),
ast.Symbol(x_, rank=(ast.Sum(ast.Prod(i_, dim), d), ))))
assign_x = ast.Assign(ast.Symbol(x, rank=(d, )), 0)
loop_x = ast.For(init=ast.Decl("unsigned int", i_, 0),
cond=ast.Less(i_, xndof),
incr=ast.Incr(i_, 1),
body=[incr_x])
block = ast.For(init=ast.Decl("unsigned int", d, 0),
cond=ast.Less(d, dim),
incr=ast.Incr(d, 1),
body=[assign_x, loop_x])
loop = ast.c_for(k, ndof, ast.Block([block] + ass_exp, open_scope=True))
user_args = []
user_init = []
for _, arg in expression._user_args:
if arg.shape == (1, ):
user_args.append(ast.Decl("double *", "%s_" % arg.name))
user_init.append(
ast.FlatBlock("const double %s = *%s_;" %
(arg.name, arg.name)))
else:
user_args.append(ast.Decl("double *", arg.name))
kernel_code = ast.FunDecl(
"void", "expression_kernel", [
ast.Decl("double", ast.Symbol(A, (nfdof, ))),
ast.Decl("double*", x_)
] + user_args, ast.Block(user_init + [init, loop], open_scope=False))
coefficients = [coords]
for _, arg in expression._user_args:
coefficients.append(GlobalWrapper(arg))
return op2.Kernel(kernel_code,
kernel_code.name), False, tuple(coefficients)
def run(diffusivity, current_time, dt, endtime, **kwargs):
op2.init(**kwargs)
# Set up finite element problem
T = FiniteElement("Lagrange", "triangle", 1)
V = VectorElement("Lagrange", "triangle", 1)
p = TrialFunction(T)
q = TestFunction(T)
t = Coefficient(T)
u = Coefficient(V)
diffusivity = 0.1
M = p * q * dx
adv_rhs = (q * t + dt * dot(grad(q), u) * t) * dx
d = -dt * diffusivity * dot(grad(q), grad(p)) * dx
diff_matrix = M - 0.5 * d
diff_rhs = action(M + 0.5 * d, t)
# Generate code for mass and rhs assembly.
mass = compile_form(M, "mass")[0]
adv_rhs = compile_form(adv_rhs, "adv_rhs")[0]
diff_matrix = compile_form(diff_matrix, "diff_matrix")[0]
diff_rhs = compile_form(diff_rhs, "diff_rhs")[0]
# Set up simulation data structures
valuetype = np.float64
nodes, coords, elements, elem_node = read_triangle(kwargs['mesh'])
num_nodes = nodes.size
sparsity = op2.Sparsity((elem_node, elem_node), 1, "sparsity")
mat = op2.Mat(sparsity, valuetype, "mat")
tracer_vals = np.asarray([0.0] * num_nodes, dtype=valuetype)
tracer = op2.Dat(nodes, 1, tracer_vals, valuetype, "tracer")
b_vals = np.asarray([0.0] * num_nodes, dtype=valuetype)
b = op2.Dat(nodes, 1, b_vals, valuetype, "b")
velocity_vals = np.asarray([1.0, 0.0] * num_nodes, dtype=valuetype)
velocity = op2.Dat(nodes, 2, velocity_vals, valuetype, "velocity")
# Set initial condition
i_cond_code = """
void i_cond(double *c, double *t)
{
double i_t = 0.01; // Initial time
double A = 0.1; // Normalisation
double D = 0.1; // Diffusivity
double pi = 3.141459265358979;
double x = c[0]-0.5;
double y = c[1]-0.5;
double r = sqrt(x*x+y*y);
if (r<0.25)
*t = A*(exp((-(r*r))/(4*D*i_t))/(4*pi*D*i_t));
else
*t = 0.0;
}
"""
i_cond = op2.Kernel(i_cond_code, "i_cond")
op2.par_loop(i_cond, nodes, coords(op2.IdentityMap, op2.READ),
tracer(op2.IdentityMap, op2.WRITE))
zero_dat_code = """
void zero_dat(double *dat)
{
*dat = 0.0;
}
"""
zero_dat = op2.Kernel(zero_dat_code, "zero_dat")
# Assemble and solve
have_advection = True
have_diffusion = True
def timestep_iteration():
# Advection
if have_advection:
tic('advection')
tic('assembly')
mat.zero()
op2.par_loop(
mass, elements(3, 3),
mat((elem_node[op2.i[0]], elem_node[op2.i[1]]), op2.INC),
coords(elem_node, op2.READ))
op2.par_loop(zero_dat, nodes, b(op2.IdentityMap, op2.WRITE))
op2.par_loop(adv_rhs, elements(3), b(elem_node[op2.i[0]], op2.INC),
coords(elem_node, op2.READ),
tracer(elem_node, op2.READ),
velocity(elem_node, op2.READ))
toc('assembly')
tic('solve')
op2.solve(mat, tracer, b)
toc('solve')
toc('advection')
# Diffusion
if have_diffusion:
tic('diffusion')
tic('assembly')
mat.zero()
op2.par_loop(
diff_matrix, elements(3, 3),
mat((elem_node[op2.i[0]], elem_node[op2.i[1]]), op2.INC),
coords(elem_node, op2.READ))
op2.par_loop(zero_dat, nodes, b(op2.IdentityMap, op2.WRITE))
op2.par_loop(diff_rhs, elements(3), b(elem_node[op2.i[0]],
op2.INC),
coords(elem_node, op2.READ),
tracer(elem_node, op2.READ))
toc('assembly')
tic('solve')
op2.solve(mat, tracer, b)
toc('solve')
toc('diffusion')
# Perform 1 iteration to warm up plan cache then reset initial condition
timestep_iteration()
op2.par_loop(i_cond, nodes, coords(op2.IdentityMap, op2.READ),
tracer(op2.IdentityMap, op2.WRITE))
reset()
# Timed iteration
t1 = clock()
while current_time < endtime:
timestep_iteration()
current_time += dt
runtime = clock() - t1
print "/fluidity :: %f" % runtime
summary('profile_pyop2_%s_%s.csv' %
(opt['mesh'].split('/')[-1], opt['backend']))
def k1_min_to_global(cls):
k = """
void k(unsigned int *x, unsigned int *g) { if (*x < *g) *g = *x; }
"""
return op2.Kernel(k, "k")
def make_blas_kernels(self, Vf, Vc):
"""
Interpolation and restriction kernels between CG/DG
tensor product spaces on quads and hexes.
Works by tabulating the coarse 1D Lagrange basis
functions in the (Nf+1)-by-(Nc+1) matrix Jhat,
and using the fact that the 2d/3d tabulation is the
tensor product J = kron(Jhat, kron(Jhat, Jhat))
"""
ndim = Vf.ufl_domain().topological_dimension()
zf = self.get_nodes_1d(Vf)
zc = self.get_nodes_1d(Vc)
Jnp = self.barycentric(zc, zf)
# Declare array shapes to be used as literals inside the kernels, I follow to the m-by-n convention.
(mx, nx) = Jnp.shape
(my, ny) = (mx, nx) if ndim >= 2 else (1, 1)
(mz, nz) = (mx, nx) if ndim >= 3 else (1, 1)
nscal = Vf.value_size # number of components
mxyz = mx * my * mz # dim of Vf scalar element
nxyz = nx * ny * nz # dim of Vc scalar element
lwork = nscal * max(mx, nx) * max(my, ny) * max(
mz, nz) # size for work arrays
Jhat = np.ascontiguousarray(Jnp.T) # BLAS uses FORTRAN ordering
# Pass the 1D tabulation as hexadecimal string
JX = ', '.join(map(float.hex, np.asarray(Jhat).flatten()))
# The Kronecker product routines assume 3D shapes, so in 2D we pass one instead of Jhat
JY = "JX" if ndim >= 2 else "&one"
JZ = "JX" if ndim >= 3 else "&one"
# Common kernel to compute y = kron(A3, kron(A2, A1)) * x
# Vector and tensor field genearalization from Deville, Fischer, and Mund section 8.3.1.
kronmxv_code = """
#include <petscsys.h>
#include <petscblaslapack.h>
static void kronmxv(int tflag,
PetscBLASInt mx, PetscBLASInt my, PetscBLASInt mz,
PetscBLASInt nx, PetscBLASInt ny, PetscBLASInt nz, PetscBLASInt nel,
PetscScalar *A1, PetscScalar *A2, PetscScalar *A3,
PetscScalar *x , PetscScalar *y){
/*
Kronecker matrix-vector product
y = op(A) * x, A = kron(A3, kron(A2, A1))
where:
op(A) = transpose(A) if tflag>0 else A
op(A1) is mx-by-nx,
op(A2) is my-by-ny,
op(A3) is mz-by-nz,
x is (nx*ny*nz)-by-nel,
y is (mx*my*mz)-by-nel.
Important notes:
The input data in x is destroyed in the process.
Need to allocate nel*max(mx, nx)*max(my, ny)*max(mz, nz) memory for both x and y.
*/
PetscBLASInt m,n,k,s,p,lda;
char TA1, TA2, TA3;
char tran='T', notr='N';
PetscScalar zero=0.0E0, one=1.0E0;
if(tflag>0){
TA1 = tran;
TA2 = notr;
}else{
TA1 = notr;
TA2 = tran;
}
TA3 = TA2;
m = mx; k = nx; n = ny*nz*nel;
lda = (tflag>0)? nx : mx;
BLASgemm_(&TA1, ¬r, &m,&n,&k, &one, A1,&lda, x,&k, &zero, y,&m);
p = 0; s = 0;
m = mx; k = ny; n = my;
lda = (tflag>0)? ny : my;
for(PetscBLASInt i=0; i<nz*nel; i++){
BLASgemm_(¬r, &TA2, &m,&n,&k, &one, y+p,&m, A2,&lda, &zero, x+s,&m);
p += m*k;
s += m*n;
}
p = 0; s = 0;
m = mx*my; k = nz; n = mz;
lda = (tflag>0)? nz : mz;
for(PetscBLASInt i=0; i<nel; i++){
BLASgemm_(¬r, &TA3, &m,&n,&k, &one, x+p,&m, A3,&lda, &zero, y+s,&m);
p += m*k;
s += m*n;
}
return;
}
"""
# UFL elements order the component DoFs related to the same node contiguously.
# We transpose before and after the multiplcation times J to have each component
# stored contiguously as a scalar field, thus reducing the number of dgemm calls.
# We could benefit from loop tiling for the transpose, but that makes the code
# more complicated.
prolong_code = f"""
{kronmxv_code}
void prolongation(PetscScalar *restrict y, const PetscScalar *restrict x){{
PetscScalar JX[{mx}*{nx}] = {{ {JX} }};
PetscScalar t0[{lwork}], t1[{lwork}];
PetscScalar one=1.0E0;
for({IntType_c} j=0; j<{nxyz}; j++)
for({IntType_c} i=0; i<{nscal}; i++)
t0[j + {nxyz}*i] = x[i + {nscal}*j];
kronmxv(0, {mx},{my},{mz}, {nx},{ny},{nz}, {nscal}, JX,{JY},{JZ}, t0,t1);
for({IntType_c} j=0; j<{mxyz}; j++)
for({IntType_c} i=0; i<{nscal}; i++)
y[i + {nscal}*j] = t1[j + {mxyz}*i];
return;
}}
"""
restrict_code = f"""
{kronmxv_code}
void restriction(PetscScalar *restrict y, const PetscScalar *restrict x,
const PetscScalar *restrict w){{
PetscScalar JX[{mx}*{nx}] = {{ {JX} }};
PetscScalar t0[{lwork}], t1[{lwork}];
PetscScalar one=1.0E0;
for({IntType_c} j=0; j<{mxyz}; j++)
for({IntType_c} i=0; i<{nscal}; i++)
t0[j + {mxyz}*i] = x[i + {nscal}*j] * w[i + {nscal}*j];
kronmxv(1, {nx},{ny},{nz}, {mx},{my},{mz}, {nscal}, JX,{JY},{JZ}, t0,t1);
for({IntType_c} j=0; j<{nxyz}; j++)
for({IntType_c} i=0; i<{nscal}; i++)
y[i + {nscal}*j] += t1[j + {nxyz}*i];
return;
}}
"""
self.prolong_kernel = op2.Kernel(prolong_code,
"prolongation",
ldargs=["-lpetsc"])
self.restrict_kernel = op2.Kernel(restrict_code,
"restriction",
ldargs=["-lpetsc"])
nodes = op2.Set(nnode, "nodes")
edges = op2.Set(nedge, "edges")
ppedge = op2.Map(edges, nodes, 2, pp, "ppedge")
p_A = op2.Dat(edges, data=A, name="p_A")
p_r = op2.Dat(nodes, data=r, name="p_r")
p_u = op2.Dat(nodes, data=u, name="p_u")
p_du = op2.Dat(nodes, data=du, name="p_du")
alpha = op2.Global(1, data=1.0, name="alpha", dtype=fp_type)
beta = op2.Global(1, data=1.0, name="beta", dtype=fp_type)
res = op2.Kernel(
"""void res(%(t)s *A, %(t)s *u, %(t)s *du, const %(t)s *beta){
*du += (*beta)*(*A)*(*u);
}""" % {'t': "double" if fp_type == np.float64 else "float"}, "res")
update = op2.Kernel(
"""
void update(%(t)s *r, %(t)s *du, %(t)s *u, %(t)s *u_sum, %(t)s *u_max) {
*u += *du + alpha * (*r);
*du = %(z)s;
*u_sum += (*u)*(*u);
*u_max = *u_max > *u ? *u_max : *u;
}""" % {
't': "double" if fp_type == np.float64 else "float",
'z': "0.0" if fp_type == np.float64 else "0.0f"
}, "update")
for iter in xrange(0, NITER):
def mass():
init = FlatBlock("""
double CG1[3][6] = { { 0.09157621, 0.09157621, 0.81684757,
0.44594849, 0.44594849, 0.10810302 },
{ 0.09157621, 0.81684757, 0.09157621,
0.44594849, 0.10810302, 0.44594849 },
{ 0.81684757, 0.09157621, 0.09157621,
0.10810302, 0.44594849, 0.44594849 } };
double d_CG1[3][6][2] = { { { 1., 0. },
{ 1., 0. },
{ 1., 0. },
{ 1., 0. },
{ 1., 0. },
{ 1., 0. } },
{ { 0., 1. },
{ 0., 1. },
{ 0., 1. },
{ 0., 1. },
{ 0., 1. },
{ 0., 1. } },
{ { -1.,-1. },
{ -1.,-1. },
{ -1.,-1. },
{ -1.,-1. },
{ -1.,-1. },
{ -1.,-1. } } };
double w[6] = { 0.05497587, 0.05497587, 0.05497587, 0.11169079,
0.11169079, 0.11169079 };
double c_q0[6][2][2];
for(int i_g = 0; i_g < 6; i_g++)
{
for(int i_d_0 = 0; i_d_0 < 2; i_d_0++)
{
for(int i_d_1 = 0; i_d_1 < 2; i_d_1++)
{
c_q0[i_g][i_d_0][i_d_1] = 0.0;
for(int q_r_0 = 0; q_r_0 < 3; q_r_0++)
{
c_q0[i_g][i_d_0][i_d_1] += c0[q_r_0][i_d_0] * d_CG1[q_r_0][i_g][i_d_1];
};
};
};
};
for(int i_g = 0; i_g < 6; i_g++)
""")
assembly = Incr(Symbol("localTensor", ("i_r_0", "i_r_1")),
FlatBlock("ST0 * w[i_g]"))
assembly = Block([
FlatBlock(
"double ST0 = 0.0;\nST0 += CG1[i_r_0][i_g] * CG1[i_r_1][i_g] * (c_q0[i_g][0][0] * \
c_q0[i_g][1][1] + -1 * c_q0[i_g][0][1] * c_q0[i_g][1][0]);\n"
), assembly
],
open_scope=True)
assembly = c_for("i_r_0", 3, c_for("i_r_1", 3, assembly))
kernel_code = FunDecl("void", "mass", [
Decl("double", Symbol("localTensor", (3, 3))),
Decl("double*", c_sym("c0[2]"))
], Block([init, assembly], open_scope=False))
return op2.Kernel(kernel_code, "mass")
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 test_invalid_mode(self, elements, elem_node, mat, mode):
"""Mat args can only have modes WRITE and INC."""
with pytest.raises(ModeValueError):
op2.par_loop(op2.Kernel("", "dummy"), elements,
mat(mode, (elem_node[op2.i[0]], elem_node[op2.i[1]])))
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 test_wo(self, backend, elems, x):
"""Set a Dat to a scalar value with op2.WRITE."""
kernel_wo = """void kernel_wo(unsigned int* x) { *x = 42; }"""
op2.par_loop(op2.Kernel(kernel_wo, "kernel_wo"), elems, x(op2.WRITE))
assert all(map(lambda x: x == 42, x.data))
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 make_extruded_coords(extruded_topology,
base_coords,
ext_coords,
layer_height,
extrusion_type='uniform',
kernel=None):
"""
Given either a kernel or a (fixed) layer_height, compute an
extruded coordinate field for an extruded mesh.
:arg extruded_topology: an :class:`~.ExtrudedMeshTopology` to extrude
a coordinate field for.
:arg base_coords: a :class:`~.Function` to read the base
coordinates from.
:arg ext_coords: a :class:`~.Function` to write the extruded
coordinates into.
:arg layer_height: the height for each layer. Either a scalar,
where layers will be equi-spaced at the specified height, or a
1D array of variable layer heights to use through the extrusion.
:arg extrusion_type: the type of extrusion to use. Predefined
options are either "uniform" (creating equi-spaced layers by
extruding in the (n+1)dth direction), "radial" (creating
equi-spaced layers by extruding in the outward direction from
the origin) or "radial_hedgehog" (creating equi-spaced layers
by extruding coordinates in the outward cell-normal
direction, needs a P1dgxP1 coordinate field).
:arg kernel: an optional kernel to carry out coordinate extrusion.
The kernel signature (if provided) is::
void kernel(double **base_coords, double **ext_coords,
double *layer_height, int layer)
The kernel iterates over the cells of the mesh and receives as
arguments the coordinates of the base cell (to read), the
coordinates on the extruded cell (to write to), the fixed layer
height, and the current cell layer.
"""
_, vert_space = ext_coords.function_space().ufl_element().sub_elements(
)[0].sub_elements()
if kernel is None and not (vert_space.degree() == 1
and vert_space.family()
in ['Lagrange', 'Discontinuous Lagrange']):
raise RuntimeError(
'Extrusion of coordinates is only possible for a P1 or P1dg interval unless a custom kernel is provided'
)
layer_height = numpy.atleast_1d(numpy.array(layer_height, dtype=RealType))
if layer_height.ndim > 1:
raise RuntimeError('Extrusion layer height should be 1d or scalar')
if layer_height.size > 1:
layer_height = numpy.cumsum(numpy.concatenate(([0], layer_height)))
layer_heights = layer_height.size
layer_height = op2.Global(layer_heights, layer_height, dtype=RealType)
if kernel is not None:
op2.ParLoop(kernel,
ext_coords.cell_set,
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
base_coords.dat(op2.READ, base_coords.cell_node_map()),
layer_height(op2.READ),
pass_layer_arg=True,
is_loopy_kernel=True).compute()
return
ext_fe = create_element(ext_coords.ufl_element())
ext_shape = ext_fe.index_shape
base_fe = create_element(base_coords.ufl_element())
base_shape = base_fe.index_shape
data = []
data.append(lp.GlobalArg("ext_coords", dtype=ScalarType, shape=ext_shape))
data.append(lp.GlobalArg("base_coords", dtype=ScalarType,
shape=base_shape))
data.append(
lp.GlobalArg("layer_height", dtype=RealType, shape=(layer_heights, )))
data.append(lp.ValueArg('layer'))
base_coord_dim = base_coords.function_space().value_size
# Deal with tensor product cells
adim = len(ext_shape) - 2
# handle single or variable layer heights
if layer_heights == 1:
height_var = "layer_height[0] * (layer + l)"
else:
height_var = "layer_height[layer + l]"
def _get_arity_axis_inames(_base):
return tuple(_base + str(i) for i in range(adim))
def _get_lp_domains(_inames, _extents):
domains = []
for idx, extent in zip(_inames, _extents):
inames = isl.make_zero_and_vars([idx])
domains.append(((inames[0].le_set(inames[idx])) &
(inames[idx].lt_set(inames[0] + extent))))
return domains
if extrusion_type == 'uniform':
domains = []
dd = _get_arity_axis_inames('d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(('c', 'l'), (base_coord_dim, 2)))
instructions = """
ext_coords[{dd}, l, c] = base_coords[{dd}, c]
ext_coords[{dd}, l, {base_coord_dim}] = ({hv})
""".format(dd=', '.join(dd),
base_coord_dim=base_coord_dim,
hv=height_var)
ast = lp.make_function(domains,
instructions,
data,
name="pyop2_kernel_uniform_extrusion",
target=lp.CTarget(),
seq_dependencies=True,
silenced_warnings=["summing_if_branches_ops"])
elif extrusion_type == 'radial':
domains = []
dd = _get_arity_axis_inames('d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(
_get_lp_domains(('c', 'k', 'l'), (base_coord_dim, ) * 2 + (2, )))
instructions = """
<{RealType}> tt[{dd}] = 0
<{RealType}> bc[{dd}] = 0
for k
bc[{dd}] = real(base_coords[{dd}, k])
tt[{dd}] = tt[{dd}] + bc[{dd}] * bc[{dd}]
end
tt[{dd}] = sqrt(tt[{dd}])
ext_coords[{dd}, l, c] = base_coords[{dd}, c] + base_coords[{dd}, c] * ({hv}) / tt[{dd}]
""".format(RealType=RealType, dd=', '.join(dd), hv=height_var)
ast = lp.make_function(domains,
instructions,
data,
name="pyop2_kernel_radial_extrusion",
target=lp.CTarget(),
seq_dependencies=True,
silenced_warnings=["summing_if_branches_ops"])
elif extrusion_type == 'radial_hedgehog':
# Only implemented for interval in 2D and triangle in 3D.
# gdim != tdim already checked in ExtrudedMesh constructor.
tdim = base_coords.ufl_domain().ufl_cell().topological_dimension()
if tdim not in [1, 2]:
raise NotImplementedError(
"Hedgehog extrusion not implemented for %s" %
base_coords.ufl_domain().ufl_cell())
# tdim == 1:
#
# normal is:
# (0 -1) (x2 - x1)
# (1 0) (y2 - y1)
#
# tdim == 2:
# normal is
# v0 x v1
#
# /\
# v0/ \
# / \
# /------\
# v1
domains = []
dd = _get_arity_axis_inames('d')
_dd = _get_arity_axis_inames('_d')
domains.extend(_get_lp_domains(dd, ext_shape[:adim]))
domains.extend(_get_lp_domains(_dd, ext_shape[:adim]))
domains.extend(
_get_lp_domains(('c0', 'c1', 'c2', 'c3', 'k', 'l'),
(base_coord_dim, ) * 5 + (2, )))
# Formula for normal, n
n_1_1 = """
n[0] = -bc[1, 1] + bc[0, 1]
n[1] = bc[1, 0] - bc[0, 0]
"""
n_2_1 = """
v0[c3] = bc[1, c3] - bc[0, c3]
v1[c3] = bc[2, c3] - bc[0, c3]
n[0] = v0[1] * v1[2] - v0[2] * v1[1]
n[1] = v0[2] * v1[0] - v0[0] * v1[2]
n[2] = v0[0] * v1[1] - v0[1] * v1[0]
"""
n_2_2 = """
v0[c3] = bc[0, 1, c3] - bc[0, 0, c3]
v1[c3] = bc[1, 0, c3] - bc[0, 0, c3]
n[0] = v0[1] * v1[2] - v0[2] * v1[1]
n[1] = v0[2] * v1[0] - v0[0] * v1[2]
n[2] = v0[0] * v1[1] - v0[1] * v1[0]
"""
n_dict = {1: {1: n_1_1}, 2: {1: n_2_1, 2: n_2_2}}
instructions = """
<{RealType}> dot = 0
<{RealType}> norm = 0
<{RealType}> v0[c2] = 0
<{RealType}> v1[c2] = 0
<{RealType}> n[c2] = 0
<{RealType}> x[c2] = 0
<{RealType}> bc[{_dd}, c1] = real(base_coords[{_dd}, c1])
for {_dd}
x[c1] = x[c1] + bc[{_dd}, c1]
end
{ninst}
for k
dot = dot + x[k] * n[k]
norm = norm + n[k] * n[k]
end
norm = sqrt(norm)
norm = -norm if dot < 0 else norm
ext_coords[{dd}, l, c0] = base_coords[{dd}, c0] + n[c0] * ({hv}) / norm
""".format(RealType=RealType,
dd=', '.join(dd),
_dd=', '.join(_dd),
ninst=n_dict[tdim][adim],
hv=height_var)
ast = lp.make_function(domains,
instructions,
data,
name="pyop2_kernel_radial_hedgehog_extrusion",
target=lp.CTarget(),
seq_dependencies=True,
silenced_warnings=["summing_if_branches_ops"])
else:
raise NotImplementedError('Unsupported extrusion type "%s"' %
extrusion_type)
kernel = op2.Kernel(ast, ast.name)
op2.ParLoop(kernel,
ext_coords.cell_set,
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
base_coords.dat(op2.READ, base_coords.cell_node_map()),
layer_height(op2.READ),
pass_layer_arg=True,
is_loopy_kernel=True).compute()
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 transfer_kernel(Pk, P1):
"""Compile a kernel that will map between Pk and P1.
:returns: a PyOP2 kernel.
The prolongation maps a solution in P1 into Pk using the natural
embedding. The restriction maps a residual in the dual of Pk into
the dual of P1 (it is the dual of the prolongation), computed
using linearity of the test function.
"""
# Mapping of a residual in Pk into a residual in P1
from coffee import base as coffee
from tsfc.coffee import generate as generate_coffee, SCALAR_TYPE
from tsfc.parameters import default_parameters
from gem import gem, impero_utils as imp
# Pk should be at least the same size as P1
assert Pk.finat_element.space_dimension(
) >= P1.finat_element.space_dimension()
# In the general case we should compute this by doing:
# numpy.linalg.solve(Pkmass, PkP1mass)
Pke = Pk.finat_element._element
P1e = P1.finat_element._element
# TODO, rework to use finat.
matrix = numpy.dot(Pke.dual.to_riesz(P1e.get_nodal_basis()),
P1e.get_coeffs().T).T
Vout, Vin = P1, Pk
weights = gem.Literal(matrix)
name = "Pk_P1_mapper"
funargs = []
assert Vin.shape == Vout.shape
shape = (P1e.space_dimension() * Vout.value_size,
Pke.space_dimension() * Vin.value_size)
outarg = coffee.Decl(SCALAR_TYPE, coffee.Symbol("A", rank=shape))
i = gem.Index()
j = gem.Index()
k = gem.Index()
indices = i, j, k
A = gem.Variable("A", shape)
outgem = [
gem.Indexed(
gem.reshape(A, (P1e.space_dimension(), Vout.value_size),
(Pke.space_dimension(), Vin.value_size)), (i, k, j, k))
]
funargs.append(outarg)
expr = gem.Indexed(weights, (i, j))
outgem, = imp.preprocess_gem(outgem)
ir = imp.compile_gem([(outgem, expr)], indices)
index_names = [(i, "i"), (j, "j"), (k, "k")]
body = generate_coffee(ir, index_names, default_parameters()["precision"])
function = coffee.FunDecl("void",
name,
funargs,
body,
pred=["static", "inline"])
return op2.Kernel(function, name=function.name)
def dg_injection_kernel(Vf, Vc, ncell):
from firedrake import Tensor, AssembledVector, TestFunction, TrialFunction
from firedrake.slate.slac import compile_expression
macro_builder = MacroKernelBuilder(ScalarType_c, ncell)
f = ufl.Coefficient(Vf)
macro_builder.set_coefficients([f])
macro_builder.set_coordinates(Vf.mesh())
Vfe = create_element(Vf.ufl_element())
macro_quadrature_rule = make_quadrature(
Vfe.cell, estimate_total_polynomial_degree(ufl.inner(f, f)))
index_cache = {}
parameters = default_parameters()
integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell")
macro_cfg = dict(interface=macro_builder,
ufl_cell=Vf.ufl_cell(),
integration_dim=integration_dim,
entity_ids=entity_ids,
index_cache=index_cache,
quadrature_rule=macro_quadrature_rule,
scalar_type=parameters["scalar_type"])
fexpr, = fem.compile_ufl(f, **macro_cfg)
X = ufl.SpatialCoordinate(Vf.mesh())
C_a, = fem.compile_ufl(X, **macro_cfg)
detJ = ufl_utils.preprocess_expression(abs(
ufl.JacobianDeterminant(f.ufl_domain())),
complex_mode=complex_mode)
macro_detJ, = fem.compile_ufl(detJ, **macro_cfg)
Vce = create_element(Vc.ufl_element())
coarse_builder = firedrake_interface.KernelBuilder("cell", "otherwise", 0,
ScalarType_c)
coarse_builder.set_coordinates(Vc.mesh())
argument_multiindices = (Vce.get_indices(), )
argument_multiindex, = argument_multiindices
return_variable, = coarse_builder.set_arguments((ufl.TestFunction(Vc), ),
argument_multiindices)
integration_dim, entity_ids = lower_integral_type(Vce.cell, "cell")
# Midpoint quadrature for jacobian on coarse cell.
quadrature_rule = make_quadrature(Vce.cell, 0)
coarse_cfg = dict(interface=coarse_builder,
ufl_cell=Vc.ufl_cell(),
integration_dim=integration_dim,
entity_ids=entity_ids,
index_cache=index_cache,
quadrature_rule=quadrature_rule,
scalar_type=parameters["scalar_type"])
X = ufl.SpatialCoordinate(Vc.mesh())
K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh()),
complex_mode=complex_mode)
C_0, = fem.compile_ufl(X, **coarse_cfg)
K, = fem.compile_ufl(K, **coarse_cfg)
i = gem.Index()
j = gem.Index()
C_0 = gem.Indexed(C_0, (j, ))
C_0 = gem.index_sum(C_0, quadrature_rule.point_set.indices)
C_a = gem.Indexed(C_a, (j, ))
X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), C_a))
K_ij = gem.Indexed(K, (i, j))
K_ij = gem.index_sum(K_ij, quadrature_rule.point_set.indices)
X_a = gem.index_sum(gem.Product(K_ij, X_a), (j, ))
C_0, = quadrature_rule.point_set.points
C_0 = gem.Indexed(gem.Literal(C_0), (i, ))
# fine quad points in coarse reference space.
X_a = gem.Sum(C_0, gem.Product(gem.Literal(-1), X_a))
X_a = gem.ComponentTensor(X_a, (i, ))
# Coarse basis function evaluated at fine quadrature points
phi_c = fem.fiat_to_ufl(
Vce.point_evaluation(0, X_a, (Vce.cell.get_dimension(), 0)), 0)
tensor_indices = tuple(gem.Index(extent=d) for d in f.ufl_shape)
phi_c = gem.Indexed(phi_c, argument_multiindex + tensor_indices)
fexpr = gem.Indexed(fexpr, tensor_indices)
quadrature_weight = macro_quadrature_rule.weight_expression
expr = gem.Product(gem.IndexSum(gem.Product(phi_c, fexpr), tensor_indices),
gem.Product(macro_detJ, quadrature_weight))
quadrature_indices = macro_builder.indices + macro_quadrature_rule.point_set.indices
reps = spectral.Integrals([expr], quadrature_indices,
argument_multiindices, parameters)
assignments = spectral.flatten([(return_variable, reps)], index_cache)
return_variables, expressions = zip(*assignments)
expressions = impero_utils.preprocess_gem(expressions,
**spectral.finalise_options)
assignments = list(zip(return_variables, expressions))
impero_c = impero_utils.compile_gem(assignments,
quadrature_indices +
argument_multiindex,
remove_zeros=True)
index_names = []
def name_index(index, name):
index_names.append((index, name))
if index in index_cache:
for multiindex, suffix in zip(index_cache[index],
string.ascii_lowercase):
name_multiindex(multiindex, name + suffix)
def name_multiindex(multiindex, name):
if len(multiindex) == 1:
name_index(multiindex[0], name)
else:
for i, index in enumerate(multiindex):
name_index(index, name + str(i))
name_multiindex(quadrature_indices, 'ip')
for multiindex, name in zip(argument_multiindices, ['j', 'k']):
name_multiindex(multiindex, name)
index_names.extend(zip(macro_builder.indices, ["entity"]))
body = generate_coffee(impero_c, index_names, ScalarType_c)
retarg = ast.Decl(ScalarType_c,
ast.Symbol("R", rank=(Vce.space_dimension(), )))
local_tensor = coarse_builder.local_tensor
local_tensor.init = ast.ArrayInit(
numpy.zeros(Vce.space_dimension(), dtype=ScalarType_c))
body.children.insert(0, local_tensor)
args = [retarg] + macro_builder.kernel_args + [
macro_builder.coordinates_arg, coarse_builder.coordinates_arg
]
# Now we have the kernel that computes <f, phi_c>dx_c
# So now we need to hit it with the inverse mass matrix on dx_c
u = TrialFunction(Vc)
v = TestFunction(Vc)
expr = Tensor(ufl.inner(u, v) * ufl.dx).inv * AssembledVector(
ufl.Coefficient(Vc))
Ainv, = compile_expression(expr, coffee=True)
Ainv = Ainv.kinfo.kernel
A = ast.Symbol(local_tensor.sym.symbol)
R = ast.Symbol("R")
body.children.append(
ast.FunCall(Ainv.name, R, coarse_builder.coordinates_arg.sym, A))
from coffee.base import Node
assert isinstance(Ainv._code, Node)
return op2.Kernel(ast.Node([
Ainv._code,
ast.FunDecl("void",
"pyop2_kernel_injection_dg",
args,
body,
pred=["static", "inline"])
]),
name="pyop2_kernel_injection_dg",
cpp=True,
include_dirs=Ainv._include_dirs,
headers=Ainv._headers)
def compile_c_kernel(expression, to_pts, to_element, fs, coords):
"""Produce a :class:`PyOP2.Kernel` from the c expression provided."""
coords_space = coords.function_space()
coords_element = coords_space.fiat_element
names = {v[0] for v in expression._user_args}
X = coords_element.tabulate(0, to_pts).values()[0]
# Produce C array notation of X.
X_str = "{{"+"},\n{".join([",".join(map(str, x)) for x in X.T])+"}}"
A = utils.unique_name("A", names)
X = utils.unique_name("X", names)
x_ = utils.unique_name("x_", names)
k = utils.unique_name("k", names)
d = utils.unique_name("d", names)
i_ = utils.unique_name("i", names)
# x is a reserved name.
x = "x"
if "x" in names:
raise ValueError("cannot use 'x' as a user-defined Expression variable")
ass_exp = [ast.Assign(ast.Symbol(A, (k,), ((len(expression.code), i),)),
ast.FlatBlock("%s" % code))
for i, code in enumerate(expression.code)]
vals = {
"X": X,
"x": x,
"x_": x_,
"k": k,
"d": d,
"i": i_,
"x_array": X_str,
"dim": coords_space.dim,
"xndof": coords_element.space_dimension(),
# FS will always either be a functionspace or
# vectorfunctionspace, so just accessing dim here is safe
# (we don't need to go through ufl_element.value_shape())
"nfdof": to_element.space_dimension() * numpy.prod(fs.dim, dtype=int),
"ndof": to_element.space_dimension(),
"assign_dim": numpy.prod(expression.value_shape(), dtype=int)
}
init = ast.FlatBlock("""
const double %(X)s[%(ndof)d][%(xndof)d] = %(x_array)s;
double %(x)s[%(dim)d];
const double pi = 3.141592653589793;
""" % vals)
block = ast.FlatBlock("""
for (unsigned int %(d)s=0; %(d)s < %(dim)d; %(d)s++) {
%(x)s[%(d)s] = 0;
for (unsigned int %(i)s=0; %(i)s < %(xndof)d; %(i)s++) {
%(x)s[%(d)s] += %(X)s[%(k)s][%(i)s] * %(x_)s[%(i)s][%(d)s];
};
};
""" % vals)
loop = ast.c_for(k, "%(ndof)d" % vals, ast.Block([block] + ass_exp,
open_scope=True))
user_args = []
user_init = []
for _, arg in expression._user_args:
if arg.shape == (1, ):
user_args.append(ast.Decl("double *", "%s_" % arg.name))
user_init.append(ast.FlatBlock("const double %s = *%s_;" %
(arg.name, arg.name)))
else:
user_args.append(ast.Decl("double *", arg.name))
kernel_code = ast.FunDecl("void", "expression_kernel",
[ast.Decl("double", ast.Symbol(A, (int("%(nfdof)d" % vals),))),
ast.Decl("double**", x_)] + user_args,
ast.Block(user_init + [init, loop],
open_scope=False))
coefficients = [coords]
for _, arg in expression._user_args:
coefficients.append(GlobalWrapper(arg))
return op2.Kernel(kernel_code, kernel_code.name), False, tuple(coefficients)
def make_extruded_coords(extruded_topology,
base_coords,
ext_coords,
layer_height,
extrusion_type='uniform',
kernel=None):
"""
Given either a kernel or a (fixed) layer_height, compute an
extruded coordinate field for an extruded mesh.
:arg extruded_topology: an :class:`~.ExtrudedMeshTopology` to extrude
a coordinate field for.
:arg base_coords: a :class:`~.Function` to read the base
coordinates from.
:arg ext_coords: a :class:`~.Function` to write the extruded
coordinates into.
:arg layer_height: an equi-spaced height for each layer.
:arg extrusion_type: the type of extrusion to use. Predefined
options are either "uniform" (creating equi-spaced layers by
extruding in the (n+1)dth direction), "radial" (creating
equi-spaced layers by extruding in the outward direction from
the origin) or "radial_hedgehog" (creating equi-spaced layers
by extruding coordinates in the outward cell-normal
direction, needs a P1dgxP1 coordinate field).
:arg kernel: an optional kernel to carry out coordinate extrusion.
The kernel signature (if provided) is::
void kernel(double **base_coords, double **ext_coords,
int **layer, double *layer_height)
The kernel iterates over the cells of the mesh and receives as
arguments the coordinates of the base cell (to read), the
coordinates on the extruded cell (to write to), the layer number
of each cell and the fixed layer height.
"""
_, vert_space = ext_coords.function_space().ufl_element().sub_elements(
)[0].sub_elements()
if kernel is None and not (vert_space.degree() == 1
and vert_space.family()
in ['Lagrange', 'Discontinuous Lagrange']):
raise RuntimeError(
'Extrusion of coordinates is only possible for a P1 or P1dg interval unless a custom kernel is provided'
)
if kernel is not None:
pass
elif extrusion_type == 'uniform':
kernel = op2.Kernel(
"""
void uniform_extrusion_kernel(double **base_coords,
double **ext_coords,
int **layer,
double *layer_height) {
for ( int d = 0; d < %(base_map_arity)d; d++ ) {
for ( int c = 0; c < %(base_coord_dim)d; c++ ) {
ext_coords[2*d][c] = base_coords[d][c];
ext_coords[2*d+1][c] = base_coords[d][c];
}
ext_coords[2*d][%(base_coord_dim)d] = *layer_height * (layer[0][0]);
ext_coords[2*d+1][%(base_coord_dim)d] = *layer_height * (layer[0][0] + 1);
}
}""" % {
'base_map_arity': base_coords.cell_node_map().arity,
'base_coord_dim': base_coords.function_space().dim
}, "uniform_extrusion_kernel")
elif extrusion_type == 'radial':
kernel = op2.Kernel(
"""
void radial_extrusion_kernel(double **base_coords,
double **ext_coords,
int **layer,
double *layer_height) {
for ( int d = 0; d < %(base_map_arity)d; d++ ) {
double norm = 0.0;
for ( int c = 0; c < %(base_coord_dim)d; c++ ) {
norm += base_coords[d][c] * base_coords[d][c];
}
norm = sqrt(norm);
for ( int c = 0; c < %(base_coord_dim)d; c++ ) {
ext_coords[2*d][c] = base_coords[d][c] * (1 + (*layer_height * layer[0][0])/norm);
ext_coords[2*d+1][c] = base_coords[d][c] * (1 + (*layer_height * (layer[0][0]+1))/norm);
}
}
}""" % {
'base_map_arity': base_coords.cell_node_map().arity,
'base_coord_dim': base_coords.function_space().dim
}, "radial_extrusion_kernel")
elif extrusion_type == 'radial_hedgehog':
# Only implemented for interval in 2D and triangle in 3D.
# gdim != tdim already checked in ExtrudedMesh constructor.
if base_coords.ufl_domain().ufl_cell().topological_dimension() not in [
1, 2
]:
raise NotImplementedError(
"Hedgehog extrusion not implemented for %s" %
base_coords.ufl_domain().ufl_cell())
kernel = op2.Kernel(
"""
void radial_hedgehog_extrusion_kernel(double **base_coords,
double **ext_coords,
int **layer,
double *layer_height) {
double v0[%(base_coord_dim)d];
double v1[%(base_coord_dim)d];
double n[%(base_coord_dim)d];
double x[%(base_coord_dim)d] = {0};
double dot = 0.0;
double norm = 0.0;
int i, c, d;
if (%(base_coord_dim)d == 2) {
/*
* normal is:
* (0 -1) (x2 - x1)
* (1 0) (y2 - y1)
*/
n[0] = -(base_coords[1][1] - base_coords[0][1]);
n[1] = base_coords[1][0] - base_coords[0][0];
} else if (%(base_coord_dim)d == 3) {
/*
* normal is
* v0 x v1
*
* /\\
* v0/ \\
* / \\
* /------\\
* v1
*/
for (i = 0; i < 3; ++i) {
v0[i] = base_coords[1][i] - base_coords[0][i];
v1[i] = base_coords[2][i] - base_coords[0][i];
}
n[0] = v0[1] * v1[2] - v0[2] * v1[1];
n[1] = v0[2] * v1[0] - v0[0] * v1[2];
n[2] = v0[0] * v1[1] - v0[1] * v1[0];
}
for (i = 0; i < %(base_map_arity)d; ++i) {
for (c = 0; c < %(base_coord_dim)d; ++c) {
x[c] += base_coords[i][c];
}
}
for (i = 0; i < %(base_coord_dim)d; ++i) {
dot += x[i] * n[i];
norm += n[i] * n[i];
}
/*
* Make inward-pointing normals point out
*/
norm = sqrt(norm);
norm *= (dot < 0 ? -1 : 1);
for (d = 0; d < %(base_map_arity)d; ++d) {
for (c = 0; c < %(base_coord_dim)d; ++c ) {
ext_coords[2*d][c] = base_coords[d][c] + n[c] * layer_height[0] * layer[0][0] / norm;
ext_coords[2*d+1][c] = base_coords[d][c] + n[c] * layer_height[0] * (layer[0][0] + 1)/ norm;
}
}
}""" % {
'base_map_arity': base_coords.cell_node_map().arity,
'base_coord_dim': base_coords.function_space().dim
}, "radial_hedgehog_extrusion_kernel")
else:
raise NotImplementedError('Unsupported extrusion type "%s"' %
extrusion_type)
# Dat to hold layer number
import firedrake.functionspace as fs
layer_fs = fs.FunctionSpace(extruded_topology, 'DG', 0)
layers = extruded_topology.layers
layer = op2.Dat(
layer_fs.dof_dset,
np.repeat(np.arange(layers - 1, dtype=np.int32),
extruded_topology.cell_set.total_size).reshape(
layers - 1,
extruded_topology.cell_set.total_size).T.ravel(),
dtype=np.int32)
height = op2.Global(1, layer_height, dtype=float)
op2.par_loop(kernel, ext_coords.cell_set,
base_coords.dat(op2.READ, base_coords.cell_node_map()),
ext_coords.dat(op2.WRITE, ext_coords.cell_node_map()),
layer(op2.READ, layer_fs.cell_node_map()), height(op2.READ))
