Source code for irksome.getForm
from functools import reduce
from operator import mul
import numpy
from firedrake import Function, TestFunction, split
from ufl import diff
from ufl.algorithms import expand_derivatives
from ufl.classes import Zero
from ufl.constantvalue import as_ufl
from .tools import ConstantOrZero, MeshConstant, replace, getNullspace, AI
from .deriv import TimeDerivative # , apply_time_derivatives
from .bcs import BCStageData, bc2space, stage2spaces4bc
[docs]
def getForm(F, butch, t, dt, u0, bcs=None, bc_type=None, 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:`Function` on the Real space over the same mesh as
`u0`. This serves as a variable referring to the current time.
:arg dt: a :class:`Function` on the Real space over the same mesh as
`u0`. This serves as a variable referring to the current time step.
The user may adjust this value between time steps.
: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
"""
if bc_type is None:
bc_type = "DAE"
v = F.arguments()[0]
V = v.function_space()
msh = V.mesh()
assert V == u0.function_space()
MC = MeshConstant(msh)
c = numpy.array([MC.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, MC) 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, MC) for aa in arow] for arow in bA1],
dtype=object)
if bA1inv is not None:
A1inv = numpy.array([[ConstantOrZero(aa, MC) 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 = []
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([MC.Constant(0) for mi in u0_mult_np],
dtype=object)
def bc2gcur(bc, i):
gorig = as_ufl(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, MC)/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], MC) / 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:
for i in range(num_stages):
Vsp = bc2space(bc, V)
Vbigi = stage2spaces4bc(bc, V, Vbig, i)
gcur = bc2gcur(bc, i)
gdat = BCStageData(Vsp, gcur, u0, u0_mult, i, t, dt)
bcnew.append(bc.reconstruct(V=Vbigi, g=gdat))
nspnew = getNullspace(V, Vbig, butch, nullspace)
return Fnew, w, bcnew, nspnew