(*
%use "Matrix/lu_solver.fsh";;
*)

(* 

This program takes a non-singular square matrix a and a vector b of
the same length and solves the equation 

	a x = b

using an LU decomposition of the matrix. The actual code of the solver
(repeated below) is given by

let lu_solver_pr b x a = 
let n = rows_var a in
new p = diagvector n 
in 
  lu_decompose a p ;
  x := permute p b ; 
  solve_lower_pr a x ;
  solve_upper_pr a x 
end
;;


*)


let rows x =  lendim #x ;;
let cols x = lendim (preddim #x) ;;

(* 
best_row : var [float] -> int -> int 

best_row a k finds the first row number >= k of the matrix whose
kth entry has the largest absolute value. Could be generalised. 
*)

let best_row a k = 
new r = k in 
  for k+1 <= j< rows a do
    if fabs a[j,k] > (fabs a[r,k])
    then r:= j
    else skip
  done
return r
;;


(* 
swap : var [a] -> int -> int -> comm
exchanges to subarrays of an array
*)

let swap (a: var b) i j = 
new t = sub(a,i) in 
  sub(a,i) := sub(a,j) ;
  sub(a,j) := t 
end 
;;


(*
install_pivot : var [float] -> int -> comm

install_pivot a k finds best_row a k and swaps that row with the kth.
This corresponds to the "exchange" of Molinari. In addition, the
permutation is recorded in the vector p which should be initialised to
be the diagonal [0,1,2,3,4...] *)


let install_pivot (a:var b) p k = 
  new r = best_row a k in
    if not (r=k) 
    then swap a k r ;
	 swap p k r
    else skip 
  end
;;


(*
decompose1 : var [float] -> int -> comm

performs one pivot
*)


let decompose1 (a:var b) k = 
let n = rows a in 
new 
  #et = float_shape 
in 
  for k+1 <= i < n do
    et := a[i,k] / a[k,k] ;
    a[i,k] := et ;
    for k+1 <= j < n do 
      a[i,j]:= a[i,j] - et * a[k,j]
    done
  done
end
;;


let lu_decompose (a:var b) p = 
let n = rows a in
  for k< n - 1 do 
    install_pivot a p k ;
    decompose1 a k 
  done
 ;;

let permute_pr (y:var d) (p:var c) (b:var d) = 
for k<lendim #p do
   y[k] := b[p[k]]
done 
;;

let permute_sh p_sh b_sh = b_sh ;;

let permute = proc2fun2 permute_pr permute_sh
;;

let solve_lower_pr (a:var [float]) (y:var [float]) =
let n = rows a in
for k<n do
  for j<k do 
    y[k] := y[k] - a[k,j] * y[j] 
  done
done
;;
 
let solve_upper_pr (a:var [float]) (y:var [float])  = 
let n = rows a in 
new #kr = int_shape in 
  for k<n do
    kr := n-k-1 ;
    for kr +1 <= j<n do
      y[kr] := y[kr] - a[kr,j] * y[j] 
    done ;
    y[kr] := y[kr] / a[kr,kr] 
  done
end
;;

let lu_solver_pr b x (a: var c) = 
let n = rows a in
new p = diagvector n
in 
  lu_decompose a p ;
  x := permute p b ; 
  solve_lower_pr a x ;
  solve_upper_pr a x 
end
;;

let lu_solver_sh bsh ash = bsh 
;;

let lu_solver a b = proc2fun (lu_solver_pr b) (lu_solver_sh #b) a
;;