Source Code for Module hedge.timestep.rk4

  1  """LSERK ODE timestepper.""" 
  2   
  3  from __future__ import division 
  4   
  5  __copyright__ = "Copyright (C) 2007 Andreas Kloeckner" 
  6   
  7  __license__ = """ 
  8  This program is free software: you can redistribute it and/or modify 
  9  it under the terms of the GNU General Public License as published by 
 10  the Free Software Foundation, either version 3 of the License, or 
 11  (at your option) any later version. 
 12   
 13  This program is distributed in the hope that it will be useful, 
 14  but WITHOUT ANY WARRANTY; without even the implied warranty of 
 15  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
 16  GNU General Public License for more details. 
 17   
 18  You should have received a copy of the GNU General Public License 
 19  along with this program.  If not, see U{http://www.gnu.org/licenses/}. 
 20  """ 
 21   
 22   
 23   
 24   
 25  import numpy 
 26  import numpy.linalg as la 
 27  from pytools import memoize 
 28  from hedge.timestep.base import TimeStepper 
 29   
 30   
 31   
 32   
 33  _RK4A = [0.0, 
 34          -567301805773 /1357537059087, 
 35          -2404267990393/2016746695238, 
 36          -3550918686646/2091501179385, 
 37          -1275806237668/ 842570457699, 
 38          ] 
 39   
 40  _RK4B = [1432997174477/ 9575080441755, 
 41          5161836677717 /13612068292357, 
 42          1720146321549 / 2090206949498, 
 43          3134564353537 / 4481467310338, 
 44          2277821191437 /14882151754819, 
 45          ] 
 46   
 47  _RK4C = [0.0, 
 48          1432997174477/9575080441755, 
 49          2526269341429/6820363962896, 
 50          2006345519317/3224310063776, 
 51          2802321613138/2924317926251, 
 52          1, 
 53          ] 
 54   
 55   
 56   
 57   
 58 -class RK4TimeStepper(TimeStepper): 
 59      dt_fudge_factor = 1 
 60   
 61 -    def __init__(self, allow_jit=True): 
 62          from pytools.log import IntervalTimer, EventCounter 
 63          self.timer = IntervalTimer( 
 64                  "t_rk4", "Time spent doing algebra in RK4") 
 65          self.flop_counter = EventCounter( 
 66                  "n_flops_rk4", "Floating point operations performed in RK4") 
 67          self.coeffs = zip(_RK4A, _RK4B, _RK4C) 
 68   
 69          self.allow_jit = allow_jit 
 70   
 71 -    def add_instrumentation(self, logmgr): 
 72          logmgr.add_quantity(self.timer) 
 73          logmgr.add_quantity(self.flop_counter) 
 74   
 75 -    def __call__(self, y, t, dt, rhs): 
 76          try: 
 77              self.residual 
 78          except AttributeError: 
 79              self.residual = 0*rhs(t, y) 
 80              from hedge.tools import count_dofs, has_data_in_numpy_arrays 
 81              self.dof_count = count_dofs(self.residual) 
 82   
 83              self.use_jit = self.allow_jit and has_data_in_numpy_arrays( 
 84                      y, allow_objarray_levels=1) 
 85   
 86          if self.use_jit: 
 87              from hedge.tools import numpy_linear_comb 
 88   
 89              for a, b, c in self.coeffs: 
 90                  this_rhs = rhs(t + c*dt, y) 
 91   
 92                  sub_timer = self.timer.start_sub_timer() 
 93                  self.residual = numpy_linear_comb([(a, self.residual), (dt, this_rhs)]) 
 94                  del this_rhs 
 95                  y = numpy_linear_comb([(1, y), (b, self.residual)]) 
 96                  sub_timer.stop().submit() 
 97          else: 
 98              for a, b, c in self.coeffs: 
 99                  this_rhs = rhs(t + c*dt, y) 
100   
101                  sub_timer = self.timer.start_sub_timer() 
102                  self.residual = a*self.residual + dt*this_rhs 
103                  del this_rhs 
104                  y = y + b * self.residual 
105                  sub_timer.stop().submit() 
106   
107          self.flop_counter.add(len(self.coeffs)*self.dof_count*5) 
108   
109          return y 
110