package owl-base

  1. Overview
  2. Docs
package owl-base
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file owl_base_slicing.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
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
# 1 "src/base/core/owl_base_slicing.ml"
(*
 * OWL - OCaml Scientific and Engineering Computing
 * Copyright (c) 2016-2019 Liang Wang <liang.wang@cl.cam.ac.uk>
 *)

open Bigarray

open Owl_types


(* convert from a list of slice definition to array for internal use *)
let sdlist_to_sdarray axis =
  List.map (function
    | I i -> I_ i
    | L i -> L_ (Array.of_list i)
    | R i -> R_ (Array.of_list i)
  ) axis
  |> Array.of_list


(* return true if slicing (all R_) or false if fancy indexing (has L_) *)
let is_basic_slicing = Array.for_all (function R_ _ -> true | _ -> false)


(* check the validity of the slice definition, also re-format slice definition,
   axis: slice definition;
   shp: shape of the original ndarray;
 *)
let check_slice_definition axis shp =
  let axis_len = Array.length axis in
  let shp_len = Array.length shp in
  assert (axis_len <= shp_len);
  (* add missing definition on higher dimensions *)
  let axis =
    if axis_len < shp_len then
      let suffix = Array.make (shp_len - axis_len) (R_ [||]) in
      Array.append axis suffix
    else axis
  in
  (* re-format slice definition, note I_ will be replaced with L_ *)
  Array.map2 (fun i n ->
    match i with
    | I_ x -> (
        let x = if x >= 0 then x else n + x in
        assert (x < n);
        R_ [|x;x;1|]
      )
    | L_ x -> (
        let is_cont = ref true in
        if Array.length x <> n then is_cont := false;
        let x = Array.mapi (fun i j ->
          let j = if j >= 0 then j else n + j in
          assert (j < n);
          if i <> j then is_cont := false;
          j
        ) x
        in
        if !is_cont = true then R_ [|0;n-1;1|] else L_ x
      )
    | R_ x -> (
        match Array.length x with
        | 0 -> R_ [|0;n-1;1|]
        | 1 -> (
            let a = if x.(0) >= 0 then x.(0) else n + x.(0) in
            assert (a < n);
            R_ [|a;a;1|]
          )
        | 2 -> (
            let a = if x.(0) >= 0 then x.(0) else n + x.(0) in
            let b = if x.(1) >= 0 then x.(1) else n + x.(1) in
            let c = if a <= b then 1 else -1 in
            assert (not (a >= n || b >= n));
            R_ [|a;b;c|]
          )
        | 3 -> (
            let a = if x.(0) >= 0 then x.(0) else n + x.(0) in
            let b = if x.(1) >= 0 then x.(1) else n + x.(1) in
            let c = x.(2) in
            assert (not (a >= n || b >= n || c = 0));
            assert (not ((a < b && c < 0) || (a > b && c > 0)));
            R_ [|a;b;c|]
          )
        | _ -> failwith "check_slice_definition: error"
      )
  ) axis shp


(* calculate the minimum continuous block size and its corresponding dimension
   axis: slice definition;
   shp: shape of the original ndarray;
 *)
let calc_continuous_blksz axis shp =
  let slice_sz = Owl_utils.calc_slice shp in
  let ssz = ref 1 in
  let d = ref 0 in
  let _ = try
    for l = Array.length shp - 1 downto -1 do
      (* note: d is actually the corresponding dimension of continuous block
         plus one; also note the loop is down to -1 so the lowest dimension is
         also considered, in which case the whole array is copied. *)
      d := l + 1;
      if l < 0 then failwith "stop";
      match axis.(l) with
      | I_ _  -> failwith "stop" (* never reached *)
      | L_ _  -> failwith "stop"
      | R_ x  -> (
          if x.(0) = 0 && x.(1) = shp.(l) - 1 && x.(2) = 1 then
            ssz := slice_sz.(l)
          else failwith "stop"
        )
    done
  with _exn -> ()
  in !d, !ssz


(* calculat the shape according to the slice definition
   axis: slice definition
 *)
let calc_slice_shape axis =
  Array.map (function
    | I_ _x -> 1 (* never reached *)
    | L_ x -> Array.length x
    | R_ x ->
        let a, b, c = x.(0), x.(1), x.(2) in
        Stdlib.(abs ((b - a) / c)) + 1
  ) axis


(* recursively copy the continuous block, stop at its corresponding dimension d
   a: slice definition
   d: the corresponding dimension of continuous block + 1
   j: current dimension index
   i: current index of the data for copying
   f: copy function of the continuous block
 *)
let rec __foreach_continuous_blk a d j i f =
  if j = d then f i
  else (
    match a.(j) with
    | I_ _ -> ( (* never reache here *) )
    | L_ x -> (
        Array.iter (fun k ->
          i.(j) <- k;
          __foreach_continuous_blk a d (j + 1) i f
        ) x
      )
    | R_ x -> (
        let k = ref x.(0) in
        if x.(2) > 0 then (
          while !k <= x.(1)  do
            i.(j) <- !k;
            k := !k + x.(2);
            __foreach_continuous_blk a d (j + 1) i f
          done
        )
        else (
          while !k >= x.(1)  do
            i.(j) <- !k;
            k := !k + x.(2);
            __foreach_continuous_blk a d (j + 1) i f
          done
        )
      )
  )


(* a : slice definition, same rank as original ndarray
   d : the corresponding dimension of the continuous block +1
   f : the copy function for the continuous block
 *)
let _foreach_continuous_blk a d f =
  let i = Array.(make (length a) 0) in
  __foreach_continuous_blk a d 0 i f


(* reshape inputs in order to optimise the slicing performance *)
let optimise_input_shape axis x y =
  let n = Genarray.num_dims x in
  let sx = Genarray.dims x in
  let sy = Genarray.dims y in
  let dim = ref (n - 1) in
  let acx = ref 1 in
  let acy = ref 1 in
  (try
    for i = !dim downto 0 do
      match axis.(i) with
      | R_ a ->
          if a.(0) = 0 && a.(1) = sx.(i) - 1 && a.(2) = 1 then (
            acx := !acx * sx.(i);
            acy := !acy * sy.(i);
            dim := i;
          )
          else failwith "stop"
      | _    -> failwith "stop"
    done
  with _exn -> ());
  if n - !dim > 1 then (
    (* can be optimised *)
    let axis' = Array.sub axis 0 (!dim + 1) in
    let sx' = Array.sub sx 0 (!dim + 1) in
    let sy' = Array.sub sy 0 (!dim + 1) in
    sx'.(!dim) <- !acx;
    sy'.(!dim) <- !acy;
    let x' = reshape x sx' in
    let y' = reshape y sy' in
    axis', x', y'
  )
  else
    (* cannot be optimised *)
    axis, x, y


(* ends here *)
OCaml

Innovation. Community. Security.

GitHub Discord Twitter Peertube RSS
About
Changelog Releases Industrial Users Academic Users Why OCaml
Ecosystem
Packages Community Events OCaml Planet Jobs
Resources
Install OCaml Get Started Platform Tools Language Manual Standard Library API Books Exercises Papers OCaml Playground Logo
Policies
Carbon Footprint Governance Privacy Code of Conduct