Source code for irksome.discontinuous_galerkin_stepper
from functools import reduce
from FIAT import (Bernstein, DiscontinuousElement,
DiscontinuousLagrange,
IntegratedLegendre, Lagrange,
make_quadrature, ufc_simplex)
from operator import mul
from ufl.classes import Zero
from ufl.constantvalue import as_ufl
from .bcs import stage2spaces4bc
from .manipulation import extract_terms, strip_dt_form
from .stage import getBits
from .tools import getNullspace, replace
import numpy as np
from firedrake import as_vector, Constant, dot, TestFunction, Function, NonlinearVariationalProblem as NLVP, NonlinearVariationalSolver as NLVS
from firedrake.dmhooks import pop_parent, push_parent
[docs]
def getFormDiscGalerkin(F, L, Q, t, dt, u0, bcs=None, nullspace=None):
"""Given a time-dependent variational form, trial and test spaces, and
a quadrature rule, produce UFL for the Discontinuous Galerkin-in-Time method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg L: A :class:`FIAT.FiniteElement` for the test and trial functions in time
:arg Q: A :class:`FIAT.QuadratureRule` for the time integration
: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 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 (possibly 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 four parts:
- Fnew, the :class:`Form` corresponding to the DG-in-Time discretized problem
- UU, the :class:`Function` representing the stages to be solved for
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the Galerkin-in-time solution,
- 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object
that represents the nullspace of the coupled system
"""
assert Q.ref_el.get_spatial_dimension() == 1
assert L.get_reference_element() == Q.ref_el
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
vecconst = Constant
num_fields = len(V)
num_stages = L.space_dimension()
Vbig = reduce(mul, (V for _ in range(num_stages)))
VV = TestFunction(Vbig)
UU = Function(Vbig)
qpts = Q.get_points()
qwts = Q.get_weights()
tabulate_basis = L.tabulate(1, qpts)
basis_vals = tabulate_basis[(0,)]
basis_dvals = tabulate_basis[(1,)]
element = L
if isinstance(element, DiscontinuousElement):
element = element._element
# sort dofs geometrically by entity location
edofs = element.entity_dofs()
perm = [*edofs[0][0], *edofs[1][0], *edofs[0][1]]
basis_vals = basis_vals[perm]
basis_dvals = basis_dvals[perm]
# mass matrix later for BC
mmat = np.multiply(basis_vals, qwts) @ basis_vals.T
# L2 projector
proj = Constant(np.linalg.solve(mmat, np.multiply(basis_vals, qwts)))
u0bits, vbits, VVbits, UUbits = getBits(num_stages, num_fields,
u0, UU, v, VV)
split_form = extract_terms(F)
dtless = strip_dt_form(split_form.time)
Fnew = Zero()
basis_vals = vecconst(basis_vals)
basis_dvals = vecconst(basis_dvals)
qpts = vecconst(qpts.reshape((-1,)))
qwts = vecconst(qwts)
# Terms with time derivatives
for i in range(num_stages):
repl = {}
for j in range(num_fields):
repl[vbits[j]] = VVbits[i][j]
for ii in np.ndindex(vbits[j].ufl_shape):
repl[vbits[j][ii]] = VVbits[i][j][ii]
F_i = replace(dtless, repl)
# now loop over quadrature points
for q in range(len(qpts)):
repl = {t: t + dt * qpts[q]}
for k in range(num_fields):
d_tosub = sum(basis_dvals[ell, q] * UUbits[ell][k] for ell in range(num_stages)) / dt
repl[u0bits[k]] = d_tosub
for ii in np.ndindex(u0bits[k].ufl_shape):
d_tosub = sum(basis_dvals[ell, q] * UUbits[ell][k][ii]
for ell in range(num_stages)) / dt
repl[u0bits[k][ii]] = d_tosub
Fnew += dt * qwts[q] * basis_vals[i, q] * replace(F_i, repl)
# jump terms
repl = {}
for k in range(num_fields):
repl[u0bits[k]] = UUbits[0][k] - u0bits[k]
repl[vbits[k]] = VVbits[0][k]
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[0][k][ii] - u0bits[k][ii]
repl[vbits[k][ii]] = VVbits[0][k][ii]
Fnew += replace(dtless, repl)
# handle the rest of the terms
for i in range(num_stages):
repl = {}
for j in range(num_fields):
repl[vbits[j]] = VVbits[i][j]
for ii in np.ndindex(vbits[j].ufl_shape):
repl[vbits[j][ii]] = VVbits[i][j][ii]
F_i = replace(split_form.remainder, repl)
# now loop over quadrature points
for q in range(len(qpts)):
repl = {t: t + dt * qpts[q]}
for k in range(num_fields):
tosub = sum(basis_vals[ell, q] * UUbits[ell][k] for ell in range(num_stages))
repl[u0bits[k]] = tosub
for ii in np.ndindex(u0bits[k].ufl_shape):
tosub = sum(basis_vals[ell, q] * UUbits[ell][k][ii]
for ell in range(num_stages))
repl[u0bits[k][ii]] = tosub
Fnew += dt * qwts[q] * basis_vals[i, q] * replace(F_i, repl)
# Oh, honey, is it the boundary conditions?
if bcs is None:
bcs = []
bcsnew = []
for bc in bcs:
bcarg = as_ufl(bc._original_arg)
bcblah_at_qp = np.zeros((len(qpts),), dtype="O")
for q in range(len(qpts)):
tcur = t + qpts[q] * dt
bcblah_at_qp[q] = replace(bcarg, {t: tcur})
bc_func_for_stages = dot(proj, as_vector(bcblah_at_qp))
for i in range(num_stages):
Vbigi = stage2spaces4bc(bc, V, Vbig, i)
bcsnew.append(bc.reconstruct(V=Vbigi, g=bc_func_for_stages[i]))
return Fnew, UU, bcsnew, getNullspace(V, Vbig, num_stages, nullspace)
[docs]
class DiscontinuousGalerkinTimeStepper:
"""Front-end class for advancing a time-dependent PDE via a Discontinuous Galerkin
in time method
:arg F: A :class:`ufl.Form` instance describing the semi-discrete problem
F(t, u; v) == 0, where `u` is the unknown
:class:`firedrake.Function and `v` is the
:class:firedrake.TestFunction`.
:arg order: an integer indicating the order of the DG space to use
(with order == 0 corresponding to DG(0)-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 u0: A :class:`firedrake.Function` containing the current
state of the problem to be solved.
:arg bcs: An iterable of :class:`firedrake.DirichletBC` containing
the strongly-enforced boundary conditions. Irksome will
manipulate these to obtain boundary conditions for each
stage of the method.
:arg basis_type: A string indicating the finite element family (either
`'Lagrange'` or `'Bernstein'`) or the Lagrange variant for the
test/trial spaces. Defaults to equispaced Lagrange elements.
:arg quadrature: A :class:`FIAT.QuadratureRule` indicating the quadrature
to be used in time, defaulting to GL with order+1 points
:arg solver_parameters: A :class:`dict` of solver parameters that
will be used in solving the algebraic problem associated
with each time step.
:arg appctx: An optional :class:`dict` containing application context.
This gets included with particular things that Irksome will
pass into the nonlinear solver so that, say, user-defined preconditioners
have access to it.
: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
"""
def __init__(self, F, order, t, dt, u0, bcs=None, basis_type=None,
quadrature=None,
solver_parameters=None, appctx=None, nullspace=None):
assert order >= 0
self.u0 = u0
self.orig_bcs = bcs
self.t = t
self.dt = dt
self.order = order
self.basis_type = basis_type
V = u0.function_space()
self.num_fields = len(V)
ufc_line = ufc_simplex(1)
if order == 0:
self.el = DiscontinuousLagrange(ufc_line, 0)
elif basis_type == "Bernstein":
self.el = DiscontinuousElement(Bernstein(ufc_line, order))
elif basis_type == "integral":
self.el = DiscontinuousElement(IntegratedLegendre(ufc_line, order))
else:
# Let recursivenodes handle the general case
variant = None if basis_type == "Lagrange" else basis_type
self.el = DiscontinuousElement(Lagrange(ufc_line, order, variant=variant))
if quadrature is None:
quadrature = make_quadrature(ufc_line, order+1)
self.quadrature = quadrature
assert np.size(quadrature.get_points()) >= order+1
self.num_steps = 0
self.num_nonlinear_iterations = 0
self.num_linear_iterations = 0
bigF, UU, bigBCs, bigNSP = \
getFormDiscGalerkin(F, self.el,
quadrature, t, dt, u0, bcs, nullspace)
self.UU = UU
self.bigBCs = bigBCs
problem = NLVP(bigF, UU, bigBCs)
appctx_irksome = {"F": F,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"nullspace": nullspace}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
push_parent(u0.function_space().dm, UU.function_space().dm)
self.solver = NLVS(problem,
appctx=appctx,
solver_parameters=solver_parameters,
nullspace=bigNSP)
pop_parent(u0.function_space().dm, UU.function_space().dm)
[docs]
def advance(self):
"""Advances the system from time `t` to time `t + dt`.
Note: overwrites the value `u0`."""
push_parent(self.u0.function_space().dm, self.UU.function_space().dm)
self.solver.solve()
pop_parent(self.u0.function_space().dm, self.UU.function_space().dm)
self.num_steps += 1
self.num_nonlinear_iterations += self.solver.snes.getIterationNumber()
self.num_linear_iterations += self.solver.snes.getLinearSolveIterations()
u0 = self.u0
u0bits = u0.subfunctions
UUs = self.UU.subfunctions
for i, u0bit in enumerate(u0bits):
u0bit.assign(UUs[self.num_fields*(self.order)+i])
[docs]
def solver_stats(self):
return (self.num_steps, self.num_nonlinear_iterations, self.num_linear_iterations)