def __init__(self, a, L, u, bcs=None, aP=None,
form_compiler_parameters=None,
constant_jacobian=True):
r"""
:param a: the bilinear form
:param L: the linear form
:param u: the :class:`.Function` to solve for
:param bcs: the boundary conditions (optional)
:param aP: an optional operator to assemble to precondition
the system (if not provided a preconditioner may be
computed from ``a``)
:param dict form_compiler_parameters: parameters to pass to the form
compiler (optional)
:param constant_jacobian: (optional) flag indicating that the
Jacobian is constant (i.e. does not depend on
varying fields). If your Jacobian can change, set
this flag to ``False``.
"""
# In the linear case, the Jacobian is the equation LHS.
J = a
# Jacobian is checked in superclass, but let's check L here.
if L is 0: # noqa: F632
F = ufl_expr.action(J, u)
else:
if not isinstance(L, (ufl.Form, slate.slate.TensorBase)):
raise TypeError("Provided RHS is a '%s', not a Form or Slate Tensor" % type(L).__name__)
if len(L.arguments()) != 1:
raise ValueError("Provided RHS is not a linear form")
F = ufl_expr.action(J, u) - L
super(LinearVariationalProblem, self).__init__(F, u, bcs, J, aP,
form_compiler_parameters=form_compiler_parameters,
is_linear=True)
self._constant_jacobian = constant_jacobian
def __init__(self, a, L, u, bcs=None, aP=None,
form_compiler_parameters=None,
constant_jacobian=True):
r"""
:param a: the bilinear form
:param L: the linear form
:param u: the :class:`.Function` to solve for
:param bcs: the boundary conditions (optional)
:param aP: an optional operator to assemble to precondition
the system (if not provided a preconditioner may be
computed from ``a``)
:param dict form_compiler_parameters: parameters to pass to the form
compiler (optional)
:param constant_jacobian: (optional) flag indicating that the
Jacobian is constant (i.e. does not depend on
varying fields). If your Jacobian can change, set
this flag to ``False``.
"""
# In the linear case, the Jacobian is the equation LHS.
J = a
# Jacobian is checked in superclass, but let's check L here.
if L is 0: # noqa: F632
F = ufl_expr.action(J, u)
else:
if not isinstance(L, (ufl.Form, slate.slate.TensorBase)):
raise TypeError("Provided RHS is a '%s', not a Form or Slate Tensor" % type(L).__name__)
if len(L.arguments()) != 1:
raise ValueError("Provided RHS is not a linear form")
F = ufl_expr.action(J, u) - L
super(LinearVariationalProblem, self).__init__(F, u, bcs, J, aP,
form_compiler_parameters=form_compiler_parameters,
is_linear=True)
self._constant_jacobian = constant_jacobian
def __init__(self, *args, bcs=None, J=None, Jp=None, method="topological", V=None, is_linear=False, Jp_eq_J=False):
from firedrake.variational_solver import check_pde_args, is_form_consistent
if isinstance(args[0], ufl.classes.Equation):
# initial construction from equation
eq = args[0]
u = args[1]
sub_domain = args[2]
if V is None:
V = eq.lhs.arguments()[0].function_space()
bcs = solving._extract_bcs(bcs)
# Jp_eq_J is progressively evaluated as the tree is constructed
self.Jp_eq_J = Jp is None and all([bc.Jp_eq_J for bc in bcs])
# linear
if isinstance(eq.lhs, ufl.Form) and isinstance(eq.rhs, ufl.Form):
J = eq.lhs
Jp = Jp or J
if eq.rhs == 0:
F = ufl_expr.action(J, u)
else:
if not isinstance(eq.rhs, (ufl.Form, slate.slate.TensorBase)):
raise TypeError("Provided BC RHS is a '%s', not a Form or Slate Tensor" % type(eq.rhs).__name__)
if len(eq.rhs.arguments()) != 1:
raise ValueError("Provided BC RHS is not a linear form")
F = ufl_expr.action(J, u) - eq.rhs
self.is_linear = True
# nonlinear
else:
if eq.rhs != 0:
raise TypeError("RHS of a nonlinear form equation has to be 0")
F = eq.lhs
J = J or ufl_expr.derivative(F, u)
Jp = Jp or J
self.is_linear = False
# Check form style consistency
is_form_consistent(self.is_linear, bcs)
# Argument checking
check_pde_args(F, J, Jp)
# EquationBCSplit objects for `F`, `J`, and `Jp`
self._F = EquationBCSplit(F, u, sub_domain, bcs=[bc if isinstance(bc, DirichletBC) else bc._F for bc in bcs], method=method, V=V)
self._J = EquationBCSplit(J, u, sub_domain, bcs=[bc if isinstance(bc, DirichletBC) else bc._J for bc in bcs], method=method, V=V)
self._Jp = EquationBCSplit(Jp, u, sub_domain, bcs=[bc if isinstance(bc, DirichletBC) else bc._Jp for bc in bcs], method=method, V=V)
elif all(isinstance(args[i], EquationBCSplit) for i in range(3)):
# reconstruction for splitting `solving_utils.split`
self.Jp_eq_J = Jp_eq_J
self.is_linear = is_linear
self._F = args[0]
self._J = args[1]
self._Jp = args[2]
else:
raise TypeError("Wrong EquationBC arguments")
def __init__(self, *args, bcs=None, J=None, Jp=None, method="topological", V=None, is_linear=False, Jp_eq_J=False):
from firedrake.variational_solver import check_pde_args, is_form_consistent
if isinstance(args[0], ufl.classes.Equation):
# initial construction from equation
eq = args[0]
u = args[1]
sub_domain = args[2]
if V is None:
V = eq.lhs.arguments()[0].function_space()
bcs = solving._extract_bcs(bcs)
# Jp_eq_J is progressively evaluated as the tree is constructed
self.Jp_eq_J = Jp is None and all([bc.Jp_eq_J for bc in bcs])
# linear
if isinstance(eq.lhs, ufl.Form) and isinstance(eq.rhs, ufl.Form):
J = eq.lhs
Jp = Jp or J
if eq.rhs == 0:
F = ufl_expr.action(J, u)
else:
if not isinstance(eq.rhs, (ufl.Form, slate.slate.TensorBase)):
raise TypeError("Provided BC RHS is a '%s', not a Form or Slate Tensor" % type(eq.rhs).__name__)
if len(eq.rhs.arguments()) != 1:
raise ValueError("Provided BC RHS is not a linear form")
F = ufl_expr.action(J, u) - eq.rhs
self.is_linear = True
# nonlinear
else:
if eq.rhs != 0:
raise TypeError("RHS of a nonlinear form equation has to be 0")
F = eq.lhs
J = J or ufl_expr.derivative(F, u)
Jp = Jp or J
self.is_linear = False
# Check form style consistency
is_form_consistent(self.is_linear, bcs)
# Argument checking
check_pde_args(F, J, Jp)
# EquationBCSplit objects for `F`, `J`, and `Jp`
self._F = EquationBCSplit(F, u, sub_domain, bcs=[bc if isinstance(bc, DirichletBC) else bc._F for bc in bcs], method=method, V=V)
self._J = EquationBCSplit(J, u, sub_domain, bcs=[bc if isinstance(bc, DirichletBC) else bc._J for bc in bcs], method=method, V=V)
self._Jp = EquationBCSplit(Jp, u, sub_domain, bcs=[bc if isinstance(bc, DirichletBC) else bc._Jp for bc in bcs], method=method, V=V)
elif all(isinstance(args[i], EquationBCSplit) for i in range(3)):
# reconstruction for splitting `solving_utils.split`
self.Jp_eq_J = Jp_eq_J
self.is_linear = is_linear
self._F = args[0]
self._J = args[1]
self._Jp = args[2]
else:
raise TypeError("Wrong EquationBC arguments")
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 reconstruct(self, field=None, V=None, subu=None, u=None, row_field=None, col_field=None, action_x=None, use_split=False):
subu = subu or self.u
row_field = row_field or field
col_field = col_field or field
# define W and form
if field is None:
# Returns self
W = self._function_space
form = self.f
else:
assert V is not None, "`V` can not be `None` when `field` is not `None`"
W = self.as_subspace(field, V, use_split)
if W is None:
return
rank = len(self.f.arguments())
splitter = ExtractSubBlock()
if rank == 1:
form = splitter.split(self.f, argument_indices=(row_field, ))
elif rank == 2:
form = splitter.split(self.f, argument_indices=(row_field, col_field))
if u is not None:
form = replace(form, {self.u: u})
if action_x is not None:
assert len(form.arguments()) == 2, "rank of self.f must be 2 when using action_x parameter"
form = ufl_expr.action(form, action_x)
ebc = EquationBCSplit(form, subu, self.sub_domain, method=self.method, V=W)
for bc in self.bcs:
if isinstance(bc, DirichletBC):
ebc.add(bc.reconstruct(V=W, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, use_split=use_split))
elif isinstance(bc, EquationBCSplit):
bc_temp = bc.reconstruct(field=field, V=V, subu=subu, u=u, row_field=row_field, col_field=col_field, action_x=action_x, use_split=use_split)
# Due to the "if index", bc_temp can be None
if bc_temp is not None:
ebc.add(bc_temp)
return ebc
def reconstruct(self, field=None, V=None, subu=None, u=None, row_field=None, col_field=None, action_x=None, use_split=False):
subu = subu or self.u
row_field = row_field or field
col_field = col_field or field
# define W and form
if field is None:
# Returns self
W = self._function_space
form = self.f
else:
assert V is not None, "`V` can not be `None` when `field` is not `None`"
W = self.as_subspace(field, V, use_split)
if W is None:
return
rank = len(self.f.arguments())
splitter = ExtractSubBlock()
if rank == 1:
form = splitter.split(self.f, argument_indices=(row_field, ))
elif rank == 2:
form = splitter.split(self.f, argument_indices=(row_field, col_field))
if u is not None:
form = replace(form, {self.u: u})
if action_x is not None:
assert len(form.arguments()) == 2, "rank of self.f must be 2 when using action_x parameter"
form = ufl_expr.action(form, action_x)
ebc = EquationBCSplit(form, subu, self.sub_domain, method=self.method, V=W)
for bc in self.bcs:
if isinstance(bc, DirichletBC):
ebc.add(bc.reconstruct(V=W, g=bc.function_arg, sub_domain=bc.sub_domain, method=bc.method, use_split=use_split))
elif isinstance(bc, EquationBCSplit):
bc_temp = bc.reconstruct(field=field, V=V, subu=subu, u=u, row_field=row_field, col_field=col_field, action_x=action_x, use_split=use_split)
# Due to the "if index", bc_temp can be None
if bc_temp is not None:
ebc.add(bc_temp)
return ebc
def _rhs(self):
from firedrake.assemble import assemble
u = function.Function(self.trial_space)
b = function.Function(self.test_space)
expr = -action(self.A.a, u)
return u, functools.partial(assemble,
expr,
tensor=b,
assembly_type="residual"), b
def new_snes_ctx(pc, op, bcs, mat_type, fcp=None):
""" Create a new SNES contex for nested preconditioning
"""
from firedrake.variational_solver import NonlinearVariationalProblem
from firedrake.function import Function
from firedrake.ufl_expr import action
from firedrake.dmhooks import get_appctx
from firedrake.solving_utils import _SNESContext
dm = pc.getDM()
old_appctx = get_appctx(dm).appctx
u = Function(op.arguments()[-1].function_space())
F = action(op, u)
nprob = NonlinearVariationalProblem(F,
u,
bcs=bcs,
J=op,
form_compiler_parameters=fcp)
nctx = _SNESContext(nprob, mat_type, mat_type, old_appctx)
return nctx
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 _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 _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
