QCaml/MOBasis/MOBasis.ml

283 lines
8.0 KiB
OCaml

open Lacaml.D
open Util
open Constants
(** One-electron orthogonal basis set, corresponding to Molecular Orbitals. *)
module HF = HartreeFock
module Si = Simulation
type mo_type =
| RHF | ROHF | UHF | CASSCF | Projected
| Natural of string
| Localized of string
type t =
{
simulation : Simulation.t; (* Simulation which produced the MOs *)
mo_type : mo_type; (* Kind of MOs (RHF, CASSCF, Localized...) *)
mo_occupation : Vec.t; (* Occupation numbers *)
mo_coef : Mat.t; (* Matrix of the MO coefficients in the AO basis *)
eN_ints : NucInt.t lazy_t; (* Electron-nucleus potential integrals *)
ee_ints : ERI.t lazy_t; (* Electron-electron potential integrals *)
kin_ints : KinInt.t lazy_t; (* Kinetic energy integrals *)
one_e_ints : Mat.t lazy_t; (* Kinetic energy integrals *)
}
let size t =
Mat.dim2 t.mo_coef
let simulation t = t.simulation
let mo_type t = t.mo_type
let ao_basis t = Si.ao_basis t.simulation
let mo_occupation t = t.mo_occupation
let mo_coef t = t.mo_coef
let eN_ints t = Lazy.force t.eN_ints
let ee_ints t = Lazy.force t.ee_ints
let kin_ints t = Lazy.force t.kin_ints
let two_e_ints t = Lazy.force t.ee_ints
let one_e_ints t = Lazy.force t.one_e_ints
let mo_energies t =
let f =
let m_C = mo_coef t in
let m_N = Mat.of_diag @@ mo_occupation t in
let m_P =
gemm m_C @@ (gemm m_N m_C ~transb:`T)
in
match t.mo_type with
| RHF -> Fock.make_rhf ~density:m_P (ao_basis t)
| ROHF -> Fock.make_uhf ~density_same:m_P ~density_other:m_P (ao_basis t)
| _ -> failwith "Not implemented"
in
let m_F = Fock.fock f in
Vec.init (size t) (fun i -> m_F.{i,i})
let mo_matrix_of_ao_matrix ~mo_coef ao_matrix =
xt_o_x ~x:mo_coef ~o:ao_matrix
let ao_matrix_of_mo_matrix ~mo_coef ~ao_overlap mo_matrix =
let sc = gemm ao_overlap mo_coef in
x_o_xt ~x:sc ~o:mo_matrix
let four_index_transform ~mo_coef eri_ao =
let ao_num = Mat.dim1 mo_coef in
let mo_num = Mat.dim2 mo_coef in
let eri_mo =
ERI.create ~size:mo_num `Dense
in
let mo_num_2 = mo_num * mo_num in
let ao_num_2 = ao_num * ao_num in
let ao_mo_num = ao_num * mo_num in
let range_mo = list_range 1 mo_num in
let range_ao = list_range 1 ao_num in
let u = Mat.create mo_num_2 mo_num
and o = Mat.create ao_num ao_num_2
and p = Mat.create ao_num_2 mo_num
and q = Mat.create ao_mo_num mo_num
in
if Parallel.master then Printf.eprintf "4-idx transformation \n%!";
let task delta =
Mat.fill u 0.;
List.iter (fun l ->
if abs_float mo_coef.{l,delta} > epsilon then
begin
let jk = ref 0 in
List.iter (fun k ->
List.iter (fun j ->
incr jk;
ERI.get_chem_all_i eri_ao ~j ~k ~l
|> Array.iteri (fun i x -> o.{i+1,!jk} <- x)
) range_ao
) range_ao;
(* o_i_jk *)
let p =
gemm ~transa:`T ~c:p o mo_coef
(* p_jk_alpha = \sum_i o_i_jk c_i_alpha *)
in
let p' =
Bigarray.reshape_2 (Bigarray.genarray_of_array2 p) ao_num ao_mo_num
(* p_j_kalpha *)
in
let q =
gemm ~transa:`T ~c:q p' mo_coef
(* q_kalpha_beta = \sum_j p_j_kalpha c_j_beta *)
in
let q' =
Bigarray.reshape_2 (Bigarray.genarray_of_array2 q) ao_num mo_num_2
(* q_k_alphabeta = \sum_j p_j_kalpha c_j_beta *)
in
ignore @@
gemm ~transa:`T ~beta:1. ~alpha:mo_coef.{l,delta} ~c:u q' mo_coef ;
(* u_alphabeta_gamma = \sum_k q_k_alphabeta c_k_gamma *)
end
) range_ao;
let u =
Bigarray.reshape
(Bigarray.genarray_of_array2 u)
[| mo_num ; mo_num ; mo_num |]
|> Bigarray.array3_of_genarray
in
let result = ref [] in
List.iter (fun gamma ->
List.iter (fun beta ->
List.iter (fun alpha ->
let x = u.{alpha,beta,gamma} in
if x <> 0. then
result := (alpha, beta, gamma, delta, x) :: !result;
) (list_range 1 beta)
) range_mo
) (list_range 1 delta);
Array.of_list !result
in
let n = ref 0 in
Stream.of_list range_mo
|> Farm.run ~f:task ~ordered:false
|> Stream.iter (fun l ->
if Parallel.master then (Printf.eprintf "\r%d / %d%!" !n mo_num; incr n);
Array.iter (fun (alpha, beta, gamma, delta, x) ->
ERI.set_chem eri_mo alpha beta gamma delta x) l);
if Parallel.master then Printf.eprintf "\n";
Parallel.broadcast (lazy eri_mo)
let make ~simulation ~mo_type ~mo_occupation ~mo_coef () =
let ao_basis =
Si.ao_basis simulation
in
let eN_ints = lazy (
AOBasis.eN_ints ao_basis
|> NucInt.matrix
|> mo_matrix_of_ao_matrix ~mo_coef
|> NucInt.of_matrix
)
and kin_ints = lazy (
AOBasis.kin_ints ao_basis
|> KinInt.matrix
|> mo_matrix_of_ao_matrix ~mo_coef
|> KinInt.of_matrix
)
and ee_ints = lazy (
AOBasis.ee_ints ao_basis
|> four_index_transform ~mo_coef
)
in
let one_e_ints = lazy (
Mat.add (NucInt.matrix @@ Lazy.force eN_ints)
(KinInt.matrix @@ Lazy.force kin_ints) )
in
{ simulation ; mo_type ; mo_occupation ; mo_coef ;
eN_ints ; ee_ints ; kin_ints ; one_e_ints }
let of_hartree_fock hf =
let mo_coef = HF.eigenvectors hf in
let simulation = HF.simulation hf in
let mo_occupation = HF.occupation hf in
let mo_type =
match HF.kind hf with
| HartreeFock.RHF -> RHF
| HartreeFock.ROHF -> ROHF
| HartreeFock.UHF -> UHF
in
make ~simulation ~mo_type ~mo_occupation ~mo_coef ()
let of_mo_basis simulation other =
let mo_coef =
let basis = Simulation.ao_basis simulation in
let basis_other = ao_basis other in
let m_S =
Overlap.(matrix @@ of_basis_pair
(AOBasis.basis basis)
(AOBasis.basis basis_other) )
in
let m_X = AOBasis.ortho basis in
(* Project other vectors in the current basis *)
let m_C =
gemm m_S @@ mo_coef other
in
(* Append dummy vectors to the input vectors *)
let result =
let vecs = Mat.to_col_vecs m_X in
Array.iteri (fun i v -> if (i < Array.length vecs) then vecs.(i) <- v)
(Mat.to_col_vecs m_C) ;
Mat.of_col_vecs vecs
in
(* Gram-Schmidt Orthonormalization *)
gemm m_X @@ (Util.qr_ortho @@ gemm ~transa:`T m_X result)
in
let mo_occupation =
let occ = mo_occupation other in
Vec.init (Mat.dim2 mo_coef) (fun i ->
if (i <= Vec.dim occ) then occ.{i}
else 0.)
in
make ~simulation ~mo_type:Projected ~mo_occupation ~mo_coef ()
let pp_mo ?(start=1) ?finish ppf t =
let open Lacaml.Io in
let rows = Mat.dim1 t.mo_coef
and cols = Mat.dim2 t.mo_coef
in
let finish =
match finish with
| None -> cols
| Some x -> x
in
let rec aux first =
if (first > finish) then ()
else
begin
Format.fprintf ppf "@[<v>@[<v4>@[<h>%s@;" "Eigenvalues:";
Array.iteri (fun i x ->
if (i+1 >= first) && (i+1 <= first+4 ) then
Format.fprintf ppf "%12f@ " x)
(Vec.to_array @@ mo_energies t);
Format.fprintf ppf "@]@;";
Format.fprintf ppf "@[%a@]"
(Lacaml.Io.pp_lfmat
~row_labels:
(Array.init rows (fun i -> Printf.sprintf "%d " (i + 1)))
~col_labels:
(Array.init (min 5 (cols-first+1)) (fun i -> Printf.sprintf "-- %d --" (i + first) ))
~print_right:false
~print_foot:false
() ) (lacpy ~ac:first ~n:(min 5 (cols-first+1)) (t.mo_coef)) ;
Format.fprintf ppf "@]@;@;@]";
aux (first+5)
end
in
aux start