mirror of
https://gitlab.com/scemama/QCaml.git
synced 2024-12-06 20:43:33 +01:00
Accelerated FourIdxStorage
This commit is contained in:
parent
15390759f9
commit
63f0b379de
@ -57,7 +57,12 @@ let make ?guess:(guess=`Huckel) ?max_scf:(max_scf=64) ?level_shift:(level_shift=
|
||||
in
|
||||
x.Fock.fock, x.Fock.core, x.Fock.coulomb, x.Fock.exchange
|
||||
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 *)
|
||||
let m_F =
|
||||
let m_SC =
|
||||
|
@ -4,7 +4,7 @@ type index_pair = { first : int ; second : int }
|
||||
|
||||
|
||||
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
|
||||
|
||||
type t =
|
||||
@ -24,49 +24,85 @@ let key_of_indices ~r1 ~r2 =
|
||||
f p q
|
||||
|
||||
|
||||
let dense_index i j size =
|
||||
(j-1)*size + i
|
||||
|
||||
|
||||
let get_four_index ~r1 ~r2 t =
|
||||
match t.four_index with
|
||||
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
|
||||
Bigarray.Genarray.get t [| i; j; k; l|]
|
||||
| Sparse t -> let key = key_of_indices ~r1 ~r2 in
|
||||
try Hashtbl.find t key
|
||||
| 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 set_four_index ~r1 ~r2 ~value t =
|
||||
match t.four_index with
|
||||
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
|
||||
Bigarray.Genarray.set t [| i; j; k; l|] value;
|
||||
Bigarray.Genarray.set t [| k; j; i; l|] value;
|
||||
Bigarray.Genarray.set t [| i; l; k; j|] value;
|
||||
Bigarray.Genarray.set t [| k; l; i; j|] value;
|
||||
Bigarray.Genarray.set t [| j; i; l; k|] value;
|
||||
Bigarray.Genarray.set t [| j; k; l; i|] value;
|
||||
Bigarray.Genarray.set t [| l; i; j; k|] value;
|
||||
Bigarray.Genarray.set t [| l; k; j; i|] value;
|
||||
| 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
|
||||
)
|
||||
|
||||
| Sparse t -> let key = key_of_indices ~r1 ~r2 in
|
||||
Hashtbl.replace t key value
|
||||
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
|
||||
Hashtbl.replace a key value
|
||||
|
||||
|
||||
let increment_four_index ~r1 ~r2 ~value x =
|
||||
match x.four_index with
|
||||
| Dense t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
|
||||
t.{i,j,k,l} <- t.{i,j,k,l} +. value;
|
||||
t.{k,j,i,l} <- t.{k,j,i,l} +. value;
|
||||
t.{i,l,k,j} <- t.{i,l,k,j} +. value;
|
||||
t.{k,l,i,j} <- t.{k,l,i,j} +. value;
|
||||
t.{j,i,l,k} <- t.{j,i,l,k} +. value;
|
||||
t.{j,k,l,i} <- t.{j,k,l,i} +. value;
|
||||
t.{l,i,j,k} <- t.{l,i,j,k} +. value;
|
||||
t.{l,k,j,i} <- t.{l,k,j,i} +. value
|
||||
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)
|
||||
)
|
||||
|
||||
| Sparse t -> let key = key_of_indices ~r1 ~r2 in
|
||||
| Sparse a -> let key = key_of_indices ~r1 ~r2 in
|
||||
let old_value =
|
||||
try Hashtbl.find t key
|
||||
try Hashtbl.find a key
|
||||
with Not_found -> 0.
|
||||
in
|
||||
Hashtbl.replace t key (old_value +. value)
|
||||
Hashtbl.replace a key (old_value +. value)
|
||||
|
||||
let get ~r1 ~r2 =
|
||||
get_four_index ~r1 ~r2
|
||||
@ -83,9 +119,9 @@ let create ~size sparsity =
|
||||
match sparsity with
|
||||
| `Dense ->
|
||||
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
|
||||
Bigarray.Genarray.fill result 0.;
|
||||
Bigarray.Array2.fill result 0.;
|
||||
Dense result
|
||||
|
||||
| `Sparse ->
|
||||
|
Loading…
Reference in New Issue
Block a user