# -*- coding: utf-8 -*-
"""This module defines different MultiStageScheme classes based on the
Rush Larsen expicit integration schemes which can be passed to a
PointIntegralSolver
"""
# Copyright (C) 2014 Johan Hake
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
import functools
import ufl
from dolfin.function.constant import Constant
from dolfin.function.function import Function
from dolfin.fem.formmanipulations import derivative
from dolfin.multistage.factorize import extract_tested_expressions
from dolfin.fem.form import Form
from dolfin.multistage.multistagescheme import (MultiStageScheme,
_check_form, _time_dependent_expressions,
_replace_dict_time_dependent_expression, safe_action,
safe_adjoint)
from dolfin import DOLFIN_EPS
from dolfin import TrialFunction
import ufl.algorithms
from ufl.algorithms import (expand_derivatives, expand_indices,
extract_coefficients)
def _rush_larsen_step(rhs_exprs, diff_rhs_exprs, linear_terms,
system_size, y0, stage_solution, dt, time, a, c,
v, DX, time_dep_expressions):
# If we need to replace the original solution with stage solution
repl = None
if stage_solution is not None:
if system_size > 1:
repl = {y0: stage_solution}
else:
repl = {y0[0]: stage_solution}
# If we have time dependent expressions
if time_dep_expressions and abs(float(c)) > DOLFIN_EPS:
time_ = time
time = time + dt * float(c)
repl.update(_replace_dict_time_dependent_expression(time_dep_expressions,
time_, dt, float(c)))
repl[time_] = time
# If all terms are linear (using generalized=True) we add a safe
# guard to the linearized term. See below
safe_guard = sum(linear_terms) == system_size
# Add componentwise contribution to rl form
rl_ufl_form = ufl.zero()
num_rl_steps = 0
for ind in range(system_size):
# forward euler step
fe_du_i = rhs_exprs[ind] * dt * float(a)
# If exact integration
if linear_terms[ind]:
num_rl_steps += 1
# Rush Larsen step Safeguard the divisor: let's hope
# diff_rhs_exprs[ind] is never 1.0e-16! Let's get rid of
# this when the conditional fixes land properly in UFL.
eps = Constant(1.0e-16)
rl_du_i = rhs_exprs[ind] / (diff_rhs_exprs[ind] + eps) * (
ufl.exp(diff_rhs_exprs[ind] * dt) - 1.0)
# If safe guard
if safe_guard:
du_i = ufl.conditional(ufl.lt(abs(diff_rhs_exprs[ind]), 1e-8), fe_du_i, rl_du_i)
else:
du_i = rl_du_i
else:
du_i = fe_du_i
# If we should replace solution in form with stage solution
if repl:
du_i = ufl.replace(du_i, repl)
rl_ufl_form += (y0[ind] + du_i) * v[ind]
return rl_ufl_form * DX
def _find_linear_terms(rhs_exprs, u):
"""Help function that takes a list of rhs expressions and return a
list of bools determining what component, rhs_exprs[i], is linear
wrt u[i].
"""
uu = [Constant(1.0) for _ in rhs_exprs]
if len(rhs_exprs) > 1:
repl = {u: ufl.as_vector(uu)}
else:
repl = {u: uu[0]}
linear_terms = []
for i, ui in enumerate(uu):
comp_i_s = expand_indices(ufl.replace(rhs_exprs[i], repl))
linear_terms.append(ui in extract_coefficients(comp_i_s) and
ui not in extract_coefficients(
expand_derivatives(ufl.diff(comp_i_s, ui))))
return linear_terms
def _rush_larsen_scheme_generator(rhs_form, solution, time, order, generalized):
"""Generates a list of forms and solutions for a given Butcher
tableau
*Arguments*
rhs_form (ufl.Form)
A UFL form representing the rhs for a time differentiated equation
solution (_Function_)
The prognostic variable
time (_Constant_)
A Constant holding the time at the start of the time step
order (int)
The order of the scheme
generalized (bool)
If True generate a generalized Rush Larsen scheme, linearizing all
components.
"""
DX = _check_form(rhs_form)
if DX != ufl.dP:
raise TypeError("Expected a form with a Pointintegral.")
# Create time step
dt = Constant(0.1)
# Get test function
# arguments = rhs_form.arguments()
# coefficients = rhs_form.coefficients()
# Get time dependent expressions
time_dep_expressions = _time_dependent_expressions(rhs_form, time)
# Extract rhs expressions from form
rhs_integrand = rhs_form.integrals()[0].integrand()
rhs_exprs, v = extract_tested_expressions(rhs_integrand)
vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1
system_size = v.ufl_shape[0] if vector_rhs else 1
# Fix for indexing of v for scalar expressions
v = v if vector_rhs else [v]
# Extract linear terms if not using generalized Rush Larsen
if not generalized:
linear_terms = _find_linear_terms(rhs_exprs, solution)
else:
linear_terms = [True for _ in range(system_size)]
# Wrap the rhs expressions into a ufl vector type
rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
rhs_jac = ufl.diff(rhs_exprs, solution)
# Takes time!
if vector_rhs:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
for ind in range(system_size)]
else:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
solution = [solution]
ufl_stage_forms = []
dolfin_stage_forms = []
dt_stage_offsets = []
# Stage solutions (3 per order rhs, linearized, and final step)
# If 2nd order the final step for 1 step is a stage
if order == 1:
stage_solutions = []
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs, linear_terms,
system_size, solution, None, dt, time, 1.0,
0.0, v, DX, time_dep_expressions)
elif order == 2:
# Stage solution for order 2
if vector_rhs:
stage_solutions = [Function(solution.function_space(), name="y_1/2")]
else:
stage_solutions = [Function(solution[0].function_space(), name="y_1/2")]
stage_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 0.5,
0.0, v, DX,
time_dep_expressions)
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, stage_solutions[0],
dt, time, 1.0, 0.5, v, DX,
time_dep_expressions)
ufl_stage_forms.append([stage_form])
dolfin_stage_forms.append([Form(stage_form)])
# Get last stage form
last_stage = Form(rl_ufl_form)
human_form = "%srush larsen %s" % ("generalized " if generalized else "",
str(order))
return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
stage_solutions, dt, dt_stage_offsets, human_form, None
def _rush_larsen_scheme_generator_tlm(rhs_form, solution, time, order,
generalized, perturbation):
"""Generates a list of forms and solutions for the tangent
linearisation of the Rush-Larsen scheme
*Arguments*
rhs_form (ufl.Form)
A UFL form representing the rhs for a time differentiated equation
solution (_Function_)
The prognostic variable
time (_Constant_)
A Constant holding the time at the start of the time step
order (int)
The order of the scheme
generalized (bool)
If True generate a generalized Rush Larsen scheme, linearizing all
components.
perturbation (Function)
The vector on which we compute the tangent linear action.
"""
DX = _check_form(rhs_form)
if DX != ufl.dP:
raise TypeError("Expected a form with a Pointintegral.")
# Create time step
dt = Constant(0.1)
# Get test function
# arguments = rhs_form.arguments()
# coefficients = rhs_form.coefficients()
# Get time dependent expressions
time_dep_expressions = _time_dependent_expressions(rhs_form, time)
# Extract rhs expressions from form
rhs_integrand = rhs_form.integrals()[0].integrand()
rhs_exprs, v = extract_tested_expressions(rhs_integrand)
vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1
system_size = v.ufl_shape[0] if vector_rhs else 1
# Fix for indexing of v for scalar expressions
v = v if vector_rhs else [v]
# Extract linear terms if not using generalized Rush Larsen
if not generalized:
linear_terms = _find_linear_terms(rhs_exprs, solution)
else:
linear_terms = [True for _ in range(system_size)]
# Wrap the rhs expressions into a ufl vector type
rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
rhs_jac = ufl.diff(rhs_exprs, solution)
# Takes time!
if vector_rhs:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
for ind in range(system_size)]
soln = solution
else:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
solution = [solution]
soln = solution[0]
ufl_stage_forms = []
dolfin_stage_forms = []
dt_stage_offsets = []
trial = TrialFunction(soln.function_space())
# Stage solutions (3 per order rhs, linearized, and final step)
# If 2nd order the final step for 1 step is a stage
if order == 1:
stage_solutions = []
# Fetch the original step
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 1.0,
0.0, v, DX,
time_dep_expressions)
rl_ufl_form = safe_action(derivative(rl_ufl_form, soln, trial), perturbation)
elif order == 2:
# Stage solution for order 2
fn_space = soln.function_space()
stage_solutions = [Function(fn_space, name="y_1/2"),
Function(fn_space, name="y_dot_1/2"),
Function(fn_space, name="y_1")]
y_half_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 0.5,
0.0, v, DX,
time_dep_expressions)
y_dot_half_form = safe_action(derivative(y_half_form, soln, trial), perturbation)
y_one_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, stage_solutions[0],
dt, time, 1.0, 0.5, v, DX,
time_dep_expressions)
rl_ufl_form = safe_action(derivative(y_one_form, soln, trial), perturbation) + \
safe_action(derivative(y_one_form, stage_solutions[0], trial), stage_solutions[1])
ufl_stage_forms.append([y_half_form])
ufl_stage_forms.append([y_dot_half_form])
ufl_stage_forms.append([y_one_form])
dolfin_stage_forms.append([Form(y_half_form)])
dolfin_stage_forms.append([Form(y_dot_half_form)])
dolfin_stage_forms.append([Form(y_one_form)])
# Get last stage form
last_stage = Form(rl_ufl_form)
human_form = "%srush larsen %s" % ("generalized " if generalized else "",
str(order))
return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
stage_solutions, dt, dt_stage_offsets, human_form, perturbation
def _rush_larsen_scheme_generator_adm(rhs_form, solution, time, order,
generalized, perturbation):
"""Generates a list of forms and solutions for the adjoint
linearisation of the Rush-Larsen scheme
*Arguments*
rhs_form (ufl.Form)
A UFL form representing the rhs for a time differentiated equation
solution (_Function_)
The prognostic variable
time (_Constant_)
A Constant holding the time at the start of the time step
order (int)
The order of the scheme
generalized (bool)
If True generate a generalized Rush Larsen scheme, linearizing all
components.
perturbation (Function)
The vector on which we compute the adjoint action.
"""
DX = _check_form(rhs_form)
if DX != ufl.dP:
raise TypeError("Expected a form with a Pointintegral.")
# Create time step
dt = Constant(0.1)
# Get test function
# arguments = rhs_form.arguments()
# coefficients = rhs_form.coefficients()
# Get time dependent expressions
time_dep_expressions = _time_dependent_expressions(rhs_form, time)
# Extract rhs expressions from form
rhs_integrand = rhs_form.integrals()[0].integrand()
rhs_exprs, v = extract_tested_expressions(rhs_integrand)
vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1
system_size = v.ufl_shape[0] if vector_rhs else 1
# Fix for indexing of v for scalar expressions
v = v if vector_rhs else [v]
# Extract linear terms if not using generalized Rush Larsen
if not generalized:
linear_terms = _find_linear_terms(rhs_exprs, solution)
else:
linear_terms = [True for _ in range(system_size)]
# Wrap the rhs expressions into a ufl vector type
rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
rhs_jac = ufl.diff(rhs_exprs, solution)
# Takes time!
if vector_rhs:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
for ind in range(system_size)]
soln = solution
else:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
solution = [solution]
soln = solution[0]
ufl_stage_forms = []
dolfin_stage_forms = []
dt_stage_offsets = []
trial = TrialFunction(soln.function_space())
# Stage solutions (3 per order rhs, linearized, and final step)
# If 2nd order the final step for 1 step is a stage
if order == 1:
stage_solutions = []
# Fetch the original step
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 1.0,
0.0, v, DX,
time_dep_expressions)
# If this is commented out, we don't get NaNs. Yhy is
# solution a list of length zero anyway?
rl_ufl_form = safe_action(safe_adjoint(derivative(rl_ufl_form, soln, trial)),
perturbation)
elif order == 2:
# Stage solution for order 2
fn_space = soln.function_space()
stage_solutions = [Function(fn_space, name="y_1/2"),
Function(fn_space, name="y_1"),
Function(fn_space, name="y_bar_1/2")]
y_half_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 0.5,
0.0, v, DX,
time_dep_expressions)
y_one_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, stage_solutions[0],
dt, time, 1.0, 0.5, v, DX,
time_dep_expressions)
y_bar_half_form = safe_action(safe_adjoint(derivative(y_one_form,
stage_solutions[0], trial)), perturbation)
rl_ufl_form = safe_action(safe_adjoint(derivative(y_one_form, soln, trial)), perturbation) + \
safe_action(safe_adjoint(derivative(y_half_form, soln, trial)), stage_solutions[2])
ufl_stage_forms.append([y_half_form])
ufl_stage_forms.append([y_one_form])
ufl_stage_forms.append([y_bar_half_form])
dolfin_stage_forms.append([Form(y_half_form)])
dolfin_stage_forms.append([Form(y_one_form)])
dolfin_stage_forms.append([Form(y_bar_half_form)])
# Get last stage form
last_stage = Form(rl_ufl_form)
human_form = "%srush larsen %s" % ("generalized " if generalized else "",
str(order))
return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
stage_solutions, dt, dt_stage_offsets, human_form, perturbation
[docs]class RushLarsenScheme(MultiStageScheme):
def __init__(self, rhs_form, solution, time, order, generalized,
generator=_rush_larsen_scheme_generator):
self._rhs_form = rhs_form
self._solution = solution
self._t = time
self._order = order
self._generalized = generalized
# FIXME: What with bcs?
bcs = []
time = time or Constant(0.0)
if order not in [1, 2]:
raise ValueError("Expected order to be either 1 or 2")
rhs_form, ufl_stage_forms, linear_terms, dofin_stage_forms, last_stage, \
stage_solutions, dt, dt_stage_offsets, human_form, contraction = \
generator(rhs_form, solution, time, order, generalized)
self.linear_terms = linear_terms
# All stages are explicit
jacobian_indices = [-1 for _ in stage_solutions]
# Highjack a and b to hold parameters for RushLarsen scheme generation
MultiStageScheme.__init__(self, rhs_form, ufl_stage_forms,
dofin_stage_forms, last_stage,
stage_solutions, solution, time, dt,
dt_stage_offsets, jacobian_indices,
order, self.__class__.__name__,
human_form, bcs, contraction)
[docs] def to_tlm(self, perturbation):
r"""Return another RushLarsenScheme that implements the tangent
linearisation of the ODE solver.
This takes \dot{y_n} (the derivative of y_n with respect to a
parameter) and computes \dot{y_n+1} (the derivative of y_n+1
with respect to that parameter).
"""
generator = functools.partial(_rush_larsen_scheme_generator_tlm,
perturbation=perturbation)
new_solution = self._solution.copy()
new_form = ufl.replace(self._rhs_form, {self._solution: new_solution})
return RushLarsenScheme(new_form, new_solution, self._t, self._order,
self._generalized, generator=generator)
[docs] def to_adm(self, perturbation):
r"""Return another RushLarsenScheme that implements the adjoint
linearisation of the ODE solver.
This takes \bar{y_n+1} (the derivative of y_n+1 with respect to a
parameter) and computes \bar{y_n} (the derivative of y_n
with respect to that parameter).
"""
generator = functools.partial(_rush_larsen_scheme_generator_adm,
perturbation=perturbation)
new_solution = self._solution.copy()
new_form = ufl.replace(self._rhs_form, {self._solution: new_solution})
return RushLarsenScheme(new_form, new_solution, self._t,
self._order, self._generalized,
generator=generator)
[docs]class RL1(RushLarsenScheme):
"""First order Rush Larsen Scheme
"""
def __init__(self, rhs_form, solution, t=None):
RushLarsenScheme.__init__(self, rhs_form, solution, t, 1, False)
[docs]class RL2(RushLarsenScheme):
"""Second order Rush Larsen Scheme
"""
def __init__(self, rhs_form, solution, t=None):
RushLarsenScheme.__init__(self, rhs_form, solution, t, 2, generalized=False)
[docs]class GRL1(RushLarsenScheme):
"""
First order generalized Rush Larsen Scheme
"""
def __init__(self, rhs_form, solution, t=None):
RushLarsenScheme.__init__(self, rhs_form, solution, t, 1, True)
[docs]class GRL2(RushLarsenScheme):
"""Second order generalized Rush Larsen Scheme
"""
def __init__(self, rhs_form, solution, t=None):
RushLarsenScheme.__init__(self, rhs_form, solution, t, 2, True)
__all__ = [name for name, attr in list(globals().items())
if isinstance(attr, type) and issubclass(attr, MultiStageScheme)]
__all__.append("MultiStageScheme")