open Util open Lacaml.D open Simulation let make ?guess:(guess=`Huckel) ?max_scf:(max_scf=64) ?level_shift:(level_shift=0.1) ?threshold_SCF:(threshold_SCF=1.e-5) simulation = (* Number of occupied MOs *) let nocc = simulation.electrons.Electrons.n_alpha in let nuclear_repulsion = simulation.nuclear_repulsion in (* Initial guess *) let guess = Guess.make ~guess simulation in (* Orthogonalization matrix *) let m_X = Lazy.force simulation.overlap_ortho in (* Overlap matrix *) let m_S = Lazy.force simulation.overlap in let m_T = Lazy.force simulation.kin_ints and m_V = Lazy.force simulation.eN_ints 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 (* SCF iterations *) let rec loop nSCF iterations energy_prev m_C diis = (* 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 m_F, m_Hc, m_J, m_K = let x = Fock.make ~density:m_P simulation in x.Fock.fock, x.Fock.core, x.Fock.coulomb, x.Fock.exchange 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 let converged = nSCF = max_scf || error < threshold_SCF in let gap = eigenvalues.{nocc+1} -. eigenvalues.{nocc}; 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 diis else let iterations = List.rev ( (energy, error, gap) :: iterations ) |> Array.of_list in { HartreeFock_type. 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 m_J; exchange_energy = 0.5 *. Mat.gemm_trace m_P m_K; } 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 diis