QCaml/Utils/FourIdxStorage.ml

694 lines
18 KiB
OCaml

open Util
open Lacaml.D
open Constants
let max_index = 1 lsl 14
type index_pair = { first : int ; second : int }
type array2 = (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t
type storage_t =
| Dense of array2
| Svd of Vec.t * array2 * array2
| Sparse of (int, float) Hashtbl.t
type t =
{
size : int ;
two_index : array2;
two_index_anti : array2;
three_index : array2;
three_index_anti : array2;
four_index : storage_t ;
}
let key_of_indices ~r1 ~r2 =
let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
let f i k =
let p, r =
if i <= k then i, k else k, i
in p + (r*(r-1))/2
in
let p = f i k and q = f j l in
f p q
let check_bounds r1 r2 t =
let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
let size = t.size in
assert ( (i lor j lor k lor l) > 0 );
assert ( i <= size && j <= size && k <= size && l <= size )
let dense_index i j size =
(j-1)*size + i
let sym_index i j =
if i < j then
(j*(j-1))/2 + i
else
(i*(i-1))/2 + j
let unsafe_get_four_index ~r1 ~r2 t =
let open Bigarray.Array2 in
let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
if i=k then
if j=l then
unsafe_get t.two_index i j
else
unsafe_get t.three_index (dense_index j l t.size) i
else if j=l then
unsafe_get t.three_index (dense_index i k t.size) j
else if i=l then
if k=j then
unsafe_get t.two_index_anti i j
else
unsafe_get t.three_index_anti (dense_index j k t.size) i
else if j=k then
unsafe_get t.three_index_anti (dense_index i l t.size) j
else if i=j then
if k=l then
unsafe_get t.two_index_anti i k
else
unsafe_get t.three_index_anti (dense_index k l t.size) i
else if k=l then
(* <ij|kk> *)
unsafe_get t.three_index_anti (dense_index i j t.size) k
else
match t.four_index with
| Dense a -> unsafe_get a (dense_index i k t.size) (sym_index j l)
| Svd _ -> assert false
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
try Hashtbl.find a key
with Not_found -> 0.
let get_four_index ~r1 ~r2 t =
check_bounds r1 r2 t;
unsafe_get_four_index ~r1 ~r2 t
let unsafe_set_four_index ~r1 ~r2 ~value t =
let open Bigarray.Array2 in
let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
let () =
if i=k then
begin
if j=l then
begin
unsafe_set t.two_index i j value;
unsafe_set t.two_index j i value;
unsafe_set t.three_index (dense_index i i t.size) j value;
end;
unsafe_set t.three_index (dense_index j l t.size) i value;
unsafe_set t.three_index (dense_index l j t.size) i value;
end
else if j=l then
begin
unsafe_set t.three_index (dense_index i k t.size) j value;
unsafe_set t.three_index (dense_index k i t.size) j value;
end
else if i=l then
begin
if j=k then
begin
unsafe_set t.two_index_anti i j value;
unsafe_set t.two_index_anti j i value;
unsafe_set t.three_index_anti (dense_index i i t.size) j value;
end;
unsafe_set t.three_index_anti (dense_index j k t.size) i value;
unsafe_set t.three_index_anti (dense_index k j t.size) i value;
end
else if j=k then
begin
unsafe_set t.three_index_anti (dense_index i l t.size) j value;
unsafe_set t.three_index_anti (dense_index l i t.size) j value;
end
else if i=j then
begin
if k=l then
begin
unsafe_set t.two_index_anti i k value;
unsafe_set t.two_index_anti k i value;
unsafe_set t.three_index_anti (dense_index i i t.size) k value;
end;
unsafe_set t.three_index_anti (dense_index k l t.size) i value;
unsafe_set t.three_index_anti (dense_index l k t.size) i value;
end
else if k=l then
(* <ij|kk> *)
begin
unsafe_set t.three_index_anti (dense_index i j t.size) k value;
unsafe_set t.three_index_anti (dense_index j i t.size) k value;
end
in
match t.four_index with
| Dense a -> let ik = (dense_index i k t.size)
and jl = (dense_index j l t.size)
and ki = (dense_index k i t.size)
and lj = (dense_index l j t.size)
and ik_s = (sym_index i k)
and jl_s = (sym_index j l)
in
begin
unsafe_set a ik jl_s value;
unsafe_set a ki jl_s value;
unsafe_set a jl ik_s value;
unsafe_set a lj ik_s value;
end
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
Hashtbl.replace a key value
| Svd _ -> assert false
let set_four_index ~r1 ~r2 ~value t =
check_bounds r1 r2 t;
unsafe_set_four_index ~r1 ~r2 ~value t
let unsafe_increment_four_index ~r1 ~r2 ~value t =
let updated_value =
value +. unsafe_get_four_index ~r1 ~r2 t
in
unsafe_set_four_index ~r1 ~r2 ~value:updated_value t
let increment_four_index ~r1 ~r2 ~value t =
check_bounds r1 r2 t;
unsafe_increment_four_index ~r1 ~r2 ~value t
let get ~r1 ~r2 a =
get_four_index ~r1 ~r2 a
let set ~r1 ~r2 ~value =
match classify_float value with
| FP_normal -> set_four_index ~r1 ~r2 ~value
| FP_zero
| FP_subnormal -> fun _ -> ()
| FP_infinite
| FP_nan ->
let msg =
Printf.sprintf "FourIdxStorage.ml : set : r1 = (%d,%d) ; r2 = (%d,%d)"
r1.first r1.second r2.first r2.second
in
raise (Invalid_argument msg)
let increment ~r1 ~r2 =
increment_four_index ~r1 ~r2
let create ~size ?(temp_dir="/dev/shm") sparsity =
assert (size < max_index);
let two_index =
SharedMemory.create ~temp_dir Float64 [| size ; size |]
|> Bigarray.array2_of_genarray
in
let two_index_anti =
SharedMemory.create ~temp_dir Float64 [| size ; size |]
|> Bigarray.array2_of_genarray
in
let three_index =
SharedMemory.create ~temp_dir Float64 [| size * size ; size |]
|> Bigarray.array2_of_genarray
in
let three_index_anti =
SharedMemory.create ~temp_dir Float64 [| size * size ; size |]
|> Bigarray.array2_of_genarray
in
let four_index =
match sparsity with
| `Dense -> let result =
SharedMemory.create ~temp_dir Float64 [| size*size ; (size*(size+1))/2 |]
|> Bigarray.array2_of_genarray
in
Dense result
| `Sparse -> let result = Hashtbl.create (size*size+13) in
Sparse result
in
{ size ; two_index ; two_index_anti ; three_index ; three_index_anti ; four_index }
let size t = t.size
(** TODO : remove epsilons *)
let get_chem t i j k l = get ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } t
let get_phys t i j k l = get ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } t
let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t
let increment_chem t i j k l value = increment ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
let increment_phys t i j k l value = increment ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t
type element = (** Element for the stream *)
{
i_r1: int ;
j_r2: int ;
k_r1: int ;
l_r2: int ;
value: float
}
let get_phys_all_i d ~j ~k ~l =
Vec.init d.size (fun i -> get_phys d i j k l)
let get_chem_all_i d ~j ~k ~l =
Vec.init d.size (fun i -> get_chem d i j k l)
let get_phys_all_ij d ~k ~l =
Mat.init_cols d.size d.size (fun i j -> get_phys d i j k l)
let get_chem_all_ij d ~k ~l =
if k = l then
let result =
Mat.col d.three_index k
|> Bigarray.genarray_of_array1
in
Bigarray.reshape_2 result d.size d.size
else
match d.four_index with
| Dense a ->
let kl = sym_index k l in
let result =
Mat.col a kl
|> Bigarray.genarray_of_array1
in
Bigarray.reshape_2 result d.size d.size
| Sparse a ->
Mat.init_cols d.size d.size (fun i j -> get_chem d i j k l)
| Svd _ -> assert false
let to_stream d =
let i = ref 0
and j = ref 1
and k = ref 1
and l = ref 1
in
let f_dense _ =
incr i;
if !i > !k then begin
i := 1;
incr j;
if !j > !l then begin
j := 1;
incr k;
if !k > !l then begin
k := 1;
incr l;
end;
end;
end;
if !l <= d.size then
Some { i_r1 = !i ; j_r2 = !j ;
k_r1 = !k ; l_r2 = !l ;
value = get_phys d !i !j !k !l
}
else
None
in
Stream.from f_dense
(** Write all integrals to a file with the <ij|kl> convention *)
let to_file ?(cutoff=Constants.integrals_cutoff) ~filename data =
let oc = open_out filename in
to_stream data
|> Stream.iter (fun {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} ->
if (abs_float value > cutoff) then
Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_r1 j_r2 k_r1 l_r2 value);
close_out oc
let of_file ~size ~sparsity filename =
let result = create ~size sparsity in
let ic = Scanf.Scanning.open_in filename in
let rec read_line () =
let result =
try
Some (Scanf.bscanf ic " %d %d %d %d %f" (fun i j k l v ->
set_phys result i j k l v))
with End_of_file -> None
in
match result with
| Some () -> (read_line [@tailcall]) ()
| None -> ()
in
read_line ();
Scanf.Scanning.close_in ic;
result
let to_list data =
let s =
to_stream data
in
let rec append accu =
let d =
try Some (Stream.next s) with
| Stream.Failure -> None
in
match d with
| None -> List.rev accu
| Some d -> (append [@tailcall]) (d :: accu)
in
append []
let broadcast t =
let size =
Parallel.InterNode.broadcast (lazy t.size)
in
let bufsize = size * size * size in
let stream = to_stream t in
let rec iterate () =
let buffer =
Parallel.InterNode.broadcast (lazy (
if Stream.peek stream = None then None else
Some (Array.init bufsize (fun _ ->
try Some (Stream.next stream)
with _ -> None
))
) )
in
match buffer with
| None -> ()
| Some buffer ->
begin
if not Parallel.InterNode.master then
Array.iter (fun x ->
match x with
| Some {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} ->
set_phys t i_r1 j_r2 k_r1 l_r2 value
| None -> () ) buffer;
(iterate [@tailcall]) ()
end
in iterate ();
t
let four_index_transform_dense_sparse ds coef source =
let mo_num = Mat.dim2 coef in
let destination = create ~size:mo_num ds in
let ao_num = Mat.dim1 coef in
let mo_num_2 = mo_num * mo_num in
let ao_num_2 = ao_num * ao_num in
let ao_mo_num = ao_num * mo_num in
let range_mo = list_range 1 mo_num in
let range_ao = list_range 1 ao_num in
let u = Mat.create mo_num_2 mo_num
and o = Mat.create ao_num ao_num_2
and p = Mat.create ao_num_2 mo_num
and q = Mat.create ao_mo_num mo_num
in
if Parallel.master then Printf.eprintf "4-idx transformation \n%!";
let task delta =
Mat.fill u 0.;
List.iter (fun l ->
if abs_float coef.{l,delta} > epsilon then
begin
let jk = ref 1 in
List.iter (fun k ->
get_chem_all_ij source ~k ~l
|> lacpy ~b:o ~bc:!jk
|> ignore;
jk := !jk + ao_num;
) range_ao;
(* o_i_jk *)
let p =
gemm ~transa:`T ~c:p o coef
(* p_jk_alpha = \sum_i o_i_jk c_i_alpha *)
in
let p' =
Bigarray.reshape_2 (Bigarray.genarray_of_array2 p) ao_num ao_mo_num
(* p_j_kalpha *)
in
let q =
gemm ~transa:`T ~c:q p' coef
(* q_kalpha_beta = \sum_j p_j_kalpha c_j_beta *)
in
let q' =
Bigarray.reshape_2 (Bigarray.genarray_of_array2 q) ao_num mo_num_2
(* q_k_alphabeta = \sum_j p_j_kalpha c_j_beta *)
in
ignore @@
gemm ~transa:`T ~beta:1. ~alpha:coef.{l,delta} ~c:u q' coef ;
(* u_alphabeta_gamma = \sum_k q_k_alphabeta c_k_gamma *)
end
) range_ao;
let u =
Bigarray.reshape
(Bigarray.genarray_of_array2 u)
[| mo_num ; mo_num ; mo_num |]
|> Bigarray.array3_of_genarray
in
let result = ref [] in
List.iter (fun gamma ->
List.iter (fun beta ->
List.iter (fun alpha ->
let x = u.{alpha,beta,gamma} in
if abs_float x > Constants.integrals_cutoff then
result := (alpha, beta, gamma, delta, x) :: !result;
) (list_range 1 beta)
) range_mo
) (list_range 1 delta);
Array.of_list !result
in
let n = ref 0 in
Stream.of_list range_mo
|> Farm.run ~f:task ~ordered:true ~comm:Parallel.Node.comm
|> Stream.iter (fun l ->
if Parallel.master then (incr n ; Printf.eprintf "\r%d / %d%!" !n mo_num);
Array.iter (fun (alpha, beta, gamma, delta, x) ->
set_chem destination alpha beta gamma delta x) l);
if Parallel.master then Printf.eprintf "\n%!";
broadcast destination
let svd_of_dense t =
let four_index =
let f =
match t.four_index with
| Dense a -> a
| _ -> invalid_arg "svd_of_dense expects a Dense storage"
in
let size = t.size in
let b = Mat.create (size*size) (size*size) in
for k = 1 to size do
for l = 1 to size do
let ac = sym_index k l in
lacpy f ~n:1 ~ac ~b ~bc:(dense_index k l size)
|> ignore;
lacpy f ~n:1 ~ac ~b ~bc:(dense_index l k size)
|> ignore
done
done;
let d, u, v = gesdd b in
let rank =
let w = copy d in
scal (1. /. d.{1}) w;
let rec aux i =
if i = Vec.dim w then i else
if w.{i} < epsilon_float then i else
aux (i+1)
in aux 1
in
let d = copy d ~n:rank in
let u = lacpy ~n:rank u in
let v = lacpy ~m:rank v in
Svd (d,u,v)
in
{ t with four_index }
let dense_of_svd t =
let four_index =
let d, u, v =
match t.four_index with
| Svd (d, u, v) -> d, u, v
| _ -> invalid_arg "dense_of_svd expects a Svd storage"
in
let f =
gemm u @@ gemm (Mat.of_diag d) v
in
let size = t.size in
let b = Mat.create (size*size) ((size*(size+1))/2) in
for k = 1 to size do
for l = 1 to k do
let bc = sym_index k l in
lacpy f ~ac:(dense_index l k size) ~n:1 ~bc ~b
|> ignore
done
done;
Dense b
in
{ t with four_index }
let four_index_transform_svd coef source =
let t0 = Unix.gettimeofday () in
let svd = svd_of_dense source in
Printf.printf "Computed SVD in %f seconds\n%!" (Unix.gettimeofday () -. t0);
let u0, d0, v0 =
match svd.four_index with
| Svd (d, u, v) -> u, d, (Mat.transpose_copy v)
| _ -> assert false
in
let t0 = Unix.gettimeofday () in
let ao_num = Mat.dim1 coef in
let mo_num = Mat.dim2 coef in
let svd_num = Vec.dim d0 in
let destination = create ~size:mo_num `Dense in
let mo_num_2 = mo_num * mo_num in
if Parallel.master then Printf.eprintf "4-idx transformation \n%!";
let u_vecs = Mat.to_col_vecs u0
and v_vecs = Mat.to_col_vecs v0
in
Printf.printf "%d %d %d %d\n" (Mat.dim1 u0) (Mat.dim2 u0) (Mat.dim1 v0) (Mat.dim2 v0);
let task a =
let o =
Bigarray.reshape_2 (Bigarray.genarray_of_array1 u_vecs.(a-1)) ao_num ao_num
in
let u =
let p = gemm ~transa:`T coef @@ gemm o coef in
Bigarray.reshape_1 (Bigarray.genarray_of_array2 p) mo_num_2
in
let o =
Bigarray.reshape_2 (Bigarray.genarray_of_array1 v_vecs.(a-1)) ao_num ao_num
in
let v =
let p = gemm ~transa:`T coef @@ gemm o coef in
Bigarray.reshape_1 (Bigarray.genarray_of_array2 p) mo_num_2
in
(u, v)
in
(*
let four_index = Svd (d,u'',v'') in
let t = make_from_four_index
{ source with
size = mo_num;
two_index;
two_index_anti;
three_index;
three_index_anti;
four_index
}
*)
let n = ref 0 in
let u_list = ref []
and v_list = ref []
in
Util.list_range 1 svd_num
|> Stream.of_list
|> Farm.run ~f:task ~ordered:true ~comm:Parallel.Node.comm
|> Stream.iter (fun (u,v) ->
if Parallel.master then (incr n ; Printf.eprintf "\r%d / %d%!" !n mo_num);
u_list := u :: !u_list;
v_list := v :: !v_list);
let u = Mat.of_col_vecs_list (List.rev !u_list)
and v = Mat.of_col_vecs_list (List.rev !v_list)
in
Printf.printf "Computed 4-idx in %f seconds\n%!" (Unix.gettimeofday () -. t0);
let result =
let p = gemm u @@ gemm ~transb:`T (Mat.of_diag d0) v in
Bigarray.reshape (Bigarray.genarray_of_array2 p) [| mo_num ; mo_num ; mo_num ; mo_num |]
in
for a=1 to mo_num do
for b=a to mo_num do
for c=1 to mo_num do
for d=c to mo_num do
let x = result.{a,b,c,d} in
if abs_float x > Constants.integrals_cutoff then
set_chem destination a b c d x;
done
done
done
done;
Printf.printf "Computed 4-idx in %f seconds\n%!" (Unix.gettimeofday () -. t0);
if Parallel.master then Printf.eprintf "\n%!";
broadcast destination
let four_index_transform coef source =
match source.four_index with
| Sparse _ -> four_index_transform_dense_sparse `Sparse coef source
| Svd _ -> assert false (* four_index_transform_svd coef source *)
| Dense _ -> four_index_transform_dense_sparse `Dense coef source