10
1
mirror of https://gitlab.com/scemama/QCaml.git synced 2024-12-23 04:43:32 +01:00
QCaml/common/lib/util.ml

286 lines
6.9 KiB
OCaml
Raw Normal View History

2024-01-17 10:30:24 +01:00
(** Utility functions *)
2021-01-28 00:34:26 +01:00
external erf_float : float -> float
= "erf_float_bytecode" "erf_float" [@@unboxed] [@@noalloc]
2018-02-24 23:57:38 +01:00
2020-12-28 01:08:55 +01:00
external erfc_float : float -> float = "erfc_float_bytecode" "erfc_float" [@@unboxed] [@@noalloc]
2018-02-01 22:39:23 +01:00
2021-01-28 00:34:26 +01:00
external gamma_float : float -> float
= "gamma_float_bytecode" "gamma_float" [@@unboxed] [@@noalloc]
2018-02-24 23:57:38 +01:00
2019-03-03 01:43:04 +01:00
external popcnt : int64 -> int32 = "popcnt_bytecode" "popcnt"
2020-12-28 01:08:55 +01:00
[@@unboxed] [@@noalloc]
2019-03-03 01:43:04 +01:00
2020-12-28 01:08:55 +01:00
let popcnt i = (popcnt [@inlined] ) i |> Int32.to_int
2019-03-03 01:43:04 +01:00
2019-11-13 11:07:39 +01:00
external trailz : int64 -> int32 = "trailz_bytecode" "trailz" "int"
2020-12-28 01:08:55 +01:00
[@@unboxed] [@@noalloc]
2019-03-03 01:43:04 +01:00
2020-12-28 01:08:55 +01:00
let trailz i = trailz i |> Int32.to_int
2019-03-03 01:43:04 +01:00
2019-11-13 11:07:39 +01:00
external leadz : int64 -> int32 = "leadz_bytecode" "leadz" "int"
2020-12-28 01:08:55 +01:00
[@@unboxed] [@@noalloc]
2018-02-24 23:57:38 +01:00
2020-12-28 01:08:55 +01:00
let leadz i = leadz i |> Int32.to_int
2018-01-17 15:56:57 +01:00
2018-02-01 22:53:00 +01:00
2019-04-05 15:53:15 +02:00
2020-02-17 19:45:53 +01:00
let memo_float_of_int =
Array.init 64 float_of_int
let float_of_int_fast i =
2020-12-28 01:08:55 +01:00
if Int.logand i 63 = i then
memo_float_of_int.(i)
else
float_of_int i
2020-02-17 19:45:53 +01:00
2019-04-05 15:53:15 +02:00
let factmax = 150
2018-01-17 15:56:57 +01:00
let fact_memo =
let rec aux accu_l accu = function
2019-09-10 18:39:14 +02:00
| 0 -> (aux [@tailcall]) [1.] 1. 1
2018-02-03 19:01:30 +01:00
| i when (i = factmax) ->
2020-12-28 01:08:55 +01:00
let x = (float_of_int factmax) *. accu in
List.rev (x::accu_l)
2018-02-03 19:01:30 +01:00
| i -> let x = (float_of_int i) *. accu in
2020-12-28 01:08:55 +01:00
(aux [@tailcall]) (x::accu_l) x (i+1)
2018-01-17 15:56:57 +01:00
in
aux [] 0. 0
|> Array.of_list
let fact = function
| i when (i < 0) ->
2020-12-28 01:08:55 +01:00
raise (Invalid_argument "Argument of factorial should be non-negative")
2018-01-17 15:56:57 +01:00
| i when (i > 150) ->
2020-12-28 01:08:55 +01:00
raise (Invalid_argument "Result of factorial is infinite")
2018-01-17 15:56:57 +01:00
| i -> fact_memo.(i)
2020-02-17 19:45:53 +01:00
let binom =
2024-01-17 10:30:24 +01:00
let memo =
2020-12-28 01:08:55 +01:00
let m = Array.make_matrix 64 64 0 in
2020-03-26 16:24:41 +01:00
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
2020-02-17 19:45:53 +01:00
in
2020-03-26 16:24:41 +01:00
let rec f n k =
assert (k >= 0);
assert (n >= k);
if k = 0 || k = n then
1
2024-01-17 10:30:24 +01:00
else if n < 64 then
2020-12-28 01:08:55 +01:00
memo.(n).(k)
2020-02-17 19:45:53 +01:00
else
2020-03-26 16:24:41 +01:00
f (n-1) (k-1) + f (n-1) k
in f
2020-02-17 19:45:53 +01:00
let binom_float n k =
binom n k
|> float_of_int_fast
2019-02-19 17:36:07 +01:00
2020-03-26 16:24:41 +01:00
2018-01-17 15:56:57 +01:00
let rec pow a = function
| 0 -> 1.
| 1 -> a
| 2 -> a *. a
| 3 -> a *. a *. a
| -1 -> 1. /. a
2018-02-24 23:57:38 +01:00
| n when n > 0 ->
2020-12-28 01:08:55 +01:00
let b = pow a (n / 2) in
b *. b *. (if n mod 2 = 0 then 1. else a)
2019-09-10 18:39:14 +02:00
| n when n < 0 -> (pow [@tailcall]) (1./.a) (-n)
2018-02-24 23:57:38 +01:00
| _ -> assert false
2018-01-22 23:19:24 +01:00
let chop f g =
2018-02-24 23:57:38 +01:00
if (abs_float f) < Constants.epsilon then 0.
2018-01-22 23:19:24 +01:00
else f *. (g ())
2018-02-01 22:19:23 +01:00
2021-01-01 16:39:33 +01:00
exception Not_implemented of string
2024-01-17 10:30:24 +01:00
let not_implemented string =
2021-01-01 16:39:33 +01:00
raise (Not_implemented string)
2018-02-01 22:19:23 +01:00
2020-02-17 19:45:53 +01:00
2020-12-28 01:08:55 +01:00
let of_some = function
| Some a -> a
| None -> assert false
2024-01-17 10:30:24 +01:00
let incomplete_gamma ~alpha x =
2020-12-28 01:08:55 +01:00
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 =
2024-01-17 10:30:24 +01:00
if x >= 1. +. a then 1. -. q_gamma x
2020-12-28 01:08:55 +01:00
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 =
2024-01-17 10:30:24 +01:00
if x < 1. +. a then 1. -. p_gamma x
2020-12-28 01:08:55 +01:00
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
let boys_function ~maxm t =
2020-09-26 12:02:53 +02:00
assert (t >= 0.);
match maxm with
| 0 ->
2020-12-28 01:08:55 +01:00
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
| _ ->
2020-12-28 01:08:55 +01:00
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
2024-01-17 10:30:24 +01:00
try
let fmax =
let t_inv = sqrt (1. /. t) in
let n = float_of_int maxm in
let dm = 0.5 +. n in
2020-03-27 17:58:50 +01:00
let f = (pow t_inv power_t_inv ) in
2018-06-27 13:13:59 +02:00
match classify_float f with
2020-12-28 01:08:55 +01:00
| 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)
2020-03-27 17:58:50 +01:00
done;
result
2020-12-28 01:08:55 +01:00
with Exit -> result
end
2019-03-01 10:30:02 +01:00
2018-03-22 00:29:14 +01:00
let list_some l =
List.filter (function None -> false | _ -> true) l
2020-03-26 17:43:11 +01:00
|> List.rev_map (function Some x -> x | _ -> assert false)
|> List.rev
2018-03-22 00:29:14 +01:00
2024-01-17 10:30:24 +01:00
let list_range first last =
2019-02-27 14:56:59 +01:00
if last < first then [] else
2020-12-28 01:08:55 +01:00
let rec aux accu = function
| 0 -> first :: accu
| i -> (aux [@tailcall]) ( (first+i)::accu ) (i-1)
in
aux [] (last-first)
2019-02-27 14:56:59 +01:00
let list_pack n l =
2020-02-05 23:32:55 +01:00
assert (n>=0);
2019-02-27 14:56:59 +01:00
let rec aux i accu1 accu2 = function
2020-12-28 01:08:55 +01:00
| [] -> if accu1 = [] then
List.rev accu2
2024-01-17 10:30:24 +01:00
else
2020-12-28 01:08:55 +01:00
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
2019-02-27 14:56:59 +01:00
in
aux (n-1) [] [] l
2018-06-28 14:43:24 +02:00
2024-01-17 10:30:24 +01:00
let array_range first last =
2020-12-28 01:08:55 +01:00
if last < first then [| |] else
Array.init (last-first+1) (fun i -> i+first)
2024-01-17 10:30:24 +01:00
let array_sum a =
2020-12-28 01:08:55 +01:00
Array.fold_left ( +. ) 0. a
2024-01-17 10:30:24 +01:00
let array_product a =
2020-12-28 01:08:55 +01:00
Array.fold_left ( *. ) 1. a
2024-01-17 10:30:24 +01:00
let seq_range first last =
2023-04-24 13:05:01 +02:00
Seq.init (last-first) (fun i -> i+first)
2024-01-17 10:30:24 +01:00
let seq_to_list seq =
let rec aux accu xs =
2023-04-24 13:05:01 +02:00
match Seq.uncons xs with
| Some (x, xs) -> aux (x::accu) xs
| None -> List.rev accu
2019-03-20 23:10:53 +01:00
in
2023-04-24 13:05:01 +02:00
aux [] seq
let seq_fold f init seq =
Seq.fold_left f init seq
2024-01-17 10:30:24 +01:00
let pp_float_array ppf a =
2020-12-28 01:08:55 +01:00
Format.fprintf ppf "@[<2>[@ ";
Array.iter (fun f -> Format.fprintf ppf "@[%10f@]@ " f) a;
Format.fprintf ppf "]@]"
2018-03-15 15:25:49 +01:00
2024-01-17 10:30:24 +01:00
let pp_float_array_size ppf a =
2018-03-20 14:11:31 +01:00
Format.fprintf ppf "@[<2>@[ %d:@[<2>" (Array.length a);
2018-03-15 15:25:49 +01:00
Array.iter (fun f -> Format.fprintf ppf "@[%10f@]@ " f) a;
2018-03-20 14:11:31 +01:00
Format.fprintf ppf "]@]@]"
2018-03-15 15:25:49 +01:00
2024-01-17 10:30:24 +01:00
let pp_float_2darray ppf a =
2018-03-20 14:11:31 +01:00
Format.fprintf ppf "@[<2>[@ ";
Array.iter (fun f -> Format.fprintf ppf "@[%a@]@ " pp_float_array f) a;
Format.fprintf ppf "]@]"
2024-01-17 10:30:24 +01:00
let pp_float_2darray_size ppf a =
2018-03-20 14:11:31 +01:00
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 "]@]@]"
2019-02-20 19:43:16 +01:00
let pp_bitstring n ppf bs =
String.init n (fun i -> if (Z.testbit bs i) then '+' else '-')
2020-12-28 01:08:55 +01:00
|> Format.fprintf ppf "@[<h>%s@]"