replace thetamethod by rungekutta

CHANGELOG.rst

@@ -9,6 +9,22 @@ features in inverse chronological order.
 New in v7.0 (in development)
 ----------------------------
 
+- New: Runge-Kutta solver
+
+  The :class:`nutils.solver.rungekutta` implements the Runge-Kutta method.
+  Butcher-tableau values can be specified via the arguments ``rkmatrix``,
+  ``rkweights`` and ``rknodes``. Additionaly, the methods ``trapezoidal``,
+  ``gausslegendre4`` and ``gausslegendre6`` provide preset values for well
+  known Runge-Kutta schemes.
+
+  For the moment all schemes are handled implicitly. In future the underlying
+  :class:`nutils.solver.newton` method will be enhanced to recognize and
+  optimize for one way dependencies, at which point triangular Butcher-tableaux
+  will automatically become explicit.
+
+  The :func:`nutils.solver.thetamethod` remains as a special case, along with
+  its presets ``impliciteuler`` and ``cranknicolson``.
+
 - Changed: locate arguments
 
   The :func:`nutils.topology.Topology.locate` method now allows ``tol`` to be
@@ -157,7 +173,7 @@ New in v7.0 (in development)
 
 - New thetamethod argument ``historysuffix`` deprecates ``target0``
 
-  To explicitly refer to the history state in :func:`nutils.solver.thetamethod`
+  To explicitly refer to the history state in ``nutils.solver.thetamethod``
   and its derivatives ``impliciteuler`` and ``cranknicolson``, instead of
   specifiying the target through the ``target0`` parameter, the new argument
   ``historysuffix`` specifies only the suffix to be added to the main target.
@@ -440,7 +456,7 @@ Release date: `2019-06-11 <https://github.com/evalf/nutils/releases/tag/v5.0>`_.
 
 - Thetamethod time target
 
-  The :class:`nutils.solver.thetamethod` class, as well as its special
+  The ``nutils.solver.thetamethod`` class, as well as its special
   cases ``impliciteuler`` and ``cranknicolson``, now have a
   ``timetarget`` argument to specify that the formulation contains a
   time variable::

nutils/solver.py

@@ -658,8 +658,8 @@ def resume(self, history):
 
 
 @iterable.single_or_multiple
-class thetamethod(cache.Recursion, length=1, version=1):
-  '''solve time dependent problem using the theta method
+class rungekutta(cache.Recursion, length=1, version=1):
+  '''solve time dependent problem using the Runge Kutta method
 
   Parameters
   ----------
@@ -671,10 +671,13 @@ class thetamethod(cache.Recursion, length=1, version=1):
       Initial time step, will scale up as residual decreases
   lhs0 : :class:`numpy.ndarray`
       Coefficient vector, starting point of the iterative procedure.
-  theta : :class:`float`
-      Theta value (theta=1 for implicit Euler, theta=0.5 for Crank-Nicolson)
-  residual0 : :class:`nutils.evaluable.AsEvaluableArray`
-      Optional additional residual component evaluated in previous timestep
+  rkmatrix : :class:`numpy.ndarray`
+      Runge-Kutta matrix, or `a` coefficients of the Butcher tableau.
+  rkweights : :class:`numpy.ndarray`
+      Runge-Kutta weights, or `b` coefficient or the Butcher tableau.
+  rknodes : :class:`numpy.ndarray` or `None`
+      Runge-Kutta nodes, or `c` coefficient or the Butcher tableau. Defaults to
+      `rknodes = rkmatrix.sum(axis=1)` if left unspecified.
   constrain : :class:`numpy.ndarray` with dtype :class:`bool` or :class:`float`
       Equal length to ``lhs0``, masks the free vector entries as ``False``
       (boolean) or NaN (float). In the remaining positions the values of
@@ -686,11 +689,9 @@ class thetamethod(cache.Recursion, length=1, version=1):
       Defines the values for :class:`nutils.function.Argument` objects in
       `residual`.  The ``target`` should not be present in ``arguments``.
       Optional.
-  timetarget : :class:`str`
+  timetarget : :class:`str` or `None`
       Name of the :class:`nutils.function.Argument` that represents time.
       Optional.
-  time0 : :class:`float`
-      The intial time.  Default: ``0.0``.
 
   Yields
   ------
@@ -699,42 +700,84 @@ class thetamethod(cache.Recursion, length=1, version=1):
   '''
 
   @types.apply_annotations
-  def __init__(self, target, residual:integraltuple, inertia:optionalintegraltuple, timestep:types.strictfloat, theta:types.strictfloat, lhs0:types.arraydata=None, target0:types.strictstr=None, constrain:arrayordict=None, newtontol:types.strictfloat=1e-10, arguments:argdict={}, newtonargs:types.frozendict={}, timetarget:types.strictstr='_thetamethod_time', time0:types.strictfloat=0., historysuffix:types.strictstr='0'):
+  def __init__(self, target, residual:integraltuple, inertia:optionalintegraltuple, timestep:types.strictfloat, rkmatrix:types.arraydata, rkweights:types.arraydata, rknodes:types.arraydata=None, lhs0:types.arraydata=None, constrain:arrayordict=None, newtontol:types.strictfloat=1e-10, arguments:argdict={}, newtonargs:types.frozendict={}, timetarget:types.strictstr=None):
     super().__init__()
     if len(residual) != len(inertia):
       raise Exception('length of residual and inertia do no match')
-    for inert, res in zip(inertia, residual):
-      if inert and inert.shape != res.shape:
-        raise ValueError('expected `inertia` with shape {} but got {}'.format(res.shape, inert.shape))
-    self.target = target
+    if not all(inert is None or inert.shape == res.shape for inert, res in zip(inertia, residual)):
+      raise ValueError('expected `inertia` with shape {} but got {}'.format(res.shape, inert.shape))
+
+    # butcher tableau
+    B = numpy.asarray(rkweights)
+    assert B.ndim == 1
+    assert B.sum() == 1
+    nrk, = B.shape
+    A = numpy.asarray(rkmatrix)
+    assert A.shape == (nrk, nrk)
+    if rknodes is None:
+      C = A.sum(axis=1)
+    else:
+      C = numpy.asarray(rknodes)
+      assert C.shape == (nrk,)
+
+    dt = timestep * evaluable.Argument('_rk_timescale', ())
+    if timetarget:
+      tt = evaluable.Argument(timetarget, ()) # time target for beginning of current timestep
+
+    argobjs = _argobjs(residual + inertia) # all the relevant dof targets
+    rktargets = {t: ['_rk{}_{}'.format(i, t) for i in range(nrk)] for t in target}
+    rkargobjs = {(t,i): evaluable.Argument(rktargets[t][i], argobjs[t].shape) for t in target for i in range(nrk)} # all the RK arguments
+
+    residuals = [] # all combined residual blocks
+    for n, (inert, res) in enumerate(zip(inertia, residual), start=1):
+      lincond = inert is None \
+        and timetarget not in res.arguments \
+        and not any(evaluable.derivative(res, argobjs[t]).simplified.arguments for t in target)
+      if lincond:
+        log.info('identified residual #{} as a linear condition'.format(n))
+        # Runge-Kutta produces update vectors k1, k2, .. such that, for all i:
+        # I,t(tn + ci h, yn + aij kj) + I,y(tn + ci h, yn + aij kj) ki/h +
+        # R(tn + ci h, yn + aij kj) = 0. This means that for blocks where I = 0
+        # and R is linear and independent of t: R(yn + aij kj) = 0. Since in
+        # general bj != Σi aij, this condition does not carry over to yn+1. We
+        # therefore modify the condition in this scenario to R(yn + ki) = 0,
+        # which is consistent if R(yn) = 0 (true for n>1) and aij nonsingular.
+      for i in range(nrk):
+        resi = res
+        if inert is not None:
+          resi += evaluable.replace_arguments(inert, {t: rkargobjs[t,i] for t in target}) / dt # TODO replace by below, fix inflation issue
+          #for t in target:
+            #resi += (evaluable.derivative(inert, argobjs[t]) * rkargobjs[t,i]).sum(inert.ndim + numpy.arange(argobjs[t].ndim)) / dt
+          if timetarget:
+            resi += evaluable.derivative(inert, tt)
+        # target replacement table for RK level i
+        si = {t: argobjs[t] + (rkargobjs[t,i] if lincond # <- consistent modification
+                 else util.sum(rkargobjs[t,j] * A[i,j] for j in range(nrk))) for t in target}
+        if timetarget:
+          si[timetarget] = tt + dt * C[i]
+        residuals.append(evaluable.replace_arguments(resi, si))
+
     self.newtonargs = newtonargs
     self.newtontol = newtontol
-    self.timestep = timestep
-    self.timetarget = timetarget
     self.lhs0, self.constrain = _parse_lhs_cons(lhs0, constrain, target, _argobjs(residual+inertia), arguments)
-    self.lhs0[timetarget] = numpy.array(time0)
-    if target0 is None:
-      self.old_new = [(t+historysuffix, t) for t in target]
-    elif len(target) == 1:
-      warnings.deprecation('target0 is deprecated; use historysuffix instead (target0=target+historysuffix)')
-      self.old_new = [(target0, target[0])]
-    else:
-      raise Exception('target0 is not supported in combination with multiple targets; use historysuffix instead')
-    self.old_new.append((timetarget+historysuffix, timetarget))
-    subs0 = {new: evaluable.Argument(old, self.lhs0[new].shape) for old, new in self.old_new}
-    dt = evaluable.Argument(timetarget, ()) - subs0[timetarget]
-    self.residuals = tuple(res * theta + evaluable.replace_arguments(res, subs0) * (1-theta) + ((inert - evaluable.replace_arguments(inert, subs0)) / dt if inert else 0) for res, inert in zip(residual, inertia))
-    self.jacobians = _derivative(self.residuals, target)
-
-  def _step(self, lhs0, dt):
-    arguments = lhs0.copy()
-    arguments.update((old, lhs0[new]) for old, new in self.old_new)
-    arguments[self.timetarget] = lhs0[self.timetarget] + dt
+    for t in target:
+      self.constrain.update(dict.fromkeys(rktargets[t], self.constrain.pop(t)))
+    self.target = tuple(trk for t in target for trk in rktargets[t]) # order must match residuals because of constraints
+    self.increments = [(t, tuple(zip(rktargets[t], B))) for t in target]
+    if timetarget:
+      self.increments.append((timetarget, (('_rk_timescale', timestep),)))
+      self.lhs0.setdefault(timetarget, 0.)
+    self.residuals = tuple(residuals)
+    self.jacobians = _derivative(residuals, self.target)
+
+  def _step(self, lhs0, timescale):
     try:
-      return newton(self.target, residual=self.residuals, jacobian=self.jacobians, constrain=self.constrain, arguments=arguments, **self.newtonargs).solve(tol=self.newtontol)
+      lhs = newton(self.target, residual=self.residuals, jacobian=self.jacobians, constrain=self.constrain, arguments=dict(_rk_timescale=timescale, **lhs0), **self.newtonargs).solve(tol=self.newtontol)
     except (SolverError, matrix.MatrixError) as e:
-      log.error('error: {}; retrying with timestep {}'.format(e, dt/2))
-      return self._step(self._step(lhs0, dt/2), dt/2)
+      log.error('error: {}; retrying with {} timestep'.format(e, timescale/2))
+      return self._step(self._step(lhs0, timescale/2), timescale/2)
+    else:
+      return {t: sum([lhs[ti] * bi for ti, bi in increments], lhs[t]) for t, increments in self.increments}
 
   def resume(self, history):
     if history:
@@ -743,9 +786,55 @@ def resume(self, history):
       lhs = self.lhs0
       yield lhs
     while True:
-      lhs = self._step(lhs, self.timestep)
+      lhs = self._step(lhs, 1.)
       yield lhs
 
+trapezoidal = functools.partial(rungekutta,
+  rkmatrix=[[0., 0.], [.5, .5]],
+  rkweights=[.5, .5],
+  rknodes=None)
+
+gausslegendre4 = functools.partial(rungekutta,
+  rkmatrix=[[.25, .25-numpy.sqrt(3)/6], [.25+numpy.sqrt(3)/6, .25]],
+  rkweights=[.5, .5],
+  rknodes=None)
+
+gausslegendre6 = functools.partial(rungekutta,
+  rkmatrix=[[5/36, 2/9-numpy.sqrt(15)/15, 5/36-numpy.sqrt(15)/30], [5/36+numpy.sqrt(15)/24, 2/9, 5/36-numpy.sqrt(15)/24], [5/36+numpy.sqrt(15)/30, 2/9+numpy.sqrt(15)/15, 5/36]],
+  rkweights=[5/18, 4/9, 5/18],
+  rknodes=None)
+
+def thetamethod(*args, theta, **kwargs):
+  '''Solve time dependent problem using the theta method.
+
+  Special case of Runge-Kutta with a Butcher tableau equal to::
+
+      θ | θ
+      --+--
+        | 1
+
+  Parameters
+  ----------
+  theta : :class:`float`
+      Value of the Runge-Kutta matrix; 1 corresponds to implicit Euler, 0.5
+      corresponds to Crank-Nicolson.
+  *args, **kwargs :
+      Derived from :class:`rungekutta`.
+
+  Yields
+  ------
+  :class:`numpy.ndarray`
+      Coefficient vector for all timesteps after the initial condition.
+  '''
+
+  if 'time0' in kwargs:
+    warnings.deprecation('thetamethod time0 argument is deprecated, specify via `arguments` instead')
+    if 'timetarget' in kwargs:
+      kwargs.setdefault('arguments', {})[kwargs['timetarget']] = kwargs.pop('time0')
+    # otherwise the time is not reacheable and we can simply ignore the argument
+
+  return rungekutta(*args, rkmatrix=[[theta]], rkweights=[1.], rknodes=None, **kwargs)
+
 impliciteuler = functools.partial(thetamethod, theta=1)
 cranknicolson = functools.partial(thetamethod, theta=0.5)
 

tests/test_solver.py

@@ -254,7 +254,7 @@ def check(self, method, theta):
       with self.subTest(i=i):
         uactual = next(uactualiter)
         self.assertAllAlmostEqual(uactual, udesired)
-        udesired += timestep*(theta*numpy.sin((i+1)*timestep)+(1-theta)*numpy.sin(i*timestep))
+        udesired += timestep*(numpy.sin((i+theta)*timestep))
 
   def test_impliciteuler(self):
     self.check(solver.impliciteuler, theta=1)