;;; Science Collection
;;; ode-initval/rk4.rkt
;;; Copyright (c) 2004-2011 M. Douglas Williams
;;;
;;; This file is part of the Science Collection.
;;;
;;; The Science Collection 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 3 of the License
;;; or (at your option) any later version.
;;;
;;; The Science Collection is distributed in the hope that it will be useful,
;;; but WITHOUT 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 the Science Collection.  If not, see
;;; <http://www.gnu.org/licenses/>.
;;;
;;; -------------------------------------------------------------------
;;;
;;; Version  Date      Description
;;; 3.0.1    07/01/08  Added header.  (Doug Williams)
;;; 4.0.0    08/16/11  Changed the header and restructured the code. (MDW)

(require scheme/unsafe/ops)

(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.
    (vector-copy! y0 0 y)
    (if dydt-in
        (vector-copy! k 0 dydt-in)
        (ode-system-function-eval system t y0 k))
    (for ((i (in-range dim)))
      (unsafe-vector-set!
       y i (unsafe-fl* (unsafe-fl/ h 6.0) (unsafe-vector-ref k i)))
      (unsafe-vector-set!
       ytmp i (unsafe-fl+ (unsafe-vector-ref y0 i)
                          (unsafe-fl* 0.5 (unsafe-fl* h (unsafe-vector-ref k i))))))
    ;; k2 step
    (ode-system-function-eval system (unsafe-fl+ t (unsafe-fl* 0.5 h)) ytmp k)
    (for ((i (in-range dim)))
      (unsafe-vector-set!
       y i (unsafe-fl+ (unsafe-vector-ref y i)
                       (unsafe-fl* (unsafe-fl/ h 3.0) (unsafe-vector-ref k i))))
      (unsafe-vector-set!
       ytmp i (unsafe-fl+ (unsafe-vector-ref y0 i)
                          (unsafe-fl* 0.5 (unsafe-fl* h (unsafe-vector-ref k i))))))
    ;; k3 step
    (ode-system-function-eval system (unsafe-fl+ t (unsafe-fl* 0.5 h)) ytmp k)
    (for ((i (in-range dim)))
      (unsafe-vector-set!
       y i (unsafe-fl+ (unsafe-vector-ref y i)
                       (unsafe-fl* (unsafe-fl/ h 3.0) (unsafe-vector-ref k i))))
      (unsafe-vector-set!
       ytmp i (unsafe-fl+ (unsafe-vector-ref y0 i) (unsafe-fl* h (unsafe-vector-ref k i)))))
    ;; k4 step, error estimate, and final sum
    (for ((i (in-range dim)))
      (unsafe-vector-set!
       y i (unsafe-fl+ (unsafe-vector-ref y i)
                       (unsafe-fl* (unsafe-fl/ h 6.0) (unsafe-vector-ref k i))))
      (vector-set!
       y-err i (unsafe-fl* h (unsafe-vector-ref y i)))
      (vector-set!
       y i (unsafe-fl+ (unsafe-vector-ref y i) (unsafe-vector-ref y0 i)))
      (when dydt-out
        (unsafe-vector-set! dydt-out i (unsafe-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)))
    (for ((i (in-range dim)))
      (unsafe-vector-set! k i 0.0)
      (unsafe-vector-set! y0 i 0.0)
      (unsafe-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))