extra/math/runge-kutta/runge-kutta.factor

   1 USING: kernel accessors combinators sequences sequences.generalizations
   2 arrays math math.vectors ;
   3 IN: math.runge-kutta
   4
   5 : runge-kutta-2-transform ( rk1 tx..n delta -- tx..n' delta )
   6     [ swap [ [ v*n ] keep prefix 1/2 v*n ] dip v+ 2 v*n ] keep ;
   7
   8 : runge-kutta-3-transform ( rk2 tx..n delta -- tx..n' delta )
   9     runge-kutta-2-transform ;
  10
  11 : runge-kutta-4-transform ( rk3 tx..n delta -- tx..n' delta )
  12     [ swapd [ v*n ] keep prefix v+ ] keep ;
  13
  14 : (runge-kutta) ( delta tx..n dx..n/dt -- rk )
  15     swapd dup length>> [ cleave ] dip narray swap v*n
  16     ; inline
  17
  18 : runge-kutta-differentials ( dx..n/dt -- seq )
  19     '[ _ (runge-kutta) ] ;
  20
  21 : runge-kutta-transforms ( tx..n delta dx..n/dt -- seq )
  22     spin
  23     [ { [ ]
  24       [ runge-kutta-2-transform ]
  25       [ runge-kutta-3-transform ]
  26       [ runge-kutta-4-transform ] } ] dip
  27     '[ _ runge-kutta-differentials compose ] map
  28     [ 2curry ] 2with map ;
  29
  30 : increment-time ( delta tx..n -- t+dtx..n )
  31     [ swap [ 0 ] 2dip '[ _ + ] change-nth ] keep ; inline
  32
  33 : increment-state-by-approximation ( rk4 t+dtx..n -- t+dtx'..n' )
  34     swap 0 prefix v+ ;
  35
  36 : (runge-kutta-4) ( dx..n/dt delta tx..n -- tx..n' )
  37     [
  38         ! set up the set of 4 equations
  39         ! NOTE differential functions and timestep are curried
  40         runge-kutta-transforms [ [ dup ] compose ] map
  41
  42         ! using concat instead causes slow down by an order of magnitude
  43         [ ] [ compose ] reduce
  44
  45         ! this will always produce 4 outputs with a duplication of the last result
  46         ! NOTE the dup is kept in the transform so that function
  47         ! can be used in higher-order estimation
  48         call( -- rk1 rk2 rk3 rk4 rk4 ) drop 4array
  49
  50         ! make array of zeroes of appropriate length to reduce into
  51         dup first length>> 0 <array>
  52
  53         ! reduce the results to the estimated change for the timestep
  54         [ v+ ] reduce
  55         1/6 v*n
  56     ] 2keep
  57     increment-time
  58     increment-state-by-approximation
  59     ;
  60
  61 : time-limit-predicate ( t-limit -- quot: ( tx..n -- ? ) )
  62     '[ dup first _ <= ] ; inline
  63
  64 : <runge-kutta-4> ( dxn..n/dt delta initial-state t-limit -- seq )
  65     time-limit-predicate [ [ (runge-kutta-4) ] [ 3drop f ] if ] compose 2with follow ; inline