def test_is_zero_simple_vector_expressions():
mesh = UnitSquareMesh(4, 4)
V = FunctionSpace(mesh, 'CG', 1)
v = TestFunction(V)
u = TrialFunction(V)
check_is_zero(dot(as_vector([Zero(), u]), as_vector([Zero(), v])), 1)
check_is_zero(dot(as_vector([Zero(), u]), as_vector([v, Zero()])), 0)
def test_is_zero_simple_scalar_expressions():
mesh = UnitSquareMesh(4, 4)
V = FunctionSpace(mesh, 'CG', 1)
v = TestFunction(V)
u = TrialFunction(V)
check_is_zero(Zero() * u * v, 0)
check_is_zero(Constant(0) * u * v, 1)
check_is_zero(Zero() * v, 0)
check_is_zero(Constant(0) * v, 1)
check_is_zero(0 * u * v, 0)
check_is_zero(0 * v, 0)
check_is_zero(ufl.sin(0) * v, 0)
check_is_zero(ufl.cos(0) * v, 1)
def argument(self, arg):
"""
Argument is UFL lingo for test (num=0) and trial (num=1) functions
"""
# Do not make several new args for the same original arg
if arg in self._cache:
return self._cache[arg]
# Get some data about the argument and get our target index
V = arg.function_space()
N = V.num_sub_spaces()
num = arg.number()
idx_wanted = [self._index_v, self._index_u][num]
new_arg = []
for idx in range(N):
# Construct non-coupled argument
Vi = V.sub(idx).collapse()
a = dolfin.function.argument.Argument(Vi, num, part=arg.part())
indices = numpy.ndindex(a.ufl_shape)
if idx == idx_wanted:
# This is a wanted index
new_arg.extend(a[I] for I in indices)
else:
# This index should be pruned
new_arg.extend(Zero() for _ in indices)
new_arg = as_vector(new_arg)
self._cache[arg] = new_arg
return new_arg
def construct_modified_terminal(mt, terminal):
"""Construct a modified terminal given terminal modifiers from an
analysed modified terminal and a terminal."""
expr = terminal
if mt.reference_value:
expr = ReferenceValue(expr)
dim = expr.ufl_domain().topological_dimension()
for n in range(mt.local_derivatives):
# Return zero if expression is trivially constant. This has to
# happen here because ReferenceGrad has no access to the
# topological dimension of a literal zero.
if is_cellwise_constant(expr):
expr = Zero(expr.ufl_shape + (dim, ), expr.ufl_free_indices,
expr.ufl_index_dimensions)
else:
expr = ReferenceGrad(expr)
# No need to apply restrictions to ConstantValue terminals
if not isinstance(expr, ConstantValue):
if mt.restriction == '+':
expr = PositiveRestricted(expr)
elif mt.restriction == '-':
expr = NegativeRestricted(expr)
return expr
def argument(self, obj):
Q = obj.ufl_function_space()
dom = Q.ufl_domain()
sub_elements = obj.ufl_element().sub_elements()
# If not a mixed element, do nothing
if not isinstance(obj.ufl_element(), MixedElement):
return obj
# Split into sub-elements, creating appropriate space for each
args = []
for i, sub_elem in enumerate(sub_elements):
Q_i = FunctionSpace(dom, sub_elem)
a = Argument(Q_i, obj.number(), part=obj.part())
indices = [()]
for m in a.ufl_shape:
indices = [(k + (j, )) for k in indices for j in range(m)]
if (i == self.idx[obj.number()]):
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
return as_vector(args)
def zero(self, o):
fi = o.ufl_free_indices
fid = o.ufl_index_dimensions
mapped_fi = tuple(self.index(Index(count=i)) for i in fi)
paired_fid = [(mapped_fi[pos], fid[pos]) for pos, a in enumerate(fi)]
new_fi, new_fid = zip(*tuple(sorted(paired_fid)))
return Zero(o.ufl_shape, new_fi, new_fid)
def modified_terminal(self, o):
mt = analyse_modified_terminal(o)
t = mt.terminal
r = mt.restriction
if isinstance(t, Argument) and r != self.restrictions[t.number()]:
return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions)
else:
return o
def coefficient(self, o):
# Define dw/dw := d/ds [w + s v] = v
# Return corresponding argument if we can find o among w
do = self._w2v.get(o)
if do is not None:
return do
# Look for o among coefficient derivatives
dos = self._cd.get(o)
if dos is None:
# If o is not among coefficient derivatives, return
# do/dw=0
do = Zero(o.ufl_shape)
return do
else:
# Compute do/dw_j = do/dw_h : v.
# Since we may actually have a tuple of oprimes and vs in a
# 'mixed' space, sum over them all to get the complete inner
# product. Using indices to define a non-compound inner product.
# Example:
# (f:g) -> (dfdu:v):g + f:(dgdu:v)
# shape(dfdu) == shape(f) + shape(v)
# shape(f) == shape(g) == shape(dfdu : v)
# Make sure we have a tuple to match the self._v tuple
if not isinstance(dos, tuple):
dos = (dos, )
if len(dos) != len(self._v):
error(
"Got a tuple of arguments, expecting a matching tuple of coefficient derivatives."
)
dosum = Zero(o.ufl_shape)
for do, v in zip(dos, self._v):
so, oi = as_scalar(do)
rv = len(v.ufl_shape)
oi1 = oi[:-rv]
oi2 = oi[-rv:]
prod = so * v[oi2]
if oi1:
dosum += as_tensor(prod, oi1)
else:
dosum += prod
return dosum
def conditional(self, f): # TODO: Is this right? What about non-scalars?
c, a, b = f.operands()
s = f.shape()
ufl_assert(s == (), "TODO: Assuming scalar valued expressions.")
_0 = Zero()
_1 = IntValue(1)
da = conditional(c, _1, _0)
db = conditional(c, _0, _1)
return (None, da, db)
def coefficient_derivative(self, o, expr, coefficients, arguments, cds):
# If we're only taking a derivative wrt part of an argument in
# a mixed space other bits might come back as zero. We want to
# propagate a zero in that case.
argument, = arguments
if all(isinstance(a, Zero) for a in argument.ufl_operands):
return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions)
else:
return self.reuse_if_untouched(o, expr, coefficients, arguments, cds)
def argument(self, o):
from ufl import split
from firedrake import MixedFunctionSpace, FunctionSpace
V = o.function_space()
if len(V) == 1:
# Not on a mixed space, just return ourselves.
return o
if o in self._arg_cache:
return self._arg_cache[o]
V_is = V.split()
indices = self.blocks[o.number()]
try:
indices = tuple(indices)
nidx = len(indices)
except TypeError:
# Only one index provided.
indices = (indices, )
nidx = 1
if nidx == 1:
W = V_is[indices[0]]
W = FunctionSpace(W.mesh(), W.ufl_element())
a = (Argument(W, o.number(), part=o.part()), )
else:
W = MixedFunctionSpace([V_is[i] for i in indices])
a = split(Argument(W, o.number(), part=o.part()))
args = []
for i in range(len(V_is)):
if i in indices:
c = indices.index(i)
a_ = a[c]
if len(a_.ufl_shape) == 0:
args += [a_]
else:
args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
else:
args += [
Zero()
for j in numpy.ndindex(V_is[i].ufl_element().value_shape())
]
return self._arg_cache.setdefault(o, as_vector(args))
def getFormStage(F,
butch,
u0,
t,
dt,
bcs=None,
splitting=None,
nullspace=None):
"""Given a time-dependent variational form and a
:class:`ButcherTableau`, produce UFL for the s-stage RK method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg butch: the :class:`ButcherTableau` for the RK method being used to
advance in time.
:arg u0: a :class:`Function` referring to the state of
the PDE system at time `t`
:arg t: a :class:`Constant` referring to the current time level.
Any explicit time-dependence in F is included
:arg dt: a :class:`Constant` referring to the size of the current
time step.
:arg splitting: a callable that maps the (floating point) Butcher matrix
a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used
to vary between the classical RK formulation and Butcher's reformulation
that leads to a denser mass matrix with block-diagonal stiffness.
Only `AI` and `IA` are currently supported.
:arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof)
containing (possible time-dependent) boundary conditions imposed
on the system.
:arg nullspace: A list of tuples of the form (index, VSB) where
index is an index into the function space associated with `u`
and VSB is a :class: `firedrake.VectorSpaceBasis` instance to
be passed to a `firedrake.MixedVectorSpaceBasis` over the
larger space associated with the Runge-Kutta method
On output, we return a tuple consisting of several parts:
- Fnew, the :class:`Form`
- possibly a 4-tuple containing information needed to solve a mass matrix to update
the solution (this is empty for RadauIIA methods for which there is a trivial
update function.
- UU, the :class:`firedrake.Function` holding all the stage time values.
It lives in a :class:`firedrake.FunctionSpace` corresponding to the
s-way tensor product of the space on which the semidiscrete
form lives.
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the stages,
- 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object
that represents the nullspace of the coupled system
- `gblah`, a list of tuples of the form (f, expr, method),
where f is a :class:`firedrake.Function` and expr is a
:class:`ufl.Expr`. At each time step, each expr needs to be
re-interpolated/projected onto the corresponding f in order
for Firedrake to pick up that time-dependent boundary
conditions need to be re-applied. The
interpolation/projection is encoded in method, which is
either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending
on whether the function space for f supports interpolation or
not.
"""
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
num_stages = butch.num_stages
num_fields = len(V)
# s-way product space for the stage variables
Vbig = reduce(mul, (V for _ in range(num_stages)))
VV = TestFunction(Vbig)
UU = Function(Vbig)
vecconst = np.vectorize(Constant)
C = vecconst(butch.c)
A = vecconst(butch.A)
# set up the pieces we need to work with to do our substitutions
nsxnf = (num_stages, num_fields)
if num_fields == 1:
u0bits = np.array([u0], dtype="O")
vbits = np.array([v], dtype="O")
if num_stages == 1: # single-stage method
VVbits = np.array([[VV]], dtype="O")
UUbits = np.array([[UU]], dtype="O")
else: # multi-stage methods
VVbits = np.reshape(split(VV), nsxnf)
UUbits = np.reshape(split(UU), nsxnf)
else:
u0bits = np.array(list(split(u0)), dtype="O")
vbits = np.array(list(split(v)), dtype="O")
VVbits = np.reshape(split(VV), nsxnf)
UUbits = np.reshape(split(UU), nsxnf)
split_form = extract_terms(F)
Fnew = Zero()
# first, process terms with a time derivative. I'm
# assuming we have something of the form inner(Dt(g(u0)), v)*dx
# For each stage i, this gets replaced with
# inner((g(stages[i]) - g(u0))/dt, v)*dx
# but we have to carefully handle the cases where g indexes into
# pieces of u
dtless = strip_dt_form(split_form.time)
if splitting is None or splitting == AI:
for i in range(num_stages):
repl = {t: t + C[i] * dt}
for j in range(num_fields):
repl[u0bits[j]] = UUbits[i][j] - u0bits[j]
repl[vbits[j]] = VVbits[i][j]
# Also get replacements right for indexing.
for j in range(num_fields):
for ii in np.ndindex(u0bits[j].ufl_shape):
repl[u0bits[j][ii]] = UUbits[i][j][ii] - u0bits[j][ii]
repl[vbits[j][ii]] = VVbits[i][j][ii]
Fnew += replace(dtless, repl)
# Now for the non-time derivative parts
for i in range(num_stages):
# replace test function
repl = {}
for k in range(num_fields):
repl[vbits[k]] = VVbits[i][k]
for ii in np.ndindex(vbits[k].ufl_shape):
repl[vbits[k][ii]] = VVbits[i][k][ii]
Ftmp = replace(split_form.remainder, repl)
# replace the solution with stage values
for j in range(num_stages):
repl = {t: t + C[j] * dt}
for k in range(num_fields):
repl[u0bits[k]] = UUbits[j][k]
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[j][k][ii]
# and sum the contribution
Fnew += A[i, j] * dt * replace(Ftmp, repl)
elif splitting == IA:
Ainv = np.vectorize(Constant)(np.linalg.inv(butch.A))
# time derivative part gets inverse of Butcher matrix.
for i in range(num_stages):
repl = {}
for k in range(num_fields):
repl[vbits[k]] = VVbits[i][k]
for ii in np.ndindex(vbits[k].ufl_shape):
repl[vbits[k][ii]] = VVbits[i][k][ii]
Ftmp = replace(dtless, repl)
for j in range(num_stages):
repl = {t: t + C[j] * dt}
for k in range(num_fields):
repl[u0bits[k]] = (UUbits[j][k] - u0bits[k])
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[j][k][ii] - u0bits[k][ii]
Fnew += Ainv[i, j] * replace(Ftmp, repl)
# rest of the operator: just diagonal!
for i in range(num_stages):
repl = {t: t + C[i] * dt}
for j in range(num_fields):
repl[u0bits[j]] = UUbits[i][j]
repl[vbits[j]] = VVbits[i][j]
# Also get replacements right for indexing.
for j in range(num_fields):
for ii in np.ndindex(u0bits[j].ufl_shape):
repl[u0bits[j][ii]] = UUbits[i][j][ii]
repl[vbits[j][ii]] = VVbits[i][j][ii]
Fnew += dt * replace(split_form.remainder, repl)
else:
raise NotImplementedError("Can't handle that splitting type")
if bcs is None:
bcs = []
bcsnew = []
gblah = []
# For each BC, we need a new BC for each stage
# so we need to figure out how the function is indexed (mixed + vec)
# and then set it to have the value of the original argument at
# time t+C[i]*dt.
for bc in bcs:
if num_fields == 1: # not mixed space
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(comp)
Vbigi = lambda i: Vbig[i].sub(comp)
else:
Vsp = V
Vbigi = lambda i: Vbig[i]
else: # mixed space
sub = bc.function_space_index()
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(sub).sub(comp)
Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp)
else:
Vsp = V.sub(sub)
Vbigi = lambda i: Vbig[sub + num_fields * i]
bcarg = bc._original_arg
for i in range(num_stages):
try:
gdat = interpolate(bcarg, Vsp)
gmethod = lambda gd, gc: gd.interpolate(gc)
except: # noqa: E722
gdat = project(bcarg, Vsp)
gmethod = lambda gd, gc: gd.project(gc)
gcur = replace(bcarg, {t: t + C[i] * dt})
bcsnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain))
gblah.append((gdat, gcur, gmethod))
nspacenew = getNullspace(V, Vbig, butch, nullspace)
# For RIIA, we have an optimized update rule and don't need to
# build the variational form for doing updates.
# But something's broken with null spaces, so that's a TO-DO.
unew = Function(V)
Fupdate = inner(unew - u0, v) * dx
B = vectorize(Constant)(butch.b)
C = vectorize(Constant)(butch.c)
for i in range(num_stages):
repl = {t: t + C[i] * dt}
for k in range(num_fields):
repl[u0bits[k]] = UUbits[i][k]
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[i][k][ii]
eFFi = replace(split_form.remainder, repl)
Fupdate += dt * B[i] * eFFi
# And the BC's for the update -- just the original BC at t+dt
update_bcs = []
update_bcs_gblah = []
for bc in bcs:
if num_fields == 1: # not mixed space
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(comp)
else:
Vsp = V
else: # mixed space
sub = bc.function_space_index()
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(sub).sub(comp)
else:
Vsp = V.sub(sub)
bcarg = bc._original_arg
try:
gdat = interpolate(bcarg, Vsp)
gmethod = lambda gd, gc: gd.interpolate(gc)
except: # noqa: E722
gdat = project(bcarg, Vsp)
gmethod = lambda gd, gc: gd.project(gc)
gcur = replace(bcarg, {t: t + dt})
update_bcs.append(DirichletBC(Vsp, gdat, bc.sub_domain))
update_bcs_gblah.append((gdat, gcur, gmethod))
return (Fnew, (unew, Fupdate, update_bcs, update_bcs_gblah), UU, bcsnew,
gblah, nspacenew)
def independent_operator(self, o):
"Return a zero with the right shape and indices for operators independent of differentiation variable."
return Zero(o.ufl_shape + self._var_shape, o.ufl_free_indices,
o.ufl_index_dimensions)
def getForm(F, butch, t, dt, u0, bcs=None):
"""Given a time-dependent variational form and a
:class:`ButcherTableau`, produce UFL for the s-stage RK method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg butch: the :class:`ButcherTableau` for the RK method being used to
advance in time.
:arg t: a :class:`Constant` referring to the current time level.
Any explicit time-dependence in F is included
:arg dt: a :class:`Constant` referring to the size of the current
time step.
:arg u0: a :class:`Function` referring to the state of
the PDE system at time `t`
:arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof)
containing (possible time-dependent) boundary conditions imposed
on the system.
On output, we return a tuple consisting of four parts:
- Fnew, the :class:`Form`
- k, the :class:`firedrake.Function` holding all the stages.
It lives in a :class:`firedrake.FunctionSpace` corresponding to the
s-way tensor product of the space on which the semidiscrete
form lives.
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the stages,
- `gblah`, a list of pairs of the form (f, expr), where f is
a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`.
at each time step, each expr needs to be re-interpolated/projected
onto the corresponding f in order for Firedrake to pick up that
time-dependent boundary conditions need to be re-applied.
"""
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
A = numpy.array([[Constant(aa) for aa in arow] for arow in butch.A])
c = numpy.array([Constant(ci) for ci in butch.c])
num_stages = len(c)
num_fields = len(V)
Vbig = numpy.prod([V for i in range(num_stages)])
# Silence a warning about transfer managers when we
# coarsen coefficients in V
push_parent(V.dm, Vbig.dm)
vnew = TestFunction(Vbig)
k = Function(Vbig)
if len(V) == 1:
u0bits = [u0]
vbits = [v]
if num_stages == 1:
vbigbits = [vnew]
kbits = [k]
else:
vbigbits = split(vnew)
kbits = split(k)
else:
u0bits = split(u0)
vbits = split(v)
vbigbits = split(vnew)
kbits = split(k)
kbits_np = numpy.zeros((num_stages, num_fields), dtype="object")
for i in range(num_stages):
for j in range(num_fields):
kbits_np[i, j] = kbits[i * num_fields + j]
Ak = A @ kbits_np
Fnew = Zero()
for i in range(num_stages):
repl = {t: t + c[i] * dt}
for j, (ubit, vbit, kbit) in enumerate(zip(u0bits, vbits, kbits)):
repl[ubit] = ubit + dt * Ak[i, j]
repl[vbit] = vbigbits[num_fields * i + j]
repl[TimeDerivative(ubit)] = kbits_np[i, j]
if (len(ubit.ufl_shape) == 1):
for kk, kbitbit in enumerate(kbits_np[i, j]):
repl[TimeDerivative(ubit[kk])] = kbitbit
repl[ubit[kk]] = repl[ubit][kk]
repl[vbit[kk]] = repl[vbit][kk]
Fnew += replace(F, repl)
bcnew = []
gblah = []
if bcs is None:
bcs = []
for bc in bcs:
if isinstance(bc.domain_args[0], str):
boundary = bc.domain_args[0]
else:
boundary = ()
try:
for j in bc.sub_domain:
boundary += j
except TypeError:
boundary = (bc.sub_domain, )
gfoo = expand_derivatives(diff(bc._original_arg, t))
if len(V) == 1:
for i in range(num_stages):
gcur = replace(gfoo, {t: t + c[i] * dt})
try:
gdat = interpolate(gcur, V)
except NotImplementedError:
gdat = project(gcur, V)
gblah.append((gdat, gcur))
bcnew.append(DirichletBC(Vbig[i], gdat, boundary))
else:
sub = bc.function_space_index()
for i in range(num_stages):
gcur = replace(gfoo, {t: t + butch.c[i] * dt})
try:
gdat = interpolate(gcur, V.sub(sub))
except NotImplementedError:
gdat = project(gcur, V)
gblah.append((gdat, gcur))
bcnew.append(
DirichletBC(Vbig[sub + num_fields * i], gdat, boundary))
return Fnew, k, bcnew, gblah
def independent_terminal(self, o):
"Return a zero with the right shape for terminals independent of differentiation variable."
return Zero(o.ufl_shape + self._var_shape)
def grad(self, g):
# If we hit this type, it has already been propagated to a
# coefficient (or grad of a coefficient), # FIXME: Assert
# this! so we need to take the gradient of the variation or
# return zero. Complications occur when dealing with
# derivatives w.r.t. single components...
# Figure out how many gradients are around the inner terminal
ngrads = 0
o = g
while isinstance(o, Grad):
o, = o.ufl_operands
ngrads += 1
if not isinstance(o, FormArgument):
error("Expecting gradient of a FormArgument, not %s" %
ufl_err_str(o))
def apply_grads(f):
for i in range(ngrads):
f = Grad(f)
return f
# Find o among all w without any indexing, which makes this
# easy
for (w, v) in zip(self._w, self._v):
if o == w and isinstance(v, FormArgument):
# Case: d/dt [w + t v]
return apply_grads(v)
# If o is not among coefficient derivatives, return do/dw=0
gprimesum = Zero(g.ufl_shape)
def analyse_variation_argument(v):
# Analyse variation argument
if isinstance(v, FormArgument):
# Case: d/dt [w[...] + t v]
vval, vcomp = v, ()
elif isinstance(v, Indexed):
# Case: d/dt [w + t v[...]]
# Case: d/dt [w[...] + t v[...]]
vval, vcomp = v.ufl_operands
vcomp = tuple(vcomp)
else:
error("Expecting argument or component of argument.")
if not all(isinstance(k, FixedIndex) for k in vcomp):
error("Expecting only fixed indices in variation.")
return vval, vcomp
def compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp):
# Apply gradients directly to argument vval, and get the
# right indexed scalar component(s)
kk = indices(ngrads)
Dvkk = apply_grads(vval)[vcomp + kk]
# Place scalar component(s) Dvkk into the right tensor
# positions
if wshape:
Ejj, jj = unit_indexed_tensor(wshape, wcomp)
else:
Ejj, jj = 1, ()
gprimeterm = as_tensor(Ejj * Dvkk, jj + kk)
return gprimeterm
# Accumulate contributions from variations in different
# components
for (w, v) in zip(self._w, self._v):
# Analyse differentiation variable coefficient
if isinstance(w, FormArgument):
if not w == o:
continue
wshape = w.ufl_shape
if isinstance(v, FormArgument):
# Case: d/dt [w + t v]
return apply_grads(v)
elif isinstance(v, ListTensor):
# Case: d/dt [w + t <...,v,...>]
for wcomp, vsub in unwrap_list_tensor(v):
if not isinstance(vsub, Zero):
vval, vcomp = analyse_variation_argument(vsub)
gprimesum = gprimesum + compute_gprimeterm(
ngrads, vval, vcomp, wshape, wcomp)
else:
if wshape != ():
error("Expecting scalar coefficient in this branch.")
# Case: d/dt [w + t v[...]]
wval, wcomp = w, ()
vval, vcomp = analyse_variation_argument(v)
gprimesum = gprimesum + compute_gprimeterm(
ngrads, vval, vcomp, wshape, wcomp)
elif isinstance(
w, Indexed
): # This path is tested in unit tests, but not actually used?
# Case: d/dt [w[...] + t v[...]]
# Case: d/dt [w[...] + t v]
wval, wcomp = w.ufl_operands
if not wval == o:
continue
assert isinstance(wval, FormArgument)
if not all(isinstance(k, FixedIndex) for k in wcomp):
error(
"Expecting only fixed indices in differentiation variable."
)
wshape = wval.ufl_shape
vval, vcomp = analyse_variation_argument(v)
gprimesum = gprimesum + compute_gprimeterm(
ngrads, vval, vcomp, wshape, wcomp)
else:
error("Expecting coefficient or component of coefficient.")
# FIXME: Handle other coefficient derivatives: oprimes =
# self._cd.get(o)
if 0:
oprimes = self._cd.get(o)
if oprimes is None:
if self._cd:
# TODO: Make it possible to silence this message
# in particular? It may be good to have for
# debugging...
warning("Assuming d{%s}/d{%s} = 0." % (o, self._w))
else:
# Make sure we have a tuple to match the self._v tuple
if not isinstance(oprimes, tuple):
oprimes = (oprimes, )
if len(oprimes) != len(self._v):
error("Got a tuple of arguments, expecting a"
" matching tuple of coefficient derivatives.")
# Compute dg/dw_j = dg/dw_h : v.
# Since we may actually have a tuple of oprimes and vs
# in a 'mixed' space, sum over them all to get the
# complete inner product. Using indices to define a
# non-compound inner product.
for (oprime, v) in zip(oprimes, self._v):
error("FIXME: Figure out how to do this with ngrads")
so, oi = as_scalar(oprime)
rv = len(v.ufl_shape)
oi1 = oi[:-rv]
oi2 = oi[-rv:]
prod = so * v[oi2]
if oi1:
gprimesum += as_tensor(prod, oi1)
else:
gprimesum += prod
return gprimesum
def ConstantOrZero(x):
return Zero() if abs(complex(x)) < 1.e-10 else Constant(x)
def getForm(F,
butch,
t,
dt,
u0,
bcs=None,
bc_type="DAE",
splitting=AI,
nullspace=None):
"""Given a time-dependent variational form and a
:class:`ButcherTableau`, produce UFL for the s-stage RK method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg butch: the :class:`ButcherTableau` for the RK method being used to
advance in time.
:arg t: a :class:`Constant` referring to the current time level.
Any explicit time-dependence in F is included
:arg dt: a :class:`Constant` referring to the size of the current
time step.
:arg splitting: a callable that maps the (floating point) Butcher matrix
a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used
to vary between the classical RK formulation and Butcher's reformulation
that leads to a denser mass matrix with block-diagonal stiffness.
Some choices of function will assume that `butch.A` is invertible.
:arg u0: a :class:`Function` referring to the state of
the PDE system at time `t`
:arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof)
containing (possible time-dependent) boundary conditions imposed
on the system.
:arg bc_type: How to manipulate the strongly-enforced boundary
conditions to derive the stage boundary conditions. Should
be a string, either "DAE", which implements BCs as
constraints in the style of a differential-algebraic
equation, or "ODE", which takes the time derivative of the
boundary data and evaluates this for the stage values
:arg nullspace: A list of tuples of the form (index, VSB) where
index is an index into the function space associated with `u`
and VSB is a :class: `firedrake.VectorSpaceBasis` instance to
be passed to a `firedrake.MixedVectorSpaceBasis` over the
larger space associated with the Runge-Kutta method
On output, we return a tuple consisting of four parts:
- Fnew, the :class:`Form`
- k, the :class:`firedrake.Function` holding all the stages.
It lives in a :class:`firedrake.FunctionSpace` corresponding to the
s-way tensor product of the space on which the semidiscrete
form lives.
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the stages,
- 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object
that represents the nullspace of the coupled system
- `gblah`, a list of tuples of the form (f, expr, method),
where f is a :class:`firedrake.Function` and expr is a
:class:`ufl.Expr`. At each time step, each expr needs to be
re-interpolated/projected onto the corresponding f in order
for Firedrake to pick up that time-dependent boundary
conditions need to be re-applied. The
interpolation/projection is encoded in method, which is
either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending
on whether the function space for f supports interpolation or
not.
"""
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
c = numpy.array([Constant(ci) for ci in butch.c], dtype=object)
bA1, bA2 = splitting(butch.A)
try:
bA1inv = numpy.linalg.inv(bA1)
except numpy.linalg.LinAlgError:
bA1inv = None
try:
bA2inv = numpy.linalg.inv(bA2)
A2inv = numpy.array([[ConstantOrZero(aa) for aa in arow]
for arow in bA2inv],
dtype=object)
except numpy.linalg.LinAlgError:
raise NotImplementedError("We require A = A1 A2 with A2 invertible")
A1 = numpy.array([[ConstantOrZero(aa) for aa in arow] for arow in bA1],
dtype=object)
if bA1inv is not None:
A1inv = numpy.array([[ConstantOrZero(aa) for aa in arow]
for arow in bA1inv],
dtype=object)
else:
A1inv = None
num_stages = butch.num_stages
num_fields = len(V)
Vbig = reduce(mul, (V for _ in range(num_stages)))
vnew = TestFunction(Vbig)
w = Function(Vbig)
if len(V) == 1:
u0bits = [u0]
vbits = [v]
if num_stages == 1:
vbigbits = [vnew]
wbits = [w]
else:
vbigbits = split(vnew)
wbits = split(w)
else:
u0bits = split(u0)
vbits = split(v)
vbigbits = split(vnew)
wbits = split(w)
wbits_np = numpy.zeros((num_stages, num_fields), dtype=object)
for i in range(num_stages):
for j in range(num_fields):
wbits_np[i, j] = wbits[i * num_fields + j]
A1w = A1 @ wbits_np
A2invw = A2inv @ wbits_np
Fnew = Zero()
for i in range(num_stages):
repl = {t: t + c[i] * dt}
for j, (ubit, vbit) in enumerate(zip(u0bits, vbits)):
repl[ubit] = ubit + dt * A1w[i, j]
repl[vbit] = vbigbits[num_fields * i + j]
repl[TimeDerivative(ubit)] = A2invw[i, j]
if (len(ubit.ufl_shape) == 1):
for kk in range(len(A1w[i, j])):
repl[TimeDerivative(ubit[kk])] = A2invw[i, j][kk]
repl[ubit[kk]] = repl[ubit][kk]
repl[vbit[kk]] = repl[vbit][kk]
Fnew += replace(F, repl)
bcnew = []
gblah = []
if bcs is None:
bcs = []
if bc_type == "ODE":
assert splitting == AI, "ODE-type BC aren't implemented for this splitting strategy"
u0_mult_np = numpy.divide(1.0,
butch.c,
out=numpy.zeros_like(butch.c),
where=butch.c != 0)
u0_mult = numpy.array([ConstantOrZero(mi) / dt for mi in u0_mult_np],
dtype=object)
def bc2gcur(bc, i):
gorig = bc._original_arg
gfoo = expand_derivatives(diff(gorig, t))
return replace(gfoo, {t: t + c[i] * dt}) + u0_mult[i] * gorig
elif bc_type == "DAE":
if bA1inv is None:
raise NotImplementedError(
"Cannot have DAE BCs for this Butcher Tableau/splitting")
u0_mult_np = A1inv @ numpy.ones_like(butch.c)
u0_mult = numpy.array([ConstantOrZero(mi) / dt for mi in u0_mult_np],
dtype=object)
def bc2gcur(bc, i):
gorig = as_ufl(bc._original_arg)
gcur = 0
for j in range(num_stages):
gcur += ConstantOrZero(bA1inv[i, j]) / dt * replace(
gorig, {t: t + c[j] * dt})
return gcur
else:
raise ValueError("Unrecognised bc_type: %s", bc_type)
# This logic uses information set up in the previous section to
# set up the new BCs for either method
for bc in bcs:
if num_fields == 1: # not mixed space
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(comp)
Vbigi = lambda i: Vbig[i].sub(comp)
else:
Vsp = V
Vbigi = lambda i: Vbig[i]
else: # mixed space
sub = bc.function_space_index()
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(sub).sub(comp)
Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp)
else:
Vsp = V.sub(sub)
Vbigi = lambda i: Vbig[sub + num_fields * i]
for i in range(num_stages):
gcur = bc2gcur(bc, i)
blah = BCStageData(Vsp, gcur, u0, u0_mult, i, t, dt)
gdat, gcr, gmethod = blah.gstuff
gblah.append((gdat, gcr, gmethod))
bcnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain))
nspnew = getNullspace(V, Vbig, butch, nullspace)
return Fnew, w, bcnew, nspnew, gblah
def negative_restricted(self, o):
return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions)
