package core_bench

  1. Overview
  2. Docs
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file linear_algebra.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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
open Core

(** Vectors *)
module Vec = struct
  type t = float array [@@deriving sexp]

  let copy = Array.copy
  let create0 len = Array.create ~len 0.

  let sumsq t =
    let accum = ref 0. in
    for i = 0 to Array.length t - 1 do
      let t_i = t.(i) in
      accum := !accum +. (t_i *. t_i)
    done;
    !accum
  ;;

  let norm t = sqrt (sumsq t)

  let almost_equal ~tol t1 t2 =
    Array.length t1 = Array.length t2
    && Array.for_all2_exn t1 t2 ~f:(fun x y -> Float.( <= ) (Float.abs (x -. y)) tol)
  ;;
end

(** Matrices *)
module Mat = struct
  type t = float array array [@@deriving sexp]

  let copy t = Array.map t ~f:Array.copy

  let create0 ~rows ~cols =
    let make_row _ = Array.create ~len:cols 0. in
    Array.init rows ~f:make_row
  ;;

  let create_per_row ~rows ~cols ~f =
    let make_row _ = Array.init cols ~f in
    Array.init rows ~f:make_row
  ;;

  (** The norm of a column of a matrix. *)
  let col_norm t column =
    let accum = ref 0. in
    for i = 0 to Array.length t - 1 do
      let entry = t.(i).(column) in
      accum := !accum +. (entry *. entry)
    done;
    sqrt !accum
  ;;

  (** The inner product of columns j1 and j2 *)
  let col_inner_prod t j1 j2 =
    let accum = ref 0. in
    for i = 0 to Array.length t - 1 do
      accum := !accum +. (t.(i).(j1) *. t.(i).(j2))
    done;
    !accum
  ;;

  let get_column t j =
    let len = Array.length t in
    Array.init len ~f:(fun i -> t.(i).(j))
  ;;

  let almost_equal ~tol t1 t2 =
    Array.length t1 = Array.length t2
    && Array.for_all2_exn t1 t2 ~f:(Vec.almost_equal ~tol)
  ;;
end

let qr_in_place a =
  (* Our implementation will just do Gram-Schmidt. *)
  let m = Array.length a in
  if m = 0
  then [||], [||] (* empty QR decomposition *)
  else (
    let n = Array.length a.(0) in
    let r = Mat.create0 ~rows:n ~cols:n in
    for j = 0 to n - 1 do
      (* handle column j *)
      let alpha = Mat.col_norm a j in
      r.(j).(j) <- alpha;
      let one_over_alpha = 1. /. alpha in
      (* Rescale this column so that it's a unit vector. *)
      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 = Mat.col_inner_prod a j j2 in
        r.(j).(j2) <- c;
        (* Now, subtract c * column j from column j2. *)
        for i = 0 to m - 1 do
          a.(i).(j2) <- a.(i).(j2) -. (c *. a.(i).(j))
        done
      done
    done;
    a, r)
;;

(** [qr A] returns the QR-decomposition of [A] as a pair (Q,R). [A] must have
    at least as many rows as columns and have full rank.

    If [in_place] (default: [false]) is [true], then [A] is overwritten with [Q].
*)
let qr ?(in_place = false) a =
  let a = if in_place then a else Mat.copy a in
  qr_in_place a
;;

(** [triu_solve R b] solves R x = b where [R] is an m x m upper-triangular matrix
    and [b] is an m x 1 column vector.  *)
let triu_solve r b =
  let m = Array.length b in
  if m <> Array.length r
  then
    Or_error.error_string
      "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 Or_error.error_string "triu_solve R b requires R to be square"
  else (
    let sol = Vec.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 sol ~f:Float.is_nan
    then Or_error.error_string "triu_solve detected NaN result"
    else Ok sol)
;;

(** [mul A B] computes the matrix product [A * B].  If [transa] (default: [false])
    is [true], then we compute A' * B where A' denotes the transpose of A. *)

(* val mul : ?transa:bool -> m -> m -> m *)

(** [mul_mv A x] computes the product [A * x] (where [A] is a matrix and [x] is
    a column vector).

    [transa] is as with [mul].
*)
let mul_mv ?(transa = false) a x =
  (* we let c denote either a or a', depending on whether transa is set. *)
  let rows = Array.length a in
  if rows = 0
  then [||]
  else (
    let cols = Array.length a.(0) in
    (* (m, n, c_get) will be (rows of c, columns of c, accessor for c). *)
    let m, n, c_get =
      if transa
      then (
        let c_get i j = a.(j).(i) in
        cols, rows, c_get)
      else (
        let c_get i j = a.(i).(j) in
        rows, cols, c_get)
    in
    if n <> Array.length x then failwith "Dimension mismatch";
    let result = Vec.create0 m in
    for i = 0 to m - 1 do
      result.(i)
      <- Array.foldi x ~init:0. ~f:(fun j accum x_j -> accum +. (c_get i j *. x_j))
    done;
    result)
;;

(** [ols A b] computes the ordinary least-squares solution to A x = b.
    [A] must have at least as many rows as columns and have full rank.

    This can be used to compute solutions to non-singular square systems,
    but is somewhat sub-optimal for that purpose.

    The algorithm is to factor A = Q * R and solve R x = Q' b where Q' denotes
    the transpose of Q.
*)
let ols ?(in_place = false) a b =
  let q, r = qr ~in_place a in
  triu_solve r (mul_mv ~transa:true q b)
;;
OCaml

Innovation. Community. Security.