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 "@[@[@[%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