open Util open Constants open Lacaml.D module Si = Simulation module El = Electrons module Ao = AOBasis module Ov = Overlap let make ~guess ~max_scf ~level_shift ~threshold_SCF simulation = (* Number of occupied MOs *) let n_alfa = El.n_alfa @@ Si.electrons simulation in let nocc = n_alfa in let n_beta = El.n_beta @@ Si.electrons simulation 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 let m_T = Ao.kin_ints ao_basis |> KinInt.matrix and m_V = Ao.eN_ints ao_basis |> NucInt.matrix in (* Level shift in MO basis *) let m_LSmo = Array.init (Mat.dim2 m_X) (fun i -> if i > n_alfa then level_shift else 0.) |> Vec.of_array |> Mat.of_diag in (* SCF iterations *) let rec loop nSCF iterations energy_prev m_C m_P_a_prev m_P_b_prev fock_a_prev fock_b_prev threshold diis = (* 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 m_P_a_prev) ~density_other:(Mat.sub m_P_b 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 m_P_b_prev) ~density_other:(Mat.sub m_P_a 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 -> (* 0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) +. 0.5 *. (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j}) *) 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 -> (* 0.5 *. (m_F_mo_a.{i,j} +. m_F_mo_b.{i,j}) -. 0.5 *. (m_F_mo_b.{i,j} -. m_F_mo_a.{i,j}) *) 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 let converged = nSCF = max_scf || error < threshold_SCF in let gap = if nocc < Vec.dim eigenvalues then eigenvalues.{nocc+1} -. eigenvalues.{nocc} else 0. in let () = match energy_prev with | Some energy_prev -> Printf.eprintf "%3d %16.10f %16.10f %11.4e %10.4f\n%!" nSCF energy (energy -. energy_prev) error gap | None -> Printf.eprintf "%3d %16.10f %16s %11.4e %10.4f\n%!" nSCF energy "" error gap in if not converged then loop (nSCF+1) ( (energy, error, gap) :: iterations) (Some energy) m_C m_P_a m_P_b (Some fock_a) (Some fock_b) error diis else let iterations = List.rev ( (energy, error, gap) :: iterations ) |> Array.of_list in HartreeFock_type.(ROHF { simulation; nocc; guess ; eigenvectors = m_C ; eigenvalues ; energy ; nuclear_repulsion; iterations ; kin_energy = Mat.gemm_trace m_P m_T; eN_energy = Mat.gemm_trace m_P m_V; coulomb_energy = 0.5 *. (Mat.gemm_trace m_P_a m_J_a) +. 0.5 *. (Mat.gemm_trace m_P_b m_J_b); exchange_energy = 0.5 *. (Mat.gemm_trace m_P_a m_K_a) +. 0.5 *. (Mat.gemm_trace m_P_b m_K_b); occupation = Mat.copy_diag m_P; }) in (* Guess coefficients *) 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 let m_C = gemm m_X m_C' in let diis = DIIS.make () in loop 1 [] None m_C m_C m_C None None threshold_SCF diis