DIIS (#29)

* SCF loop rewritten with DIIS option

* DIIS acceleration working, but not compatible with uncertain values

* remove floating point type restrictions from tests

* Committing clang-format changes

* DIIS default on

Co-authored-by: Ryan Richard <ryanmrichard1@gmail.com>

* update test settings for DIIS default

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Ryan Richard <ryanmrichard1@gmail.com>

Author: 3 people

src/scf/driver/driver.hpp

@@ -30,7 +30,9 @@ inline void load_modules(pluginplay::ModuleManager& mm) {
 inline void set_defaults(pluginplay::ModuleManager& mm) {
     mm.change_submod("Loop", "Electronic energy", "Electronic energy");
     mm.change_submod("Loop", "Density matrix", "Density matrix builder");
-    mm.change_submod("Loop", "Guess update", "Diagonalization Fock update");
+    mm.change_submod("Loop", "Diagonalizer",
+                     "Generalized eigensolve via Eigen");
+    mm.change_submod("Loop", "Fock matrix builder", "Fock matrix builder");
     mm.change_submod("Loop", "One-electron Fock operator",
                      "Restricted One-Electron Fock op");
     mm.change_submod("Loop", "Fock operator", "Restricted Fock Op");

src/scf/driver/scf_loop.cpp

@@ -20,36 +20,52 @@
 namespace scf::driver {
 namespace {
 
+// Comparison between convergence metrics based on floating point type
 struct Kernel {
-    explicit Kernel(parallelzone::runtime::RuntimeView rv) : m_rv(rv) {}
+    using rv_t = parallelzone::runtime::RuntimeView;
+
+    explicit Kernel(rv_t rv) : m_rv(rv) {}
+
     template<typename FloatType>
     auto run(const tensorwrapper::buffer::BufferBase& a, double tol) {
         tensorwrapper::allocator::Eigen<FloatType> allocator(m_rv);
         const auto& eigen_a = allocator.rebind(a);
         return tensorwrapper::types::fabs(eigen_a.get_data(0)) < FloatType(tol);
     }
 
-    parallelzone::runtime::RuntimeView m_rv;
+    rv_t m_rv;
+};
+
+// Pull out nuclear-nuclear interaction term, if there is one.
+struct GrabNuclear : chemist::qm_operator::OperatorVisitor {
+    using V_nn_type = simde::type::V_nn_type;
+
+    GrabNuclear() : chemist::qm_operator::OperatorVisitor(false) {}
+
+    void run(const V_nn_type& V_nn) { m_pv = &V_nn; }
+
+    const V_nn_type* m_pv = nullptr;
 };
 
 const auto desc = R"(
 )";
 
 } // namespace
 
+using tensor_t        = simde::type::tensor;
+using cmos_t          = simde::type::cmos;
+using diis_t          = tensorwrapper::diis::DIIS;
+using density_t       = simde::type::decomposable_e_density;
+using fock_pt         = simde::FockOperator<density_t>;
+using density_pt      = simde::aos_rho_e_aos<cmos_t>;
+using v_nn_pt         = simde::charge_charge_interaction;
+using fock_matrix_pt  = simde::aos_f_e_aos;
+using diagonalizer_pt = simde::GeneralizedEigenSolve;
+using s_pt            = simde::aos_s_e_aos;
 using simde::type::electronic_hamiltonian;
-using simde::type::hamiltonian;
-using simde::type::op_base_type;
-
-using density_t      = simde::type::decomposable_e_density;
-using fock_pt        = simde::FockOperator<density_t>;
-using density_pt     = simde::aos_rho_e_aos<simde::type::cmos>;
-using v_nn_pt        = simde::charge_charge_interaction;
-using fock_matrix_pt = simde::aos_f_e_aos;
-using s_pt           = simde::aos_s_e_aos;
 
 template<typename WfType>
-using egy_pt = simde::eval_braket<WfType, hamiltonian, WfType>;
+using egy_pt = simde::eval_braket<WfType, simde::type::hamiltonian, WfType>;
 
 template<typename WfType>
 using elec_egy_pt = simde::eval_braket<WfType, electronic_hamiltonian, WfType>;
@@ -60,15 +76,7 @@ using pt = simde::Optimize<egy_pt<WfType>, WfType>;
 template<typename WfType>
 using update_pt = simde::UpdateGuess<WfType>;
 
-struct GrabNuclear : chemist::qm_operator::OperatorVisitor {
-    using V_nn_type = simde::type::V_nn_type;
-
-    GrabNuclear() : chemist::qm_operator::OperatorVisitor(false) {}
-
-    void run(const V_nn_type& V_nn) { m_pv = &V_nn; }
-
-    const V_nn_type* m_pv = nullptr;
-};
+using tensorwrapper::utilities::floating_point_dispatch;
 
 MODULE_CTOR(SCFLoop) {
     using wf_type = simde::type::rscf_wf;
@@ -81,47 +89,59 @@ MODULE_CTOR(SCFLoop) {
     add_input<double>("density tolerance").set_default(1.0E-6);
     add_input<double>("gradient tolerance").set_default(1.0E-6);
 
+    const std::size_t diis_sample_default = 5;
+    add_input<bool>("DIIS").set_default(true);
+    add_input<std::size_t>("DIIS max samples").set_default(diis_sample_default);
+
     add_submodule<elec_egy_pt<wf_type>>("Electronic energy");
     add_submodule<density_pt>("Density matrix");
-    add_submodule<update_pt<wf_type>>("Guess update");
+    add_submodule<s_pt>("Overlap matrix builder");
+    add_submodule<fock_matrix_pt>("Fock matrix builder");
+    add_submodule<diagonalizer_pt>("Diagonalizer");
     add_submodule<fock_pt>("One-electron Fock operator");
     add_submodule<fock_pt>("Fock operator");
-    add_submodule<fock_matrix_pt>("Fock matrix builder");
     add_submodule<v_nn_pt>("Charge-charge");
-    add_submodule<s_pt>("Overlap matrix builder");
 }
 
 MODULE_RUN(SCFLoop) {
-    using wf_type               = simde::type::rscf_wf;
-    using density_op_type       = simde::type::rho_e<simde::type::cmos>;
+    using wf_type         = simde::type::rscf_wf;
+    using density_op_type = simde::type::rho_e<cmos_t>;
+
     const auto&& [braket, psi0] = pt<wf_type>::unwrap_inputs(inputs);
-    // TODO: Assert bra == ket == psi0
-    const auto& H       = braket.op();
-    const auto& H_elec  = H.electronic_hamiltonian();
-    const auto& H_core  = H_elec.core_hamiltonian();
-    const auto& aos     = psi0.orbitals().from_space();
     const auto max_iter = inputs.at("max iterations").value<unsigned int>();
     const auto e_tol    = inputs.at("energy tolerance").value<double>();
     const auto dp_tol   = inputs.at("density tolerance").value<double>();
     const auto g_tol    = inputs.at("gradient tolerance").value<double>();
 
-    auto& egy_mod     = submods.at("Electronic energy");
-    auto& density_mod = submods.at("Density matrix");
-    auto& update_mod  = submods.at("Guess update");
-    auto& fock_mod    = submods.at("One-electron Fock operator");
-    auto& Fock_mod    = submods.at("Fock operator");
-    auto& V_nn_mod    = submods.at("Charge-charge");
-    // TODO: should be split off into orbital gradient module
-    auto& F_mod = submods.at("Fock matrix builder");
-    auto& S_mod = submods.at("Overlap matrix builder");
-
-    // Step 1: Nuclear-nuclear repulsion
+    auto& egy_mod          = submods.at("Electronic energy");
+    auto& density_mod      = submods.at("Density matrix");
+    auto& diagonalizer_mod = submods.at("Diagonalizer");
+    auto& fock_mod         = submods.at("One-electron Fock operator");
+    auto& Fock_mod         = submods.at("Fock operator");
+    auto& V_nn_mod         = submods.at("Charge-charge");
+    auto& F_mod            = submods.at("Fock matrix builder");
+    auto& S_mod            = submods.at("Overlap matrix builder");
+
+    // Unwrap Braket and trial wavefunction
+    // TODO: Assert bra == ket == psi0
+    const auto& H      = braket.op();
+    const auto& H_elec = H.electronic_hamiltonian();
+    const auto& H_core = H_elec.core_hamiltonian();
+    const auto& aos    = psi0.orbitals().from_space();
+
+    // DIIS settings
+    const auto diis_on = inputs.at("DIIS").value<bool>();
+    const auto diis_max_samples =
+      inputs.at("DIIS max samples").value<std::size_t>();
+    diis_t diis(diis_max_samples);
+
+    // Nuclear-nuclear repulsion
     GrabNuclear visitor;
     H.visit(visitor);
     bool has_nn = (visitor.m_pv != nullptr);
 
     // TODO: Clean up charges class to make this easier...
-    simde::type::tensor e_nuclear(0.0);
+    tensor_t e_nuclear(0.0);
     if(has_nn) {
         const auto& V_nn       = *visitor.m_pv;
         const auto n_lhs       = V_nn.lhs_particle().as_nuclei();
@@ -143,109 +163,121 @@ MODULE_RUN(SCFLoop) {
     chemist::braket::BraKet s_mn(aos, simde::type::s_e_type{}, aos);
     const auto& S = S_mod.run_as<s_pt>(s_mn);
 
-    wf_type psi_old = psi0;
-    simde::type::tensor e_old;
-
-    density_op_type rho_hat(psi_old.orbitals(), psi_old.occupations());
-    chemist::braket::BraKet P_mn(aos, rho_hat, aos);
-    const auto& P = density_mod.run_as<density_pt>(P_mn);
-    density_t rho_old(P, psi_old.orbitals());
-
-    auto& logger = get_runtime().logger();
+    // For convergence checking
+    wf_type psi_old;
+    density_t rho_old;
+    tensor_t e_old;
+    tensor_t F_old;
 
+    // Initialize loop
     unsigned int iter = 0;
+    auto& logger      = get_runtime().logger();
     while(iter < max_iter) {
-        // Step 2: Old density is used to create the corresponding Fock operator
-        // TODO: Make easier to go from many-electron to one-electron
-        // TODO: template fock_pt on Hamiltonian type and only pass H_elec
-        const auto& f_old = fock_mod.run_as<fock_pt>(H, rho_old);
-
-        // Step 3: Fock operator used to compute new wavefunction/density
-        const auto& psi_new =
-          update_mod.run_as<update_pt<wf_type>>(f_old, psi_old);
+        // Step 1: Generate trial wavefunction
+        wf_type psi;
+        if(iter > 0) {
+            // Diagonalize the Fock matrix
+            const auto&& [evalues, evectors] =
+              diagonalizer_mod.run_as<diagonalizer_pt>(F_old, S);
+
+            // Construct new trial wavefunction
+            cmos_t cmos(evalues, aos, evectors);
+            psi = wf_type(psi_old.orbital_indices(), cmos);
+        } else {
+            // Use trial wavefunction provided from initial guess
+            psi = psi0;
+        }
 
-        density_op_type rho_hat_new(psi_new.orbitals(), psi_new.occupations());
-        chemist::braket::BraKet P_mn_new(aos, rho_hat_new, aos);
-        const auto& P_new = density_mod.run_as<density_pt>(P_mn_new);
+        // Step 2: Construct electronic density
+        density_op_type rho_hat(psi.orbitals(), psi.occupations());
+        chemist::braket::BraKet P_mn(aos, rho_hat, aos);
+        const auto& P = density_mod.run_as<density_pt>(P_mn);
+        density_t rho(P, psi.orbitals());
 
-        density_t rho_new(P_new, psi_new.orbitals());
+        // Step 3: Construct new Fock matrix
+        const auto& f_hat = fock_mod.run_as<fock_pt>(H, rho);
+        chemist::braket::BraKet f_mn(aos, f_hat, aos);
+        auto F = F_mod.run_as<fock_matrix_pt>(f_mn);
 
         // Step 4: New electronic energy
         // Step 4a: New Fock operator to new electronic Hamiltonian
-        // TODO: Should just be H_core + F;
-        const auto& F_new = Fock_mod.run_as<fock_pt>(H, rho_new);
+        // TODO: Should just be H_core + F_hat;
+        const auto& F_hat = Fock_mod.run_as<fock_pt>(H, rho);
         electronic_hamiltonian H_new;
         for(std::size_t i = 0; i < H_core.size(); ++i)
             H_new.emplace_back(H_core.coefficient(i),
                                H_core.get_operator(i).clone());
-        for(std::size_t i = 0; i < F_new.size(); ++i)
-            H_new.emplace_back(F_new.coefficient(i),
-                               F_new.get_operator(i).clone());
+        for(std::size_t i = 0; i < F_hat.size(); ++i)
+            H_new.emplace_back(F_hat.coefficient(i),
+                               F_hat.get_operator(i).clone());
 
         // Step 4b: New electronic hamiltonian to new electronic energy
-        chemist::braket::BraKet H_00(psi_new, H_new, psi_new);
-        auto e_new = egy_mod.run_as<elec_egy_pt<wf_type>>(H_00);
+        chemist::braket::BraKet H_00(psi, H_new, psi);
+        auto e = egy_mod.run_as<elec_egy_pt<wf_type>>(H_00);
 
         auto e_msg = "SCF iteration = " + std::to_string(iter) + ":";
-        e_msg += "  Electronic Energy = " + e_new.to_string();
+        e_msg += "  Electronic Energy = " + e.to_string();
         logger.log(e_msg);
 
-        bool converged = false;
         // Step 5: Converged?
+        bool converged = false;
         if(iter > 0) {
             // Change in the energy
-            simde::type::tensor de;
-            de("") = e_new("") - e_old("");
+            tensor_t de;
+            de("") = e("") - e_old("");
 
             // Change in the density
-            simde::type::tensor dp;
+            tensor_t dp;
             const auto& P_old = rho_old.value();
-            dp("m,n")         = rho_new.value()("m,n") - P_old("m,n");
+            dp("m,n")         = rho.value()("m,n") - P_old("m,n");
             auto dp_norm      = tensorwrapper::operations::infinity_norm(dp);
 
             // Orbital gradient: FPS-SPF
-            // TODO: module satisfying BraKet(aos, Commutator(F,P), aos)
-            const auto& f_new = fock_mod.run_as<fock_pt>(H, rho_new);
-            chemist::braket::BraKet F_mn(aos, f_new, aos);
-            const auto& F_matrix = F_mod.run_as<fock_matrix_pt>(F_mn);
-
-            auto grad = commutator(F_matrix, P_new, S);
-
-            simde::type::tensor grad_norm;
+            auto grad = commutator(F, P, S);
+            tensor_t grad_norm;
             grad_norm("") = grad("m,n") * grad("n,m");
 
-            Kernel k(get_runtime());
+            // Log convergence metrics
+            logger.log("  dE = " + de.to_string());
+            logger.log("  dP = " + dp_norm.to_string());
+            logger.log("  dG = " + grad_norm.to_string());
 
-            using tensorwrapper::utilities::floating_point_dispatch;
+            // Check for convergence
+            Kernel k(get_runtime());
             auto e_conv = floating_point_dispatch(k, de.buffer(), e_tol);
             auto g_conv = floating_point_dispatch(k, grad_norm.buffer(), g_tol);
             auto dp_conv = floating_point_dispatch(k, dp_norm.buffer(), dp_tol);
+            if(e_conv && g_conv && dp_conv) converged = true;
 
-            logger.log("  dE = " + de.to_string());
-            logger.log("  dP = " + dp_norm.to_string());
-            logger.log("  dG = " + grad_norm.to_string());
+            // If using DIIS and not converged, extrapolate new Fock matrix
+            if(diis_on && !converged) { F = diis.extrapolate(F, grad); }
+        } else if(diis_on) {
+            // For DIIS, still need to the orbital gradient
+            auto grad = commutator(F, P, S);
 
-            if(e_conv && g_conv && dp_conv) converged = true;
+            // If using DIIS, extrapolate new Fock matrix
+            F = diis.extrapolate(F, grad);
         }
 
         // Step 6: Not converged so reset
-        e_old   = e_new;
-        psi_old = psi_new;
-        rho_old = rho_new;
+        e_old   = e;
+        psi_old = psi;
+        rho_old = rho;
+        F_old   = F;
         if(converged) break;
         ++iter;
     }
     if(iter == max_iter) throw std::runtime_error("SCF failed to converge");
 
-    simde::type::tensor e_total;
+    tensor_t e_total;
 
     // e_nuclear is a double. This hack converts it to udouble (if needed)
     tensorwrapper::allocator::Eigen<double> dalloc(get_runtime());
     using tensorwrapper::types::udouble;
     tensorwrapper::allocator::Eigen<udouble> ualloc(get_runtime());
 
     if(ualloc.can_rebind(e_old.buffer())) {
-        simde::type::tensor temp(e_old);
+        tensor_t temp(e_old);
         auto val = dalloc.rebind(e_nuclear.buffer()).get_data(0);
         ualloc.rebind(temp.buffer()).set_data(0, val);
         e_nuclear = temp;