QCaml/common/lib/smallarray.ml

104 lines
2.2 KiB
OCaml

type t = { dim1 : int;
dim2 : int;
dim3 : int;
data : float array;
}
let dim1 t = t.dim1
let dim2 t = t.dim2
let dim3 t = t.dim3
let data t = t.data
let at t i j k =
k + t.dim3*(j + t.dim2*i)
let from t l =
let i = l / (t.dim3*t.dim2) in
let l' = l - i*t.dim3*t.dim2 in
let j = l' / t.dim3 in
let k = l' - j * t.dim3
in (i,j,k)
let init dim1 dim2 dim3 f =
let t = { dim1 ; dim2 ; dim3 ; data = Array.create_float (dim1 * dim2 * dim3) } in
let l = ref 0 in
for i=0 to dim1-1 do
for j=0 to dim2-1 do
for k=0 to dim3-1 do
t.data.(!l) <- f i j k;
incr l
done
done
done;
t
let make dim1 dim2 dim3 v =
{ dim1 ; dim2 ; dim3 ;
data = Array.make (dim1 * dim2 * dim3) v
}
let get t i j k =
Array.get t.data (at t i j k)
let unsafe_get t i j k =
Array.unsafe_get t.data (at t i j k)
let set t i j k v =
Array.set t.data (at t i j k) v
let unsafe_set t i j k v =
Array.unsafe_set t.data (at t i j k) v
let iter3 i j f t =
assert (0 <= i && i < t.dim1);
assert (0 <= j && j < t.dim2);
let start = t.dim3*(j+i*t.dim2) in
for k=0 to t.dim3-1 do
f (Array.unsafe_get t.data (k+start))
done
let iteri3 i j f t =
assert (0 <= i && i < t.dim1);
assert (0 <= j && j < t.dim2);
let start = t.dim3*(j+i*t.dim2) in
for k=0 to t.dim3-1 do
f k (Array.unsafe_get t.data (k+start))
done
let to_array t =
Array.init t.dim1 (fun i ->
Array.init t.dim2 (fun j ->
Array.init t.dim3 (fun k ->
unsafe_get t i j k)))
let of_array a =
init (Array.length a) (Array.length a.(0)) (Array.length a.(0).(0))
(fun i j k -> a.(i).(j).(k) )
let sub1 i t =
let n = t.dim2 * t.dim3 in
let shift = i * n in
{ dim1 = 1 ; dim2 = t.dim2 ; dim3 = t.dim3 ;
data = Array.init n (fun k -> Array.unsafe_get t.data (shift+k)) }
let sum t =
Array.fold_left (+.) 0. t.data
let sum_sub1 k t =
let n = t.dim2 * t.dim3 in
let shift = k * n in
let accu = ref 0. in
for i=shift to shift+n-1 do
accu := !accu +. Array.unsafe_get t.data i
done;
!accu
let scale f t =
{ dim1 = t.dim1 ; dim2 = t.dim2 ; dim3 = t.dim3 ;
data = Array.map (fun x -> f *. x) t.data }
let ( .@() ) t (i,j,k) = get t i j k
let ( .@() <- ) t (i,j,k) v = set t i j k v