diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index e357a47d6742adf5a4e2029487861d0743215a6f..8305715617f8d536d562db21b025f3c140e0c9b8 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 """ @@ -405,4 +406,341 @@ 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, _extra_bootstrap=False): + """ + :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.extra_bootstrap = _extra_bootstrap + + self.component_id = component_id + + # Declare variables + self.step = var('

step') + self.function = var('' + component_id) + self.history = \ + [var('

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

t_n_minus_' + str(i)) + for i in range(hist_length - 1, 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_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 + # include the *next* point in time. + if not self.static_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): + 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() + + unknowns.add(rhs_next_var) + + # Update history + history = self.history + [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 + from pymbolic.mapper.distributor import DistributeMapper as DistMap + solve_expression = collapse_constants( + 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) + + # {{{ 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. + for i in range(self.hist_length - 1): + 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 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={ + "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 IMPLICIT RK method.""" + + equations = [] + unknowns = set() + knowns = set() + rhs_var_to_unknown = {} + + 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 + + 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") + + cb(rhs_var, self.eval_rhs(self.t, self.state)) + cb(self.history[0], rhs_var) + + if not self.static_dt: + cb(self.time_history[0], self.t) + + # Stage loop + rhss = [var("rk_rhs_" + str(i)) for i in range(len(rk_tableau))] + for stage_num, (c, coeffs) in enumerate(rk_tableau): + 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) + + # 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)) + + 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) + + # 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)): + 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 diff --git a/leap/multistep/multirate/__init__.py b/leap/multistep/multirate/__init__.py index a8481f19e6db212dc5ae81a666a7098eb29b0310..ebe695d82ce65888730758ea905992708762f213 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]_. @@ -415,13 +423,29 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): # {{{ 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 emit_small_rk_step(self, cb, name_prefix, name_gen, rhss_on_entry, + imex=False): + """Emit a single step of an IMEX-RK method.""" + + 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))] @@ -433,111 +457,189 @@ class MultiRateMultiStepMethodBuilder(MethodBuilder): 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]) + "{name_prefix}_rk_{comp_name}_rhs{irhs}" + .format( + name_prefix=name_prefix, + comp_name=comp_name, + irhs=irhs)) - else: - component_state_ests = {} + for istage in range(nstages): - for icomp, (comp_name, component_rhss) in enumerate( - zip(self.component_names, self.rhss)): - - if not self.is_ode_component[comp_name]: - continue + 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]: + continue - contribs = [] - for irhs, rhs in enumerate(component_rhss): + 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" + 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(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))) - contribs.append(state_contrib_var) + if rhs.is_implicit: + unknowns.add(stage_rhss[comp_name, irhs][istage]) + contribs.append(_linear_comb(coeffs, + stage_rhss[comp_name, irhs])) + else: + # First stage. + if len(coeffs) == 0: + cb(state_contrib_var, 0) + else: + contribs.append(_linear_comb(coeffs, + stage_rhss[comp_name, irhs])) + 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))) - cb(state_contrib_var, - _linear_comb(coeffs, stage_rhss[comp_name, irhs])) + state_expr = ( + var("" + comp_name) + + (self.dt/self.nsubsteps) * sum(contribs)) - state_var = var( - name_gen( - "state_{comp_name}_st{istage}" - .format(comp_name=comp_name, istage=istage))) + if comp_name in self.state_filters: + state_expr = self.state_filters[comp_name](state_expr) - 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) - cb(state_var, state_expr) + # At this point, we have all the ODE state estimates evaluated. - component_state_ests[comp_name] = state_var + # {{{ evaluate the non-ODE RHSs - # At this point, we have all the ODE state estimates evaluated. + for comp_name, component_rhss in zip( + self.component_names, self.rhss): + if self.is_ode_component[comp_name]: + continue - # {{{ evaluate the non-ODE RHSs + contribs = [] - for comp_name, component_rhss in zip( - self.component_names, self.rhss): - if 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) - contribs = [] + contribs.append(var(rhs.func_name)( + t=self.t + (c/self.nsubsteps) * self.dt, + **kwargs)) - 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) + state_var = var( + name_gen( + "state_{comp_name}_st{istage}" + .format(comp_name=comp_name, istage=istage))) - contribs.append(var(rhs.func_name)( - t=self.t + (c/self.nsubsteps) * self.dt, - **kwargs)) + cb(state_var, sum(contribs)) - state_var = var( - name_gen( - "state_{comp_name}_st{istage}" - .format(comp_name=comp_name, istage=istage))) + component_state_ests[comp_name] = state_var - 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): - # {{{ evaluate the ODE RHSs + if not self.is_ode_component[comp_name]: + continue - for comp_name, component_rhss in zip( - self.component_names, self.rhss): + equations = [] + rhs_var_to_unknown = {} - 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) + # 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 len(coeffs) == 0: + cb(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[ + stage_rhss[comp_name, + irhs][istage]] = unkvar + solve_expression = 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() - 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) + else: cb(stage_rhss[comp_name, irhs][istage], var(rhs.func_name)( t=self.t + (c/self.nsubsteps) * self.dt, **kwargs)) - # }}} + # }}} component_state_ests = {} @@ -711,7 +813,22 @@ 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, imex=False) + else: + # Otherwise, run the IMEX RK bootstrapper. + self.emit_small_rk_step(cb, name_prefix, name_gen, + current_rhss, imex=True) cb(self.bootstrap_step, self.bootstrap_step + 1) @@ -724,7 +841,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 +899,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 +982,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 +1030,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 +1064,80 @@ 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() + + 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 +1326,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/leap/rk/__init__.py b/leap/rk/__init__.py index e11393c76814bb51feec3c45d5caa2fb24159b71..de70f9331e98a3055388627730053f3304a3d16a 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 @@ -396,6 +400,32 @@ class SimpleButcherTableauMethodBuilder(ButcherTableauMethodBuilder): }) +class ImplicitButcherTableauMethodBuilder(ButcherTableauMethodBuilder): + def __init__(self, component_id, state_filter_name=None, + rhs_func_name=None): + super(ImplicitButcherTableauMethodBuilder, 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 +447,9 @@ class BackwardEulerMethodBuilder(SimpleButcherTableauMethodBuilder): .. automethod:: __init__ .. automethod:: generate """ - c = (0,) + c = (1,) - a_explicit = ( + a_implicit = ( (1,), ) @@ -524,6 +554,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 +665,22 @@ ORDER_TO_RK_METHOD_BUILDER = { 5: RK5MethodBuilder, } +IMPLICIT_ORDER_TO_RK_METHOD_BUILDER = { + 1: BackwardEulerMethodBuilder, + 2: DIRK2MethodBuilder, + 3: DIRK3MethodBuilder, + 4: DIRK4MethodBuilder, + 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 b4884b964ae986aac10e67e59a3c45df7b7a2814..b3e24546f5759a005eb7461f9d2335928c61fcaa 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)) diff --git a/test/test_implicit_ab.py b/test/test_implicit_ab.py new file mode 100755 index 0000000000000000000000000000000000000000..704be320df024a4b9f68161139d0c8e901b327ed --- /dev/null +++ b/test/test_implicit_ab.py @@ -0,0 +1,128 @@ +#! /usr/bin/env python +from __future__ import division, with_statement + +__copyright__ = "Copyright (C) 2020 Cory Mikida" + +__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) + + +@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] + for static_dt in [True, False] + ] + [ + (AdamsMoultonMethodBuilder("y", order, hist_length=order+1, + static_dt=static_dt), order, static_dt) + 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): + 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 + 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 + 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__]) diff --git a/test/test_multirate.py b/test/test_multirate.py index 045ce075358ed2d11a0734299aa2b8f281ce0a73..c2470557b74e4cf53eaf4bb9df8d7fe16a7f1998 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, _extra_bootstrap=True) + 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-12 + + +@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()):