Source file fft.ml

(* gsl-ocaml - OCaml interface to GSL                       *)
(* Copyright (©) 2002-2012 - Olivier Andrieu                *)
(* Distributed under the terms of the GPL version 3         *)

open Gsl_complex

let () = Error.init ()

exception Wrong_layout

let _ = Callback.register_exception "mlgsl_layout_exn" Wrong_layout

type layout = Real | Halfcomplex | Halfcomplex_rad2 | Complex
type fft_array = { mutable layout : layout; data : float array }

let check_layout layout a = if a.layout <> layout then raise Wrong_layout

module Real = struct
  type workspace
  type wavetable

  external alloc_workspace : int -> workspace
    = "ml_gsl_fft_real_workspace_alloc"

  external alloc_wavetable : int -> wavetable
    = "ml_gsl_fft_real_wavetable_alloc"

  external free_workspace : workspace -> unit = "ml_gsl_fft_real_workspace_free"
  external free_wavetable : wavetable -> unit = "ml_gsl_fft_real_wavetable_free"

  let make_workspace size =
    let ws = alloc_workspace size in
    Gc.finalise free_workspace ws;
    ws

  let make_wavetable size =
    let wt = alloc_wavetable size in
    Gc.finalise free_wavetable wt;
    wt

  external transform :
    ?stride:int -> fft_array -> wavetable -> workspace -> unit
    = "ml_gsl_fft_real_transform"

  external transform_rad2 : ?stride:int -> fft_array -> unit
    = "ml_gsl_fft_real_radix2_transform"

  external ex_unpack : ?stride:int -> float array -> float array -> unit
    = "ml_gsl_fft_real_unpack"

  let unpack ?(stride = 1) r_arr =
    if r_arr.layout <> Real then raise Wrong_layout;
    let c_arr = Array.make (2 * Array.length r_arr.data / stride) 0. in
    ex_unpack ~stride r_arr.data c_arr;
    { layout = Complex; data = c_arr }
end

module Halfcomplex = struct
  type wavetable

  external alloc_wavetable : int -> wavetable
    = "ml_gsl_fft_halfcomplex_wavetable_alloc"

  external free_wavetable : wavetable -> unit
    = "ml_gsl_fft_halfcomplex_wavetable_free"

  let make_wavetable size =
    let wt = alloc_wavetable size in
    Gc.finalise free_wavetable wt;
    wt

  external transform :
    ?stride:int -> fft_array -> wavetable -> Real.workspace -> unit
    = "ml_gsl_fft_halfcomplex_transform"

  external transform_rad2 : ?stride:int -> fft_array -> unit
    = "ml_gsl_fft_halfcomplex_radix2_transform"

  external backward :
    ?stride:int -> fft_array -> wavetable -> Real.workspace -> unit
    = "ml_gsl_fft_halfcomplex_backward"

  external backward_rad2 : ?stride:int -> fft_array -> unit
    = "ml_gsl_fft_halfcomplex_radix2_backward"

  external inverse :
    ?stride:int -> fft_array -> wavetable -> Real.workspace -> unit
    = "ml_gsl_fft_halfcomplex_inverse"

  external inverse_rad2 : ?stride:int -> fft_array -> unit
    = "ml_gsl_fft_halfcomplex_radix2_inverse"

  external ex_unpack : ?stride:int -> float array -> float array -> unit
    = "ml_gsl_fft_halfcomplex_unpack"

  external ex_unpack_rad2 : ?stride:int -> float array -> float array -> unit
    = "ml_gsl_fft_halfcomplex_unpack_rad2"

  let unpack ?(stride = 1) hc_arr =
    match hc_arr.layout with
    | Halfcomplex ->
        let c_arr = Array.make (2 * Array.length hc_arr.data / stride) 0. in
        ex_unpack ~stride hc_arr.data c_arr;
        { layout = Complex; data = c_arr }
    | Halfcomplex_rad2 ->
        let c_arr = Array.make (2 * Array.length hc_arr.data / stride) 0. in
        ex_unpack_rad2 ~stride hc_arr.data c_arr;
        { layout = Complex; data = c_arr }
    | _ -> raise Wrong_layout
end

module Complex = struct
  type workspace
  type wavetable
  type direction = Forward | Backward

  external alloc_workspace : int -> workspace
    = "ml_gsl_fft_complex_workspace_alloc"

  external alloc_wavetable : int -> wavetable
    = "ml_gsl_fft_complex_wavetable_alloc"

  external free_workspace : workspace -> unit
    = "ml_gsl_fft_complex_workspace_free"

  external free_wavetable : wavetable -> unit
    = "ml_gsl_fft_complex_wavetable_free"

  let make_workspace size =
    let ws = alloc_workspace size in
    Gc.finalise free_workspace ws;
    ws

  let make_wavetable size =
    let wt = alloc_wavetable size in
    Gc.finalise free_wavetable wt;
    wt

  external forward :
    ?stride:int -> complex_array -> wavetable -> workspace -> unit
    = "ml_gsl_fft_complex_forward"

  external forward_rad2 : ?dif:bool -> ?stride:int -> complex_array -> unit
    = "ml_gsl_fft_complex_rad2_forward"

  external transform :
    ?stride:int -> complex_array -> wavetable -> workspace -> direction -> unit
    = "ml_gsl_fft_complex_transform"

  external transform_rad2 :
    ?dif:bool -> ?stride:int -> complex_array -> direction -> unit
    = "ml_gsl_fft_complex_rad2_transform"

  external backward :
    ?stride:int -> complex_array -> wavetable -> workspace -> unit
    = "ml_gsl_fft_complex_backward"

  external backward_rad2 : ?dif:bool -> ?stride:int -> complex_array -> unit
    = "ml_gsl_fft_complex_rad2_backward"

  external inverse :
    ?stride:int -> complex_array -> wavetable -> workspace -> unit
    = "ml_gsl_fft_complex_inverse"

  external inverse_rad2 : ?dif:bool -> ?stride:int -> complex_array -> unit
    = "ml_gsl_fft_complex_rad2_inverse"
end

let unpack = function
  | { layout = Real } as f -> (Real.unpack f).data
  | ({ layout = Halfcomplex } | { layout = Halfcomplex_rad2 }) as f ->
      (Halfcomplex.unpack f).data
  | { layout = Complex; data = d } -> d

let hc_mult ({ data = a } as fa) ({ data = b } as fb) =
  check_layout Halfcomplex fa;
  check_layout Halfcomplex fb;
  let len = Array.length a in
  if Array.length b <> len then invalid_arg "hc_mult: array sizes differ";
  a.(0) <- a.(0) *. b.(0);
  for i = 1 to pred len / 2 do
    let a_re = a.((2 * i) - 1) in
    let a_im = a.(2 * i) in
    let b_re = b.((2 * i) - 1) in
    let b_im = b.(2 * i) in
    a.((2 * i) - 1) <- (a_re *. b_re) -. (a_im *. b_im);
    a.(2 * i) <- (a_re *. b_im) +. (a_im *. b_re)
  done;
  if len mod 2 = 0 then a.(pred len) <- a.(pred len) *. b.(pred len)

let hc_mult_rad2 ({ data = a } as fa) ({ data = b } as fb) =
  check_layout Halfcomplex_rad2 fa;
  check_layout Halfcomplex_rad2 fb;
  let len = Array.length a in
  if Array.length b <> len then invalid_arg "hc_mult_rad2: array sizes differ";
  a.(0) <- a.(0) *. b.(0);
  for i = 1 to pred len / 2 do
    let a_re = a.(i) in
    let a_im = a.(len - i) in
    let b_re = b.(i) in
    let b_im = b.(len - i) in
    a.(i) <- (a_re *. b_re) -. (a_im *. b_im);
    a.(len - i) <- (a_re *. b_im) +. (a_im *. b_re)
  done;
  a.(len / 2) <- a.(len / 2) *. b.(len / 2)