From 90d1f7c3d098d693bef891bb4f0fbdbca88a260c Mon Sep 17 00:00:00 2001 From: atb1995 Date: Thu, 24 Oct 2024 11:00:27 +0100 Subject: [PATCH 01/10] Increment form for implicit RK added and tested --- .../explicit_runge_kutta.py | 6 +- .../implicit_runge_kutta.py | 161 ++++++++++++++---- .../model/test_time_discretisation.py | 9 +- 3 files changed, 134 insertions(+), 42 deletions(-) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 64e9545e0..834b0f569 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -96,13 +96,9 @@ def __init__(self, domain, butcher_matrix, field_name=None, solver_parameters=solver_parameters, limiter=limiter, options=options) self.butcher_matrix = butcher_matrix - self.nbutcher = int(np.shape(self.butcher_matrix)[0]) + self.nStages = int(np.shape(self.butcher_matrix)[0]) self.increment_form = increment_form - @property - def nStages(self): - return self.nbutcher - def setup(self, equation, apply_bcs=True, *active_labels): """ Set up the time discretisation based on the equation. diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 18b2368bb..e4bbd8bf7 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -3,7 +3,7 @@ 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 @@ -56,7 +56,8 @@ class ImplicitRungeKutta(TimeDiscretisation): # --------------------------------------------------------------------------- def __init__(self, domain, butcher_matrix, field_name=None, - solver_parameters=None, limiter=None, options=None,): + increment_form=True, solver_parameters=None, + limiter=None, options=None,): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -66,6 +67,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. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -80,6 +84,7 @@ def __init__(self, domain, butcher_matrix, field_name=None, limiter=limiter, options=options) self.butcher_matrix = butcher_matrix self.nStages = int(np.shape(self.butcher_matrix)[1]) + self.increment_form = increment_form def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -94,6 +99,7 @@ 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)] + self.xs = [Function(self.fs) for i in range(self.nStages)] def lhs(self): return super().lhs @@ -101,23 +107,90 @@ def lhs(self): 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.""" + # 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 + # dt*(a_s1*F(y_1) + a_s2*F(y_2)+ ... + a_{s,s-1}*F(y_{s-1})) + print(stage) + 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 + 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 + # 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.increment_form: + 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) + + else: + 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): @@ -128,31 +201,43 @@ 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.limiter is not None: - self.limiter.apply(self.x1) - - 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))]) + if self.increment_form: + for i in range(stage): + self.x1.assign(self.x1 + self.butcher_matrix[stage, i]*self.dt*self.k[i]) + + if self.limiter is not None: + self.limiter.apply(self.x1) + + 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() + + self.k[stage].assign(self.x_out) else: - self.xnph = self.x1 + self.butcher_matrix[stage, stage]*self.dt*self.x_out - solver = self.solvers[stage] - solver.solve() + if (stage > 0): + self.x_out.assign(self.xs[stage-1]) + solver = self.solvers[stage] + solver.solve() - self.k[stage].assign(self.x_out) + 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.increment_form: + 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]) + else: + self.final_solver.solve() + x_out.assign(self.x_out) if self.limiter is not None: self.limiter.apply(x_out) @@ -168,14 +253,17 @@ 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, - limiter=None, options=None): + def __init__(self, domain, field_name=None, increment_form=True, + solver_parameters=None, limiter=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. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -187,6 +275,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, + increment_form=increment_form, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -202,14 +291,17 @@ 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, - limiter=None, options=None): + def __init__(self, domain, field_name=None, increment_form=True, + solver_parameters=None, limiter=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. + increment_form (bool, optional): whether to write the RK scheme in + "increment form", solving for increments rather than updated + fields. Defaults to True. solver_parameters (dict, optional): dictionary of parameters to pass to the underlying solver. Defaults to None. limiter (:class:`Limiter` object, optional): a limiter to apply to @@ -221,5 +313,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, + increment_form=increment_form, solver_parameters=solver_parameters, limiter=limiter, options=options) diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index 5146107b2..ea5cff559 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -10,7 +10,8 @@ def run(timestepper, tmax, f_end): @pytest.mark.parametrize( "scheme", ["ssprk3_increment", "TrapeziumRule", "ImplicitMidpoint", - "QinZhang", "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", + "QinZhang_increment", "QinZhang_predictor", + "RK4", "Heun", "BDF2", "TR_BDF2", "AdamsBashforth", "Leapfrog", "AdamsMoulton", "AdamsMoulton", "ssprk3_predictor"]) def test_time_discretisation(tmpdir, scheme, tracer_setup): if (scheme == "AdamsBashforth"): @@ -35,8 +36,10 @@ def test_time_discretisation(tmpdir, scheme, tracer_setup): transport_scheme = TrapeziumRule(domain) elif scheme == "ImplicitMidpoint": transport_scheme = ImplicitMidpoint(domain) - elif scheme == "QinZhang": - transport_scheme = QinZhang(domain) + elif scheme == "QinZhang_increment": + transport_scheme = QinZhang(domain, increment_form=True) + elif scheme == "QinZhang_predictor": + transport_scheme = QinZhang(domain, increment_form=False) elif scheme == "RK4": transport_scheme = RK4(domain) elif scheme == "Heun": From 8b21447ddc2e605f4ccf6074e739d9c2dbbd9409 Mon Sep 17 00:00:00 2001 From: Alex Brown <81297297+atb1995@users.noreply.github.com> Date: Wed, 6 Nov 2024 12:22:44 +0000 Subject: [PATCH 02/10] Update gusto/time_discretisation/implicit_runge_kutta.py Co-authored-by: Thomas Bendall --- gusto/time_discretisation/implicit_runge_kutta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index ef095e84d..2e70fd555 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -126,7 +126,7 @@ def res(self, stage): 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 + # Loop through stages up to s-1 and calculate sum # 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( From 3414a27a70bfd82c921841bc98a4ad8a2000dd06 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 6 Nov 2024 12:48:06 +0000 Subject: [PATCH 03/10] Changes in response to review --- .../time_discretisation/implicit_runge_kutta.py | 16 ++++++---------- gusto/time_discretisation/time_discretisation.py | 6 +++--- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 2e70fd555..5f2b628e6 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -31,7 +31,9 @@ class ImplicitRungeKutta(TimeDiscretisation): For each i = 1, s in an s stage method we have the intermediate solutions: \n y_i = y^n + dt*(a_i1*k_1 + a_i2*k_2 + ... + a_ii*k_i) \n - We compute the gradient at the intermediate location, k_i = F(y_i) \n + For the increment form we compute the gradient at the \n + intermediate location, k_i = F(y_i), whilst for the \n + predictor form we solve for each intermediate solution y_i. \n At the last stage, compute the new solution by: \n y^{n+1} = y^n + dt*(b_1*k_1 + b_2*k_2 + .... + b_s*k_s) @@ -110,14 +112,8 @@ def setup(self, equation, apply_bcs=True, *active_labels): 'Runge-Kutta formulation is not implemented' ) - def lhs(self): - return super().lhs - - def rhs(self): - return super().rhs - def res(self, stage): - """Set up the discretisation's residual for a given stage.""" + """Set up the residual for the predictor formulation for a given stage.""" # Add time derivative terms y_s - y^n for stage s mass_form = self.residual.label_map( lambda t: t.has_label(time_derivative), @@ -150,7 +146,7 @@ def res(self, stage): @property def final_res(self): - """Set up the discretisation's final residual.""" + """Set up the final residual fpr the predictor formulation.""" # Add time derivative terms y^{n+1} - y^n mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop) @@ -194,7 +190,7 @@ def solver(self, stage): @cached_property def final_solver(self): - """Set up a solver for the final solve to evaluate time level n+1.""" + """Set up a solver for the final solve for the predictor formulation 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__ diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index df108a615..ee3fb1574 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -5,7 +5,7 @@ operator F. """ -from abc import ABCMeta, abstractmethod, abstractproperty +from abc import ABCMeta, abstractmethod import math from firedrake import (Function, TestFunction, TestFunctions, DirichletBC, @@ -261,7 +261,7 @@ def setup(self, equation, apply_bcs=True, *active_labels): def nlevels(self): return 1 - @abstractproperty + @property def lhs(self): """Set up the discretisation's left hand side (the time derivative).""" l = self.residual.label_map( @@ -271,7 +271,7 @@ def lhs(self): return l.form - @abstractproperty + @property def rhs(self): """Set up the time discretisation's right hand side.""" r = self.residual.label_map( From f812c9801dade14c09646bc7fec6a2136605d436 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Fri, 8 Nov 2024 09:47:14 +0000 Subject: [PATCH 04/10] Changes in response to code review --- gusto/time_discretisation/implicit_runge_kutta.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 5f2b628e6..794093e4e 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -146,7 +146,7 @@ def res(self, stage): @property def final_res(self): - """Set up the final residual fpr the predictor formulation.""" + """Set up the final residual for the predictor formulation.""" # Add time derivative terms y^{n+1} - y^n mass_form = self.residual.label_map(lambda t: t.has_label(time_derivative), map_if_false=drop) @@ -190,7 +190,10 @@ def solver(self, stage): @cached_property def final_solver(self): - """Set up a solver for the final solve for the predictor formulation to evaluate time level n+1.""" + """ + Set up a solver for the final solve for the predictor + formulation 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__ From 249dfd3600e7193b1d8aaaa5535de443964fc65f Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 18 Dec 2024 10:01:37 +0000 Subject: [PATCH 05/10] Intermediate stage testing Linear vs Nonlinear solvers, no difference in results or timings --- .../explicit_runge_kutta.py | 90 ++++++++++++++++++- 1 file changed, 89 insertions(+), 1 deletion(-) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 8746db9ea..2d24df97d 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -4,7 +4,8 @@ from enum import Enum from firedrake import (Function, Constant, NonlinearVariationalProblem, - NonlinearVariationalSolver) + NonlinearVariationalSolver, LinearVariationalProblem, + LinearVariationalSolver, TrialFunction) from firedrake.fml import replace_subject, all_terms, drop, keep, Term from firedrake.utils import cached_property from firedrake.formmanipulation import split_form @@ -41,6 +42,7 @@ class RungeKuttaFormulation(Enum): increment = 1595712 predictor = 8234900 + predictor_lin_solve = 1386823 linear = 269207 @@ -141,6 +143,11 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.rk_formulation == RungeKuttaFormulation.predictor: self.field_i = [Function(self.fs) for _ in range(self.nStages+1)] + elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: + self.field_i = [Function(self.fs) for _ in range(self.nStages+1)] + # self.field_i_inc = [Function(self.fs) for _ in range(self.nStages+1)] + self.df_trial = TrialFunction(self.fs) + self.df = Function(self.fs) elif self.rk_formulation == RungeKuttaFormulation.increment: self.k = [Function(self.fs) for _ in range(self.nStages)] elif self.rk_formulation == RungeKuttaFormulation.linear: @@ -171,6 +178,22 @@ def solver(self): ) solver_list.append(solver) return solver_list + elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: + solver_list = [] + + for stage in range(self.nStages): + # setup linear solver using lhs and rhs defined in derived class + problem = LinearVariationalProblem( + self.lhs[stage].form, self.rhs[stage].form, + self.df, bcs=self.bcs + ) + solver_name = self.field_name+self.__class__.__name__+str(stage) + solver = LinearVariationalSolver( + problem, solver_parameters=self.solver_parameters, + options_prefix=solver_name + ) + solver_list.append(solver) + return solver_list elif self.rk_formulation == RungeKuttaFormulation.linear: problem = NonlinearVariationalProblem( @@ -222,6 +245,17 @@ def lhs(self): return lhs_list + elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: + lhs_list = [] + for stage in range(self.nStages): + l = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.df_trial, self.idx), + map_if_false=drop) + lhs_list.append(l) + + return lhs_list + if self.rk_formulation == RungeKuttaFormulation.linear: l = self.residual.label_map( lambda t: t.has_label(time_derivative), @@ -289,6 +323,32 @@ def rhs(self): return rhs_list + elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: + rhs_list = [] + + for stage in range(self.nStages): + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) + + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t) + + for i in range(1, stage+1): + r_i = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) + ) + + r -= self.butcher_matrix[stage, i]*self.dt*r_i + + rhs_list.append(r) + + return rhs_list + elif self.rk_formulation == RungeKuttaFormulation.linear: r = self.residual.label_map( @@ -391,6 +451,34 @@ def solve_stage(self, x0, stage): if self.limiter is not None: self.limiter.apply(self.x1) + elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: + # Set initial field + if stage == 0: + self.field_i[0].assign(x0) + + # Use previous stage value as a first guess (otherwise may not converge) + self.field_i[stage+1].assign(self.field_i[stage]) + + # Update field_i for physics / limiters + for evaluate in self.evaluate_source: + # TODO: not implemented! Here we need to evaluate the m-th term + # in the i-th RHS with field_m + raise NotImplementedError( + 'Physics not implemented with RK schemes that use the ' + + 'predictor form') + if self.limiter is not None: + self.limiter.apply(self.field_i[stage]) + + # Obtain field_ip1 = field_n - dt* sum_m{a_im*F[field_m]} + self.solver[stage].solve() + + self.field_i[stage+1].assign(self.field_i[0] + self.df) + + if (stage == self.nStages - 1): + self.x1.assign(self.field_i[stage+1]) + if self.limiter is not None: + self.limiter.apply(self.x1) + elif self.rk_formulation == RungeKuttaFormulation.linear: # Set combined index of stage and subcycle From ff9f524382ad414fa137cbe8f8d5314bd8d67857 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 18 Dec 2024 10:04:00 +0000 Subject: [PATCH 06/10] Revert "Intermediate stage testing Linear vs Nonlinear solvers, no difference in results or timings" This reverts commit 249dfd3600e7193b1d8aaaa5535de443964fc65f. --- .../explicit_runge_kutta.py | 90 +------------------ 1 file changed, 1 insertion(+), 89 deletions(-) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 2d24df97d..8746db9ea 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -4,8 +4,7 @@ from enum import Enum from firedrake import (Function, Constant, NonlinearVariationalProblem, - NonlinearVariationalSolver, LinearVariationalProblem, - LinearVariationalSolver, TrialFunction) + NonlinearVariationalSolver) from firedrake.fml import replace_subject, all_terms, drop, keep, Term from firedrake.utils import cached_property from firedrake.formmanipulation import split_form @@ -42,7 +41,6 @@ class RungeKuttaFormulation(Enum): increment = 1595712 predictor = 8234900 - predictor_lin_solve = 1386823 linear = 269207 @@ -143,11 +141,6 @@ def setup(self, equation, apply_bcs=True, *active_labels): if self.rk_formulation == RungeKuttaFormulation.predictor: self.field_i = [Function(self.fs) for _ in range(self.nStages+1)] - elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: - self.field_i = [Function(self.fs) for _ in range(self.nStages+1)] - # self.field_i_inc = [Function(self.fs) for _ in range(self.nStages+1)] - self.df_trial = TrialFunction(self.fs) - self.df = Function(self.fs) elif self.rk_formulation == RungeKuttaFormulation.increment: self.k = [Function(self.fs) for _ in range(self.nStages)] elif self.rk_formulation == RungeKuttaFormulation.linear: @@ -178,22 +171,6 @@ def solver(self): ) solver_list.append(solver) return solver_list - elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: - solver_list = [] - - for stage in range(self.nStages): - # setup linear solver using lhs and rhs defined in derived class - problem = LinearVariationalProblem( - self.lhs[stage].form, self.rhs[stage].form, - self.df, bcs=self.bcs - ) - solver_name = self.field_name+self.__class__.__name__+str(stage) - solver = LinearVariationalSolver( - problem, solver_parameters=self.solver_parameters, - options_prefix=solver_name - ) - solver_list.append(solver) - return solver_list elif self.rk_formulation == RungeKuttaFormulation.linear: problem = NonlinearVariationalProblem( @@ -245,17 +222,6 @@ def lhs(self): return lhs_list - elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: - lhs_list = [] - for stage in range(self.nStages): - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.df_trial, self.idx), - map_if_false=drop) - lhs_list.append(l) - - return lhs_list - if self.rk_formulation == RungeKuttaFormulation.linear: l = self.residual.label_map( lambda t: t.has_label(time_derivative), @@ -323,32 +289,6 @@ def rhs(self): return rhs_list - elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: - rhs_list = [] - - for stage in range(self.nStages): - r = self.residual.label_map( - all_terms, - map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) - - r = r.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t) - - for i in range(1, stage+1): - r_i = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) - ) - - r -= self.butcher_matrix[stage, i]*self.dt*r_i - - rhs_list.append(r) - - return rhs_list - elif self.rk_formulation == RungeKuttaFormulation.linear: r = self.residual.label_map( @@ -451,34 +391,6 @@ def solve_stage(self, x0, stage): if self.limiter is not None: self.limiter.apply(self.x1) - elif self.rk_formulation == RungeKuttaFormulation.predictor_lin_solve: - # Set initial field - if stage == 0: - self.field_i[0].assign(x0) - - # Use previous stage value as a first guess (otherwise may not converge) - self.field_i[stage+1].assign(self.field_i[stage]) - - # Update field_i for physics / limiters - for evaluate in self.evaluate_source: - # TODO: not implemented! Here we need to evaluate the m-th term - # in the i-th RHS with field_m - raise NotImplementedError( - 'Physics not implemented with RK schemes that use the ' - + 'predictor form') - if self.limiter is not None: - self.limiter.apply(self.field_i[stage]) - - # Obtain field_ip1 = field_n - dt* sum_m{a_im*F[field_m]} - self.solver[stage].solve() - - self.field_i[stage+1].assign(self.field_i[0] + self.df) - - if (stage == self.nStages - 1): - self.x1.assign(self.field_i[stage+1]) - if self.limiter is not None: - self.limiter.apply(self.x1) - elif self.rk_formulation == RungeKuttaFormulation.linear: # Set combined index of stage and subcycle From eeafdd9b31ba535f02db48bd7b22cf4ff004b09a Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 18 Dec 2024 11:25:40 +0000 Subject: [PATCH 07/10] Removing lhs and rhs and adding in res --- .../explicit_runge_kutta.py | 81 ++--- .../multi_level_schemes.py | 293 ++++++++++-------- .../time_discretisation.py | 194 ++++++------ 3 files changed, 297 insertions(+), 271 deletions(-) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index 8746db9ea..3f738c4af 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -161,7 +161,7 @@ def solver(self): for stage in range(self.nStages): # setup linear solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem( - self.lhs[stage].form - self.rhs[stage].form, + self.res[stage].form, self.field_i[stage+1], bcs=self.bcs ) solver_name = self.field_name+self.__class__.__name__+str(stage) @@ -174,7 +174,7 @@ def solver(self): elif self.rk_formulation == RungeKuttaFormulation.linear: problem = NonlinearVariationalProblem( - self.lhs - self.rhs[0], self.x1, bcs=self.bcs + self.res[0], self.x1, bcs=self.bcs ) solver_name = self.field_name+self.__class__.__name__ solver = NonlinearVariationalSolver( @@ -184,7 +184,7 @@ def solver(self): # Set up problem for final step problem_last = NonlinearVariationalProblem( - self.lhs - self.rhs[1], self.x1, bcs=self.bcs + self.res[1], self.x1, bcs=self.bcs ) solver_name = self.field_name+self.__class__.__name__+'_last' solver_last = NonlinearVariationalSolver( @@ -200,54 +200,21 @@ def solver(self): ) @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" + def res(self): + """Set up the discretisation's residual.""" if self.rk_formulation == RungeKuttaFormulation.increment: - l = self.residual.label_map( + residual = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x_out, self.idx), map_if_false=drop) - - return l.form - - elif self.rk_formulation == RungeKuttaFormulation.predictor: - lhs_list = [] - for stage in range(self.nStages): - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.field_i[stage+1], self.idx), - map_if_false=drop) - lhs_list.append(l) - - return lhs_list - - if self.rk_formulation == RungeKuttaFormulation.linear: - l = self.residual.label_map( - lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x1, self.idx), - map_if_false=drop) - - return l.form - - else: - raise NotImplementedError( - 'Runge-Kutta formulation is not implemented' - ) - - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side.""" - - if self.rk_formulation == RungeKuttaFormulation.increment: r = self.residual.label_map( all_terms, map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map( + residual += r.label_map( lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=lambda t: -1*t) + map_if_true=drop) # If there are no active labels, we may have no terms at this point # So that we can still do xnp1 = xn, put in a zero term here @@ -259,19 +226,22 @@ def rhs(self): # Drop label from this map_if_true=lambda t: time_derivative.remove(t), map_if_false=drop) - r += null_term + residual += null_term - return r.form + return residual.form elif self.rk_formulation == RungeKuttaFormulation.predictor: - rhs_list = [] - + residual_list = [] for stage in range(self.nStages): + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.field_i[stage+1], self.idx), + map_if_false=drop) r = self.residual.label_map( all_terms, map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) - r = r.label_map( + residual -= r.label_map( lambda t: t.has_label(time_derivative), map_if_true=keep, map_if_false=lambda t: -self.butcher_matrix[stage, 0]*self.dt*t) @@ -283,14 +253,16 @@ def rhs(self): map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) ) - r -= self.butcher_matrix[stage, i]*self.dt*r_i - - rhs_list.append(r) + residual += self.butcher_matrix[stage, i]*self.dt*r_i + residual_list.append(residual) - return rhs_list - - elif self.rk_formulation == RungeKuttaFormulation.linear: + return residual_list + if self.rk_formulation == RungeKuttaFormulation.linear: + time_term = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x1, self.idx), + map_if_false=drop) r = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x0, old_idx=self.idx), @@ -329,8 +301,9 @@ def rhs(self): map_if_true=keep, map_if_false=lambda t: -self.dt*t ) - - return r_all_but_last.form, r.form + res = time_term - r + res_all_but_last = time_term - r_all_but_last + return res_all_but_last.form, res.form else: raise NotImplementedError( diff --git a/gusto/time_discretisation/multi_level_schemes.py b/gusto/time_discretisation/multi_level_schemes.py index a11671cc4..2227c4bb6 100644 --- a/gusto/time_discretisation/multi_level_schemes.py +++ b/gusto/time_discretisation/multi_level_schemes.py @@ -63,58 +63,57 @@ def nlevels(self): return 2 @property - def lhs0(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res0(self): + """Set up the discretisation's residual for initial BDF step.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt * t + ) - @property - def rhs0(self): - """Set up the time discretisation's right hand side for inital BDF step.""" r = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) + residual -= r - return r.form + return residual.form @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: (2/3)*self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: (2 / 3) * self.dt * t + ) - @property - def rhs(self): - """Set up the time discretisation's right hand side for BDF2 steps.""" xn = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) xnm1 = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.xnm1, old_idx=self.idx), - map_if_false=drop) - - r = (4/3.) * xn - (1/3.) * xnm1 + map_if_false=drop + ) - return r.form + residual -= (4 / 3) * xn - (1 / 3) * xnm1 + return residual.form @property def solver0(self): """Set up the problem and the solver for initial BDF step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -122,8 +121,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for BDF2 steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.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) @@ -162,38 +161,53 @@ def nlevels(self): return 2 @property - def rhs0(self): - """Set up the discretisation's right hand side for initial forward euler step.""" + def res0(self): + """Set up the discretisation's residual for initial forward euler step.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) - - return r.form + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * t + ) - @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(Leapfrog, self).lhs + residual -= r + return residual.form @property def rhs(self): - """Set up the discretisation's right hand side for leapfrog steps.""" + """Set up the discretisation's residual for leapfrog steps.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + r = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_false=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.xnm1, old_idx=self.idx), - map_if_false=lambda t: -2.0*self.dt*t) + map_if_false=replace_subject(self.x1, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.xnm1, old_idx=self.idx), + map_if_false=lambda t: -2.0 * self.dt * t + ) - return r.form + residual -= r + return residual.form @property def solver0(self): """Set up the problem and the solver for initial forward euler step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -201,8 +215,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for leapfrog steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.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) @@ -291,42 +305,59 @@ def nlevels(self): return self.order @property - def rhs0(self): - """Set up the discretisation's right hand side for initial forward euler step.""" + def res0(self): + """Set up the discretisation's residual for initial forward euler step.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * t + ) - return r.form + residual -= r + return residual.form @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(AdamsBashforth, self).lhs - - @property - def rhs(self): - """Set up the discretisation's right hand side for Adams Bashforth steps.""" - r = self.residual.label_map(all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.b[-1]*self.dt*t) - for n in range(self.nlevels-1): - rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.x[n], old_idx=self.idx)) - rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*self.b[n]*t) - r += rtemp - return r.form + def res(self): + """Set up the discretisation's residual for Adams Bashforth steps.""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + residual -= r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.b[-1] * self.dt * t + ) + for n in range(self.nlevels - 1): + rtemp = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.x[n], old_idx=self.idx) + ) + rtemp = rtemp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * self.b[n] * t + ) + residual -= rtemp + return residual.form @property def solver0(self): """Set up the problem and the solverfor initial forward euler step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -334,8 +365,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for Adams Bashforth steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.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) @@ -427,63 +458,67 @@ def nlevels(self): return self.order @property - def rhs0(self): + def res0(self): """ - Set up the discretisation's right hand side for initial trapezoidal - step. + Set up the discretisation's residual for initial trapezoidal step. """ + residual = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: 0.5 * self.dt * t + ) r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -0.5*self.dt*t) + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -0.5 * self.dt * t + ) - return r.form + residual -= r + return residual.form @property - def lhs0(self): - """ - Set up the time discretisation's right hand side for initial - trapezoidal step. - """ - l = self.residual.label_map( + def res(self): + """Set up the time discretisation's residual for Adams Moulton steps.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: 0.5*self.dt*t) - return l.form - - @property - def lhs(self): - """Set up the time discretisation's right hand side for Adams Moulton steps.""" - l = self.residual.label_map( + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.bl * self.dt * t + ) + r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.bl*self.dt*t) - return l.form - - @property - def rhs(self): - """Set up the discretisation's right hand side for Adams Moulton steps.""" - r = self.residual.label_map(all_terms, - map_if_true=replace_subject(self.x[-1], old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.br[-1]*self.dt*t) - for n in range(self.nlevels-1): - rtemp = self.residual.label_map(lambda t: t.has_label(time_derivative), - map_if_true=drop, - map_if_false=replace_subject(self.x[n], old_idx=self.idx)) - rtemp = rtemp.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*self.br[n]*t) - r += rtemp - return r.form + map_if_true=replace_subject(self.x[-1], old_idx=self.idx) + ) + residual -= r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.br[-1] * self.dt * t + ) + for n in range(self.nlevels - 1): + rtemp = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=drop, + map_if_false=replace_subject(self.x[n], old_idx=self.idx) + ) + rtemp = rtemp.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -self.dt * self.br[n] * t + ) + residual -= rtemp + return residual.form @property def solver0(self): """Set up the problem and the solver for initial trapezoidal step.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs0-self.rhs0, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res0, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"0" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -491,8 +526,8 @@ def solver0(self): @property def solver(self): """Set up the problem and the solver for Adams Moulton steps.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.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) diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index 79627ff76..3796cc11c 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -325,33 +325,32 @@ def nlevels(self): return 1 @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x_out, old_idx=self.idx), - map_if_false=drop) - - return l.form - - @property - def rhs(self): - """Set up the time discretisation's right hand side.""" + map_if_false=drop + ) r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) r = r.label_map( lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -self.dt*t) + map_if_false=lambda t: -self.dt * t + ) + + residual -= r - return r.form + return residual.form @cached_property def solver(self): """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ solver = NonlinearVariationalSolver( problem, @@ -461,20 +460,33 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.solver_parameters['snes_type'] = 'newtonls' @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual""" + residual = self.residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_false=drop + ) + + r = self.residual.label_map( + all_terms, + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) + + r = r.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, self.idx), - map_if_false=drop) + map_if_false=lambda t: -self.dt * t + ) - return l.form + residual -= r + + return residual.form @cached_property def solver(self): """Set up the problem and the solver.""" - # setup linear solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs - self.rhs, self.x_out, bcs=self.bcs) + # setup linear solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.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) @@ -547,25 +559,27 @@ def __init__(self, domain, field_name=None, solver_parameters=None, limiter=limiter, options=options) @property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.dt*t) - - return l.form - - @property - def rhs(self): - """Set up the time discretisation's right hand side.""" + map_if_true=replace_subject( + self.x_out, old_idx=self.idx + ) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.dt*t + ) r = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.x1, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) + + residual -= r - return r.form + return residual.form @wrapper_apply def apply(self, x_out, x_in): @@ -639,26 +653,27 @@ def __init__(self, domain, theta, field_name=None, self.theta = theta @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: self.theta*self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: self.theta * self.dt * t + ) - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side.""" r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -(1-self.theta)*self.dt*t) - - return r.form + map_if_true=replace_subject(self.x1, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -(1 - self.theta) * self.dt * t + ) + residual -= r + return residual.form @wrapper_apply def apply(self, x_out, x_in): @@ -756,60 +771,63 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.xn = Function(self.fs) @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative) for the TR stage.""" - l = self.residual.label_map( + def res(self): + """Set up the discretisation's residual for the TR stage.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.xnpg, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: 0.5*self.gamma*self.dt*t) - - return l.form + map_if_true=replace_subject(self.xnpg, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: 0.5 * self.gamma * self.dt * t + ) - @cached_property - def rhs(self): - """Set up the time discretisation's right hand side for the TR stage.""" r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.xn, old_idx=self.idx)) - r = r.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: -0.5*self.gamma*self.dt*t) + map_if_true=replace_subject(self.xn, old_idx=self.idx) + ) + r = r.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: -0.5 * self.gamma * self.dt * t + ) + residual -= r - return r.form + return residual.form @cached_property - def lhs_bdf2(self): - """Set up the discretisation's left hand side (the time derivative) for - the BDF2 stage.""" - l = self.residual.label_map( + def res_bdf2(self): + """Set up the discretisation's residual for the BDF2 stage.""" + residual = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x_out, old_idx=self.idx)) - l = l.label_map(lambda t: t.has_label(time_derivative), - map_if_false=lambda t: ((1.0-self.gamma)/(2.0-self.gamma))*self.dt*t) - - return l.form + map_if_true=replace_subject(self.x_out, old_idx=self.idx) + ) + residual = residual.label_map( + lambda t: t.has_label(time_derivative), + map_if_false=lambda t: (1.0 - self.gamma) / (2.0 - self.gamma) * self.dt * t + ) - @cached_property - def rhs_bdf2(self): - """Set up the time discretisation's right hand side for the BDF2 stage.""" xn = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.xn, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) xnpg = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.xnpg, old_idx=self.idx), - map_if_false=drop) + map_if_false=drop + ) - r = (1.0/(self.gamma*(2.0-self.gamma)))*xnpg - ((1.0-self.gamma)**2/(self.gamma*(2.0-self.gamma)))*xn + r = (1.0 / (self.gamma * (2.0 - self.gamma))) * xnpg - \ + ((1.0 - self.gamma) ** 2 / (self.gamma * (2.0 - self.gamma))) * xn - return r.form + residual -= r + return residual.form @cached_property def solver_tr(self): """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.xnpg, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res, self.xnpg, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"_tr" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @@ -817,8 +835,8 @@ def solver_tr(self): @cached_property def solver_bdf2(self): """Set up the problem and the solver.""" - # setup solver using lhs and rhs defined in derived class - problem = NonlinearVariationalProblem(self.lhs_bdf2-self.rhs_bdf2, self.x_out, bcs=self.bcs) + # setup solver using residual (res) defined in derived class + problem = NonlinearVariationalProblem(self.res_bdf2, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__+"_bdf2" return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) From 9b560ebd6d7551698f1cb75fb1f13144f3931ae5 Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 18 Dec 2024 13:28:56 +0000 Subject: [PATCH 08/10] Added augmentation to IMEX & Implicit RK, as well as multilevel --- gusto/time_discretisation/imex_runge_kutta.py | 42 +++++++++++++------ .../implicit_runge_kutta.py | 13 +++++- .../multi_level_schemes.py | 23 +++++++--- 3 files changed, 58 insertions(+), 20 deletions(-) diff --git a/gusto/time_discretisation/imex_runge_kutta.py b/gusto/time_discretisation/imex_runge_kutta.py index fcaefe1ec..f97851199 100644 --- a/gusto/time_discretisation/imex_runge_kutta.py +++ b/gusto/time_discretisation/imex_runge_kutta.py @@ -61,7 +61,7 @@ class IMEXRungeKutta(TimeDiscretisation): def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -82,10 +82,13 @@ def __init__(self, domain, butcher_imp, butcher_exp, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, field_name=field_name, solver_parameters=nonlinear_solver_parameters, - options=options) + options=options, augmentation=augmentation) self.butcher_imp = butcher_imp self.butcher_exp = butcher_exp self.nStages = int(np.shape(self.butcher_imp)[1]) @@ -269,7 +272,7 @@ class IMEX_Euler(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -286,13 +289,16 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ butcher_imp = np.array([[0., 0.], [0., 1.], [0., 1.]]) butcher_exp = np.array([[0., 0.], [1., 0.], [1., 0.]]) super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_ARS3(IMEXRungeKutta): @@ -313,7 +319,7 @@ class IMEX_ARS3(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -330,6 +336,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ g = (3. + np.sqrt(3.))/6. butcher_imp = np.array([[0., 0., 0.], [0., g, 0.], [0., 1-2.*g, g], [0., 0.5, 0.5]]) @@ -338,7 +347,7 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_ARK2(IMEXRungeKutta): @@ -359,7 +368,7 @@ class IMEX_ARK2(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -376,6 +385,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ g = 1. - 1./np.sqrt(2.) d = 1./(2.*np.sqrt(2.)) @@ -385,7 +397,7 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_SSP3(IMEXRungeKutta): @@ -404,7 +416,7 @@ class IMEX_SSP3(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -421,6 +433,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ g = 1. - (1./np.sqrt(2.)) butcher_imp = np.array([[g, 0., 0.], [1-2.*g, g, 0.], [0.5-g, 0., g], [(1./6.), (1./6.), (2./3.)]]) @@ -428,7 +443,7 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) class IMEX_Trap2(IMEXRungeKutta): @@ -447,7 +462,7 @@ class IMEX_Trap2(IMEXRungeKutta): """ def __init__(self, domain, field_name=None, linear_solver_parameters=None, nonlinear_solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -464,6 +479,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ e = 0. butcher_imp = np.array([[0., 0., 0., 0.], [e, 0., 0., 0.], [0.5, 0., 0.5, 0.], [0.5, 0., 0., 0.5], [0.5, 0., 0., 0.5]]) @@ -471,4 +489,4 @@ def __init__(self, domain, field_name=None, super().__init__(domain, butcher_imp, butcher_exp, field_name, linear_solver_parameters=linear_solver_parameters, nonlinear_solver_parameters=nonlinear_solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, augmentation=augmentation) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 700ced5fa..238501152 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -79,6 +79,9 @@ def __init__(self, domain, butcher_matrix, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ super().__init__(domain, field_name=field_name, solver_parameters=solver_parameters, @@ -280,6 +283,9 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ butcher_matrix = np.array([[0.5], [1.]]) super().__init__(domain, butcher_matrix, field_name, @@ -301,7 +307,7 @@ class QinZhang(ImplicitRungeKutta): """ def __init__(self, domain, field_name=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, options=None): + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -317,9 +323,12 @@ def __init__(self, domain, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to 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) + options=options, augmentation=augmentation) diff --git a/gusto/time_discretisation/multi_level_schemes.py b/gusto/time_discretisation/multi_level_schemes.py index 2227c4bb6..fd3936a61 100644 --- a/gusto/time_discretisation/multi_level_schemes.py +++ b/gusto/time_discretisation/multi_level_schemes.py @@ -17,7 +17,7 @@ class MultilevelTimeDiscretisation(TimeDiscretisation): """Base class for multi-level timesteppers""" def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -32,12 +32,16 @@ def __init__(self, domain, field_name=None, solver_parameters=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") super().__init__(domain=domain, field_name=field_name, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) self.initial_timesteps = 0 @abstractproperty @@ -252,7 +256,7 @@ class AdamsBashforth(MultilevelTimeDiscretisation): y^(n+1) = y^n + dt*(b_0*F[y^(n)] + b_1*F[y^(n-1)] + b_2*F[y^(n-2)] + b_3*F[y^(n-3)] + b_4*F[y^(n-4)]) """ def __init__(self, domain, order, field_name=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -266,6 +270,9 @@ def __init__(self, domain, order, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. Raises: ValueError: if order is not provided, or is in incorrect range. @@ -278,7 +285,7 @@ def __init__(self, domain, order, field_name=None, super().__init__(domain, field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) self.order = order @@ -402,7 +409,8 @@ class AdamsMoulton(MultilevelTimeDiscretisation): y^(n+1) = y^n + dt*(b_0*F[y^(n+1)] + b_1*F[y^(n)] + b_2*F[y^(n-1)] + b_3*F[y^(n-2)]) """ def __init__(self, domain, order, field_name=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -416,6 +424,9 @@ def __init__(self, domain, order, field_name=None, options to either be passed to the spatial discretisation, or to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. + augmentation (:class:`Augmentation`): allows the equation solved in + this time discretisation to be augmented, for instances with + extra terms of another auxiliary variable. Defaults to None. Raises: ValueError: if order is not provided, or is in incorrect range. @@ -431,7 +442,7 @@ def __init__(self, domain, order, field_name=None, super().__init__(domain, field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) self.order = order From 81ba663026f0f127a57a90cd91286440bc2a0f6a Mon Sep 17 00:00:00 2001 From: Alex Brown <81297297+atb1995@users.noreply.github.com> Date: Wed, 18 Dec 2024 14:03:22 +0000 Subject: [PATCH 09/10] Update gusto/time_discretisation/implicit_runge_kutta.py Co-authored-by: Thomas Bendall --- gusto/time_discretisation/implicit_runge_kutta.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 238501152..e494868da 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -157,7 +157,7 @@ def final_res(self): 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 + # Loop through stages up to s-1 and calculate/sum # 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( From 9fa02b5b807f0a3bae08c369ed5ec71e5722678d Mon Sep 17 00:00:00 2001 From: atb1995 Date: Wed, 18 Dec 2024 14:06:31 +0000 Subject: [PATCH 10/10] Removal of lhs & rhs + updated comments --- gusto/time_discretisation/imex_runge_kutta.py | 12 +----------- gusto/time_discretisation/implicit_runge_kutta.py | 2 +- 2 files changed, 2 insertions(+), 12 deletions(-) diff --git a/gusto/time_discretisation/imex_runge_kutta.py b/gusto/time_discretisation/imex_runge_kutta.py index f97851199..c76a7f04c 100644 --- a/gusto/time_discretisation/imex_runge_kutta.py +++ b/gusto/time_discretisation/imex_runge_kutta.py @@ -130,16 +130,6 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.xs = [Function(self.fs) for i in range(self.nStages)] - @cached_property - def lhs(self): - """Set up the discretisation's left hand side (the time derivative).""" - return super(IMEXRungeKutta, self).lhs - - @cached_property - def rhs(self): - """Set up the discretisation's right hand side (the time derivative).""" - return super(IMEXRungeKutta, self).rhs - def res(self, stage): """Set up the discretisation's residual for a given stage.""" # Add time derivative terms y_s - y^n for stage s @@ -229,7 +219,7 @@ def solvers(self): @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 + # setup solver using residual (res) 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.linear_solver_parameters, options_prefix=solver_name) diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index e494868da..bcce30a92 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -197,7 +197,7 @@ def final_solver(self): Set up a solver for the final solve for the predictor formulation to evaluate time level n+1. """ - # setup solver using lhs and rhs defined in derived class + # setup solver using residual (res) 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)