;;; Matrix arithmetic. -*- Mode:Common-Lisp; Package:MATH; Base:8. -*-

;;;                           RESTRICTED RIGHTS LEGEND

;;;Use, duplication, or disclosure by the Government is subject to
;;;restrictions as set forth in subdivision (c)(1)(ii) of the Rights in
;;;Technical Data and Computer Software clause at 52.227-7013.

;;;                     TEXAS INSTRUMENTS INCORPORATED.
;;;                              P.O. BOX 2909
;;;                           AUSTIN, TEXAS 78769
;;;                                 MS 2151

;;; Copyright (C) 1985-1989 Texas Instruments Incorporated. All rights reserved.

;;Revised:
;;  8/06/87 DNG - Fixed INVERT-MATRIX to use / instead of QUOTIENT so that it 
;;		works right for integers.  [SPR 5674]

(export '(multiply-matrices invert-matrix transpose-matrix determinant decompose
			    solve list-2d-array fill-2d-array))

;;; Convert a 2d array into a list of lists of the elements
(defun list-2d-array (array)
  (check-arg array
	     (and (arrayp array)
		  (= (array-rank array) #o2))
	     "a two-dimensional array")
  (do ((i #o0 (1+ i))
       (dim-1 (array-dimension array #o0))
       (dim-2 (array-dimension array #o1))
       (list ()))
      ((>= i dim-1)
       (nreverse (the list list)))
    (push
      (do ((j #o0 (1+ j))
	   (list ()))
	  ((>= j dim-2)
	   (nreverse (the list list)))
	(push (aref array i j) list))
      list))) 

;;; Fill up a 2d array from a list.  Like FILLARRAY, the lists can wrap around as needed.
(defun fill-2d-array (array list)
  (check-arg array
	     (and (arrayp array)
		  (= (array-rank array) #o2))
	     "a two-dimensional array")
  (do ((i #o0 (1+ i))
       (dim-1 (array-dimension array #o0))
       (dim-2 (array-dimension array #o1))
       (l list (cdr l))
       (sublist))
      ((>= i dim-1))
    (and (null l)
	 (setq l list))
    (setq sublist (car l))
    (do ((j #o0 (1+ j))
	 (l sublist (cdr l)))
	((>= j dim-2))
      (and (null l)
	   (setq l sublist))
      (setf (aref array i j) (car l))))) 

;;; Multiply two matrices into a third.  
;;; A 1d array of dimension N is treated as a Nx1 array.
(defun multiply-matrices (matrix-1 matrix-2 &optional matrix-3 &aux saved-matrix-3)
  "Multiply matrices MATRIX-1 and MATRIX-2, storing into MATRIX-3 if supplied.
If MATRIX-3 is not supplied, then a new (ART-Q type) array is returned, else
MATRIX-3 must have exactly the right dimensions for holding the result of the multiplication.
Both MATRIX-1 and MATRIX-2 must be either one- or two-diimensional.
The first dimension of MATRIX-2 must equal the second dimension of MATRIX-1, unless MATRIX-1
is one-dimensional, when the first dimensions must match (thus allowing multiplications of the
form VECTOR x MATRIX)"
  (check-arg matrix-1 1-or-2d-arrayp "a one- or two-dimensional array")
  (check-arg matrix-2 1-or-2d-arrayp "a one- or two-dimensional array")
  (check-arg matrix-3
	     (or (null matrix-3)
		 (1-or-2d-arrayp matrix-3))
	     "a one- or two-dimensional array or NIL")
  (let ((dim-1-1 (if (= (array-rank matrix-1) #o1) #o1 (array-dimension matrix-1 #o0)))
	(dim-1-2 (if (= (array-rank matrix-1) #o1)
		     (array-dimension matrix-1 #o0)
		     (array-dimension matrix-1 #o1)))
	(dim-2-1 (array-dimension matrix-2 #o0))
	(dim-2-2 (if (= (array-rank matrix-2) #o1) #o1 (array-dimension matrix-2 #o1)))
	(dim-3-1 (when matrix-3
		   (if (= (array-rank matrix-3) #o1) #o1 (array-dimension matrix-3 #o0))))
	(dim-3-2 (when matrix-3
		   (if (= (array-rank matrix-3) #o1)
		       (array-dimension matrix-3 #o0)
		       (array-dimension matrix-3 #o1)))))
    (unless (= dim-2-1 dim-1-2)
      (ferror nil "The ~~Dx~D matrix ~S and the ~Dx~D matrix ~S cannot be multiplied~"
	      dim-1-1 dim-1-2 matrix-1 dim-2-1 dim-2-2 matrix-2))
    (if matrix-3
	(if (and (= dim-1-1 dim-3-1)
		 (= dim-2-2 dim-3-2))
	    ;; We have a destination; see if it's the same as one of the sources.
	    ;; If it is, substitute a temporary for the destination.  We only check
	    ;; EQness, not displacements.
	    (when (let ((forwarded (follow-structure-forwarding matrix-3)))
		    (or (eq forwarded (follow-structure-forwarding matrix-1))
			(eq forwarded (follow-structure-forwarding matrix-2))))
	      (setq saved-matrix-3 matrix-3
		    matrix-3 (make-array (array-dimensions matrix-3) :type (array-type matrix-3))))
	    (ferror nil
		    "The ~~Dx~D matrix ~S is not the right size for multiplying the
~Dx~D matrix ~S and the
~Dx~D matrix ~S.~"
		    dim-3-1 dim-3-2 matrix-3 dim-1-1 dim-1-2 matrix-1 dim-2-1 dim-2-2 matrix-2))
	;; We don't make a 1xn matrix here since the user probably wants a vector result.
	(setq matrix-3
	      (make-array (if (= (array-rank matrix-1) #o1)
			      dim-2-2
			      (list dim-1-1 dim-2-2)))))
    ;; Make indirect arrays to any vectors, so can use ar-2 everywhere below.
    (let ((mat-1 (if (= #o2 (array-rank matrix-1))
		     matrix-1
		     (make-array (list dim-1-1 dim-2-1)
				 :type (array-type matrix-1)
				 :displaced-to matrix-1)))
	  (mat-2 (if (= #o2 (array-rank matrix-2))
		     matrix-2
		     (make-array (list dim-2-1 dim-2-2)
				 :type (array-type matrix-2)
				 :displaced-to matrix-2)))
	  (mat-3 (if (= #o2 (array-rank matrix-3))
		     matrix-3
		     (make-array (list dim-1-1 dim-2-2)
				 :type (array-type matrix-3)
				 :displaced-to matrix-3))))
      ;; Do the actual multiplication
      (dotimes (i dim-1-1)
	(dotimes (j dim-2-2)
	  (setf (aref mat-3 i j)
		(do ((k #o0 (1+ k))
		     (sum #o0 (+ sum (* (aref mat-1 i k) (aref mat-2 k j)))))
		    ((>= k dim-2-1)
		     sum)))))
      ;; Try to get rid of any temporary arrays we made
      (when (neq matrix-3 mat-3)
	(return-array (prog1
			mat-3
			(setf mat-3 nil))))
      (when (neq matrix-2 mat-2)
	(return-array (prog1
			mat-2
			(setf mat-2 nil))))
      (when (neq matrix-1 mat-1)
	(return-array (prog1
			mat-1
			(setf mat-1 nil))))
      ;; If we substituted a temporary above, copy the result into the saved
      ;; destination and return the temporary to free storage.
      (when saved-matrix-3
	(copy-array-contents matrix-3 saved-matrix-3)
	(return-array (prog1
			matrix-3
			(setf matrix-3 nil)))
	(setq matrix-3 saved-matrix-3))))
  matrix-3) 

;;; Gauss-Jordan inversion
(defun invert-matrix (matrix &optional into-matrix)
  "If matrix is singular returns nil, else returns its inverse.
   If into-matrix is supplied, inverse is returned in it,
    otherwise a new array is created."
  (let* ((dims (array-dimensions matrix))
	 (rows (car dims))
	 (cols (cadr dims)))
    (unless (= rows cols)
      (error () "Matrix ~s not square" matrix))
    (if (null into-matrix)
	(setq into-matrix (make-array dims :type (array-type matrix)))
	(let ((dims-into (array-dimensions into-matrix)))
	  (unless (and (= rows (car dims-into)) (= cols (cadr dims-into)))
	    (error () "into-matrix ~s does not have same dimensions as matrix ~s" into-matrix
		   matrix))))
    (let ((xcols (+ rows cols))
	  (vv (make-array rows)))
      (dotimes (r rows)
	(let ((v (make-array xcols :initial-value #o0 :type (array-type matrix))))
	  (setf (aref v (+ r rows)) #o1)
	  (dotimes (c cols)
	    (setf (aref v c) (aref matrix r c)))
	  (setf (aref vv r) v)))
      (dotimes (r rows)
	(let ((imax #o-1)
	      (vmax #o-1))
	  (do ((i r (1+ i)))
	      ((= i rows))
	    (let ((v (abs (aref (aref vv i) r))))
	      (cond
		((> v vmax) (setq vmax v) (setq imax i)))))
	  (cond ((zerop vmax) (return-from invert-matrix ())))
	  (cond ((not (= r imax))
		 (let ((exch-temp (aref vv r)))
		   (setf (aref vv r) (aref vv imax))
		   (setf (aref vv imax) exch-temp)))))
	(let ((pivot-row (aref vv r)))
	  (do ((d (aref pivot-row r))
	       (i (1+ r) (1+ i)))
	      ((= i xcols))
	    (setf (aref pivot-row i)
		  (/ (aref pivot-row i) d)))
	  (do ((i (1+ r) (1+ i)))
	      ((= i rows))
	    (let ((row (aref vv i)))
	      (do ((d (aref row r))
		   (j (1+ r) (1+ j)))
		  ((= j xcols))
		(setf (aref row j) (- (aref row j) (* d (aref pivot-row j)))))))))
      (do ((r (1- rows) (1- r)))
	  ((= r #o0))
	(let ((pivot-row (aref vv r)))
	  (do ((i (1- r) (1- i)))
	      ((< i #o0))
	    (let ((row (aref vv i)))
	      (do ((d (aref row r))
		   (j rows (1+ j)))
		  ((= j xcols))
		(setf (aref row j) (- (aref row j) (* d (aref pivot-row j)))))))))
      (dotimes (r rows)
	(let ((v (aref vv r)))
	  (dotimes (c cols)
	    (setf (aref into-matrix r c) (aref v (+ c cols)))))))
    into-matrix))  

(defun matrix-norm (matrix)
  "Returns the norm of matrix.
   The norm is the maximum over the rows of the sum of the abs of the columns."
  (let* ((dim (array-dimensions matrix))
         (rows (car dim))
         (cols (cadr dim))
         (m 0))
    (dotimes (r rows)
      (let ((s 0))
        (dotimes (c cols)
          (incf s (float (abs (aref matrix r c)))))
        (if (> s m) (setf m s))))
    m))

(defun invert-matrix-iterate (matrix &optional into-matrix)
  "If matrix is singular returns nil, else returns the inverse of matrix.
   Uses iterative improvement until no further improvement is possible."
  (let* ((dim (array-dimensions matrix))
	 (rows (car dim))
	 (cols (cadr dim)))
    (unless (= rows cols)
      (error () "matrix ~s not square" matrix))
    (prog ((d (invert-matrix matrix into-matrix))
	   (f (make-array dim))
	   (dt (make-array dim))
	   n
	   nt)
	  (unless d (return nil))
	  (dotimes (r rows)
	    (dotimes (c cols)
	      (let ((s (if (= r c) #o1 #o0)))
		(dotimes (i rows)
		  (decf s (* (aref matrix r i) (aref d i c))))
		(setf (aref f r c) s))))
	  (setq n (matrix-norm f))
       la
	  (dotimes (r rows)
	    (dotimes (c cols)
	      (let ((s #o0))
		(dotimes (i rows)
		  (incf s (* (aref d r i) (aref f i c))))
		(setf (aref dt r c) (+ s (aref d r c))))))
	  (dotimes (r rows)
	    (dotimes (c cols)
	      (let ((s (if (= r c) #o1 #o0)))
		(dotimes (i rows)
		  (decf s (* (aref matrix r i) (aref dt i c))))
		(setf (aref f r c) s))))
	  (setq nt (matrix-norm f))
	  (cond
	    ((< nt n) (setq n nt) (setq nt d) (setq d dt) (setq dt nt) (go la))
	    ((null into-matrix) (return d))
	    ((eq d into-matrix) (return d))
	    (t
	     (dotimes (r rows)
	       (dotimes (c cols)
		 (setf (aref into-matrix r c) (aref d r c))))
	     (return into-matrix)))))) 

(defun transpose-matrix (matrix &optional into-matrix &aux dim-1 dim-2)
  (check-arg matrix
	     (and (arrayp matrix)
		  (= (array-rank matrix) #o2))
	     "a two-dimensional array")
  (setq dim-1 (array-dimension matrix #o0)
	dim-2 (array-dimension matrix #o1))
  (if into-matrix
      (or
	(and (eq dim-1 (array-dimension into-matrix #o1))
	     (eq dim-2 (array-dimension into-matrix #o0)))
	(ferror nil "~s wrong dimensions for transpose of ~s" into-matrix matrix))
      (setq into-matrix (make-array (list dim-2 dim-1) :type (array-type matrix))))
  (if (eq matrix into-matrix)			;special case
      (dotimes (i dim-1)
	(do ((j i (1+ j)))
	    ((>= j dim-1)
	     nil)
	  (setf (aref matrix j i)
		(prog1
		  (aref matrix i j)
		  (setf (aref matrix i j) (aref matrix j i))))))
      (dotimes (i dim-1)
	(dotimes (j dim-2)
	  (setf (aref into-matrix j i) (aref matrix i j)))))
  into-matrix) 

;Determinant, based on the facts that the determinant of a triangular
;matrix is the product of the diagonal elements, and the determinant of
;the product of two matrices is the product of the determinants.
(defun determinant (matrix)
  (condition-case ()
      (multiple-value-bind (lu ps) (decompose matrix nil nil t)
	(do ((i (1- (array-length ps)) (1- i))
	     (det 1 (* det (aref lu (aref ps i) i))))
	    ((minusp i)
	     (if (and (rationalp det) (= (denominator det) 1))
		 (setq det (floor det)))
	     (if (minusp (permutation-sign ps)) (setq det (- det)))
	     (return-array ps)
	     (return-array lu)
	     det)))
    (singular-matrix 0)))

;Note that this trashes its argument
(defun permutation-sign (ps)
  (loop with sign = #o1
	for i from #o0 below (array-total-size ps)
	as j = (aref ps i)
	when (/= i j)				;Found a cycle, determine its length-1
	do (loop as k = (aref ps j)
		 do (setf (aref ps j) j)
		 (setq sign (- sign))
		 (setq j k)
		 until (= j i))
	finally (return sign))) 

;;; Linear equation solving.   

;;; The functions below are useful for solving systems of simultaneous
;;; linear equations.  They are taken from the text "Computer Solution of
;;; Linear Algebraic Systems", by Forsythe and Moler, Prentice-Hall 1967.
;;; 
;;; The function DECOMPOSE takes a square matrix A (N by N elements) and
;;; returns a square matrix holding the LU decomposition of A.  The
;;; function finds the unique solution of L * U = A, where L is a lower
;;; triangular matrix with 1's along the diagonal, and U is an upper
;;; triangular matrix.  The function returns a square matrix holding L-I+U;
;;; that is, the lower triangle not including the diagonal holds L, and the
;;; rest holds U, with the 1's along the diagonal of L not actually stored.
;;; (Note: the LU decomposition exists uniquely only if all of the
;;; principle minor matrices made from the first K rows and columns are
;;; non-singular; see Forsythe and Moler, Theorem 9.2.)
;;; 
;;; The function SOLVE takes the LU decomposition of A, and a vector of
;;; solutions of the equations B, and returns X where A * X = B.
;;; 
;;; DECOMPOSE uses partial pivoting.  Rather than actually moving the
;;; elements of the array from one row to another, it returns a permutation
;;; array PS telling how the rows of LU are permuted.  The PS array must
;;; then be passed into SOLVE so that it can interpret LU properly.
;;;
;;; Iterative improvement is not yet implemented.

;;; Utility functions and macros.

(defun 1d-arrayp (array)
   (and (arrayp array)
	(= (array-rank array) 1)))

(defun 2d-arrayp (array)
   (and (arrayp array)
	(= (array-rank array) 2)))

(defun 1-or-2d-arrayp (array)
  (and (arrayp array)
       (or (= (array-rank array) 1)
	   (= (array-rank array) 2))))

(defmacro exchange (f1 f2)
   (let ((v1 (gensym))
	 (v2 (gensym)))
     `(let ((,v1 ,f1)
	    (,v2 ,f2))
	(setf ,f1 ,v2)
	(setf ,f2 ,v1))))


;;; DECOMPOSE
;;; A is an N by N array.
;;; Two values are returned: LU and PS.
;;; The caller may provide arrays to be used for LU and PS by passing
;;; the optional arguments; otherwise, new arrays will be allocated.
;;; If the same array is passed as A and LU, A is overwriten with
;;; the decomposition correctly.
;;; The condition SINGULAR is raised if the matrix is singular.

;;; The additional optional argument, ALLOW-RATIONALS, is a temporary
;;; kludge until rational numbers are really installed.  It causes integer
;;; matrix elements to be treated as rationals rather than floats.
;;; The bug with this is that rationals don't automatically change back
;;; to integers when the denominator is 1.
(defun decompose (a &optional lu ps allow-rationals &aux n)
  (declare (values lu ps))
  ;; Prepare arguments.
  (check-arg a 2d-arrayp "a two-dimensional array")
  (setq n (array-dimension a #o0))
  (check-arg a (= n (array-dimension a #o1)) "a square array")
  (if lu
      (check-arg lu 2d-arrayp "a two-dimensional array")
      (setq lu (make-array (list n n) :type (array-type a))))
  (if ps
      (check-arg ps 1d-arrayp "a one-dimensional array")
      (setq ps (make-array n)))
  (let ((scales (make-array n)))
    ;; Init PS to the identity, LU to A, and SCALES to the reciprocal
    ;; of the largest-magnitude element on a given row.
    (dotimes (i n)
      (setf (aref ps i) i)
      (let ((normrow #o0))
	(dotimes (j n)
	  (let ((aij (aref a i j)))
	    (setf (aref lu i j) aij)
	    (setq normrow (max (abs aij) normrow))))
	(if (zerop normrow)
	    (error 'singular-matrix "The matrix ~s is singular: it has a zero row." a))
	(setf (aref scales i) (/ 1.0s0 normrow))))
    ;; Gaussian elimination with partial pivoting.
    (dotimes (k (- n #o1))
      ;; Find the pivot index.
      (let ((pivotindex nil)
	    (biggest #o0))
	(do ((i k (1+ i)))
	    ((>= i n)
	     nil)
	  (let ((size (* (abs (aref lu (aref ps i) k)) (aref scales (aref ps i)))))
	    (cond
	      ((> size biggest) (setq biggest size) (setq pivotindex i)))))
	(if (zerop biggest)
	    (error 'singular-matrix "The matrix ~s is singular: SOLVE will divide by zero." a))
	(exchange (aref ps pivotindex) (aref ps k)))
      ;; Do the elimination with that pivoting.
      (let* ((psk (aref ps k)) (pivot (aref lu psk k)))
	;; For fixnum arguments, get an approximate answer, rather than the wrong
	;; answer, unless ALLOW-RATIONALS is specified, in which case get the
	;; exact answer slowly.  this is most useful for determinant where we know
	;; that the rationals will eventually cancel.
	(and (integerp pivot) (/= pivot #o1) (/= pivot #o-1)
	     (setq pivot (if allow-rationals
			     (rational pivot)
			     (float pivot))))
	(do ((i (1+ k) (1+ i)))
	    ((>= i n)
	     nil)
	  (let ((psi (aref ps i)))
	    (let ((mult (/ (aref lu psi k) pivot)))
	      (setf (aref lu psi k) mult)
	      (if (not (zerop mult))
		  (do ((j (1+ k) (1+ j)))
		      ((>= j n)
		       nil)
		    (setf (aref lu psi j) (- (aref lu psi j) (* mult (aref lu psk j)))))))))))
    (if (zerop (aref lu (aref ps (1- n)) (1- n)))
	(error 'singular-matrix "The matrix ~s is singular: SOLVE will divide by zero." a))
    (return-array (prog1
		    scales
		    (setf scales nil))))
  (values lu ps)) 

;;; SOLVE
;;; LU is the N by N LU-decomposition of A.
;;; PS is the N-long permutation vector for LU.  B is an N-long array
;;; of solutions to the equations.
;;; The returned value is X: the solution of A * X = B.
;;; The caller may provide the array to be used as X by passing the optional
;;; argument.

(defun solve (lu ps b &optional x &aux n)
  ;; Prepare arguments.
  (check-arg lu 2d-arrayp "a two-dimensional array")
  (setq n (array-dimension lu #o0))
  (or (= n (array-dimension lu #o1))
      (ferror () "The first argument must be a square array."))
  (check-arg ps 1d-arrayp "a one-dimensional array")
  (check-arg b 1d-arrayp "a one-dimensional array")
  (if x
      (check-arg x 1d-arrayp "a one-dimensional array")
      (setq x (make-array n :type (array-type b))))
  (dotimes (i n)
    (let ((psi (aref ps i))
	  (dot #o0))
      (dotimes (j i)
	(setq dot (+ dot (* (aref lu psi j) (aref x j)))))
      (setf (aref x i) (- (aref b psi) dot))))
  (do ((i (1- n) (1- i)))
      ((< i #o0))
    (let ((psi (aref ps i))
	  (dot #o0))
      (do ((j (1+ i) (1+ j)))
	  ((>= j n))
	(setq dot (+ dot (* (aref lu psi j) (aref x j)))))
      (setf (aref x i)
	    (/ (- (aref x i) dot)
	       (let ((d (aref lu psi i))) 
		 ;; Compensate for fixnum division--take this out when
		 ;; there are "rational" numbers.
		 (and (integerp d) (/= d #o1) (/= d #o-1) (setq d (float d)))
		 d)))))
  x)