Source

papl / src / PaplSampler.ml

(*
  Copyright (c) 2012 Anders Lau Olsen.
  See LICENSE file for terms and conditions.
*)
type rng_t = PaplRandom.rng_t

type 'a pair_t = 'a * 'a

type 'a t = 'a BatEnum.t

type range_t = float * float

type int_range_t = int * int

exception RangeError of string * (float * float)

exception IntRangeError of string * (int * int)

let next_fun s =
  let obj = BatEnum.to_object s in
    fun () -> obj#next

let pi = BatFloat.pi

let map2 f sa sb =
  let next_a = next_fun sa in
  let next_b = next_fun sb in
  let next () =
    let a = next_a () in
    let b = next_b () in
      f a b
  in
    BatEnum.from next

let product2 sa sb = BatEnum.combine (sa, sb)

let product3 sa sb sc =
  let next_a = next_fun sa in
  let next_b = next_fun sb in
  let next_c = next_fun sc in
  let next () =
    let a = next_a () in
    let b = next_b () in
    let c = next_c () in
      (a, b, c)
  in
    BatEnum.from next

let pair s = product2 s s

let intersperse sx total sy =
  let n = ref total in
  let cnt = ref 0 in
  let next_x = next_fun sx in
  let rec next () =
    if !cnt == 0 then
      let x = next_x () in
        cnt := !n;
        x
    else begin
      cnt := !cnt - 1;
      match BatEnum.get sy with
          None ->
            n := 0;
            cnt := 0;
            next ()
        | Some y ->
            y
    end
  in BatEnum.from next

let rec constrain sampler constr =
  BatEnum.filter (PaplConstraint.accept constr) sampler

let constrain_option sampler constr =
  BatEnum.map
    (fun x ->
       if PaplConstraint.accept constr x
       then Some x
       else None)
    sampler

let rec constrain_with_default sampler constr default =
  let next_sample = next_fun sampler in
  let rec next () =
    let x = next_sample () in
      if PaplConstraint.reject constr x then
        match BatEnum.get default with
            None -> next ()
          | Some y -> y
      else
        x
  in BatEnum.from next

let rec round_robin samplers =
  let ss = ref samplers in
  let acc = ref [] in
  let rec next () =
    match BatEnum.get !ss, !acc with
        None, [] -> raise BatEnum.No_more_elements
      | None, _ ->
          ss := BatList.backwards !acc;
          acc := [];
          next ()
      | Some s, _ ->
          match BatEnum.get s with
              None -> next ()
            | Some x ->
                acc := s :: !acc;
                x
  in BatEnum.from next

(*
  The values are inserted in a balanced binary tree with each node storing the
  total weight of its values. The nodes are stored in an array as is usually
  done for heaps.
*)
let distribute_by_weight_helper make_sampler pairs =
  let tree = Array.of_list pairs in
  let left i = 2 * i + 1 in
  let right i = 2 * i + 2 in
  let get i = tree.(i) in
  let valid i = i < Array.length tree in
  let rec build i =
    if valid i then begin
      let wl = build (left i) in
      let wr = build (right i) in
      let (w, x) = get i in
      let w' = w +. wl +. wr in
        tree.(i) <- (w', x);
        w'
    end
    else 0.0
  in
  let max = build 0 in
  let position_sampler = make_sampler max in
  let weight i = if valid i then fst (tree.(i)) else 0.0 in
  let rec lookup i p =
    let wl = weight (left i) in
    if p < wl then
      lookup (left i) p
    else
      let p' = p -. wl in
      let wr = weight (right i) in
        if p' < wr then
          lookup (right i) p'
        else
          snd (get i) in
  let next_position = next_fun position_sampler in
  let next () = lookup 0 (next_position ()) in
    if Array.length tree == 0 then
      invalid_arg "Empty enum of pairs given."
    else
      BatEnum.from next

let distribute_by_weight ?rng pairs =
  distribute_by_weight_helper (PaplRandom.enum_float ?rng) pairs

let uniform_helper enum_fun (a, b) =
  if a > b then BatEnum.empty ()
  else BatEnum.map ((+.) a) (enum_fun (b -. a))

let get_uniform_helper float_fun (a, b) =
  if a > b then
    raise (RangeError ("a > b", (a, b)))
  else a +. float_fun (b -. a)

let get_uniform ?rng =
  let float_fun = PaplRandom.float ?rng in
    fun range -> get_uniform_helper float_fun range

let uniform ?rng range = uniform_helper (PaplRandom.enum_float ?rng) range

let uniform_int_helper enum_fun (a, b) =
  if a >= b then BatEnum.empty ()
  else if a = b - 1 then BatEnum.repeat a
  else BatEnum.map ((+) a) (enum_fun (b - a))

let uniform_int ?rng range =
  uniform_int_helper (PaplRandom.enum_int ?rng) range

(* Given a function that returns pairs (x, y), construct a caching function that
   returns first x and then y.
*)
let sequence_pairs pair_fun =
  let result = ref None in
  let next () =
    match !result with
      | None ->
          let (x, y) = pair_fun () in
            result := Some y;
            x
      | Some y ->
          result := None;
          y
  in next

(* Normally distributed numbers by the Box-Muller transform.

   http://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform

   Boost Random uses this algorithm and so does Python's random module.
*)
let box_muller_gaussian float_fun =
  let r () = float_fun 1.0 in (* Range [0, 1) *)
  let pairs () =
    let u1, u2 = r (), r () in
    let rho = sqrt (-2. *. log (1. -. u1)) in
    let a = 2. *. pi *. u2 in
      (rho *. cos a, rho *. sin a)
  in let next = sequence_pairs pairs in
    BatEnum.from next

let gaussian_helper = box_muller_gaussian

let gaussian ?rng () =
  gaussian_helper (PaplRandom.float ?rng)

let mean_and_variance stream =
  let xs = BatList.of_enum stream in
  let n = List.length xs in
    if n < 2 then
      invalid_arg "stream is shorter than 2\n"
    else
      let n' = float_of_int n in
      let sum = BatList.fsum xs in
      let mean = sum /. n' in
      let variance =
        let sum_of_squares =
          List.fold_right (fun x acc -> acc +. x *. x) xs 0.
        in (sum_of_squares -. n' *. mean *. mean) /. (n' -. 1.)
      in (mean, variance)

(*

We should add a bunch of useful sampler, perhaps here, perhaps in their own
modules. Things we can sample:

  - Configuration space boxes [done: see PaplQ].
  - Rotation matrices
  - Transformation matrices
  - Sphere surfaces
  - Sphere volumes
  - Triangle surfaces
  - Polygon surfaces
  - Deterministic samplers for spheres, boxes, etc.

Some of these things we should leave out or keep in their own modules. We want a
compact and precise library, not something stuffed with every arbitrary
math-like function that other people are also writing all the time. We want to
show that we can integrate things. We don't want to provide all the things that
can be integrated. Or at least not at this stage. It might be worthwhile to aim
for a larger library if the library had a group of users.

We really want to keep things separate. There should pretty much be zero
concrete types anywhere in the core set of modules.

*)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.