;;; PLT Scheme Science Collection
;;; ode-initval/rk4.ss
;;; Copyright (c) 2004-2008 M. Douglas Williams
;;;
;;; This library is free software; you can redistribute it and/or
;;; modify it under the terms of the GNU Lesser General Public
;;; License as published by the Free Software Foundation; either
;;; version 2.1 of the License, or (at your option) any later version.
;;;
;;; This library is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Lesser General Public License for more details.
;;;
;;; You should have received a copy of the GNU Lesser General Public
;;; License along with this library; if not, write to the Free
;;; Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
;;; 02111-1307 USA.
;;;
;;; -------------------------------------------------------------------
;;;
;;; Version  Date      Description
;;; 3.0.1    07/01/08  Added header.  (Doug Williams)

(define-values (struct:rk4-state
                rk4-state-constructor
                rk4-state?
                rk4-state-field-ref
                set-rk4-state-field!)
  (make-struct-type 'rk4-state #f 3 0))

(define (make-rk4-state dim)
  (rk4-state-constructor
   (make-vector dim 0.0)
   (make-vector dim 0.0)
   (make-vector dim 0.0)))

(define rk4-state-k
  (make-struct-field-accessor rk4-state-field-ref 0 'k))
(define set-rk4-state-k!
  (make-struct-field-mutator set-rk4-state-field! 0 'k))

(define rk4-state-y0
  (make-struct-field-accessor rk4-state-field-ref 1 'y0))
(define set-rk4-state-y0!
  (make-struct-field-mutator set-rk4-state-field! 1 'y0))

(define rk4-state-ytmp
  (make-struct-field-accessor rk4-state-field-ref 2 'ytmp))
(define set-rk4-state-ytmp!
  (make-struct-field-mutator set-rk4-state-field! 2 'ytmp))

(define (rk4-apply state dim t h y y-err dydt-in dydt-out system)
  (let ((k (rk4-state-k state))
        (y0 (rk4-state-y0 state))
        (ytmp (rk4-state-ytmp state)))
    ;; Copy the starting value.  We will write over the y[] vector,
    ;; using it for scratch and then filling it with the final result.
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y0 i (vector-ref y i)))
    (if dydt-in
        (do ((i 0 (+ i 1)))
            ((= i dim) (void))
          (vector-set! k i (vector-ref dydt-in i)))
        (ode-system-function-eval system t y0 k))
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (* (/ h 6.0) (vector-ref k i)))
      (vector-set! ytmp i (+ (vector-ref y0 i) (* 0.5 h (vector-ref k i)))))
    ;; k2 step
    (ode-system-function-eval system (+ t (* 0.5 h)) ytmp k)
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (+ (vector-ref y i) (* (/ h 3.0) (vector-ref k i))))
      (vector-set! ytmp i (+ (vector-ref y0 i) (* 0.5 h (vector-ref k i)))))
    ;; k3 step
    (ode-system-function-eval system (+ t (* 0.5 h)) ytmp k)
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (+ (vector-ref y i) (* (/ h 3.0) (vector-ref k i))))
      (vector-set! ytmp i (+ (vector-ref y0 i) (* h (vector-ref k i)))))
    ;; k4 step, error estimate, and final sum
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! y i (+ (vector-ref y i) (* (/ h 6.0) (vector-ref k i))))
      (vector-set! y-err i (* h (vector-ref y i)))
      (vector-set! y i (+ (vector-ref y i) (vector-ref y0 i)))
      (when dydt-out
        (vector-set! dydt-out i (vector-ref k i))))))

(define (rk4-reset state dim)
  (let ((k (rk4-state-k state))
        (y0 (rk4-state-y0 state))
        (ytmp (rk4-state-ytmp state)))
    (do ((i 0 (+ i 1)))
        ((= i dim) (void))
      (vector-set! k i 0.0)
      (vector-set! y0 i 0.0)
      (vector-set! ytmp i 0.0))))

(define (rk4-order state)
  4)

(define rk4-ode-type
  (make-ode-step-type
   "rk4"
   #t
   #f
   make-rk4-state
   rk4-apply
   rk4-reset
   rk4-order))