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 _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 _Abcs(self):
"""A function storing the action of the operator on a zero Function
satisfying the BCs.
Used in the presence of BCs.
"""
b = function.Function(self._W)
for bc in self.A.bcs:
bc.apply(b)
from firedrake.assemble import _assemble
if isinstance(self.A.a, slate.TensorBase):
return _assemble(self.A.a * slate.AssembledVector(b))
else:
return _assemble(ufl.action(self.A.a, b))
def backward_solve(A, b, x, reconstruct_fields):
"""Returns a sequence of linear algebra contexts containing
Slate expressions for backwards substitution.
:arg A: a `slate.Tensor` corresponding to the
mixed UFL operator.
:arg b: a `firedrake.Function` corresponding
to the right-hand side.
:arg x: a `firedrake.Function` corresponding
to the solution.
:arg reconstruct_fields: a `tuple` of indices denoting
which fields to reconstruct.
"""
if not isinstance(A, slate.Tensor):
raise ValueError("Left-hand operator must be a Slate Tensor")
all_fields = list(range(len(A.arg_function_spaces[0])))
nfields = len(all_fields)
reconstruct_fields = as_tuple(reconstruct_fields)
_A = A.blocks
_b = b.split()
_x = x.split()
# Ordering matters
systems = []
# Reconstruct one unknown from one determined field:
#
# | A_ee A_ef || x_e | | b_e |
# | || | = | |
# | A_fe A_ff || x_f | | b_f |
#
# where x_f is known from a previous computation.
# Returns the system:
#
# A_ee x_e = b_e - A_ef * x_f.
if nfields == 2:
id_e, = reconstruct_fields
id_f, = [idx for idx in all_fields if idx != id_e]
A_ee = _A[id_e, id_e]
A_ef = _A[id_e, id_f]
b_e = slate.AssembledVector(_b[id_e])
x_f = slate.AssembledVector(_x[id_f])
r_e = b_e - A_ef * x_f
local_system = LAContext(lhs=A_ee, rhs=r_e, field_idx=(id_e,))
systems.append(local_system)
# Reconstruct two unknowns from one determined field:
#
# | A_e0e0 A_e0e1 A_e0f || x_e0 | | b_e0 |
# | A_e1e0 A_e1e1 A_e1f || x_e1 | = | b_e1 |
# | A_fe0 A_fe1 A_ff || x_f | | b_f |
#
# where x_f is the known field. Returns two systems to be
# solved in order (determined from the reverse order of indices
# e0 and e1):
#
# Solve for e1 first (obtained from eliminating x_e0):
#
# S_e1 x_e1 = r_e1
#
# where
#
# S_e1 = A_e1e1 - A_e1e0 * A_e0e0.inv * A_e0e1, and
# r_e1 = b_e1 - A_e1e0 * A_e0e0.inv * b_e0
# - (A_e1f - A_e1e0 * A_e0e0.inv * A_e0f) * x_f,
#
# And then solve for x_e0 given x_f and x_e1:
#
# A_e0e0 x_e0 = b_e0 - A_e0e1 * x_e1 - A_e0f * x_f.
elif nfields == 3:
if len(reconstruct_fields) != nfields - 1:
raise NotImplementedError("Implemented for 1 determined field")
# Order of reconstruction doesn't need to be in order
# of increasing indices
id_e0, id_e1 = reconstruct_fields
id_f, = [idx for idx in all_fields if idx not in reconstruct_fields]
A_e0e0 = _A[id_e0, id_e0]
A_e0e1 = _A[id_e0, id_e1]
A_e1e0 = _A[id_e1, id_e0]
A_e1e1 = _A[id_e1, id_e1]
A_e0f = _A[id_e0, id_f]
A_e1f = _A[id_e1, id_f]
x_e1 = slate.AssembledVector(_x[id_e1])
x_f = slate.AssembledVector(_x[id_f])
b_e0 = slate.AssembledVector(_b[id_e0])
b_e1 = slate.AssembledVector(_b[id_e1])
# Solve for e1
Sf = A_e1f - A_e1e0 * A_e0e0.inv * A_e0f
S_e1 = A_e1e1 - A_e1e0 * A_e0e0.inv * A_e0e1
r_e1 = b_e1 - A_e1e0 * A_e0e0.inv * b_e0 - Sf * x_f
systems.append(LAContext(lhs=S_e1, rhs=r_e1, field_idx=(id_e1,)))
# Solve for e0
r_e0 = b_e0 - A_e0e1 * x_e1 - A_e0f * x_f
systems.append(LAContext(lhs=A_e0e0, rhs=r_e0, field_idx=(id_e0,)))
else:
msg = "Not implemented for systems with %s fields" % nfields
raise NotImplementedError(msg)
return systems
def condense_and_forward_eliminate(A, b, elim_fields):
"""Returns Slate expressions for the operator and
right-hand side vector after eliminating specified
unknowns.
:arg A: a `slate.Tensor` corresponding to the
mixed UFL operator.
:arg b: a `firedrake.Function` corresponding
to the right-hand side.
:arg elim_fields: a `tuple` of indices denoting
which fields to eliminate.
"""
if not isinstance(A, slate.Tensor):
raise ValueError("Left-hand operator must be a Slate Tensor")
# Ensures field indices are in increasing order
elim_fields = list(as_tuple(elim_fields))
elim_fields.sort()
all_fields = list(range(len(A.arg_function_spaces[0])))
condensed_fields = list(set(all_fields) - set(elim_fields))
condensed_fields.sort()
_A = A.blocks
_b = slate.AssembledVector(b).blocks
# NOTE: Does not support non-contiguous field elimination
e_idx0 = elim_fields[0]
e_idx1 = elim_fields[-1]
f_idx0 = condensed_fields[0]
f_idx1 = condensed_fields[-1]
# Finite element systems where static condensation
# is possible have the general form:
#
# | A_ee A_ef || x_e | | b_e |
# | || | = | |
# | A_fe A_ff || x_f | | b_f |
#
# where subscript `e` denotes the coupling with fields
# that will be eliminated, and `f` denotes the condensed
# fields.
Aff = _A[f_idx0:f_idx1 + 1, f_idx0:f_idx1 + 1]
Aef = _A[e_idx0:e_idx1 + 1, f_idx0:f_idx1 + 1]
Afe = _A[f_idx0:f_idx1 + 1, e_idx0:e_idx1 + 1]
Aee = _A[e_idx0:e_idx1 + 1, e_idx0:e_idx1 + 1]
bf = _b[f_idx0:f_idx1 + 1]
be = _b[e_idx0:e_idx1 + 1]
# The reduced operator and right-hand side are:
# S = A_ff - A_fe * A_ee.inv * A_ef
# r = b_f - A_fe * A_ee.inv * b_e
# as show in Slate:
S = Aff - Afe * Aee.inv * Aef
r = bf - Afe * Aee.inv * be
field_idx = [idx for idx in range(f_idx0, f_idx1)]
return LAContext(lhs=S, rhs=r, field_idx=field_idx)