Cory implicit adams

.github/workflows/ci.yml

@@ -52,7 +52,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

leap/multistep/__init__.py

@@ -1,10 +1,10 @@
-"""Adams-Bashforth ODE solvers."""
+"""Adams-Bashforth and Adams-Moulton ODE solvers."""
 
 
 __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 @@
 .. autoclass:: AdamsIntegrationFunctionFamily
 .. autoclass:: AdamsMonomialIntegrationFunctionFamily
 .. autoclass:: AdamsBashforthMethodBuilder
+.. autoclass:: AdamsMoultonMethodBuilder
 """
 
 
@@ -416,4 +417,341 @@ def rk_bootstrap(self, cb):
 
 # }}}
 
+# {{{ am method
+
+
+class AdamsMoultonMethodBuilder(MethodBuilder):
+    """
+    User-supplied context:
+        <state> + component_id: The value that is integrated
+        <func> + 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:`AdamsIntegrationFunctionFamily`
+            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 = AdamsMonomialIntegrationFunctionFamily(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("<p>step")
+        self.function = var("<func>" + component_id)
+        self.history = \
+            [var("<p>f_n_minus_" + str(i)) for i in range(hist_length - 1, 0, -1)]
+
+        if not self.static_dt:
+            self.time_history = [
+                    var("<p>t_n_minus_" + str(i))
+                    for i in range(hist_length - 1, 0, -1)]
+
+        self.state = var("<state>" + component_id)
+        self.t = var("<t>")
+        self.dt = var("<dt>")
+
+        if state_filter_name is not None:
+            self.state_filter = var("<func>" + 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("<builtin>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_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)
+
+            # 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