Commit 3ca117c0 authored by Philippe Veber's avatar Philippe Veber
Browse files

tk: new Discrete_pd module

parent 8ac66b28
open Core_kernel
type t = {
n : int ;
shift : int ;
weights : float array ;
}
let is_leaf dpd i = i >= dpd.shift
let init n ~f =
let shift = Float.(to_int (2. ** round_up (log (float n) /. log 2.))) - 1 in
let m = shift + n in
let weights = Array.create ~len:m 0. in
for i = 0 to n - 1 do
weights.(shift + i) <- f (i)
done ;
for i = shift - 1 downto 0 do
if 2 * i + 1 < m then weights.(i) <- weights.(2 * i + 1) ;
if 2 * i + 2 < m then weights.(i) <- weights.(i) +. weights.(2 * i + 2)
done ;
{ n ; shift ; weights }
let draw dpd rng =
let x = dpd.weights.(0) *. Gsl.Rng.uniform rng in
let rec loop acc i =
if is_leaf dpd i then i
else if Float.( >= ) (acc +. dpd.weights.(2 * i + 1)) x then
loop acc (2 * i + 1)
else loop (acc +. dpd.weights.(2 * i + 1)) (2 * i + 2)
in
loop 0. 0 - dpd.shift
let update dpd i w_i =
let m = Array.length dpd.weights in
let j = i + dpd.shift in
dpd.weights.(j) <- w_i ;
let rec loop k =
dpd.weights.(k) <- dpd.weights.(2 * k + 1) ;
if 2 * k + 2 < m then dpd.weights.(k) <- dpd.weights.(k) +. dpd.weights.(2 * k + 2) ;
if k > 0 then loop ((k - 1) / 2)
in
loop ((j - 1) / 2)
let total_weight dpd = dpd.weights.(0)
let demo ~n ~ncat =
let rng = Gsl.Rng.(make (default ())) in
let probs = Array.init ncat ~f:(fun _ -> Gsl.Rng.uniform rng) in
let sum = Array.fold probs ~init:0. ~f:( +. ) in
let pd = init ncat ~f:(fun _ -> 0.) in
let counts = Array.create ~len:ncat 0 in
Array.iteri probs ~f:(update pd) ;
for _ = 1 to n do
let k = draw pd rng in
counts.(k) <- counts.(k) + 1
done ;
Array.map probs ~f:(fun x -> x /. sum),
Array.map counts ~f:(fun k -> float k /. float n)
(* Updatable discrete probability distribution *)
type t
val init : int -> f:(int -> float) -> t
val update : t -> int -> float -> unit
val draw : t -> Gsl.Rng.t -> int
val total_weight : t -> float
val demo : n:int -> ncat:int -> float array * float array
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment