9
1
mirror of https://github.com/QuantumPackage/qp2.git synced 2024-11-15 01:43:38 +01:00
qp2/ocaml/Input_ao_basis.ml
2019-01-25 11:39:31 +01:00

469 lines
13 KiB
OCaml

open Qptypes;;
open Qputils;;
open Core;;
module Ao_basis : sig
type t =
{ ao_basis : AO_basis_name.t;
ao_num : AO_number.t ;
ao_prim_num : AO_prim_number.t array;
ao_prim_num_max : AO_prim_number.t;
ao_nucl : Nucl_number.t array;
ao_power : Symmetry.Xyz.t array;
ao_coef : AO_coef.t array;
ao_expo : AO_expo.t array;
ao_cartesian : bool;
} [@@deriving sexp]
;;
val read : unit -> t option
val to_string : t -> string
val to_basis : t -> Basis.t
val reorder : t -> t
val ordering : t -> int array
val write : t -> unit
val to_md5 : t -> MD5.t
val to_rst : t -> Rst_string.t
end = struct
type t =
{ ao_basis : AO_basis_name.t;
ao_num : AO_number.t ;
ao_prim_num : AO_prim_number.t array;
ao_prim_num_max : AO_prim_number.t;
ao_nucl : Nucl_number.t array;
ao_power : Symmetry.Xyz.t array;
ao_coef : AO_coef.t array;
ao_expo : AO_expo.t array;
ao_cartesian : bool;
} [@@deriving sexp]
;;
let get_default = Qpackage.get_ezfio_default "ao_basis";;
let read_ao_basis () =
Ezfio.get_ao_basis_ao_basis ()
|> AO_basis_name.of_string
;;
let read_ao_num () =
Ezfio.get_ao_basis_ao_num ()
|> AO_number.of_int
;;
let read_ao_prim_num () =
Ezfio.get_ao_basis_ao_prim_num ()
|> Ezfio.flattened_ezfio
|> Array.map ~f:AO_prim_number.of_int
;;
let read_ao_prim_num_max () =
Ezfio.get_ao_basis_ao_prim_num ()
|> Ezfio.flattened_ezfio
|> Array.fold ~f:(fun x y -> if x>y then x else y) ~init:0
|> AO_prim_number.of_int
;;
let read_ao_nucl () =
let nmax = Nucl_number.get_max () in
Ezfio.get_ao_basis_ao_nucl ()
|> Ezfio.flattened_ezfio
|> Array.map ~f:(fun x-> Nucl_number.of_int ~max:nmax x)
;;
let read_ao_power () =
let x = Ezfio.get_ao_basis_ao_power () in
let dim = x.Ezfio.dim.(0) in
let data = Ezfio.flattened_ezfio x in
let result = Array.init dim ~f:(fun x -> "") in
for i=1 to dim
do
if (data.(i-1) > 0) then
result.(i-1) <- result.(i-1)^"x"^(Int.to_string data.(i-1));
if (data.(dim+i-1) > 0) then
result.(i-1) <- result.(i-1)^"y"^(Int.to_string data.(dim+i-1));
if (data.(2*dim+i-1) > 0) then
result.(i-1) <- result.(i-1)^"z"^(Int.to_string data.(2*dim+i-1));
done;
Array.map ~f:Symmetry.Xyz.of_string result
;;
let read_ao_coef () =
Ezfio.get_ao_basis_ao_coef ()
|> Ezfio.flattened_ezfio
|> Array.map ~f:AO_coef.of_float
;;
let read_ao_expo () =
Ezfio.get_ao_basis_ao_expo ()
|> Ezfio.flattened_ezfio
|> Array.map ~f:AO_expo.of_float
;;
let read_ao_cartesian () =
if not (Ezfio.has_ao_basis_ao_cartesian ()) then
get_default "ao_cartesian"
|> Bool.of_string
|> Ezfio.set_ao_basis_ao_cartesian
;
Ezfio.get_ao_basis_ao_cartesian ()
;;
let to_long_basis b =
let ao_num = AO_number.to_int b.ao_num in
let gto_array = Array.init (AO_number.to_int b.ao_num)
~f:(fun i ->
let s = Symmetry.Xyz.to_symmetry b.ao_power.(i) in
let ao_prim_num = AO_prim_number.to_int b.ao_prim_num.(i) in
let prims = List.init ao_prim_num ~f:(fun j ->
let prim = { GaussianPrimitive.sym = s ;
GaussianPrimitive.expo = b.ao_expo.(ao_num*j+i)
}
in
let coef = b.ao_coef.(ao_num*j+i) in
(prim,coef)
) in
Gto.of_prim_coef_list prims
)
in
let rec do_work accu sym gto nucl =
match (sym, gto, nucl) with
| (s::srest, g::grest, n::nrest) ->
do_work ((s,g,n)::accu) srest grest nrest
| ([],[],[]) -> List.rev accu
| _ -> assert false
in
do_work []
(Array.to_list b.ao_power)
(Array.to_list gto_array)
(Array.to_list b.ao_nucl)
;;
let to_basis b =
to_long_basis b
|> Long_basis.to_basis
;;
let to_md5 b =
let short_basis = to_basis b in
Basis.to_md5 short_basis
;;
let write_md5 b =
to_md5 b
|> MD5.to_string
|> Ezfio.set_ao_basis_ao_md5
;;
let write_ao_basis name =
AO_basis_name.to_string name
|> Ezfio.set_ao_basis_ao_basis
;;
let write b =
let { ao_basis ;
ao_num ;
ao_prim_num ;
ao_prim_num_max ;
ao_nucl ;
ao_power ;
ao_coef ;
ao_expo ;
ao_cartesian ;
} = b
in
write_md5 b ;
write_ao_basis ao_basis;
let ao_num = AO_number.to_int ao_num
and ao_prim_num_max = AO_prim_number.to_int ao_prim_num_max
in
let ao_prim_num =
Array.to_list ao_prim_num
|> List.map ~f:AO_prim_number.to_int
in
Ezfio.set_ao_basis_ao_prim_num (Ezfio.ezfio_array_of_list
~rank:1 ~dim:[| ao_num |] ~data:ao_prim_num) ;
let ao_nucl =
Array.to_list ao_nucl
|> List.map ~f:Nucl_number.to_int
in
Ezfio.set_ao_basis_ao_nucl(Ezfio.ezfio_array_of_list
~rank:1 ~dim:[| ao_num |] ~data:ao_nucl) ;
let ao_power =
let l = Array.to_list ao_power in
List.concat [
(List.map ~f:(fun a -> Positive_int.to_int a.Symmetry.Xyz.x) l) ;
(List.map ~f:(fun a -> Positive_int.to_int a.Symmetry.Xyz.y) l) ;
(List.map ~f:(fun a -> Positive_int.to_int a.Symmetry.Xyz.z) l) ]
in
Ezfio.set_ao_basis_ao_power(Ezfio.ezfio_array_of_list
~rank:2 ~dim:[| ao_num ; 3 |] ~data:ao_power) ;
Ezfio.set_ao_basis_ao_cartesian(ao_cartesian);
let ao_coef =
Array.to_list ao_coef
|> List.map ~f:AO_coef.to_float
in
Ezfio.set_ao_basis_ao_coef(Ezfio.ezfio_array_of_list
~rank:2 ~dim:[| ao_num ; ao_prim_num_max |] ~data:ao_coef) ;
let ao_expo =
Array.to_list ao_expo
|> List.map ~f:AO_expo.to_float
in
Ezfio.set_ao_basis_ao_expo(Ezfio.ezfio_array_of_list
~rank:2 ~dim:[| ao_num ; ao_prim_num_max |] ~data:ao_expo) ;
;;
let read () =
if (Ezfio.has_ao_basis_ao_basis ()) then
begin
let result =
{ ao_basis = read_ao_basis ();
ao_num = read_ao_num () ;
ao_prim_num = read_ao_prim_num ();
ao_prim_num_max = read_ao_prim_num_max ();
ao_nucl = read_ao_nucl ();
ao_power = read_ao_power ();
ao_coef = read_ao_coef () ;
ao_expo = read_ao_expo () ;
ao_cartesian = read_ao_cartesian () ;
}
in
to_md5 result
|> MD5.to_string
|> Ezfio.set_ao_basis_ao_md5 ;
Some result
end
else
None
;;
let ordering b =
let ordered_basis =
to_basis b
|> Long_basis.of_basis
|> Array.of_list
and unordered_basis =
to_long_basis b
|> Array.of_list
in
let find x a =
let rec find x a i =
if i = Array.length a then
find2 x a 0
else
if a.(i) = Some x then
(a.(i) <- None ; i)
else
find x a (i+1)
and find2 (s,g,n) a i =
if i = Array.length a then -1
else
match a.(i) with
| None -> find2 (s,g,n) a (i+1)
| Some (s', g', n') ->
if s <> s' || n <> n' then find2 (s,g,n) a (i+1)
else
let lc = List.map ~f:(fun (prim, _) -> prim) g.Gto.lc
and lc' = List.map ~f:(fun (prim, _) -> prim) g'.Gto.lc
in
if lc <> lc' then find2 (s,g,n) a (i+1) else (a.(i) <- None ; i)
in
find x a 0
in
let search_array = Array.map ~f:(fun i -> Some i) unordered_basis in
Array.map ~f:(fun x -> find x search_array) ordered_basis
;;
let of_long_basis long_basis name ao_cartesian =
let ao_num = List.length long_basis |> AO_number.of_int in
let ao_prim_num =
List.map long_basis ~f:(fun (_,g,_) -> List.length g.Gto.lc
|> AO_prim_number.of_int )
|> Array.of_list
and ao_nucl =
List.map long_basis ~f:(fun (_,_,n) -> n)
|> Array.of_list
and ao_power =
List.map ~f:(fun (x,_,_) -> x) long_basis
|> Array.of_list
in
let ao_prim_num_max = Array.fold ~init:0 ~f:(fun s x ->
if AO_prim_number.to_int x > s then AO_prim_number.to_int x else s)
ao_prim_num
|> AO_prim_number.of_int
in
let gtos =
List.map long_basis ~f:(fun (_,x,_) -> x)
in
let create_expo_coef ec =
let coefs =
begin match ec with
| `Coefs -> List.map gtos ~f:(fun x->
List.map x.Gto.lc ~f:(fun (_,coef) -> AO_coef.to_float coef) )
| `Expos -> List.map gtos ~f:(fun x->
List.map x.Gto.lc ~f:(fun (prim,_) -> AO_expo.to_float
prim.GaussianPrimitive.expo) )
end
in
let rec get_n n accu = function
| [] -> List.rev accu
| h::tail ->
let y =
begin match List.nth h n with
| Some x -> x
| None -> 0.
end
in
get_n n (y::accu) tail
in
let rec build accu = function
| n when n=(AO_prim_number.to_int ao_prim_num_max) -> accu
| n -> build ( accu @ (get_n n [] coefs) ) (n+1)
in
build [] 0
in
let ao_coef = create_expo_coef `Coefs
|> Array.of_list
|> Array.map ~f:AO_coef.of_float
and ao_expo = create_expo_coef `Expos
|> Array.of_list
|> Array.map ~f:AO_expo.of_float
in
{ ao_basis = name ;
ao_num ; ao_prim_num ; ao_prim_num_max ; ao_nucl ;
ao_power ; ao_coef ; ao_expo ; ao_cartesian }
;;
let reorder b =
let order = ordering b in
let f a = Array.init (Array.length a) ~f:(fun i -> a.(order.(i))) in
let ao_prim_num_max = AO_prim_number.to_int b.ao_prim_num_max
and ao_num = AO_number.to_int b.ao_num in
let ao_coef =
Array.init ao_prim_num_max ~f:(fun i ->
f @@ Array.init ao_num ~f:(fun j -> b.ao_coef.(i*ao_num + j) )
) |> Array.to_list |> Array.concat
in
let ao_expo =
Array.init ao_prim_num_max ~f:(fun i ->
f @@ Array.init ao_num ~f:(fun j -> b.ao_expo.(i*ao_num + j) )
) |> Array.to_list |> Array.concat
in
{ b with
ao_prim_num = f b.ao_prim_num ;
ao_nucl = f b.ao_nucl ;
ao_power = f b.ao_power ;
ao_coef ;
ao_expo ;
}
;;
let to_rst b =
let print_sym =
let l = List.init (Array.length b.ao_power) ~f:(
fun i -> ( (i+1),b.ao_nucl.(i),b.ao_power.(i) ) )
in
let rec do_work = function
| [] -> []
| (i,n,x)::tail ->
(Printf.sprintf " %5d %6d %-8s\n" i (Nucl_number.to_int n)
(Symmetry.Xyz.to_string x)
)::(do_work tail)
in do_work l
|> String.concat
in
let short_basis = to_basis b in
Printf.sprintf "
Name of the AO basis ::
ao_basis = %s
Cartesian coordinates (6d,10f,...) ::
ao_cartesian = %s
Basis set (read-only) ::
%s
======= ========= ===========
Basis Nucleus Symmetries
======= ========= ===========
%s
======= ========= ===========
" (AO_basis_name.to_string b.ao_basis)
(Bool.to_string b.ao_cartesian)
(Basis.to_string short_basis
|> String.split ~on:'\n'
|> List.map ~f:(fun x-> " "^x)
|> String.concat ~sep:"\n"
) print_sym
|> Rst_string.of_string
;;
let read_rst s =
let s = Rst_string.to_string s
|> String.split ~on:'\n'
in
let rec extract_basis = function
| [] -> failwith "Error in basis set"
| line :: tail ->
let line = String.strip line in
if line = "Basis set (read-only) ::" then
String.concat tail ~sep:"\n"
else
extract_basis tail
in
extract_basis s
;;
let to_string b =
Printf.sprintf "
ao_basis = %s
ao_num = %s
ao_prim_num = %s
ao_prim_num_max = %s
ao_nucl = %s
ao_power = %s
ao_coef = %s
ao_expo = %s
ao_cartesian = %s
md5 = %s
"
(AO_basis_name.to_string b.ao_basis)
(AO_number.to_string b.ao_num)
(b.ao_prim_num |> Array.to_list |> List.map
~f:(AO_prim_number.to_string) |> String.concat ~sep:", " )
(AO_prim_number.to_string b.ao_prim_num_max)
(b.ao_nucl |> Array.to_list |> List.map ~f:Nucl_number.to_string |>
String.concat ~sep:", ")
(b.ao_power |> Array.to_list |> List.map ~f:(fun x->
"("^(Symmetry.Xyz.to_string x)^")" )|> String.concat ~sep:", ")
(b.ao_coef |> Array.to_list |> List.map ~f:AO_coef.to_string
|> String.concat ~sep:", ")
(b.ao_expo |> Array.to_list |> List.map ~f:AO_expo.to_string
|> String.concat ~sep:", ")
(b.ao_cartesian |> Bool.to_string)
(to_md5 b |> MD5.to_string )
;;
end