def __init__(self, a, row_bcs=[], col_bcs=[], fc_params=None, appctx=None):
self.a = a
self.aT = a.T if isinstance(self.a, slate.TensorBase) else adjoint(a)
self.fc_params = fc_params
self.appctx = appctx
self.row_bcs = row_bcs
self.col_bcs = col_bcs
# create functions from test and trial space to help
# with 1-form assembly
test_space, trial_space = [
a.arguments()[i].function_space() for i in (0, 1)
]
from firedrake import function
self._y = function.Function(test_space)
self._x = function.Function(trial_space)
# These are temporary storage for holding the BC
# values during matvec application. _xbc is for
# the action and ._ybc is for transpose.
if len(self.row_bcs) > 0:
self._xbc = function.Function(trial_space)
if len(self.col_bcs) > 0:
self._ybc = function.Function(test_space)
# Get size information from template vecs on test and trial spaces
trial_vec = trial_space.dof_dset.layout_vec
test_vec = test_space.dof_dset.layout_vec
self.col_sizes = trial_vec.getSizes()
self.row_sizes = test_vec.getSizes()
self.block_size = (test_vec.getBlockSize(), trial_vec.getBlockSize())
if isinstance(self.a, slate.TensorBase):
self.action = self.a * slate.AssembledVector(self._x)
self.actionT = self.aT * slate.AssembledVector(self._y)
else:
self.action = action(self.a, self._x)
self.actionT = action(self.aT, self._y)
from firedrake.assemble import create_assembly_callable
self._assemble_action = create_assembly_callable(
self.action,
tensor=self._y,
form_compiler_parameters=self.fc_params)
self._assemble_actionT = create_assembly_callable(
self.actionT,
tensor=self._x,
form_compiler_parameters=self.fc_params)
def initialize(self, pc):
self.options_prefix = f"{pc.getOptionsPrefix()}custom_primal_dual_"
print(self.options_prefix)
_, P = pc.getOperators()
context = P.getPythonContext()
appctx = self.get_appctx(pc)
test, trial = context.a.arguments()
appctx["V"] = test.function_space()
if "state" in appctx and "velocity_space" in appctx:
appctx["u_k"] = appctx["state"].split()[appctx["velocity_space"]]
problem = appctx["problem"]
linear_form = self.form(appctx, problem)
self.a = self.assemble_ele_schur(linear_form)
AA = Tensor(self.a)
A = AA.blocks
schur = A[1, 0] * A[0, 0].inv * A[0, 1]
Q = A[1, 1]
bcs = appctx["bcs"] if "bcs" in appctx else context.row_bcs
self.A_schur = allocate_matrix(
schur,
bcs=bcs,
form_compiler_parameters=context.fc_params,
mat_type="aij",
options_prefix=self.options_prefix)
self._assemble_form_schur = create_assembly_callable(
schur,
tensor=self.A_schur,
bcs=bcs,
form_compiler_parameters=context.fc_params,
mat_type="aij")
self._assemble_form_schur()
self.ksp_schur = self.setup_ksp(self.A_schur.petscmat, pc)
self.A_Q = allocate_matrix(Q,
bcs=bcs,
form_compiler_parameters=context.fc_params,
mat_type="aij",
options_prefix=self.options_prefix)
self._assemble_form_Q = create_assembly_callable(
Q,
tensor=self.A_Q,
bcs=bcs,
form_compiler_parameters=context.fc_params,
mat_type="aij")
self._assemble_form_Q()
self.ksp_Q = self.setup_ksp(self.A_Q.petscmat, pc)
def initialize(self, pc):
self.options_prefix = f"{pc.getOptionsPrefix()}custom_{self.prefix}_"
print(self.options_prefix)
_, P = pc.getOperators()
context = P.getPythonContext()
appctx = self.get_appctx(pc)
test, trial = context.a.arguments()
appctx["V"] = test.function_space()
if "state" in appctx and "velocity_space" in appctx:
appctx["u_k"] = appctx["state"].split()[appctx["velocity_space"]]
problem = appctx["problem"]
linear_form = self.form(appctx, problem)
self.a = self.assemble_ele_schur(linear_form)
bcs = appctx["bcs"] if "bcs" in appctx else context.row_bcs
self.A = allocate_matrix(self.a,
bcs=bcs,
form_compiler_parameters=context.fc_params,
mat_type="aij",
options_prefix=self.options_prefix)
self._assemble_form = create_assembly_callable(
self.a,
tensor=self.A,
bcs=bcs,
form_compiler_parameters=context.fc_params,
mat_type="aij")
self._assemble_form()
self.ksp = self.setup_ksp(self.A.petscmat, pc)
def _assemble_pjac(self):
from firedrake.assemble import create_assembly_callable
return create_assembly_callable(self.Jp,
tensor=self._pjac,
bcs=self.bcs_Jp,
form_compiler_parameters=self.fcp,
mat_type=self.pmat_type)
def local_solver_calls(self, A, rhs, x, elim_fields):
"""Provides solver callbacks for inverting local operators
and reconstructing eliminated fields.
:arg A: A Slate Tensor containing the mixed bilinear form.
:arg rhs: A firedrake function for the right-hand side.
:arg x: A firedrake function for the solution.
:arg elim_fields: An iterable of eliminated field indices
to recover.
"""
from firedrake.slate.static_condensation.la_utils import backward_solve
from firedrake.assemble import create_assembly_callable
fields = x.split()
systems = backward_solve(A, rhs, x, reconstruct_fields=elim_fields)
local_solvers = []
for local_system in systems:
Ae = local_system.lhs
be = local_system.rhs
i, = local_system.field_idx
local_solve = Ae.solve(be, decomposition="PartialPivLU")
solve_call = create_assembly_callable(
local_solve,
tensor=fields[i],
form_compiler_parameters=self.cxt.fc_params)
local_solvers.append(solve_call)
return local_solvers
def __init__(self, problem, mat_type, pmat_type, appctx=None,
pre_jacobian_callback=None, pre_function_callback=None,
options_prefix=None):
from firedrake.assemble import create_assembly_callable
if pmat_type is None:
pmat_type = mat_type
self.mat_type = mat_type
self.pmat_type = pmat_type
self.options_prefix = options_prefix
matfree = mat_type == 'matfree'
pmatfree = pmat_type == 'matfree'
self._problem = problem
self._pre_jacobian_callback = pre_jacobian_callback
self._pre_function_callback = pre_function_callback
self.fcp = problem.form_compiler_parameters
# Function to hold current guess
self._x = problem.u
if appctx is None:
appctx = {}
# A split context will already get the full state.
# TODO, a better way of doing this.
# Now we don't have a temporary state inside the snes
# context we could just require the user to pass in the
# full state on the outside.
appctx.setdefault("state", self._x)
appctx.setdefault("form_compiler_parameters", self.fcp)
self.appctx = appctx
self.matfree = matfree
self.pmatfree = pmatfree
self.F = problem.F
self.J = problem.J
if mat_type != pmat_type or problem.Jp is not None:
# Need separate pmat if either Jp is different or we want
# a different pmat type to the mat type.
if problem.Jp is None:
self.Jp = self.J
else:
self.Jp = problem.Jp
else:
# pmat_type == mat_type and Jp is None
self.Jp = None
self._assemble_residual = create_assembly_callable(self.F,
tensor=self._F,
form_compiler_parameters=self.fcp)
self._jacobian_assembled = False
self._splits = {}
self._coarse = None
self._fine = None
self._nullspace = None
self._nullspace_T = None
self._near_nullspace = None
def _assemble_diagonal(self):
from firedrake.assemble import create_assembly_callable
return create_assembly_callable(
self.a,
tensor=self._diagonal,
form_compiler_parameters=self.fc_params,
diagonal=True)
def local_solver_calls(self, A, rhs, x, elim_fields):
"""Provides solver callbacks for inverting local operators
and reconstructing eliminated fields.
:arg A: A Slate Tensor containing the mixed bilinear form.
:arg rhs: A firedrake function for the right-hand side.
:arg x: A firedrake function for the solution.
:arg elim_fields: An iterable of eliminated field indices
to recover.
"""
from firedrake.slate.static_condensation.la_utils import backward_solve
from firedrake.assemble import create_assembly_callable
fields = x.split()
systems = backward_solve(A, rhs, x, reconstruct_fields=elim_fields)
local_solvers = []
for local_system in systems:
Ae = local_system.lhs
be = local_system.rhs
i, = local_system.field_idx
local_solve = Ae.solve(be, decomposition="PartialPivLU")
solve_call = create_assembly_callable(
local_solve,
tensor=fields[i],
form_compiler_parameters=self.cxt.fc_params)
local_solvers.append(solve_call)
return local_solvers
def _reconstruction_calls(self, split_mixed_op, split_trace_op):
"""This generates the reconstruction calls for the unknowns using the
Lagrange multipliers.
:arg split_mixed_op: a ``dict`` of split forms that make up the broken
mixed operator from the original problem.
:arg split_trace_op: a ``dict`` of split forms that make up the trace
contribution in the hybridized mixed system.
"""
from firedrake.assemble import create_assembly_callable
# We always eliminate the velocity block first
id0, id1 = (self.vidx, self.pidx)
# TODO: When PyOP2 is able to write into mixed dats,
# the reconstruction expressions can simplify into
# one clean expression.
A = Tensor(split_mixed_op[(id0, id0)])
B = Tensor(split_mixed_op[(id0, id1)])
C = Tensor(split_mixed_op[(id1, id0)])
D = Tensor(split_mixed_op[(id1, id1)])
K_0 = Tensor(split_trace_op[(0, id0)])
K_1 = Tensor(split_trace_op[(0, id1)])
# Split functions and reconstruct each bit separately
split_residual = self.broken_residual.split()
split_sol = self.broken_solution.split()
g = AssembledVector(split_residual[id0])
f = AssembledVector(split_residual[id1])
sigma = split_sol[id0]
u = split_sol[id1]
lambdar = AssembledVector(self.trace_solution)
M = D - C * A.inv * B
R = K_1.T - C * A.inv * K_0.T
u_rec = M.solve(f - C * A.inv * g - R * lambdar,
decomposition="PartialPivLU")
self._sub_unknown = create_assembly_callable(
u_rec, tensor=u, form_compiler_parameters=self.ctx.fc_params)
sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar,
decomposition="PartialPivLU")
self._elim_unknown = create_assembly_callable(
sigma_rec,
tensor=sigma,
form_compiler_parameters=self.ctx.fc_params)
def assemble_firedrake(dual_lin, bcs=[]):
matrix = allocate_matrix(dual_lin, bcs=bcs, mat_type="aij")
_assemble_form = create_assembly_callable(
dual_lin, tensor=matrix, bcs=bcs, mat_type="aij")
_assemble_form()
ai, aj, av = matrix.petscmat.getValuesCSR()
matrix_scipy = sp.sparse.csr_matrix((av, aj, ai))
return matrix_scipy
def _reconstruction_calls(self, split_mixed_op, split_trace_op):
"""This generates the reconstruction calls for the unknowns using the
Lagrange multipliers.
:arg split_mixed_op: a ``dict`` of split forms that make up the broken
mixed operator from the original problem.
:arg split_trace_op: a ``dict`` of split forms that make up the trace
contribution in the hybridized mixed system.
"""
from firedrake.assemble import create_assembly_callable
# We always eliminate the velocity block first
id0, id1 = (self.vidx, self.pidx)
# TODO: When PyOP2 is able to write into mixed dats,
# the reconstruction expressions can simplify into
# one clean expression.
A = Tensor(split_mixed_op[(id0, id0)])
B = Tensor(split_mixed_op[(id0, id1)])
C = Tensor(split_mixed_op[(id1, id0)])
D = Tensor(split_mixed_op[(id1, id1)])
K_0 = Tensor(split_trace_op[(0, id0)])
K_1 = Tensor(split_trace_op[(0, id1)])
# Split functions and reconstruct each bit separately
split_residual = self.broken_residual.split()
split_sol = self.broken_solution.split()
g = AssembledVector(split_residual[id0])
f = AssembledVector(split_residual[id1])
sigma = split_sol[id0]
u = split_sol[id1]
lambdar = AssembledVector(self.trace_solution)
M = D - C * A.inv * B
R = K_1.T - C * A.inv * K_0.T
u_rec = M.solve(f - C * A.inv * g - R * lambdar,
decomposition="PartialPivLU")
self._sub_unknown = create_assembly_callable(u_rec,
tensor=u,
form_compiler_parameters=self.ctx.fc_params)
sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar,
decomposition="PartialPivLU")
self._elim_unknown = create_assembly_callable(sigma_rec,
tensor=sigma,
form_compiler_parameters=self.ctx.fc_params)
def __init__(self, a, row_bcs=[], col_bcs=[],
fc_params=None, appctx=None):
self.a = a
self.aT = adjoint(a)
self.fc_params = fc_params
self.appctx = appctx
self.row_bcs = row_bcs
self.col_bcs = col_bcs
# create functions from test and trial space to help
# with 1-form assembly
test_space, trial_space = [
a.arguments()[i].function_space() for i in (0, 1)
]
from firedrake import function
self._y = function.Function(test_space)
self._x = function.Function(trial_space)
# These are temporary storage for holding the BC
# values during matvec application. _xbc is for
# the action and ._ybc is for transpose.
if len(self.row_bcs) > 0:
self._xbc = function.Function(trial_space)
if len(self.col_bcs) > 0:
self._ybc = function.Function(test_space)
# Get size information from template vecs on test and trial spaces
trial_vec = trial_space.dof_dset.layout_vec
test_vec = test_space.dof_dset.layout_vec
self.col_sizes = trial_vec.getSizes()
self.row_sizes = test_vec.getSizes()
self.block_size = (test_vec.getBlockSize(), trial_vec.getBlockSize())
self.action = action(self.a, self._x)
self.actionT = action(self.aT, self._y)
from firedrake.assemble import create_assembly_callable
self._assemble_action = create_assembly_callable(self.action, tensor=self._y,
form_compiler_parameters=self.fc_params)
self._assemble_actionT = create_assembly_callable(self.actionT, tensor=self._x,
form_compiler_parameters=self.fc_params)
def _rhs(self):
from firedrake.assemble import create_assembly_callable
u = function.Function(self.trial_space)
b = function.Function(self.test_space)
if isinstance(self.A.a, slate.TensorBase):
expr = -self.A.a * slate.AssembledVector(u)
else:
expr = -ufl.action(self.A.a, u)
return u, create_assembly_callable(expr, tensor=b), b
def initialize(self, pc):
from firedrake.assemble import allocate_matrix, create_assembly_callable
_, P = pc.getOperators()
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
opc = pc
context = P.getPythonContext()
prefix = pc.getOptionsPrefix()
options_prefix = prefix + "assembled_"
# It only makes sense to preconditioner/invert a diagonal
# block in general. That's all we're going to allow.
if not context.on_diag:
raise ValueError("Only makes sense to invert diagonal block")
mat_type = PETSc.Options().getString(options_prefix + "mat_type",
"aij")
self.P = allocate_matrix(context.a,
bcs=context.row_bcs,
form_compiler_parameters=context.fc_params,
mat_type=mat_type,
options_prefix=options_prefix)
self._assemble_P = create_assembly_callable(
context.a,
tensor=self.P,
bcs=context.row_bcs,
form_compiler_parameters=context.fc_params,
mat_type=mat_type)
self._assemble_P()
self.P.force_evaluation()
# Transfer nullspace over
Pmat = self.P.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
# Internally, we just set up a PC object that the user can configure
# however from the PETSc command line. Since PC allows the user to specify
# a KSP, we can do iterative by -assembled_pc_type ksp.
pc = PETSc.PC().create(comm=opc.comm)
pc.incrementTabLevel(1, parent=opc)
pc.setOptionsPrefix(options_prefix)
pc.setOperators(Pmat, Pmat)
pc.setFromOptions()
pc.setUp()
self.pc = pc
def _build_up_solver(self):
"""Constructs the solver for the velocity-pressure increments."""
from firedrake.assemble import create_assembly_callable
# strong no-slip boundary conditions on the top
# and bottom of the atmospheric domain
bcs = [
DirichletBC(self._Wmixed.sub(0), 0.0, "bottom"),
DirichletBC(self._Wmixed.sub(0), 0.0, "top")
]
un, pn, bn = self._state.split()
utest, ptest = TestFunctions(self._Wmixed)
u, p = TrialFunctions(self._Wmixed)
def outward(arg):
return cross(self._khat, arg)
# Linear gravity wave system for the velocity and pressure
# increments (buoyancy has been eliminated in the discrete
# equations since there is no orography)
a_up = (
ptest * p + self._dt_half_c2 * ptest * div(u) -
self._dt_half * div(utest) * p +
(dot(utest, u) + self._dt_half * dot(utest, self._f * outward(u)) +
self._omega_N2 * dot(utest, self._khat) * dot(u, self._khat))
) * dx
L_up = (dot(utest, un) +
self._dt_half * dot(utest, self._f * outward(un)) +
self._dt_half * dot(utest, self._khat * bn) + ptest * pn) * dx
# Set up linear solver
up_problem = LinearVariationalProblem(a_up, L_up, self._up, bcs=bcs)
params = self._params
solver = LinearVariationalSolver(up_problem,
solver_parameters=params,
options_prefix='up-implicit-solver')
self.up_solver = solver
r = action(a_up, self._up) - L_up
self._assemble_upr = create_assembly_callable(r,
tensor=self._up_residual)
def assembleK(self):
ctx = self.ctx
mat_type = PETSc.Options().getString(
self.prefix + "assembled_mat_type", "aij")
self.K = allocate_matrix(ctx.a,
bcs=ctx.row_bcs,
form_compiler_parameters=ctx.fc_params,
mat_type=mat_type)
self._assemble_K = create_assembly_callable(
ctx.a,
tensor=self.K,
bcs=ctx.row_bcs,
form_compiler_parameters=ctx.fc_params,
mat_type=mat_type)
self._assemble_K()
self.mat_type = mat_type
self.K.force_evaluation()
def initialize(self, pc):
from firedrake.assemble import allocate_matrix, create_assembly_callable
_, P = pc.getOperators()
context = P.getPythonContext()
prefix = pc.getOptionsPrefix()
# It only makes sense to preconditioner/invert a diagonal
# block in general. That's all we're going to allow.
if not context.on_diag:
raise ValueError("Only makes sense to invert diagonal block")
mat_type = PETSc.Options().getString(prefix + "assembled_mat_type", "aij")
self.P = allocate_matrix(context.a, bcs=context.row_bcs,
form_compiler_parameters=context.fc_params,
mat_type=mat_type)
self._assemble_P = create_assembly_callable(context.a, tensor=self.P,
bcs=context.row_bcs,
form_compiler_parameters=context.fc_params,
mat_type=mat_type)
self._assemble_P()
self.mat_type = mat_type
self.P.force_evaluation()
# Transfer nullspace over
Pmat = self.P.petscmat
Pmat.setNullSpace(P.getNullSpace())
Pmat.setTransposeNullSpace(P.getTransposeNullSpace())
# Internally, we just set up a PC object that the user can configure
# however from the PETSc command line. Since PC allows the user to specify
# a KSP, we can do iterative by -assembled_pc_type ksp.
pc = PETSc.PC().create()
pc.setOptionsPrefix(prefix+"assembled_")
pc.setOperators(Pmat, Pmat)
pc.setUp()
pc.setFromOptions()
self.pc = pc
def initialize(self, pc):
from firedrake.assemble import allocate_matrix, create_assembly_callable
_, P = pc.getOperators()
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
opc = pc
appctx = self.get_appctx(pc)
fcp = appctx.get("form_compiler_parameters")
V = get_function_space(pc.getDM())
if len(V) == 1:
V = FunctionSpace(V.mesh(), V.ufl_element())
else:
V = MixedFunctionSpace([V_ for V_ in V])
test = TestFunction(V)
trial = TrialFunction(V)
if P.type == "python":
context = P.getPythonContext()
# It only makes sense to preconditioner/invert a diagonal
# block in general. That's all we're going to allow.
if not context.on_diag:
raise ValueError("Only makes sense to invert diagonal block")
prefix = pc.getOptionsPrefix()
options_prefix = prefix + self._prefix
mat_type = PETSc.Options().getString(options_prefix + "mat_type", "aij")
(a, bcs) = self.form(pc, test, trial)
self.P = allocate_matrix(a, bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type,
options_prefix=options_prefix)
self._assemble_P = create_assembly_callable(a, tensor=self.P,
bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type)
self._assemble_P()
self.P.force_evaluation()
# Transfer nullspace over
Pmat = self.P.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
# Internally, we just set up a PC object that the user can configure
# however from the PETSc command line. Since PC allows the user to specify
# a KSP, we can do iterative by -assembled_pc_type ksp.
pc = PETSc.PC().create(comm=opc.comm)
pc.incrementTabLevel(1, parent=opc)
# We set a DM and an appropriate SNESContext on the constructed PC so one
# can do e.g. multigrid or patch solves.
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.solving_utils import _SNESContext
dm = opc.getDM()
octx = get_appctx(dm)
oproblem = octx._problem
nproblem = NonlinearVariationalProblem(oproblem.F, oproblem.u, bcs, J=a, form_compiler_parameters=fcp)
nctx = _SNESContext(nproblem, mat_type, mat_type, octx.appctx)
push_appctx(dm, nctx)
pc.setDM(dm)
pc.setOptionsPrefix(options_prefix)
pc.setOperators(Pmat, Pmat)
pc.setFromOptions()
pc.setUp()
self.pc = pc
pop_appctx(dm)
def __init__(self, Vc, Vf):
meshc = Vc.mesh()
meshf = Vf.mesh()
self.Vc = Vc
self.Vf = Vf
P1c = VectorFunctionSpace(meshc, "CG", 1)
FBc = VectorFunctionSpace(meshc, "FacetBubble", 3)
P1f = VectorFunctionSpace(meshf, "CG", 1)
FBf = VectorFunctionSpace(meshf, "FacetBubble", 3)
self.p1c = Function(P1c)
self.fbc = Function(FBc)
self.p1f = Function(P1f)
self.fbf = Function(FBf)
self.rhs = Function(FBc)
trial = TrialFunction(FBc)
test = TestFunction(FBc)
n = FacetNormal(meshc)
a = inner(inner(trial, n), inner(test, n)) + inner(
trial - inner(trial, n) * n, test - inner(test, n) * n)
a = a * ds + avg(a) * dS
A = assemble(a).M.handle.getDiagonal()
# A is diagonal, so "solve" by dividing by the diagonal.
ainv = A.copy()
ainv.reciprocal()
self.ainv = ainv
L = inner(inner(self.fbc, n) / 0.625, inner(test, n)) + inner(
self.fbc - inner(self.fbc, n) * n, test - inner(test, n) * n)
L = L * ds + avg(L) * dS
# This is the guts of assemble(L, tensor=self.rhs), but saves a little time
self.assemble_rhs = create_assembly_callable(L, tensor=self.rhs)
# TODO: Parameterise these by the element definition.
# These are correct for P1 + FB in 3D.
# Dof layout for the combined element on the reference cell is:
# 4 vertex dofs, then 4 face dofs.
"""
These two kernels perform a change of basis from a nodal to a
hierachichal basis ("split") and then back from a hierachichal
basis to a nodal basis ("combine"). The nodal functionspace is
given by P1FB, the hierachichal one is given by two spaces: P1 and
FB. The splitting kernel is defined so that
[a b] * [p1fb_1; ...;p1fb_8] = [p1_1;...;p1_4;fb_1;...;fb_4]
and the combine kernel is defined so that
[a;b] * [p1_1;...;p1_4;fb_1;...;fb_4] = [p1fb_1;...;p1fb_8]
Notation: [x;y]: vertically stack x and y, [x y]: horizontally stacked
"""
self.split_kernel = op2.Kernel(
"""
void split(double p1[12], double fb[12], const double both[24]) {
// for (int i = 0; i < 12; i++) {
// p1[i] = 0;
// fb[i] = 0;
// }
double a[8][4] = {{1., 0., 0., 0.},
{0., 1., 0., 0.},
{0., 0., 1., 0.},
{0., 0., 0., 1.},
{0., 0., 0., 0.},
{0., 0., 0., 0.},
{0., 0., 0., 0.},
{0., 0., 0., 0.}};
double b[8][4] = {{0., -1./3, -1./3, -1./3},
{-1./3, 0., -1./3, -1./3},
{-1./3, -1./3, 0., -1./3},
{-1./3, -1./3, -1./3, 0.},
{1., 0., 0., 0.},
{0., 1., 0., 0.},
{0., 0., 1., 0.},
{0., 0., 0., 1.}};
for (int k = 0; k < 8; k++) {
for (int i = 0; i < 4; i++) {
for (int j = 0; j < 3; j++) {
p1[3 * i + j] += a[k][i] * both[3 * k + j];
fb[3 * i + j] += b[k][i] * both[3 * k + j];
}
}
}
}""", "split")
self.split_kernel_adj = op2.Kernel(
"""
void splitadj(const double p1[12], const double fb[12], double both[24]) {
//for (int i = 0; i < 24; i++) both[i] = 0;
double a[8][4] = {{1., 0., 0., 0.},
{0., 1., 0., 0.},
{0., 0., 1., 0.},
{0., 0., 0., 1.},
{0., 0., 0., 0.},
{0., 0., 0., 0.},
{0., 0., 0., 0.},
{0., 0., 0., 0.}};
double b[8][4] = {{0., -1./3, -1./3, -1./3},
{-1./3, 0., -1./3, -1./3},
{-1./3, -1./3, 0., -1./3},
{-1./3, -1./3, -1./3, 0.},
{1., 0., 0., 0.},
{0., 1., 0., 0.},
{0., 0., 1., 0.},
{0., 0., 0., 1.}};
for (int k = 0; k < 8; k++) {
for (int i = 0; i < 4; i++) {
for (int j = 0; j < 3; j++) {
both[3 * k + j] += a[k][i] * p1[3*i+j];
both[3 * k + j] += b[k][i] * fb[3*i+j];
}
}
}
}""", "splitadj")
self.combine_kernel = op2.Kernel(
"""
void combine(const double p1[12], const double fb[12], double both[24]) {
//for (int i = 0; i < 24; i++) both[i] = 0;
double a[4][8] = {{1., 0., 0., 0., 0., 1./3., 1./3., 1./3.},
{0., 1., 0., 0., 1./3., 0., 1./3., 1./3.},
{0., 0., 1., 0., 1./3., 1./3., 0., 1./3.},
{0., 0., 0., 1., 1./3., 1./3., 1./3., 0. }};
double b[4][8] = {{0., 0., 0., 0., 1., 0., 0., 0.},
{0., 0., 0., 0., 0., 1., 0., 0.},
{0., 0., 0., 0., 0., 0., 1., 0.},
{0., 0., 0., 0., 0., 0., 0., 1.}};
for (int k = 0; k < 4; k++) {
for (int i = 0; i < 8; i++) {
for (int j = 0; j < 3; j++) {
both[3 * i + j] += a[k][i] * p1[3 * k + j];
both[3 * i + j] += b[k][i] * fb[3 * k + j];
}
}
}
}
""", "combine")
self.combine_kernel_adj = op2.Kernel(
"""
void combineadj(const double both[24], double p1[12], double fb[12]) {
//for (int i = 0; i < 12; i++) {
// p1[i] = 0;
// fb[i] = 0;
//}
double a[4][8] = {{1., 0., 0., 0., 0., 1./3., 1./3., 1./3.},
{0., 1., 0., 0., 1./3., 0., 1./3., 1./3.},
{0., 0., 1., 0., 1./3., 1./3., 0., 1./3.},
{0., 0., 0., 1., 1./3., 1./3., 1./3., 0. }};
double b[4][8] = {{0., 0., 0., 0., 1., 0., 0., 0.},
{0., 0., 0., 0., 0., 1., 0., 0.},
{0., 0., 0., 0., 0., 0., 1., 0.},
{0., 0., 0., 0., 0., 0., 0., 1.}};
for (int k = 0; k < 4; k++) {
for (int i = 0; i < 8; i++) {
for (int j = 0; j < 3; j++) {
p1[3 * k + j] += a[k][i] * both[3 * i + j];
fb[3 * k + j] += b[k][i] * both[3 * i + j];
}
}
}
}
""", "combineadj")
self.count_kernel = op2.Kernel(
"""
void count(double both[24], double fb[12], double p1[12]) {
for (int i = 0; i < 24; i++) {
both[i] += 1;
}
for (int i = 0; i < 12; i++) {
fb[i] += 1;
p1[i] += 1;
}
}
""", "count")
self.countp1f = Function(self.p1f.function_space())
self.countfbf = Function(self.fbf.function_space())
self.countvf = Function(self.Vf)
self.countp1c = Function(self.p1c.function_space())
self.countfbc = Function(self.fbc.function_space())
self.countvc = Function(self.Vc)
op2.par_loop(self.count_kernel,
self.countp1f.ufl_domain().cell_set,
self.countvf.dat(op2.INC, self.countvf.cell_node_map()),
self.countfbf.dat(op2.INC, self.countfbf.cell_node_map()),
self.countp1f.dat(op2.INC, self.countp1f.cell_node_map()))
op2.par_loop(self.count_kernel,
self.countp1c.ufl_domain().cell_set,
self.countvc.dat(op2.INC, self.countvc.cell_node_map()),
self.countfbc.dat(op2.INC, self.countfbc.cell_node_map()),
self.countp1c.dat(op2.INC, self.countp1c.cell_node_map()))
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from ufl.algorithms.map_integrands import map_integrand_dags
from firedrake import (FunctionSpace, TrialFunction, TrialFunctions,
TestFunction, Function, BrokenElement,
MixedElement, FacetNormal, Constant,
DirichletBC, Projector)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import ArgumentReplacer, split_form
# Extract the problem context
prefix = pc.getOptionsPrefix()
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
V = test.function_space()
if V.mesh().cell_set._extruded:
# TODO: Merge FIAT branch to support TPC trace elements
raise NotImplementedError("Not implemented on extruded meshes.")
# Break the function spaces and define fully discontinuous spaces
broken_elements = [BrokenElement(Vi.ufl_element()) for Vi in V]
elem = MixedElement(broken_elements)
V_d = FunctionSpace(V.mesh(), elem)
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
# Replace the problems arguments with arguments defined
# on the new discontinuous spaces
replacer = ArgumentReplacer(arg_map)
new_form = map_integrand_dags(replacer, context.a)
# Create the space of approximate traces.
# The vector function space will have a non-empty value_shape
W = next(v for v in V if bool(v.ufl_element().value_shape()))
if W.ufl_element().family() in ["Raviart-Thomas", "RTCF"]:
tdegree = W.ufl_element().degree() - 1
else:
tdegree = W.ufl_element().degree()
# NOTE: Once extruded is ready, we will need to be aware of this
# and construct the appropriate trace space for the HDiv element
TraceSpace = FunctionSpace(V.mesh(), "HDiv Trace", tdegree)
# NOTE: For extruded, we will need to add "on_top" and "on_bottom"
trace_conditions = [
DirichletBC(TraceSpace, Constant(0.0), "on_boundary")
]
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_rhs = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_rhs = Function(V)
# Create the symbolic Schur-reduction
Atilde = Tensor(new_form)
gammar = TestFunction(TraceSpace)
n = FacetNormal(V.mesh())
# Vector trial function will have a non-empty ufl_shape
sigma = next(f for f in TrialFunctions(V_d) if bool(f.ufl_shape))
# NOTE: Once extruded is ready, this will change slightly
# to include both horizontal and vertical interior facets
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_rhs,
tensor=self.schur_rhs,
form_compiler_parameters=context.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullsp = P.getNullSpace()
if nullsp.handle != 0:
new_vecs = get_trace_nullspace_vecs(K * Atilde.inv, nullsp, V, V_d,
TraceSpace)
tr_nullsp = PETSc.NullSpace().create(vectors=new_vecs,
comm=pc.comm)
Smat.setNullSpace(tr_nullsp)
# Set up the KSP for the system of Lagrange multipliers
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.setOptionsPrefix(prefix + "trace_")
ksp.setTolerances(rtol=1e-13)
ksp.setOperators(Smat)
ksp.setUp()
ksp.setFromOptions()
self.ksp = ksp
# Now we construct the local tensors for the reconstruction stage
# TODO: Add support for mixed tensors and these variables
# become unnecessary
split_forms = split_form(new_form)
A = Tensor(next(sf.form for sf in split_forms if sf.indices == (0, 0)))
B = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 0)))
C = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 1)))
trial = TrialFunction(
FunctionSpace(V.mesh(), BrokenElement(W.ufl_element())))
K_local = Tensor(gammar('+') * ufl.dot(trial, n) * ufl.dS)
# Split functions and reconstruct each bit separately
sigma_h, u_h = self.broken_solution.split()
g, f = self.broken_rhs.split()
# Pressure reconstruction
M = B * A.inv * B.T + C
u_sol = M.inv * f + M.inv * (
B * A.inv * K_local.T * self.trace_solution - B * A.inv * g)
self._assemble_pressure = create_assembly_callable(
u_sol, tensor=u_h, form_compiler_parameters=context.fc_params)
# Velocity reconstruction
sigma_sol = A.inv * g + A.inv * (B.T * u_h -
K_local.T * self.trace_solution)
self._assemble_velocity = create_assembly_callable(
sigma_sol,
tensor=sigma_h,
form_compiler_parameters=context.fc_params)
# Set up the projector for projecting the broken solution
# into the unbroken finite element spaces
# NOTE: Tolerance here matters!
sigma_b, _ = self.broken_solution.split()
sigma_u, _ = self.unbroken_solution.split()
self.projector = Projector(sigma_b,
sigma_u,
solver_parameters={
"ksp_type": "cg",
"ksp_rtol": 1e-13
})
def __init__(
self,
problem,
mat_type,
pmat_type,
appctx=None,
pre_jacobian_callback=None,
pre_function_callback=None,
post_jacobian_callback=None,
post_function_callback=None,
options_prefix=None,
transfer_manager=None,
):
from firedrake.assemble import create_assembly_callable
from firedrake.bcs import DirichletBC
if pmat_type is None:
pmat_type = mat_type
self.mat_type = mat_type
self.pmat_type = pmat_type
self.options_prefix = options_prefix
matfree = mat_type == "matfree"
pmatfree = pmat_type == "matfree"
self._problem = problem
self._pre_jacobian_callback = pre_jacobian_callback
self._pre_function_callback = pre_function_callback
self._post_jacobian_callback = post_jacobian_callback
self._post_function_callback = post_function_callback
self.fcp = problem.form_compiler_parameters
# Function to hold current guess
self._x = problem.u
# Function to hold time derivative
self._xdot = problem.udot
# Constant to hold time shift value
self._shift = problem.shift
# Constant to hold current time value
self._time = problem.time
if appctx is None:
appctx = {}
# A split context will already get the full state.
# TODO, a better way of doing this.
# Now we don't have a temporary state inside the snes
# context we could just require the user to pass in the
# full state on the outside.
appctx.setdefault("state", self._x)
appctx.setdefault("form_compiler_parameters", self.fcp)
self.appctx = appctx
self.matfree = matfree
self.pmatfree = pmatfree
self.F = problem.F
self.J = problem.J
# For Jp to equal J, bc.Jp must equal bc.J for all EquationBC objects.
Jp_eq_J = problem.Jp is None and all(bc.Jp_eq_J for bc in problem.bcs)
if mat_type != pmat_type or not Jp_eq_J:
# Need separate pmat if either Jp is different or we want
# a different pmat type to the mat type.
if problem.Jp is None:
self.Jp = self.J
else:
self.Jp = problem.Jp
else:
# pmat_type == mat_type and Jp_eq_J
self.Jp = None
self.bcs_F = [
bc if isinstance(bc, DirichletBC) else bc._F for bc in problem.bcs
]
self.bcs_J = [
bc if isinstance(bc, DirichletBC) else bc._J for bc in problem.bcs
]
self.bcs_Jp = [
bc if isinstance(bc, DirichletBC) else bc._Jp for bc in problem.bcs
]
self._assemble_residual = create_assembly_callable(
self.F,
tensor=self._F,
bcs=self.bcs_F,
form_compiler_parameters=self.fcp)
self._jacobian_assembled = False
self._splits = {}
self._coarse = None
self._fine = None
self._nullspace = None
self._nullspace_T = None
self._near_nullspace = None
self._transfer_manager = transfer_manager
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, Function, DirichletBC, dx, \
assemble, Mesh, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
nu = context.appctx["nu"]
dt = context.appctx["dt"]
mass = (1.0 / nu) * p * q * dx
stiffness = inner(grad(p), grad(q)) * dx
velid = context.appctx["velocity_space"]
self.bcs = context.appctx["PCDbc"]
opts = PETSc.Options()
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
bcs=self.bcs,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
u = context.appctx["u"]
fp = (1.0 / nu) * ((1.0 / dt) * p + dot(u, grad(p))) * q * dx
Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=Fp,
form_compiler_parameters=context.fc_params,
mat_type=Fp_mat_type)
self.Fp = Fp
self._assemble_Fp()
self.workspace = [Fp.petscmat.createVecLeft() for i in (0, 1)]
self.tmp = Function(Q)
def gradient(
model, mesh, comm, c, receivers, guess, residual, output=False, save_adjoint=False
):
"""Discrete adjoint with secord-order in time fully-explicit timestepping scheme
with implementation of a Perfectly Matched Layer (PML) using
CG FEM with or without higher order mass lumping (KMV type elements).
Parameters
----------
model: Python `dictionary`
Contains model options and parameters
mesh: Firedrake.mesh object
The 2D/3D triangular mesh
comm: Firedrake.ensemble_communicator
The MPI communicator for parallelism
c: Firedrake.Function
The velocity model interpolated onto the mesh nodes.
receivers: A :class:`spyro.Receivers` object.
Contains the receiver locations and sparse interpolation methods.
guess: A list of Firedrake functions
Contains the forward wavefield at a set of timesteps
residual: array-like [timesteps][receivers]
The difference between the observed and modeled data at
the receivers
output: boolean
optional, write the adjoint to disk (only for debugging)
save_adjoint: A list of Firedrake functions
Contains the adjoint at all timesteps
Returns
-------
dJdc_local: A Firedrake.Function containing the gradient of
the functional w.r.t. `c`
adjoint: Optional, a list of Firedrake functions containing the adjoint
"""
method = model["opts"]["method"]
degree = model["opts"]["degree"]
dim = model["opts"]["dimension"]
dt = model["timeaxis"]["dt"]
tf = model["timeaxis"]["tf"]
nspool = model["timeaxis"]["nspool"]
fspool = model["timeaxis"]["fspool"]
PML = model["BCs"]["status"]
if PML:
Lx = model["mesh"]["Lx"]
Lz = model["mesh"]["Lz"]
lx = model["BCs"]["lx"]
lz = model["BCs"]["lz"]
x1 = 0.0
x2 = Lx
a_pml = lx
z1 = 0.0
z2 = -Lz
c_pml = lz
if dim == 3:
Ly = model["mesh"]["Ly"]
ly = model["BCs"]["ly"]
y1 = 0.0
y2 = Ly
b_pml = ly
if method == "KMV":
params = {"ksp_type": "preonly", "pc_type": "jacobi"}
elif method == "CG":
params = {"ksp_type": "cg", "pc_type": "jacobi"}
else:
raise ValueError("method is not yet supported")
element = space.FE_method(mesh, method, degree)
V = FunctionSpace(mesh, element)
qr_x, qr_s, _ = quadrature.quadrature_rules(V)
nt = int(tf / dt) # number of timesteps
receiver_locations = model["acquisition"]["receiver_locations"]
dJ = Function(V, name="gradient")
if dim == 2:
z, x = SpatialCoordinate(mesh)
elif dim == 3:
z, x, y = SpatialCoordinate(mesh)
if PML:
Z = VectorFunctionSpace(V.ufl_domain(), V.ufl_element())
if dim == 2:
W = V * Z
u, pp = TrialFunctions(W)
v, qq = TestFunctions(W)
u_np1, pp_np1 = Function(W).split()
u_n, pp_n = Function(W).split()
u_nm1, pp_nm1 = Function(W).split()
elif dim == 3:
W = V * V * Z
u, psi, pp = TrialFunctions(W)
v, phi, qq = TestFunctions(W)
u_np1, psi_np1, pp_np1 = Function(W).split()
u_n, psi_n, pp_n = Function(W).split()
u_nm1, psi_nm1, pp_nm1 = Function(W).split()
if dim == 2:
(sigma_x, sigma_z) = damping.functions(
model, V, dim, x, x1, x2, a_pml, z, z1, z2, c_pml
)
(Gamma_1, Gamma_2) = damping.matrices_2D(sigma_z, sigma_x)
elif dim == 3:
sigma_x, sigma_y, sigma_z = damping.functions(
model,
V,
dim,
x,
x1,
x2,
a_pml,
z,
z1,
z2,
c_pml,
y,
y1,
y2,
b_pml,
)
Gamma_1, Gamma_2, Gamma_3 = damping.matrices_3D(sigma_x, sigma_y, sigma_z)
# typical CG in N-d
else:
u = TrialFunction(V)
v = TestFunction(V)
u_nm1 = Function(V)
u_n = Function(V)
u_np1 = Function(V)
if output:
outfile = helpers.create_output_file("adjoint.pvd", comm, 0)
t = 0.0
# -------------------------------------------------------
m1 = ((u - 2.0 * u_n + u_nm1) / Constant(dt ** 2)) * v * dx(rule=qr_x)
a = c * c * dot(grad(u_n), grad(v)) * dx(rule=qr_x) # explicit
nf = 0
if model["BCs"]["outer_bc"] == "non-reflective":
nf = c * ((u_n - u_nm1) / dt) * v * ds(rule=qr_s)
FF = m1 + a + nf
if PML:
X = Function(W)
B = Function(W)
if dim == 2:
pml1 = (sigma_x + sigma_z) * ((u - u_n) / dt) * v * dx(rule=qr_x)
pml2 = sigma_x * sigma_z * u_n * v * dx(rule=qr_x)
pml3 = c * c * inner(grad(v), dot(Gamma_2, pp_n)) * dx(rule=qr_x)
FF += pml1 + pml2 + pml3
# -------------------------------------------------------
mm1 = (dot((pp - pp_n), qq) / Constant(dt)) * dx(rule=qr_x)
mm2 = inner(dot(Gamma_1, pp_n), qq) * dx(rule=qr_x)
dd = inner(qq, grad(u_n)) * dx(rule=qr_x)
FF += mm1 + mm2 + dd
elif dim == 3:
pml1 = (sigma_x + sigma_y + sigma_z) * ((u - u_n) / dt) * v * dx(rule=qr_x)
uuu1 = (-v * psi_n) * dx(rule=qr_x)
pml2 = (
(sigma_x * sigma_y + sigma_x * sigma_z + sigma_y * sigma_z)
* u_n
* v
* dx(rule=qr_x)
)
dd1 = c * c * inner(grad(v), dot(Gamma_2, pp_n)) * dx(rule=qr_x)
FF += pml1 + pml2 + dd1 + uuu1
# -------------------------------------------------------
mm1 = (dot((pp - pp_n), qq) / dt) * dx(rule=qr_x)
mm2 = inner(dot(Gamma_1, pp_n), qq) * dx(rule=qr_x)
pml4 = inner(qq, grad(u_n)) * dx(rule=qr_x)
FF += mm1 + mm2 + pml4
# -------------------------------------------------------
pml3 = (sigma_x * sigma_y * sigma_z) * phi * u_n * dx(rule=qr_x)
mmm1 = (dot((psi - psi_n), phi) / dt) * dx(rule=qr_x)
mmm2 = -c * c * inner(grad(phi), dot(Gamma_3, pp_n)) * dx(rule=qr_x)
FF += mmm1 + mmm2 + pml3
else:
X = Function(V)
B = Function(V)
lhs_ = lhs(FF)
rhs_ = rhs(FF)
A = assemble(lhs_, mat_type="matfree")
solver = LinearSolver(A, solver_parameters=params)
# Define gradient problem
m_u = TrialFunction(V)
m_v = TestFunction(V)
mgrad = m_u * m_v * dx(rule=qr_x)
uuadj = Function(V) # auxiliarly function for the gradient compt.
uufor = Function(V) # auxiliarly function for the gradient compt.
ffG = 2.0 * c * dot(grad(uuadj), grad(uufor)) * m_v * dx(rule=qr_x)
G = mgrad - ffG
lhsG, rhsG = lhs(G), rhs(G)
gradi = Function(V)
grad_prob = LinearVariationalProblem(lhsG, rhsG, gradi)
if method == "KMV":
grad_solver = LinearVariationalSolver(
grad_prob,
solver_parameters={
"ksp_type": "preonly",
"pc_type": "jacobi",
"mat_type": "matfree",
},
)
elif method == "CG":
grad_solver = LinearVariationalSolver(
grad_prob,
solver_parameters={
"mat_type": "matfree",
},
)
assembly_callable = create_assembly_callable(rhs_, tensor=B)
rhs_forcing = Function(V) # forcing term
if save_adjoint:
adjoint = [Function(V, name="adjoint_pressure") for t in range(nt)]
for step in range(nt - 1, -1, -1):
t = step * float(dt)
rhs_forcing.assign(0.0)
# Solver - main equation - (I)
# B = assemble(rhs_, tensor=B)
assembly_callable()
f = receivers.apply_receivers_as_source(rhs_forcing, residual, step)
# add forcing term to solve scalar pressure
B0 = B.sub(0)
B0 += f
# AX=B --> solve for X = B/Aˆ-1
solver.solve(X, B)
if PML:
if dim == 2:
u_np1, pp_np1 = X.split()
elif dim == 3:
u_np1, psi_np1, pp_np1 = X.split()
psi_nm1.assign(psi_n)
psi_n .assign(psi_np1)
pp_nm1.assign(pp_n)
pp_n .assign(pp_np1)
else:
u_np1.assign(X)
# only compute for snaps that were saved
if step % fspool == 0:
# compute the gradient increment
uuadj.assign(u_np1)
uufor.assign(guess.pop())
grad_solver.solve()
dJ += gradi
u_nm1.assign(u_n)
u_n.assign(u_np1)
if step % nspool == 0:
if output:
outfile.write(u_n, time=t)
if save_adjoint:
adjoint.append(u_n)
helpers.display_progress(comm, t)
if save_adjoint:
return dJ, adjoint
else:
return dJ
def initialize(self, pc):
from firedrake import TestFunction, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
from firedrake.interpolation import Interpolator
from firedrake.solving_utils import _SNESContext
from firedrake.matrix_free.operators import ImplicitMatrixContext
_, P = pc.getOperators()
appctx = self.get_appctx(pc)
fcp = appctx.get("form_compiler_parameters")
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
ctx = dmhooks.get_appctx(pc.getDM())
if ctx is None:
raise ValueError("No context found.")
if not isinstance(ctx, _SNESContext):
raise ValueError("Don't know how to get form from %r", ctx)
prefix = pc.getOptionsPrefix()
options_prefix = prefix + self._prefix
opts = PETSc.Options()
# Handle the fine operator if type is python
if P.getType() == "python":
ictx = P.getPythonContext()
if ictx is None:
raise ValueError("No context found on matrix")
if not isinstance(ictx, ImplicitMatrixContext):
raise ValueError("Don't know how to get form from %r", ictx)
fine_operator = ictx.a
fine_bcs = ictx.row_bcs
if fine_bcs != ictx.col_bcs:
raise NotImplementedError("Row and column bcs must match")
fine_mat_type = opts.getString(options_prefix + "mat_type",
parameters["default_matrix_type"])
self.fine_op = allocate_matrix(fine_operator,
bcs=fine_bcs,
form_compiler_parameters=fcp,
mat_type=fine_mat_type,
options_prefix=options_prefix)
self._assemble_fine_op = create_assembly_callable(
fine_operator,
tensor=self.fine_op,
bcs=fine_bcs,
form_compiler_parameters=fcp,
mat_type=fine_mat_type)
self._assemble_fine_op()
fine_petscmat = self.fine_op.petscmat
else:
fine_petscmat = P
# Transfer fine operator null space
fine_petscmat.setNullSpace(P.getNullSpace())
fine_transpose_nullspace = P.getTransposeNullSpace()
if fine_transpose_nullspace.handle != 0:
fine_petscmat.setTransposeNullSpace(fine_transpose_nullspace)
# Handle the coarse operator
coarse_options_prefix = options_prefix + "mg_coarse"
coarse_mat_type = opts.getString(coarse_options_prefix + "mat_type",
parameters["default_matrix_type"])
get_coarse_space = appctx.get("get_coarse_space", None)
if not get_coarse_space:
raise ValueError(
"Need to provide a callback which provides the coarse space.")
coarse_space = get_coarse_space()
get_coarse_operator = appctx.get("get_coarse_operator", None)
if not get_coarse_operator:
raise ValueError(
"Need to provide a callback which provides the coarse operator."
)
coarse_operator = get_coarse_operator()
coarse_space_bcs = appctx.get("coarse_space_bcs", None)
# These should be callbacks which return the relevant nullspaces
get_coarse_nullspace = appctx.get("get_coarse_op_nullspace", None)
get_coarse_transpose_nullspace = appctx.get(
"get_coarse_op_transpose_nullspace", None)
self.coarse_op = allocate_matrix(coarse_operator,
bcs=coarse_space_bcs,
form_compiler_parameters=fcp,
mat_type=coarse_mat_type,
options_prefix=coarse_options_prefix)
self._assemble_coarse_op = create_assembly_callable(
coarse_operator,
tensor=self.coarse_op,
bcs=coarse_space_bcs,
form_compiler_parameters=fcp)
self._assemble_coarse_op()
coarse_opmat = self.coarse_op.petscmat
# Set nullspace if provided
if get_coarse_nullspace:
nsp = get_coarse_nullspace()
coarse_opmat.setNullSpace(nsp.nullspace(comm=pc.comm))
if get_coarse_transpose_nullspace:
tnsp = get_coarse_transpose_nullspace()
coarse_opmat.setTransposeNullSpace(tnsp.nullspace(comm=pc.comm))
interp_petscmat = appctx.get("interpolation_matrix", None)
if interp_petscmat is None:
# Create interpolation matrix from coarse space to fine space
fine_space = ctx.J.arguments()[0].function_space()
interpolator = Interpolator(TestFunction(coarse_space), fine_space)
interpolation_matrix = interpolator.callable()
interp_petscmat = interpolation_matrix.handle
# We set up a PCMG object that uses the constructed interpolation
# matrix to generate the restriction/prolongation operators.
# This is a two-level multigrid preconditioner.
pcmg = PETSc.PC().create(comm=pc.comm)
pcmg.incrementTabLevel(1, parent=pc)
pcmg.setType(pc.Type.MG)
pcmg.setOptionsPrefix(options_prefix)
pcmg.setMGLevels(2)
pcmg.setMGCycleType(pc.MGCycleType.V)
pcmg.setMGInterpolation(1, interp_petscmat)
pcmg.setOperators(A=fine_petscmat, P=fine_petscmat)
# Create new appctx
self._ctx_ref = self.new_snes_ctx(pc, coarse_operator,
coarse_space_bcs, coarse_mat_type,
fcp)
coarse_solver = pcmg.getMGCoarseSolve()
coarse_solver.setOperators(A=coarse_opmat, P=coarse_opmat)
# coarse space dm
coarse_dm = coarse_space.dm
coarse_solver.setDM(coarse_dm)
coarse_solver.setDMActive(False)
pcmg.setDM(coarse_dm)
pcmg.setFromOptions()
self.pc = pcmg
self._dm = coarse_dm
with dmhooks.add_hooks(coarse_dm,
self,
appctx=self._ctx_ref,
save=False):
coarse_solver.setFromOptions()
def _rhs(self):
from firedrake.assemble import create_assembly_callable
u = function.Function(self.trial_space)
b = function.Function(self.test_space)
expr = -action(self.A.a, u)
return u, create_assembly_callable(expr, tensor=b), b
def assembler(self):
from firedrake.assemble import create_assembly_callable
return create_assembly_callable(self.rhs_form, tensor=self.residual,
form_compiler_parameters=self.form_compiler_parameters)
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p*q*dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q))*dx + Constant(1e-6)*p*q*dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix+"Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix+"Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix+"Fp_mat_type", "matfree")
Mp = assemble(mass, form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type)
Kp = assemble(stiffness, form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type)
Mp.force_evaluation()
Kp.force_evaluation()
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create()
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create()
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx
self.Re = Re
self.Fp = allocate_matrix(fp, form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp = create_assembly_callable(fp, tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
self.Fp.force_evaluation()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
shapes = (V[self.vidx].finat_element.space_dimension(),
np.prod(V[self.vidx].shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = Function(V[self.vidx])
par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)},
is_loopy_kernel=True)
instructions = """
for i, j
vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j]
end
"""
self.average_kernel = (domain, instructions)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d),
trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h
+ gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc")
if bc.function_space().index != self.vidx:
raise NotImplementedError("Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain])
trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand*measure
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(TraceSpace)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
def initialize(self, pc):
from firedrake.assemble import allocate_matrix, create_assembly_callable
_, P = pc.getOperators()
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
opc = pc
appctx = self.get_appctx(pc)
fcp = appctx.get("form_compiler_parameters")
V = get_function_space(pc.getDM())
if len(V) == 1:
V = FunctionSpace(V.mesh(), V.ufl_element())
else:
V = MixedFunctionSpace([V_ for V_ in V])
test = TestFunction(V)
trial = TrialFunction(V)
if P.type == "python":
context = P.getPythonContext()
# It only makes sense to preconditioner/invert a diagonal
# block in general. That's all we're going to allow.
if not context.on_diag:
raise ValueError("Only makes sense to invert diagonal block")
prefix = pc.getOptionsPrefix()
options_prefix = prefix + self._prefix
mat_type = PETSc.Options().getString(options_prefix + "mat_type", "aij")
(a, bcs) = self.form(pc, test, trial)
self.P = allocate_matrix(a, bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type,
options_prefix=options_prefix)
self._assemble_P = create_assembly_callable(a, tensor=self.P,
bcs=bcs,
form_compiler_parameters=fcp,
mat_type=mat_type)
self._assemble_P()
self.P.force_evaluation()
# Transfer nullspace over
Pmat = self.P.petscmat
Pmat.setNullSpace(P.getNullSpace())
tnullsp = P.getTransposeNullSpace()
if tnullsp.handle != 0:
Pmat.setTransposeNullSpace(tnullsp)
# Internally, we just set up a PC object that the user can configure
# however from the PETSc command line. Since PC allows the user to specify
# a KSP, we can do iterative by -assembled_pc_type ksp.
pc = PETSc.PC().create(comm=opc.comm)
pc.incrementTabLevel(1, parent=opc)
# We set a DM and an appropriate SNESContext on the constructed PC so one
# can do e.g. multigrid or patch solves.
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.solving_utils import _SNESContext
dm = opc.getDM()
octx = get_appctx(dm)
oproblem = octx._problem
nproblem = NonlinearVariationalProblem(oproblem.F, oproblem.u, bcs, J=a, form_compiler_parameters=fcp)
nctx = _SNESContext(nproblem, mat_type, mat_type, octx.appctx)
push_appctx(dm, nctx)
self._ctx_ref = nctx
pc.setDM(dm)
pc.setOptionsPrefix(options_prefix)
pc.setOperators(Pmat, Pmat)
pc.setFromOptions()
pc.setUp()
self.pc = pc
pop_appctx(dm)
def forward(
model,
mesh,
comm,
c,
excitations,
wavelet,
receivers,
source_num=0,
output=False,
):
"""Secord-order in time fully-explicit scheme
with implementation of a Perfectly Matched Layer (PML) using
CG FEM with or without higher order mass lumping (KMV type elements).
Parameters
----------
model: Python `dictionary`
Contains model options and parameters
mesh: Firedrake.mesh object
The 2D/3D triangular mesh
comm: Firedrake.ensemble_communicator
The MPI communicator for parallelism
c: Firedrake.Function
The velocity model interpolated onto the mesh.
excitations: A list Firedrake.Functions
wavelet: array-like
Time series data that's injected at the source location.
receivers: A :class:`spyro.Receivers` object.
Contains the receiver locations and sparse interpolation methods.
source_num: `int`, optional
The source number you wish to simulate
output: `boolean`, optional
Whether or not to write results to pvd files.
Returns
-------
usol: list of Firedrake.Functions
The full field solution at `fspool` timesteps
usol_recv: array-like
The solution interpolated to the receivers at all timesteps
"""
method = model["opts"]["method"]
degree = model["opts"]["degree"]
dim = model["opts"]["dimension"]
dt = model["timeaxis"]["dt"]
tf = model["timeaxis"]["tf"]
nspool = model["timeaxis"]["nspool"]
fspool = model["timeaxis"]["fspool"]
PML = model["BCs"]["status"]
excitations.current_source = source_num
if PML:
Lx = model["mesh"]["Lx"]
Lz = model["mesh"]["Lz"]
lx = model["BCs"]["lx"]
lz = model["BCs"]["lz"]
x1 = 0.0
x2 = Lx
a_pml = lx
z1 = 0.0
z2 = -Lz
c_pml = lz
if dim == 3:
Ly = model["mesh"]["Ly"]
ly = model["BCs"]["ly"]
y1 = 0.0
y2 = Ly
b_pml = ly
nt = int(tf / dt) # number of timesteps
if method == "KMV":
params = {"ksp_type": "preonly", "pc_type": "jacobi"}
elif (method == "CG" and mesh.ufl_cell() != quadrilateral
and mesh.ufl_cell() != hexahedron):
params = {"ksp_type": "cg", "pc_type": "jacobi"}
elif method == "CG" and (mesh.ufl_cell() == quadrilateral
or mesh.ufl_cell() == hexahedron):
params = {"ksp_type": "preonly", "pc_type": "jacobi"}
else:
raise ValueError("method is not yet supported")
element = space.FE_method(mesh, method, degree)
V = FunctionSpace(mesh, element)
qr_x, qr_s, _ = quadrature.quadrature_rules(V)
if dim == 2:
z, x = SpatialCoordinate(mesh)
elif dim == 3:
z, x, y = SpatialCoordinate(mesh)
if PML:
Z = VectorFunctionSpace(V.ufl_domain(), V.ufl_element())
if dim == 2:
W = V * Z
u, pp = TrialFunctions(W)
v, qq = TestFunctions(W)
u_np1, pp_np1 = Function(W).split()
u_n, pp_n = Function(W).split()
u_nm1, pp_nm1 = Function(W).split()
elif dim == 3:
W = V * V * Z
u, psi, pp = TrialFunctions(W)
v, phi, qq = TestFunctions(W)
u_np1, psi_np1, pp_np1 = Function(W).split()
u_n, psi_n, pp_n = Function(W).split()
u_nm1, psi_nm1, pp_nm1 = Function(W).split()
if dim == 2:
sigma_x, sigma_z = damping.functions(model, V, dim, x, x1, x2,
a_pml, z, z1, z2, c_pml)
Gamma_1, Gamma_2 = damping.matrices_2D(sigma_z, sigma_x)
pml1 = ((sigma_x + sigma_z) * ((u - u_nm1) / Constant(2.0 * dt)) *
v * dx(rule=qr_x))
elif dim == 3:
sigma_x, sigma_y, sigma_z = damping.functions(
model,
V,
dim,
x,
x1,
x2,
a_pml,
z,
z1,
z2,
c_pml,
y,
y1,
y2,
b_pml,
)
Gamma_1, Gamma_2, Gamma_3 = damping.matrices_3D(
sigma_x, sigma_y, sigma_z)
# typical CG FEM in 2d/3d
else:
u = TrialFunction(V)
v = TestFunction(V)
u_nm1 = Function(V)
u_n = Function(V)
u_np1 = Function(V)
if output:
outfile = helpers.create_output_file("forward.pvd", comm, source_num)
t = 0.0
# -------------------------------------------------------
m1 = ((u - 2.0 * u_n + u_nm1) / Constant(dt**2)) * v * dx(rule=qr_x)
a = c * c * dot(grad(u_n), grad(v)) * dx(rule=qr_x) # explicit
nf = 0
if model["BCs"]["outer_bc"] == "non-reflective":
nf = c * ((u_n - u_nm1) / dt) * v * ds(rule=qr_s)
FF = m1 + a + nf
if PML:
X = Function(W)
B = Function(W)
if dim == 2:
pml2 = sigma_x * sigma_z * u_n * v * dx(rule=qr_x)
pml3 = inner(pp_n, grad(v)) * dx(rule=qr_x)
FF += pml1 + pml2 + pml3
# -------------------------------------------------------
mm1 = (dot((pp - pp_n), qq) / Constant(dt)) * dx(rule=qr_x)
mm2 = inner(dot(Gamma_1, pp_n), qq) * dx(rule=qr_x)
dd = c * c * inner(grad(u_n), dot(Gamma_2, qq)) * dx(rule=qr_x)
FF += mm1 + mm2 + dd
elif dim == 3:
pml1 = ((sigma_x + sigma_y + sigma_z) *
((u - u_n) / Constant(dt)) * v * dx(rule=qr_x))
pml2 = (
(sigma_x * sigma_y + sigma_x * sigma_z + sigma_y * sigma_z) *
u_n * v * dx(rule=qr_x))
pml3 = (sigma_x * sigma_y * sigma_z) * psi_n * v * dx(rule=qr_x)
pml4 = inner(pp_n, grad(v)) * dx(rule=qr_x)
FF += pml1 + pml2 + pml3 + pml4
# -------------------------------------------------------
mm1 = (dot((pp - pp_n), qq) / Constant(dt)) * dx(rule=qr_x)
mm2 = inner(dot(Gamma_1, pp_n), qq) * dx(rule=qr_x)
dd1 = c * c * inner(grad(u_n), dot(Gamma_2, qq)) * dx(rule=qr_x)
dd2 = -c * c * inner(grad(psi_n), dot(Gamma_3, qq)) * dx(rule=qr_x)
FF += mm1 + mm2 + dd1 + dd2
# -------------------------------------------------------
mmm1 = (dot((psi - psi_n), phi) / Constant(dt)) * dx(rule=qr_x)
uuu1 = (-u_n * phi) * dx(rule=qr_x)
FF += mmm1 + uuu1
else:
X = Function(V)
B = Function(V)
lhs_ = lhs(FF)
rhs_ = rhs(FF)
A = assemble(lhs_, mat_type="matfree")
solver = LinearSolver(A, solver_parameters=params)
usol = [Function(V, name="pressure") for t in range(nt) if t % fspool == 0]
usol_recv = []
save_step = 0
assembly_callable = create_assembly_callable(rhs_, tensor=B)
rhs_forcing = Function(V)
for step in range(nt):
rhs_forcing.assign(0.0)
assembly_callable()
f = excitations.apply_source(rhs_forcing, wavelet[step])
B0 = B.sub(0)
B0 += f
solver.solve(X, B)
if PML:
if dim == 2:
u_np1, pp_np1 = X.split()
elif dim == 3:
u_np1, psi_np1, pp_np1 = X.split()
psi_nm1.assign(psi_n)
psi_n.assign(psi_np1)
pp_nm1.assign(pp_n)
pp_n.assign(pp_np1)
else:
u_np1.assign(X)
usol_recv.append(receivers.interpolate(
u_np1.dat.data_ro_with_halos[:]))
if step % fspool == 0:
usol[save_step].assign(u_np1)
save_step += 1
if step % nspool == 0:
assert (
norm(u_n) < 1
), "Numerical instability. Try reducing dt or building the mesh differently"
if output:
outfile.write(u_n, time=t, name="Pressure")
if t > 0:
helpers.display_progress(comm, t)
u_nm1.assign(u_n)
u_n.assign(u_np1)
t = step * float(dt)
usol_recv = helpers.fill(usol_recv, receivers.is_local, nt,
receivers.num_receivers)
usol_recv = utils.communicate(usol_recv, comm)
return usol, usol_recv
def _rhs(self):
from firedrake.assemble import create_assembly_callable
u = function.Function(self.trial_space)
b = function.Function(self.test_space)
expr = -action(self.A.a, u)
return u, create_assembly_callable(expr, tensor=b), b
def __init__(self, problem, mat_type, pmat_type, appctx=None, pre_jacobian_callback=None, pre_function_callback=None):
from firedrake.assemble import allocate_matrix, create_assembly_callable
if pmat_type is None:
pmat_type = mat_type
self.mat_type = mat_type
self.pmat_type = pmat_type
matfree = mat_type == 'matfree'
pmatfree = pmat_type == 'matfree'
self._problem = problem
self._pre_jacobian_callback = pre_jacobian_callback
self._pre_function_callback = pre_function_callback
fcp = problem.form_compiler_parameters
# Function to hold current guess
self._x = problem.u
if appctx is None:
appctx = {}
if matfree or pmatfree:
# A split context will already get the full state.
# TODO, a better way of doing this.
# Now we don't have a temporary state inside the snes
# context we could just require the user to pass in the
# full state on the outside.
appctx.setdefault("state", self._x)
self.appctx = appctx
self.matfree = matfree
self.pmatfree = pmatfree
self.F = problem.F
self.J = problem.J
self._jac = allocate_matrix(self.J, bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=mat_type,
appctx=appctx)
self._assemble_jac = create_assembly_callable(self.J,
tensor=self._jac,
bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=mat_type)
self.is_mixed = self._jac.block_shape != (1, 1)
if mat_type != pmat_type or problem.Jp is not None:
# Need separate pmat if either Jp is different or we want
# a different pmat type to the mat type.
if problem.Jp is None:
self.Jp = self.J
else:
self.Jp = problem.Jp
self._pjac = allocate_matrix(self.Jp, bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=pmat_type,
appctx=appctx)
self._assemble_pjac = create_assembly_callable(self.Jp,
tensor=self._pjac,
bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=pmat_type)
else:
# pmat_type == mat_type and Jp is None
self.Jp = None
self._pjac = self._jac
self._F = function.Function(self.F.arguments()[0].function_space())
self._assemble_residual = create_assembly_callable(self.F,
tensor=self._F,
form_compiler_parameters=fcp)
self._jacobian_assembled = False
self._splits = {}
self._coarse = None
self._fine = None
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p * q * dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q)) * dx + Constant(1e-6) * p * q * dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mp.force_evaluation()
Kp.force_evaluation()
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0 / Re * inner(grad(p), grad(q)) * dx + inner(u0,
grad(p)) * q * dx
self.Re = Re
self.Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
self.Fp.force_evaluation()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
def __init__(self,
problem,
mat_type,
pmat_type,
appctx=None,
pre_jacobian_callback=None,
pre_function_callback=None,
options_prefix=None):
from firedrake.assemble import allocate_matrix, create_assembly_callable
if pmat_type is None:
pmat_type = mat_type
self.mat_type = mat_type
self.pmat_type = pmat_type
matfree = mat_type == 'matfree'
pmatfree = pmat_type == 'matfree'
self._problem = problem
self._pre_jacobian_callback = pre_jacobian_callback
self._pre_function_callback = pre_function_callback
fcp = problem.form_compiler_parameters
# Function to hold current guess
self._x = problem.u
if appctx is None:
appctx = {}
if matfree or pmatfree:
# A split context will already get the full state.
# TODO, a better way of doing this.
# Now we don't have a temporary state inside the snes
# context we could just require the user to pass in the
# full state on the outside.
appctx.setdefault("state", self._x)
self.appctx = appctx
self.matfree = matfree
self.pmatfree = pmatfree
self.F = problem.F
self.J = problem.J
self._jac = allocate_matrix(self.J,
bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=mat_type,
appctx=appctx,
options_prefix=options_prefix)
self._assemble_jac = create_assembly_callable(
self.J,
tensor=self._jac,
bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=mat_type)
self.is_mixed = self._jac.block_shape != (1, 1)
if mat_type != pmat_type or problem.Jp is not None:
# Need separate pmat if either Jp is different or we want
# a different pmat type to the mat type.
if problem.Jp is None:
self.Jp = self.J
else:
self.Jp = problem.Jp
self._pjac = allocate_matrix(self.Jp,
bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=pmat_type,
appctx=appctx,
options_prefix=options_prefix)
self._assemble_pjac = create_assembly_callable(
self.Jp,
tensor=self._pjac,
bcs=problem.bcs,
form_compiler_parameters=fcp,
mat_type=pmat_type)
else:
# pmat_type == mat_type and Jp is None
self.Jp = None
self._pjac = self._jac
self._F = function.Function(self.F.arguments()[0].function_space())
self._assemble_residual = create_assembly_callable(
self.F, tensor=self._F, form_compiler_parameters=fcp)
self._jacobian_assembled = False
self._splits = {}
self._coarse = None
self._fine = None
def initialize(self, pc):
""" Set up the problem context. Takes the original
mixed problem and transforms it into the equivalent
hybrid-mixed system.
A KSP object is created for the Lagrange multipliers
on the top/bottom faces of the mesh cells.
"""
from firedrake import (FunctionSpace, Function, Constant,
FiniteElement, TensorProductElement,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, interval, MixedElement,
BrokenElement)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
from ufl.cell import TensorProductCell
# Extract PC context
prefix = pc.getOptionsPrefix() + "vert_hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
# Magically determine which spaces are vector and scalar valued
for i, Vi in enumerate(V):
# Vector-valued spaces will have a non-empty value_shape
if Vi.ufl_element().value_shape():
self.vidx = i
else:
self.pidx = i
Vv = V[self.vidx]
Vp = V[self.pidx]
# Create the space of approximate traces in the vertical.
# NOTE: Technically a hack since the resulting space is technically
# defined in cell interiors, however the degrees of freedom will only
# be geometrically defined on edges. Arguments will only be used in
# surface integrals
deg, _ = Vv.ufl_element().degree()
# Assumes a tensor product cell (quads, triangular-prisms, cubes)
if not isinstance(Vp.ufl_element().cell(), TensorProductCell):
raise NotImplementedError(
"Currently only implemented for tensor product discretizations"
)
# Only want the horizontal cell
cell, _ = Vp.ufl_element().cell()._cells
DG = FiniteElement("DG", cell, deg)
CG = FiniteElement("CG", interval, 1)
Vv_tr_element = TensorProductElement(DG, CG)
Vv_tr = FunctionSpace(mesh, Vv_tr_element)
# Break the spaces
broken_elements = MixedElement(
[BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up relevant functions
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(Vv_tr)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up transfer kernels to and from the broken velocity space
# NOTE: Since this snippet of code is used in a couple places in
# in Gusto, might be worth creating a utility function that is
# is importable and just called where needed.
shapes = {
"i": Vv.finat_element.space_dimension(),
"j": np.prod(Vv.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self.weight = Function(Vv)
par_loop(weight_kernel, dx, {"w": (self.weight, INC)})
# Averaging kernel
self.average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# Original mixed operator replaced with "broken" arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(Vv_tr)
n = FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# Again, assumes tensor product structure. Why use this if you
# don't have some form of vertical extrusion?
Kform = gammar('+') * jump(sigma, n=n) * dS_h
# Here we deal with boundary conditions
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError(
"Dirichlet conditions for scalar variable not supported. Use a weak bc."
)
if bc.function_space().index != self.vidx:
raise NotImplementedError(
"Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(as_tuple(subdom, int))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * dot(sigma, n)
measures = []
trace_subdomains = []
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({"top": ds_t, "bottom": ds_b}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand * measure
else:
trace_subdomains = ["top", "bottom"]
trace_bcs = [
DirichletBC(Vv_tr, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(Vv_tr)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("vert_trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vv_tr)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
def __init__(self, a, row_bcs=[], col_bcs=[], fc_params=None, appctx=None):
self.a = a
self.aT = adjoint(a)
self.fc_params = fc_params
self.appctx = appctx
# Collect all DirichletBC instances including
# DirichletBCs applied to an EquationBC.
# all bcs (DirichletBC, EquationBCSplit)
self.bcs = row_bcs
self.bcs_col = col_bcs
self.row_bcs = tuple(bc for bc in itertools.chain(*row_bcs)
if isinstance(bc, DirichletBC))
self.col_bcs = tuple(bc for bc in itertools.chain(*col_bcs)
if isinstance(bc, DirichletBC))
# create functions from test and trial space to help
# with 1-form assembly
test_space, trial_space = [
a.arguments()[i].function_space() for i in (0, 1)
]
from firedrake import function
self._y = function.Function(test_space)
self._x = function.Function(trial_space)
# These are temporary storage for holding the BC
# values during matvec application. _xbc is for
# the action and ._ybc is for transpose.
if len(self.bcs) > 0:
self._xbc = function.Function(trial_space)
if len(self.col_bcs) > 0:
self._ybc = function.Function(test_space)
# Get size information from template vecs on test and trial spaces
trial_vec = trial_space.dof_dset.layout_vec
test_vec = test_space.dof_dset.layout_vec
self.col_sizes = trial_vec.getSizes()
self.row_sizes = test_vec.getSizes()
self.block_size = (test_vec.getBlockSize(), trial_vec.getBlockSize())
self.action = action(self.a, self._x)
self.actionT = action(self.aT, self._y)
from firedrake.assemble import create_assembly_callable
# For assembling action(f, self._x)
self.bcs_action = []
for bc in self.bcs:
if isinstance(bc, DirichletBC):
self.bcs_action.append(bc)
elif isinstance(bc, EquationBCSplit):
self.bcs_action.append(bc.reconstruct(action_x=self._x))
self._assemble_action = create_assembly_callable(
self.action,
tensor=self._y,
bcs=self.bcs_action,
form_compiler_parameters=self.fc_params)
# For assembling action(adjoint(f), self._y)
# Sorted list of equation bcs
self.objs_actionT = []
for bc in self.bcs:
self.objs_actionT += bc.sorted_equation_bcs()
self.objs_actionT.append(self)
# Each par_loop is to run with appropriate masks on self._y
self._assemble_actionT = []
# Deepest EquationBCs first
for bc in self.bcs:
for ebc in bc.sorted_equation_bcs():
self._assemble_actionT.append(
create_assembly_callable(
action(adjoint(ebc.f), self._y),
tensor=self._xbc,
bcs=None,
form_compiler_parameters=self.fc_params))
# Domain last
self._assemble_actionT.append(
create_assembly_callable(
self.actionT,
tensor=self._x if len(self.bcs) == 0 else self._xbc,
bcs=None,
form_compiler_parameters=self.fc_params))
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(Smat)
c_ksp.setUp()
c_ksp.setFromOptions()
self.condensed_ksp = c_ksp
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A, self.residual,
self.solution,
elim_fields)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement(
[ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
shapes = (V[self.vidx].finat_element.space_dimension(),
np.prod(V[self.vidx].shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = Function(V[self.vidx])
par_loop((domain, instructions),
ufl.dx, {"w": (self.weight, INC)},
is_loopy_kernel=True)
instructions = """
for i, j
vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j]
end
"""
self.average_kernel = (domain, instructions)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h +
gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError(
"Dirichlet conditions for scalar variable not supported. Use a weak bc"
)
if bc.function_space().index != self.vidx:
raise NotImplementedError(
"Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(
as_tuple(subdom, numbers.Integral))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({
"top": ufl.ds_t,
"bottom": ufl.ds_b
}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand * measure
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
with timed_region("HybridOperatorAssembly"):
self._assemble_S()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(TraceSpace)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
def assembler(self):
from firedrake.assemble import create_assembly_callable
return create_assembly_callable(self.rhs_form, tensor=self.residual,
form_compiler_parameters=self.form_compiler_parameters)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, assemble)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.cxt.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up the KSP for the hdiv residual projection
hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_")
# HDiv mass operator
p = TrialFunction(V[self.vidx])
q = TestFunction(V[self.vidx])
mass = ufl.dot(p, q)*ufl.dx
# TODO: Bcs?
M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params)
M.force_evaluation()
Mmat = M.petscmat
hdiv_mass_ksp.setOperators(Mmat)
hdiv_mass_ksp.setUp()
hdiv_mass_ksp.setFromOptions()
self.hdiv_mass_ksp = hdiv_mass_ksp
# Storing the result of A.inv * r, where A is the HDiv
# mass matrix and r is the HDiv residual
self._primal_r = Function(V[self.vidx])
tau = TestFunction(V_d[self.vidx])
self._assemble_broken_r = create_assembly_callable(
ufl.dot(self._primal_r, tau)*ufl.dx,
tensor=self.broken_residual.split()[self.vidx],
form_compiler_parameters=self.cxt.fc_params)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d),
trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.cxt.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# We zero out the contribution of the trace variables on the exterior
# boundary. Extruded cells will have both horizontal and vertical
# facets
if mesh.cell_set._extruded:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary"),
DirichletBC(TraceSpace, Constant(0.0), "bottom"),
DirichletBC(TraceSpace, Constant(0.0), "top")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS_h +
gammar('+') * ufl.dot(sigma, n) * ufl.dS_v)
else:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.cxt.row_bcs:
raise NotImplementedError("Strong BCs not currently handled. Try imposing them weakly.")
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_residual,
tensor=self.schur_rhs,
form_compiler_parameters=self.cxt.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullspace = create_schur_nullspace(P, -K * Atilde,
V, V_d, TraceSpace,
pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
# NOTE: The projection stage *might* be replaced by a Fortin
# operator. We may want to allow the user to specify if they
# wish to use a Fortin operator over a projection, or vice-versa.
# In a future add-on, we can add a switch which chooses either
# the Fortin reconstruction or the usual KSP projection.
# Set up the projection KSP
hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_')
# Reuse the mass operator from the hdiv_mass_ksp
hdiv_projection_ksp.setOperators(Mmat)
# Construct the RHS for the projection stage
self._projection_rhs = Function(V[self.vidx])
self._assemble_projection_rhs = create_assembly_callable(
ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx,
tensor=self._projection_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Finalize ksp setup
hdiv_projection_ksp.setUp()
hdiv_projection_ksp.setFromOptions()
self.hdiv_projection_ksp = hdiv_projection_ksp
def initialize(self, pc):
"""Set up the problem context. This takes the incoming
three-field system and constructs the static
condensation operators using Slate expressions.
A KSP is created for the reduced system. The eliminated
variables are recovered via back-substitution.
"""
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.bcs import DirichletBC
from firedrake.function import Function
from firedrake.functionspace import FunctionSpace
from firedrake.interpolation import interpolate
prefix = pc.getOptionsPrefix() + "condensed_field_"
A, P = pc.getOperators()
self.cxt = A.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("Context must be an ImplicitMatrixContext")
self.bilinear_form = self.cxt.a
# Retrieve the mixed function space
W = self.bilinear_form.arguments()[0].function_space()
if len(W) > 3:
raise NotImplementedError("Only supports up to three function spaces.")
elim_fields = PETSc.Options().getString(pc.getOptionsPrefix()
+ "pc_sc_eliminate_fields",
None)
if elim_fields:
elim_fields = [int(i) for i in elim_fields.split(',')]
else:
# By default, we condense down to the last field in the
# mixed space.
elim_fields = [i for i in range(0, len(W) - 1)]
condensed_fields = list(set(range(len(W))) - set(elim_fields))
if len(condensed_fields) != 1:
raise NotImplementedError("Cannot condense to more than one field")
c_field, = condensed_fields
# Need to duplicate a space which is NOT
# associated with a subspace of a mixed space.
Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element())
bcs = []
cxt_bcs = self.cxt.row_bcs
for bc in cxt_bcs:
if bc.function_space().index != c_field:
raise NotImplementedError("Strong BC set on unsupported space")
if isinstance(bc.function_arg, Function):
bc_arg = interpolate(bc.function_arg, Vc)
else:
# Constants don't need to be interpolated
bc_arg = bc.function_arg
bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain))
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
self.c_field = c_field
self.condensed_rhs = Function(Vc)
self.residual = Function(W)
self.solution = Function(W)
# Get expressions for the condensed linear system
A_tensor = Tensor(self.bilinear_form)
reduced_sys = self.condensed_system(A_tensor, self.residual, elim_fields)
S_expr = reduced_sys.lhs
r_expr = reduced_sys.rhs
# Construct the condensed right-hand side
self._assemble_Srhs = create_assembly_callable(
r_expr,
tensor=self.condensed_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Allocate and set the condensed operator
self.S = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S = create_assembly_callable(
S_expr,
tensor=self.S,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S()
Smat = self.S.petscmat
# If a different matrix is used for preconditioning,
# assemble this as well
if A != P:
self.cxt_pc = P.getPythonContext()
P_tensor = Tensor(self.cxt_pc.a)
P_reduced_sys = self.condensed_system(P_tensor,
self.residual,
elim_fields)
S_pc_expr = P_reduced_sys.lhs
self.S_pc_expr = S_pc_expr
# Allocate and set the condensed operator
self.S_pc = allocate_matrix(S_expr,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type,
options_prefix=prefix,
appctx=self.get_appctx(pc))
self._assemble_S_pc = create_assembly_callable(
S_pc_expr,
tensor=self.S_pc,
bcs=bcs,
form_compiler_parameters=self.cxt.fc_params,
mat_type=mat_type)
self._assemble_S_pc()
Smat_pc = self.S_pc.petscmat
else:
self.S_pc_expr = S_expr
Smat_pc = Smat
# Get nullspace for the condensed operator (if any).
# This is provided as a user-specified callback which
# returns the basis for the nullspace.
nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
if nullspace is not None:
nsp = nullspace(Vc)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Create a SNESContext for the DM associated with the trace problem
self._ctx_ref = self.new_snes_ctx(pc,
S_expr,
bcs,
mat_type,
self.cxt.fc_params,
options_prefix=prefix)
# Push new context onto the dm associated with the condensed problem
c_dm = Vc.dm
# Set up ksp for the condensed problem
c_ksp = PETSc.KSP().create(comm=pc.comm)
c_ksp.incrementTabLevel(1, parent=pc)
# Set the dm for the condensed solver
c_ksp.setDM(c_dm)
c_ksp.setDMActive(False)
c_ksp.setOptionsPrefix(prefix)
c_ksp.setOperators(A=Smat, P=Smat_pc)
self.condensed_ksp = c_ksp
with dmhooks.add_hooks(c_dm, self,
appctx=self._ctx_ref,
save=False):
c_ksp.setFromOptions()
# Set up local solvers for backwards substitution
self.local_solvers = self.local_solver_calls(A_tensor,
self.residual,
self.solution,
elim_fields)