diff --git a/doc/reference.rst b/doc/reference.rst index dbb7ffb03ee0d5155e1936de83337f530a97b86e..5df74cb4e73d353c459df5c6e20ff4d00dbafb7b 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -24,6 +24,11 @@ Analysis .. automodule:: leap.step_matrix .. automodule:: leap.stability +Support for Implicit Methods +---------------------------- + +.. automodule:: leap.implicit + Transformation -------------- diff --git a/examples/fully-implicit-rk/fully-implicit-rk.py b/examples/fully-implicit-rk/fully-implicit-rk.py new file mode 100755 index 0000000000000000000000000000000000000000..b3fa4ef22eb669ab0b1657e9b0fdeb2957575a45 --- /dev/null +++ b/examples/fully-implicit-rk/fully-implicit-rk.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python +"""Fully implicit Runge-Kutta timestepper""" + +__copyright__ = "Copyright (C) 2020 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 logging +import sys + +import numpy as np +import pytest # noqa +from pymbolic import var + +from leap.rk import ButcherTableauMethodBuilder + + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class FullyImplicitOrder4RKMethodBuilder(ButcherTableauMethodBuilder): + """ + Source: E. Hairer, S. P. Nørsett, and G. Wanner; + Solving Ordinary Differential Equations I; Table 7.3 + """ + + c = (0.21132486540518713, 0.7886751345948129) + + a_implicit = ( + (0.25, -0.038675134594812866), + (0.5386751345948129, 0.25)) + + output_coeffs = (0.5, 0.5) + + recycle_last_stage_coeff_set_names = () + + order = 4 + + def __init__(self, component_id): + super().__init__(component_id) + self.rhs_func_name = f"{component_id}" + + def generate(self): + 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, + }) + + +def solver_hook(insn): + """Replaces the *AssignImplicit* statement generated in + *FullyImplicitOrder4RKMethodBuilder* with a function call to a solver. + """ + args = [] + args.extend(insn.input_expressions["times"]) + args.extend(insn.input_expressions["incrs"]) + for coeff_row in insn.input_expressions["coeffs"]: + args.extend(coeff_row) + return var("solver")(*args) + + +class ScipyFullyImplicitRKSolverImpl(object): + """A solver implementation based on *scipy.optimize.root*""" + + def __init__(self, n, rhs): + self.rhs = rhs + self.n = n + + def __call__(self, *args): + from scipy.optimize import root + n = self.n + assert len(args) == 2*n + n*n + + times = np.array(args[:n]) + incrs = np.array(args[n:2*n]) + coeffs = np.array(args[2*n:2*n+n*n]).reshape(n, n) + + def f(x): + result = x - np.array([ + self.rhs(time, incr + (x * coeff_row).sum()) + for time, incr, coeff_row in zip(times, incrs, coeffs)]) + logger.debug("residual: %e", np.max(np.abs(result))) + return result + + guess = np.array([self.rhs(times[0], incrs[0])] * n) + return root(f, x0=np.array(guess)).x + + +def test_fully_implicit_rk(plot_solution=False): + component_id = "y" + mbuilder = FullyImplicitOrder4RKMethodBuilder(component_id) + code = mbuilder.generate() + logger.info("DAG before solver hook:\n%s", code) + from leap.implicit import replace_AssignImplicit + code = replace_AssignImplicit(code, solver_hook) + logger.info("DAG after solver hook:\n%s", code) + + h = -0.5 + y_0 = 1.0 + + def rhs(t, y): + return h * y + + def soln(t): + return y_0 * np.exp(h * t) + + from pytools.convergence import EOCRecorder + eocrec = EOCRecorder() + + solver = ScipyFullyImplicitRKSolverImpl(2, rhs) + + for n in range(1, 5): + dt = 2**(-n) + t = 0.0 + final_t = 1 + + from dagrt.exec_numpy import NumpyInterpreter + interp = NumpyInterpreter(code, + function_map={"solver": solver}) + + interp.set_up(t_start=t, dt_start=dt, context={component_id: y_0}) + + 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) + times.append(event.t) + + assert abs(times[-1] - final_t) < 1e-10 + times = np.array(times) + + if plot_solution: + import matplotlib.pyplot as pt + pt.plot(times, values, label="comp") + pt.plot(times, soln(times), label="true") + pt.legend() + pt.show() + + error = abs(values[-1]-soln(final_t)) + eocrec.add_data_point(dt, error) + + print(eocrec.pretty_print()) + + assert eocrec.order_estimate() >= mbuilder.order - 0.1 + + +if __name__ == "__main__": + if len(sys.argv) > 1: + exec(sys.argv[1]) + else: + from pytest import main + main([__file__]) + +# vim: filetype=pyopencl:fdm=marker diff --git a/examples/imex/imex.py b/examples/imex/imex.py index d88878b90d2d04f3775afba49fff7a73add8459a..960b351ceff20ea68b2cfd96171a2856b65f53f5 100755 --- a/examples/imex/imex.py +++ b/examples/imex/imex.py @@ -56,31 +56,30 @@ class KapsProblem(object): _atol = 1.0e-5 -def solver(f, j, t, u_n, x, c): +def solver(f, j, t, u_n, c): """Kennedy and Carpenter, page 15""" import numpy.linalg as nla - I = np.eye(len(u_n)) + I_ = np.eye(len(u_n)) u = u_n while True: - M = I - c * j(t, u) - r = -(u - u_n) + x + c * f(t=t, y=u) + M = I_ - c * j(t, u) + r = u_n + c * f(t=t, y=u) - u d = nla.solve(M, r) u = u + d if 0.005 * _atol >= nla.norm(d): return f(t=t, y=u) -def solver_hook(solve_expr, solve_var, solver_id, guess): - from dagrt.expression import match, substitute +def solver_hook(insn): + assert len(insn.assignees) == 1 + from dagrt.expression import substitute - pieces = match("unk - rhs(t=t, y=y + sub_y + coeff*unk)", - solve_expr, - bound_variable_names=["y"], - pre_match={"unk": solve_var}) + pieces = dict( + t=insn.input_expressions["times"][0], + u_n=insn.input_expressions["incrs"][0], + c=insn.input_expressions["coeffs"][0][0]) - pieces["guess"] = guess - - return substitute("solver(t, y, sub_y, coeff)", pieces) + return substitute("solver(t, u_n, c)", pieces) def run(): @@ -94,7 +93,7 @@ def run(): code = mgen.generate() from leap.implicit import replace_AssignImplicit - code = replace_AssignImplicit(code, {"solve": solver_hook}) + code = replace_AssignImplicit(code, solver_hook) IMEXIntegrator = PythonCodeGenerator("IMEXIntegrator").get_class(code) # Set up the problem and run the method. @@ -118,7 +117,8 @@ def run(): t = event.t y = event.state_component - print("Error: " + str(np.linalg.norm(y - problem.exact(t)))) + print("Tolerance:", _atol) + print("Error:", np.linalg.norm(y - problem.exact(t))) if __name__ == "__main__": diff --git a/examples/implicit_euler/implicit_euler.py b/examples/implicit_euler/implicit_euler.py index bc33f320cbf88cbbdc978aed12eb1c03b75b786b..3d23b2d088053542554fad0f87aebb1502db25cb 100755 --- a/examples/implicit_euler/implicit_euler.py +++ b/examples/implicit_euler/implicit_euler.py @@ -5,7 +5,6 @@ from __future__ import division from leap import MethodBuilder from dagrt.language import DAGCode, CodeBuilder from pymbolic import var -from pymbolic.primitives import CallWithKwargs __copyright__ = "Copyright (C) 2014 Matt Wala" @@ -38,18 +37,24 @@ class ImplicitEulerMethodBuilder(MethodBuilder): rhs_func: The right hand side function """ - SOLVER_EXPRESSION_ID = 0 - def __init__(self, component_id): self.component_id = component_id - self.dt = var('
') - self.t = var('') - self.state = var('' + component_id) - self.rhs_func = var('' + component_id) + self.dt = var("
") + self.t = var("") + self.state = var("" + component_id) + self.rhs_func = var("" + component_id) def generate(self, solver_hook): """Return code that implements the implicit Euler method for the single - state component supported.""" + state component supported. + + The solver hook should generate an expression for the solution *u* to + the implicit equation:: + + u - rhs(t + dt, y + dt * u) = 0 + + The inputs to the solver are *t*, *dt*, and *y*. + """ with CodeBuilder(name="primary") as cb: self._make_primary(cb) @@ -59,26 +64,23 @@ class ImplicitEulerMethodBuilder(MethodBuilder): initial_phase="primary") from leap.implicit import replace_AssignImplicit - - return replace_AssignImplicit(code, {self.SOLVER_EXPRESSION_ID: solver_hook}) + return replace_AssignImplicit(code, solver_hook) def _make_primary(self, builder): """Add code to drive the primary stage.""" - solve_component = var('next_state') - solve_expression = solve_component - self.state - \ - self.dt * CallWithKwargs( - function=self.rhs_func, - parameters=(), - kw_parameters={ - 't': self.t + self.dt, - self.component_id: solve_component - }) - - builder.assign_implicit_1(self.state, solve_component, - solve_expression, self.state, - self.SOLVER_EXPRESSION_ID) - - builder.yield_state(self.state, self.component_id, - self.t + self.dt, 'final') + rhs = var("rhs_{self.component}") + + builder.assign_implicit( + assignees=(rhs,), + input_expressions=dict( + t=self.t, + dt=self.dt, + y=self.state), + params=dict()) + + builder.assign(self.state, self.state + self.dt * rhs) + builder.assign(self.t, self.t + self.dt) + + builder.yield_state(self.state, self.component_id, self.t, "final") diff --git a/examples/implicit_euler/test_implicit_euler.py b/examples/implicit_euler/test_implicit_euler.py index cf235cc087d3716f8d22b9935d50a2804873bff0..a862203bf4ec836dfaaaa1e3c1acd33a4ea53131 100755 --- a/examples/implicit_euler/test_implicit_euler.py +++ b/examples/implicit_euler/test_implicit_euler.py @@ -42,16 +42,19 @@ def python_method_impl_codegen(code, **kwargs): return codegen.get_class(code)(**kwargs) -def solver(f, t, h, y, guess): +def solver(f, t, dt, y): from scipy.optimize import newton - return newton(lambda u: u-y-h*f(t=t, y=u), guess) + guess = f(t, y) + return newton(lambda u: u - f(t=t+dt, y=y+dt*u), guess) -def solver_hook(expr, var, solver_id, guess): - from dagrt.expression import match, substitute - pieces = match("unk-y-h*f(t=t,y=unk)", expr, pre_match={"unk": var}) - pieces["guess"] = guess - return substitute("solver(t,h,y,guess)", pieces) +def solver_hook(insn): + from dagrt.expression import substitute + pieces = dict( + t=insn.input_expressions["t"], + dt=insn.input_expressions["dt"], + y=insn.input_expressions["y"]) + return substitute("solver(t, dt, y)", pieces) @pytest.mark.parametrize("python_method_impl", diff --git a/examples/rk/run_rk_method.py b/examples/rk/run_rk_method.py index 25ba6d0296827676a9a4c0c9982250ac3ee869e3..573948082beb3447795c414a939701c751e499b2 100755 --- a/examples/rk/run_rk_method.py +++ b/examples/rk/run_rk_method.py @@ -98,5 +98,6 @@ def main(show_dag=False, plot_solution=False): orderest = eocrec.estimate_order_of_convergence()[0, 1] assert orderest > expected_order*0.95 + if __name__ == "__main__": main() diff --git a/examples/semi-implicit-rk/semi-implicit-rk.py b/examples/semi-implicit-rk/semi-implicit-rk.py index 94175c21b6886c0011f1be0aebe733d587bfd4e0..4ee5219eb2a3960143c6bb2c40ede8f46724b6c0 100755 --- a/examples/semi-implicit-rk/semi-implicit-rk.py +++ b/examples/semi-implicit-rk/semi-implicit-rk.py @@ -20,6 +20,9 @@ import logging logging.basicConfig(level=logging.INFO) +K = -10 + + class SimpleDecayProblem(object): t_start = 0 @@ -29,16 +32,16 @@ class SimpleDecayProblem(object): return np.array([1]) def exact(self, t): - return np.array([np.exp(-10 * t)]) + return np.array([np.exp(K * t)]) def __call__(self, t, y): - return -10 * y + return K * y _default_dts = 2 ** -np.array(range(5, 10), dtype=np.float64) -def get_convergence_data(method_class, problem, dts=_default_dts): +def get_convergence_data(method_class, problem, solver, dts=_default_dts): from pytools.convergence import EOCRecorder eocrec = EOCRecorder() @@ -52,6 +55,7 @@ def get_convergence_data(method_class, problem, dts=_default_dts): interp = method_class(function_map={ "f": problem, + "solver": solver, }) interp.set_up(t_start=t, dt_start=dt, context={component_id: y}) @@ -84,15 +88,19 @@ def demo_rk_implicit(): with CodeBuilder("primary") as cb: cb.assign_implicit_1( - k1, - solve_component=k1, - expression=k1 - f(t + gamma * dt, y + dt * gamma * k1), - guess=f(t, y)) + assignee=k1, + input_expressions=dict( + time=t + gamma * dt, + incr=y, + coeff=dt * gamma, + guess=f(t, y))), cb.assign_implicit_1( - k2, - solve_component=k2, - expression=k2 - f(t + dt, y + dt * ((1 - gamma) * k1 + gamma * k2)), # noqa - guess=k1) + assignee=k2, + input_expressions=dict( + time=t + dt, + incr=y + dt * (1 - gamma) * k1, + coeff=dt * gamma, + guess=k1)) cb(y, y + dt * ((1 - gamma) * k1 + gamma * k2)) cb(t, t + dt) cb.yield_state(y, "y", t, None) @@ -103,13 +111,17 @@ def demo_rk_implicit(): }, initial_phase="primary") - def solver_hook(solve_expr, unknown, solver_id, guess): - from dagrt.expression import match, substitute - pieces = match( - "k - rhs(time, y + dt * (c0 + c1 * k))", - solve_expr, - pre_match={"k": unknown}) - return substitute("-10 * (dt * c0 + y) / (10 * dt * c1 + 1)", pieces) + def solver_hook(insn): + incr = insn.input_expressions["incr"] + coeff = insn.input_expressions["coeff"] + + from dagrt.expression import substitute + pieces = {"coeff": coeff, "incr": incr} + return substitute("solver(incr, coeff)", pieces) + + def solver(incr, coeff): + # The implicit system for y' = Ky has a closed-form solution. + return K*incr / (1 - K*coeff) from leap.implicit import replace_AssignImplicit code = replace_AssignImplicit(code, solver_hook) @@ -117,7 +129,7 @@ def demo_rk_implicit(): from dagrt.codegen import PythonCodeGenerator IRKMethodBuilder = PythonCodeGenerator("IRKMethodBuilder").get_class(code) - eocrec = get_convergence_data(IRKMethodBuilder, SimpleDecayProblem()) + eocrec = get_convergence_data(IRKMethodBuilder, SimpleDecayProblem(), solver) print(eocrec.pretty_print()) diff --git a/leap/implicit.py b/leap/implicit.py index ee10a648c02bad853cea7ea3982d8d0d8052fc39..4bd678f54fc47a86e01f855ae750d9a0d1ae5e43 100644 --- a/leap/implicit.py +++ b/leap/implicit.py @@ -1,9 +1,5 @@ -from __future__ import division - -"""Implicit solver utilities""" - __copyright__ = """ -Copyright (C) 2014, 2015 Matt Wala +Copyright (C) 2014, 2015, 2020 Matt Wala """ __license__ = """ @@ -26,32 +22,38 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -import six + +__doc__ = """ +.. autofunction:: replace_AssignImplicit + +""" def replace_AssignImplicit(dag, solver_hooks): - """ - :arg dag: The :class:`DAGCode` instance + """Replace :class:`~dagrt.language.AssignImplicit` instructions with + explicit assignments. + + :arg dag: The :class:`~dagrt.language.DAGCode` instance :arg solver_hooks: Either a callable, or a map from solver names to functions that generate solver calls. A solver hook should have the signature:: - def solver_hook(expr, var, id, **kwargs): + def solver_hook(insn): - where: - * *expr* is the expression passed to the AssignImplicit instruction - * *var* is the name of the unknown - * *id* is the *solver_id* field of the AssignImplicit instruction - * any other arguments are passed in *kwargs*. - """ + where *insn* is the AssignImplicit instruction to be replaced. + + The solver hook should return a symbolic expression for a function call. + This expression gets turned into an assignment instruction. - if six.callable(solver_hooks): + :returns: a :class:`~dagrt.language.DAGCode` without implicit assignments + """ + if callable(solver_hooks): hook = solver_hooks from collections import defaultdict solver_hooks = defaultdict(lambda: hook) - from dagrt.language import Assign, AssignImplicit + from dagrt.language import AssignFunctionCall, AssignImplicit new_phases = {} @@ -64,25 +66,17 @@ def replace_AssignImplicit(dag, solver_hooks): new_statements.append(stmt) continue - if len(stmt.assignees) != 1: - from dagrt.utils import TODO - raise TODO("Implement lowering for AssignImplicit statements " - "returning multiple values.") - - expression = stmt.expressions[0] - solve_variable = stmt.solve_variables[0] - solver_id = stmt.solver_id - other_params = stmt.other_params - solver_hook = solver_hooks[stmt.solver_id] - solver_expression = solver_hook(expression, solve_variable, - solver_id, **other_params) - - new_statements.append( - Assign( - assignee=stmt.assignees[0], - assignee_subscript=(), - expression=solver_expression, + solver_expression = solver_hook(stmt) + + from pymbolic.primitives import Call, CallWithKwargs + assert isinstance(solver_expression, (Call, CallWithKwargs)) + new_statements.append(AssignFunctionCall( + assignees=stmt.assignees, + function_id=solver_expression.function.name, + parameters=solver_expression.parameters, + kw_parameters=getattr( + solver_expression, "kw_parameters", {}), id=stmt.id, condition=stmt.condition, depends_on=stmt.depends_on)) diff --git a/leap/rk/__init__.py b/leap/rk/__init__.py index e11393c76814bb51feec3c45d5caa2fb24159b71..24776650d161e3b3647c80326a6a9c9d710478dd 100644 --- a/leap/rk/__init__.py +++ b/leap/rk/__init__.py @@ -27,15 +27,20 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +import dataclasses import six import numpy as np -from leap import MethodBuilder, TwoOrderAdaptiveMethodBuilderMixin -from dagrt.language import CodeBuilder, DAGCode +from typing import Dict, Union +from dagrt.language import CodeBuilder, DAGCode +from pymbolic.primitives import Variable, Expression from pymbolic import var -__doc__ = """ +from leap import MethodBuilder, TwoOrderAdaptiveMethodBuilderMixin + + +__doc__ = r""" .. data:: ORDER_TO_RK_METHOD_BUILDER A dictionary mapping desired order of accuracy to a corresponding RK method @@ -60,6 +65,57 @@ Adaptive/Embedded Methods .. autoclass:: ODE23MethodBuilder .. autoclass:: ODE45MethodBuilder +.. _rk-implicit-interface: + +Implicit Interface for RK Methods +--------------------------------- + +Implicit assignments generated for RK methods describe a system of equations to +be solved in a way similar to a Butcher tableau. + +Let :math:`n` be the number of assignees. The system of :math:`n` equations to +be solved has the following form. Let: + +- :math:`y_i` be the :math:`i`-th unknown, a state component; +- :math:`f_i` be the :math:`i`-th right-hand side function; +- :math:`t_i` be the (scalar) :math:`i`-th solve time; +- :math:`\delta_i` be a known state component value; +- and :math:`c_{ij}` be a scalar multiplier. + +The implicit system is given by: + +.. math:: + + y_i = f_i\left(t_i, \delta_i + \sum_{i=0}^{n-1} c_{ij} y_i\right), + \quad i \in \{0, \ldots, n-1\}. + +The solved value of :math:`y_i` should be assigned to the variable +*assignees[i]*. How to obtain the input values :math:`f_i`, :math:`t_i`, +:math:`\delta_i`, and :math:`c_{ij}` is explained below. + +The attributes of the generated :class:`~dagrt.language.AssignImplicit` +statement are as follows. + +``params`` contains *stage_coeff_set_names*, a tuple of length +:math:`n`. *stage_coeff_set_names[i]* is a string that determines the stage +coefficient set from which *assignees[i]* is derived, and also the meaning of +the right-hand side :math:`f_i` in the above system. + +``input_expressions`` contains the following: + +- *times* and *incrs* are tuples of length :math:`n` containing + expressions for solve time and state increments, respectively. + *times[i]* corresponds to :math:`t_i` and *incrs[i]* corresponds to + :math:`\delta_i` in the above system. +- *coeffs* is a tuple of tuples containing Butcher tableau weights for + the unknowns. *coeffs[i][j]* corresponds to :math:`c_{ij}` in the + above system. +- *guess* is an (optionally present) tuple of length :math:`n` + containing previously evaluated right-hand side values that may be + useful as initial guesses for a solver. +- *state* and *t*, which may be used for generating initial guesses if + *guess* is unavailable or not desired. + """ @@ -104,8 +160,36 @@ def _is_last_stage_same_as_output(c, coeff_sets, output_stage_coefficients): # {{{ fully general butcher tableau to code +SymbolicExpr = Union[Expression, float] + + +@dataclasses.dataclass +class _ImplicitEquationData: + times: Dict[Variable, SymbolicExpr] = dataclasses.field(default_factory=dict) + coeff_sets: Dict[Variable, Dict[Variable, SymbolicExpr]] = \ + dataclasses.field(default_factory=dict) + state_increments: Dict[Variable, SymbolicExpr] = \ + dataclasses.field(default_factory=dict) + stage_coeff_set_names: Dict[Variable, str] = \ + dataclasses.field(default_factory=dict) + + def build_tableau(self, unknowns): + tableau = [] + for unknown in unknowns: + row = [] + coeff_set = self.coeff_sets.setdefault(unknown, {}) + for rhs_var in unknowns: + row.append(coeff_set.get(rhs_var, 0)) + tableau.append(tuple(row)) + return tuple(tableau) + + class ButcherTableauMethodBuilder(MethodBuilder): - """Infrastructure to generate code from butcher tableaux.""" + """Infrastructure to generate code from a Butcher tableau. + + See :ref:`implicit-interface` for the details of the implicit interface + generated by this class. + """ @property def c(self): @@ -127,9 +211,9 @@ class ButcherTableauMethodBuilder(MethodBuilder): self.component_id = component_id - self.dt = var('
') - self.t = var('') - self.state = var('' + component_id) + self.dt = var("
") + self.t = var("") + self.state = var("" + component_id) if state_filter_name is not None: self.state_filter = var("" + state_filter_name) @@ -187,17 +271,9 @@ class ButcherTableauMethodBuilder(MethodBuilder): # }}} stage_rhs_vars = {} - rhs_var_to_unknown = {} for name in stage_coeff_set_names: stage_rhs_vars[name] = [ - cb.fresh_var('rhs_%s_s%d' % (name, i)) for i in range(nstages)] - - # These are rhss if they are not yet known and pending an implicit solve. - for i, rhsvar in enumerate(stage_rhs_vars[name]): - unkvar = cb.fresh_var('unk_%s_s%d' % (name, i)) - rhs_var_to_unknown[rhsvar] = unkvar - - knowns = set() + cb.fresh_var("rhs_%s_s%d" % (name, i)) for i in range(nstages)] # {{{ stage loop @@ -205,106 +281,146 @@ class ButcherTableauMethodBuilder(MethodBuilder): 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) + # The set of known variables. + knowns = set() + + # Data associated with equations for unknown variables. + eq_data = _ImplicitEquationData() + # The list of unknown variables for which there is an equation. + curr_solved_for_unknowns = [] + # The set of unknown variables participating in the current system. + curr_needed_unknowns = set() + + # Used for generating initial guesses for implicit solves. + # FIXME: This is not necessarily a smart initial guess. + most_recent_stage_rhss = {} + for name in stage_coeff_set_names: + if name in last_rhss: + most_recent_stage_rhss[name] = last_rhss[name] 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) + # Gather used vars. + rhs_var_to_coeff = {} + for src_name in stage_coeff_set_names: + coeffs = stage_coeff_sets[src_name][istage] + for src_istage, coeff in enumerate(coeffs): + if coeff: + rhs_var = stage_rhs_vars[src_name][src_istage] + rhs_var_to_coeff[rhs_var] = coeff + + if set(rhs_var_to_coeff) <= knowns: + # RHS is explicitly assigned. + 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])): + # Use the first-same-as-last property to skip the + # function call. + rhs_expr = last_rhss[name] + else: + # Emit the RHS as a function call. + state_est = state + dt * sum( + coeff * rhs_var + for rhs_var, coeff in rhs_var_to_coeff.items()) - if istage + 1 == nstages: - last_state_est_var_valid = False + if self.state_filter is not None: + state_est = self.state_filter(state_est) - 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}) + 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] + knowns.add(my_rhs) + cb(my_rhs, rhs_expr) + most_recent_stage_rhss[name] = my_rhs + else: + # RHS is implicitly assigned. + eq_data.times[my_rhs] = t + c * dt + + coeffs = {} + state_incr = state + for rhs_var, coeff in rhs_var_to_coeff.items(): + if rhs_var in knowns: + state_incr += dt * coeff * rhs_var + else: + coeffs[rhs_var] = dt * coeff + eq_data.state_increments[my_rhs] = state_incr + eq_data.coeff_sets[my_rhs] = coeffs + eq_data.stage_coeff_set_names[my_rhs] = name + + curr_needed_unknowns |= set(rhs_var_to_coeff) - knowns + curr_needed_unknowns.add(my_rhs) + curr_solved_for_unknowns.append(my_rhs) + + # Emit a solve if we have an equation for each pending variable. + if ( + curr_solved_for_unknowns + and ( + set(curr_solved_for_unknowns) + == curr_needed_unknowns)): + + # {{{ emit solve + + # Generate a guess (only if one is available at zero cost). + guess = [] + guess_dict = {"state": state, "t": t} + for unk in curr_solved_for_unknowns: + scset = eq_data.stage_coeff_set_names[unk] + if scset not in most_recent_stage_rhss: + break + guess.append(most_recent_stage_rhss[scset]) + else: + guess_dict["guess"] = tuple(guess) - from pymbolic import substitute - subst_dict = dict( - (rhs_var.name, rhs_var_to_unknown[rhs_var]) - for rhs_var in unknowns) + def as_tuple(dict_): + return tuple(dict_[unk] + for unk in curr_solved_for_unknowns) # noqa 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() - - # }}} + assignees=curr_solved_for_unknowns, + input_expressions=dict( + times=as_tuple(eq_data.times), + incrs=as_tuple(eq_data.state_increments), + coeffs=eq_data.build_tableau( + curr_solved_for_unknowns), + **guess_dict), + params=dict( + stage_coeff_set_names=as_tuple( + eq_data.stage_coeff_set_names)), + solver_id="solver") + + # Update most_recent_stage_rhss. + for rhs in curr_solved_for_unknowns: + scset = eq_data.stage_coeff_set_names[rhs] + most_recent_stage_rhss[scset] = rhs + + knowns |= curr_needed_unknowns + curr_needed_unknowns.clear() + del curr_solved_for_unknowns[:] + + # }}} + + # All unknowns should have been accounted for. + assert knowns == set(rhs_var + for rhs_var_list in stage_rhs_vars.values() + for rhs_var in rhs_var_list) + assert not curr_solved_for_unknowns + assert not curr_needed_unknowns + + del knowns + del eq_data + del most_recent_stage_rhss + del curr_needed_unknowns + del curr_solved_for_unknowns # Compute solution estimates. estimate_vars = [ @@ -504,7 +620,7 @@ class RK4MethodBuilder(SimpleButcherTableauMethodBuilder): class RK5MethodBuilder(SimpleButcherTableauMethodBuilder): """ - Source: J. C. Butcher, Numerical MethodBuilders for Ordinary Differential + Source: J. C. Butcher, Numerical Methods for Ordinary Differential Equations, 2nd ed., pages 94 - 99 """ diff --git a/leap/rk/imex.py b/leap/rk/imex.py index b4884b964ae986aac10e67e59a3c45df7b7a2814..46e3fb49e42af425ce50d6a4740628f19324b751 100644 --- a/leap/rk/imex.py +++ b/leap/rk/imex.py @@ -35,7 +35,7 @@ THE SOFTWARE. __doc__ = """ -IMEX MethodBuilders +IMEX Methods ------------ .. autoclass :: KennedyCarpenterIMEXRungeKuttaMethodBuilderBase @@ -59,6 +59,8 @@ class KennedyCarpenterIMEXRungeKuttaMethodBuilderBase( .. automethod:: __init__ .. automethod:: generate + Follows the :ref:`rk-implicit-interface`. + """ @property diff --git a/leap/version.py b/leap/version.py index 5cea8857d9550183c2caabd0948032faac48e0ae..82ea6fb128de4adc8cddd1928ee3eb063bd70aa0 100644 --- a/leap/version.py +++ b/leap/version.py @@ -1,2 +1,2 @@ -VERSION = (2020, 1) +VERSION = (2020, 2) VERSION_TEXT = ".".join(str(i) for i in VERSION) diff --git a/requirements.txt b/requirements.txt index e22f333cbd530c0259a7df4fcb84ac4496c7482c..c9503d2afc659e9d3333555cb26ae43ca765e368 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,3 @@ git+git://github.com/inducer/pytools git+git://github.com/inducer/pymbolic -git+https://gitlab.tiker.net/inducer/dagrt.git \ No newline at end of file +git+https://gitlab.tiker.net/inducer/dagrt.git@new-implicit-interface diff --git a/setup.py b/setup.py index d6f1a27fc8421c66de1cb0c7836ababcde8efcec..d26ce5b4c0c3f7b06e8f46c6109dad950cf6bb0c 100644 --- a/setup.py +++ b/setup.py @@ -37,13 +37,13 @@ def main(): packages=find_packages(), - python_requires="~=3.6", + python_requires="~=3.7", install_requires=[ "numpy>=1.5", "pytools>=2014.1", "pymbolic>=2014.1", "pytest>=2.3", - "dagrt>=2019.4", + "dagrt>=2020.2", "mako", "six", ], diff --git a/test/test_imex.py b/test/test_imex.py index 8e5b1d973885baca30e65ca56186c790575033a2..b72c959f5cc14e80f0413bc33138a677f89c83bc 100755 --- a/test/test_imex.py +++ b/test/test_imex.py @@ -41,12 +41,18 @@ def solver(f, t, sub_y, coeff, guess): 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 +def solver_hook(insn): + from dagrt.expression import substitute + + assert len(insn.assignees) == 1 + assert insn.params["stage_coeff_set_names"][0] == "implicit" + + pieces = dict( + t=insn.input_expressions["times"][0], + sub_y=insn.input_expressions["incrs"][0], + coeff=insn.input_expressions["coeffs"][0][0], + guess=insn.input_expressions["guess"][0]) - 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) @@ -61,7 +67,7 @@ def test_convergence(python_method_impl, problem, method, expected_order): code = method.generate() from leap.implicit import replace_AssignImplicit - code = replace_AssignImplicit(code, {"solve": solver_hook}) + code = replace_AssignImplicit(code, solver_hook) from pytools.convergence import EOCRecorder eocrec = EOCRecorder() @@ -129,9 +135,8 @@ def test_adaptive(python_method_impl, problem, method): generator = method("y", atol=atol) code = generator.generate() - #sgen = ScipySolverGenerator(*generator.implicit_expression()) from leap.implicit import replace_AssignImplicit - code = replace_AssignImplicit(code, {"solve": solver_hook}) + code = replace_AssignImplicit(code, solver_hook) from functools import partial interp = python_method_impl(code, function_map={