QCaml/common/util.org

#+begin_src elisp tangle: no :results none :exports none
(setq pwd (file-name-directory buffer-file-name))
(setq name (file-name-nondirectory (substring buffer-file-name 0 -4)))
(setq lib (concat pwd "lib/"))
(setq testdir (concat pwd "test/"))
(setq mli (concat lib name ".mli"))
(setq ml  (concat lib name ".ml"))
(setq c  (concat lib name ".c"))
(setq test-ml  (concat testdir name ".ml"))
(org-babel-tangle)
#+end_src 

* Util
  :PROPERTIES:
  :header-args: :noweb yes :comments both
  :END:

  Utility functions.

  
** Test header                                                     :noexport:

   #+begin_src ocaml :tangle (eval test-ml) :exports none
open Common.Util
open Alcotest
   #+end_src

** External C functions

   | ~erf_float~   | Error function ~erf~ from =libm=                |
   | ~erfc_float~  | Complementary error function ~erfc~ from =libm= |
   | ~gamma_float~ | Gamma function ~gamma~ from =libm=              |
   | ~popcnt~      | ~popcnt~ instruction                            |
   | ~trailz~      | ~ctz~ instruction                               |
   | ~leadz~       | ~bsf~ instruction                               |
   
   #+begin_src c :tangle (eval c) :exports none 
#include <math.h>
#include <caml/mlvalues.h>
#include <caml/alloc.h>
   #+end_src

*** Erf

    #+begin_src c :tangle (eval c) :exports none 
CAMLprim value erf_float_bytecode(value x) {
  return copy_double(erf(Double_val(x)));
}  

CAMLprim double erf_float(double x) {
  return erf(x);
}
    #+end_src

    #+begin_src ocaml :tangle (eval mli)
external erf_float : float -> float = "erf_float_bytecode" "erf_float" [@@unboxed] [@@noalloc]
    #+end_src

    #+begin_src ocaml :tangle (eval ml) :exports none
external erf_float : float -> float = "erf_float_bytecode" "erf_float" [@@unboxed] [@@noalloc]
    #+end_src

*** Erfc

    #+begin_src c :tangle (eval c) :exports none
CAMLprim value erfc_float_bytecode(value x) {
  return copy_double(erfc(Double_val(x)));
}  

CAMLprim double erfc_float(double x) {
  return erfc(x);
}
    #+end_src

    #+begin_src ocaml :tangle (eval mli)
external erfc_float : float -> float = "erfc_float_bytecode" "erfc_float" [@@unboxed] [@@noalloc]
    #+end_src

    #+begin_src ocaml :tangle (eval ml) :exports none
external erfc_float : float -> float = "erfc_float_bytecode" "erfc_float" [@@unboxed] [@@noalloc]
    #+end_src

*** Gamma

    #+begin_src c :tangle (eval c) :exports none
CAMLprim value gamma_float_bytecode(value x) {
  return copy_double(tgamma(Double_val(x)));
}  


CAMLprim double gamma_float(double x) {
  return tgamma(x);
}
    #+end_src

    #+begin_src ocaml :tangle (eval mli)
external gamma_float : float -> float = "gamma_float_bytecode" "gamma_float" [@@unboxed] [@@noalloc]
    #+end_src

    #+begin_src ocaml :tangle (eval ml) :exports none
external gamma_float : float -> float = "gamma_float_bytecode" "gamma_float" [@@unboxed] [@@noalloc]
    #+end_src

*** Popcnt

    #+begin_src c :tangle (eval c) :exports none
CAMLprim int32_t popcnt(int64_t i) {
  return __builtin_popcountll (i);
}


CAMLprim value popcnt_bytecode(value i) {
  return caml_copy_int32(__builtin_popcountll (Int64_val(i)));
}
    #+end_src

    #+begin_src ocaml :tangle (eval mli)
val popcnt : int64 -> int
    #+end_src

    #+begin_src ocaml :tangle (eval ml) :exports none
external popcnt : int64 -> int32 = "popcnt_bytecode" "popcnt"
[@@unboxed] [@@noalloc]

let popcnt i = (popcnt [@inlined] ) i |> Int32.to_int 
    #+end_src

*** Trailz

    #+begin_src c :tangle (eval c) :exports none
CAMLprim int32_t trailz(int64_t i) {
  return __builtin_ctzll (i);
}


CAMLprim value trailz_bytecode(value i) {
  return caml_copy_int32(__builtin_ctzll (Int64_val(i)));
}
    #+end_src

    #+begin_src ocaml :tangle (eval mli)
val trailz : int64 -> int
    #+end_src

    #+begin_src ocaml :tangle (eval ml) :exports none
external trailz : int64 -> int32 = "trailz_bytecode" "trailz" "int"
[@@unboxed] [@@noalloc]

let trailz i = trailz i |> Int32.to_int 
    #+end_src

*** Leadz

    #+begin_src c :tangle (eval c) :exports none
CAMLprim int32_t leadz(int64_t i) {
  return __builtin_clzll(i);
}


CAMLprim value leadz_bytecode(value i) {
  return caml_copy_int32(__builtin_clzll (Int64_val(i)));
}
    #+end_src

    #+begin_src ocaml :tangle (eval mli)
val leadz : int64 -> int
    #+end_src

    #+begin_src ocaml :tangle (eval ml) :exports none
external leadz : int64 -> int32 = "leadz_bytecode" "leadz" "int"
[@@unboxed] [@@noalloc]

let leadz i = leadz i |> Int32.to_int 
    #+end_src

*** Test

    #+begin_src ocaml :tangle (eval test-ml) :exports none
let test_external () =
  check (float 1.e-15) "erf" 0.842700792949715 (erf_float 1.0);
  check (float 1.e-15) "erf" 0.112462916018285 (erf_float 0.1);
  check (float 1.e-15) "erf" (-0.112462916018285) (erf_float (-0.1));
  check (float 1.e-15) "erfc" 0.157299207050285 (erfc_float 1.0);
  check (float 1.e-15) "erfc" 0.887537083981715 (erfc_float 0.1);
  check (float 1.e-15) "erfc" (1.112462916018285) (erfc_float (-0.1));
  check (float 1.e-14) "gamma" (1.77245385090552) (gamma_float 0.5);
  check (float 1.e-14) "gamma" (9.51350769866873) (gamma_float (0.1));
  check (float 1.e-14) "gamma" (-3.54490770181103) (gamma_float (-0.5));
  check int "popcnt" 6 (popcnt @@ Int64.of_int 63);
  check int "popcnt" 8 (popcnt @@ Int64.of_int 299605);
  check int "popcnt" 1 (popcnt @@ Int64.of_int 65536);
  check int "popcnt" 0 (popcnt @@ Int64.of_int 0);
  check int "trailz" 3 (trailz @@ Int64.of_int 8);
  check int "trailz" 2 (trailz @@ Int64.of_int 12);
  check int "trailz" 0 (trailz @@ Int64.of_int 1);
  check int "trailz" 64 (trailz @@ Int64.of_int 0);
  check int "leadz" 60 (leadz @@ Int64.of_int 8);
  check int "leadz" 60 (leadz @@ Int64.of_int 12);
  check int "leadz" 63 (leadz @@ Int64.of_int 1);
  check int "leadz" 64 (leadz @@ Int64.of_int 0);
  ()
    #+end_src
    
** General functions

   #+begin_src ocaml :tangle (eval mli)
val fact : int -> float
(* @raise Invalid_argument for negative arguments or arguments >100. *)
val binom       : int -> int -> int
val binom_float : int -> int -> float

val chop              : float -> (unit -> float) -> float
val pow               : float -> int -> float
val float_of_int_fast : int -> float

val not_implemented : unit -> 'a
val of_some : 'a option -> 'a
   #+end_src

   | ~fact~              | Factorial function.                                                                                                       |
   | ~binom~             | Binomial coefficient. ~binom n k~ = $C_n^k$                                                                               |
   | ~binom_float~       | float variant of ~binom~                                                                                                  |
   | ~pow~               | Fast implementation of the power function for small integer powers                                                        |
   | ~chop~              | In ~chop a f~, evaluate ~f~ only if the absolute value of ~a~ is larger than ~Constants.epsilon~, and return ~a *. f ()~. |
   | ~float_of_int_fast~ | Faster implementation of float_of_int for small positive ints                                                             |
   | ~not_implemented~   | Fails with error message if some functionality is not implemented                                                         |
   | ~of_some~           | Extracts the value of an option                                                                                           |

   #+begin_src ocaml :tangle (eval ml) :exports none
let memo_float_of_int =
  Array.init 64 float_of_int

let float_of_int_fast i =
  if Int.logand i 63 = i then
    memo_float_of_int.(i)
  else
    float_of_int i


let factmax = 150

let fact_memo =
  let rec aux accu_l accu = function
    | 0 ->  (aux [@tailcall]) [1.] 1. 1
    | i when (i = factmax) ->
        let x = (float_of_int factmax) *. accu in
        List.rev (x::accu_l)
    | i ->  let x = (float_of_int i) *. accu in
        (aux [@tailcall]) (x::accu_l) x (i+1)
  in
  aux [] 0. 0
  |> Array.of_list

let fact = function
  | i when (i < 0) ->
      raise (Invalid_argument "Argument of factorial should be non-negative")
  | i when (i > 150) ->
      raise (Invalid_argument "Result of factorial is infinite")
  | i -> fact_memo.(i)


let binom =
  let memo = 
    let m = Array.make_matrix 64 64 0 in
    for n=0 to Array.length m - 1 do
      m.(n).(0) <- 1;
      m.(n).(n) <- 1;
      for k=1 to (n - 1) do
        m.(n).(k) <- m.(n-1).(k-1) + m.(n-1).(k)
      done
    done;
    m
  in
  let rec f n k =
    assert (k >= 0);
    assert (n >= k);
    if k = 0 || k = n then
      1
    else if n < 64 then 
      memo.(n).(k)
    else
      f (n-1) (k-1) + f (n-1) k
  in f


let binom_float n k =
  binom n k
  |> float_of_int_fast


let rec pow a = function
  | 0 -> 1.
  | 1 -> a
  | 2 -> a *. a
  | 3 -> a *. a *. a
  | -1 -> 1. /. a
  | n when  n > 0  ->
      let b = pow a (n / 2) in
      b *. b *. (if n mod 2 = 0 then 1. else a)
  | n when  n < 0  -> (pow [@tailcall]) (1./.a) (-n)
  | _ -> assert false


let chop f g =
  if (abs_float f) < Constants.epsilon then 0.
  else f *. (g ())


let not_implemented () = 
  failwith "Not implemented"


let of_some = function
  | Some a -> a
  | None   -> assert false

   #+end_src


   #+begin_src ocaml :tangle (eval test-ml) :exports none
let test_general () = 
  check int "of_some_of_int_fast" 1 (of_some (Some 1)) ;
  check int "binom" 35 (binom 7 4);
  check (float 1.e-15) "fact" 5040. (fact 7);
  check (float 1.e-15) "binom_float" 35.0 (binom_float 7 4);
  check (float 1.e-15) "pow" 729.0 (pow 3.0 6);
  check (float 1.e-15) "float_of_int_fast" 10.0 (float_of_int_fast 10);
  ()
   #+end_src

** Functions related to the Boys function

   #+begin_src ocaml :tangle (eval mli)
val incomplete_gamma : alpha:float -> float -> float
(* @raise Failure when the calculation doesn't converge. *)
   #+end_src

   The lower [[https://en.wikipedia.org/wiki/Incomplete_gamma_function][Incomplete Gamma function]] is implemented :
   \[
   \gamma(\alpha,x)  =  \int_0^x e^{-t}  t^{\alpha-1} dt
   \]

   p: $\frac{1}{\Gamma(\alpha)} \int_0^x e^{-t} t^{\alpha-1} dt$

   q: $\frac{1}{\Gamma(\alpha)} \int_x^\infty e^{-t} t^{\alpha-1} dt$

   Reference : Haruhiko Okumura: C-gengo niyoru saishin algorithm jiten
   (New Algorithm handbook in C language) (Gijyutsu hyouron sha,
   Tokyo, 1991) p.227 [in Japanese] 


   #+begin_src ocaml :tangle (eval ml) :exports none
let incomplete_gamma ~alpha x = 
  assert (alpha >= 0.);
  assert (x >= 0.);
  let a = alpha in
  let a_inv = 1./. a in
  let gf = gamma_float alpha in
  let loggamma_a = log gf in
  let rec p_gamma x =
    if x >= 1. +. a then 1. -. q_gamma x 
    else if x = 0. then 0.
    else
      let rec pg_loop prev res term k =
        if k > 1000. then failwith "p_gamma did not converge."
        else if prev = res then res
        else
          let term = term *. x /. (a +. k) in
          (pg_loop [@tailcall]) res (res +. term) term (k +. 1.)
      in
      let r0 =  exp (a *. log x -. x -. loggamma_a) *. a_inv in
      pg_loop min_float r0 r0 1.

  and q_gamma x =
    if x < 1. +. a then 1. -. p_gamma x 
    else
      let rec qg_loop prev res la lb w k =
        if k > 1000. then failwith "q_gamma did not converge."
        else if prev = res then res
        else
          let k_inv = 1. /. k in
          let kma = (k -. 1. -. a) *. k_inv in
          let la, lb =
            lb, kma *. (lb -. la) +. (k +. x) *. lb *. k_inv
          in
          let w = w *. kma  in
          let prev, res = res, res +. w /. (la *. lb) in
          (qg_loop [@tailcall]) prev res la lb w (k +. 1.)
      in
      let w = exp (a *. log x -. x -. loggamma_a) in
      let lb = (1. +. x -. a) in
      qg_loop min_float (w /. lb) 1. lb w 2.0
  in
  gf *. p_gamma x 
   #+end_src

   #+begin_src ocaml :tangle (eval mli)
val boys_function : maxm:int -> float -> float array
   #+end_src

   The [[https://link.springer.com/article/10.1007/s10910-005-9023-3][Generalized Boys function]] is implemented, 
 ~maxm~ is the maximum total angular momentum.
  
   \[
   F_m(x) = \frac{\gamma(m+1/2,x)}{2x^{m+1/2}}
   \]
   where $\gamma$ is the incomplete gamma function.

   - $F_0(0.) = 1$ 
   - $F_0(t) = \frac{\sqrt{\pi}}{2\sqrt{t}} \text{erf} ( \sqrt{t} )$ 
   - $F_m(0.) = \frac{1}{2m+1}$ 
   - $F_m(t)  = \frac{\gamma{m+1/2,t}}{2t^{m+1/2}}$
   - $F_m(t)  = \frac{ 2t\, F_{m+1}(t) + e^{-t} }{2m+1}$ 

   #+begin_src ocaml :tangle (eval ml) :exports none
let boys_function ~maxm t =
  assert (t >= 0.);
  match maxm with
  | 0 ->
      begin
        if t = 0. then [| 1. |] else
          let sq_t = sqrt t in
          [| (Constants.sq_pi_over_two /. sq_t) *.  erf_float sq_t |]
      end
  | _ ->
      begin
        assert (maxm > 0);
        let result =
          Array.init (maxm+1) (fun m -> 1. /. float_of_int (2*m+1))
        in
        let power_t_inv = (maxm+maxm+1)  in
        try  
          let fmax = 
            let t_inv = sqrt (1. /. t) in
            let n = float_of_int maxm in
            let dm = 0.5 +. n in
            let f  = (pow t_inv power_t_inv ) in
            match classify_float f with
            | FP_normal -> (incomplete_gamma ~alpha:dm t) *. 0.5 *. f
            | FP_zero
            | FP_subnormal -> 0.
            | _ -> raise Exit
          in
          let emt = exp (-. t) in
          result.(maxm) <- fmax;
          for n=maxm-1 downto 0 do
            result.(n) <- ( (t+.t) *. result.(n+1) +. emt) *. result.(n)
          done;
          result
        with Exit -> result
      end
   #+end_src

   #+begin_src ocaml :tangle (eval test-ml) :exports none
let test_boys () =
  check (float 1.e-15) "incomplete_gamma" 0.0 (incomplete_gamma ~alpha:0.5 0.);
  check (float 1.e-15) "incomplete_gamma" 1.114707979049507 (incomplete_gamma ~alpha:0.5 0.4);
  check (float 1.e-15) "incomplete_gamma" 1.4936482656248544 (incomplete_gamma ~alpha:0.5 1.);
  check (float 1.e-15) "incomplete_gamma" 1.7724401246392805 (incomplete_gamma ~alpha:0.5 10.);
  check (float 1.e-15) "incomplete_gamma" 1.7724538509055159 (incomplete_gamma ~alpha:0.5 100.);

  check (float 1.e-15) "boys" 1.0 (boys_function ~maxm:0 0.).(0);
  check (float 1.e-15) "boys" 0.2 (boys_function ~maxm:2 0.).(2);
  check (float 1.e-15) "boys" (1./.3.) (boys_function ~maxm:2 0.).(1);
  check (float 1.e-15) "boys" 0.8556243918921488 (boys_function ~maxm:0 0.5).(0);
  check (float 1.e-15) "boys" 0.14075053682591263 (boys_function ~maxm:2 0.5).(2);
  check (float 1.e-15) "boys" 0.00012711171070276764 (boys_function ~maxm:3 15.).(3);
  ()
   #+end_src

** List functions
   
   #+begin_src ocaml :tangle (eval mli)
val list_some  : 'a option list -> 'a list
val list_range : int -> int -> int list
val list_pack  : int -> 'a list -> 'a list list
   #+end_src

   | ~list_some~  | Filters out all ~None~ elements of the list, and returns the elements without the ~Some~ |
   | ~list_range~ | Creates a list of consecutive integers                                                   |
   | ~list_pack~  | ~list_pack n l~ Creates a list of ~n~-elements lists                                     |

   #+begin_src ocaml :tangle (eval ml) :exports none
let list_some l =
  List.filter (function None -> false | _ -> true) l
  |> List.rev_map (function Some x -> x | _ -> assert false)
  |> List.rev


let list_range first last = 
  if last < first then [] else
    let rec aux accu = function
      | 0 -> first :: accu
      | i -> (aux [@tailcall]) ( (first+i)::accu ) (i-1)
    in
    aux [] (last-first)


let list_pack n l =
  assert (n>=0);
  let rec aux i accu1 accu2 = function
    | [] -> if accu1 = [] then
          List.rev accu2
        else 
          List.rev ((List.rev accu1) :: accu2)
    | a :: rest ->
        match i with
        | 0 -> (aux [@tailcall]) (n-1)  [] ((List.rev (a::accu1)) :: accu2) rest
        | _ -> (aux [@tailcall]) (i-1)  (a::accu1) accu2 rest
  in
  aux (n-1) [] [] l

   #+end_src

   #+begin_src ocaml :tangle (eval test-ml) :exports none
let test_list () =
  check bool "list_range" true ([ 2; 3; 4 ] = list_range 2 4);
  check bool "list_some" true ([ 2; 3; 4 ] =
                               list_some ([ None ; Some 2 ; None ; Some 3 ; None ; None ; Some 4]) );
  check bool "list_pack" true (list_pack 3 (list_range 1 20) =
                               [[1; 2; 3]; [4; 5; 6]; [7; 8; 9]; [10; 11; 12]; [13; 14; 15];
                                [16; 17; 18]; [19; 20]]);
  ()
   #+end_src

** Array functions
   
   #+begin_src ocaml :tangle (eval mli)
val array_range   : int -> int -> int array
val array_sum     : float array -> float
val array_product : float array -> float
   #+end_src

   | ~array_range~   | Creates an array of consecutive integers             |
   | ~array_sum~     | Returns the sum of all the elements of the array     |
   | ~array_product~ | Returns the product of all the elements of the array |

   #+begin_src ocaml :tangle (eval ml) :exports none
let array_range first last = 
  if last < first then [| |] else
    Array.init (last-first+1) (fun i -> i+first)


let array_sum a = 
  Array.fold_left ( +. ) 0. a


let array_product a = 
  Array.fold_left ( *. ) 1. a
   #+end_src

   #+begin_src ocaml :tangle (eval test-ml) :exports none
let test_array () =
  check bool "array_range" true ([| 2; 3; 4 |] = array_range 2 4);
  check (float 1.e-15) "array_sum" 9. (array_sum [| 2.; 3.; 4. |]);
  check (float 1.e-15) "array_product" 24. (array_product [| 2.; 3.; 4. |]);
  ()
   #+end_src

** Stream functions
   
   #+begin_src ocaml :tangle (eval mli)
val stream_range   : int -> int -> int Stream.t
val stream_to_list : 'a Stream.t -> 'a list
val stream_fold    : ('a -> 'b -> 'a) -> 'a -> 'b Stream.t -> 'a
   #+end_src

   | ~stream_range~   | Creates a stream returning consecutive integers |
   | ~stream_to_list~ | Read a stream and put items in a list           |
   | ~stream_fold~    | Apply a fold to the elements of the stream      |

   #+begin_src ocaml :tangle (eval ml) :exports none
let stream_range first last = 
  Stream.from (fun i ->
    let result = i+first in
    if result <= last then
      Some result
    else None
  )

let stream_to_list stream = 
  let rec aux accu =
    let new_accu = 
      try
        Some (Stream.next stream :: accu)
      with Stream.Failure -> None
    in
    match new_accu with
    | Some new_accu -> (aux [@tailcall]) new_accu
    | None -> accu
  in List.rev @@ aux []


let stream_fold f init stream =
  let rec aux accu =
    let new_accu = 
      try
        let element = Stream.next stream in
        Some (f accu element)
      with Stream.Failure -> None
    in  
    match new_accu with
    | Some new_accu -> (aux [@tailcall]) new_accu
    | None -> accu
  in
  aux init

   #+end_src

** Printers

   #+begin_src ocaml :tangle (eval mli)
val pp_float_array_size   : Format.formatter -> float array -> unit
val pp_float_array        : Format.formatter -> float array -> unit
val pp_float_2darray_size : Format.formatter -> float array array -> unit
val pp_float_2darray      : Format.formatter -> float array array -> unit
val pp_bitstring          : int -> Format.formatter -> Z.t -> unit
   #+end_src

   | ~pp_float_array~        | Printer for float arrays                             |
   | ~pp_float_array_size~   | Printer for float arrays with size                   |
   | ~pp_float_2darray~      | Printer for matrices                                 |
   | ~pp_float_2darray_size~ | Printer for matrices with size                       |
   | ~pp_bitstring~          | Printer for bit strings (used by ~Bitstring~ module) |

   #+begin_example
pp_float_array_size:
[ 6:   1.000000   1.732051   1.732051   1.000000   1.732051   1.000000 ]

pp_float_array: 
[    1.000000   1.732051   1.732051   1.000000   1.732051   1.000000 ]

pp_float_2darray_size
[
  2:[ 6:  1.000000   1.732051   1.732051   1.000000   1.732051   1.000000 ]
    [ 4:  1.000000   2.000000   3.000000   4.000000 ] ]

pp_float_2darray:
[ [   1.000000   1.732051   1.732051   1.000000   1.732051   1.000000 ]
  [   1.000000   2.000000   3.000000   4.000000 ] ]

pp_bitstring 14:
+++++------+--
   #+end_example

   #+begin_src ocaml :tangle (eval ml) :exports none
let pp_float_array ppf a = 
  Format.fprintf ppf "@[<2>[@ ";
  Array.iter (fun f -> Format.fprintf ppf "@[%10f@]@ " f) a;
  Format.fprintf ppf "]@]"

let pp_float_array_size ppf a = 
  Format.fprintf ppf "@[<2>@[ %d:@[<2>" (Array.length a);
  Array.iter (fun f -> Format.fprintf ppf "@[%10f@]@ " f) a;
  Format.fprintf ppf "]@]@]"

let pp_float_2darray ppf a = 
  Format.fprintf ppf "@[<2>[@ ";
  Array.iter (fun f -> Format.fprintf ppf "@[%a@]@ " pp_float_array f) a;
  Format.fprintf ppf "]@]"

let pp_float_2darray_size ppf a = 
  Format.fprintf ppf "@[<2>@[ %d:@[" (Array.length a);
  Array.iter (fun f -> Format.fprintf ppf "@[%a@]@ " pp_float_array_size f) a;
  Format.fprintf ppf "]@]@]"

let pp_bitstring n ppf bs =
  String.init n (fun i -> if (Z.testbit bs i) then '+' else '-')
  |> Format.fprintf ppf "@[<h>%s@]"
   #+end_src


** Test footer                                                     :noexport:

   #+begin_src ocaml :tangle (eval test-ml) :exports none
let tests = [
  "External", `Quick, test_external;
  "General" , `Quick, test_general;
  "Boys"    , `Quick, test_boys;
  "List"    , `Quick, test_list;
  "Array"   , `Quick, test_array;
]
   #+end_src