Separated 4, 3 and 2-idx integrals

This commit is contained in:
Anthony Scemama 2019-04-05 14:33:31 +02:00
parent a0b84249ec
commit 0e34c3adca
2 changed files with 181 additions and 77 deletions

View File

@ -6,15 +6,20 @@ 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 (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t
| Dense of array2
| Sparse of (int, float) Hashtbl.t
type t =
{
size : int ;
four_index : storage_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 =
@ -28,85 +33,164 @@ let key_of_indices ~r1 ~r2 =
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 get_four_index ~r1 ~r2 t =
match t.four_index with
| Dense a -> (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 );
Bigarray.Array2.unsafe_get a (dense_index i j size) (dense_index k l size)
)
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
try Hashtbl.find a key
with Not_found -> 0.
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
let set_four_index ~r1 ~r2 ~value t =
match t.four_index with
| Dense a -> (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 ij = (dense_index i j size)
and kl = (dense_index k l size)
and il = (dense_index i l size)
and kj = (dense_index k j size)
and ji = (dense_index j i size)
and lk = (dense_index l k size)
and li = (dense_index l i size)
and jk = (dense_index j k size)
in
let open Bigarray.Array2 in
unsafe_set a ij kl value;
unsafe_set a kj il value;
unsafe_set a il kj value;
unsafe_set a kl ij value;
unsafe_set a ji lk value;
unsafe_set a li jk value;
unsafe_set a jk li value;
unsafe_set a lk ji value
)
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
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
Hashtbl.replace a key value
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
let increment_four_index ~r1 ~r2 ~value t =
match t.four_index with
| Dense a -> (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 ij = (dense_index i j size)
and kl = (dense_index k l size)
and il = (dense_index i l size)
and kj = (dense_index k j size)
and ji = (dense_index j i size)
and lk = (dense_index l k size)
and li = (dense_index l i size)
and jk = (dense_index j k size)
in
let open Bigarray.Array2 in
unsafe_set a ij kl (value +. unsafe_get a ij kl) ;
unsafe_set a kj il (value +. unsafe_get a kj il) ;
unsafe_set a il kj (value +. unsafe_get a il kj) ;
unsafe_set a kl ij (value +. unsafe_get a kl ij) ;
unsafe_set a ji lk (value +. unsafe_get a ji lk) ;
unsafe_set a li jk (value +. unsafe_get a li jk) ;
unsafe_set a jk li (value +. unsafe_get a jk li) ;
unsafe_set a lk ji (value +. unsafe_get a lk ji)
)
else if j=k then
unsafe_get t.three_index_anti (dense_index i l t.size) j
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
let old_value =
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 j t.size) (dense_index k l t.size)
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
try Hashtbl.find a key
with Not_found -> 0.
in
Hashtbl.replace a key (old_value +. value)
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
if i=k then
if j=l then
begin
unsafe_set t.two_index i j value;
unsafe_set t.two_index j i value;
end
else
begin
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
if j=k then
begin
unsafe_set t.two_index_anti i j value;
unsafe_set t.two_index_anti j i value;
end
else
begin
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
if k=l then
begin
unsafe_set t.two_index_anti i k value;
unsafe_set t.two_index_anti k i value;
end
else
(* <ii|kl> *)
begin
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
else
match t.four_index with
| Dense a -> let ij = (dense_index i j t.size)
and kl = (dense_index k l t.size)
and il = (dense_index i l t.size)
and kj = (dense_index k j t.size)
and ji = (dense_index j i t.size)
and lk = (dense_index l k t.size)
and li = (dense_index l i t.size)
and jk = (dense_index j k t.size)
in
begin
unsafe_set a ij kl value;
unsafe_set a kj il value;
unsafe_set a il kj value;
unsafe_set a kl ij value;
unsafe_set a ji lk value;
unsafe_set a li jk value;
unsafe_set a jk li value;
unsafe_set a lk ji value
end
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
Hashtbl.replace a key value
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
@ -124,21 +208,39 @@ let set ~r1 ~r2 ~value =
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 Float64 [| size*size ; size*size |]
SharedMemory.create ~temp_dir Float64 [| size*size ; size*size |]
|> Bigarray.array2_of_genarray
in Dense result
in
Dense result
| `Sparse -> let result = Hashtbl.create (size*size+13) in
Sparse result
in
{ size ; four_index }
{ size ; two_index ; two_index_anti ; three_index ; three_index_anti ; four_index }
let size t = t.size

View File

@ -6,7 +6,7 @@ let nuclei_file : string option ref = ref None
let speclist = [
( "-b" , Arg.String (fun x -> basis_file := Some x),
"File containing the atomic basis set") ;
( "-c" , Arg.String (fun x -> nuclei_file := Some x),
( "-x" , Arg.String (fun x -> nuclei_file := Some x),
"File containing the nuclear coordinates") ;
( "-o" , Arg.String (fun x -> out_file := Some x) ,
"Output file") ;
@ -23,7 +23,7 @@ let run ~out =
| Some x -> x
and nuclei_file =
match !nuclei_file with
| None -> raise (Invalid_argument "Coordinate file should be specified with -c")
| None -> raise (Invalid_argument "Coordinate file should be specified with -x")
| Some x -> x
in
@ -41,12 +41,14 @@ let run ~out =
let eN_ints = AOBasis.eN_ints ao_basis in
let kin_ints = AOBasis.kin_ints ao_basis in
let ee_ints = AOBasis.ee_ints ao_basis in
let f12_ints = AOBasis.f12_ints ao_basis in
Overlap.to_file ~filename:(out_file^".overlap") overlap;
NucInt.to_file ~filename:(out_file^".nuc") eN_ints;
KinInt.to_file ~filename:(out_file^".kin") kin_ints;
ERI.to_file ~filename:(out_file^".eri") ee_ints;
(*
let f12_ints = AOBasis.f12_ints ao_basis in
F12.to_file ~filename:(out_file^".f12") f12_ints;
*)
()