Source file conv.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
open Printf
open Bigarray
include Sexplib0.Sexp_conv
open Sexp

type bigstring = (char, int8_unsigned_elt, c_layout) Array1.t
type float32_vec = (float, float32_elt, fortran_layout) Array1.t
type float64_vec = (float, float64_elt, fortran_layout) Array1.t
type vec = float64_vec
type float32_mat = (float, float32_elt, fortran_layout) Array2.t
type float64_mat = (float, float64_elt, fortran_layout) Array2.t
type mat = float64_mat

let sexp_of_float_vec vec =
  let lst_ref = ref [] in
  for i = Array1.dim vec downto 1 do
    lst_ref := sexp_of_float vec.{i} :: !lst_ref
  done;
  List !lst_ref
;;

let sexp_of_bigstring (bstr : bigstring) =
  let n = Array1.dim bstr in
  let str = Bytes.create n in
  for i = 0 to n - 1 do
    Bytes.set str i bstr.{i}
  done;
  Atom (Bytes.unsafe_to_string str)
;;

let sexp_of_float32_vec (vec : float32_vec) = sexp_of_float_vec vec
let sexp_of_float64_vec (vec : float64_vec) = sexp_of_float_vec vec
let sexp_of_vec (vec : vec) = sexp_of_float_vec vec

let sexp_of_float_mat mat =
  let m = Array2.dim1 mat in
  let n = Array2.dim2 mat in
  let lst_ref = ref [] in
  (* It's surprising that we serialize [Fortran_layout] matrices in row-major order. I can
     only speculate that it was chosen for readability. The cache performance is
     irrelevant because people who care won't serialize to sexp. *)
  for row = n downto 1 do
    for col = m downto 1 do
      lst_ref := sexp_of_float mat.{col, row} :: !lst_ref
    done
  done;
  List (sexp_of_int m :: sexp_of_int n :: !lst_ref)
;;

let sexp_of_float32_mat (mat : float32_mat) = sexp_of_float_mat mat
let sexp_of_float64_mat (mat : float64_mat) = sexp_of_float_mat mat
let sexp_of_mat (mat : mat) = sexp_of_float_mat mat
let bigstring_sexp_grammar : bigstring Sexplib0.Sexp_grammar.t = { untyped = String }

let bigstring_of_sexp sexp =
  match sexp with
  | Atom str ->
    let len = String.length str in
    let bstr = Array1.create char c_layout len in
    for i = 0 to len - 1 do
      bstr.{i} <- str.[i]
    done;
    bstr
  | List _ -> of_sexp_error "bigstring_of_sexp: atom needed" sexp
;;

let float_vec_of_sexp empty_float_vec create_float_vec sexp =
  match sexp with
  | List [] -> empty_float_vec
  | List lst ->
    let len = List.length lst in
    let res = create_float_vec len in
    let rec loop i = function
      | [] -> res
      | h :: t ->
        res.{i} <- float_of_sexp h;
        loop (i + 1) t
    in
    loop 1 lst
  | Atom _ -> of_sexp_error "float_vec_of_sexp: list needed" sexp
;;

let create_float32_vec = Array1.create float32 fortran_layout
let create_float64_vec = Array1.create float64 fortran_layout
let empty_float32_vec = create_float32_vec 0
let empty_float64_vec = create_float64_vec 0
let float32_vec_of_sexp = float_vec_of_sexp empty_float32_vec create_float32_vec
let float64_vec_of_sexp = float_vec_of_sexp empty_float64_vec create_float64_vec
let vec_of_sexp = float_vec_of_sexp empty_float64_vec create_float64_vec
let vec_sexp_grammar : _ Sexplib0.Sexp_grammar.t = { untyped = List (Many Float) }
let float32_vec_sexp_grammar = vec_sexp_grammar
let float64_vec_sexp_grammar = vec_sexp_grammar

let check_too_much_data sexp data res =
  if data = [] then res else of_sexp_error "float_mat_of_sexp: too much data" sexp
;;

let float_mat_of_sexp create_float_mat sexp =
  match sexp with
  | List (sm :: sn :: data) ->
    let m = int_of_sexp sm in
    let n = int_of_sexp sn in
    let res = create_float_mat m n in
    if m = 0 || n = 0
    then check_too_much_data sexp data res
    else (
      let rec loop_cols col data =
        let vec = Array2.slice_right res col in
        let rec loop_rows row = function
          | [] -> of_sexp_error "float_mat_of_sexp: not enough data" sexp
          | h :: t ->
            vec.{row} <- float_of_sexp h;
            if row = m
            then if col = n then check_too_much_data sexp t res else loop_cols (col + 1) t
            else loop_rows (row + 1) t
        in
        loop_rows 1 data
      in
      loop_cols 1 data)
  | List _ -> of_sexp_error "float_mat_of_sexp: list too short" sexp
  | Atom _ -> of_sexp_error "float_mat_of_sexp: list needed" sexp
;;

let create_float32_mat = Array2.create float32 fortran_layout
let create_float64_mat = Array2.create float64 fortran_layout
let float32_mat_of_sexp = float_mat_of_sexp create_float32_mat
let float64_mat_of_sexp = float_mat_of_sexp create_float64_mat
let mat_of_sexp = float_mat_of_sexp create_float64_mat

let mat_sexp_grammar : _ Sexplib0.Sexp_grammar.t =
  { untyped = List (Cons (Integer, Cons (Integer, Many Float))) }
;;

let float32_mat_sexp_grammar = mat_sexp_grammar
let float64_mat_sexp_grammar = mat_sexp_grammar
let string_of__of__sexp_of to_sexp x = Sexp.to_string (to_sexp x)

let of_string__of__of_sexp of_sexp s =
  try
    let sexp = Sexp.of_string s in
    of_sexp sexp
  with
  | e ->
    failwith
      (sprintf "of_string failed on %s with %s" s (Sexp.to_string_hum (sexp_of_exn e)))
;;