10
1
mirror of https://gitlab.com/scemama/QCaml.git synced 2024-08-29 07:51:59 +02:00

Accelerated FourIdxStorage

This commit is contained in:
Anthony Scemama 2018-06-01 10:07:17 +02:00
parent 15390759f9
commit 63f0b379de
2 changed files with 74 additions and 33 deletions

View File

@ -57,7 +57,12 @@ let make ?guess:(guess=`Huckel) ?max_scf:(max_scf=64) ?level_shift:(level_shift=
in in
x.Fock.fock, x.Fock.core, x.Fock.coulomb, x.Fock.exchange x.Fock.fock, x.Fock.core, x.Fock.coulomb, x.Fock.exchange
in in
(*
debug_matrix "Fock" m_F;
debug_matrix "Coulomb" m_J;
debug_matrix "Exchange" m_K;
debug_matrix "HCore" m_Hc;
*)
(* Add level shift in AO basis *) (* Add level shift in AO basis *)
let m_F = let m_F =
let m_SC = let m_SC =

View File

@ -4,7 +4,7 @@ type index_pair = { first : int ; second : int }
type storage_t = type storage_t =
| Dense of (float, Bigarray.float32_elt, Bigarray.fortran_layout) Bigarray.Genarray.t | Dense of (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t
| Sparse of (int, float) Hashtbl.t | Sparse of (int, float) Hashtbl.t
type t = type t =
@ -24,49 +24,85 @@ let key_of_indices ~r1 ~r2 =
f p q f p q
let dense_index i j size =
(j-1)*size + i
let get_four_index ~r1 ~r2 t = let get_four_index ~r1 ~r2 t =
match t.four_index with match t.four_index with
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in | Dense a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
Bigarray.Genarray.get t [| i; j; k; l|] let size = t.size in
| Sparse t -> let key = key_of_indices ~r1 ~r2 in assert ( (i lor j lor k lor l) > 0 );
try Hashtbl.find t key 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. with Not_found -> 0.
let set_four_index ~r1 ~r2 ~value t = let set_four_index ~r1 ~r2 ~value t =
match t.four_index with match t.four_index with
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in | Dense a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
Bigarray.Genarray.set t [| i; j; k; l|] value; let size = t.size in
Bigarray.Genarray.set t [| k; j; i; l|] value; assert ( (i lor j lor k lor l) > 0 );
Bigarray.Genarray.set t [| i; l; k; j|] value; assert ( i <= size && j <= size && k <= size && l <= size);
Bigarray.Genarray.set t [| k; l; i; j|] value; let ij = (dense_index i j size)
Bigarray.Genarray.set t [| j; i; l; k|] value; and kl = (dense_index k l size)
Bigarray.Genarray.set t [| j; k; l; i|] value; and il = (dense_index i l size)
Bigarray.Genarray.set t [| l; i; j; k|] value; and kj = (dense_index k j size)
Bigarray.Genarray.set t [| l; k; j; i|] value; 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
)
| Sparse t -> let key = key_of_indices ~r1 ~r2 in | Sparse a -> let key = key_of_indices ~r1 ~r2 in
Hashtbl.replace t key value Hashtbl.replace a key value
let increment_four_index ~r1 ~r2 ~value x = let increment_four_index ~r1 ~r2 ~value t =
match x.four_index with match t.four_index with
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in | Dense a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
t.{i,j,k,l} <- t.{i,j,k,l} +. value; let size = t.size in
t.{k,j,i,l} <- t.{k,j,i,l} +. value; assert ( (i lor j lor k lor l) > 0 );
t.{i,l,k,j} <- t.{i,l,k,j} +. value; assert ( i <= size && j <= size && k <= size && l <= size);
t.{k,l,i,j} <- t.{k,l,i,j} +. value; let ij = (dense_index i j size)
t.{j,i,l,k} <- t.{j,i,l,k} +. value; and kl = (dense_index k l size)
t.{j,k,l,i} <- t.{j,k,l,i} +. value; and il = (dense_index i l size)
t.{l,i,j,k} <- t.{l,i,j,k} +. value; and kj = (dense_index k j size)
t.{l,k,j,i} <- t.{l,k,j,i} +. value 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)
)
| Sparse t -> let key = key_of_indices ~r1 ~r2 in | Sparse a -> let key = key_of_indices ~r1 ~r2 in
let old_value = let old_value =
try Hashtbl.find t key try Hashtbl.find a key
with Not_found -> 0. with Not_found -> 0.
in in
Hashtbl.replace t key (old_value +. value) Hashtbl.replace a key (old_value +. value)
let get ~r1 ~r2 = let get ~r1 ~r2 =
get_four_index ~r1 ~r2 get_four_index ~r1 ~r2
@ -83,9 +119,9 @@ let create ~size sparsity =
match sparsity with match sparsity with
| `Dense -> | `Dense ->
let result = let result =
Bigarray.Genarray.create Float32 Bigarray.fortran_layout [| size ; size ; size ; size |] Bigarray.Array2.create Float64 Bigarray.fortran_layout (size*size) (size*size)
in in
Bigarray.Genarray.fill result 0.; Bigarray.Array2.fill result 0.;
Dense result Dense result
| `Sparse -> | `Sparse ->