fft.ml

fft,ocaml语言实现,objective functional language

(* A Staged Implementation of FFT

   Author:  Walid Taha, Oleg Kiselyov, Kedar N. Swadi
   Date:    Thu Jul  8 20:24:46 CDT 2004

   Problem:    See "A Methodology for Generating Verified 
                    Combinatorial Circuits"
*)

(* Step 1 *)
let pi = 4.0 *. atan(1.0)

(*elided*)
let w dir n j = (* exp(-2PI dir/N)^j dir=1.0 for the forward transform *)
  let theta = dir *. ((float_of_int (-2 * j)) *. pi) /. (float_of_int n) in
  ((cos theta), (sin theta)) 

let add ((r1,i1), (r2, i2)) = ((r1 +. r2), (i1 +. i2))
(*elided*)
let sub ((r1,i1), (r2, i2)) = ((r1 -. r2), (i1 -. i2))

(*elided*)
let mult ((r1, i1), (r2, i2)) =
  let rp = (r1 *. r2) -. (i1 *. i2) in
  let ip = (r1 *. i2) +. (r2 *. i1) in
  (rp, ip)

let rec  split l =
  match l with 
    [] -> ([], [])
  | x::y::xs -> let (a,b) = split xs in (x::a, y::b)

let rec merge dir l1 l2 =
  let n = 2 * List.length l1 in
  let rec mg l1 l2 j =
    match (l1, l2) with
      (x::xs, y::ys) ->
        let z1 = mult (w dir n j, y) in
        let zx = add (x, z1) in
        let zy = sub (x, z1) in
        let (a,b) = (mg xs ys (j+1)) in (zx::a, zy::b)
    | _ -> ([], []) in
  let (a,b) = mg l1 l2 0 in (a @ b)

let rec fft dir l =
  if (List.length l = 1) then 
    l
  else
    let (e,o) = split l in
    let y0 = fft dir e  in
    let y1 = fft dir o  in
    merge dir y0 y1

(* Step 2a: The correctness should be obvious  *)

(* List of size n with (1,0) at position l *)
let rec impulse n l = 
   if n = 0 then []
   else let t = if l = 0 then (1.0,0.0) else (0.0,0.0)
   in t::(impulse (n-1) (l-1))

let () = assert ([(0., 0.); (4., 0.); (0., 0.); (0., 0.)] = 
                fft (-1.0) (fft 1.0 (impulse 4 1)))

let () = assert ([(0., 0.); (8., 0.); (0., 0.); (0., 0.); (0., 0.); 
                 (0., 0.); (0., 0.); (0., 0.)] = 
                fft (-1.0) (fft 1.0 (impulse 8 1)))

let test3_u = fft (-1.0) (fft 1.0 (impulse 16 1))


(* -------------------------------------------------------------------- *)
(* Step 2b: Converting the program into monadic form                    *)

let (ret,bind,y_sm, run) = 
  let ret a = fun s k -> k s a in 
  let bind a f = 
    fun s k -> a s (fun s' b -> f b s' k) in
  let rec y_sm f = f (fun x s k -> y_sm f x s k) in
  let run f size nums =
(* convert list to array code *)
    let list2array l arr = 
      let rec lfta l m =
        match l with
          [] -> (arr)
        | ((x,y)::z) ->
            let sm = m+1 in
            let ssm = m+2 in
            (((arr).(m) <- x;
              (arr).(sm) <- y;
              ((lfta z (ssm)))))
      in lfta l 0 in
(* convert input array to list format *)
    let array2list arr n k =
      let rec gl m l =
        if m = n then k l arr
        else
          let sm = m+1 in
          (let x1 = (arr).(m) in
          let y1 = (arr).(sm) in
          (gl (m+2) (l @ [(x1, y1)])))
      in gl 0 [] in
    let run1 f arr = f [] (fun s x -> list2array x arr) in
     (array2list nums size (fun l arr2 -> run1 (f l) arr2 )) in
  (ret,bind,y_sm,run)

let merge_m dir l1 l2 =
  let n = 2 * List.length l1 in
  let rec mg l1 l2 j =
    match (l1, l2) with
      (x::xs, y::ys) ->
        bind (ret (mult (w dir n j, y))) (fun z1 ->
        bind (ret (add (x, z1))) (fun zx ->
        bind (ret (sub (x, z1))) (fun zy ->
        bind (mg xs ys (j+1)) (fun (a,b) ->
        ret (zx::a, zy::b)))))
    | _ -> ret ([], []) in
    bind (mg l1 l2 0) (fun (a,b) ->
    ret (a @ b))

let rec fft_m dir f l =
  if (List.length l = 1) then ret l
  else
    let (e,o)  = split l in
    bind (f e) (fun y0 ->
    bind (f o) (fun y1 ->
    merge_m dir y0 y1))

(*tests *)

let rec impulse n l = 
   let arr = Array.make (2 * n) 0.0 in
   let () = arr.(2 * l) <- 1.0 in
   arr

let () = assert ([|0.; 0.; 4.; 0.; 0.; 0.; 0.; 0.;|] = 
                 let s1 = run (y_sm (fft_m (1.0))) 8 (impulse 4 1)  in
                 let s2 = run (y_sm (fft_m (-1.0))) 8 s1 in
                 s2)
      
let () = assert ([|0.; 0.; 8.; 0.; 0.; 0.; 0.; 0.; 0.; 0.; 
                   0.; 0.; 0.; 0.; 0.; 0.;|] = 
                 let s1 = run (y_sm (fft_m (1.0))) 16 (impulse 8 1)  in
                 let s2 = run (y_sm (fft_m (-1.0))) 16 s1 in
                 s2)

let test3_u = run (y_sm (fft_m (-1.0))) 32
                   (run (y_sm (fft_m 1.0)) 32 (impulse 16 1))

(* -------------------------------------------------------------------- *)
(* Step 3,4: Staging the monadic program                                *)

let (retS,retN,bind,y_sm,run) = 
  let retS a = fun s k -> k s a in 
  let retN a = fun s k -> .<let z = .~a in .~(k s .<z>.)>. in 
  let bind a f = 
    fun s k -> a s (fun s' b -> f b s' k) in
  let rec y_sm f =  f (fun x s k -> y_sm f x s k) in
  let run f size nums =
    (* convert list to array code *)
    let list2array_s l arr = 
      let rec lfta l m =
        match l with
          [] -> (arr)
        | ((x,y)::z) ->
            let sm = m+1 in
        let ssm = m+2 in
        (.<((.~arr).(m) <- .~x;
            (.~arr).(sm) <- .~y;
            (.~(lfta z (ssm))))>.)
      in lfta l 0 in
    (* convert input array to list format *)
    let array2list_s arr n k =
      let rec gl m l =
        if m = n then k l arr
        else
          let sm = m+1 in
          (.<let x1 = (.~arr).(m) in
          let y1 = (.~arr).(sm) in
          .~(gl (m+2) (l @ [(.<x1>., .<y1>.)]))>.)
      in gl 0 [] in
    let run1 f a = f [] (fun s x -> list2array_s x a) in
     (array2list_s nums size (fun l arr2 -> run1 (f l) arr2 )) in
    (retS,retN,bind,y_sm,run)

(* Standard auxiliary monadic operators *)

let liftM f = fun x -> retS (f x)

let compM a b =
  fun x -> bind (b x) (fun nx -> a nx)

let liftcM op (x,y) = 
  bind (op x) (fun nx ->
  bind (op y) (fun ny ->
  retS (nx,ny)))

let w_s dir n j = (* exp(-2PI dir/N)^j dir=1.0 for the forward transform *)
  let theta = dir *. ((float_of_int (-2 * j)) *. pi) /. (float_of_int n) in
  let ct = cos theta in
  let st = sin theta in
  (.<ct>., .<st>.)

let add_s ((r1,i1), (r2, i2)) = ((.<.~r1 +. .~r2>.), (.<.~i1 +. .~i2>.))
(*elided*)
let sub_s ((r1,i1), (r2, i2)) = ((.<.~r1 -. .~r2>.), (.<.~i1 -. .~i2>.))

(*elided*)
let mult_s ((r1, i1), (r2, i2)) =
  let rp = .<(.~r1 *. .~r2) -. (.~i1 *. .~i2)>. in
  let ip = .<(.~r1 *. .~i2) +. (.~r2 *. .~i1)>. in
  (rp, ip)

let merge_ms dir l1 l2 =
  let n = 2 * List.length l1 in
  let rec mg l1 l2 j =
    match (l1, l2) with
      (x::xs, y::ys) ->
        bind (retS (mult_s (w_s dir n j, y))) (fun z1 ->
        bind (retS (add_s (x, z1))) (fun zx ->
        bind (retS (sub_s (x, z1))) (fun zy ->
        bind (mg xs ys (j+1)) (fun (a,b) ->
        retS (zx::a, zy::b)))))
    | _ -> retS ([], []) in
    bind (mg l1 l2 0) (fun (a,b) ->
    retS (a @ b))

let fft_ms dir f l =
  if (List.length l = 1) then retS l
  else
    let (e,o)  = split l in
    bind (f e) (fun y0 ->
    bind (f o) (fun y1 ->
    merge_ms dir y0 y1))

(* Testing the generated code *)
let s1m = .<fun x -> .~(run (y_sm (fft_ms 1.0))  (8 * 2) .<x>.);0>.
(*
let _ = Trx.C.addToLibrary (Trx.C.toC s1m) "." "test";;
*)

let test_ms dir arr =
  let n = Array.length arr in
  let tranc = .<fun x -> .~(run (y_sm (fft_ms dir))  n .<x>.)>. in
  (.!tranc) (arr)

let () = assert ([|0.; 0.; 4.; 0.; 0.; 0.; 0.; 0.;|] = 
                test_ms (-1.0) (test_ms 1.0 (impulse 4 1)))
let () = assert ([|0.; 0.;  8.00000000000512; 0.;  0.; 0.;  0.; 0.;
		   0.; 0.;
		   -5.11946041115152184e-12; 0.; 0.; 0.; 0.; 0.|] = 
                 test_ms (-1.0) (test_ms 1.0 (impulse 8 1)))
let test3_ms = test_ms (-1.0) (test_ms 1.0 (impulse 16 1))
 
let merge_ms dir l1 l2 =
  let n = 2 * List.length l1 in
  let rec mg l1 l2 j =
    match (l1, l2) with
      (x::xs, y::ys) ->
        bind ((liftcM retN) (mult_s (w_s dir n j, y))) (fun z1 ->
        bind (retS (add_s (x, z1))) (fun zx ->
        bind (retS (sub_s (x, z1))) (fun zy ->
        bind (mg xs ys (j+1)) (fun (a,b) ->
        retS (zx::a, zy::b)))))
    | _ -> retS ([], []) in
    bind (mg l1 l2 0) (fun (a,b) ->
    retS (a @ b))

let fft_ms dir f l =
  if (List.length l = 1) then retS l
  else
    let (e,o)  = split l in
    bind (f e) (fun y0 ->
    bind (f o) (fun y1 ->
    merge_ms dir y0 y1))

let s1m = .<fun x -> .~(run (y_sm (fft_ms 1.0))  16 .<x>.);0>.
(*
let _ = Trx.C.addToLibrary (Trx.C.toC s1m) "." "test2";;
*)
(* Step 4b: Changing some of the  implicit retS -> retN *) 

(* We still have some code duplication. So, we notice that 
   the code above had a few _implicit_ retS:
   (bind (retS x)) (fun x ->
   (bind (retS y)) (fun y ->
  Note that bind . retS is the identity...
 So, we replace those retS with RetN:
*)

let merge_ms1 dir l1 l2 =
  let n = 2 * List.length l1 in
  let rec mg l1 l2 j =
    match (l1, l2) with
      (x::xs, y::ys) ->
   bind ((liftcM retN) x) (fun x ->
   bind ((liftcM retN) y) (fun y ->