;;
 ;; Runge-Kutta integration of ODEs
 ;; Based on Haskell code by Uwe Hollerbach <uh@alumni.caltech.edu>
 ;;
 ;; dy
 ;; -- = f(t, y)
 ;; dt
 ;;
 ;; y_{n+1} = y_n + h sum_{i=1}^s b_i k_i
 ;; k_i = f(t_n + c_i h, y_n + h sum_{j=1}^s a_{ij} k_j)
 ;; ;; "Butcher Tableau" is
 ;;
 ;; c_1  a_11 a_12 ... a_1s
 ;; c_2  a_21 a_22 ... a_2s
 ;; ...  ...
 ;; c_s  a_s1 a_s2 ... a_ss
 ;;      b_1  b_2  ... b_s
 ;;
 ;; This module implements a method that can do a generic tableau, then
 ;; specializes with different tableaux to let the user pick a specific
 ;; method. Adaptive step-size methods, see below, add a row of d_j
 ;; coefficients and use that to report the error:
 ;;
 ;; e_{n+1} = h sum_{i=1}^s d_i k_i
 ;;
 ;; non-adaptive solvers:
 ;;	rkfe, rk3, rk4a, rk4b
 ;;
 ;; adaptive solvers:
 ;;	rkhe, rkbs, rkf45, rkck, rkdp, rkf78, rkv65
 ;;
 ;; auxiliary non-adaptive solvers (error estimators from the adaptive ones):
 ;;	rkhe_aux, rkbs_aux, rkf45_aux, rkck_aux, rkdp_aux, rkf78_aux, rkv65_aux
 ;;
 ;; use rk4[ab] if you don't need an adaptive solver, rkdp or rkv65 if you do;
 ;; or use what you need if you're an expert.
 ;;
 ;; DO NOT USE rkfe EXCEPT FOR DEMONSTRATIONS OF INSTABILITY!
 ;; (Or if you're an expert.)
 ;;
 ;; Reference: E. Hairer, S. P. Norsett, G. Wanner,
 ;; Solving Ordinary Differential Equations I: Nonstiff Problems
 ;; (second revised edition, 1993).
 ;;

(module runge-kutta

	(
         (core1 k-sum gen-ks )
	 (core2 k-sum gen-ks )
	 make-rkfe make-rk3 make-rk4a make-rk4b
	 make-he make-he-aux
	 make-bs make-bs-aux
	 make-rkf45 make-rkf45-aux
	 diffs ratlist->reals ratlist->RCL ratlists->RCLs R //
	 )

	(import scheme chicken)
	(require-library data-structures srfi-1 numbers)
	(import (only srfi-1 fold filter-map)
		(only data-structures compose)
		(only numbers inexact->exact exact->inexact denominator numerator gcd + - * /))


(define (foldl1 f lst)
  (cond ((pair? lst) 
	 (if (pair? (cdr lst))
	     (fold f (f (car lst) (cadr lst)) (cddr lst))
	     (car lst)))
	(else (error 'foldl1  "insufficient arguments"))))


;; Store a list of rational numbers as a common denominator, then a
;; list of numerators, all stored as doubles: this lets us write the
;; values in the source as exact rational numbers, yet we can take
;; advantage of not having to do too many conversions and also of not
;; having to multiply by 0 or 1 in some cases.


;; ratToRCL :: [rational] -> RCL 

(define (ratlist->RCL rs)
  (if (null? rs)  '(1.0 ())
      (let* ((ds (map denominator rs))
	     (dp (foldl1 * ds))
	     (ns (map (compose numerator (lambda (x) (* dp x))) rs))
	     (g  (fold gcd dp ns)))
	`(,(exact->inexact (/ dp g)) 
	  ,(map (compose exact->inexact (lambda (x) (/ x g))) ns)))))


;; ratToRCLs :: [[rational]] -> [RCL] 
(define (ratlists->RCLs x) (map ratlist->RCL x))

;; ratToReals :: [rational] -> [real] *)
(define (ratlist->reals x) (map exact->inexact x))


;; Helper function to sum a list of K_i, skipping unnecessary multiplications and additions 
;; sc_fn: real * 'a -> 'a
;; sum_fn: 'a * 'a -> 'a

(define (k-sum sc-fn sum-fn h)
  (define (mscale s v)
    (cond ((fp= s 0.0) #f)
	  ((fp= s 1.0) v)
	  (else (sc-fn s v))))
  (lambda (d+ns ks)
    (let ((d (car d+ns)) (ns (cadr d+ns)))
      (sc-fn (fp/ h d) (foldl1 sum-fn (filter-map mscale ns ks)))
      )))

;; Helper function to generate a list of K_i 
;; ksum_fn,sum_fn: 'a * 'a -> 'a
;; der_fn: real * 'a -> 'a

(define (gen-ks ksum-fn sum-fn der-fn h tn+yn ks cs as)
  (if (and (null? cs) (null? as)) ks
      (let* (
	     (tn (car tn+yn))
	     (yn (cadr tn+yn))
	     (yn1 (if (null? ks) yn
		      (let ((a (car as)))
			(sum-fn yn (ksum-fn a ks))))))
	(let ((ar (cdr as))
	      (c  (car cs))
	      (cr (cdr cs)))
	  (let ((ks1 (list (der-fn (fp+ tn (fp* c h)) yn1))))
	    (gen-ks ksum-fn sum-fn der-fn h tn+yn (append ks ks1) cr ar)
	    )))))


	  
;;    This is the first core routine: it does not get used directly, only
;;    partial applications of it; see below.
;;
;;    Its arguments are:
;;
;;      c table (specified internally)
;;      a table (specified internally)
;;      b table (specified internally)
;;
;;    user-specified arguments:
;;
;;      scale function to scale a Y state vector ::
;; 	(real * a -> a)
;;
;;      sum function to add two Y state vectors ::
;; 	(a * a -> a)
;;
;;      derivative function F ::
;; 	(real * a -> a)
;;
;;      step size H ::
;; 	real
;;
;;      current state (T,Y) ::
;; 	(real, a)
;;
;;      and the return value is the new state (T_new,Y_new)
;; 



;;  cl: real list, al: RCL list, bl: RCL
;;  sc_fn:  real * 'a -> 'a, 
;;  sum_fn: 'a * 'a -> 'a,
;;  der_fn: real * 'a -> 'a)
;;  h: real 
;;  old as (tn: real,yn: 'a)

(define-syntax core1
  (syntax-rules ()
    [(_  cl al bl)
     (lambda (sc-fn sum-fn der-fn)
       (lambda (h)
	 (let ((ksum (k-sum sc-fn sum-fn h)))
	   (lambda (tn yn)
	     (let ((ks (gen-ks ksum sum-fn der-fn h (list tn yn) '() cl al)))
	       (sum-fn yn (ksum bl ks)))
	     ))))]
    ))

;;   This is the second core routine, analogous to the previous one.
;;   The difference is that this gets an additional internal table arg,
;;   and it returns a 3-tuple instead of a 2-tuple: (tnew,ynew,enew),
;;   where enew is the error state vector
;;
;;   e_{n+1} = h sum_{i=1}^s (b_i - b'_i) k_i


;; dl: RCL

(define-syntax core2
  (syntax-rules ()
    [(_  cl al bl dl)
     (lambda (sc-fn sum-fn der-fn)
       (lambda (h)
	 (let ((ksum (k-sum sc-fn sum-fn h)))
	   (lambda (tn yn)
	     (let ((ks (gen-ks ksum sum-fn der-fn h (list tn yn) '() cl al)))
	       (list (sum-fn yn (ksum bl ks)) (ksum dl ks))
	       )))))]
    ))


;;   Some specific explicit methods, taken from
;;   "List of Runge-Kutta methods" at Wikipedia

(define R (lambda (x) x))
(define // /)



;; forward Euler: unconditionally unstable: don't use this! 

(define-syntax make-rkfe
  (syntax-rules ()
    [(_)
     (let ((cs_fe (ratlist->reals `(,[R 0])))
	   (as_fe (ratlists->RCLs `(())))
	   (bs_fe (ratlist->RCL   `(,[R 1]))))
       (core1 cs_fe as_fe bs_fe))]))

;; Kutta's third-order method: 
(define-syntax make-rk3
  (syntax-rules ()
    [(_)
     (let ((cs_rk3 (ratlist->reals `(,[R 0] ,[// 1 2] ,[R 1])))
	   (as_rk3 (ratlists->RCLs `(() (,[// 1 2]) (,[R -1] ,[R 2]))))
	   (bs_rk3 (ratlist->RCL   `(,[// 1 6] ,[// 2 3] ,[// 1 6]))))
       (core1 cs_rk3 as_rk3 bs_rk3))]))

;; Classic fourth-order method 
(define-syntax make-rk4a
  (syntax-rules ()
    [(_)
     (let ((cs_rk4a (ratlist->reals `(,[R 0] ,[// 1 2] ,[// 1 2] ,[R 1])))
	   (as_rk4a (ratlists->RCLs `(() (,[// 1 2]) (,[R 0] ,[// 1 2]) (,[R 0] ,[R 0] ,[R 1]))))
	   (bs_rk4a (ratlist->RCL   `(,[// 1 6] ,[// 1 3] ,[// 1 3] ,[// 1 6]))))
       (core1 cs_rk4a as_rk4a bs_rk4a))]))

;; Kutta's other fourth-order method... "The first [above] is more popular,
;; the second is more precise." (Hairer, Norsett, Wanner) 
(define-syntax make-rk4b
  (syntax-rules ()
    [(_)
     (let ((cs_rk4b (ratlist->reals `(,[R 0] ,[// 1 3] ,[// 2 3] ,[R 1])))
	   (as_rk4b (ratlists->RCLs `(() (,[// 1 3]) (,[// -1 3] ,[R 1]) (,[R 1] ,[R -1] ,[R 1]))))
	   (bs_rk4b (ratlist->RCL   `(,[// 1 8] ,[// 3 8] ,[// 3 8] ,[// 1 8]))))
       (core1 cs_rk4b as_rk4b bs_rk4b))]))


;;   Some adaptive-stepsize methods, also from Wikipedia; more from
;;   HNW.  These don't auto-adapt, but they do allow the user to make
;;   a somewhat intelligent decision about what the step size ought to
;;   be at each step.


;;   A helper function to take the difference of two lists of
;;   rationals: we don't want to use zipWith (-) because that gives
;;   only the head where both lists have entries; we want implicit
;;   zeros at the end, as far as is necessary.

(define (diffs a b)
  (cond ((and (null? a) (null? b))  '())
	((and (pair? a) (null? b))  a)
	((and (null? a) (pair? b))  (map - b))
	(else (cons (+ (car a) (- (car b))) (diffs (cdr a) (cdr b))))))
	
;; Heun-Euler, order 2/1 

(define-syntax cs_he (syntax-rules () [(_) (ratlist->reals `(,[R 0] ,[R 1]))]))
(define-syntax as_he (syntax-rules () [(_) (ratlists->RCLs `(() (,[R 1])))]))

(define-syntax make-he
  (syntax-rules ()
    [(_)
     (let ((cs (cs_he))
	   (as (as_he))
	   ;; second order coefficients
	   (r1_he `(,[// 1 2] ,[// 1 2]))
	   ;; first-order coefficients
	   (r2_he `(,[R 1])))
       (let ((bs (ratlist->RCL r1_he))
	     (ds (ratlist->RCL (diffs r1_he r2_he))))
	 (core2 cs as bs ds)))]
    ))


(define-syntax make-he-aux
  (syntax-rules ()
    [(_)
     (let ((r2_he `(,[R 1])))
       (let ((cs (cs_he))
	     (as (as_he))
	     (bs (ratlist->RCL r2_he)))
	 (core1 cs as bs)))]
    ))



;; Bogacki-Shampine, order 3/2 

(define-syntax cs_bs (syntax-rules () [(_) (ratlist->reals  `(,[R 0] ,[// 1 2] ,[// 3 4] ,[R 1]))]))
(define-syntax as_bs (syntax-rules () [(_) (ratlists->RCLs `(() 
							     (,[// 1 2]) 
							     (,[R 0]  ,[// 3 4]) 
							     (,[// 2 9] ,[// 1 3] ,[// 4 9])))]))

(define-syntax make-bs
  (syntax-rules ()
    [(_)
     (let ((cs (cs_bs))
	   (as (as_bs))
	   ;; third-order coefficients
	   (r1_bs  `(,[// 2 9] ,[// 1 3] ,[// 4 9]))
	   ;; second-order coefficients
	   (r2_bs `(,[// 7 24] ,[// 1 4] ,[// 1 3] ,[// 1 8])))
       (let ((bs (ratlist->RCL r1_bs))
	     (ds (ratlist->RCL (diffs r1_bs r2_bs))))
	 (core2 cs as bs ds)))]
    ))

(define-syntax make-bs-aux
  (syntax-rules ()
    [(_)
     (let ((cs (cs_bs))
	   (as (as_bs))
	   (r2_bs `(,[// 7 24] ,[// 1 4] ,[// 1 3] ,[// 1 8])))
       (let ((bs (ratlist->RCL r2_bs)))
	 (core1 cs as bs)))]
    ))

;; Runge-Kutta-Fehlberg, order 4/5 
(define-syntax cs_rkf45
  (syntax-rules () [(_)  (ratlist->reals  `(,[R 0] ,[// 1 4] ,[// 3 8] ,[// 12 13] ,[R 1] ,[// 1 2]))]))
(define-syntax as_rkf45 
   (syntax-rules () [(_) 
		     (ratlists->RCLs `(() 
				       (,[// 1 4]) 
				       (,[// 3 32]      ,[// 9 32])
				       (,[// 1932 2197] ,[// -7200 2197] ,[// 7296 2197])
				       (,[// 439 216]   ,[R -8]        ,[// 3680 513]   ,[// -845 4104])
				       (,[// -8 27]     ,[R 2]         ,[// -3544 2565] ,[// 1859 4104] ,[// -11 40])))]))



(define-syntax make-rkf45
  (syntax-rules ()
    [(_)
     (let ((cs (cs_rkf45))
	   (as (as_rkf45))
	   ;; fourth-order coefficients
	   (r1_rkf  `(,[// 25 216] ,[R 0] ,[// 1408 2565] ,[// 2197 4104] ,[// -1 5]))
	   ;; fifth-order coefficients 
	   (r2_rkf  `(,[// 16 135] ,[R 0] ,[// 6656 12825] ,[// 28561 56430] ,[// -9 50] ,[// 2 55])))
       (let ((bs (ratlist->RCL r1_rkf))
	     (ds (ratlist->RCL (diffs r1_rkf r2_rkf))))
	 (core2 cs as bs ds)))]
    ))

(define-syntax make-rkf45-aux
  (syntax-rules ()
    [(_)
     (let ((cs (cs_rkf45))
	   (as (as_rkf45))
	   (r2_rkf  `(,[// 16 135] ,[R 0] ,[// 6656 12825] ,[// 28561 56430] ,[// -9 50] ,[// 2 55])))
       (let ((bs (ratlist->RCL r2_rkf)))
	 (core1 cs as bs)))]
    ))



)