Source file pdfs.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
let poisson ~lambda k =
  if lambda <=. 0.0 then invalid_arg "Pdfs.poisson"
  else
    let log_lambda = log lambda in
    exp @@ ((float_of_int k *. log_lambda) -. lambda -. Specfun.log_factorial k)

let poisson_ln ~lambda k =
  if lambda <=. 0.0 then invalid_arg "Pdfs.poisson_ln"
  else
    let log_lambda = log lambda in
    Log_space.unsafe_cast
      ((float_of_int k *. log_lambda) -. lambda -. Specfun.log_factorial k)

let pi = acos ~-.1.0

let gaussian ~mean ~std x =
  if std <=. 0.0 then invalid_arg "Pdfs.gaussian" ;
  let normalizer = 1.0 /. (std *. (sqrt @@ (2.0 *. pi))) in
  let delta = (x -. mean) /. std in
  let delta_sq = delta *. delta in
  normalizer *. (exp @@ (~-.0.5 *. delta_sq))

let gaussian_ln ~mean ~std x =
  if std <=. 0.0 then invalid_arg "Pdfs.gaussian_ln" ;
  let normalizer = -.log std -. (0.5 *. log (2. *. pi)) in
  let delta = (x -. mean) /. std in
  let delta_sq = delta *. delta in
  Log_space.unsafe_cast (normalizer -. (0.5 *. delta_sq))

let exponential ~rate x = rate *. exp (~-.rate *. x)

let exponential_ln ~rate x = Log_space.unsafe_cast (log rate -. (rate *. x))

let geometric ~p k =
  if p <=. 0.0 || p >. 1.0 then invalid_arg "Pdfs.geometric" ;
  if k < 0 then invalid_arg "Pdfs.geometric" ;
  ((1. -. p) ** float_of_int k) *. p

let geometric_ln ~p k =
  if p <=. 0.0 || p >. 1.0 then invalid_arg "Pdfs.geometric_ln" ;
  if k < 0 then invalid_arg "Pdfs.geometric_ln" ;
  Log_space.unsafe_cast (log p +. (float_of_int k *. log (1. -. p)))

let uniform { Stats_intf.min; max } (x : float) =
  let delta = max -. min in
  if delta >=. 0. then if x >=. min && x <=. max then 1. /. delta else 0.0
  else invalid_arg "uniform"

let uniform_ln { Stats_intf.min; max } (x : float) =
  let delta = max -. min in
  if delta >=. 0. then
    if x >=. min && x <=. max then Log_space.unsafe_cast ~-.(log delta)
    else Log_space.zero
  else invalid_arg "uniform_ln"

let rectangle_ln ~min ~max (x : float array) =
  if Array.length min <> Array.length max then
    invalid_arg "rectangle_ln: Array.length min <> Array.length max" ;
  if Array.length min = 0 then invalid_arg "rectangle_ln: min has zero length" ;
  let deltas = Array.map2 (fun x y -> y -. x) min max in
  if Array.exists (fun delta -> delta <. 0.0) deltas then
    invalid_arg "rectangle_ln: invalid box, some dimension is empty" ;
  let exception Out_of_range in
  let acc = ref 0.0 in
  try
    for i = 0 to Array.length min - 1 do
      let xi = x.(i) in
      if xi >=. min.(i) || xi <=. max.(i) then acc := !acc -. log deltas.(i)
      else raise Out_of_range
    done ;
    Log_space.unsafe_cast !acc
  with Out_of_range -> Log_space.zero

let rectangle ~min ~max x = Log_space.to_float (rectangle_ln ~min ~max x)

let binomial_ln ~p ~n k =
  if p <=. 0.0 || p >. 1.0 then invalid_arg "Pdfs.binomial_ln" ;
  if n < 0 then invalid_arg "Pdfs.binomial_ln" ;
  if k < 0 || k > n then invalid_arg "Pdfs.binomial_ln" ;
  (*
    log (bincoeff n k) + k log p + (n-k) log (1-p)
    log (bincoeff n k) = log (n!) - (log (k!) + log ((n-k)!))
   *)
  let log_bincoeff =
    let open Specfun in
    log_factorial n -. log_factorial k -. log_factorial (n - k)
  in
  log_bincoeff
  +. (float_of_int k *. log p)
  +. (float_of_int (n - k) *. log (1. -. p))
  |> Log_space.unsafe_cast

let binomial ~p ~n k = Log_space.to_float (binomial_ln ~p ~n k)

let gamma_ln ~shape ~scale x =
  if shape <= 0 then invalid_arg "Pdfs.gamma_ln" ;
  if scale <=. 0.0 then invalid_arg "Pdfs.gamma_ln" ;
  let fshape = float_of_int shape in
  ((fshape -. 1.) *. log x)
  -. (x /. scale)
  -. (fshape *. log scale)
  -. Specfun.log_factorial (shape - 1)
  |> Log_space.unsafe_cast

let gamma ~shape ~scale x = Log_space.to_float (gamma_ln ~shape ~scale x)

let tup2 p1 p2 (a, b) = p1 a *. p2 b

let tup3 p1 p2 p3 (a, b, c) = p1 a *. p2 b *. p3 c

let tup4 p1 p2 p3 p4 (a, b, c, d) = p1 a *. p2 b *. p3 c *. p4 d

let tup5 p1 p2 p3 p4 p5 (a, b, c, d, e) = p1 a *. p2 b *. p3 c *. p4 d *. p5 e

let tup6 p1 p2 p3 p4 p5 p6 (a, b, c, d, e, f) =
  p1 a *. p2 b *. p3 c *. p4 d *. p5 e *. p6 f

let ( *: ) = Log_space.mul

let tup2_ln p1 p2 (a, b) = p1 a *: p2 b

let tup3_ln p1 p2 p3 (a, b, c) = p1 a *: p2 b *: p3 c

let tup4_ln p1 p2 p3 p4 (a, b, c, d) = p1 a *: p2 b *: p3 c *: p4 d

let tup5_ln p1 p2 p3 p4 p5 (a, b, c, d, e) =
  p1 a *: p2 b *: p3 c *: p4 d *: p5 e

let tup6_ln p1 p2 p3 p4 p5 p6 (a, b, c, d, e, f) =
  p1 a *: p2 b *: p3 c *: p4 d *: p5 e *: p6 f

let mixture_ln coeffs log_pdfs =
  if Array.length coeffs = 0 then invalid_arg "mixture_ln: empty coeffs" ;
  let sum = ref 0.0 in
  for i = 0 to Array.length coeffs - 1 do
    let c = coeffs.(i) in
    if c <. 0.0 then invalid_arg "mixture_ln: negative weight" ;
    sum := !sum +. coeffs.(i)
  done ;
  if abs_float (1. -. !sum) >. 0.001 then
    invalid_arg "mixture_ln: unnormalized weights" ;
  let log_coeffs = Array.map Log_space.of_float coeffs in
  fun x ->
    let open Log_space in
    let acc = ref one in
    for i = 0 to Array.length coeffs - 1 do
      acc := mul !acc (mul log_coeffs.(i) (log_pdfs i x))
    done ;
    !acc

let mixture coeffs pdfs x =
  mixture_ln coeffs (fun i x -> Log_space.of_float (pdfs i x)) x
  |> Log_space.to_float