-
Notifications
You must be signed in to change notification settings - Fork 12
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Increment form for implicit RK added and tested #566
Changes from 4 commits
90d1f7c
93e20d2
8f8d084
4e72bb2
8b21447
3414a27
1c7f86e
f812c98
5250ad5
704ba3f
249dfd3
ff9f524
eeafdd9
d1afa20
8e46d7d
9b560eb
81ba663
9fa02b5
e4f36e3
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,14 +3,15 @@ | |
import numpy as np | ||
|
||
from firedrake import (Function, split, NonlinearVariationalProblem, | ||
NonlinearVariationalSolver) | ||
NonlinearVariationalSolver, Constant) | ||
from firedrake.fml import replace_subject, all_terms, drop | ||
from firedrake.utils import cached_property | ||
|
||
from gusto.core.labels import time_derivative | ||
from gusto.time_discretisation.time_discretisation import ( | ||
TimeDiscretisation, wrapper_apply | ||
) | ||
from gusto.time_discretisation.explicit_runge_kutta import RungeKuttaFormulation | ||
|
||
|
||
__all__ = ["ImplicitRungeKutta", "ImplicitMidpoint", "QinZhang"] | ||
|
@@ -56,6 +57,7 @@ class ImplicitRungeKutta(TimeDiscretisation): | |
# --------------------------------------------------------------------------- | ||
|
||
def __init__(self, domain, butcher_matrix, field_name=None, | ||
rk_formulation=RungeKuttaFormulation.increment, | ||
solver_parameters=None, options=None,): | ||
""" | ||
Args: | ||
|
@@ -66,6 +68,9 @@ def __init__(self, domain, butcher_matrix, field_name=None, | |
discretisation. | ||
field_name (str, optional): name of the field to be evolved. | ||
Defaults to None. | ||
rk_formulation (:class:`RungeKuttaFormulation`, optional): | ||
an enumerator object, describing the formulation of the Runge- | ||
Kutta scheme. Defaults to the increment form. | ||
solver_parameters (dict, optional): dictionary of parameters to | ||
pass to the underlying solver. Defaults to None. | ||
options (:class:`AdvectionOptions`, optional): an object containing | ||
|
@@ -78,6 +83,7 @@ def __init__(self, domain, butcher_matrix, field_name=None, | |
options=options) | ||
self.butcher_matrix = butcher_matrix | ||
self.nStages = int(np.shape(self.butcher_matrix)[1]) | ||
self.rk_formulation = rk_formulation | ||
|
||
def setup(self, equation, apply_bcs=True, *active_labels): | ||
""" | ||
|
@@ -91,31 +97,108 @@ def setup(self, equation, apply_bcs=True, *active_labels): | |
|
||
super().setup(equation, apply_bcs, *active_labels) | ||
|
||
self.k = [Function(self.fs) for i in range(self.nStages)] | ||
if self.rk_formulation == RungeKuttaFormulation.predictor: | ||
self.xs = [Function(self.fs) for _ in range(self.nStages)] | ||
elif self.rk_formulation == RungeKuttaFormulation.increment: | ||
self.k = [Function(self.fs) for _ in range(self.nStages)] | ||
elif self.rk_formulation == RungeKuttaFormulation.linear: | ||
raise NotImplementedError( | ||
'Linear Implicit Runge-Kutta formulation is not implemented' | ||
) | ||
else: | ||
raise NotImplementedError( | ||
'Runge-Kutta formulation is not implemented' | ||
) | ||
|
||
def lhs(self): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do you know we don't set up the
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I have removed lhs & rhs and stopped them being an abstract property There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Instead they are now just a property There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for making this change. @jshipton are you happy with this? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I have now removed all lhs & rhs. Each time discretisation just has a res (residual). |
||
return super().lhs | ||
|
||
def rhs(self): | ||
return super().rhs | ||
|
||
def solver(self, stage): | ||
residual = self.residual.label_map( | ||
lambda t: t.has_label(time_derivative), | ||
map_if_true=drop, | ||
map_if_false=replace_subject(self.xnph, self.idx), | ||
) | ||
def res(self, stage): | ||
"""Set up the discretisation's residual for a given stage.""" | ||
atb1995 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
# Add time derivative terms y_s - y^n for stage s | ||
mass_form = self.residual.label_map( | ||
lambda t: t.has_label(time_derivative), | ||
map_if_false=drop) | ||
residual += mass_form.label_map(all_terms, | ||
replace_subject(self.x_out, self.idx)) | ||
residual = mass_form.label_map(all_terms, | ||
map_if_true=replace_subject(self.x_out, old_idx=self.idx)) | ||
residual -= mass_form.label_map(all_terms, | ||
map_if_true=replace_subject(self.x1, old_idx=self.idx)) | ||
# Loop through stages up to s-1 and calcualte/sum | ||
atb1995 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
# dt*(a_s1*F(y_1) + a_s2*F(y_2)+ ... + a_{s,s-1}*F(y_{s-1})) | ||
for i in range(stage): | ||
r_imp = self.residual.label_map( | ||
lambda t: not t.has_label(time_derivative), | ||
map_if_true=replace_subject(self.xs[i], old_idx=self.idx), | ||
map_if_false=drop) | ||
r_imp = r_imp.label_map( | ||
all_terms, | ||
map_if_true=lambda t: Constant(self.butcher_matrix[stage, i])*self.dt*t) | ||
residual += r_imp | ||
# Calculate and add on dt*a_ss*F(y_s) | ||
r_imp = self.residual.label_map( | ||
lambda t: not t.has_label(time_derivative), | ||
map_if_true=replace_subject(self.x_out, old_idx=self.idx), | ||
map_if_false=drop) | ||
r_imp = r_imp.label_map( | ||
all_terms, | ||
map_if_true=lambda t: Constant(self.butcher_matrix[stage, stage])*self.dt*t) | ||
residual += r_imp | ||
return residual.form | ||
|
||
@property | ||
def final_res(self): | ||
"""Set up the discretisation's final residual.""" | ||
# Add time derivative terms y^{n+1} - y^n | ||
atb1995 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), | ||
map_if_false=drop) | ||
residual = mass_form.label_map(all_terms, | ||
map_if_true=replace_subject(self.x_out, old_idx=self.idx)) | ||
residual -= mass_form.label_map(all_terms, | ||
map_if_true=replace_subject(self.x1, old_idx=self.idx)) | ||
# Loop through stages up to s-1 and calcualte/sum | ||
atb1995 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
# dt*(b_1*F(y_1) + b_2*F(y_2) + .... + b_s*F(y_s)) | ||
for i in range(self.nStages): | ||
r_imp = self.residual.label_map( | ||
lambda t: not t.has_label(time_derivative), | ||
map_if_true=replace_subject(self.xs[i], old_idx=self.idx), | ||
map_if_false=drop) | ||
r_imp = r_imp.label_map( | ||
all_terms, | ||
map_if_true=lambda t: Constant(self.butcher_matrix[self.nStages, i])*self.dt*t) | ||
residual += r_imp | ||
return residual.form | ||
|
||
problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) | ||
def solver(self, stage): | ||
if self.rk_formulation == RungeKuttaFormulation.increment: | ||
residual = self.residual.label_map( | ||
lambda t: t.has_label(time_derivative), | ||
map_if_true=drop, | ||
map_if_false=replace_subject(self.xnph, self.idx), | ||
) | ||
mass_form = self.residual.label_map( | ||
lambda t: t.has_label(time_derivative), | ||
map_if_false=drop) | ||
residual += mass_form.label_map(all_terms, | ||
replace_subject(self.x_out, self.idx)) | ||
|
||
problem = NonlinearVariationalProblem(residual.form, self.x_out, bcs=self.bcs) | ||
|
||
elif self.rk_formulation == RungeKuttaFormulation.predictor: | ||
problem = NonlinearVariationalProblem(self.res(stage), self.x_out, bcs=self.bcs) | ||
|
||
solver_name = self.field_name+self.__class__.__name__ + "%s" % (stage) | ||
return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, | ||
options_prefix=solver_name) | ||
return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) | ||
|
||
@cached_property | ||
def final_solver(self): | ||
"""Set up a solver for the final solve to evaluate time level n+1.""" | ||
# setup solver using lhs and rhs defined in derived class | ||
problem = NonlinearVariationalProblem(self.final_res, self.x_out, bcs=self.bcs) | ||
solver_name = self.field_name+self.__class__.__name__ | ||
return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) | ||
|
||
@cached_property | ||
def solvers(self): | ||
|
@@ -126,32 +209,48 @@ def solvers(self): | |
|
||
def solve_stage(self, x0, stage): | ||
self.x1.assign(x0) | ||
for i in range(stage): | ||
self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) | ||
if self.rk_formulation == RungeKuttaFormulation.increment: | ||
for i in range(stage): | ||
self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) | ||
|
||
if self.idx is None and len(self.fs) > 1: | ||
self.xnph = tuple([self.dt*self.butcher_matrix[stage, stage]*a + b | ||
for a, b in zip(split(self.x_out), split(self.x1))]) | ||
else: | ||
self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out | ||
solver = self.solvers[stage] | ||
# Set initial guess for solver | ||
if (stage > 0): | ||
self.x_out.assign(self.k[stage-1]) | ||
if self.idx is None and len(self.fs) > 1: | ||
self.xnph = tuple( | ||
self.dt * self.butcher_matrix[stage, stage] * a + b | ||
for a, b in zip(split(self.x_out), split(self.x1)) | ||
) | ||
else: | ||
self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out | ||
|
||
solver = self.solvers[stage] | ||
|
||
solver.solve() | ||
# Set initial guess for solver | ||
if (stage > 0): | ||
self.x_out.assign(self.k[stage-1]) | ||
|
||
self.k[stage].assign(self.x_out) | ||
solver.solve() | ||
self.k[stage].assign(self.x_out) | ||
|
||
elif self.rk_formulation == RungeKuttaFormulation.predictor: | ||
if (stage > 0): | ||
self.x_out.assign(self.xs[stage-1]) | ||
solver = self.solvers[stage] | ||
solver.solve() | ||
|
||
self.xs[stage].assign(self.x_out) | ||
|
||
@wrapper_apply | ||
def apply(self, x_out, x_in): | ||
|
||
self.x_out.assign(x_in) | ||
for i in range(self.nStages): | ||
self.solve_stage(x_in, i) | ||
|
||
x_out.assign(x_in) | ||
for i in range(self.nStages): | ||
x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) | ||
if self.rk_formulation == RungeKuttaFormulation.increment: | ||
x_out.assign(x_in) | ||
for i in range(self.nStages): | ||
x_out.assign(x_out + self.butcher_matrix[self.nStages, i]*self.dt*self.k[i]) | ||
elif self.rk_formulation == RungeKuttaFormulation.predictor: | ||
self.final_solver.solve() | ||
x_out.assign(self.x_out) | ||
|
||
|
||
class ImplicitMidpoint(ImplicitRungeKutta): | ||
|
@@ -164,14 +263,18 @@ class ImplicitMidpoint(ImplicitRungeKutta): | |
k0 = F[y^n + 0.5*dt*k0] \n | ||
y^(n+1) = y^n + dt*k0 \n | ||
""" | ||
def __init__(self, domain, field_name=None, solver_parameters=None, | ||
options=None): | ||
def __init__(self, domain, field_name=None, | ||
rk_formulation=RungeKuttaFormulation.increment, | ||
solver_parameters=None, options=None): | ||
""" | ||
Args: | ||
domain (:class:`Domain`): the model's domain object, containing the | ||
mesh and the compatible function spaces. | ||
field_name (str, optional): name of the field to be evolved. | ||
Defaults to None. | ||
rk_formulation (:class:`RungeKuttaFormulation`, optional): | ||
an enumerator object, describing the formulation of the Runge- | ||
Kutta scheme. Defaults to the increment form. | ||
solver_parameters (dict, optional): dictionary of parameters to | ||
pass to the underlying solver. Defaults to None. | ||
options (:class:`AdvectionOptions`, optional): an object containing | ||
|
@@ -181,6 +284,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, | |
""" | ||
butcher_matrix = np.array([[0.5], [1.]]) | ||
super().__init__(domain, butcher_matrix, field_name, | ||
rk_formulation=rk_formulation, | ||
solver_parameters=solver_parameters, | ||
options=options) | ||
|
||
|
@@ -196,14 +300,18 @@ class QinZhang(ImplicitRungeKutta): | |
k1 = F[y^n + 0.5*dt*k0 + 0.25*dt*k1] \n | ||
y^(n+1) = y^n + 0.5*dt*(k0 + k1) \n | ||
""" | ||
def __init__(self, domain, field_name=None, solver_parameters=None, | ||
options=None): | ||
def __init__(self, domain, field_name=None, | ||
rk_formulation=RungeKuttaFormulation.increment, | ||
solver_parameters=None, options=None): | ||
""" | ||
Args: | ||
domain (:class:`Domain`): the model's domain object, containing the | ||
mesh and the compatible function spaces. | ||
field_name (str, optional): name of the field to be evolved. | ||
Defaults to None. | ||
rk_formulation (:class:`RungeKuttaFormulation`, optional): | ||
an enumerator object, describing the formulation of the Runge- | ||
Kutta scheme. Defaults to the increment form. | ||
solver_parameters (dict, optional): dictionary of parameters to | ||
pass to the underlying solver. Defaults to None. | ||
options (:class:`AdvectionOptions`, optional): an object containing | ||
|
@@ -213,5 +321,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None, | |
""" | ||
butcher_matrix = np.array([[0.25, 0], [0.5, 0.25], [0.5, 0.5]]) | ||
super().__init__(domain, butcher_matrix, field_name, | ||
rk_formulation=rk_formulation, | ||
solver_parameters=solver_parameters, | ||
options=options) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is it worth us making the
predictor
andincrement
forms clear in the docstrings?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I've made it a bit more clear, describing what we are solving for