open Lacaml.D open Util open Constants type hartree_fock_data = { iteration : int ; coefficients : Mat.t option ; eigenvalues : Vec.t option ; error : float option ; diis : DIIS.t option ; energy : float option ; density : Mat.t option ; density_a : Mat.t option ; density_b : Mat.t option ; fock : Fock.t option ; fock_a : Fock.t option ; fock_b : Fock.t option ; } type hartree_fock_kind = | RHF (** Restricted Hartree-Fock *) | ROHF (** Restricted Open-shell Hartree-Fock *) | UHF (** Unrestricted Hartree-Fock *) type t = { kind : hartree_fock_kind; simulation : Simulation.t; guess : Guess.t; data : hartree_fock_data option lazy_t array; nocc : int ; } let empty = { iteration = 0 ; coefficients = None ; eigenvalues = None ; error = None ; diis = None ; energy = None ; density = None ; density_a = None ; density_b = None ; fock = None ; fock_a = None ; fock_b = None ; } module Si = Simulation module El = Electrons module Ao = AOBasis module Ov = Overlap let kind t = t.kind let simulation t = t.simulation let guess t = t.guess let nocc t = t.nocc let n_iterations t = Array.fold_left (fun accu x -> match Lazy.force x with | Some x -> accu + 1 | None -> accu ) 0 t.data let last_iteration t = of_some @@ Lazy.force (t.data.(n_iterations t - 1)) let eigenvectors t = let data = last_iteration t in of_some data.coefficients let eigenvalues t = let data = last_iteration t in of_some data.eigenvalues let density t = let data = last_iteration t in match kind t with | RHF -> of_some data.density | ROHF -> Mat.add (of_some data.density_a) (of_some data.density_b) | _ -> failwith "Not implemented" let occupation t = density t |> Mat.copy_diag ~n:(Mat.dim2 @@ eigenvectors t) let energy t = let data = last_iteration t in of_some data.energy let nuclear_repulsion t = Si.nuclear_repulsion (simulation t) let ao_basis t = Si.ao_basis (simulation t) let kin_energy t = let m_T = ao_basis t |> Ao.kin_ints |> KinInt.matrix in let m_P = density t in Mat.gemm_trace m_P m_T let eN_energy t = let m_V = ao_basis t |> Ao.eN_ints |> NucInt.matrix in let m_P = density t in Mat.gemm_trace m_P m_V let coulomb_energy t = let data = last_iteration t in match kind t with | RHF -> let m_P = of_some data.density in let fock = of_some data.fock in let m_J = Fock.coulomb fock in 0.5 *. Mat.gemm_trace m_P m_J | ROHF -> let m_P_a = of_some data.density_a in let m_P_b = of_some data.density_b in let fock_a = of_some data.fock_a in let fock_b = of_some data.fock_b in let m_J_a = Fock.coulomb fock_a in let m_J_b = Fock.coulomb fock_b in 0.5 *. ( (Mat.gemm_trace m_P_a m_J_a) +. (Mat.gemm_trace m_P_b m_J_b) ) | _ -> failwith "Not implemented" let exchange_energy t = let data = last_iteration t in match kind t with | RHF -> let m_P = of_some data.density in let fock = of_some data.fock in let m_K = Fock.exchange fock in 0.5 *. Mat.gemm_trace m_P m_K | ROHF -> let m_P_a = of_some data.density_a in let m_P_b = of_some data.density_b in let fock_a = of_some data.fock_a in let fock_b = of_some data.fock_b in let m_K_a = Fock.exchange fock_a in let m_K_b = Fock.exchange fock_b in 0.5 *. ( (Mat.gemm_trace m_P_a m_K_a) +. (Mat.gemm_trace m_P_b m_K_b) ) | _ -> failwith "Not implemented" let make ?kind ?guess:(guess=`Huckel) ?max_scf:(max_scf=64) ?level_shift:(level_shift=0.1) ?threshold_SCF:(threshold_SCF=1.e-8) simulation = (* Number of occupied MOs *) let n_alfa, n_beta = El.n_alfa @@ Si.electrons simulation, El.n_beta @@ Si.electrons simulation in let nocc = n_alfa in let kind = match kind with | Some kind -> kind | None -> if (n_alfa = n_beta) then RHF else ROHF in let nuclear_repulsion = Si.nuclear_repulsion simulation in let ao_basis = Si.ao_basis simulation in (* Initial guess *) let guess = Guess.make ~nocc ~guess ao_basis in (* Orthogonalization matrix *) let m_X = Ao.ortho ao_basis in (* Overlap matrix *) let m_S = Ao.overlap ao_basis |> Ov.matrix in (* Level shift in MO basis *) let m_LSmo = Array.init (Mat.dim2 m_X) (fun i -> if i > nocc then level_shift else 0.) |> Vec.of_array |> Mat.of_diag in (* Guess coefficients *) let m_C = let m_H = match guess with | Guess.Hcore m_H -> m_H | Guess.Huckel m_H -> m_H in let m_Hmo = xt_o_x m_H m_X in let m_C', _ = diagonalize_symm m_Hmo in gemm m_X m_C' in (* A single SCF iteration *) let scf_iteration_rhf data = let nSCF = data.iteration + 1 and m_C = of_some data.coefficients and m_P_prev = data.density and fock_prev = data.fock and diis = match data.diis with | Some diis -> diis | None -> DIIS.make () and threshold = match data.error with | Some error -> error | None -> threshold_SCF *. 2. in (* Density matrix over nocc occupied MOs *) let m_P = gemm ~alpha:2. ~transb:`T ~k:nocc m_C m_C in (* Fock matrix in AO basis *) let fock = match fock_prev, m_P_prev, threshold > 100. *. threshold_SCF with | Some fock_prev, Some m_P_prev, true -> let threshold = 1.e-8 in Fock.make_rhf ~density:(Mat.sub m_P m_P_prev) ~threshold ao_basis |> Fock.add fock_prev | _ -> Fock.make_rhf ~density:m_P ao_basis in let m_F, m_Hc, m_J, m_K = let x = fock in Fock.(fock x, core x, coulomb x, exchange x) in (* Add level shift in AO basis *) let m_F = let m_SC = gemm m_S m_C in gemm m_SC (gemm m_LSmo m_SC ~transb:`T) |> Mat.add m_F in (* Fock matrix in orthogonal basis *) let m_F_ortho = xt_o_x m_F m_X in let error_fock = let fps = gemm m_F (gemm m_P m_S) and spf = gemm m_S (gemm m_P m_F) in xt_o_x (Mat.sub fps spf) m_X in let diis = DIIS.append ~p:(Mat.as_vec m_F_ortho) ~e:(Mat.as_vec error_fock) diis in let m_F_diis = let x = Bigarray.genarray_of_array1 (DIIS.next diis) in Bigarray.reshape_2 x (Mat.dim1 m_F_ortho) (Mat.dim2 m_F_ortho) in (* MOs in orthogonal MO basis *) let m_C', eigenvalues = diagonalize_symm m_F_diis in (* MOs in AO basis *) let m_C = gemm m_X m_C' in (* Hartree-Fock energy *) let energy = nuclear_repulsion +. 0.5 *. Mat.gemm_trace m_P (Mat.add m_Hc m_F) in (* Convergence criterion *) let error = error_fock |> Mat.as_vec |> amax |> abs_float in { empty with iteration = nSCF ; eigenvalues = Some eigenvalues ; coefficients = Some m_C ; error = Some error ; diis = Some diis ; energy = Some energy ; density = Some m_P ; fock = Some fock ; } in let scf_iteration_rohf data = let nSCF = data.iteration + 1 and m_C = of_some data.coefficients and m_P_a_prev = data.density_a and m_P_b_prev = data.density_b and fock_a_prev = data.fock_a and fock_b_prev = data.fock_b and diis = match data.diis with | Some diis -> diis | None -> DIIS.make () and threshold = match data.error with | Some error -> error | None -> threshold_SCF *. 2. in (* Density matrix *) let m_P_a = gemm ~alpha:1. ~transb:`T ~k:n_alfa m_C m_C in let m_P_b = gemm ~alpha:1. ~transb:`T ~k:n_beta m_C m_C in let m_P = Mat.add m_P_a m_P_b in (* Fock matrix in AO basis *) let fock_a = match fock_a_prev, threshold > 100. *. threshold_SCF with | Some fock_a_prev, true -> let threshold = 1.e-8 in Fock.make_uhf ~density_same:(Mat.sub m_P_a @@ of_some m_P_a_prev) ~density_other:(Mat.sub m_P_b @@ of_some m_P_b_prev) ~threshold ao_basis |> Fock.add fock_a_prev | _ -> Fock.make_uhf ~density_same:m_P_a ~density_other:m_P_b ao_basis in let fock_b = match fock_b_prev, threshold > 100. *. threshold_SCF with | Some fock_b_prev, true -> let threshold = 1.e-8 in Fock.make_uhf ~density_same:(Mat.sub m_P_b @@ of_some m_P_b_prev) ~density_other:(Mat.sub m_P_a @@ of_some m_P_a_prev) ~threshold ao_basis |> Fock.add fock_b_prev | _ -> Fock.make_uhf ~density_same:m_P_b ~density_other:m_P_a ao_basis in let m_F_a, m_Hc_a, m_J_a, m_K_a = let x = fock_a in Fock.(fock x, core x, coulomb x, exchange x) in let m_F_b, m_Hc_b, m_J_b, m_K_b = let x = fock_b in Fock.(fock x, core x, coulomb x, exchange x) in let m_F_mo_a = xt_o_x ~o:m_F_a ~x:m_C in let m_F_mo_b = xt_o_x ~o:m_F_b ~x:m_C in let m_F_mo = let space k = if k <= n_beta then `Core else if k <= n_alfa then `Active else `Virtual in Array.init (Mat.dim2 m_F_mo_a) (fun i -> let i = i+1 in Array.init (Mat.dim1 m_F_mo_a) (fun j -> let j = j+1 in match (space i), (space j) with | `Core , `Core -> 0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) -. (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j}) | `Active , `Core | `Core , `Active -> m_F_mo_b.{i,j} | `Core , `Virtual | `Virtual , `Core | `Active , `Active -> 0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) | `Virtual , `Active | `Active , `Virtual -> m_F_mo_a.{i,j} | `Virtual , `Virtual -> 0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) +. (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j}) ) ) |> Mat.of_array in let m_SC = gemm m_S m_C in let m_F = x_o_xt ~x:m_SC ~o:m_F_mo in (* Add level shift in AO basis *) let m_F = x_o_xt ~x:m_SC ~o:m_LSmo |> Mat.add m_F in (* Fock matrix in orthogonal basis *) let m_F_ortho = xt_o_x m_F m_X in let error_fock = let fps = gemm m_F (gemm m_P m_S) and spf = gemm m_S (gemm m_P m_F) in xt_o_x (Mat.sub fps spf) m_X in let diis = DIIS.append ~p:(Mat.as_vec m_F_ortho) ~e:(Mat.as_vec error_fock) diis in let m_F_diis = let x = Bigarray.genarray_of_array1 (DIIS.next diis) in Bigarray.reshape_2 x (Mat.dim1 m_F_ortho) (Mat.dim2 m_F_ortho) in (* MOs in orthogonal MO basis *) let m_C', eigenvalues = diagonalize_symm m_F_diis in (* MOs in AO basis *) let m_C = gemm m_X m_C' in (* Hartree-Fock energy *) let energy = nuclear_repulsion +. 0.5 *. ( Mat.gemm_trace m_P_a (Mat.add m_Hc_a m_F_a) +. Mat.gemm_trace m_P_b (Mat.add m_Hc_b m_F_b) ) in (* Convergence criterion *) let error = error_fock |> Mat.as_vec |> amax |> abs_float in { empty with iteration = nSCF ; eigenvalues = Some eigenvalues ; coefficients = Some m_C ; error = Some error ; diis = Some diis ; energy = Some energy ; density_a = Some m_P_a ; density_b = Some m_P_b ; fock_a = Some fock_a ; fock_b = Some fock_b ; } in let scf_iteration data = match kind with | RHF -> scf_iteration_rhf data | ROHF -> scf_iteration_rohf data | _ -> failwith "Not implemented" in let array_data = let storage = Array.make max_scf None in let rec iteration = function | 0 -> Some (scf_iteration { empty with coefficients = Some m_C }) | n -> begin match storage.(n) with | Some result -> Some result | None -> begin let data = (iteration (n-1)) in match data with | None -> None | Some data -> begin (** Check convergence *) let converged, error = match data.error with | None -> false, 0. | Some error -> (data.iteration = max_scf || error < threshold_SCF), error in if converged then None else begin storage.(n-1) <- Some { empty with iteration = data.iteration; energy = data.energy ; eigenvalues = data.eigenvalues ; error = data.error ; }; storage.(n) <- Some (scf_iteration data); storage.(n); end end end end in Array.init max_scf (fun i -> lazy (iteration i)) in { kind; simulation; guess ; data = array_data; nocc; } let linewidth = 60 let pp_iterations ppf t = Format.fprintf ppf "@[%4s%s@]@." "" (Printing.line linewidth); Format.fprintf ppf "@[%4s@[%5s@]@,@[%16s@]@,@[%16s@]@,@[%11s@]@]@." "" "#" "HF energy " "Convergence" "HOMO-LUMO"; Format.fprintf ppf "@[%4s%s@]@." "" (Printing.line linewidth); let nocc = nocc t in Array.iter (fun data -> let data = Lazy.force data in match data with | None -> () | Some data -> let e = of_some data.eigenvalues in let gap = e.{nocc+1} -. e.{nocc} in begin Format.fprintf ppf "@[%4s@[%5d@]@,@[%16.8f@]@,@[%16.4e@]@,@[%11.4f@]@]@." "" (data.iteration) (of_some data.energy) (of_some data.error) gap; end ) t.data; Format.fprintf ppf "@[%4s%s@]@." "" (Printing.line linewidth) let pp_summary ppf t = let print text value = Format.fprintf ppf "@[@[%30s@]@,@[%16.10f@]@]@;" text value; and line () = Format.fprintf ppf "@[ %s @]@;" (Printing.line (linewidth-4)); in let ha_to_ev = Constants.ha_to_ev in let e = eigenvalues t in Format.fprintf ppf "@[%s@]@;" (Printing.line ~c:'=' linewidth); Format.fprintf ppf "@["; print "One-electron energy" (kin_energy t +. eN_energy t) ; print "Kinetic" (kin_energy t) ; print "Potential" (eN_energy t) ; line () ; print "Two-electron energy" (coulomb_energy t +. exchange_energy t) ; print "Coulomb" (coulomb_energy t) ; print "Exchange" (exchange_energy t) ; line (); print "HF HOMO" (ha_to_ev *. e.{nocc t}); print "HF LUMO" (ha_to_ev *. e.{nocc t + 1}); print "HF LUMO-LUMO" (ha_to_ev *. (e.{nocc t + 1} -. e.{nocc t })); line (); print "Electronic energy" (energy t -. nuclear_repulsion t) ; print "Nuclear repulsion" (nuclear_repulsion t) ; print "Hartree-Fock energy" (energy t) ; Format.fprintf ppf "@]"; Format.fprintf ppf "@[%s@]@;" (Printing.line ~c:'=' linewidth) let pp_hf ppf t = Format.fprintf ppf "@.@[%s@]@." (Printing.line ~c:'=' 70); Format.fprintf ppf "@[%34s %-34s@]@." (match t.kind with | UHF -> "Unrestricted" | RHF -> "Restricted" | ROHF -> "Restricted Open-shell") "Hartree-Fock"; Format.fprintf ppf "@[%s@]@.@." (Printing.line ~c:'=' 70); Format.fprintf ppf "@[%a@]@." pp_iterations t; Format.fprintf ppf "@[@;%a@]@." pp_summary t