mirror of
https://gitlab.com/scemama/QCaml.git
synced 2024-12-27 14:53:32 +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
|
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 =
|
||||||
|
@ -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 ->
|
||||||
|
Loading…
Reference in New Issue
Block a user