diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d2e24ec..a315f31 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -53,7 +53,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: "Main Script" run: | - EXTRA_INSTALL="numpy" + EXTRA_INSTALL="numpy scipy" sudo apt update sudo apt install gfortran-7 liblapack-dev libblas-dev sudo ln -sf /usr/bin/gfortran-7 /usr/bin/gfortran diff --git a/leap/implicit.py b/leap/implicit.py index bc9d008..a61dea4 100644 --- a/leap/implicit.py +++ b/leap/implicit.py @@ -86,3 +86,27 @@ def solver_hook(expr, var, id, **kwargs): new_phases[phase_name] = phase.copy(statements=new_statements) return dag.copy(phases=new_phases) + + +def generate_solve(cb, unknowns, equations, var_to_unknown, guess): + + assert len(unknowns) == len(equations) + assignees = [unk.name for unk in unknowns] + + from pymbolic import substitute + subst_dict = { + rhs_var.name: var_to_unknown[rhs_var] + for rhs_var in unknowns} + + cb.assign_implicit( + assignees=assignees, + solve_components=[ + var_to_unknown[unk].name + for unk in unknowns], + expressions=[ + substitute(eq, subst_dict) + for eq in equations], + + other_params={ + "guess": guess}, + solver_id="solve") diff --git a/leap/multistep/__init__.py b/leap/multistep/__init__.py index 660dbbc..6847d9a 100644 --- a/leap/multistep/__init__.py +++ b/leap/multistep/__init__.py @@ -4,7 +4,7 @@ __copyright__ = """ Copyright (C) 2007 Andreas Kloeckner Copyright (C) 2014, 2015 Matt Wala -Copyright (C) 2015 Cory Mikida +Copyright (C) 2015, 2020 Cory Mikida """ __license__ = """ @@ -36,7 +36,9 @@ __doc__ = """ .. autoclass:: AdamsIntegrationFunctionFamily .. autoclass:: AdamsMonomialIntegrationFunctionFamily +.. autoclass:: AdamsMethodBuilder .. autoclass:: AdamsBashforthMethodBuilder +.. autoclass:: AdamsMoultonMethodBuilder """ @@ -206,9 +208,10 @@ def emit_adams_extrapolation(cb, name_gen, # }}} -# {{{ ab method +# {{{ adams method + -class AdamsBashforthMethodBuilder(MethodBuilder): +class AdamsMethodBuilder(MethodBuilder): """ User-supplied context: + component_id: The value that is integrated @@ -265,6 +268,7 @@ def __init__(self, component_id, function_family=None, state_filter_name=None, self.t = var("") self.dt = var("
") + self.state_filter_name = state_filter_name if state_filter_name is not None: self.state_filter = var("" + state_filter_name) else: @@ -274,62 +278,16 @@ 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") - # Initialization with CodeBuilder(name="initialization") as cb_init: cb_init(self.step, 1) # Primary with CodeBuilder(name="primary") as cb_primary: - - if not self.static_dt: - time_history_data = self.time_history + [self.t] - 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+1, 0+1)) # noqa pylint:disable=invalid-unary-operand-type - dt_factor = self.dt - t_end = 1 - - cb_primary(rhs_var, self.eval_rhs(self.t, self.state)) - history = self.history + [rhs_var] - - ab_sum = emit_adams_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) - 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) + self.generate_primary(cb_primary) if self.hist_length == 1: # The first order method requires no bootstrapping. @@ -367,6 +325,100 @@ def eval_rhs(self, t, y): parameters=(), kw_parameters={"t": t, self.component_id: y}) + def rotate_and_yield(self, cb, hist, time_hist): + for i in range(self.hist_length - 1): + cb(self.history[i], hist[i + 1]) + + if not self.static_dt: + cb(self.time_history[i], time_hist[i + 1]) + + cb(self.t, self.t + self.dt) + cb.yield_state(expression=self.state, + component_id=self.component_id, + time_id="", time=self.t) + + def set_up_time_data(self, cb, new_t): + """Creates full snapshot of time points involved in the current + Adams step, synthesized from the existing time history along with + the next time point to be used (new_t). Returns time_data + (a list of the time values themselves) and relevant_times + (an *array* of those same time values centered around self.t + for use in emit_adams_integration). In the case of static timestep, + both of these are replaced with integer lists. + :arg new_t: the most recent time point that will be used + in performing the Adams integration step. This + should be either self.t (the current time, in the + case of Adams-Bashforth) or self.t + self.dt (the + next time point, in Adams-Moulton) + """ + from pytools import UniqueNameGenerator + name_gen = UniqueNameGenerator() + array = var("array") + if not self.static_dt: + time_data = self.time_history + [new_t] + time_data_var = var(name_gen("time_data")) + cb(time_data_var, array(self.hist_length)) + for i in range(self.hist_length): + cb(time_data_var[i], time_data[i] - self.t) + + relevant_times = time_data_var + t_end = self.dt + dt_factor = 1 + + else: + if new_t == self.t: + relevant_times = list(range(-self.hist_length+1, 0+1)) # noqa pylint:disable=invalid-unary-operand-type + time_data = list(range(-self.hist_length+1, 0+1)) # noqa pylint:disable=invalid-unary-operand-type + elif new_t == self.t + self.dt: + # In implicit mode, the vector of times + # passed to adams_integration must + # include the *next* point in time. + relevant_times = list(range(-self.hist_length+2, 0+2)) # noqa pylint:disable=invalid-unary-operand-type + time_data = list(range(-self.hist_length+2, 0+2)) # noqa pylint:disable=invalid-unary-operand-type + else: + raise ValueError("Invalid time point specified for Adams step") + dt_factor = self.dt + t_end = 1 + + return time_data, relevant_times, dt_factor, t_end + + def generate_primary(self, cb): + raise NotImplementedError() + + def rk_bootstrap(self, cb): + raise NotImplementedError() + +# }}} + + +# {{{ ab method + +class AdamsBashforthMethodBuilder(AdamsMethodBuilder): + def generate_primary(self, cb): + rhs_var = var("rhs_var") + from pytools import UniqueNameGenerator + name_gen = UniqueNameGenerator() + + time_history_data, time_hist, \ + dt_factor, t_end = self.set_up_time_data(cb, self.t) + + cb(rhs_var, self.eval_rhs(self.t, self.state)) + history = self.history + [rhs_var] + + ab_sum = emit_adams_integration( + cb, 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) + cb(self.state, state_est) + + # Rotate history and time history. + self.rotate_and_yield(cb, history, time_history_data) + def rk_bootstrap(self, cb): """Initialize the timestepper with an RK method.""" @@ -385,35 +437,151 @@ def rk_bootstrap(self, cb): 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_coeff_set_names = ("explicit",) + stage_coeff_sets = {"explicit": rk_method.a_explicit} + estimate_coeff_set_names = ("main",) + estimate_coeff_sets = {"main": rk_coeffs} + rhs_funcs = {"explicit": var(""+self.component_id)} + + # Traverse RK stage loop of appropriate order and update state. + rk = rk_method(self.component_id, self.state_filter_name) + cb = rk.generate_butcher_init(cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets) + cb, rhss, est_vars = rk.generate_butcher_primary(cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets) + + # Assign the value of the new state. + cb(self.state, est_vars[0]) + +# }}} - # 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) +# {{{ am method - 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)) +class AdamsMoultonMethodBuilder(AdamsMethodBuilder): + def generate_primary(self, cb): - state_est = self.state + self.dt * rk_comb + from pytools import UniqueNameGenerator + name_gen = UniqueNameGenerator() + rhs_next_var = var("rhs_next_var") + rhs_var_to_unknown = {} + unkvar = cb.fresh_var("unk") + rhs_var_to_unknown[rhs_next_var] = unkvar + + # In implicit mode, the vector of times + # passed to adams_integration must + # include the *next* point in time. + time_data, relevant_times, \ + dt_factor, t_end = self.set_up_time_data(cb, self.t + self.dt) + + # Implicit setup - rhs_next_var is an unknown, needs implicit solve. + equations = [] + unknowns = set() + knowns = set() + + unknowns.add(rhs_next_var) + + # Create RHS vector for Adams setup, + # including RHS value to be implicitly + # solved for + rhss = self.history + [rhs_next_var] + + # Set up the actual Adams-Moulton step. + am_sum = emit_adams_integration( + cb, name_gen, + self.function_family, + relevant_times, rhss, + 0, t_end) + + state_est = self.state + dt_factor * am_sum + + # 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.assign, cb.fresh_var) + equations.append(solve_expression) + + # {{{ emit solve if possible + + if unknowns and len(unknowns) == len(equations): + from leap.implicit import generate_solve + generate_solve(cb, unknowns, equations, rhs_var_to_unknown, self.state) + elif not unknowns: + raise ValueError("Adams-Moulton implicit timestep has no unknowns") + elif len(unknowns) > len(equations): + raise ValueError("Adams-Moulton implicit timestep has more unknowns " + "than equations") + elif len(unknowns) < len(equations): + raise ValueError("Adams-Moulton implicit timestep has more equations " + "than unknowns") + else: + raise AssertionError() + + del equations[:] + knowns.update(unknowns) + unknowns.clear() + + # }}} + + # Update the state now that we've solved. if self.state_filter is not None: state_est = self.state_filter(state_est) + cb(self.state, state_est) + + # Add new RHS and time to history and rotate. + history = self.history + [rhs_next_var] + self.rotate_and_yield(cb, history, time_data) + + def rk_bootstrap(self, cb): + """Initialize the timestepper with an IMPLICIT RK method.""" + + from leap.rk import IMPLICIT_ORDER_TO_RK_METHOD_BUILDER + rk_method = IMPLICIT_ORDER_TO_RK_METHOD_BUILDER[self.function_family.order] + rk_coeffs = rk_method.output_coeffs + stage_coeff_set_names = ("implicit",) + stage_coeff_sets = {"implicit": rk_method.a_implicit} + estimate_coeff_set_names = ("main",) + estimate_coeff_sets = {"main": rk_coeffs} + rhs_funcs = {"implicit": var(""+self.component_id)} + + # Traverse RK stage loop of appropriate order and update state. + rk = rk_method(self.component_id, self.state_filter_name) + cb = rk.generate_butcher_init(cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets) + cb, rhss, est_vars = rk.generate_butcher_primary(cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets) # Assign the value of the new state. - cb(self.state, state_est) + cb(self.state, est_vars[0]) + + # 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(len(self.history)): + + with cb.if_(self.step, "==", i + 1): + 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/rk/__init__.py b/leap/rk/__init__.py index 9b5d254..f4c1b7b 100644 --- a/leap/rk/__init__.py +++ b/leap/rk/__init__.py @@ -4,6 +4,7 @@ __copyright__ = """ Copyright (C) 2007-2013 Andreas Kloeckner Copyright (C) 2014, 2015 Matt Wala +Copyright (C) 2020 Cory Mikida """ __license__ = """ @@ -39,8 +40,12 @@ A dictionary mapping desired order of accuracy to a corresponding RK method builder. +.. data:: IMPLICIT_ORDER_TO_RK_METHOD_BUILDER + + A dictionary mapping desired order of accuracy to a corresponding implicit + RK method builder. + .. autoclass:: ForwardEulerMethodBuilder -.. autoclass:: BackwardEulerMethodBuilder .. autoclass:: MidpointMethodBuilder .. autoclass:: HeunsMethodBuilder .. autoclass:: RK3MethodBuilder @@ -64,6 +69,15 @@ .. autoclass:: SSPRKMethodBuilder .. autoclass:: SSPRK22MethodBuilder .. autoclass:: SSPRK33MethodBuilder + +Implicit Methods +----------------------------------------- + +.. autoclass:: BackwardEulerMethodBuilder +.. autoclass:: DIRK2MethodBuilder +.. autoclass:: DIRK3MethodBuilder +.. autoclass:: DIRK4MethodBuilder +.. autoclass:: DIRK5MethodBuilder """ @@ -119,6 +133,10 @@ def c(self): def a_explicit(self): raise NotImplementedError + @property + def a_implicit(self): + raise NotImplementedError + @property def output_coeffs(self): raise NotImplementedError @@ -153,48 +171,85 @@ def generate_butcher(self, stage_coeff_set_names, stage_coeff_sets, rhs_funcs, :arg estimate_coeffs_sets: a mapping from estimate coefficient set names to cofficients. """ - - from pymbolic import var - comp = self.component_id - - dt = self.dt - t = self.t - state = self.state - - nstages = len(self.c) - # {{{ check coefficients for plausibility for name in stage_coeff_set_names: - for istage in range(nstages): + for istage in range(len(self.c)): coeff_sum = sum(stage_coeff_sets[name][istage]) assert abs(coeff_sum - self.c[istage]) < 1e-12, ( name, istage, coeff_sum, self.c[istage]) # }}} - # {{{ initialization + cb_init = CodeBuilder(name="initialization") + cb_init = self.generate_butcher_init(cb_init, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets) + cb_primary = CodeBuilder(name="primary") + cb_primary, stage_rhs_vars, estimate_vars, = self.generate_butcher_primary( + cb_primary, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets) + cb_primary = self.generate_butcher_finish(cb_primary, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, + estimate_coeff_set_names, + estimate_coeff_sets, + stage_rhs_vars, + estimate_vars) + return DAGCode( + phases={ + "initial": cb_init.as_execution_phase(next_phase="primary"), + "primary": cb_primary.as_execution_phase(next_phase="primary") + }, + initial_phase="initial") - last_rhss = {} + # {{{ initialization - with CodeBuilder(name="initialization") as cb: - for name in stage_coeff_set_names: - if ( - name in self.recycle_last_stage_coeff_set_names - and _is_first_stage_same_as_last_stage( - self.c, stage_coeff_sets[name])): - last_rhss[name] = var("

last_rhs_" + name) - cb(last_rhss[name], rhs_funcs[name](t=t, **{comp: state})) + def generate_butcher_init(self, cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, estimate_coeff_set_names, + estimate_coeff_sets): + self.last_rhss = {} + from pymbolic import var + for name in stage_coeff_set_names: + if ( + name in self.recycle_last_stage_coeff_set_names + and _is_first_stage_same_as_last_stage( + self.c, stage_coeff_sets[name])): + self.last_rhss[name] = var("

last_rhs_" + name) + cb(self.last_rhss[name], rhs_funcs[name]( + t=self.t, **{self.component_id: self.state})) cb_init = cb - # }}} + return cb_init + + # }}} + + # {{{ stage loop + + def generate_butcher_primary(self, cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, estimate_coeff_set_names, + estimate_coeff_sets): + + last_state_est_var = cb.fresh_var("last_state_est") + last_state_est_var_valid = False + + equations = [] + unknowns = set() + + comp = self.component_id + nstages = len(self.c) + dt = self.dt + t = self.t + last_rhss = self.last_rhss stage_rhs_vars = {} rhs_var_to_unknown = {} for name in stage_coeff_set_names: stage_rhs_vars[name] = [ - cb.fresh_var(f"rhs_{name}_s{i}") for i in range(nstages)] + cb.fresh_var(f"rhs_{name}_s{i}") for i in range(len(self.c))] # These are rhss if they are not yet known and pending an implicit solve. for i, rhsvar in enumerate(stage_rhs_vars[name]): @@ -203,166 +258,155 @@ def generate_butcher(self, stage_coeff_set_names, stage_coeff_sets, rhs_funcs, knowns = set() - # {{{ stage loop - - last_state_est_var = cb.fresh_var("last_state_est") - last_state_est_var_valid = False - - with CodeBuilder(name="primary") as cb: - equations = [] - unknowns = set() - - def make_known(v): - unknowns.discard(v) - knowns.add(v) - - for istage in range(nstages): - for name in stage_coeff_set_names: - c = self.c[istage] - my_rhs = stage_rhs_vars[name][istage] - - if ( - name in self.recycle_last_stage_coeff_set_names - and istage == 0 - and _is_first_stage_same_as_last_stage( - self.c, stage_coeff_sets[name])): - cb(my_rhs, last_rhss[name]) - make_known(my_rhs) - - else: - is_implicit = False - - state_increment = 0 - for src_name in stage_coeff_set_names: - coeffs = stage_coeff_sets[src_name][istage] - for src_istage, coeff in enumerate(coeffs): - rhsval = stage_rhs_vars[src_name][src_istage] - if rhsval not in knowns: - unknowns.add(rhsval) - is_implicit = True - - state_increment += dt * coeff * rhsval - - state_est = state + state_increment - if (self.state_filter is not None - and not ( - # reusing last output state - c == 0 - and all( - len(stage_coeff_sets[src_name][istage]) == 0 - for src_name in stage_coeff_set_names))): - state_est = self.state_filter(state_est) - - if is_implicit: - rhs_expr = rhs_funcs[name]( - t=t + c*dt, **{comp: state_est}) - - from dagrt.expression import collapse_constants - solve_expression = collapse_constants( - my_rhs - rhs_expr, - list(unknowns) + [self.state], - cb.assign, cb.fresh_var) - equations.append(solve_expression) - - if istage + 1 == nstages: - last_state_est_var_valid = False - - else: - if istage + 1 == nstages: - cb(last_state_est_var, state_est) - state_est = last_state_est_var - last_state_est_var_valid = True - - rhs_expr = rhs_funcs[name]( - t=t + c*dt, **{comp: state_est}) - - cb(my_rhs, rhs_expr) - make_known(my_rhs) - - # {{{ 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 = { - 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": state}, - solver_id="solve") - - del equations[:] - knowns.update(unknowns) - unknowns.clear() - - # }}} - - # Compute solution estimates. - estimate_vars = [ - cb.fresh_var("est_"+name) - for name in estimate_coeff_set_names] - - for iest, name in enumerate(estimate_coeff_set_names): - out_coeffs = estimate_coeff_sets[name] + for istage in range(len(self.c)): + for name in stage_coeff_set_names: + c = self.c[istage] + my_rhs = stage_rhs_vars[name][istage] if ( - last_state_est_var_valid - and # noqa: W504 - _is_last_stage_same_as_output(self.c, - stage_coeff_sets, out_coeffs)): - state_est = last_state_est_var + name in self.recycle_last_stage_coeff_set_names + and istage == 0 + and _is_first_stage_same_as_last_stage( + self.c, stage_coeff_sets[name])): + cb(my_rhs, last_rhss[name]) + knowns.add(my_rhs) else: + is_implicit = False + state_increment = 0 for src_name in stage_coeff_set_names: - state_increment += sum( - coeff * stage_rhs_vars[src_name][src_istage] - for src_istage, coeff in enumerate(out_coeffs)) + coeffs = stage_coeff_sets[src_name][istage] + for src_istage, coeff in enumerate(coeffs): + rhsval = stage_rhs_vars[src_name][src_istage] + if rhsval not in knowns: + unknowns.add(rhsval) + is_implicit = True + + state_increment += dt * coeff * rhsval + + state_est = self.state + state_increment + if (self.state_filter is not None + and not ( + # reusing last output state + c == 0 + and all( + len(stage_coeff_sets[src_name][istage]) == 0 + for src_name in stage_coeff_set_names))): + state_est = self.state_filter(state_est) - state_est = state + dt*state_increment + if is_implicit: + rhs_expr = rhs_funcs[name]( + t=t + c*dt, **{comp: state_est}) - if self.state_filter is not None: - state_est = self.state_filter(state_est) + from dagrt.expression import collapse_constants + solve_expression = collapse_constants( + my_rhs - rhs_expr, + list(unknowns) + [self.state], + cb.assign, cb.fresh_var) + equations.append(solve_expression) - cb( - estimate_vars[iest], - state_est) + if istage + 1 == nstages: + last_state_est_var_valid = False - # This updates . - self.finish(cb, estimate_coeff_set_names, estimate_vars) + else: + if istage + 1 == nstages: + cb(last_state_est_var, state_est) + state_est = last_state_est_var + last_state_est_var_valid = True + + rhs_expr = rhs_funcs[name]( + t=t + c*dt, **{comp: state_est}) + + cb(my_rhs, rhs_expr) + knowns.add(my_rhs) + + # {{{ emit solve if we have any unknowns/equations + + if unknowns and len(unknowns) == len(equations): + from leap.implicit import generate_solve + generate_solve(cb, unknowns, equations, rhs_var_to_unknown, + self.state) + elif not unknowns: + # we have an explicit Runge-Kutta method + pass + elif len(unknowns) > len(equations): + raise ValueError("Runge-Kutta implicit timestep has more " + "unknowns than equations") + elif len(unknowns) < len(equations): + raise ValueError("Runge-Kutta implicit timestep has more " + "equations than unknowns") + else: + raise AssertionError() - # These updates have to happen *after* finish because before we - # don't yet know whether finish will accept the new state. - for name in stage_coeff_set_names: - if ( - name in self.recycle_last_stage_coeff_set_names - and _is_first_stage_same_as_last_stage( - self.c, stage_coeff_sets[name])): - cb(last_rhss[name], stage_rhs_vars[name][-1]) + del equations[:] + knowns.update(unknowns) + unknowns.clear() + + # }}} + + # Compute solution estimates. + estimate_vars = [ + cb.fresh_var("est_"+name) + for name in estimate_coeff_set_names] + + for iest, name in enumerate(estimate_coeff_set_names): + out_coeffs = estimate_coeff_sets[name] + + if ( + last_state_est_var_valid + and # noqa: W504 + _is_last_stage_same_as_output(self.c, + stage_coeff_sets, out_coeffs)): + state_est = last_state_est_var + + else: + state_increment = 0 + for src_name in stage_coeff_set_names: + state_increment += sum( + coeff * stage_rhs_vars[src_name][src_istage] + for src_istage, coeff in enumerate(out_coeffs)) + + state_est = self.state + dt*state_increment + + if self.state_filter is not None: + state_est = self.state_filter(state_est) + + cb( + estimate_vars[iest], + state_est) cb_primary = cb - # }}} + return cb_primary, stage_rhs_vars, estimate_vars - return DAGCode( - phases={ - "initial": cb_init.as_execution_phase(next_phase="primary"), - "primary": cb_primary.as_execution_phase(next_phase="primary") - }, - initial_phase="initial") + # }}} + + # {{{ update state and time + + def generate_butcher_finish(self, cb, stage_coeff_set_names, + stage_coeff_sets, rhs_funcs, estimate_coeff_set_names, + estimate_coeff_sets, stage_rhs_vars, estimate_vars): + + last_rhss = self.last_rhss + + # This updates . + self.finish(cb, estimate_coeff_set_names, estimate_vars) + + # These updates have to happen *after* finish because before we + # don't yet know whether finish will accept the new state. + for name in stage_coeff_set_names: + if ( + name in self.recycle_last_stage_coeff_set_names + and _is_first_stage_same_as_last_stage( + self.c, stage_coeff_sets[name])): + cb(last_rhss[name], stage_rhs_vars[name][-1]) + + cb_primary = cb + + return cb_primary + + # }}} def finish(self, cb, estimate_names, estimate_vars): cb(self.state, estimate_vars[0]) @@ -416,22 +460,6 @@ class ForwardEulerMethodBuilder(SimpleButcherTableauMethodBuilder): recycle_last_stage_coeff_set_names = () -class BackwardEulerMethodBuilder(SimpleButcherTableauMethodBuilder): - """ - .. automethod:: __init__ - .. automethod:: generate - """ - c = (0,) - - a_explicit = ( - (1,), - ) - - output_coeffs = (1,) - - recycle_last_stage_coeff_set_names = () - - class MidpointMethodBuilder(SimpleButcherTableauMethodBuilder): """ .. automethod:: __init__ @@ -539,14 +567,178 @@ class RK5MethodBuilder(SimpleButcherTableauMethodBuilder): # }}} +# {{{ implicit butcher tableau methods + + +class ImplicitButcherTableauMethodBuilder(ButcherTableauMethodBuilder): + def __init__(self, component_id, state_filter_name=None, + rhs_func_name=None): + super().__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 BackwardEulerMethodBuilder(ImplicitButcherTableauMethodBuilder): + """ + .. automethod:: __init__ + .. automethod:: generate + """ + c = (1,) + + a_implicit = ( + (1,), + ) + + output_coeffs = (1,) + + recycle_last_stage_coeff_set_names = () + + +class DIRK2MethodBuilder(ImplicitButcherTableauMethodBuilder): + """ + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 72, eqn 221 + + .. automethod:: __init__ + .. automethod:: generate + """ + + _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): + """ + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 77, eqn 229 & eqn 230 + + .. automethod:: __init__ + .. automethod:: generate + """ + + _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): + """ + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 78, eqn 232 + + .. automethod:: __init__ + .. automethod:: generate + """ + + _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): + """ + Source: Kennedy & Carpenter: Diagonally Implicit Runge-Kutta + Methods for Ordinary Differential Equations. A Review + pp. 98, Table 24 + + .. automethod:: __init__ + .. automethod:: generate + """ + + 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 = () + + +IMPLICIT_ORDER_TO_RK_METHOD_BUILDER = { + 1: BackwardEulerMethodBuilder, + 2: DIRK2MethodBuilder, + 3: DIRK3MethodBuilder, + 4: DIRK4MethodBuilder, + 5: DIRK5MethodBuilder, + } + +# }}} + + # {{{ Embedded Runge-Kutta schemes base class + class EmbeddedButcherTableauMethodBuilder( ButcherTableauMethodBuilder, TwoOrderAdaptiveMethodBuilderMixin): """ User-supplied context: - + component_id: The value that is integrated - + component_id: The right hand side function + + component_id: The value that is integrated + + component_id: The right hand side function """ @property diff --git a/test/test_ab.py b/test/test_adams.py similarity index 73% rename from test/test_ab.py rename to test/test_adams.py index afea0f2..5e3fe73 100755 --- a/test/test_ab.py +++ b/test/test_adams.py @@ -28,7 +28,7 @@ import sys import pytest -from leap.multistep import AdamsBashforthMethodBuilder +from leap.multistep import AdamsBashforthMethodBuilder, AdamsMoultonMethodBuilder from utils import ( # noqa python_method_impl_interpreter as pmi_int, @@ -53,6 +53,24 @@ def test_ab_accuracy(python_method_impl, method, expected_order, plot_solution=plot_solution) +@pytest.mark.parametrize(("method", "expected_order"), [ + (AdamsMoultonMethodBuilder("y", order, static_dt=static_dt), order) + for order in [1, 3, 5] + for static_dt in [True, False] + ] + [ + (AdamsMoultonMethodBuilder("y", order, hist_length=order+1, + static_dt=static_dt), order) + for order in [1, 3, 5] + for static_dt in [True, False] + ]) +def test_am_accuracy(python_method_impl, method, expected_order, + show_dag=False, plot_solution=False): + from utils import check_simple_convergence + check_simple_convergence(method=method, method_impl=python_method_impl, + expected_order=expected_order, show_dag=show_dag, + plot_solution=plot_solution, implicit=True) + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1]) diff --git a/test/test_rk.py b/test/test_rk.py index 6786476..876c30c 100755 --- a/test/test_rk.py +++ b/test/test_rk.py @@ -35,6 +35,9 @@ RK3MethodBuilder, RK4MethodBuilder, RK5MethodBuilder, LSRK4MethodBuilder, SSPRK22MethodBuilder, SSPRK33MethodBuilder, + BackwardEulerMethodBuilder, DIRK2MethodBuilder, + DIRK3MethodBuilder, DIRK4MethodBuilder, + DIRK5MethodBuilder, ) from leap.rk.imex import KennedyCarpenterIMEXARK4MethodBuilder import numpy as np @@ -77,6 +80,21 @@ def test_rk_accuracy(python_method_impl, method, expected_order, expected_order=expected_order, show_dag=show_dag, plot_solution=plot_solution) + +@pytest.mark.parametrize(("method", "expected_order"), [ + (BackwardEulerMethodBuilder("y"), 1), + (DIRK2MethodBuilder("y"), 2), + (DIRK3MethodBuilder("y"), 3), + (DIRK4MethodBuilder("y"), 4), + (DIRK5MethodBuilder("y"), 5), + ]) +def test_implicit_rk_accuracy(python_method_impl, method, expected_order, + show_dag=False, plot_solution=False): + from utils import check_simple_convergence + check_simple_convergence(method=method, method_impl=python_method_impl, + expected_order=expected_order, show_dag=show_dag, + plot_solution=plot_solution, implicit=True) + # }}} diff --git a/test/utils.py b/test/utils.py index 60d3ac3..6bec194 100644 --- a/test/utils.py +++ b/test/utils.py @@ -38,6 +38,20 @@ def python_method_impl_codegen(code, **kwargs): # }}} +def 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 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 execute_and_return_single_result(python_method_impl, code, initial_context=None, max_steps=1): if initial_context is None: @@ -97,7 +111,7 @@ def __call__(self, t, y): def check_simple_convergence(method, method_impl, expected_order, problem=None, dts=_default_dts, - show_dag=False, plot_solution=False): + show_dag=False, plot_solution=False, implicit=False): if problem is None: problem = DefaultProblem() @@ -109,6 +123,10 @@ def check_simple_convergence(method, method_impl, expected_order, from dagrt.language import show_dependency_graph show_dependency_graph(code) + if implicit: + from leap.implicit import replace_AssignImplicit + code = replace_AssignImplicit(code, {"solve": solver_hook}) + from pytools.convergence import EOCRecorder eocrec = EOCRecorder() @@ -117,9 +135,16 @@ def check_simple_convergence(method, method_impl, expected_order, y = problem.initial() final_t = problem.t_end - interp = method_impl(code, function_map={ - "" + component_id: problem, - }) + if implicit: + from functools import partial + interp = method_impl(code, function_map={ + "" + component_id: problem, + "solver": partial(solver, problem), + }) + else: + interp = method_impl(code, function_map={ + "" + component_id: problem, + }) interp.set_up(t_start=t, dt_start=dt, context={component_id: y}) times = [] @@ -150,7 +175,7 @@ def check_simple_convergence(method, method_impl, expected_order, print(eocrec.pretty_print()) orderest = eocrec.estimate_order_of_convergence()[0, 1] - assert orderest > expected_order * 0.9 + assert orderest > expected_order * 0.89 # vim: foldmethod=marker