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
|
(* Code under Apache License 2.0 - Jane Street Group, LLC <opensource@janestreet.com> *)
let col_norm a column =
let acc = ref 0. in
for i = 0 to Array.length a - 1 do
let entry = a.(i).(column) in
acc := !acc +. (entry *. entry)
done ;
sqrt !acc
let col_inner_prod t j1 j2 =
let acc = ref 0. in
for i = 0 to Array.length t - 1 do
acc := !acc +. (t.(i).(j1) *. t.(i).(j2))
done ;
!acc
let qr_in_place a =
let m = Array.length a in
if m = 0 then ([||], [||])
else
let n = Array.length a.(0) in
let r = Array.make_matrix n n 0. in
for j = 0 to n - 1 do
let alpha = col_norm a j in
r.(j).(j) <- alpha ;
let one_over_alpha = 1. /. alpha in
for i = 0 to m - 1 do
a.(i).(j) <- a.(i).(j) *. one_over_alpha
done ;
for j2 = j + 1 to n - 1 do
let c = col_inner_prod a j j2 in
r.(j).(j2) <- c ;
for i = 0 to m - 1 do
a.(i).(j2) <- a.(i).(j2) -. (c *. a.(i).(j))
done
done
done ;
(a, r)
let qr ?(in_place = false) a =
let a = if in_place then a else Array.map Array.copy a in
qr_in_place a
let mul_mv ?(trans = false) a x =
let rows = Array.length a in
if rows = 0 then [||]
else
let cols = Array.length a.(0) in
let m, n, get =
if trans then
let get i j = a.(j).(i) in
(cols, rows, get)
else
let get i j = a.(i).(j) in
(rows, cols, get)
in
if n <> Array.length x then failwith "Dimension mismatch" ;
let result = Array.make m 0. in
for i = 0 to m - 1 do
let v, _ =
Array.fold_left
(fun (acc, j) x -> (acc +. (get i j *. x), succ j))
(0., 0) x
in
result.(i) <- v
done ;
result
let is_nan v = match classify_float v with FP_nan -> true | _ -> false
let error_msg msg = Error (`Msg msg)
let triu_solve r b =
let m = Array.length b in
if m <> Array.length r then
error_msg
"triu_solve R b requires R to be square with same number of rows as b"
else if m = 0 then Ok [||]
else if m <> Array.length r.(0) then
error_msg "triu_solve R b requires R to be a square"
else
let sol = Array.copy b in
for i = m - 1 downto 0 do
sol.(i) <- sol.(i) /. r.(i).(i) ;
for j = 0 to i - 1 do
sol.(j) <- sol.(j) -. (r.(j).(i) *. sol.(i))
done
done ;
if Array.exists is_nan sol then
error_msg "triu_solve detected NaN result"
else Ok sol
let ols ?(in_place = false) a b =
let q, r = qr ~in_place a in
triu_solve r (mul_mv ~trans:true q b)
let make_lr_inputs responder predictors m =
Array.init (Array.length m) (fun i -> Array.map (fun a -> a m.(i)) predictors),
Array.init (Array.length m) (fun i -> responder m.(i))
let r_square m responder predictors r =
let predictors_matrix, responder_vector =
make_lr_inputs responder predictors m
in
let sum_responder = Array.fold_left ( +. ) 0. responder_vector in
let mean = sum_responder /. float (Array.length responder_vector) in
let tot_ss = ref 0. in
let res_ss = ref 0. in
let predicted i =
let x = ref 0. in
for j = 0 to Array.length r - 1 do
x := !x +. (predictors_matrix.(i).(j) *. r.(j))
done ;
!x
in
for i = 0 to Array.length responder_vector - 1 do
tot_ss := !tot_ss +. ((responder_vector.(i) -. mean) ** 2.) ;
res_ss := !res_ss +. ((responder_vector.(i) -. predicted i) ** 2.)
done ;
1. -. (!res_ss /. !tot_ss)
let ols responder predictors m =
let matrix, vector = make_lr_inputs responder predictors m in
match ols ~in_place:true matrix vector with
| Ok estimates ->
let r_square = r_square m responder predictors estimates in
Ok (estimates, r_square)
| Error _ as err -> err
|
