Fully Explicit Embedded Adams w/Adaptive Timestep

.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/__init__.py

@@ -1,7 +1,10 @@
 """Leap root module"""
 
 
-__copyright__ = "Copyright (C) 2014 Andreas Kloeckner"
+__copyright__ = """
+Copyright (C) 2014 Andreas Kloeckner
+CopyRight (C) 2020 Cory Mikida
+"""
 
 __license__ = """
 Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -82,6 +85,7 @@ def implicit_expression(self, expression_tag=None):
 
 # {{{ two-order adaptivity
 
+
 class TwoOrderAdaptiveMethodBuilderMixin(MethodBuilder):
     """
     This class expected the following members to be defined: state, t, dt.
@@ -104,6 +108,10 @@ def __init__(self, atol=0, rtol=0, max_dt_growth=None, min_dt_shrinkage=None):
     def finish_nonadaptive(self, cb, high_order_estimate, low_order_estimate):
         raise NotImplementedError()
 
+    def finish_nonadaptive_adams(self, cb, high_order_estimate, low_order_estimate,
+                                 hist, time_hist):
+        raise NotImplementedError()
+
     def finish_adaptive(self, cb, high_order_estimate, low_order_estimate):
         from pymbolic import var
         from pymbolic.primitives import Comparison, LogicalOr, Max, Min
@@ -152,6 +160,57 @@ def norm(expr):
                Min((0.9 * self.dt * rel_error ** (-1 / self.high_order),
                     self.max_dt_growth * self.dt)))
 
+    def finish_adaptive_adams(self, cb, high_order_estimate, low_order_estimate,
+                              hist, time_hist):
+        from pymbolic import var
+        from pymbolic.primitives import Comparison, LogicalOr, Max, Min
+        from dagrt.expression import IfThenElse
+
+        norm_start_state = var("norm_start_state")
+        norm_end_state = var("norm_end_state")
+        rel_error_raw = var("rel_error_raw")
+        rel_error = var("rel_error")
+
+        def norm(expr):
+            return var("<builtin>norm_2")(expr)
+
+        cb(norm_start_state, norm(self.state))
+        cb(norm_end_state, norm(low_order_estimate))
+        cb(rel_error_raw, norm(high_order_estimate - low_order_estimate)
+                / (var("<builtin>len")(self.state) ** 0.5
+                    * (
+                        self.atol + self.rtol
+                        * Max((norm_start_state, norm_end_state))
+                        )))
+
+        cb(rel_error, IfThenElse(Comparison(rel_error_raw, "==", 0),
+                                 1.0e-14, rel_error_raw))
+
+        with cb.if_(LogicalOr((Comparison(rel_error, ">", 1),
+                               var("<builtin>isnan")(rel_error)))):
+
+            with cb.if_(var("<builtin>isnan")(rel_error)):
+                cb(self.dt, self.min_dt_shrinkage * self.dt)
+            with cb.else_():
+                cb(self.dt, Max((0.9 * self.dt
+                    * rel_error ** (-1 / self.low_order),
+                    self.min_dt_shrinkage * self.dt)))
+
+            with cb.if_(self.t + self.dt, "==", self.t):
+                cb.raise_(TimeStepUnderflow)
+            with cb.else_():
+                cb.fail_step()
+
+        with cb.else_():
+            # This updates <t>: <dt> should not be set before this is called.
+            self.finish_nonadaptive_adams(cb, high_order_estimate,
+                                          low_order_estimate,
+                                          hist, time_hist)
+
+            cb(self.dt,
+               Min((0.9 * self.dt * rel_error ** (-1 / self.high_order),
+                    self.max_dt_growth * self.dt)))
+
 # }}}
 
 

leap/implicit.py

@@ -86,3 +86,28 @@ 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):
+
+    # got a square system, let's solve
+    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],
+
+            # TODO: Could supply a starting guess
+            other_params={
+                "guess": guess},
+            solver_id="solve")