mirror of https://gitlab.com/scemama/QCaml.git
331 lines
9.8 KiB
OCaml
331 lines
9.8 KiB
OCaml
open Lacaml.D
|
|
|
|
let epsilon = Constants.epsilon
|
|
|
|
type index_value =
|
|
{
|
|
index: int;
|
|
value: float
|
|
}
|
|
|
|
type sparse_vector =
|
|
{
|
|
n: int;
|
|
v: index_value array
|
|
}
|
|
|
|
type t =
|
|
| Dense of Vec.t
|
|
| Sparse of sparse_vector
|
|
|
|
|
|
let is_sparse = function
|
|
| Sparse _ -> true
|
|
| Dense _ -> false
|
|
|
|
|
|
let is_dense = function
|
|
| Sparse _ -> false
|
|
| Dense _ -> true
|
|
|
|
|
|
exception Found of float
|
|
|
|
let get = function
|
|
| Dense v -> (fun i -> v.{i})
|
|
| Sparse { n ; v } -> (fun i ->
|
|
if i < 1 || i > n then invalid_arg "index out of bounds";
|
|
try
|
|
Array.iter (fun {index ; value} ->
|
|
if index=i then
|
|
raise (Found value)) v;
|
|
raise Not_found
|
|
with
|
|
| Not_found -> 0.
|
|
| Found x -> x
|
|
)
|
|
|
|
|
|
|
|
let dim = function
|
|
| Dense v -> Vec.dim v
|
|
| Sparse {n ; v} -> n
|
|
|
|
|
|
let sparse_of_dense ?(threshold=epsilon) = function
|
|
| Sparse _ -> invalid_arg "Expected a dense vector"
|
|
| Dense v ->
|
|
let rec aux accu = function
|
|
| 0 -> accu |> Array.of_list
|
|
| i ->
|
|
let x = v.{i} in
|
|
if abs_float x < threshold then
|
|
(aux [@tailcall]) accu (i-1)
|
|
else
|
|
(aux [@tailcall]) ({index=i ; value=x}::accu) (i-1)
|
|
in
|
|
let n = Vec.dim v in
|
|
Sparse { n ; v=aux [] n }
|
|
|
|
|
|
let rec to_assoc_list ?(threshold=epsilon) = function
|
|
| Sparse {n ; v} -> Array.map (fun {index ; value} -> (index, value)) v |> Array.to_list
|
|
| Dense v -> to_assoc_list @@ sparse_of_dense ~threshold (Dense v)
|
|
|
|
|
|
let dense_of_sparse = function
|
|
| Dense _ -> invalid_arg "Expected a sparse vector"
|
|
| Sparse {n ; v} ->
|
|
let v' = Vec.make0 n in
|
|
Array.iter (fun {index ; value} -> v'.{index} <- value) v;
|
|
Dense v'
|
|
|
|
|
|
let dense_of_vec v = Dense v
|
|
|
|
|
|
let sparse_of_vec ?(threshold=epsilon) v =
|
|
dense_of_vec v
|
|
|> sparse_of_dense ~threshold
|
|
|
|
|
|
let sparse_of_assoc_list n v =
|
|
Sparse { n ;
|
|
v = List.map (fun (index, value) -> {index ; value}) v
|
|
|> Array.of_list
|
|
}
|
|
|
|
|
|
let rec to_vec = function
|
|
| Dense v -> v
|
|
| Sparse v -> dense_of_sparse (Sparse v) |> to_vec
|
|
|
|
|
|
|
|
let scale ?(threshold=epsilon) x = function
|
|
| Dense v -> let v' = copy v in (scal x v'; Dense v')
|
|
| Sparse {n ; v} ->
|
|
Sparse {n ; v = Array.map (fun {index ; value} -> { index ; value=x *. value} ) v }
|
|
|
|
|
|
let rec neg = function
|
|
| Dense v -> Dense (Vec.neg v)
|
|
| Sparse {n ; v} ->
|
|
Sparse {n ; v = Array.map (fun {index ; value} -> {index ; value = -. value}) v}
|
|
|
|
|
|
|
|
let axpy ?(threshold=epsilon) ?(alpha=1.) x y =
|
|
if dim x <> dim y then
|
|
invalid_arg "Inconsistent dimensions";
|
|
|
|
match x, y with
|
|
| Dense x , Dense y -> Dense (let y = copy y in axpy ~alpha x y ; y)
|
|
| Sparse {n ; v}, Dense y ->
|
|
begin
|
|
let v' = copy y in
|
|
Array.iter (fun {index ; value} -> v'.{index} <- v'.{index} +. alpha *. value) v;
|
|
sparse_of_vec ~threshold v'
|
|
end
|
|
| Dense x , Sparse {n ; v} ->
|
|
begin
|
|
let v' = copy x in
|
|
scal alpha v';
|
|
Array.iter (fun {index ; value} -> v'.{index} <- v'.{index} +. value) v;
|
|
sparse_of_vec ~threshold v'
|
|
end
|
|
| Sparse {n ; v}, Sparse {n=n' ; v=v'} ->
|
|
begin
|
|
let rec aux accu k l =
|
|
match k < Array.length v, l < Array.length v' with
|
|
| true, true ->
|
|
begin
|
|
let {index=i ; value=x} = v.(k)
|
|
and {index=j ; value=y} = v'.(l)
|
|
in
|
|
match compare i j with
|
|
| -1 ->
|
|
let z = alpha *. x in
|
|
let new_accu =
|
|
if abs_float z > threshold then
|
|
{index=i ; value=z} :: accu
|
|
else
|
|
accu
|
|
in (aux [@tailcall]) new_accu (k+1) l
|
|
| 1 ->
|
|
let new_accu =
|
|
if abs_float y > threshold then
|
|
{index=j ; value=y} :: accu
|
|
else
|
|
accu
|
|
in (aux [@tailcall]) new_accu k (l+1)
|
|
| 0 ->
|
|
let z = alpha *. x +. y in
|
|
let new_accu =
|
|
if abs_float z > threshold then
|
|
{index=i ; value=z} :: accu
|
|
else
|
|
accu
|
|
in (aux [@tailcall]) new_accu (k+1) (l+1)
|
|
| _ -> assert false
|
|
end
|
|
|
|
| true, false ->
|
|
let {index=i ; value=x} = v.(k) in
|
|
(aux [@tailcall]) ({index=i ; value=alpha *. x}::accu) (k+1) l
|
|
|
|
| false, true ->
|
|
(aux [@tailcall]) (v'.(l)::accu) k (l+1)
|
|
|
|
| false, false -> {n ; v=List.rev accu |> Array.of_list}
|
|
in
|
|
Sparse (aux [] 0 0)
|
|
end
|
|
|
|
let add = axpy ~alpha:1.
|
|
|
|
let sub ?(threshold=epsilon) x y = add ~threshold x @@ neg y
|
|
|
|
let pp ppf = function
|
|
| Dense m -> Util.pp_float_array ppf @@ Vec.to_array m
|
|
| Sparse {n ; v} ->
|
|
begin
|
|
Format.fprintf ppf "@[[ %d | " n;
|
|
Array.iter (fun {index ; value} -> Format.fprintf ppf "@[(%d, %f); @]" index value) v;
|
|
Format.fprintf ppf "]@]"
|
|
end
|
|
|
|
|
|
|
|
let dot v v' =
|
|
|
|
let d_d v v' =
|
|
dot v v'
|
|
in
|
|
|
|
let d_sp v' {n ; v} =
|
|
if n <> Vec.dim v' then
|
|
invalid_arg "Inconsistent dimensions";
|
|
Array.fold_left (fun accu {index ; value} -> accu +. value *. v'.{index}) 0. v
|
|
in
|
|
|
|
let sp_sp {n ; v} {n=n' ; v=v'} =
|
|
if n <> n' then
|
|
invalid_arg "Inconsistent dimensions";
|
|
|
|
let rec aux accu k l =
|
|
match Array.length v > k, Array.length v' > l with
|
|
| true, true ->
|
|
let {index=i ; value=x} = v.(k)
|
|
and {index=j ; value=y} = v'.(l)
|
|
in
|
|
begin
|
|
match compare i j with
|
|
| -1 -> (aux [@tailcall]) accu (k+1) l
|
|
| 1 -> (aux [@tailcall]) accu k (l+1)
|
|
| 0 -> (aux [@tailcall]) (accu +. x *. y) (k+1) (l+1)
|
|
| _ -> assert false
|
|
end
|
|
| _ -> accu
|
|
in
|
|
aux 0. 0 0
|
|
in
|
|
|
|
match v, v' with
|
|
| (Dense v), (Dense v') -> d_d v v'
|
|
| (Sparse v), (Sparse v') -> sp_sp v v'
|
|
| (Dense v), (Sparse v') -> d_sp v v'
|
|
| (Sparse v), (Dense v') -> d_sp v' v
|
|
|
|
|
|
let norm v = sqrt @@ dot v v
|
|
|
|
let test_case () =
|
|
|
|
let x1 = Vec.map (fun x -> if abs_float x < 0.6 then 0. else x) (Vec.random 100)
|
|
and x2 = Vec.map (fun x -> if abs_float x < 0.3 then 0. else x) (Vec.random 100)
|
|
in
|
|
let x3 = Vec.map (fun x -> 2. *. x) x1
|
|
and x4 = Vec.add x1 x2
|
|
and x5 = Vec.sub x1 x2
|
|
and x6 =
|
|
let v = copy x2 in
|
|
Lacaml.D.axpy ~alpha:3. x1 v;
|
|
v
|
|
in
|
|
|
|
let v1 = dense_of_vec x1
|
|
and v2 = dense_of_vec x2
|
|
and v3 = dense_of_vec x3
|
|
and v4 = dense_of_vec x4
|
|
and v5 = dense_of_vec x5
|
|
and v6 = dense_of_vec x6
|
|
in
|
|
|
|
let v1_s = sparse_of_vec x1
|
|
and v2_s = sparse_of_vec x2
|
|
and v3_s = sparse_of_vec x3
|
|
and v4_s = sparse_of_vec x4
|
|
and v5_s = sparse_of_vec x5
|
|
and v6_s = sparse_of_vec x6
|
|
in
|
|
|
|
let zero = dense_of_vec (Vec.make0 100)
|
|
and zero_s = sparse_of_vec (Vec.make0 100)
|
|
in
|
|
|
|
let test_conversion () =
|
|
Alcotest.(check bool) "sparse -> dense 1" true (dense_of_sparse v1_s = v1 );
|
|
Alcotest.(check bool) "sparse -> dense 2" true (dense_of_sparse v2_s = v2 );
|
|
Alcotest.(check bool) "dense -> sparse 1" true (sparse_of_dense v1 = v1_s);
|
|
Alcotest.(check bool) "dense -> sparse 2" true (sparse_of_dense v2 = v2_s);
|
|
in
|
|
let test_operations () =
|
|
Alcotest.(check bool) "dense scale" true (scale 2. v1 = v3);
|
|
Alcotest.(check bool) "sparse scale" true (scale 2. v1_s = v3_s);
|
|
|
|
Alcotest.(check bool) "dense dense add" true (add v1 v2 = v4);
|
|
Alcotest.(check bool) "dense sparse add" true (add v1 v2_s = v4_s);
|
|
Alcotest.(check bool) "sparse dense add" true (add v1_s v2 = v4_s);
|
|
Alcotest.(check bool) "sparse dense add" true (add v1 v2_s = v4_s);
|
|
Alcotest.(check bool) "sparse sparse add" true (add v1_s v2_s = v4_s);
|
|
|
|
Alcotest.(check bool) "dense dense sub" true (sub v1 v2 = v5);
|
|
Alcotest.(check bool) "dense sparse sub" true (sub v1 v2_s = v5_s);
|
|
Alcotest.(check bool) "sparse dense sub" true (sub v1_s v2 = v5_s);
|
|
Alcotest.(check bool) "sparse dense sub" true (sub v1 v2_s = v5_s);
|
|
Alcotest.(check bool) "sparse sparse sub" true (sub v1_s v2_s = v5_s);
|
|
|
|
Alcotest.(check bool) "dense dense sub" true (sub v1 v1 = zero);
|
|
Alcotest.(check bool) "dense sparse sub" true (sub v1 v1_s = zero_s);
|
|
Alcotest.(check bool) "sparse dense sub" true (sub v1_s v1 = zero_s);
|
|
Alcotest.(check bool) "sparse sparse sub" true (sub v1_s v1_s = zero_s);
|
|
|
|
Alcotest.(check bool) "dense dense axpy" true (axpy ~alpha:3. v1 v2 = v6);
|
|
Alcotest.(check bool) "dense sparse axpy" true (sub ~threshold:1.e-12 (axpy ~alpha:3. v1 v2_s) v6_s = zero_s);
|
|
Alcotest.(check bool) "sparse dense axpy" true (sub ~threshold:1.e-12 (axpy ~alpha:3. v1_s v2) v6_s = zero_s);
|
|
|
|
Alcotest.(check bool) "sparse sparse axpy" true (sub ~threshold:1.e-12 (axpy ~alpha:3. v1_s v2_s) v6_s = zero_s);
|
|
in
|
|
let test_dot () =
|
|
let d1d2 = Lacaml.D.dot x1 x2
|
|
and d1d1 = Lacaml.D.dot x1 x1
|
|
and d2d2 = Lacaml.D.dot x2 x2
|
|
in
|
|
Alcotest.(check (float 1.e-10)) "sparse x dense 1" (dot v1_s v2 ) d1d2;
|
|
Alcotest.(check (float 1.e-10)) "sparse x dense 2" (dot v1_s v1 ) d1d1;
|
|
Alcotest.(check (float 1.e-10)) "sparse x dense 3" (dot v2_s v2 ) d2d2;
|
|
Alcotest.(check (float 1.e-10)) "dense x sparse 1" (dot v1 v2_s) d1d2;
|
|
Alcotest.(check (float 1.e-10)) "dense x sparse 2" (dot v1 v1_s) d1d1;
|
|
Alcotest.(check (float 1.e-10)) "dense x sparse 3" (dot v2 v2_s) d2d2;
|
|
Alcotest.(check (float 1.e-10)) "sparse x sparse 1" (dot v1_s v2_s) d1d2;
|
|
Alcotest.(check (float 1.e-10)) "sparse x sparse 2" (dot v1_s v1_s) d1d1;
|
|
Alcotest.(check (float 1.e-10)) "sparse x sparse 3" (dot v2_s v2_s) d2d2;
|
|
in
|
|
[
|
|
"Conversion", `Quick, test_conversion;
|
|
"Operations", `Quick, test_operations;
|
|
"Dot product", `Quick, test_dot;
|
|
]
|
|
|