def update(self, w):
if not self.separate_wind:
return
if isinstance(w, Function):
self.wind.assign(w)
else:
with self.wind.dat.vec_wo as wind_v:
w.copy(wind_v)
# Need to update your coarse grid mates, too.
# There's probably a better way to do this.
dm = self.V.dm
wind = self.wind
while get_level(get_function_space(dm).mesh())[1] > 0:
cdm = dm.coarsen()
cctx = get_appctx(cdm)
if cctx is None:
break
cwind = cctx.F._cache['coefficient_mapping'][wind]
inject(wind, cwind)
dm = cdm
wind = cwind
def error_of_flood_estimate(mesh, p, T):
# mixed function space
h, lvl = get_level(mesh)
v_h = FunctionSpace(mesh, "DG", p)
v_mu = FunctionSpace(mesh, "DG", p)
v_mv = FunctionSpace(mesh, "DG", p)
V = v_h*v_mu*v_mv
# parameters
if p == 0:
parameters["flooddrake"].update({"eps2": 8e-3})
if p == 1:
parameters["flooddrake"].update({"eps2": 3.5e-3})
# setup free surface depth
g = Function(V)
x = SpatialCoordinate(V.mesh())
g.sub(0).interpolate(conditional(x[0] < 20, 1.0/9.8, parameters["flooddrake"]["eps2"]))
# setup bed
bed = Function(V)
# setup state
state = State(V, g, bed)
# setup source (is only a depth function)
source = Function(v_h)
# timestep
solution = Timestepper(V, state.bed, source, 0.05)
solution.stepper(0, T, state.w, 0.5)
return solution
def coarsen(dm, comm):
"""Callback to coarsen a DM.
:arg DM: The DM to coarsen.
:arg comm: The communicator for the new DM (ignored)
This transfers a coarse application context over to the coarsened
DM (if found on the input DM).
"""
from firedrake.mg.utils import get_level
from firedrake.mg.ufl_utils import coarsen
V = get_function_space(dm)
if V is None:
raise RuntimeError("No functionspace found on DM")
hierarchy, level = get_level(V.mesh())
if level < 1:
raise RuntimeError("Cannot coarsen coarsest DM")
if hasattr(V, "_coarse"):
cdm = V._coarse.dm
else:
V._coarse = firedrake.FunctionSpace(hierarchy[level - 1],
V.ufl_element())
cdm = V._coarse.dm
ctx = get_appctx(dm)
if ctx is not None:
set_appctx(cdm, coarsen(ctx))
# Necessary for MG inside a fieldsplit in a SNES.
cdm.setKSPComputeOperators(
firedrake.solving_utils._SNESContext.compute_operators)
return cdm
def coarsen(dm, comm):
"""Callback to coarsen a DM.
:arg DM: The DM to coarsen.
:arg comm: The communicator for the new DM (ignored)
This transfers a coarse application context over to the coarsened
DM (if found on the input DM).
"""
from firedrake.mg.utils import get_level
V = get_function_space(dm)
if V is None:
raise RuntimeError("No functionspace found on DM")
hierarchy, level = get_level(V.mesh())
if level < 1:
raise RuntimeError("Cannot coarsen coarsest DM")
if hasattr(V, "_coarse"):
cdm = V._coarse.dm
else:
coarsen = get_ctx_coarsener(dm)
V._coarse = coarsen(V, coarsen)
cdm = V._coarse.dm
transfer = get_transfer_operators(dm)
push_transfer_operators(cdm, *transfer)
coarsen = get_ctx_coarsener(dm)
push_ctx_coarsener(cdm, coarsen)
ctx = get_appctx(dm)
if ctx is not None:
push_appctx(cdm, coarsen(ctx, coarsen))
# Necessary for MG inside a fieldsplit in a SNES.
cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators)
V._coarse._fine = V
return cdm
def coarsen(dm, comm):
"""Callback to coarsen a DM.
:arg DM: The DM to coarsen.
:arg comm: The communicator for the new DM (ignored)
This transfers a coarse application context over to the coarsened
DM (if found on the input DM).
"""
from firedrake.mg.utils import get_level
from firedrake.mg.ufl_utils import coarsen
V = get_function_space(dm)
if V is None:
raise RuntimeError("No functionspace found on DM")
hierarchy, level = get_level(V.mesh())
if level < 1:
raise RuntimeError("Cannot coarsen coarsest DM")
if hasattr(V, "_coarse"):
cdm = V._coarse.dm
else:
V._coarse = firedrake.FunctionSpace(hierarchy[level - 1], V.ufl_element())
cdm = V._coarse.dm
ctx = get_appctx(dm)
if ctx is not None:
set_appctx(cdm, coarsen(ctx))
# Necessary for MG inside a fieldsplit in a SNES.
cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators)
return cdm
def form_function(cls, snes, X, F):
"""Form the residual for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg F: the residual at X (a Vec)
"""
from firedrake.assemble import assemble
dm = snes.getDM()
ctx = dm.getAppCtx()
_, lvl = utils.get_level(dm)
# FIXME: Think about case where DM is refined but we don't
# have a hierarchy of problems better.
if len(ctx._problems) == 1:
lvl = -1
problem = ctx._problems[lvl]
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._xs[lvl].dat.vec as v:
if v != X:
X.copy(v)
assemble(ctx.Fs[lvl],
tensor=ctx._Fs[lvl],
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
for bc in problem.bcs:
bc.zero(ctx._Fs[lvl])
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._Fs[lvl].dat.vec_ro as v:
v.copy(F)
def form_function(cls, snes, X, F):
"""Form the residual for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg F: the residual at X (a Vec)
"""
from firedrake.assemble import assemble
dm = snes.getDM()
ctx = dm.getAppCtx()
_, lvl = utils.get_level(dm)
# FIXME: Think about case where DM is refined but we don't
# have a hierarchy of problems better.
if len(ctx._problems) == 1:
lvl = -1
problem = ctx._problems[lvl]
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._xs[lvl].dat.vec as v:
if v != X:
X.copy(v)
assemble(ctx.Fs[lvl], tensor=ctx._Fs[lvl],
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
for bc in problem.bcs:
bc.zero(ctx._Fs[lvl])
# F may not be the same vector as self._F, so copy
# residual out to F.
with ctx._Fs[lvl].dat.vec_ro as v:
v.copy(F)
def prolong(self, coarse, fine):
# Rebuild without any indices
V = FunctionSpace(fine.ufl_domain(),
fine.function_space().ufl_element())
key = V.dim()
firsttime = self.bcs.get(key, None) is None
gamma = self.gamma
_, level = get_level(fine.function_space().mesh())
nu = self.nus[level]
if firsttime:
bcs = self.fix_coarse_boundaries(V)
u = TrialFunction(V)
v = TestFunction(V)
tildeu, rhs = Function(V), Function(V)
if self.element == "sv":
A = assemble(nu * inner(grad(u), grad(v)) * dx +
gamma * inner(div(u), div(v)) * dx,
bcs=bcs,
mat_type=self.patchparams["mat_type"])
bform = nu * inner(grad(rhs), grad(v)) * dx + gamma * inner(
div(rhs), div(v)) * dx
else:
A = assemble(nu * inner(grad(u), grad(v)) * dx +
gamma * inner(cell_avg(div(u)), div(v)) * dx,
bcs=bcs,
mat_type=self.patchparams["mat_type"])
bform = nu * inner(grad(rhs), grad(v)) * dx + gamma * inner(
cell_avg(div(rhs)), div(v)) * dx
b = Function(V)
solver = LinearSolver(A,
solver_parameters=self.patchparams,
options_prefix="prolongation")
self.bcs[key] = bcs
self.solver[key] = solver
self.rhs[key] = tildeu, rhs
self.tensors[key] = A, b, bform
else:
bcs = self.bcs[key]
solver = self.solver[key]
A, b, bform = self.tensors[key]
tildeu, rhs = self.rhs[key]
prolong(coarse, rhs)
b = assemble(bform, bcs=bcs, tensor=b)
# # Could do
# #solver.solve(tildeu, b)
# # but that calls a lot of SNES and KSP overhead.
# # We know we just want to apply the PC:
with solver.inserted_options():
with b.dat.vec_ro as rhsv:
with tildeu.dat.vec_wo as x:
solver.ksp.pc.apply(rhsv, x)
fine.assign(rhs - tildeu)
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 test_get_level():
mesh = UnitSquareMesh(5, 5)
L = 2
refinements_per_level = 2
MH = MeshHierarchy(mesh, L, refinements_per_level=refinements_per_level)
for i, mesh in enumerate(MH):
assert get_level(mesh)[1] == i
def test_get_skip_level():
mesh = UnitSquareMesh(5, 5)
L = 2
refinements_per_level = 2
MH = MeshHierarchy(mesh, L, refinements_per_level=refinements_per_level)
for i, mesh in enumerate(MH._unskipped_hierarchy):
assert get_level(mesh)[1] == Fraction(i, refinements_per_level)
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 _dm(self):
from firedrake.mg.utils import get_level
dm = self.dof_dset.dm
_, level = get_level(self.mesh())
dmhooks.attach_hooks(dm, level=level,
sf=self.mesh()._plex.getPointSF(),
section=None)
# Remember the function space so we can get from DM back to FunctionSpace.
dmhooks.set_function_space(dm, self)
return dm
def _dm(self):
from firedrake.mg.utils import get_level
dm = self.dof_dset.dm
_, level = get_level(self.mesh())
dmhooks.attach_hooks(dm, level=level,
sf=self.mesh().topology_dm.getPointSF(),
section=self._shared_data.global_numbering)
# Remember the function space so we can get from DM back to FunctionSpace.
dmhooks.set_function_space(dm, self)
return dm
def _dm(self):
from firedrake.mg.utils import get_level
dm = self.dof_dset.dm
_, level = get_level(self.mesh())
dmhooks.attach_hooks(dm, level=level,
sf=self.mesh()._plex.getPointSF(),
section=self._shared_data.global_numbering)
# Remember the function space so we can get from DM back to FunctionSpace.
dmhooks.set_function_space(dm, self)
return dm
def coarsen(dm, comm):
"""Callback to coarsen a DM.
:arg DM: The DM to coarsen.
:arg comm: The communicator for the new DM (ignored)
This transfers a coarse application context over to the coarsened
DM (if found on the input DM).
"""
from firedrake.mg.utils import get_level
V = get_function_space(dm)
hierarchy, level = get_level(V.mesh())
if level < 1:
raise RuntimeError("Cannot coarsen coarsest DM")
coarsen = get_ctx_coarsener(dm)
Vc = coarsen(V, coarsen)
cdm = Vc.dm
transfer = get_transfer_operators(dm)
parent = get_parent(dm)
add_hook(parent,
setup=partial(push_parent, cdm, parent),
teardown=partial(pop_parent, cdm, parent),
call_setup=True)
add_hook(parent,
setup=partial(push_transfer_operators, cdm, transfer),
teardown=partial(pop_transfer_operators, cdm, transfer),
call_setup=True)
if len(V) > 1:
for V_, Vc_ in zip(V, Vc):
transfer = get_transfer_operators(V_.dm)
add_hook(parent,
setup=partial(push_parent, Vc_.dm, parent),
teardown=partial(pop_parent, Vc_.dm, parent),
call_setup=True)
add_hook(parent,
setup=partial(push_transfer_operators, Vc_.dm, transfer),
teardown=partial(pop_transfer_operators, Vc_.dm,
transfer),
call_setup=True)
add_hook(parent,
setup=partial(push_ctx_coarsener, cdm, coarsen),
teardown=partial(pop_ctx_coarsener, cdm, coarsen),
call_setup=True)
ctx = get_appctx(dm)
if ctx is not None:
cctx = coarsen(ctx, coarsen)
add_hook(parent,
setup=partial(push_appctx, cdm, cctx),
teardown=partial(pop_appctx, cdm, cctx),
call_setup=True)
# Necessary for MG inside a fieldsplit in a SNES.
cdm.setKSPComputeOperators(
firedrake.solving_utils._SNESContext.compute_operators)
return cdm
def aug_jacobian(X, J, ctx):
mh, level = get_level(ctx._x.ufl_domain())
if case == 4 or case == 5:
BTWBlevel = BTWB_dict[level]
if level == nref:
Jsub = J.getNestSubMatrix(0, 0)
rmap, cmap = Jsub.getLGMap()
Jsub.axpy(1, BTWBlevel, Jsub.Structure.SUBSET_NONZERO_PATTERN)
Jsub.setLGMap(rmap, cmap)
else:
rmap, cmap = J.getLGMap()
J.axpy(1, BTWBlevel, J.Structure.SUBSET_NONZERO_PATTERN)
J.setLGMap(rmap, cmap)
def __call__(self, pc):
from firedrake.mg.utils import get_level
from firedrake.cython.mgimpl import get_entity_renumbering
dmf = pc.getDM()
ctx = pc.getAttr("ctx")
mf = ctx._x.ufl_domain()
(mh, level) = get_level(mf)
coarse_to_fine_cell_map = mh.coarse_to_fine_cells[level - 1]
(_,
firedrake_to_plex) = get_entity_renumbering(dmf, mf._cell_numbering,
"cell")
mc = mh[level - 1]
(_, coarse_firedrake_to_plex) = get_entity_renumbering(
mc._topology_dm, mc._cell_numbering, "cell")
patches = []
tdim = mf.topological_dimension()
for i, fine_firedrake in enumerate(coarse_to_fine_cell_map):
# there are d+1 many coarse cells that all map to the same fine cells.
# We only want to build the patch once, so skip repitions
if coarse_firedrake_to_plex[i] % (tdim + 1) != 0:
continue
# we need to convert firedrake cell numbering to plex cell numbering
fine_plex = [firedrake_to_plex[ff] for ff in fine_firedrake]
entities = []
for fp in fine_plex:
(pts, _) = dmf.getTransitiveClosure(fp, True)
for pt in pts:
value = dmf.getLabelValue("prolongation", pt)
if not (value > -1 and value <= level):
entities.append(pt)
iset = PETSc.IS().createGeneral(unique(entities),
comm=PETSc.COMM_SELF)
patches.append(iset)
piterset = PETSc.IS().createStride(size=len(patches),
first=0,
step=1,
comm=PETSc.COMM_SELF)
return (patches, piterset)
def test_inject_to_any_level():
M = MeshHierarchy(UnitSquareMesh(10, 10), 3)
V = FunctionSpaceHierarchy(M, 'DG', 0)
F = FunctionHierarchy(V)
# check for prolonging to top level
F[-1].interpolate(Expression("3"))
F[0].interpolate(Expression("3"))
H = FunctionHierarchy(V)
A = InjectDownToAnyLevel(0, F[-1], H)
assert get_level(A)[1] == 0
assert norm(assemble(A - F[0])) <= 0
def test_prolong_to_any_level():
M = MeshHierarchy(UnitSquareMesh(10, 10), 3)
V = FunctionSpaceHierarchy(M, 'DG', 0)
F = FunctionHierarchy(V)
# check for prolonging to top level
F[0].interpolate(Expression("3"))
F[2].interpolate(Expression("3"))
H = FunctionHierarchy(V)
A = ProlongUpToAnyLevel(2, F[0], H)
assert get_level(A)[1] == 2
assert norm(assemble(A - F[2])) <= 0
def refine(dm, comm):
"""Callback to refine a DM.
:arg DM: The DM to refine.
:arg comm: The communicator for the new DM (ignored)
"""
from firedrake.mg.utils import get_level
V = get_function_space(dm)
if V is None:
raise RuntimeError("No functionspace found on DM")
hierarchy, level = get_level(V.mesh())
if level >= len(hierarchy) - 1:
raise RuntimeError("Cannot refine finest DM")
if hasattr(V, "_fine"):
fdm = V._fine.dm
else:
V._fine = firedrake.FunctionSpace(hierarchy[level + 1], V.ufl_element())
fdm = V._fine.dm
return fdm
def form_jacobian(cls, snes, X, J, P):
"""Form the Jacobian for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
from firedrake.assemble import assemble
dm = snes.getDM()
ctx = dm.getAppCtx()
_, lvl = utils.get_level(dm)
# FIXME: Think about case where DM is refined but we don't
# have a hierarchy of problems better.
if len(ctx._problems) == 1:
lvl = -1
problem = ctx._problems[lvl]
if problem._constant_jacobian and ctx._jacobians_assembled[lvl]:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobians_assembled[lvl] = True
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._xs[lvl].dat.vec as v:
X.copy(v)
assemble(ctx.Js[lvl],
tensor=ctx._jacs[lvl],
bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
ctx._jacs[lvl].M._force_evaluation()
if ctx.Jps[lvl] is not None:
assemble(ctx.Jps[lvl],
tensor=ctx._pjacs[lvl],
bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
ctx._pjacs[lvl].M._force_evaluation()
def form_jacobian(cls, snes, X, J, P):
"""Form the Jacobian for this problem
:arg snes: a PETSc SNES object
:arg X: the current guess (a Vec)
:arg J: the Jacobian (a Mat)
:arg P: the preconditioner matrix (a Mat)
"""
from firedrake.assemble import assemble
dm = snes.getDM()
ctx = dm.getAppCtx()
_, lvl = utils.get_level(dm)
# FIXME: Think about case where DM is refined but we don't
# have a hierarchy of problems better.
if len(ctx._problems) == 1:
lvl = -1
problem = ctx._problems[lvl]
if problem._constant_jacobian and ctx._jacobians_assembled[lvl]:
# Don't need to do any work with a constant jacobian
# that's already assembled
return
ctx._jacobians_assembled[lvl] = True
# X may not be the same vector as the vec behind self._x, so
# copy guess in from X.
with ctx._xs[lvl].dat.vec as v:
X.copy(v)
assemble(ctx.Js[lvl],
tensor=ctx._jacs[lvl],
bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
ctx._jacs[lvl].M._force_evaluation()
if ctx.Jps[lvl] is not None:
assemble(ctx.Jps[lvl],
tensor=ctx._pjacs[lvl],
bcs=problem.bcs,
form_compiler_parameters=problem.form_compiler_parameters,
nest=problem._nest)
ctx._pjacs[lvl].M._force_evaluation()
def refine(dm, comm):
"""Callback to refine a DM.
:arg DM: The DM to refine.
:arg comm: The communicator for the new DM (ignored)
"""
from firedrake.mg.utils import get_level
V = get_function_space(dm)
if V is None:
raise RuntimeError("No functionspace found on DM")
hierarchy, level = get_level(V.mesh())
if level >= len(hierarchy) - 1:
raise RuntimeError("Cannot refine finest DM")
if hasattr(V, "_fine"):
fdm = V._fine.dm
else:
V._fine = firedrake.FunctionSpace(hierarchy[level + 1], V.ufl_element())
fdm = V._fine.dm
V._fine._coarse = V
return fdm
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 fix_coarse_boundaries(V):
hierarchy, level = get_level(V.mesh())
dm = V.mesh()._topology_dm
section = V.dm.getDefaultSection()
indices = []
fStart, fEnd = dm.getHeightStratum(1)
# Spin over faces, if the face is marked with a magic label
# value, it means it was in the coarse mesh.
for p in range(fStart, fEnd):
value = dm.getLabelValue("prolongation", p)
if value > -1 and value <= level:
# OK, so this is a coarse mesh face.
# Grab all the points in the closure.
closure, _ = dm.getTransitiveClosure(p)
for c in closure:
# Now add all the dofs on that point to the list
# of boundary nodes.
dof = section.getDof(c)
off = section.getOffset(c)
for d in range(dof):
indices.append(off + d)
nodelist = unique(indices).astype(IntType)
class FixedDirichletBC(DirichletBC):
def __init__(self, V, g, nodelist):
self.nodelist = nodelist
DirichletBC.__init__(self, V, g, "on_boundary")
@utils.cached_property
def nodes(self):
return self.nodelist
dim = V.mesh().topological_dimension()
bc = FixedDirichletBC(V, ufl.zero(V.ufl_element().value_shape()),
nodelist)
return bc
def standard_transfer(self, source, target, mode):
if not (source.ufl_shape[0] == 3
and "CG1" in source.ufl_element().shortstr()):
return super().standard_transfer(source, target, mode)
if mode == "prolong":
coarse = source
fine = target
elif mode == "restrict":
fine = source
coarse = target
else:
raise NotImplementedError
(mh, level) = get_level(coarse.ufl_domain())
if level not in self.transfers:
self.transfers[level] = BubbleTransfer(coarse.function_space(),
fine.function_space())
if mode == "prolong":
self.transfers[level].prolong(coarse, fine)
elif mode == "restrict":
self.transfers[level].restrict(fine, coarse)
else:
raise NotImplementedError
def __call__(self, pc):
from firedrake.mg.utils import get_level
from firedrake.cython.mgimpl import get_entity_renumbering
dmf = pc.getDM()
ctx = pc.getAttr("ctx")
mf = ctx._x.ufl_domain()
(mh, level) = get_level(mf)
coarse_to_fine_cell_map = mh.coarse_to_fine_cells[level - 1]
(_,
firedrake_to_plex) = get_entity_renumbering(dmf, mf._cell_numbering,
"cell")
patches = []
for fine_firedrake in coarse_to_fine_cell_map:
# we need to convert firedrake cell numbering to plex cell numbering
fine_plex = [firedrake_to_plex[ff] for ff in fine_firedrake]
entities = []
for fp in fine_plex:
(pts, _) = dmf.getTransitiveClosure(fp, True)
for pt in pts:
value = dmf.getLabelValue("prolongation", pt)
if not (value > -1 and value <= level):
entities.append(pt)
iset = PETSc.IS().createGeneral(unique(entities),
comm=PETSc.COMM_SELF)
patches.append(iset)
piterset = PETSc.IS().createStride(size=len(patches),
first=0,
step=1,
comm=PETSc.COMM_SELF)
return (patches, piterset)
def coarsen(dm, comm):
"""Callback to coarsen a DM.
:arg DM: The DM to coarsen.
:arg comm: The communicator for the new DM (ignored)
This transfers a coarse application context over to the coarsened
DM (if found on the input DM).
"""
from firedrake.mg.utils import get_level
V = get_function_space(dm)
hierarchy, level = get_level(V.mesh())
if level < 1:
raise RuntimeError("Cannot coarsen coarsest DM")
coarsen = get_ctx_coarsener(dm)
Vc = coarsen(V, coarsen)
cdm = Vc.dm
push_ctx_coarsener(cdm, coarsen)
ctx = get_appctx(dm)
if ctx is not None:
push_appctx(cdm, coarsen(ctx, coarsen))
# Necessary for MG inside a fieldsplit in a SNES.
cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators)
return cdm
def error_of_flood_estimate(mesh, p, T):
# mixed function space
h, lvl = get_level(mesh)
v_h = FunctionSpace(mesh, "DG", p)
v_mu = FunctionSpace(mesh, "DG", p)
v_mv = FunctionSpace(mesh, "DG", p)
V = v_h * v_mu * v_mv
# parameters
if p == 0:
parameters["flooddrake"].update({"eps2": 8e-3})
if p == 1:
parameters["flooddrake"].update({"eps2": 3.5e-3})
# setup free surface depth
g = Function(V)
x = SpatialCoordinate(V.mesh())
g.sub(0).interpolate(
conditional(x[0] < 20, 1.0 / 9.8, parameters["flooddrake"]["eps2"]))
# setup bed
bed = Function(V)
# setup state
state = State(V, g, bed)
# setup source (is only a depth function)
source = Function(v_h)
# timestep
solution = Timestepper(V, state.bed, source, 0.05)
solution.stepper(0, T, state.w, 0.5)
return solution
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 _dm(self):
from firedrake.mg.utils import get_level
dm = self.dof_dset.dm
_, level = get_level(self.mesh())
dmhooks.attach_hooks(dm, level=level)
return dm
def create_matrix(cls, dm):
ctx = dm.getAppCtx()
_, lvl = utils.get_level(dm)
return ctx._jacs[lvl]._M.handle
def create_matrix(cls, dm):
ctx = dm.getAppCtx()
_, lvl = utils.get_level(dm)
return ctx._jacs[lvl]._M.handle
def aug_jacobian(X, J, ctx):
mh, level = get_level(ctx._x.ufl_domain())
levelmesh = mh[level]
if case == 4 or case == 5:
if args.quad:
Vlevel = FunctionSpace(levelmesh, "RTCF", k)
Qlevel = FunctionSpace(levelmesh, "DQ", k - 1)
else:
if args.discretisation == "hdiv":
Vlevel = FunctionSpace(levelmesh, "BDM", k)
Qlevel = FunctionSpace(levelmesh, "DG", k - 1)
elif args.discretisation == "cg":
if dim == 2:
Vlevel = VectorFunctionSpace(levelmesh, "CG", k)
Qlevel = FunctionSpace(levelmesh, "DG", 0)
elif dim == 3:
Pklevel = FiniteElement("Lagrange", levelmesh.ufl_cell(),
2)
FBlevel = FiniteElement("FacetBubble",
levelmesh.ufl_cell(), 3)
eleulevel = VectorElement(
NodalEnrichedElement(Pklevel, FBlevel))
Vlevel = FunctionSpace(levelmesh, eleulevel)
Qlevel = FunctionSpace(levelmesh, "DG", 0)
else:
raise NotImplementedError("Only implemented for dim=2,3")
else:
raise ValueError(
"please specify hdiv or cg for --discretisation")
Zlevel = Vlevel * Qlevel
# Get B
tmpu, tmpp = TrialFunctions(Zlevel)
tmpv, tmpq = TestFunctions(Zlevel)
tmpF = -tmpq * div(tmpu) * dx
if dim == 2:
tmpbcs = [
DirichletBC(Zlevel.sub(0), Constant((0., 0.)), "on_boundary")
]
elif dim == 3:
tmpbcs = [
DirichletBC(Zlevel.sub(0), Constant((0., 0., 0.)),
"on_boundary")
]
else:
raise NotImplementedError("Only implemented for dim=2,3")
tmpa = lhs(tmpF)
M = assemble(tmpa, bcs=tmpbcs)
Blevel = M.M[1, 0].handle
# Get W
ptrial = TrialFunction(Qlevel)
ptest = TestFunction(Qlevel)
if case == 4:
Wlevel = assemble(Tensor(inner(ptrial, ptest) *
dx).inv).M[0, 0].handle
if case == 5:
Wlevel = assemble(Tensor(1.0/mu(levelmesh)*inner(ptrial, ptest)*\
dx).inv).M[0,0].handle
# Form BTWB
BTWlevel = Blevel.transposeMatMult(Wlevel)
BTWlevel *= args.gamma
BTWBlevel = BTWlevel.matMult(Blevel)
J.axpy(1, BTWBlevel, structure=J.Structure.SUBSET_NONZERO_PATTERN)
def assemble_mixed_mass_matrix(V_A, V_B):
"""
Construct the mixed mass matrix of two function spaces,
using the TrialFunction from V_A and the TestFunction
from V_B.
"""
if len(V_A) > 1 or len(V_B) > 1:
raise NotImplementedError(
"Sorry, only implemented for non-mixed spaces")
if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping(
) != "identity":
msg = """
Sorry, only implemented for affine maps for now. To do non-affine, we'd need to
import much more of the assembly engine of UFL/TSFC/etc to do the assembly on
each supermesh cell.
"""
raise NotImplementedError(msg)
mesh_A = V_A.mesh()
mesh_B = V_B.mesh()
dim = mesh_A.geometric_dimension()
assert dim == mesh_B.geometric_dimension()
assert dim == mesh_A.topological_dimension()
assert dim == mesh_B.topological_dimension()
(mh_A, level_A) = get_level(mesh_A)
(mh_B, level_B) = get_level(mesh_B)
if mesh_A is mesh_B:
def likely(cell_A):
return [cell_A]
else:
if (mh_A is None or mh_B is None) or (mh_A is not mh_B):
# No mesh hierarchy structure, call libsupermesh for
# intersection finding
intersections = intersection_finder(mesh_A, mesh_B)
likely = intersections.__getitem__
else:
# We do have a mesh hierarchy, use it
if abs(level_A - level_B) > 1:
raise NotImplementedError(
"Only works for transferring between adjacent levels for now."
)
# What are the cells of B that (probably) intersect with a given cell in A?
if level_A > level_B:
cell_map = mh_A.fine_to_coarse_cells[level_A]
def likely(cell_A):
return cell_map[cell_A]
elif level_A < level_B:
cell_map = mh_A.coarse_to_fine_cells[level_A]
def likely(cell_A):
return cell_map[cell_A]
assert V_A.value_size == V_B.value_size
orig_value_size = V_A.value_size
if V_A.value_size > 1:
V_A = firedrake.FunctionSpace(mesh_A,
V_A.ufl_element().sub_elements()[0])
if V_B.value_size > 1:
V_B = firedrake.FunctionSpace(mesh_B,
V_B.ufl_element().sub_elements()[0])
assert V_A.value_size == 1
assert V_B.value_size == 1
preallocator = PETSc.Mat().create(comm=mesh_A.comm)
preallocator.setType(PETSc.Mat.Type.PREALLOCATOR)
rset = V_B.dof_dset
cset = V_A.dof_dset
nrows = rset.layout_vec.getSizes()
ncols = cset.layout_vec.getSizes()
preallocator.setLGMap(rmap=rset.scalar_lgmap, cmap=cset.scalar_lgmap)
preallocator.setSizes(size=(nrows, ncols), bsize=1)
preallocator.setUp()
zeros = numpy.zeros((V_B.cell_node_map().arity, V_A.cell_node_map().arity),
dtype=ScalarType)
for cell_A, dofs_A in enumerate(V_A.cell_node_map().values):
for cell_B in likely(cell_A):
dofs_B = V_B.cell_node_map().values_with_halo[cell_B, :]
preallocator.setValuesLocal(dofs_B, dofs_A, zeros)
preallocator.assemble()
dnnz, onnz = get_preallocation(preallocator, nrows[0])
# Unroll from block to AIJ
dnnz = dnnz * cset.cdim
dnnz = numpy.repeat(dnnz, rset.cdim)
onnz = onnz * cset.cdim
onnz = numpy.repeat(onnz, cset.cdim)
preallocator.destroy()
assert V_A.value_size == V_B.value_size
rdim = V_B.dof_dset.cdim
cdim = V_A.dof_dset.cdim
#
# Preallocate M_AB.
#
mat = PETSc.Mat().create(comm=mesh_A.comm)
mat.setType(PETSc.Mat.Type.AIJ)
rsizes = tuple(n * rdim for n in nrows)
csizes = tuple(c * cdim for c in ncols)
mat.setSizes(size=(rsizes, csizes), bsize=(rdim, cdim))
mat.setPreallocationNNZ((dnnz, onnz))
mat.setLGMap(rmap=rset.lgmap, cmap=cset.lgmap)
# TODO: Boundary conditions not handled.
mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False)
mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True)
mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True)
mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, False)
mat.setOption(mat.Option.IGNORE_ZERO_ENTRIES, True)
mat.setUp()
evaluate_kernel_A = compile_element(ufl.Coefficient(V_A),
name="evaluate_kernel_A")
evaluate_kernel_B = compile_element(ufl.Coefficient(V_B),
name="evaluate_kernel_B")
# We only need one of these since we assume that the two meshes both have CG1 coordinates
to_reference_kernel = to_reference_coordinates(
mesh_A.coordinates.ufl_element())
if dim == 2:
reference_mesh = UnitTriangleMesh(comm=COMM_SELF)
else:
reference_mesh = UnitTetrahedronMesh(comm=COMM_SELF)
evaluate_kernel_S = compile_element(ufl.Coefficient(
reference_mesh.coordinates.function_space()),
name="evaluate_kernel_S")
V_S_A = FunctionSpace(reference_mesh, V_A.ufl_element())
V_S_B = FunctionSpace(reference_mesh, V_B.ufl_element())
M_SS = assemble(inner(TrialFunction(V_S_A), TestFunction(V_S_B)) * dx)
M_SS = M_SS.M.handle[:, :]
node_locations_A = utils.physical_node_locations(
V_S_A).dat.data_ro_with_halos
node_locations_B = utils.physical_node_locations(
V_S_B).dat.data_ro_with_halos
num_nodes_A = node_locations_A.shape[0]
num_nodes_B = node_locations_B.shape[0]
to_reference_kernel = to_reference_coordinates(
mesh_A.coordinates.ufl_element())
supermesh_kernel_str = """
#include "libsupermesh-c.h"
#include <petsc.h>
%(to_reference)s
%(evaluate_S)s
%(evaluate_A)s
%(evaluate_B)s
#define complex_mode %(complex_mode)s
#define PrintInfo(...) do { if (PetscLogPrintInfo) printf(__VA_ARGS__); } while (0)
static void print_array(PetscScalar *arr, int d)
{
for(int j=0; j<d; j++)
PrintInfo(stderr, "%%+.2f ", arr[j]);
}
static void print_coordinates(PetscScalar *simplex, int d)
{
for(int i=0; i<d+1; i++)
{
PrintInfo("\t");
print_array(&simplex[d*i], d);
PrintInfo("\\n");
}
}
#if complex_mode
static void seperate_real_and_imag(PetscScalar *simplex, double *real_simplex, double *imag_simplex, int d)
{
for(int i=0; i<d+1; i++)
{
for(int j=0; j<d; j++)
{
real_simplex[d*i+j] = creal(simplex[d*i+j]);
imag_simplex[d*i+j] = cimag(simplex[d*i+j]);
}
}
}
static void merge_back_to_simplex(PetscScalar* simplex, double* real_simplex, double* imag_simplex, int d)
{
print_coordinates(simplex,d);
for(int i=0; i<d+1; i++)
{
for(int j=0; j<d; j++)
{
simplex[d*i+j] = real_simplex[d*i+j]+imag_simplex[d*i+j]*_Complex_I;
}
}
}
#endif
int supermesh_kernel(PetscScalar* simplex_A, PetscScalar* simplex_B, PetscScalar* simplices_C, PetscScalar* nodes_A, PetscScalar* nodes_B, PetscScalar* M_SS, PetscScalar* outptr, int num_ele)
{
#define d %(dim)s
#define num_nodes_A %(num_nodes_A)s
#define num_nodes_B %(num_nodes_B)s
double simplex_ref_measure;
PrintInfo("simplex_A coordinates\\n");
print_coordinates(simplex_A, d);
PrintInfo("simplex_B coordinates\\n");
print_coordinates(simplex_B, d);
int num_elements = num_ele;
if (d == 2) simplex_ref_measure = 0.5;
else if (d == 3) simplex_ref_measure = 1.0/6;
PetscScalar R_AS[num_nodes_A][num_nodes_A];
PetscScalar R_BS[num_nodes_B][num_nodes_B];
PetscScalar coeffs_A[%(num_nodes_A)s] = {0.};
PetscScalar coeffs_B[%(num_nodes_B)s] = {0.};
PetscScalar reference_nodes_A[num_nodes_A][d];
PetscScalar reference_nodes_B[num_nodes_B][d];
#if complex_mode
double real_simplex_A[d*(d+1)];
double imag_simplex_A[d*(d+1)];
seperate_real_and_imag(simplex_A, real_simplex_A, imag_simplex_A, d);
double real_simplex_B[d*(d+1)];
double imag_simplex_B[d*(d+1)];
seperate_real_and_imag(simplex_B, real_simplex_B, imag_simplex_B, d);
double real_simplices_C[num_elements*d*(d+1)];
double imag_simplices_C[num_elements*d*(d+1)];
for (int ii=0; ii<num_elements*d*(d+1); ++ii) imag_simplices_C[ii] = 0.;
%(libsupermesh_intersect_simplices)s(real_simplex_A, real_simplex_B, real_simplices_C, &num_elements);
merge_back_to_simplex(simplex_A, real_simplex_A, imag_simplex_A, d);
merge_back_to_simplex(simplex_B, real_simplex_B, imag_simplex_B, d);
for(int s=0; s<num_elements; s++)
{
PetscScalar* simplex_C = &simplices_C[s * d * (d+1)];
double* real_simplex_C = &real_simplices_C[s * d * (d+1)];
double* imag_simplex_C = &imag_simplices_C[s * d * (d+1)];
merge_back_to_simplex(simplex_C, real_simplex_C, imag_simplex_C, d);
}
#else
%(libsupermesh_intersect_simplices)s(simplex_A, simplex_B, simplices_C, &num_elements);
#endif
PrintInfo("Supermesh consists of %%i elements\\n", num_elements);
// would like to do this
//PetscScalar MAB[%(num_nodes_A)s][%(num_nodes_B)s] = (PetscScalar (*)[%(num_nodes_B)s])outptr;
// but have to do this instead because we don't grok C
PetscScalar (*MAB)[num_nodes_A] = (PetscScalar (*)[num_nodes_A])outptr;
PetscScalar (*MSS)[num_nodes_A] = (PetscScalar (*)[num_nodes_A])M_SS; // note the underscore
for ( int i = 0; i < num_nodes_B; i++ ) {
for (int j = 0; j < num_nodes_A; j++) {
MAB[i][j] = 0.0;
}
}
for(int s=0; s<num_elements; s++)
{
PetscScalar* simplex_S = &simplices_C[s * d * (d+1)];
double simplex_S_measure;
#if complex_mode
double real_simplex_S[d*(d+1)];
double imag_simplex_S[d*(d+1)];
seperate_real_and_imag(simplex_S, real_simplex_S, imag_simplex_S, d);
%(libsupermesh_simplex_measure)s(real_simplex_S, &simplex_S_measure);
merge_back_to_simplex(simplex_S, real_simplex_S, imag_simplex_S, d);
#else
%(libsupermesh_simplex_measure)s(simplex_S, &simplex_S_measure);
#endif
PrintInfo("simplex_S coordinates with measure %%f\\n", simplex_S_measure);
print_coordinates(simplex_S, d);
PrintInfo("Start mapping nodes for V_A\\n");
PetscScalar physical_nodes_A[num_nodes_A][d];
for(int n=0; n < num_nodes_A; n++) {
PetscScalar* reference_node_location = &nodes_A[n*d];
PetscScalar* physical_node_location = physical_nodes_A[n];
for (int j=0; j < d; j++) physical_node_location[j] = 0.0;
pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location);
PrintInfo("\\tNode ");
print_array(reference_node_location, d);
PrintInfo(" mapped to ");
print_array(physical_node_location, d);
PrintInfo("\\n");
}
PrintInfo("Start mapping nodes for V_B\\n");
PetscScalar physical_nodes_B[num_nodes_B][d];
for(int n=0; n < num_nodes_B; n++) {
PetscScalar* reference_node_location = &nodes_B[n*d];
PetscScalar* physical_node_location = physical_nodes_B[n];
for (int j=0; j < d; j++) physical_node_location[j] = 0.0;
pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location);
PrintInfo("\\tNode ");
print_array(reference_node_location, d);
PrintInfo(" mapped to ");
print_array(physical_node_location, d);
PrintInfo("\\n");
}
PrintInfo("==========================================================\\n");
PrintInfo("Start pulling back dof from S into reference space for A.\\n");
for(int n=0; n < num_nodes_A; n++) {
for(int i=0; i<d; i++) reference_nodes_A[n][i] = 0.;
to_reference_coords_kernel(reference_nodes_A[n], physical_nodes_A[n], simplex_A);
PrintInfo("Pulling back ");
print_array(physical_nodes_A[n], d);
PrintInfo(" to ");
print_array(reference_nodes_A[n], d);
PrintInfo("\\n");
}
PrintInfo("Start pulling back dof from S into reference space for B.\\n");
for(int n=0; n < num_nodes_B; n++) {
for(int i=0; i<d; i++) reference_nodes_B[n][i] = 0.;
to_reference_coords_kernel(reference_nodes_B[n], physical_nodes_B[n], simplex_B);
PrintInfo("Pulling back ");
print_array(physical_nodes_B[n], d);
PrintInfo(" to ");
print_array(reference_nodes_B[n], d);
PrintInfo("\\n");
}
PrintInfo("Start evaluating basis functions of V_A at dofs for V_A on S\\n");
for(int i=0; i<num_nodes_A; i++) {
coeffs_A[i] = 1.;
for(int j=0; j<num_nodes_A; j++) {
R_AS[i][j] = 0.;
pyop2_kernel_evaluate_kernel_A(&R_AS[i][j], coeffs_A, reference_nodes_A[j]);
}
print_array(R_AS[i], num_nodes_A);
PrintInfo("\\n");
coeffs_A[i] = 0.;
}
PrintInfo("Start evaluating basis functions of V_B at dofs for V_B on S\\n");
for(int i=0; i<num_nodes_B; i++) {
coeffs_B[i] = 1.;
for(int j=0; j<num_nodes_B; j++) {
R_BS[i][j] = 0.;
pyop2_kernel_evaluate_kernel_B(&R_BS[i][j], coeffs_B, reference_nodes_B[j]);
}
print_array(R_BS[i], num_nodes_B);
PrintInfo("\\n");
coeffs_B[i] = 0.;
}
PrintInfo("Start doing the matmatmat mult\\n");
for ( int i = 0; i < num_nodes_B; i++ ) {
for (int j = 0; j < num_nodes_A; j++) {
for ( int k = 0; k < num_nodes_B; k++) {
for ( int l = 0; l < num_nodes_A; l++) {
MAB[i][j] += (simplex_S_measure/simplex_ref_measure) * R_BS[i][k] * MSS[k][l] * R_AS[j][l];
}
}
}
}
}
return num_elements;
}
""" % {
"evaluate_S":
str(evaluate_kernel_S),
"evaluate_A":
str(evaluate_kernel_A),
"evaluate_B":
str(evaluate_kernel_B),
"to_reference":
str(to_reference_kernel),
"num_nodes_A":
num_nodes_A,
"num_nodes_B":
num_nodes_B,
"libsupermesh_simplex_measure":
"libsupermesh_triangle_area"
if dim == 2 else "libsupermesh_tetrahedron_volume",
"libsupermesh_intersect_simplices":
"libsupermesh_intersect_tris_real"
if dim == 2 else "libsupermesh_intersect_tets_real",
"dim":
dim,
"complex_mode":
1 if complex_mode else 0
}
dirs = get_petsc_dir() + (sys.prefix, )
includes = ["-I%s/include" % d for d in dirs]
libs = ["-L%s/lib" % d for d in dirs]
libs = libs + ["-Wl,-rpath,%s/lib" % d
for d in dirs] + ["-lpetsc", "-lsupermesh"]
lib = load(supermesh_kernel_str,
"c",
"supermesh_kernel",
cppargs=includes,
ldargs=libs,
argtypes=[
ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp,
ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp,
ctypes.c_voidp
],
restype=ctypes.c_int)
ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS,
ctypes.addressof(lib), mat)
if orig_value_size == 1:
return mat
else:
(lrows, grows), (lcols, gcols) = mat.getSizes()
lrows *= orig_value_size
grows *= orig_value_size
lcols *= orig_value_size
gcols *= orig_value_size
size = ((lrows, grows), (lcols, gcols))
context = BlockMatrix(mat, orig_value_size)
blockmat = PETSc.Mat().createPython(size,
context=context,
comm=mat.comm)
blockmat.setUp()
return blockmat
def _dm(self):
from firedrake.mg.utils import get_level
dm = self.dof_dset.dm
_, level = get_level(self.mesh())
dmhooks.attach_hooks(dm, level=level)
return dm
def test_convergence_decerasing_nl():
mesh = UnitSquareMesh(4, 4)
L = 4
# Generate the Mesh Hierarchy and the FunctionSpaceHierarchies
Mesh_Hierarchy = GenerateMeshHierarchy(mesh, L)
FunctionSpaceHierarchies = FunctionSpaceHierarchy(Mesh_Hierarchy, 'CG', 1)
# Create Discretization Object
ensemble_hierarchy = EnsembleHierarchy()
level_to_prolong_to = 5
eps = 4e-2
n = 16
converge = 0
i = 0
Nl = []
while converge == 0:
lvlc = i
lvlf = i + 1
for j in range(n):
u_h = FunctionHierarchy(FunctionSpaceHierarchies)
u = state(u_h[lvlc], u_h[lvlf])
xc = np.random.normal(0.0, 0.05, 1)
u.state[0].interpolate(
Expression(
"exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))",
x_center=xc))
u.state[1].interpolate(
Expression(
"exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))",
x_center=xc))
u.state[0].assign(
Solve(
lvlc,
FunctionSpaceHierarchies,
u.state[0]))
u.state[1].assign(
Solve(
lvlf,
FunctionSpaceHierarchies,
u.state[1]))
u_new = state(
ProlongUpToFinestLevel(
u.state[0],
FunctionHierarchy(FunctionSpaceHierarchies)),
ProlongUpToFinestLevel(
u.state[1],
FunctionHierarchy(FunctionSpaceHierarchies)))
ensemble_hierarchy.AppendToEnsemble(u_new.state, i)
assert get_level(u_new.state[1])[1] == len(Mesh_Hierarchy) - 1
Nl.append(n)
# Calculate sample statistics
SampleStatistics(ensemble_hierarchy)
Ns = OptimalNl(ensemble_hierarchy, eps)
for j in range(i + 1):
if Ns[j] > Nl[j]:
dN = Ns[j] - Nl[j]
for k in range(dN):
u_h = FunctionHierarchy(FunctionSpaceHierarchies)
u = state(u_h[lvlc], u_h[lvlf])
xc = np.random.normal(0.0, 0.05, 1)
u.state[0].interpolate(
Expression(
"exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))",
x_center=xc))
u.state[1].interpolate(
Expression(
"exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))",
x_center=xc))
u.state[0].assign(
Solve(
lvlc,
FunctionSpaceHierarchies,
u.state[0]))
u.state[1].assign(
Solve(
lvlf,
FunctionSpaceHierarchies,
u.state[1]))
u_new = state(
ProlongUpToFinestLevel(
u.state[0],
FunctionHierarchy(FunctionSpaceHierarchies)),
ProlongUpToFinestLevel(
u.state[1],
FunctionHierarchy(FunctionSpaceHierarchies)))
ensemble_hierarchy.AppendToEnsemble(u_new.state, i)
Nl[j] = Ns[j]
# Recalculate sample statistics
SampleStatistics(ensemble_hierarchy)
converge = Convergence(ensemble_hierarchy, eps)
i += 1
assert np.all(np.diff(np.array(Nl)) <= 0) == 1
return ensemble_hierarchy
def aug_jacobian(X, J, ctx):
mh, level = get_level(ctx._x.ufl_domain())
if args.galerkin:
rmap, cmap = J.getLGMap()
levelOps[nref-level].copy(J,structure=J.Structure.DIFFERENT_NONZERO_PATTERN)
J.setLGMap(rmap, cmap)
def assemble_mixed_mass_matrix(V_A, V_B):
"""
Construct the mixed mass matrix of two function spaces,
using the TrialFunction from V_A and the TestFunction
from V_B.
"""
if len(V_A) > 1 or len(V_B) > 1:
raise NotImplementedError("Sorry, only implemented for non-mixed spaces")
if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping() != "identity":
msg = """
Sorry, only implemented for affine maps for now. To do non-affine, we'd need to
import much more of the assembly engine of UFL/TSFC/etc to do the assembly on
each supermesh cell.
"""
raise NotImplementedError(msg)
mesh_A = V_A.mesh()
mesh_B = V_B.mesh()
dim = mesh_A.geometric_dimension()
assert dim == mesh_B.geometric_dimension()
assert dim == mesh_A.topological_dimension()
assert dim == mesh_B.topological_dimension()
(mh_A, level_A) = get_level(mesh_A)
(mh_B, level_B) = get_level(mesh_B)
if mesh_A is mesh_B:
def likely(cell_A):
return [cell_A]
else:
if (mh_A is None or mh_B is None) or (mh_A is not mh_B):
# No mesh hierarchy structure, call libsupermesh for
# intersection finding
intersections = intersection_finder(mesh_A, mesh_B)
likely = intersections.__getitem__
else:
# We do have a mesh hierarchy, use it
if abs(level_A - level_B) > 1:
raise NotImplementedError("Only works for transferring between adjacent levels for now.")
# What are the cells of B that (probably) intersect with a given cell in A?
if level_A > level_B:
cell_map = mh_A.fine_to_coarse_cells[level_A]
def likely(cell_A):
return cell_map[cell_A]
elif level_A < level_B:
cell_map = mh_A.coarse_to_fine_cells[level_A]
def likely(cell_A):
return cell_map[cell_A]
assert V_A.value_size == V_B.value_size
orig_value_size = V_A.value_size
if V_A.value_size > 1:
V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements()[0])
if V_B.value_size > 1:
V_B = firedrake.FunctionSpace(mesh_B, V_B.ufl_element().sub_elements()[0])
assert V_A.value_size == 1
assert V_B.value_size == 1
preallocator = PETSc.Mat().create(comm=mesh_A.comm)
preallocator.setType(PETSc.Mat.Type.PREALLOCATOR)
rset = V_B.dof_dset
cset = V_A.dof_dset
nrows = rset.layout_vec.getSizes()
ncols = cset.layout_vec.getSizes()
preallocator.setLGMap(rmap=rset.scalar_lgmap, cmap=cset.scalar_lgmap)
preallocator.setSizes(size=(nrows, ncols), bsize=1)
preallocator.setUp()
zeros = numpy.zeros((V_B.cell_node_map().arity, V_A.cell_node_map().arity), dtype=ScalarType)
for cell_A, dofs_A in enumerate(V_A.cell_node_map().values):
for cell_B in likely(cell_A):
dofs_B = V_B.cell_node_map().values_with_halo[cell_B, :]
preallocator.setValuesLocal(dofs_B, dofs_A, zeros)
preallocator.assemble()
dnnz, onnz = get_preallocation(preallocator, nrows[0])
# Unroll from block to AIJ
dnnz = dnnz * cset.cdim
dnnz = numpy.repeat(dnnz, rset.cdim)
onnz = onnz * cset.cdim
onnz = numpy.repeat(onnz, cset.cdim)
preallocator.destroy()
assert V_A.value_size == V_B.value_size
rdim = V_B.dof_dset.cdim
cdim = V_A.dof_dset.cdim
#
# Preallocate M_AB.
#
mat = PETSc.Mat().create(comm=mesh_A.comm)
mat.setType(PETSc.Mat.Type.AIJ)
rsizes = tuple(n * rdim for n in nrows)
csizes = tuple(c * cdim for c in ncols)
mat.setSizes(size=(rsizes, csizes),
bsize=(rdim, cdim))
mat.setPreallocationNNZ((dnnz, onnz))
mat.setLGMap(rmap=rset.lgmap, cmap=cset.lgmap)
# TODO: Boundary conditions not handled.
mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False)
mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True)
mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True)
mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, False)
mat.setOption(mat.Option.IGNORE_ZERO_ENTRIES, True)
mat.setUp()
evaluate_kernel_A = compile_element(ufl.Coefficient(V_A), name="evaluate_kernel_A")
evaluate_kernel_B = compile_element(ufl.Coefficient(V_B), name="evaluate_kernel_B")
# We only need one of these since we assume that the two meshes both have CG1 coordinates
to_reference_kernel = to_reference_coordinates(mesh_A.coordinates.ufl_element())
if dim == 2:
reference_mesh = UnitTriangleMesh(comm=COMM_SELF)
else:
reference_mesh = UnitTetrahedronMesh(comm=COMM_SELF)
evaluate_kernel_S = compile_element(ufl.Coefficient(reference_mesh.coordinates.function_space()), name="evaluate_kernel_S")
V_S_A = FunctionSpace(reference_mesh, V_A.ufl_element())
V_S_B = FunctionSpace(reference_mesh, V_B.ufl_element())
M_SS = assemble(inner(TrialFunction(V_S_A), TestFunction(V_S_B)) * dx)
M_SS.force_evaluation()
M_SS = M_SS.M.handle[:, :]
node_locations_A = utils.physical_node_locations(V_S_A).dat.data_ro_with_halos
node_locations_B = utils.physical_node_locations(V_S_B).dat.data_ro_with_halos
num_nodes_A = node_locations_A.shape[0]
num_nodes_B = node_locations_B.shape[0]
to_reference_kernel = to_reference_coordinates(mesh_A.coordinates.ufl_element())
supermesh_kernel_str = """
#include "libsupermesh-c.h"
#include <petsc.h>
%(to_reference)s
%(evaluate_S)s
%(evaluate_A)s
%(evaluate_B)s
#define PrintInfo(...) do { if (PetscLogPrintInfo) printf(__VA_ARGS__); } while (0)
static void print_array(double *arr, int d)
{
for(int j=0; j<d; j++)
PrintInfo("%%+.2f ", arr[j]);
}
static void print_coordinates(double *simplex, int d)
{
for(int i=0; i<d+1; i++)
{
PrintInfo("\t");
print_array(&simplex[d*i], d);
PrintInfo("\\n");
}
}
int supermesh_kernel(double* simplex_A, double* simplex_B, double* simplices_C, double* nodes_A, double* nodes_B, double* M_SS, double* outptr)
{
#define d %(dim)s
#define num_nodes_A %(num_nodes_A)s
#define num_nodes_B %(num_nodes_B)s
double simplex_ref_measure;
PrintInfo("simplex_A coordinates\\n");
print_coordinates(simplex_A, d);
PrintInfo("simplex_B coordinates\\n");
print_coordinates(simplex_B, d);
int num_elements;
if (d == 2) simplex_ref_measure = 0.5;
else if (d == 3) simplex_ref_measure = 1.0/6;
double R_AS[num_nodes_A][num_nodes_A];
double R_BS[num_nodes_B][num_nodes_B];
double coeffs_A[%(num_nodes_A)s] = {0.};
double coeffs_B[%(num_nodes_B)s] = {0.};
double reference_nodes_A[num_nodes_A][d];
double reference_nodes_B[num_nodes_B][d];
%(libsupermesh_intersect_simplices)s(simplex_A, simplex_B, simplices_C, &num_elements);
PrintInfo("Supermesh consists of %%i elements\\n", num_elements);
// would like to do this
//double MAB[%(num_nodes_A)s][%(num_nodes_B)s] = (double (*)[%(num_nodes_B)s])outptr;
// but have to do this instead because we don't grok C
double (*MAB)[num_nodes_A] = (double (*)[num_nodes_A])outptr;
double (*MSS)[num_nodes_A] = (double (*)[num_nodes_A])M_SS; // note the underscore
for ( int i = 0; i < num_nodes_B; i++ ) {
for (int j = 0; j < num_nodes_A; j++) {
MAB[i][j] = 0;
}
}
for(int s=0; s<num_elements; s++)
{
double* simplex_S = &simplices_C[s * d * (d+1)];
double simplex_S_measure;
%(libsupermesh_simplex_measure)s(simplex_S, &simplex_S_measure);
PrintInfo("simplex_S coordinates with measure %%f\\n", simplex_S_measure);
print_coordinates(simplex_S, d);
PrintInfo("Start mapping nodes for V_A\\n");
double physical_nodes_A[num_nodes_A][d];
for(int n=0; n < num_nodes_A; n++) {
double* reference_node_location = &nodes_A[n*d];
double* physical_node_location = physical_nodes_A[n];
for (int j=0; j < d; j++) physical_node_location[j] = 0.0;
pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location);
PrintInfo("\\tNode ");
print_array(reference_node_location, d);
PrintInfo(" mapped to ");
print_array(physical_node_location, d);
PrintInfo("\\n");
}
PrintInfo("Start mapping nodes for V_B\\n");
double physical_nodes_B[num_nodes_B][d];
for(int n=0; n < num_nodes_B; n++) {
double* reference_node_location = &nodes_B[n*d];
double* physical_node_location = physical_nodes_B[n];
for (int j=0; j < d; j++) physical_node_location[j] = 0.0;
pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location);
PrintInfo("\\tNode ");
print_array(reference_node_location, d);
PrintInfo(" mapped to ");
print_array(physical_node_location, d);
PrintInfo("\\n");
}
PrintInfo("==========================================================\\n");
PrintInfo("Start pulling back dof from S into reference space for A.\\n");
for(int n=0; n < num_nodes_A; n++) {
for(int i=0; i<d; i++) reference_nodes_A[n][i] = 0.;
to_reference_coords_kernel(reference_nodes_A[n], physical_nodes_A[n], simplex_A);
PrintInfo("Pulling back ");
print_array(physical_nodes_A[n], d);
PrintInfo(" to ");
print_array(reference_nodes_A[n], d);
PrintInfo("\\n");
}
PrintInfo("Start pulling back dof from S into reference space for B.\\n");
for(int n=0; n < num_nodes_B; n++) {
for(int i=0; i<d; i++) reference_nodes_B[n][i] = 0.;
to_reference_coords_kernel(reference_nodes_B[n], physical_nodes_B[n], simplex_B);
PrintInfo("Pulling back ");
print_array(physical_nodes_B[n], d);
PrintInfo(" to ");
print_array(reference_nodes_B[n], d);
PrintInfo("\\n");
}
PrintInfo("Start evaluating basis functions of V_A at dofs for V_A on S\\n");
for(int i=0; i<num_nodes_A; i++) {
coeffs_A[i] = 1.;
for(int j=0; j<num_nodes_A; j++) {
R_AS[i][j] = 0.;
pyop2_kernel_evaluate_kernel_A(&R_AS[i][j], coeffs_A, reference_nodes_A[j]);
}
print_array(R_AS[i], num_nodes_A);
PrintInfo("\\n");
coeffs_A[i] = 0.;
}
PrintInfo("Start evaluating basis functions of V_B at dofs for V_B on S\\n");
for(int i=0; i<num_nodes_B; i++) {
coeffs_B[i] = 1.;
for(int j=0; j<num_nodes_B; j++) {
R_BS[i][j] = 0.;
pyop2_kernel_evaluate_kernel_B(&R_BS[i][j], coeffs_B, reference_nodes_B[j]);
}
print_array(R_BS[i], num_nodes_B);
PrintInfo("\\n");
coeffs_B[i] = 0.;
}
PrintInfo("Start doing the matmatmat mult\\n");
for ( int i = 0; i < num_nodes_B; i++ ) {
for (int j = 0; j < num_nodes_A; j++) {
for ( int k = 0; k < num_nodes_B; k++) {
for ( int l = 0; l < num_nodes_A; l++) {
MAB[i][j] += (simplex_S_measure/simplex_ref_measure) * R_BS[i][k] * MSS[k][l] * R_AS[j][l];
}
}
}
}
}
return num_elements;
}
""" % {
"evaluate_S": str(evaluate_kernel_S),
"evaluate_A": str(evaluate_kernel_A),
"evaluate_B": str(evaluate_kernel_B),
"to_reference": str(to_reference_kernel),
"num_nodes_A": num_nodes_A,
"num_nodes_B": num_nodes_B,
"value_size_A": V_A.value_size,
"value_size_B": V_B.value_size,
"libsupermesh_simplex_measure": "libsupermesh_triangle_area" if dim == 2 else "libsupermesh_tetrahedron_volume",
"libsupermesh_intersect_simplices": "libsupermesh_intersect_tris_real" if dim == 2 else "libsupermesh_intersect_tets_real",
"dim": dim
}
dirs = get_petsc_dir() + (os.environ["VIRTUAL_ENV"], )
includes = ["-I%s/include" % d for d in dirs]
libs = ["-L%s/lib" % d for d in dirs]
libs = libs + ["-Wl,-rpath,%s/lib" % d for d in dirs] + ["-lpetsc", "-lsupermesh"]
lib = load(supermesh_kernel_str, "c", "supermesh_kernel",
cppargs=includes,
ldargs=libs,
argtypes=[ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp],
restype=ctypes.c_int)
ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat)
if orig_value_size == 1:
return mat
else:
(lrows, grows), (lcols, gcols) = mat.getSizes()
lrows *= orig_value_size
grows *= orig_value_size
lcols *= orig_value_size
gcols *= orig_value_size
size = ((lrows, grows), (lcols, gcols))
context = BlockMatrix(mat, orig_value_size)
blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm)
blockmat.setUp()
return blockmat
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 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 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"))
