(* Matrix Multiplication and Markov Chains

   Author:  Walid Taha
   Date:    Thu Jun 26 08:26:28 CDT 2003
            Mon Mar  8 20:34:43 CST 2004 Bug fix.
            Sun Mar 14 17:43:12 CST 2004 Bug fix.

   Problem:  Given a particular Markov model (defined in detail below)
             during a transition a state can remain unchanged with
             probability "p" or move to some close-by state with
             another given probability, compute p so that there is
             a 50% chance after 2 transitions that if we start from s0
             we will still be in s0.

   Idea:     Even if we want to use numerical techniques (for their
             generality), we don't have to recompute the full matrix
             everytime.

   Credit:   Unstaged matrix operations based on those of Markus Mottl.
             With the optimizing OCaml compiler, the performance of
             OCaml is comparable to that of C.
             (See http://www.bagley.org/~doug/shootout/).
             Thanks to Robert Rosenbaum for finding sum bug.
             Thanks to Robert Rosenbaum for finding multiplication bug.

   Notes:    mkmatrix assumes n greater than 3.

*)

Trx.init_times ();;

(**************************************************************************)
(* Building a model parameterized by a probability (Unstaged)             *)
(**************************************************************************)

let mkmatrix rows cols p =
  let last_col = cols - 1
  and m = Array.make_matrix rows cols 0.0 in
  for i = 0 to rows - 1 do
    let mi = m.(i) in
    for j = 0 to last_col do 
     (* stay probability is p *)
     (* choices are max(0,i-5)...min(n-1,i+5), divide remain prob *)
     let r  = rows / 4 in
     let mn = max 0 (i-r) in
     let mx = min (rows-1) (i+r) in
     let rn = mx-mn in
     let q =(if i=j then
              p
             else if (j>=mn)&&(j<=mx) then
              (1.0 -. p)/.(float rn)
             else
              0.0) in
     mi.(j) <- q; 
    done;
  done;
  m

let rec inner_loop k v m1i m2 j =
  if k < 0 then v
  else inner_loop (k - 1) (v +. m1i.(k) *. (m2.(k).(j))) m1i m2 j

let mmult rows cols m1 m2 m3 =
  let last_col = cols - 1 and last_row = rows - 1 in
  for i = 0 to last_row do
    let m1i = m1.(i) and m3i = m3.(i) in
    for j = 0 to last_col do 
     m3i.(j) <- inner_loop last_row 0.0 m1i m2 j done;
  done

let ans1 =
 fun p ->
  let size = 30 in
  let m1 = mkmatrix size size p and m2 = mkmatrix size size p in
  let m3 = Array.make_matrix size size 0.0 in
  mmult size size m1 m2 m3;
  m3.(0).(0);;

(**************************************************************************)
(* Staged version                                                         *)
(**************************************************************************)

(* Staged version uses an option type to elimination additions, 
   subtractions, and multiplications that involve a known 0.0 value *)
   
let ( -..)  x y =
  match (x,y) with
   (None,   None)   -> None
  |(None,   Some b) -> Some .< -. .~b>.
  |(Some a, None)   -> Some a
  |(Some a, Some b) -> Some .< .~a -. .~ b>.

let ( +..)  x y =
  match (x,y) with
   (None,   None)   -> None
  |(None,   Some b) -> Some b
  |(Some a, None)   -> Some a
  |(Some a, Some b) -> Some .< .~a +. .~ b>.

let ( *..)  x y =
  match (x,y) with
   (None,   None)   -> None
  |(None,   Some b) -> None
  |(Some a, None)   -> None
  |(Some a, Some b) -> Some .< .~a *. .~ b>.

(* Staged version proper *)

let lift x = .<x>.;;

let mkmatrix rows cols p =
  let last_col = cols - 1
  and m = Array.make_matrix rows cols None in
  for i = 0 to rows - 1 do
    let mi = m.(i) in
    for j = 0 to last_col do 
     (* stay probability is p *)
     (* choices are max(0,i-5)...min(n-1,i+5), divide remain prob *)
     let r  = rows / 4 in
     let mn = max 0 (i-r) in
     let mx = min (rows-1) (i+r) in
     let rn = mx-mn in
     let q =(if i=j then
              Some p
             else if (j>=mn)&&(j<=mx) then
              Some .<(1.0 -. .~p)/.(.~(lift (float rn)))>.
             else
              None) in
     mi.(j) <- q; 
    done;
  done;
  m

let rec inner_loop k v m1i m2 j =
  if k < 0 then v
  else inner_loop (k - 1) (v +.. (m1i.(k) *.. (m2.(k).(j)))) m1i m2 j

let mmult rows cols m1 m2 m3 =
  let last_col = cols - 1 and last_row = rows - 1 in
  for i = 0 to last_row do
    let m1i = m1.(i) and m3i = m3.(i) in
    for j = 0 to last_col do 
     m3i.(j) <- inner_loop last_row None m1i m2 j done;
  done;;

let ans =
Trx.timenew "ans" (fun () ->
.<fun p ->
.~(
   let size = 30 in
   let m1 = mkmatrix size size .<p>. and m2 = mkmatrix size size .<p>. in
   let m3 = Array.make_matrix size size None in
   mmult size size m1 m2 m3;
   match m3.(0).(0) with Some q -> q
  )
>.);;

let ans2 = 
Trx.timenew "ans2" (fun () ->
.! ans);;

(**************************************************************************)
(* Simple Solver                                                          *)
(**************************************************************************)

(* A function that numerically solves for x such that f(x)=y *)

let ab x = if x>0.0 then x else -. x;;

let rec solve f y x0 y0 x1 y1 e n =
  let xb = (x0 +. x1)/.2.0 in
  let yg = f xb in
  let en = ab(y-.yg) in
  if (en<=e) then
   (n+1,xb)
  else
   if (yg<y)&&(y0<y1) then
    solve f y xb yg x1 y1 e (n+1)
   else
    solve f y x0 y0 xb yg e (n+1);;

let r1 = 
Trx.timenew "r1" 
 (fun () -> solve ans1 0.5 0.0 (ans1 0.0) 1.0 (ans1 1.0) (1E-15) 0);;

let r2 = 
Trx.timenew "r2"
 (fun () -> solve ans2 0.5 0.0 (ans2 0.0) 1.0 (ans2 1.0) (1E-15) 0);;

Trx.print_times ();;

(* Timings on my machine:

 __ r1 __________________________ 1x avg= 1.409808E+03 msec
 __ ans _________________________ 8x avg= 1.842252E+02 msec
 __ ans2 ______________________ 256x avg= 5.682198E+00 msec
 __ r2 _______________________ 4096x avg= 3.892584E-01 msec 

r1   is solving using the unstaged version
ans  is the time needed to generate for the staged version
ans2 is the time needed to compile the staged version
r2   is solving using the staged version

Notice that the sum of the last three times is still almost an order
of magnitude less than the time needed to computer r1.

If we just compare r2 to r1, we have a speedup of over 3,500 times.

*)