Skip to content

DIIS #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jwaldrop107 merged 6 commits into from
Jun 24, 2025
Merged

DIIS #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
17f4d0d
SCF loop rewritten with DIIS option
jwaldrop107 Jun 20, 2025
ba9fead
DIIS acceleration working, but not compatible with uncertain values
jwaldrop107 Jun 20, 2025
1649354
remove floating point type restrictions from tests
jwaldrop107 Jun 23, 2025
4d44ed1
Committing clang-format changes
github-actions[bot] Jun 23, 2025
81c23c1
DIIS default on
jwaldrop107 Jun 24, 2025
c5d1bc1
update test settings for DIIS default
jwaldrop107 Jun 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/scf/driver/driver.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
224 changes: 128 additions & 96 deletions src/scf/driver/scf_loop.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>;
Expand All @@ -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;
Expand All @@ -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();
Expand All @@ -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;
Expand Down
Loading