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 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
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
transfer = get_transfer_operators(dm)
push_transfer_operators(cdm, *transfer)
ctx = get_appctx(dm)
if ctx is not None:
push_appctx(cdm, coarsen(ctx))
# Necessary for MG inside a fieldsplit in a SNES.
cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators)
V._coarse._fine = V
return cdm
def test_transfer_manager_inside_coarsen(sub):
mesh = UnitSquareMesh(1, 1)
mh = MeshHierarchy(mesh, 1)
mesh = mh[-1]
V = FunctionSpace(mesh, "N1curl", 2)
Q = FunctionSpace(mesh, "P", 1)
Z = V * Q
x, y = SpatialCoordinate(mesh)
if sub:
bc_space = Z.sub(0)
else:
bc_space = V
bcdata = project(as_vector([-y, x]), bc_space)
bc = DirichletBC(Z.sub(0), bcdata, "on_boundary")
u = Function(Z)
v = TestFunction(Z)
F = inner(u, v) * dx
problem = NonlinearVariationalProblem(F, u, bcs=bc)
solver = NonlinearVariationalSolver(problem)
with dmhooks.add_hooks(Z.dm, solver, appctx=solver._ctx):
cctx = coarsen(solver._ctx, coarsen)
bc, = cctx._problem.bcs
V = bc.function_space()
mesh = V.ufl_domain()
x, y = SpatialCoordinate(mesh)
expect = project(as_vector([-y, x]), V)
assert numpy.allclose(bc.function_arg.dat.data_ro, expect.dat.data_ro)
def test_transfers_transferring():
mesh = UnitIntervalMesh(2)
mh = MeshHierarchy(mesh, 1)
mesh = mh[-1]
Vf = FunctionSpace(mesh, "P", 1)
Qf = FunctionSpace(mesh, "DP", 0)
Wf = Vf * Qf
Vcount = count()
Qcount = count()
Wcount = count()
def myinject(f, c, counter=None):
next(counter)
inject(f, c)
injectQ = partial(myinject, counter=Qcount)
injectV = partial(myinject, counter=Vcount)
injectW = partial(myinject, counter=Wcount)
f = Function(Wf)
f.assign(1)
with dmhooks.transfer_operators(Wf, inject=injectW):
c = coarsen(f, coarsen)
assert next(Wcount) == 1
for d in c.split():
assert numpy.allclose(d.dat.data_ro, 1)
f.sub(0).assign(2)
f.sub(1).assign(3)
with dmhooks.transfer_operators(
Wf.sub(0),
inject=injectV), dmhooks.transfer_operators(Wf.sub(1),
inject=injectQ):
c = coarsen(f, coarsen)
assert next(Vcount) == 1
assert next(Qcount) == 1
assert numpy.allclose(c.sub(0).dat.data_ro, 2)
assert numpy.allclose(c.sub(1).dat.data_ro, 3)
def test_transfers_transferring():
mesh = UnitIntervalMesh(2)
mh = MeshHierarchy(mesh, 1)
mesh = mh[-1]
Vf = FunctionSpace(mesh, "P", 1)
Qf = FunctionSpace(mesh, "DP", 0)
Wf = Vf*Qf
Vcount = count()
Qcount = count()
Wcount = count()
def myinject(f, c, counter=None):
next(counter)
inject(f, c)
injectQ = partial(myinject, counter=Qcount)
injectV = partial(myinject, counter=Vcount)
injectW = partial(myinject, counter=Wcount)
f = Function(Wf)
f.assign(1)
with dmhooks.transfer_operators(Wf, inject=injectW):
c = coarsen(f, coarsen)
assert next(Wcount) == 1
for d in c.split():
assert numpy.allclose(d.dat.data_ro, 1)
f.sub(0).assign(2)
f.sub(1).assign(3)
with dmhooks.transfer_operators(Wf.sub(0), inject=injectV), dmhooks.transfer_operators(Wf.sub(1), inject=injectQ):
c = coarsen(f, coarsen)
assert next(Vcount) == 1
assert next(Qcount) == 1
assert numpy.allclose(c.sub(0).dat.data_ro, 2)
assert numpy.allclose(c.sub(1).dat.data_ro, 3)
