Source file `owl_sparse_dok_matrix.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
# 1 "src/owl/sparse/owl_sparse_dok_matrix.ml"
(*
 * OWL - OCaml Scientific and Engineering Computing
 * Copyright (c) 2016-2020 Liang Wang <liang.wang@cl.cam.ac.uk>
 *)

type ('a, 'b) kind = ('a, 'b) Bigarray.kind

type ('a, 'b) t =
  { mutable m : int
  ; (* number of rows *)
    mutable n : int
  ; (* number of columns *)
    mutable k : ('a, 'b) kind
  ; (* type of sparse matrices *)
    mutable d : (int * int, 'a) Hashtbl.t (* hashtbl for storing data *)
  }

let zeros ?(density = 0.30) k m n =
  { m
  ; n
  ; k
  ; d =
      (let c = int_of_float (float_of_int (m * n) *. density) in
       Hashtbl.create c)
  }


let row_num x = x.m

let col_num x = x.n

let shape x = x.m, x.n

let numel x = x.m * x.n

let prune x r = Hashtbl.filter_map_inplace (fun _ v -> if v = r then None else Some v) x.d

let nnz x =
  let _ = prune x (Owl_const.zero x.k) in
  Hashtbl.((stats x.d).num_bindings)


let density x = float_of_int (nnz x) /. float_of_int (numel x)

let kind x = x.k

let _check_boundary i j m n =
  if i < 0 || i >= m || j < 0 || j >= n then failwith "error: index beyond the boundary"


let set x i j a =
  _check_boundary i j x.m x.n;
  let _a0 = Owl_const.zero x.k in
  match Hashtbl.mem x.d (i, j) with
  | true  -> if a <> _a0 then Hashtbl.replace x.d (i, j) a else Hashtbl.remove x.d (i, j)
  | false -> if a <> _a0 then Hashtbl.add x.d (i, j) a


let get x i j =
  _check_boundary i j x.m x.n;
  match Hashtbl.mem x.d (i, j) with
  | true  -> Hashtbl.find x.d (i, j)
  | false -> Owl_const.zero x.k


let reset x = Hashtbl.reset x.d

let copy x = { m = x.m; n = x.n; k = x.k; d = Hashtbl.copy x.d }

let iteri_nz f x = Hashtbl.iter (fun (i, j) v -> f i j v) x.d

let save x f = Owl_io.marshal_to_file x f

let load _k f = Owl_io.marshal_from_file f