def __init__(self,
F,
butcher_tableau,
t,
dt,
u0,
bcs=None,
solver_parameters=None):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
bigF, stages, bigBCs, bigBCdata = \
getForm(F, butcher_tableau, t, dt, u0, bcs)
self.stages = stages
self.bigBCs = bigBCs
self.bigBCdata = bigBCdata
problem = NLVP(bigF, stages, bigBCs)
self.solver = NLVS(problem, solver_parameters=solver_parameters)
self.ks = stages.split()
def split(self, fields):
from ufl import as_vector, replace
from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
try:
if len(field) > 1:
raise NotImplementedError("Can't split into subblock")
except TypeError:
# Just a single field, we can handle that
pass
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
subu = us[field]
vec = []
for i, u in enumerate(us):
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp,
argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if bc.function_space().index == field:
V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
bcs.append(
type(bc)(V,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(
F,
subu,
bcs=bcs,
J=J,
Jp=None,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(
type(self)(new_problem,
mat_type=self.mat_type,
pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None,
solver_parameters=None, splitting=AI,
appctx=None, nullspace=None, bc_type="DAE"):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
bigF, stages, bigBCs, bigNSP, bigBCdata = \
getForm(F, butcher_tableau, t, dt, u0, bcs, bc_type, splitting, nullspace)
self.stages = stages
self.bigBCs = bigBCs
self.bigBCdata = bigBCdata
problem = NLVP(bigF, stages, bigBCs)
appctx_irksome = {"F": F,
"butcher_tableau": butcher_tableau,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"bc_type": bc_type,
"splitting": splitting,
"nullspace": nullspace}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
push_parent(u0.function_space().dm, stages.function_space().dm)
self.solver = NLVS(problem,
appctx=appctx,
solver_parameters=solver_parameters,
nullspace=bigNSP)
pop_parent(u0.function_space().dm, stages.function_space().dm)
if self.num_stages == 1 and self.num_fields == 1:
self.ws = (stages,)
else:
self.ws = stages.split()
A1, A2 = splitting(butcher_tableau.A)
try:
self.updateb = numpy.linalg.solve(A2.T, butcher_tableau.b)
except numpy.linalg.LinAlgError:
raise NotImplementedError("A=A1 A2 splitting needs A2 invertible")
boo = numpy.zeros(self.updateb.shape, dtype=self.updateb.dtype)
boo[-1] = 1
if numpy.allclose(self.updateb, boo):
self._update = self._update_A2Tmb
else:
self._update = self._update_general
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake import NonlinearVariationalProblem as NLVP
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
val, = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
Vbc = bc.function_space()
if Vbc.parent is not None and isinstance(Vbc.parent.ufl_element(), VectorElement):
index = Vbc.parent.index
else:
index = Vbc.index
cmpt = Vbc.component
# TODO: need to test this logic
if index in field:
if len(field) == 1:
W = V
else:
W = V.sub(field_renumbering[index])
if cmpt is not None:
W = W.sub(cmpt)
bcs.append(type(bc)(W,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)