;;;;;;;;;;;;;;;;;;; -*- Mode: Lisp; Package: Macsyma -*- ;;;;;;;;;;;;;;;;;;;
;;;     (c) Copyright 1981 Massachusetts Institute of Technology         ;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(macsyma-module ndiffq)

(load-macsyma-macros numerm)

;;; Some numerical differential equation routines.

(defmfun $init_float_array (array x0 x1 &aux
				     (a (get-array array '(flonum) 1)))
  (setq x0 (float x0)
	x1 (float x1))
  (let ((n (array-dimension-n 1 a)))
    (do ((j 0 (1+ j))
	 (h (//$ (-$ x1 x0) (float (1- n))))
	 (x x0 (+$ x h)))
	((= j n) array)
      (setf (aref$ a j) x))))

(defmfun $map_float_array (ya f xa)
  (let* ((y (get-array ya '(flonum) 1))
	 (n (array-dimension-n 1 y))
	 (x (get-array xa '(flonum) 1 n)))
    (bind-tramp1$
     f f
     (do ((j 0 (1+ j)))
	 ((= j n) ya)
       (setf (aref$ y j) (fcall$ f (aref$ x j)))))))
  
;;; Runge-Kutta method for getting starting values.

(defvar runge-^]-int nil)
(defun runge-^]-int () (setq runge-^]-int t))

(defun $runge_kutta (f x y &rest higher-order)
  (let  ((runge-^]-int nil)
	 (USER-TIMESOFAR (CONS #'runge-^]-int USER-TIMESOFAR)))
    (if ($listp f)
	(if higher-order
	    (merror "Runge_Kutta handles systems of order 1 only.")
	    (let* ((fl (mapcar #'(lambda (f) (make-gtramp$ f 2)) (cdr f)))
		   (xa (get-array x '(flonum) 1))
		   (n (array-dimension-n 1 xa)))
	      (if (and ($listp y)
		       (= (length fl) (length (cdr y))))
		  (runge-kutta-1-n fl xa
				   (mapcar #'(lambda (y)
					       (get-array y '(flonum) 1 n))
					   (cdr y)))
		  (merror "Not a list of length ~M~%~M" (length fl) y))))
	(let* ((xa (get-array x '(flonum) 1))
	       (n (array-dimension-n 1 xa))
	       (ya (get-array y '(flonum) 1 n)))
	  (caseq (length higher-order)
	    ((0)
	     (bind-tramp2$
	      f f
	      (runge-kutta-1 f xa ya)))
	    ((1)
	     (bind-tramp3$
	      f f
	      (runge-kutta-2 f xa ya
			     (get-array (car higher-order) '(flonum) 1 n))))
	    (t
	     (merror "Runge_Kutta of order greater than 2 is unimplemented"))))))
  ;; return value to user.
  y)

(defvar one-half$   (//$ 1.0 2.0))
(defvar one-third$  (//$ 1.0 3.0))
(defvar one-sixth$  (//$ 1.0 6.0))
(defvar one-eighth$ (//$ 1.0 8.0))

(DEFVAR RUNGE-KUTTA-1 NIL)

(defun runge-kutta-1 (f x y)
  (do ((m-1 (1- (array-dimension-n 1 x)))
       (n 0 (1+ n))
       (x_n)(y_n)(h)(k1)(k2)(k3)(k4))
      ((= n m-1))
    (declare (fixnum n-1 n)
	     (flonum x_n y_n h k1 k2 k3 k4))
    (setq x_n (aref$ x n))
    (setq y_n (aref$ y n))
    (WHEN RUNGE-^]-INT
	  (SETQ RUNGE-^]-INT NIL)
	  (MTELL "~A steps, calculating F(~A,~A)" N X_N Y_N))
    (setq h (-$ (aref$ x (1+ n)) x_n))
    ;; Formula 25.5.10 pp 896 of Abramowitz & Stegun.
    (setq k1 (*$ h (fcall$ f x_n y_n)))
    (setq k2 (*$ h (fcall$ f
			   (+$ x_n (*$ one-half$ h))
			   (+$ y_n (*$ one-half$ k1)))))
    (setq k3 (*$ h (fcall$ f
			   (+$ x_n (*$ one-half$ h))
			   (+$ y_n (*$ one-half$ k2)))))
    (setq k4 (*$ h (fcall$ f
			   (+$ x_n h)
			   (+$ y_n k3))))
    (setf (aref$ y (1+ n))
	  (+$ y_n (*$ one-sixth$ (+$ k1 k4))
	          (*$ one-third$ (+$ k2 k3))))))

(defun runge-kutta-2 (f x y y-p)
  (do ((m-1 (1- (array-dimension-n 1 x)))
       (n 0 (1+ n))
       (x_n)(y_n)(y-p_n)(h)(k1)(k2)(k3)(k4))
      ((= n m-1))
    (declare (fixnum m-1 n)
	     (flonum x_n y_n y-p_n h k1 k2 k3 k4))
    (setq x_n (aref$ x n))
    (setq y_n (aref$ y n))
    (setq y-p_n (aref$ y-p n))
    (WHEN RUNGE-^]-INT
	  (SETQ RUNGE-^]-INT NIL)
	  (MTELL "~A steps, calculating F(~A,~A,~A)" N X_N Y_N Y-P_N))
    (setq h (-$ (aref$ x (1+ n)) x_n))
    ;; Formula 25.5.20 pp 897 of Abramowitz & Stegun.
    (setq k1 (*$ h (fcall$ f x_n y_n y-p_n)))
    (setq k2 (*$ h (fcall$ f
			   (+$ x_n (*$ one-half$ h))
			   (+$ y_n (*$ one-half$ h y-p_n)
			           (*$ one-eighth$ h k1))
			   (+$ y-p_n (*$ one-half$ k1)))))
    (setq k3 (*$ h (fcall$ f
			   (+$ x_n (*$ one-half$ h))
			   (+$ y_n (*$ one-half$ h y-p_n)
			           (*$ one-eighth$ h k1))
			   (+$ y-p_n (*$ one-half$ k2)))))
    (setq k4 (*$ h (fcall$ f
			   (+$ x_n h)
			   (+$ y_n (*$ h y-p_n)
			           (*$ one-half$ h k3))
			   (+$ y-p_n k3))))
    (setf (aref$ y (1+ n))
	  (+$ y_n (*$ h (+$ y-p_n (*$ one-sixth$ (+$ k1 k2 k3))))))
    (setf (aref$ y-p (1+ n))
	  (+$ y-p_n (+$ (*$ one-third$ (+$ k2 k3))
			(*$ one-sixth$ (+$ k1 k4)))))))

(defun runge-kutta-1-n (fl x yl
			   &aux
			   (m (array-dimension-n 1 x))
			   (d (length fl)))
  (do ((m-1 (1- m))
       (n 0 (1+ n))
       (h)
       (x_n)
       (y_n (make-array$ d))
       (K1 (make-array$ d))
       (K2 (make-array$ d))
       (K3 (make-array$ d))
       (K4 (make-array$ d))
       (ACC (make-array$ d)))
      ((= n m-1)
       (free-array$ y_n)
       (free-array$ k1)
       (free-array$ k2)
       (free-array$ k3)
       (free-array$ k4)
       (free-array$ acc)
       nil)
    (declare (fixnum m-1 n) (flonum x_n h))
    (setq x_n (aref$ x n))
    (when (= n 0)
	  (do ((l yl (cdr l))
	       (j 0 (1+ j)))
	      ((null l))
	    (setf (aref$ y_n j) (aref$ (car l) n))))
    (WHEN RUNGE-^]-INT
	  (SETQ RUNGE-^]-INT NIL)
	  (MTELL "~A steps, calculating ~M" n
		 `(($F) ,x_n ,@(listarray y_n))))
    (setq h (-$ (aref$ x (1+ n)) x_n))
    (gvapply$-x-ar$ k1 fl x_n y_n)
    (ar$*s k1 k1 h)					
    (ar$*s acc k1 one-half$)
    (ar$+ar$ acc acc y_n)
    (gvapply$-x-ar$ k2 fl (+$ x_n (*$ h one-half$)) acc)
    (ar$*s k2 k2 h)
    (ar$*s acc k2 one-half$)
    (ar$+ar$ acc acc y_n)
    (gvapply$-x-ar$ k3 fl (+$ x_n (*$ h one-half$)) acc)
    (ar$*s k3 k3 h)
    (ar$+ar$ acc k3 y_n)
    (gvapply$-x-ar$ k4 fl (+$ x_n h) acc)
    (ar$*s k4 k4 h)
    (ar$+ar$ k1 k1 k4)
    (ar$*s k1 k1 one-sixth$)
    (ar$+ar$ k2 k2 k3)
    (ar$*s k2 k2 one-third$)
    (ar$+ar$ y_n y_n k1)
    (ar$+ar$ y_n y_n k2)
    (do ((l yl (cdr l))
	 (j 0 (1+ j)))
	((null l))
      (setf (aref$ (car l) (1+ n)) (aref$ y_n j)))))