From cb564e59a49d26ef57dd8ca2f2bc01d47b000ccd Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Thu, 18 Jun 2020 17:20:19 -0500 Subject: [PATCH 01/14] Initial stab at implicit Adams method generation --- leap/multistep/__init__.py | 277 +++++++++++++++++++++++++++++++++++++ test/test_implicit_ab.py | 153 ++++++++++++++++++++ 2 files changed, 430 insertions(+) create mode 100755 test/test_implicit_ab.py diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index e357a47..2736326 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -405,4 +405,281 @@ class AdamsBashforthMethodBuilder(MethodBuilder): # }}} +# {{{ am method + + +class AdamsMoultonMethodBuilder(MethodBuilder): + """ + User-supplied context: + + component_id: The value that is integrated + + component_id: The right hand side + + .. automethod:: __init__ + .. automethod:: generate + """ + + def __init__(self, component_id, function_family=None, state_filter_name=None, + hist_length=None, static_dt=False, order=None): + """ + :arg function_family: Accepts an instance of + :class:`ABIntegrationFunctionFamily` + or an integer, in which case the classical monomial function family + with the order given by the integer is used. + :arg static_dt: If *True*, changing the timestep during time integration + is not allowed. + """ + + if function_family is not None and order is not None: + raise ValueError("may not specify both function_family and order") + + if function_family is None: + function_family = order + del order + + if isinstance(function_family, int): + function_family = ABMonomialIntegrationFunctionFamily(function_family) + + super(AdamsMoultonMethodBuilder, self).__init__() + self.function_family = function_family + + if hist_length is None: + hist_length = len(function_family) + + self.hist_length = hist_length + self.static_dt = static_dt + + self.component_id = component_id + + # Declare variables + self.step = var('

step') + self.function = var('' + component_id) + # One shorter for implicit. + self.history = \ + [var('

f_n_minus_' + str(i)) for i in range(hist_length - 2, 0, -1)] + + if not self.static_dt: + # One shorter for implicit. + self.time_history = [ + var('

t_n_minus_' + str(i)) + for i in range(hist_length - 2, 0, -1)] + + self.state = var('' + component_id) + self.t = var('') + self.dt = var('

') + + if state_filter_name is not None: + self.state_filter = var("" + state_filter_name) + else: + self.state_filter = None + + def generate(self): + """ + :returns: :class:`dagrt.language.DAGCode` + """ + from pytools import UniqueNameGenerator + name_gen = UniqueNameGenerator() + + from dagrt.language import DAGCode, CodeBuilder + + array = var("array") + rhs_var = var("rhs_var") + rhs_next_var = var("rhs_next_var") + + # Initialization + with CodeBuilder(name="initialization") as cb_init: + cb_init(self.step, 1) + + # Primary + with CodeBuilder(name="primary") as cb_primary: + + rhs_var_to_unknown = {} + unkvar = cb_primary.fresh_var('unk') + rhs_var_to_unknown[rhs_next_var] = unkvar + + # In implicit mode, the time history must also + # include the next point in time. + if not self.static_dt: + time_history_data = self.time_history + [self.t] + [self.t + self.dt] + time_hist_var = var(name_gen("time_history")) + cb_primary(time_hist_var, array(self.hist_length)) + for i in range(self.hist_length): + cb_primary(time_hist_var[i], time_history_data[i] - self.t) + + time_hist = time_hist_var + t_end = self.dt + dt_factor = 1 + + else: + time_hist = list(range(-self.hist_length+2, 0+2)) # noqa pylint:disable=invalid-unary-operand-type + dt_factor = self.dt + t_end = 1 + + # Implicit setup - rhs_next_var is an unknown, needs implicit solve. + equations = [] + unknowns = set() + knowns = set() + + def make_known(v): + unknowns.discard(v) + knowns.add(v) + + unknowns.add(rhs_next_var) + + cb_primary(rhs_var, self.eval_rhs(self.t, self.state)) + + # Update history + history = self.history + [rhs_var] + [rhs_next_var] + + # Set up the actual Adams-Moulton step. + ab_sum = emit_ab_integration( + cb_primary, name_gen, + self.function_family, + time_hist, history, + 0, t_end) + + state_est = self.state + dt_factor * ab_sum + if self.state_filter is not None: + state_est = self.state_filter(state_est) + + # Build the implicit solve expression. + from dagrt.expression import collapse_constants + solve_expression = collapse_constants( + rhs_next_var - self.eval_rhs(self.t + self.dt, state_est), + list(unknowns) + [self.state], + cb_primary.assign, cb_primary.fresh_var) + equations.append(solve_expression) + + # {{{ emit solve if possible + + if unknowns and len(unknowns) == len(equations): + # got a square system, let's solve + assignees = [unk.name for unk in unknowns] + + from pymbolic import substitute + subst_dict = dict( + (rhs_var.name, rhs_var_to_unknown[rhs_var]) + for rhs_var in unknowns) + + cb_primary.assign_implicit( + assignees=assignees, + solve_components=[ + rhs_var_to_unknown[unk].name + for unk in unknowns], + expressions=[ + substitute(eq, subst_dict) + for eq in equations], + + # TODO: Could supply a starting guess + other_params={ + "guess": self.state}, + solver_id="solve") + + del equations[:] + knowns.update(unknowns) + unknowns.clear() + + # }}} + + # Update the state now that we've solved. + cb_primary(self.state, state_est) + + # Rotate history and time history. + # In the case of implicit, history length is one less. + for i in range(self.hist_length - 2): + cb_primary(self.history[i], history[i + 1]) + + if not self.static_dt: + cb_primary(self.time_history[i], time_history_data[i + 1]) + + cb_primary(self.t, self.t + self.dt) + cb_primary.yield_state(expression=self.state, + component_id=self.component_id, + time_id='', time=self.t) + + if self.hist_length == 1: + # The first order method requires no bootstrapping. + return DAGCode( + phases={ + "initial": cb_init.as_execution_phase(next_phase="primary"), + "primary": cb_primary.as_execution_phase(next_phase="primary") + }, + initial_phase="initial") + + # Bootstrap + with CodeBuilder(name="bootstrap") as cb_bootstrap: + self.rk_bootstrap(cb_bootstrap) + cb_bootstrap(self.t, self.t + self.dt) + cb_bootstrap.yield_state(expression=self.state, + component_id=self.component_id, + time_id='', time=self.t) + cb_bootstrap(self.step, self.step + 1) + # Bootstrap length is one less because of implicit. + with cb_bootstrap.if_(self.step, "==", self.hist_length - 1): + cb_bootstrap.switch_phase("primary") + + return DAGCode( + phases={ + "initialization": cb_init.as_execution_phase("bootstrap"), + "bootstrap": cb_bootstrap.as_execution_phase("bootstrap"), + "primary": cb_primary.as_execution_phase("primary"), + }, + initial_phase="initialization") + + def eval_rhs(self, t, y): + """Return a node that evaluates the RHS at the given time and + component value.""" + from pymbolic.primitives import CallWithKwargs + return CallWithKwargs(function=self.function, + parameters=(), + kw_parameters={"t": t, self.component_id: y}) + + def rk_bootstrap(self, cb): + """Initialize the timestepper with an RK method.""" + + rhs_var = var("rhs_var") + + cb(rhs_var, self.eval_rhs(self.t, self.state)) + + # Save the current RHS to the AB history + + for i in range(len(self.history)): + with cb.if_(self.step, "==", i + 1): + cb(self.history[i], rhs_var) + + if not self.static_dt: + cb(self.time_history[i], self.t) + + from leap.rk import ORDER_TO_RK_METHOD_BUILDER + rk_method = ORDER_TO_RK_METHOD_BUILDER[self.function_family.order] + rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit)) + rk_coeffs = rk_method.output_coeffs + + # Stage loop (taken from EmbeddedButcherTableauMethodBuilder) + rhss = [var("rk_rhs_" + str(i)) for i in range(len(rk_tableau))] + for stage_num, (c, coeffs) in enumerate(rk_tableau): + if len(coeffs) == 0: + assert c == 0 + cb(rhss[stage_num], rhs_var) + else: + stage = self.state + sum(self.dt * coeff * rhss[j] + for (j, coeff) + in enumerate(coeffs)) + + if self.state_filter is not None: + stage = self.state_filter(stage) + + cb(rhss[stage_num], self.eval_rhs(self.t + c * self.dt, stage)) + + # Merge the values of the RHSs. + rk_comb = sum(coeff * rhss[j] for j, coeff in enumerate(rk_coeffs)) + + state_est = self.state + self.dt * rk_comb + if self.state_filter is not None: + state_est = self.state_filter(state_est) + + # Assign the value of the new state. + cb(self.state, state_est) + +# }}} + # vim: fdm=marker diff --git a/test/test_implicit_ab.py b/test/test_implicit_ab.py new file mode 100755 index 0000000..6681fdf --- /dev/null +++ b/test/test_implicit_ab.py @@ -0,0 +1,153 @@ +#! /usr/bin/env python +from __future__ import division, with_statement + +__copyright__ = "Copyright (C) 2014 Matt Wala" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +import numpy as np +import pytest +import sys + +from leap.multistep import AdamsMoultonMethodBuilder + +from utils import ( # noqa + python_method_impl_interpreter as pmi_int, + python_method_impl_codegen as pmi_cg) + + +def am_solver(f, t, sub_y, coeff, guess): + from scipy.optimize import root + return root(lambda unk: unk - f(t=t, y=sub_y + coeff*unk), guess).x + + +def am_solver_hook(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match("unk - rhs(t=t, y=sub_y + coeff*unk)", solve_expr, + pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, sub_y, coeff, guess)", pieces) + + +def am_solver_static(f, t, sub_y, coeff, dt, add, guess): + from scipy.optimize import root + return root(lambda unk: unk - f(t=t, y=sub_y + dt*(add + coeff*unk)), guess).x + + +def am_solver_hook_static(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match("unk - rhs(t=t, y=sub_y + dt*(add + coeff*unk))", + solve_expr, + pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, sub_y, coeff, dt, add, guess)", pieces) + + +@pytest.mark.parametrize(("method", "expected_order", "static_dt"), [ + (AdamsMoultonMethodBuilder("y", order, static_dt=static_dt), order, static_dt) + for order in [2, 3, 4, 5] + for static_dt in [True, False] + ] + [ + (AdamsMoultonMethodBuilder("y", order, hist_length=order+1, + static_dt=static_dt), order, static_dt) + for order in [2, 3, 4, 5] + for static_dt in [True, False] + ]) +def test_am_convergence(python_method_impl, method, expected_order, static_dt): + plot_solution = False + pytest.importorskip("scipy") + from utils import DefaultProblem + problem = DefaultProblem() + + component_id = method.component_id + code = method.generate() + + from leap.implicit import replace_AssignImplicit + if static_dt: + code = replace_AssignImplicit(code, {"solve": am_solver_hook_static}) + else: + code = replace_AssignImplicit(code, {"solve": am_solver_hook}) + + from pytools.convergence import EOCRecorder + eocrec = EOCRecorder() + + # default dts + dts = 2 ** -np.array(range(4, 7), dtype=np.float64) # noqa pylint:disable=invalid-unary-operand-type + + for dt in dts: + t = problem.t_start + y = problem.initial() + final_t = problem.t_end + + from functools import partial + # FIXME + if static_dt: + interp = python_method_impl(code, function_map={ + "" + component_id: problem, + "solver": partial(am_solver_static, problem), + }) + else: + interp = python_method_impl(code, function_map={ + "" + component_id: problem, + "solver": partial(am_solver, problem), + }) + interp.set_up(t_start=t, dt_start=dt, context={component_id: y}) + + times = [] + values = [] + for event in interp.run(t_end=final_t): + if isinstance(event, interp.StateComputed): + assert event.component_id == component_id + values.append(event.state_component[0]) + times.append(event.t) + + assert abs(times[-1] - final_t) / final_t < 0.1 + + times = np.array(times) + + if plot_solution: + import matplotlib.pyplot as pt + pt.plot(times, values, label="comp") + pt.plot(times, problem.exact(times), label="true") + pt.show() + + error = abs(values[-1] - problem.exact(final_t)) + eocrec.add_data_point(dt, error) + + print("------------------------------------------------------") + print("%s: expected order %d" % (method.__class__.__name__, + expected_order)) + print("------------------------------------------------------") + print(eocrec.pretty_print()) + + orderest = eocrec.estimate_order_of_convergence()[0, 1] + assert orderest > expected_order * 0.9 + + +if __name__ == "__main__": + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) -- GitLab From e4651774ee052b591528326dd7cb30fae2367feb Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 24 Jun 2020 11:11:33 -0500 Subject: [PATCH 02/14] Adds implicit DIRK schemes for Adams-Moulton bootstrapping, fixes Backward Euler tableau --- leap/rk/__init__.py | 141 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 139 insertions(+), 2 deletions(-) diff --git a/leap/rk/__init__.py b/leap/rk/__init__.py index e11393c..5c66be6 100644 --- a/leap/rk/__init__.py +++ b/leap/rk/__init__.py @@ -396,6 +396,32 @@ class SimpleButcherTableauMethodBuilder(ButcherTableauMethodBuilder): }) +class ImplicitButcherTableauMethodBuilder(ButcherTableauMethodBuilder): + def __init__(self, component_id, state_filter_name=None, + rhs_func_name=None): + super(SimpleButcherTableauMethodBuilder, self).__init__( + component_id=component_id, + state_filter_name=state_filter_name) + + if rhs_func_name is None: + rhs_func_name = ""+self.component_id + self.rhs_func_name = rhs_func_name + + def generate(self): + """ + :returns: :class:`dagrt.language.DAGCode` + """ + return self.generate_butcher( + stage_coeff_set_names=("implicit",), + stage_coeff_sets={ + "implicit": self.a_implicit}, + rhs_funcs={"implicit": var(self.rhs_func_name)}, + estimate_coeff_set_names=("main",), + estimate_coeff_sets={ + "main": self.output_coeffs, + }) + + class ForwardEulerMethodBuilder(SimpleButcherTableauMethodBuilder): """ .. automethod:: __init__ @@ -417,9 +443,9 @@ class BackwardEulerMethodBuilder(SimpleButcherTableauMethodBuilder): .. automethod:: __init__ .. automethod:: generate """ - c = (0,) + c = (1,) - a_explicit = ( + a_implicit = ( (1,), ) @@ -524,6 +550,109 @@ class RK5MethodBuilder(SimpleButcherTableauMethodBuilder): recycle_last_stage_coeff_set_names = () +class DIRK2MethodBuilder(ImplicitButcherTableauMethodBuilder): + """ + .. automethod:: __init__ + .. automethod:: generate + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 72, eqn 221 + """ + + x = (2 - np.sqrt(2))/2 + + c = (x, 1) + + a_implicit = ( + (x,), + (1 - x, x,), + ) + + output_coeffs = (1 - x, x) + + recycle_last_stage_coeff_set_names = () + + +class DIRK3MethodBuilder(ImplicitButcherTableauMethodBuilder): + """ + .. automethod:: __init__ + .. automethod:: generate + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 77, eqn 229 & eqn 230 + """ + + x = 0.4358665215 + + c = (x, (1 + x)/2, 1) + + a_implicit = ( + (x,), + ((1 - x)/2, x,), + (-3*x**2/2 + 4*x - 0.25, 3*x**2/2 - 5*x + 5/4, x,), + ) + + output_coeffs = (-3*x**2/2 + 4*x - 0.25, 3*x**2/2 - 5*x + 5/4, x) + + recycle_last_stage_coeff_set_names = () + + +class DIRK4MethodBuilder(ImplicitButcherTableauMethodBuilder): + """ + .. automethod:: __init__ + .. automethod:: generate + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 78, eqn 232 + """ + + x1 = 1.06858 + + c = (x1, 1/2, 1 - x1) + + a_implicit = ( + (x1,), + (1/2 - x1, x1,), + (2*x1, 1 - 4*x1, x1), + ) + + output_coeffs = (1/(6*(1 - 2*x1)**2), (3*(1 - 2*x1)**2 - 1)/(3*(1 - 2*x1)**2), + 1/(6*(1 - 2*x1)**2)) + + recycle_last_stage_coeff_set_names = () + + +class DIRK5MethodBuilder(ImplicitButcherTableauMethodBuilder): + """ + .. automethod:: __init__ + .. automethod:: generate + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 98, Table 24 + """ + + c = (4024571134387/14474071345096, 5555633399575/5431021154178, + 5255299487392/12852514622453, 3/20, 10449500210709/14474071345096) + + a_implicit = ( + (4024571134387/14474071345096,), + (9365021263232/12572342979331, 4024571134387/14474071345096,), + (2144716224527/9320917548702, -397905335951/4008788611757, + 4024571134387/14474071345096,), + (-291541413000/6267936762551, 226761949132/4473940808273, + -1282248297070/9697416712681, 4024571134387/14474071345096,), + (-2481679516057/4626464057815, -197112422687/6604378783090, + 3952887910906/9713059315593, 4906835613583/8134926921134, + 4024571134387/14474071345096,), + ) + + output_coeffs = (-2522702558582/12162329469185, 1018267903655/12907234417901, + 4542392826351/13702606430957, 5001116467727/12224457745473, + 1509636094297/3891594770934) + + recycle_last_stage_coeff_set_names = () + + ORDER_TO_RK_METHOD_BUILDER = { 1: ForwardEulerMethodBuilder, 2: MidpointMethodBuilder, @@ -532,6 +661,14 @@ ORDER_TO_RK_METHOD_BUILDER = { 5: RK5MethodBuilder, } +IMPLICIT_ORDER_TO_RK_METHOD_BUILDER = { + 1: BackwardEulerMethodBuilder, + 2: DIRK2MethodBuilder, + 3: DIRK3MethodBuilder, + 4: DIRK4MethodBuilder, + 5: DIRK4MethodBuilder, + } + # }}} -- GitLab From 44baab1171f81e1b50d655160ac48a64c922182c Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 24 Jun 2020 11:12:27 -0500 Subject: [PATCH 03/14] Fixes single-rate AM to use history in a way more consistent with AB and avoid extra RHS evals - also uses implicit DIRK for bootstrap --- leap/multistep/__init__.py | 124 ++++++++++++++++++++++++++----------- 1 file changed, 87 insertions(+), 37 deletions(-) diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index 2736326..62c3a11 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -453,15 +453,13 @@ class AdamsMoultonMethodBuilder(MethodBuilder): # Declare variables self.step = var('

step') self.function = var('' + component_id) - # One shorter for implicit. self.history = \ - [var('

f_n_minus_' + str(i)) for i in range(hist_length - 2, 0, -1)] + [var('

f_n_minus_' + str(i)) for i in range(hist_length - 1, 0, -1)] if not self.static_dt: - # One shorter for implicit. self.time_history = [ var('

t_n_minus_' + str(i)) - for i in range(hist_length - 2, 0, -1)] + for i in range(hist_length - 1, 0, -1)] self.state = var('' + component_id) self.t = var('') @@ -482,7 +480,6 @@ class AdamsMoultonMethodBuilder(MethodBuilder): from dagrt.language import DAGCode, CodeBuilder array = var("array") - rhs_var = var("rhs_var") rhs_next_var = var("rhs_next_var") # Initialization @@ -496,10 +493,10 @@ class AdamsMoultonMethodBuilder(MethodBuilder): unkvar = cb_primary.fresh_var('unk') rhs_var_to_unknown[rhs_next_var] = unkvar - # In implicit mode, the time history must also - # include the next point in time. + # In implicit mode, the time history must + # include the *next* point in time. if not self.static_dt: - time_history_data = self.time_history + [self.t] + [self.t + self.dt] + time_history_data = self.time_history + [self.t + self.dt] time_hist_var = var(name_gen("time_history")) cb_primary(time_hist_var, array(self.hist_length)) for i in range(self.hist_length): @@ -525,10 +522,8 @@ class AdamsMoultonMethodBuilder(MethodBuilder): unknowns.add(rhs_next_var) - cb_primary(rhs_var, self.eval_rhs(self.t, self.state)) - # Update history - history = self.history + [rhs_var] + [rhs_next_var] + history = self.history + [rhs_next_var] # Set up the actual Adams-Moulton step. ab_sum = emit_ab_integration( @@ -584,8 +579,7 @@ class AdamsMoultonMethodBuilder(MethodBuilder): cb_primary(self.state, state_est) # Rotate history and time history. - # In the case of implicit, history length is one less. - for i in range(self.hist_length - 2): + for i in range(self.hist_length - 1): cb_primary(self.history[i], history[i + 1]) if not self.static_dt: @@ -634,41 +628,84 @@ class AdamsMoultonMethodBuilder(MethodBuilder): kw_parameters={"t": t, self.component_id: y}) def rk_bootstrap(self, cb): - """Initialize the timestepper with an RK method.""" + """Initialize the timestepper with an IMPLICIT RK method.""" - rhs_var = var("rhs_var") + equations = [] + unknowns = set() + knowns = set() + rhs_var_to_unknown = {} - cb(rhs_var, self.eval_rhs(self.t, self.state)) + def make_known(v): + unknowns.discard(v) + knowns.add(v) - # Save the current RHS to the AB history + from leap.rk import IMPLICIT_ORDER_TO_RK_METHOD_BUILDER + rk_method = IMPLICIT_ORDER_TO_RK_METHOD_BUILDER[self.function_family.order] + rk_tableau = tuple(zip(rk_method.c, rk_method.a_implicit)) + rk_coeffs = rk_method.output_coeffs - for i in range(len(self.history)): - with cb.if_(self.step, "==", i + 1): - cb(self.history[i], rhs_var) + with cb.if_(self.step, "==", 1): + # Save the first RHS to the AM history + rhs_var = var("rhs_var") - if not self.static_dt: - cb(self.time_history[i], self.t) + cb(rhs_var, self.eval_rhs(self.t, self.state)) + cb(self.history[0], rhs_var) - from leap.rk import ORDER_TO_RK_METHOD_BUILDER - rk_method = ORDER_TO_RK_METHOD_BUILDER[self.function_family.order] - rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit)) - rk_coeffs = rk_method.output_coeffs + if not self.static_dt: + cb(self.time_history[0], self.t) - # Stage loop (taken from EmbeddedButcherTableauMethodBuilder) + # Stage loop rhss = [var("rk_rhs_" + str(i)) for i in range(len(rk_tableau))] for stage_num, (c, coeffs) in enumerate(rk_tableau): - if len(coeffs) == 0: - assert c == 0 - cb(rhss[stage_num], rhs_var) - else: - stage = self.state + sum(self.dt * coeff * rhss[j] - for (j, coeff) - in enumerate(coeffs)) + stage = self.state + sum(self.dt * coeff * rhss[j] + for (j, coeff) + in enumerate(coeffs)) - if self.state_filter is not None: - stage = self.state_filter(stage) + if self.state_filter is not None: + stage = self.state_filter(stage) - cb(rhss[stage_num], self.eval_rhs(self.t + c * self.dt, stage)) + # In a DIRK setting, the unknown is always the same RHS + # as the stage number. + unknowns.add(rhss[stage_num]) + unkvar = cb.fresh_var('unk_s%d' % (stage_num)) + rhs_var_to_unknown[rhss[stage_num]] = unkvar + from dagrt.expression import collapse_constants + solve_expression = collapse_constants( + rhss[stage_num] - self.eval_rhs(self.t + c*self.dt, stage), + list(unknowns) + [self.state], + cb.assign, cb.fresh_var) + equations.append(solve_expression) + + # {{{ emit solve if possible + + if unknowns and len(unknowns) == len(equations): + # got a square system, let's solve + assignees = [unk.name for unk in unknowns] + + from pymbolic import substitute + subst_dict = dict( + (rhs_var.name, rhs_var_to_unknown[rhs_var]) + for rhs_var in unknowns) + + cb.assign_implicit( + assignees=assignees, + solve_components=[ + rhs_var_to_unknown[unk].name + for unk in unknowns], + expressions=[ + substitute(eq, subst_dict) + for eq in equations], + + # TODO: Could supply a starting guess + other_params={ + "guess": self.state}, + solver_id="solve") + + del equations[:] + knowns.update(unknowns) + unknowns.clear() + + # }}} # Merge the values of the RHSs. rk_comb = sum(coeff * rhss[j] for j, coeff in enumerate(rk_coeffs)) @@ -680,6 +717,19 @@ class AdamsMoultonMethodBuilder(MethodBuilder): # Assign the value of the new state. cb(self.state, state_est) + # Save the "next" RHS to the AM history + rhs_next_var = var("rhs_next_var") + + cb(rhs_next_var, self.eval_rhs(self.t + self.dt, self.state)) + + for i in range(1, len(self.history)): + with cb.if_(self.step, "==", i): + cb(self.history[i], rhs_next_var) + + if not self.static_dt: + cb(self.time_history[i], self.t + self.dt) + + # }}} # vim: fdm=marker -- GitLab From 1d1e717f7d3e84d26a4b86f5ef8b25440fb0f2eb Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 24 Jun 2020 11:17:17 -0500 Subject: [PATCH 04/14] Fix copyright line, add first order to tests --- test/test_implicit_ab.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_implicit_ab.py b/test/test_implicit_ab.py index 6681fdf..9d04a2c 100755 --- a/test/test_implicit_ab.py +++ b/test/test_implicit_ab.py @@ -1,7 +1,7 @@ #! /usr/bin/env python from __future__ import division, with_statement -__copyright__ = "Copyright (C) 2014 Matt Wala" +__copyright__ = "Copyright (C) 2020 Cory Mikida" __license__ = """ Permission is hereby granted, free of charge, to any person obtaining a copy @@ -66,12 +66,12 @@ def am_solver_hook_static(solve_expr, solve_var, solver_id, guess): @pytest.mark.parametrize(("method", "expected_order", "static_dt"), [ (AdamsMoultonMethodBuilder("y", order, static_dt=static_dt), order, static_dt) - for order in [2, 3, 4, 5] + for order in [1, 2, 3, 4, 5] for static_dt in [True, False] ] + [ (AdamsMoultonMethodBuilder("y", order, hist_length=order+1, static_dt=static_dt), order, static_dt) - for order in [2, 3, 4, 5] + for order in [1, 2, 3, 4, 5] for static_dt in [True, False] ]) def test_am_convergence(python_method_impl, method, expected_order, static_dt): -- GitLab From 1905ae42ffa1a93fcba2aff9ebcfe7e9971c9349 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 5 Aug 2020 13:48:50 -0500 Subject: [PATCH 05/14] Simplifies solver hooks by applying distribution to solve expressions --- leap/multistep/__init__.py | 12 +++--------- test/test_implicit_ab.py | 35 +++++------------------------------ 2 files changed, 8 insertions(+), 39 deletions(-) diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index 62c3a11..593d293 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -516,10 +516,6 @@ class AdamsMoultonMethodBuilder(MethodBuilder): unknowns = set() knowns = set() - def make_known(v): - unknowns.discard(v) - knowns.add(v) - unknowns.add(rhs_next_var) # Update history @@ -538,8 +534,10 @@ class AdamsMoultonMethodBuilder(MethodBuilder): # Build the implicit solve expression. from dagrt.expression import collapse_constants + from pymbolic.mapper.distributor import DistributeMapper as DistMap solve_expression = collapse_constants( - rhs_next_var - self.eval_rhs(self.t + self.dt, state_est), + rhs_next_var - self.eval_rhs(self.t + self.dt, + DistMap()(state_est)), list(unknowns) + [self.state], cb_primary.assign, cb_primary.fresh_var) equations.append(solve_expression) @@ -635,10 +633,6 @@ class AdamsMoultonMethodBuilder(MethodBuilder): knowns = set() rhs_var_to_unknown = {} - def make_known(v): - unknowns.discard(v) - knowns.add(v) - from leap.rk import IMPLICIT_ORDER_TO_RK_METHOD_BUILDER rk_method = IMPLICIT_ORDER_TO_RK_METHOD_BUILDER[self.function_family.order] rk_tableau = tuple(zip(rk_method.c, rk_method.a_implicit)) diff --git a/test/test_implicit_ab.py b/test/test_implicit_ab.py index 9d04a2c..704be32 100755 --- a/test/test_implicit_ab.py +++ b/test/test_implicit_ab.py @@ -49,21 +49,6 @@ def am_solver_hook(solve_expr, solve_var, solver_id, guess): return substitute("solver(t, sub_y, coeff, guess)", pieces) -def am_solver_static(f, t, sub_y, coeff, dt, add, guess): - from scipy.optimize import root - return root(lambda unk: unk - f(t=t, y=sub_y + dt*(add + coeff*unk)), guess).x - - -def am_solver_hook_static(solve_expr, solve_var, solver_id, guess): - from dagrt.expression import match, substitute - - pieces = match("unk - rhs(t=t, y=sub_y + dt*(add + coeff*unk))", - solve_expr, - pre_match={"unk": solve_var}) - pieces["guess"] = guess - return substitute("solver(t, sub_y, coeff, dt, add, guess)", pieces) - - @pytest.mark.parametrize(("method", "expected_order", "static_dt"), [ (AdamsMoultonMethodBuilder("y", order, static_dt=static_dt), order, static_dt) for order in [1, 2, 3, 4, 5] @@ -84,10 +69,7 @@ def test_am_convergence(python_method_impl, method, expected_order, static_dt): code = method.generate() from leap.implicit import replace_AssignImplicit - if static_dt: - code = replace_AssignImplicit(code, {"solve": am_solver_hook_static}) - else: - code = replace_AssignImplicit(code, {"solve": am_solver_hook}) + code = replace_AssignImplicit(code, {"solve": am_solver_hook}) from pytools.convergence import EOCRecorder eocrec = EOCRecorder() @@ -101,17 +83,10 @@ def test_am_convergence(python_method_impl, method, expected_order, static_dt): final_t = problem.t_end from functools import partial - # FIXME - if static_dt: - interp = python_method_impl(code, function_map={ - "" + component_id: problem, - "solver": partial(am_solver_static, problem), - }) - else: - interp = python_method_impl(code, function_map={ - "" + component_id: problem, - "solver": partial(am_solver, problem), - }) + interp = python_method_impl(code, function_map={ + "" + component_id: problem, + "solver": partial(am_solver, problem), + }) interp.set_up(t_start=t, dt_start=dt, context={component_id: y}) times = [] -- GitLab From f9f661972ba50b43ad393a6239840247f950e998 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 5 Aug 2020 13:50:36 -0500 Subject: [PATCH 06/14] Sets up IMEX bootstrapping for IMEX MR --- leap/rk/__init__.py | 8 ++++++++ leap/rk/imex.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/leap/rk/__init__.py b/leap/rk/__init__.py index 5c66be6..2e6dca7 100644 --- a/leap/rk/__init__.py +++ b/leap/rk/__init__.py @@ -669,6 +669,14 @@ IMPLICIT_ORDER_TO_RK_METHOD_BUILDER = { 5: DIRK4MethodBuilder, } +from leap.rk.imex import (KennedyCarpenterIMEXARK3MethodBuilder, + KennedyCarpenterIMEXARK4MethodBuilder) + +IMEX_ORDER_TO_RK_METHOD_BUILDER = { + 3: KennedyCarpenterIMEXARK3MethodBuilder, + 4: KennedyCarpenterIMEXARK4MethodBuilder, + } + # }}} diff --git a/leap/rk/imex.py b/leap/rk/imex.py index b4884b9..b3e2454 100644 --- a/leap/rk/imex.py +++ b/leap/rk/imex.py @@ -213,3 +213,46 @@ class KennedyCarpenterIMEXARK4MethodBuilder( == len(low_order_coeffs) == len(high_order_coeffs) == len(c)) + + +class KennedyCarpenterIMEXARK3MethodBuilder( + KennedyCarpenterIMEXRungeKuttaMethodBuilderBase): + gamma = 1767732205903/4055673282236 + + c = [0, 2*gamma, 3/5, 1] + low_order = 3 + high_order = 4 + + high_order_coeffs = [ + 1471266399579/7840856788654, + -4482444167858/7529755066697, + 11266239266428/11593286722821, + gamma] + + low_order_coeffs = [ + 2756255671327/12835298489170, + -10771552573575/22201958757719, + 9247589265047/10645013368117, + 2193209047091/5459859503100] + + # ARK3(2)4L[2]SA-ERK (explicit) + a_explicit = [[], + [2*gamma], + [5535828885825/10492691773637, 788022342437/10882634858940], + [6485989280629/16251701735622, + -4246266847089/9704473918619, + 10755448449292/10357097424841]] + + # ARK4(3)6L[2]SA-ESDIRK (implicit) + a_implicit = [[], + [gamma, gamma], + [2746238789719/10658868560708, -640167445237/6845629431997, gamma], + [1471266399579/7840856788654, -4482444167858/7529755066697, + 11266239266428/11593286722821, gamma]] + + recycle_last_stage_coeff_set_names = ("implicit",) + + assert (len(a_explicit) == len(a_implicit) + == len(low_order_coeffs) + == len(high_order_coeffs) + == len(c)) -- GitLab From 62de9f852cc0ec3ce1dce02d0cd8cd32383b4350 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 5 Aug 2020 14:07:50 -0500 Subject: [PATCH 07/14] First pass at implicit-explicit Adams-based multi-rate --- leap/multistep/multirate/__init__.py | 437 +++++++++++++++++++++++++-- test/test_multirate.py | 327 ++++++++++++++++++++ 2 files changed, 737 insertions(+), 27 deletions(-) diff --git a/leap/multistep/multirate/__init__.py b/leap/multistep/multirate/__init__.py index a8481f1..4757f91 100644 --- a/leap/multistep/multirate/__init__.py +++ b/leap/multistep/multirate/__init__.py @@ -67,7 +67,7 @@ class MultiRateHistory(Record): """ def __init__(self, interval, func_name, arguments, order=None, rhs_policy=rhs_policy.late, invalidate_computed_state=False, - hist_length=None): + hist_length=None, is_rhs_implicit=False): """ :arg interval: An integer indicating the interval (relative to the smallest available timestep) at which this history is to be @@ -85,6 +85,8 @@ class MultiRateHistory(Record): :arg hist_length: history length. If greater than order, we use a least-squares solve rather than a linear solve to obtain the AB coefficients for this history + :arg is_rhs_implicit: If True, we use Adams-Moulton for the state + contribution from this RHS. """ super(MultiRateHistory, self).__init__( interval=interval, @@ -93,12 +95,17 @@ class MultiRateHistory(Record): order=order, rhs_policy=rhs_policy, invalidate_computed_state=invalidate_computed_state, - hist_length=hist_length) + hist_length=hist_length, + is_rhs_implicit=is_rhs_implicit) @property def history_length(self): return self.hist_length + @property + def is_implicit(self): + return self.is_rhs_implicit + class RHS(MultiRateHistory): def __init__(self, *args, **kwargs): @@ -159,7 +166,8 @@ class InconsistentHistoryError: class MultiRateMultiStepMethodBuilder(MethodBuilder): """Simultaneously timesteps multiple parts of an ODE system, - each with adjustable orders, rates, and dependencies. + each with adjustable orders, rates, and dependencies. RHS + components can also be advanced implicitly. Considerably generalizes [GearWells]_. @@ -413,7 +421,7 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): cb(self.bootstrap_step, 0) - # {{{ rk bootstrap: step + # {{{ rk bootstrap: fully explicit step def emit_small_rk_step(self, cb, name_prefix, name_gen, rhss_on_entry): """Emit a single step of an RK method.""" @@ -588,6 +596,283 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): # }}} + # {{{ rk bootstrap: IMEX step + + def emit_small_imex_rk_step(self, cb, name_prefix, name_gen, rhss_on_entry): + """Emit a single step of an IMEX-RK method.""" + + from leap.rk import (IMEX_ORDER_TO_RK_METHOD_BUILDER, + _is_first_stage_same_as_last_stage) + rk_method = IMEX_ORDER_TO_RK_METHOD_BUILDER[self.max_order] + rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit, + rk_method.a_implicit)) + rk_coeffs = rk_method.high_order_coeffs + stage_coeff_sets = { + "explicit": rk_method.a_explicit, + "implicit": rk_method.a_implicit} + + def make_stage_history(prefix): + return [var(prefix + "_stage" + str(i)) for i in range(len(rk_tableau))] + + def make_imp_stage_history(prefix): + return [var(prefix + "_stage" + str(i) + "_imp") + for i in range(len(rk_tableau))] + + stage_rhss = {} + imp_stage_rhss = {} + for comp_name, component_rhss in zip(self.component_names, self.rhss): + if not self.is_ode_component[comp_name]: + continue + + for irhs, rhs in enumerate(component_rhss): + if rhs.is_implicit: + imp_stage_rhss[comp_name, irhs] = make_imp_stage_history( + "{name_prefix}_rk_{comp_name}_rhs{irhs}" + .format( + name_prefix=name_prefix, + comp_name=comp_name, + irhs=irhs)) + else: + stage_rhss[comp_name, irhs] = make_stage_history( + "{name_prefix}_rk_{comp_name}_rhs{irhs}" + .format( + name_prefix=name_prefix, + comp_name=comp_name, + irhs=irhs)) + + for istage, (c, exp_coeffs, imp_coeffs) in enumerate(rk_tableau): + + unknowns = set() + knowns = set() + component_state_ests = {} + for comp_name, component_rhss in zip( + self.component_names, self.rhss): + if not self.is_ode_component[comp_name]: + continue + + contribs = [] + for irhs, rhs in enumerate(component_rhss): + # We still need to go through the whole solving deal here. + if rhs.is_implicit: + name = "implicit" + else: + name = "explicit" + if ( + name in rk_method.recycle_last_stage_coeff_set_names + and istage == 0 + and _is_first_stage_same_as_last_stage( + rk_method.c, stage_coeff_sets[name])): + if rhs.is_implicit: + unknowns.add(imp_stage_rhss[comp_name, irhs][istage]) + else: + cb(stage_rhss[comp_name, irhs][istage], + rhss_on_entry[comp_name, irhs]) + else: + state_contrib_var = var( + name_gen( + "state_contrib_{comp_name}_rhs{irhs}" + .format(comp_name=comp_name, irhs=irhs))) + + if rhs.is_implicit: + unknowns.add(imp_stage_rhss[comp_name, irhs][istage]) + contribs.append(_linear_comb(imp_coeffs, + imp_stage_rhss[comp_name, irhs])) + else: + # First stage. + if exp_coeffs == []: + cb(state_contrib_var, 0) + else: + contribs.append(_linear_comb(exp_coeffs, + stage_rhss[comp_name, irhs])) + cb(state_contrib_var, + _linear_comb(exp_coeffs, + stage_rhss[comp_name, irhs])) + + state_var = var( + name_gen( + "state_{comp_name}_st{istage}" + .format(comp_name=comp_name, istage=istage))) + + state_expr = ( + var("" + comp_name) + + (self.dt/self.nsubsteps) * sum(contribs)) + + if comp_name in self.state_filters: + state_expr = self.state_filters[comp_name](state_expr) + + # Avoid need for multiple solver hooks. + from pymbolic.mapper.distributor import DistributeMapper as DistMap + component_state_ests[comp_name] = DistMap()(state_expr) + + # At this point, we have all the ODE state estimates evaluated. + + # {{{ evaluate the non-ODE RHSs + + for comp_name, component_rhss in zip( + self.component_names, self.rhss): + if self.is_ode_component[comp_name]: + continue + + contribs = [] + + for irhs, rhs in enumerate(component_rhss): + kwargs = dict( + (self.comp_name_to_kwarg_name[arg_comp_name], + component_state_ests[arg_comp_name]) + for arg_comp_name in rhs.arguments) + + contribs.append(var(rhs.func_name)( + t=self.t + (c/self.nsubsteps) * self.dt, + **kwargs)) + + state_var = var( + name_gen( + "state_{comp_name}_st{istage}" + .format(comp_name=comp_name, istage=istage))) + + cb(state_var, sum(contribs)) + + component_state_ests[comp_name] = state_var + + # }}} + + # {{{ evaluate the ODE RHSs + + for comp_name, component_rhss in zip( + self.component_names, self.rhss): + + if not self.is_ode_component[comp_name]: + continue + + equations = [] + rhs_var_to_unknown = {} + + for irhs, rhs in enumerate(component_rhss): + kwargs = dict( + (self.comp_name_to_kwarg_name[arg_comp_name], + component_state_ests[arg_comp_name]) + for arg_comp_name in rhs.arguments) + if rhs.is_implicit: + # First stage: no solve required. + if imp_coeffs == []: + cb(imp_stage_rhss[comp_name, irhs][istage], + var(rhs.func_name)( + t=self.t + (c/self.nsubsteps) * self.dt, + **kwargs)) + else: + # We require a solve here. + # In a DIRK/ARK setting, the unknown is always the same + # RHS as the stage number. + unkvar = cb.fresh_var('unk_s%d' % (istage)) + rhs_var_to_unknown[ + imp_stage_rhss[comp_name, + irhs][istage]] = unkvar + solve_expression = imp_stage_rhss[ + comp_name, irhs][istage] - var(rhs.func_name)( + t=self.t + (c/self.nsubsteps) * self.dt, + **kwargs) + equations.append(solve_expression) + + # Emit solve. + + if unknowns and len(unknowns) == len(equations): + # got a square system, let's solve + assignees = [unk.name for unk in unknowns] + + from pymbolic import substitute + subst_dict = dict( + (rhs_var.name, rhs_var_to_unknown[rhs_var]) + for rhs_var in unknowns) + + cb.assign_implicit( + assignees=assignees, + solve_components=[ + rhs_var_to_unknown[unk].name + for unk in unknowns], + expressions=[ + substitute(eq, subst_dict) + for eq in equations], + + # TODO: Could supply a starting guess + other_params={ + "guess": var(''+comp_name)}, + solver_id="solve") + + del equations[:] + knowns.update(unknowns) + unknowns.clear() + + else: + cb(stage_rhss[comp_name, irhs][istage], + var(rhs.func_name)( + t=self.t + (c/self.nsubsteps) * self.dt, + **kwargs)) + + # }}} + + # Now that we've traversed both stage loops, + # we can actually sum the RHS's to get the state estimates. + component_state_ests = {} + + for icomp, (comp_name, component_rhss) in enumerate( + zip(self.component_names, self.rhss)): + + contribs = [] + for irhs, rhs in enumerate(component_rhss): + if not self.is_ode_component[comp_name]: + continue + + # Implicit, if applicable: + if rhs.is_implicit: + imp_state_contrib_var = var( + name_gen( + "state_contrib_{comp_name}_rhs{irhs}_imp" + .format(comp_name=comp_name, irhs=irhs))) + + contribs.append(imp_state_contrib_var) + + cb(imp_state_contrib_var, + _linear_comb(rk_coeffs, imp_stage_rhss[comp_name, irhs])) + else: + state_contrib_var = var( + name_gen( + "state_contrib_{comp_name}_rhs{irhs}" + .format(comp_name=comp_name, irhs=irhs))) + + contribs.append(state_contrib_var) + + cb(state_contrib_var, + _linear_comb(rk_coeffs, stage_rhss[comp_name, irhs])) + + state_var = var( + name_gen( + "state_{comp_name}_final" + .format(comp_name=comp_name))) + + state_expr = ( + var("" + comp_name) + + (self.dt/self.nsubsteps) * sum(contribs)) + if comp_name in self.state_filters: + state_expr = self.state_filters[comp_name](state_expr) + + cb(state_var, state_expr) + + component_state_ests[comp_name] = state_var + + for component_name in self.component_names: + state = component_state_ests[component_name] + if self.is_ode_component[component_name]: + cb.yield_state( + state, + component_name, self.t + self.dt/self.nsubsteps, + "bootstrap") + + cb(var(""+component_name), state) + + cb(self.t, self.t + self.dt/self.nsubsteps) + + # }}} + # {{{ rk bootstrap: overall control def emit_rk_bootstrap(self, cb): @@ -711,7 +996,20 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): cb.switch_phase("primary") cb.restart_step() - self.emit_small_rk_step(cb, name_prefix, name_gen, current_rhss) + # Check if we have any implicit RHSs. + imp_exist = 0 + for comp_name, component_rhss in zip( + self.component_names, self.rhss): + for irhs, rhs in enumerate(component_rhss): + if rhs.is_implicit: + imp_exist += 1 + + if imp_exist == 0: + # If we have no implicit RHSs, run the full explicit bootstrapper. + self.emit_small_rk_step(cb, name_prefix, name_gen, current_rhss) + else: + # Otherwise, run the IMEX RK bootstrapper. + self.emit_small_imex_rk_step(cb, name_prefix, name_gen, current_rhss) cb(self.bootstrap_step, self.bootstrap_step + 1) @@ -724,7 +1022,7 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): # {{{ main method generation - def emit_ab_method(self, cb, explainer): + def emit_adams_method(self, cb, explainer): from pytools import UniqueNameGenerator name_gen = UniqueNameGenerator() @@ -782,7 +1080,7 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): # {{{ get_state - def get_state(comp_name, isubstep): + def get_state(comp_name, isubstep, return_exp=False): states = computed_states[comp_name] # {{{ see if we've got that state ready to go @@ -865,7 +1163,15 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): if comp_name in self.state_filters: state_expr = self.state_filters[comp_name](state_expr) - cb(state_var, state_expr) + + # If we are returning this state as an implicit RHS + # argument, we need it to be an expression, not a variable. + if return_exp: + # To simplify solver hooks + from pymbolic.mapper.distributor import DistributeMapper as DistMap + return DistMap()(state_expr) + else: + cb(state_var, state_expr) # Only keep temporary state if integrates exactly # one interval ahead for the fastest right-hand side, @@ -905,26 +1211,31 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): rhs = self.rhss[comp_idx][irhs] - # {{{ get arguments together - - progress_frac = isubstep / self.nsubsteps - t_expr = self.t + self.dt * progress_frac - - kwargs = dict( - (self.comp_name_to_kwarg_name[arg_comp_name], - get_state(arg_comp_name, isubstep)) - for arg_comp_name in rhs.arguments) - - # }}} - rhs_var = var( name_gen( "rhs_{comp_name}_rhs{irhs}_sub{isubstep}" .format(comp_name=comp_name, irhs=irhs, isubstep=isubstep))) - cb(rhs_var, var(rhs.func_name)(t=t_expr, **kwargs)) + # Stuff for implicit. + equations = [] + unknowns = set() + knowns = set() + rhs_var_to_unknown = {} + states = computed_states[comp_name] + latest_state = states[-1][-1] - temp_hist_substeps[comp_name, irhs].append(isubstep) + if rhs.is_implicit: + # Implicit setup - rhs_var is an unknown, needs implicit solve. + temp_hist_vars[comp_name, irhs].append(rhs_var) + implicit_rhs_var = temp_hist_vars[comp_name, irhs][-1] + unknowns.add(implicit_rhs_var) + unkvar = cb.fresh_var('unk') + rhs_var_to_unknown[implicit_rhs_var] = unkvar + + # {{{ get arguments together + + progress_frac = isubstep / self.nsubsteps + t_expr = self.t + self.dt * progress_frac if not self.static_dt: t_var = var( @@ -934,13 +1245,85 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): comp_name=comp_name, irhs=irhs, isubstep=isubstep))) - cb(t_var, t_expr) - temp_time_vars[comp_name, irhs].append(t_var) + if rhs.is_implicit: + cb(t_var, t_expr) + temp_time_vars[comp_name, irhs].append(t_var) + else: + if rhs.is_implicit: + temp_time_vars[comp_name, irhs].append(progress_frac) + # }}} + + if rhs.is_implicit: + # TODO multivariable implicit solves + # Call to get state for a given RHS. + kwargs = dict( + (self.comp_name_to_kwarg_name[arg_comp_name], + get_state(arg_comp_name, isubstep, return_exp=True)) + for arg_comp_name in rhs.arguments) + # Build the implicit solve expression. + solve_expression = implicit_rhs_var - var(rhs.func_name)(t=t_expr, + **kwargs) + equations.append(solve_expression) else: - temp_time_vars[comp_name, irhs].append(progress_frac) + kwargs = dict( + (self.comp_name_to_kwarg_name[arg_comp_name], + get_state(arg_comp_name, isubstep)) + for arg_comp_name in rhs.arguments) + cb(rhs_var, var(rhs.func_name)(t=t_expr, **kwargs)) + + # Implicit solve(s), if needed. + if unknowns and len(unknowns) == len(equations): + # got a square system, let's solve + assignees = [unk.name for unk in unknowns] + + from pymbolic import substitute + subst_dict = dict( + (rhs_var.name, rhs_var_to_unknown[rhs_var]) + for rhs_var in unknowns) + + cb.assign_implicit( + assignees=assignees, + solve_components=[ + rhs_var_to_unknown[unk].name + for unk in unknowns], + expressions=[ + substitute(eq, subst_dict) + for eq in equations], + + other_params={ + "guess": latest_state}, + solver_id="solve") + + del equations[:] + knowns.update(unknowns) + unknowns.clear() + + # Now that we have the RHS value solved for, + # update useful temp states + for arg_comp_name in rhs.arguments: + post_solve_state = get_state(arg_comp_name, isubstep) + + temp_hist_substeps[comp_name, irhs].append(isubstep) + + if not rhs.is_implicit: + # History shuffles. + if not self.static_dt: + t_var = var( + name_gen( + "t_{comp_name}_rhs{irhs}_sub{isubstep}" + .format( + comp_name=comp_name, + irhs=irhs, + isubstep=isubstep))) + + cb(t_var, t_expr) + temp_time_vars[comp_name, irhs].append(t_var) + + else: + temp_time_vars[comp_name, irhs].append(progress_frac) - temp_hist_vars[comp_name, irhs].append(rhs_var) + temp_hist_vars[comp_name, irhs].append(rhs_var) explainer.eval_rhs( rhs_var.name, comp_name, rhs.func_name, isubstep, kwargs) @@ -1129,7 +1512,7 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): self.emit_initialization(cb_init) with CodeBuilder(name="primary") as cb_primary: - self.emit_ab_method(cb_primary, explainer) + self.emit_adams_method(cb_primary, explainer) with CodeBuilder(name="bootstrap") as cb_bootstrap: self.emit_rk_bootstrap(cb_bootstrap) diff --git a/test/test_multirate.py b/test/test_multirate.py index 045ce07..55c7841 100644 --- a/test/test_multirate.py +++ b/test/test_multirate.py @@ -319,6 +319,333 @@ def test_single_rate_identical(order=3, hist_length=3): assert diff < 1e-13 +def am_solver(f, t, dt, sub_fast, sub_slow, coeff1, coeff2, guess): + from scipy.optimize import root + return root(lambda unk: unk + (-1)*f(t=t + coeff1*dt, + fast=sub_fast + coeff2*unk, + slow=sub_slow), guess).x + + +def am_solver_hook(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match( + "unk + (-1)*f(t=t + coeff1*dt, fast=sub_fast + coeff2*unk, " + "slow=sub_slow)", solve_expr, pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, dt, sub_fast, sub_slow, " + "coeff1, coeff2, guess)", pieces) + + +def am_solver_sr(f, t, sub_y, coeff, guess): + from scipy.optimize import root + return root(lambda unk: unk - f(t=t, y=sub_y + coeff*unk), guess).x + + +def am_solver_sr_hook(solve_expr, solve_var, solver_id, guess): + from dagrt.expression import match, substitute + + pieces = match("unk - rhs(t=t, y=sub_y + coeff*unk)", solve_expr, + pre_match={"unk": solve_var}) + pieces["guess"] = guess + return substitute("solver(t, sub_y, coeff, guess)", pieces) + + +def test_implicit_single_rate_identical(order=3, hist_length=3): + from dagrt.exec_numpy import NumpyInterpreter + from leap.multistep import (AdamsMoultonMethodBuilder, + AdamsBashforthMethodBuilder) + from leap.implicit import replace_AssignImplicit + from functools import partial + + from multirate_test_systems import StiffUncoupled + # We need to use the uncoupled system to test + # IMEX. + ode = StiffUncoupled() + + t_start = 0 + dt = 0.1 + nsteps = 20 + static_dt = False + + # Compare with single-rate implicit. + + # {{{ single rate + + single_rate_method = AdamsMoultonMethodBuilder("y", order=order, + hist_length=hist_length) + single_rate_code = single_rate_method.generate() + + single_rate_code = replace_AssignImplicit(single_rate_code, + {"solve": am_solver_sr_hook}) + + def single_rate_rhs(t, y): + f, s = y + return np.array([ + ode.f2f_rhs(t, f, s)+ode.s2f_rhs(t, f, s), + ode.f2s_rhs(t, f, s)+ode.s2s_rhs(t, f, s), + ]) + + single_rate_interp = NumpyInterpreter( + single_rate_code, + function_map={"y": single_rate_rhs, + "solver": partial(am_solver_sr, single_rate_rhs)}) + + single_rate_interp.set_up(t_start=t_start, dt_start=dt, + context={"y": np.array([ + ode.soln_0(t_start), + ode.soln_1(t_start), + ])}) + + single_rate_values = {} + + for event in single_rate_interp.run(): + if isinstance(event, single_rate_interp.StateComputed): + single_rate_values[event.t] = event.state_component + + if len(single_rate_values) == nsteps: + break + + # }}} + + # {{{ single rate explicit + + single_rate_exp_method = AdamsBashforthMethodBuilder("y", order=order, + hist_length=hist_length, static_dt=static_dt) + single_rate_exp_code = single_rate_exp_method.generate() + + single_rate_exp_interp = NumpyInterpreter( + single_rate_exp_code, + function_map={"y": single_rate_rhs}) + + single_rate_exp_interp.set_up(t_start=t_start, dt_start=dt, + context={"y": np.array([ + ode.soln_0(t_start), + ode.soln_1(t_start), + ])}) + + single_rate_exp_values = {} + + for event in single_rate_exp_interp.run(): + if isinstance(event, single_rate_interp.StateComputed): + single_rate_exp_values[event.t] = event.state_component + + if len(single_rate_exp_values) == nsteps: + break + + # }}} + + # {{{ two rate + + multi_rate_method = MultiRateMultiStepMethodBuilder( + order, + ( + ( + 'dt', 'fast', '=', + MRHistory(1, "f", ("fast", "slow",), + hist_length=hist_length, + is_rhs_implicit=True), + ), + ( + 'dt', 'slow', '=', + MRHistory(1, "s", ("fast", "slow",), + rhs_policy=rhs_policy.late, hist_length=hist_length, + is_rhs_implicit=False), + ),), + static_dt=static_dt, + hist_consistency_threshold=1e-8, + early_hist_consistency_threshold=dt**order) + + multi_rate_code = multi_rate_method.generate() + + multi_rate_code = replace_AssignImplicit(multi_rate_code, + {"solve": am_solver_hook}) + + def rhs_fast(t, fast, slow): + return ode.f2f_rhs(t, fast, slow)+ode.s2f_rhs(t, fast, slow) + + def rhs_slow(t, fast, slow): + return ode.f2s_rhs(t, fast, slow)+ode.s2s_rhs(t, fast, slow) + + multi_rate_interp = NumpyInterpreter( + multi_rate_code, + function_map={"f": rhs_fast, "s": rhs_slow, + "solver": partial(am_solver, rhs_fast)}) + + multi_rate_interp.set_up(t_start=t_start, dt_start=dt, + context={ + "fast": ode.soln_0(t_start), + "slow": ode.soln_1(t_start), + }) + + multi_rate_values = {} + + for event in multi_rate_interp.run(): + if isinstance(event, single_rate_interp.StateComputed): + idx = {"fast": 0, "slow": 1}[event.component_id] + if event.t not in multi_rate_values: + multi_rate_values[event.t] = [None, None] + + multi_rate_values[event.t][idx] = event.state_component + + if len(multi_rate_values) > nsteps: + break + + # }}} + + times = sorted(single_rate_values) + single_rate_values = np.array([single_rate_values[t] for t in times]) + multi_rate_values = np.array([multi_rate_values[t] for t in times]) + single_rate_exp_values = np.array([single_rate_exp_values[t] for t in times]) + print(multi_rate_values) + print(single_rate_values) + print(single_rate_exp_values) + + single_rate_test = np.zeros((len(times), 2)) + for i in range(0, len(times)): + single_rate_test[i][0] = single_rate_values[i][0] + single_rate_test[i][1] = single_rate_exp_values[i][1] + + diff = la.norm((single_rate_values-multi_rate_values).reshape(-1)) + + print(diff) + assert diff < 1e-13 + + +@pytest.mark.parametrize(("order", "hist_length"), [(3, 3), + (3, 4), (4, 4), (4, 5)]) +@pytest.mark.parametrize("step_ratio", [1, 2, 3, 4]) +@pytest.mark.parametrize("system", ["Full", "StiffUncoupled"]) +@pytest.mark.parametrize("static_dt", [True, False]) +def test_implicit_accuracy(order, hist_length, system, static_dt, step_ratio, + plotting=False): + from dagrt.exec_numpy import NumpyInterpreter + from leap.implicit import replace_AssignImplicit + from functools import partial + + import multirate_test_systems + + test_system = getattr(multirate_test_systems, system) + ode = test_system() + + """Check that the implicit multirate timestepper has the advertised accuracy""" + + def true_f(t): + return ode.soln_0(t) + + def true_s(t): + return ode.soln_1(t) + + # {{{ construct multi-rate IMEX method + + # FIXME: no history consistency thresholding for now. + multi_rate_method = MultiRateMultiStepMethodBuilder( + order, + ( + ( + 'dt', 'fast', '=', + MRHistory(1, "f", ("fast", "slow",), + hist_length=hist_length, + is_rhs_implicit=True), + ), + ( + 'dt', 'slow', '=', + MRHistory(step_ratio, "s", ("fast", "slow",), + rhs_policy=rhs_policy.late, hist_length=hist_length, + is_rhs_implicit=False), + ),), + static_dt=static_dt) + + multi_rate_code = multi_rate_method.generate() + + multi_rate_code = replace_AssignImplicit(multi_rate_code, + {"solve": am_solver_hook}) + + def rhs_fast(t, fast, slow): + return ode.f2f_rhs(t, fast, slow)+ode.s2f_rhs(t, fast, slow) + + def rhs_slow(t, fast, slow): + return ode.f2s_rhs(t, fast, slow)+ode.s2s_rhs(t, fast, slow) + + # {{{ convergence run(s) + from pytools.convergence import EOCRecorder + eocrec = EOCRecorder() + + for n in range(4, 7): + t_start = 0 + dt = 2**(-n) + final_t = 2 + + multi_rate_interp = NumpyInterpreter( + multi_rate_code, + function_map={"f": rhs_fast, "s": rhs_slow, + "solver": partial(am_solver, rhs_fast)}) + + multi_rate_interp.set_up(t_start=t_start, dt_start=dt, + context={ + "fast": ode.soln_0(t_start), + "slow": ode.soln_1(t_start), + }) + + f_times = [] + f_values = [] + s_times = [] + s_values = [] + for event in multi_rate_interp.run(t_end=final_t): + if isinstance(event, multi_rate_interp.StateComputed): + if event.component_id == "fast": + f_times.append(event.t) + f_values.append(event.state_component[0]) + elif event.component_id == "slow": + s_times.append(event.t) + # FIXME: + if system == "Full": + s_values.append(event.state_component[0]) + else: + s_values.append(event.state_component) + else: + assert False, event.component_id + + f_times = np.array(f_times) + s_times = np.array(s_times) + f_values_true = true_f(f_times) + s_values_true = true_s(s_times) + f_err = f_values - f_values_true + s_err = s_values - s_values_true + + if plotting: + import matplotlib.pyplot as plt + plt.plot(f_times, f_values, label='IMEX') + plt.plot(f_times, f_values_true, label='True') + plt.title('Stiff Uncoupled System - Fast Component Evolution') + plt.xlabel('t') + plt.ylabel('f') + plt.legend() + plt.savefig('imex_comparison_sr2.png') + plt.clf() + + plt.plot(f_times, f_err, label='IMEX') + plt.title('Stiff Uncoupled System - Fast Component Error') + plt.xlabel('t') + plt.ylabel('f_err') + plt.savefig('imex_error_sr2.png') + plt.clf() + + error = ( + la.norm(f_err) / la.norm(f_values_true) + + # noqa: W504 + la.norm(s_err) / la.norm(s_values_true)) + + eocrec.add_data_point(dt, error) + + print(eocrec.pretty_print()) + + orderest = eocrec.estimate_order_of_convergence()[0, 1] + assert orderest > order * 0.7 + + # }}} + + @pytest.mark.parametrize("method_name", ["F", "Fqsr", "Srsf", "S"]) def test_2rab_scheme_explainers(method_name, order=3, step_ratio=3, explainer=TextualSchemeExplainer()): -- GitLab From c1b563c32b3c5d5bf877185ab7b98bb78676bbca Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 5 Aug 2020 14:57:54 -0500 Subject: [PATCH 08/14] Make minimal changes required to combine bootstrap step functions for IMEX and fully explicit --- leap/multistep/multirate/__init__.py | 299 ++++++--------------------- 1 file changed, 59 insertions(+), 240 deletions(-) diff --git a/leap/multistep/multirate/__init__.py b/leap/multistep/multirate/__init__.py index 4757f91..f651805 100644 --- a/leap/multistep/multirate/__init__.py +++ b/leap/multistep/multirate/__init__.py @@ -421,230 +421,54 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): cb(self.bootstrap_step, 0) - # {{{ rk bootstrap: fully explicit step + # {{{ rk bootstrap: step - def emit_small_rk_step(self, cb, name_prefix, name_gen, rhss_on_entry): - """Emit a single step of an RK method.""" - - from leap.rk import ORDER_TO_RK_METHOD_BUILDER - rk_method = ORDER_TO_RK_METHOD_BUILDER[self.max_order] - rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit)) - rk_coeffs = rk_method.output_coeffs - - def make_stage_history(prefix): - return [var(prefix + "_stage" + str(i)) for i in range(len(rk_tableau))] - - stage_rhss = {} - for comp_name, component_rhss in zip(self.component_names, self.rhss): - if not self.is_ode_component[comp_name]: - continue - - for irhs, rhs in enumerate(component_rhss): - stage_rhss[comp_name, irhs] = make_stage_history( - "{name_prefix}_rk_{comp_name}_rhs{irhs}" - .format( - name_prefix=name_prefix, - comp_name=comp_name, - irhs=irhs)) - - for istage, (c, coeffs) in enumerate(rk_tableau): - if len(coeffs) == 0: - assert c == 0 - for comp_name, component_rhss in zip( - self.component_names, self.rhss): - if not self.is_ode_component[comp_name]: - continue - - for irhs, rhs in enumerate(component_rhss): - cb(stage_rhss[comp_name, irhs][istage], - rhss_on_entry[comp_name, irhs]) - - else: - component_state_ests = {} - - for icomp, (comp_name, component_rhss) in enumerate( - zip(self.component_names, self.rhss)): - - if not self.is_ode_component[comp_name]: - continue - - contribs = [] - for irhs, rhs in enumerate(component_rhss): - state_contrib_var = var( - name_gen( - "state_contrib_{comp_name}_rhs{irhs}" - .format(comp_name=comp_name, irhs=irhs))) - - contribs.append(state_contrib_var) - - cb(state_contrib_var, - _linear_comb(coeffs, stage_rhss[comp_name, irhs])) - - state_var = var( - name_gen( - "state_{comp_name}_st{istage}" - .format(comp_name=comp_name, istage=istage))) - - state_expr = ( - var("" + comp_name) - + (self.dt/self.nsubsteps) * sum(contribs)) - if comp_name in self.state_filters: - state_expr = self.state_filters[comp_name](state_expr) - - cb(state_var, state_expr) - - component_state_ests[comp_name] = state_var - - # At this point, we have all the ODE state estimates evaluated. - - # {{{ evaluate the non-ODE RHSs - - for comp_name, component_rhss in zip( - self.component_names, self.rhss): - if self.is_ode_component[comp_name]: - continue - - contribs = [] - - for irhs, rhs in enumerate(component_rhss): - kwargs = dict( - (self.comp_name_to_kwarg_name[arg_comp_name], - component_state_ests[arg_comp_name]) - for arg_comp_name in rhs.arguments) - - contribs.append(var(rhs.func_name)( - t=self.t + (c/self.nsubsteps) * self.dt, - **kwargs)) - - state_var = var( - name_gen( - "state_{comp_name}_st{istage}" - .format(comp_name=comp_name, istage=istage))) - - cb(state_var, sum(contribs)) - - component_state_ests[comp_name] = state_var - - # }}} - - # {{{ evaluate the ODE RHSs - - for comp_name, component_rhss in zip( - self.component_names, self.rhss): - - if not self.is_ode_component[comp_name]: - continue - - for irhs, rhs in enumerate(component_rhss): - kwargs = dict( - (self.comp_name_to_kwarg_name[arg_comp_name], - component_state_ests[arg_comp_name]) - for arg_comp_name in rhs.arguments) - cb(stage_rhss[comp_name, irhs][istage], - var(rhs.func_name)( - t=self.t + (c/self.nsubsteps) * self.dt, - **kwargs)) - - # }}} - - component_state_ests = {} - - for icomp, (comp_name, component_rhss) in enumerate( - zip(self.component_names, self.rhss)): - - contribs = [] - for irhs, rhs in enumerate(component_rhss): - if not self.is_ode_component[comp_name]: - continue - - state_contrib_var = var( - name_gen( - "state_contrib_{comp_name}_rhs{irhs}" - .format(comp_name=comp_name, irhs=irhs))) - - contribs.append(state_contrib_var) - - cb(state_contrib_var, - _linear_comb(rk_coeffs, stage_rhss[comp_name, irhs])) - - state_var = var( - name_gen( - "state_{comp_name}_final" - .format(comp_name=comp_name))) - - state_expr = ( - var("" + comp_name) - + (self.dt/self.nsubsteps) * sum(contribs)) - if comp_name in self.state_filters: - state_expr = self.state_filters[comp_name](state_expr) - - cb(state_var, state_expr) - - component_state_ests[comp_name] = state_var - - for component_name in self.component_names: - state = component_state_ests[component_name] - if self.is_ode_component[component_name]: - cb.yield_state( - state, - component_name, self.t + self.dt/self.nsubsteps, - "bootstrap") - - cb(var(""+component_name), state) - - cb(self.t, self.t + self.dt/self.nsubsteps) - - # }}} - - # {{{ rk bootstrap: IMEX step - - def emit_small_imex_rk_step(self, cb, name_prefix, name_gen, rhss_on_entry): + def emit_small_rk_step(self, cb, name_prefix, name_gen, rhss_on_entry, + imex=False): """Emit a single step of an IMEX-RK method.""" - from leap.rk import (IMEX_ORDER_TO_RK_METHOD_BUILDER, - _is_first_stage_same_as_last_stage) - rk_method = IMEX_ORDER_TO_RK_METHOD_BUILDER[self.max_order] - rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit, - rk_method.a_implicit)) - rk_coeffs = rk_method.high_order_coeffs - stage_coeff_sets = { - "explicit": rk_method.a_explicit, - "implicit": rk_method.a_implicit} + if imex: + from leap.rk import (IMEX_ORDER_TO_RK_METHOD_BUILDER, + _is_first_stage_same_as_last_stage) + rk_method = IMEX_ORDER_TO_RK_METHOD_BUILDER[self.max_order] + rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit, + rk_method.a_implicit)) + rk_coeffs = rk_method.high_order_coeffs + nstages = len(rk_method.c) + stage_coeff_sets = { + "explicit": rk_method.a_explicit, + "implicit": rk_method.a_implicit} + else: + from leap.rk import ORDER_TO_RK_METHOD_BUILDER + rk_method = ORDER_TO_RK_METHOD_BUILDER[self.max_order] + rk_tableau = tuple(zip(rk_method.c, rk_method.a_explicit)) + rk_coeffs = rk_method.output_coeffs + nstages = len(rk_method.c) + stage_coeff_sets = { + "explicit": rk_method.a_explicit} def make_stage_history(prefix): return [var(prefix + "_stage" + str(i)) for i in range(len(rk_tableau))] - def make_imp_stage_history(prefix): - return [var(prefix + "_stage" + str(i) + "_imp") - for i in range(len(rk_tableau))] - stage_rhss = {} - imp_stage_rhss = {} for comp_name, component_rhss in zip(self.component_names, self.rhss): if not self.is_ode_component[comp_name]: continue for irhs, rhs in enumerate(component_rhss): - if rhs.is_implicit: - imp_stage_rhss[comp_name, irhs] = make_imp_stage_history( - "{name_prefix}_rk_{comp_name}_rhs{irhs}" - .format( - name_prefix=name_prefix, - comp_name=comp_name, - irhs=irhs)) - else: - stage_rhss[comp_name, irhs] = make_stage_history( - "{name_prefix}_rk_{comp_name}_rhs{irhs}" - .format( - name_prefix=name_prefix, - comp_name=comp_name, - irhs=irhs)) + stage_rhss[comp_name, irhs] = make_stage_history( + "{name_prefix}_rk_{comp_name}_rhs{irhs}" + .format( + name_prefix=name_prefix, + comp_name=comp_name, + irhs=irhs)) - for istage, (c, exp_coeffs, imp_coeffs) in enumerate(rk_tableau): + for istage in range(nstages): unknowns = set() knowns = set() component_state_ests = {} + c = rk_method.c[istage] for comp_name, component_rhss in zip( self.component_names, self.rhss): if not self.is_ode_component[comp_name]: @@ -657,13 +481,14 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): name = "implicit" else: name = "explicit" + coeffs = stage_coeff_sets[name][istage] if ( name in rk_method.recycle_last_stage_coeff_set_names and istage == 0 and _is_first_stage_same_as_last_stage( rk_method.c, stage_coeff_sets[name])): if rhs.is_implicit: - unknowns.add(imp_stage_rhss[comp_name, irhs][istage]) + unknowns.add(stage_rhss[comp_name, irhs][istage]) else: cb(stage_rhss[comp_name, irhs][istage], rhss_on_entry[comp_name, irhs]) @@ -674,18 +499,18 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): .format(comp_name=comp_name, irhs=irhs))) if rhs.is_implicit: - unknowns.add(imp_stage_rhss[comp_name, irhs][istage]) - contribs.append(_linear_comb(imp_coeffs, - imp_stage_rhss[comp_name, irhs])) + unknowns.add(stage_rhss[comp_name, irhs][istage]) + contribs.append(_linear_comb(coeffs, + stage_rhss[comp_name, irhs])) else: # First stage. - if exp_coeffs == []: + if len(coeffs) == 0: cb(state_contrib_var, 0) else: - contribs.append(_linear_comb(exp_coeffs, + contribs.append(_linear_comb(coeffs, stage_rhss[comp_name, irhs])) cb(state_contrib_var, - _linear_comb(exp_coeffs, + _linear_comb(coeffs, stage_rhss[comp_name, irhs])) state_var = var( @@ -752,10 +577,16 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): (self.comp_name_to_kwarg_name[arg_comp_name], component_state_ests[arg_comp_name]) for arg_comp_name in rhs.arguments) + # We still need to go through the whole solving deal here. + if rhs.is_implicit: + name = "implicit" + else: + name = "explicit" + coeffs = stage_coeff_sets[name][istage] if rhs.is_implicit: # First stage: no solve required. - if imp_coeffs == []: - cb(imp_stage_rhss[comp_name, irhs][istage], + if len(coeffs) == 0: + cb(stage_rhss[comp_name, irhs][istage], var(rhs.func_name)( t=self.t + (c/self.nsubsteps) * self.dt, **kwargs)) @@ -765,9 +596,9 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): # RHS as the stage number. unkvar = cb.fresh_var('unk_s%d' % (istage)) rhs_var_to_unknown[ - imp_stage_rhss[comp_name, + stage_rhss[comp_name, irhs][istage]] = unkvar - solve_expression = imp_stage_rhss[ + solve_expression = stage_rhss[ comp_name, irhs][istage] - var(rhs.func_name)( t=self.t + (c/self.nsubsteps) * self.dt, **kwargs) @@ -810,8 +641,6 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): # }}} - # Now that we've traversed both stage loops, - # we can actually sum the RHS's to get the state estimates. component_state_ests = {} for icomp, (comp_name, component_rhss) in enumerate( @@ -822,27 +651,15 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): if not self.is_ode_component[comp_name]: continue - # Implicit, if applicable: - if rhs.is_implicit: - imp_state_contrib_var = var( - name_gen( - "state_contrib_{comp_name}_rhs{irhs}_imp" - .format(comp_name=comp_name, irhs=irhs))) - - contribs.append(imp_state_contrib_var) - - cb(imp_state_contrib_var, - _linear_comb(rk_coeffs, imp_stage_rhss[comp_name, irhs])) - else: - state_contrib_var = var( - name_gen( - "state_contrib_{comp_name}_rhs{irhs}" - .format(comp_name=comp_name, irhs=irhs))) + state_contrib_var = var( + name_gen( + "state_contrib_{comp_name}_rhs{irhs}" + .format(comp_name=comp_name, irhs=irhs))) - contribs.append(state_contrib_var) + contribs.append(state_contrib_var) - cb(state_contrib_var, - _linear_comb(rk_coeffs, stage_rhss[comp_name, irhs])) + cb(state_contrib_var, + _linear_comb(rk_coeffs, stage_rhss[comp_name, irhs])) state_var = var( name_gen( @@ -1006,10 +823,12 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): if imp_exist == 0: # If we have no implicit RHSs, run the full explicit bootstrapper. - self.emit_small_rk_step(cb, name_prefix, name_gen, current_rhss) + self.emit_small_rk_step(cb, name_prefix, name_gen, + current_rhss, imex=False) else: # Otherwise, run the IMEX RK bootstrapper. - self.emit_small_imex_rk_step(cb, name_prefix, name_gen, current_rhss) + self.emit_small_rk_step(cb, name_prefix, name_gen, + current_rhss, imex=True) cb(self.bootstrap_step, self.bootstrap_step + 1) -- GitLab From 47d9b738beaae6d51d717e68aef1f6d884ae35de Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Wed, 5 Aug 2020 15:02:32 -0500 Subject: [PATCH 09/14] Minor copyright/docstring updates --- leap/multistep/__init__.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index 593d293..8f0ab86 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -1,11 +1,11 @@ -"""Adams-Bashforth ODE solvers.""" +"""Adams-Bashforth and Adams-Moulton ODE solvers.""" from __future__ import division __copyright__ = """ Copyright (C) 2007 Andreas Kloeckner Copyright (C) 2014, 2015 Matt Wala -Copyright (C) 2015 Cory Mikida +Copyright (C) 2015, 2020 Cory Mikida """ __license__ = """ @@ -37,6 +37,7 @@ from pymbolic import var __doc__ = """ .. autoclass:: AdamsBashforthMethodBuilder +.. autoclass:: AdamsMoultonMethodBuilder """ -- GitLab From 4b8c9bfc874a287cb0f38aba0f6f9e495216474a Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Fri, 7 Aug 2020 15:40:19 -0500 Subject: [PATCH 10/14] Placate flake8 by removing needless state update --- leap/multistep/multirate/__init__.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/leap/multistep/multirate/__init__.py b/leap/multistep/multirate/__init__.py index f651805..ebe695d 100644 --- a/leap/multistep/multirate/__init__.py +++ b/leap/multistep/multirate/__init__.py @@ -1118,11 +1118,6 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): knowns.update(unknowns) unknowns.clear() - # Now that we have the RHS value solved for, - # update useful temp states - for arg_comp_name in rhs.arguments: - post_solve_state = get_state(arg_comp_name, isubstep) - temp_hist_substeps[comp_name, irhs].append(isubstep) if not rhs.is_implicit: -- GitLab From dd7990a310155aaef96f5e249ef20f7128be7454 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Fri, 7 Aug 2020 15:40:49 -0500 Subject: [PATCH 11/14] Add extra bootstrap option to AM to get implicit single-rate identical test to pass --- leap/multistep/__init__.py | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index 8f0ab86..8305715 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -420,7 +420,7 @@ class AdamsMoultonMethodBuilder(MethodBuilder): """ def __init__(self, component_id, function_family=None, state_filter_name=None, - hist_length=None, static_dt=False, order=None): + hist_length=None, static_dt=False, order=None, _extra_bootstrap=False): """ :arg function_family: Accepts an instance of :class:`ABIntegrationFunctionFamily` @@ -448,6 +448,7 @@ class AdamsMoultonMethodBuilder(MethodBuilder): self.hist_length = hist_length self.static_dt = static_dt + self.extra_bootstrap = _extra_bootstrap self.component_id = component_id @@ -606,9 +607,15 @@ class AdamsMoultonMethodBuilder(MethodBuilder): component_id=self.component_id, time_id='', time=self.t) cb_bootstrap(self.step, self.step + 1) - # Bootstrap length is one less because of implicit. - with cb_bootstrap.if_(self.step, "==", self.hist_length - 1): - cb_bootstrap.switch_phase("primary") + # Bootstrap length is typically one less because of implicit, + # but if we are comparing with IMEX MRAM, we need one more + # bootstrap step. + if self.extra_bootstrap: + with cb_bootstrap.if_(self.step, "==", self.hist_length): + cb_bootstrap.switch_phase("primary") + else: + with cb_bootstrap.if_(self.step, "==", self.hist_length - 1): + cb_bootstrap.switch_phase("primary") return DAGCode( phases={ @@ -639,7 +646,12 @@ class AdamsMoultonMethodBuilder(MethodBuilder): rk_tableau = tuple(zip(rk_method.c, rk_method.a_implicit)) rk_coeffs = rk_method.output_coeffs - with cb.if_(self.step, "==", 1): + if self.extra_bootstrap: + first_save_step = 2 + else: + first_save_step = 1 + + with cb.if_(self.step, "==", first_save_step): # Save the first RHS to the AM history rhs_var = var("rhs_var") @@ -718,13 +730,17 @@ class AdamsMoultonMethodBuilder(MethodBuilder): cb(rhs_next_var, self.eval_rhs(self.t + self.dt, self.state)) for i in range(1, len(self.history)): - with cb.if_(self.step, "==", i): + if self.extra_bootstrap: + save_crit = i+1 + else: + save_crit = i + + with cb.if_(self.step, "==", save_crit): cb(self.history[i], rhs_next_var) if not self.static_dt: cb(self.time_history[i], self.t + self.dt) - # }}} # vim: fdm=marker -- GitLab From 694590908fdc80eff4e339b4a71f021847e27798 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Fri, 7 Aug 2020 15:41:41 -0500 Subject: [PATCH 12/14] Add extra bootstrap option to single-rate identical test, tweak tolerance --- test/test_multirate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_multirate.py b/test/test_multirate.py index 55c7841..c247055 100644 --- a/test/test_multirate.py +++ b/test/test_multirate.py @@ -373,7 +373,7 @@ def test_implicit_single_rate_identical(order=3, hist_length=3): # {{{ single rate single_rate_method = AdamsMoultonMethodBuilder("y", order=order, - hist_length=hist_length) + hist_length=hist_length, _extra_bootstrap=True) single_rate_code = single_rate_method.generate() single_rate_code = replace_AssignImplicit(single_rate_code, @@ -509,7 +509,7 @@ def test_implicit_single_rate_identical(order=3, hist_length=3): diff = la.norm((single_rate_values-multi_rate_values).reshape(-1)) print(diff) - assert diff < 1e-13 + assert diff < 1e-12 @pytest.mark.parametrize(("order", "hist_length"), [(3, 3), -- GitLab From b066fbcb32db80b4823422d007bf38c7bc281458 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Fri, 7 Aug 2020 16:03:24 -0500 Subject: [PATCH 13/14] PyLint fix (whoopsie) --- leap/rk/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/leap/rk/__init__.py b/leap/rk/__init__.py index 2e6dca7..3ee144e 100644 --- a/leap/rk/__init__.py +++ b/leap/rk/__init__.py @@ -399,7 +399,7 @@ class SimpleButcherTableauMethodBuilder(ButcherTableauMethodBuilder): class ImplicitButcherTableauMethodBuilder(ButcherTableauMethodBuilder): def __init__(self, component_id, state_filter_name=None, rhs_func_name=None): - super(SimpleButcherTableauMethodBuilder, self).__init__( + super(ImplicitButcherTableauMethodBuilder, self).__init__( component_id=component_id, state_filter_name=state_filter_name) -- GitLab From 266d87d29336767090bf5fc4fd54cc3ae4277d52 Mon Sep 17 00:00:00 2001 From: Cory Mikida Date: Fri, 7 Aug 2020 16:17:53 -0500 Subject: [PATCH 14/14] More PyLint fixes --- leap/rk/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/leap/rk/__init__.py b/leap/rk/__init__.py index 3ee144e..de70f93 100644 --- a/leap/rk/__init__.py +++ b/leap/rk/__init__.py @@ -115,6 +115,10 @@ class ButcherTableauMethodBuilder(MethodBuilder): def a_explicit(self): raise NotImplementedError + @property + def a_implicit(self): + raise NotImplementedError + @property def output_coeffs(self): raise NotImplementedError -- GitLab