LU decomposition

LU.tq

let s = Array.Unsafe.set
let g = Array.get
let g2 = Matrix.get
let s2 = Matrix.Unsafe.set

let swap a (i0, j0) (i1, j1) =
  let t = g2 a (i0, j0) in
  let () = g2 a (i1, j1) @ s2 a (i0, j0) in
  s2 a (i1, j1) t

let swapRow a i0 i1 =
  let m, n = Matrix.dimensions a in
  for 0 1 (n-1) () [(), j -> swap a (i0, j) (i1, j)]

(** Partial pivoting for LU decomposition to improve numerical stability for approximate numeric types. *)
let pivot p m n a j =
  let jp, t =
    for (j+1) 1 (m-1) (j, abs(g2 a (j, j))) [(jp, t), i ->
      let ab = abs(g2 a (i, j)) in
      if ab > t then i, ab else jp, t] in
  let () = Array.Unsafe.set p j jp in
  let () = if g2 a (jp, j) = 0 then panic "SingularMatrix" else () in
  if jp = j then () else swapRow a j jp

(** Representation-agnostic LU decomposition. *)
let lu a =
  let m, n = Matrix.dimensions a in
  let minmn = min m n in
  let p = Array.init minmn [_ -> 0] in
  let a = Matrix.copy a in
  let () =
    for 0 1 (minmn-1) () [(), j ->
      let () = pivot p m n a j in
      let () =
        if j >= m-1 then () else
          let recp = 1 / g2 a (j, j) in
          for (j+1) 1 (m-1) () [(), k ->
            g2 a (k, j) * recp @ s2 a (k, j)] in
      if j >= minmn - 1 then () else
        for (j+1) 1 (m-1) () [(), ii ->
          for (j+1) 1 (n-1) () [(), jj ->
            g2 a (ii, jj) - g2 a (ii, j) * g2 a (j, jj) @ s2 a (ii, jj)]]] in
  let l = Matrix.init m n [i, j ->
    compare(i, j)
    @ [ Less -> 0
      | Equal -> 1
      | Greater -> Matrix.get a (i, j) ]] in
  let u = Matrix.init m n [i, j ->
    compare(i, j)
    @ [ Less
      | Equal -> Matrix.get a (i, j)
      | Greater -> 0 ]] in
  l, u, p

let determinant a =
  let _, u, p = lu a in
  let s = p @ Array.mapi [i, n -> if i=n then 1 else -1] @ ∏ in
  let det = Matrix.diagonal u @ Vector.product in
  let () = yield(s, det) in
  s*det