Skip to content

Commit

Permalink
QED HF
Browse files Browse the repository at this point in the history
  • Loading branch information
dmejiar authored and ajaypanyala committed Oct 27, 2024
1 parent 69b8228 commit 6a5a6fc
Show file tree
Hide file tree
Showing 8 changed files with 500 additions and 12 deletions.
2 changes: 2 additions & 0 deletions exachem/scf/scf.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ include(TargetMacros)
set(SCF_SRCDIR scf)

set(SCF_INCLUDES
${SCF_SRCDIR}/scf_qed.hpp
${SCF_SRCDIR}/scf_vars.hpp
${SCF_SRCDIR}/scf_iter.hpp
${SCF_SRCDIR}/scf_main.hpp
Expand All @@ -21,6 +22,7 @@ set(SCF_INCLUDES
)

set(SCF_SRCS
${SCF_SRCDIR}/scf_qed.cpp
${SCF_SRCDIR}/scf_iter.cpp
${SCF_SRCDIR}/scf_main.cpp
${SCF_SRCDIR}/scf_gauxc.cpp
Expand Down
55 changes: 55 additions & 0 deletions exachem/scf/scf_hartree_fock.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ void exachem::scf::SCFHartreeFock::initialize(ExecutionContext& exc, ChemEnv& ch

void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_env) {
SCFCompute scf_compute;
SCFQed scf_qed;
SCFIter scf_iter;
SCFGuess scf_guess;
SCFRestart scf_restart;
Expand Down Expand Up @@ -336,6 +337,27 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
/*** compute 1-e integrals ***/
/*** =========================== ***/
scf_compute.compute_hamiltonian<TensorType>(ec, scf_vars, chem_env, ttensors, etensors);
if(chem_env.sys_data.is_qed) {
ttensors.QED_Dx = {tAO, tAO};
ttensors.QED_Dy = {tAO, tAO};
ttensors.QED_Dz = {tAO, tAO};
ttensors.QED_Qxx = {tAO, tAO};
ttensors.QED_Qxy = {tAO, tAO};
ttensors.QED_Qxz = {tAO, tAO};
ttensors.QED_Qyy = {tAO, tAO};
ttensors.QED_Qyz = {tAO, tAO};
ttensors.QED_Qzz = {tAO, tAO};
ttensors.QED_1body = {tAO, tAO};
ttensors.QED_2body = {tAO, tAO};
Tensor<TensorType>::allocate(&ec, ttensors.QED_Dx, ttensors.QED_Dy, ttensors.QED_Dz,
ttensors.QED_Qxx, ttensors.QED_Qxy, ttensors.QED_Qxz,
ttensors.QED_Qyy, ttensors.QED_Qyz, ttensors.QED_Qzz,
ttensors.QED_1body, ttensors.QED_2body);

scf_qed.compute_qed_emult_ints<TensorType>(ec, chem_env, scf_vars, ttensors);
if(chem_env.sys_data.do_qed)
scf_qed.compute_QED_1body<TensorType>(ec, chem_env, scf_vars, ttensors);
}

if(chem_env.sys_data.has_ecp) {
Tensor<TensorType> ECP{tAO, tAO};
Expand Down Expand Up @@ -653,6 +675,11 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
shell2bf, SchwarzK, max_nprim4, ttensors, etensors,
is_3c_init, do_density_fitting, xHF);

// Add QED contribution
if(chem_env.sys_data.do_qed) {
scf_qed.compute_QED_2body<TensorType>(ec, chem_env, scf_vars, ttensors);
}

#if defined(USE_GAUXC)
// Add snK contribution
if(chem_env.sys_data.do_snK) {
Expand Down Expand Up @@ -707,6 +734,7 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
auto debug = chem_env.ioptions.scf_options.debug;
if(rank == 0 && debug)
std::cout << std::fixed << std::setprecision(2) << "xcf: " << xcf_time << "s, ";
if(sys_data.is_qed && !sys_data.do_qed) { scf_vars.eqed = gauxc_exc; }
}

ehf += gauxc_exc;
Expand Down Expand Up @@ -770,6 +798,11 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
shell2bf, SchwarzK, max_nprim4, ttensors, etensors,
is_3c_init, do_density_fitting, xHF);

// Add QED contribution
if(chem_env.sys_data.do_qed) {
scf_qed.compute_QED_2body<TensorType>(ec, chem_env, scf_vars, ttensors);
}

std::tie(ehf, rmsd) = scf_iter.scf_iter_body<TensorType>(ec, chem_env, scalapack_info, iter,
scf_vars, ttensors, etensors
#if defined(USE_GAUXC)
Expand Down Expand Up @@ -920,6 +953,13 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
scf_iter.compute_2bf<TensorType>(ec, chem_env, scalapack_info, scf_vars, do_schwarz_screen,
shell2bf, SchwarzK, max_nprim4, ttensors, etensors,
is_3c_init, do_density_fitting, xHF_adjust);

// Add QED contribution;
// CHECK

if(chem_env.sys_data.do_qed) {
scf_qed.compute_QED_2body<TensorType>(ec, chem_env, scf_vars, ttensors);
}
}
else if(scf_vars.lshift > 0) {
// Remove level shift from Fock matrix
Expand Down Expand Up @@ -1001,6 +1041,20 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
Tensor<TensorType>::deallocate(ttensors.VXC_alpha);
if(chem_env.sys_data.is_unrestricted) Tensor<TensorType>::deallocate(ttensors.VXC_beta);
}
if(chem_env.sys_data.is_qed) {
scf_output.rw_mat_disk<TensorType>(ttensors.QED_Dx, files_prefix + ".QED_Dx",
chem_env.ioptions.scf_options.debug);
scf_output.rw_mat_disk<TensorType>(ttensors.QED_Dy, files_prefix + ".QED_Dy",
chem_env.ioptions.scf_options.debug);
scf_output.rw_mat_disk<TensorType>(ttensors.QED_Dz, files_prefix + ".QED_Dz",
chem_env.ioptions.scf_options.debug);
scf_output.rw_mat_disk<TensorType>(ttensors.QED_Qxx, files_prefix + ".QED_Qxx",
chem_env.ioptions.scf_options.debug);
Tensor<TensorType>::deallocate(ttensors.QED_Dx, ttensors.QED_Dy, ttensors.QED_Dz,
ttensors.QED_Qxx, ttensors.QED_Qxy, ttensors.QED_Qxz,
ttensors.QED_Qyy, ttensors.QED_Qyz, ttensors.QED_Qzz,
ttensors.QED_1body, ttensors.QED_2body);
}
#if SCF_THROTTLE_RESOURCES
ec.flush_and_sync();
#endif
Expand Down Expand Up @@ -1064,6 +1118,7 @@ void exachem::scf::SCFHartreeFock::scf_hf(ExecutionContext& exc, ChemEnv& chem_e
chem_env.sys_data.update();
if(rank == 0 && chem_env.ioptions.scf_options.debug) chem_env.sys_data.print();

chem_env.sys_data.scf_energy = ehf;
// iter not broadcasted, but fine since only rank 0 writes to json
if(rank == 0) {
chem_env.sys_data.results["output"]["SCF"]["final_energy"] = ehf;
Expand Down
1 change: 1 addition & 0 deletions exachem/scf/scf_hartree_fock.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include "exachem/scf/scf_compute.hpp"
#include "exachem/scf/scf_iter.hpp"
#include "exachem/scf/scf_outputs.hpp"
#include "exachem/scf/scf_qed.hpp"
#include "exachem/scf/scf_restart.hpp"
#include "exachem/scf/scf_taskmap.hpp"
#include <variant>
Expand Down
1 change: 1 addition & 0 deletions exachem/scf/scf_iter.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ std::tuple<TensorType, TensorType> exachem::scf::SCFIter::scf_iter_body(
std::chrono::duration_cast<std::chrono::duration<double>>((xcf_stop - xcf_start)).count();
if(rank == 0 && debug)
std::cout << std::fixed << std::setprecision(2) << "xcf: " << xcf_time << "s, ";
if(sys_data.is_qed && !sys_data.do_qed) { scf_vars.eqed = gauxc_exc; }
}

ehf += gauxc_exc;
Expand Down
82 changes: 74 additions & 8 deletions exachem/scf/scf_outputs.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,26 +26,85 @@ double exachem::scf::SCFIO::tt_trace(ExecutionContext& ec, Tensor<TensorType>& T
void exachem::scf::SCFIO::print_energies(ExecutionContext& ec, ChemEnv& chem_env,
TAMMTensors& ttensors, EigenTensors& etensors,
SCFVars& scf_vars, ScalapackInfo& scalapack_info) {
const SystemData& sys_data = chem_env.sys_data;
const SystemData& sys_data = chem_env.sys_data;
const SCFOptions& scf_options = chem_env.ioptions.scf_options;

const bool is_uhf = sys_data.is_unrestricted;
const bool is_rhf = sys_data.is_restricted;
const bool is_ks = sys_data.is_ks;
const bool is_qed = sys_data.is_qed;
const bool do_qed = sys_data.do_qed;

double nelectrons = 0.0;
double kinetic_1e = 0.0;
double NE_1e = 0.0;
double energy_1e = 0.0;
double energy_2e = 0.0;

double nelectrons = 0.0;
double kinetic_1e = 0.0;
double NE_1e = 0.0;
double energy_1e = 0.0;
double energy_2e = 0.0;
double energy_qed = 0.0;
double energy_qed_et = 0.0;
if(is_rhf) {
nelectrons = tt_trace(ec, ttensors.D_last_alpha, ttensors.S1);
kinetic_1e = tt_trace(ec, ttensors.D_last_alpha, ttensors.T1);
NE_1e = tt_trace(ec, ttensors.D_last_alpha, ttensors.V1);
energy_1e = tt_trace(ec, ttensors.D_last_alpha, ttensors.H1);
energy_2e = 0.5 * tt_trace(ec, ttensors.D_last_alpha, ttensors.F_alpha_tmp);

if(is_ks) { energy_2e += scf_vars.exc; }
if(is_ks) {
if((!is_qed) || (is_qed && do_qed)) { energy_2e += scf_vars.exc; }
}

if(is_qed) {
if(do_qed) {
energy_qed = tt_trace(ec, ttensors.D_last_alpha, ttensors.QED_1body);
energy_qed += 0.5 * tt_trace(ec, ttensors.D_last_alpha, ttensors.QED_2body);
}
else { energy_qed = scf_vars.eqed; }

Tensor<double> X_a;

#if defined(USE_SCALAPACK)
X_a = {scf_vars.tAO, scf_vars.tAO_ortho};
Tensor<double>::allocate(&ec, X_a);
if(scalapack_info.comm != MPI_COMM_NULL) {
tamm::from_block_cyclic_tensor(ttensors.X_alpha, X_a);
}
#else
X_a = ttensors.X_alpha;
#endif

energy_qed_et = 0.0;
std::vector<double> polvec = {0.0, 0.0, 0.0};
Scheduler sch{ec};
Tensor<double> ehf_tmp{scf_vars.tAO, scf_vars.tAO};
Tensor<double> QED_Qxx{scf_vars.tAO, scf_vars.tAO};
Tensor<double> QED_Qyy{scf_vars.tAO, scf_vars.tAO};
sch.allocate(ehf_tmp, QED_Qxx, QED_Qyy).execute();

for(int i = 0; i < sys_data.qed_nmodes; i++) {
polvec = scf_options.qed_polvecs[i];

// clang-format off
sch
(ehf_tmp("i", "j") = X_a("i", "k") * X_a("j", "k"))
(ehf_tmp("i", "j") -= 0.5 * ttensors.D_last_alpha("i", "j"))
(QED_Qxx("i", "j") = polvec[0] * ttensors.QED_Dx("i", "j"))
(QED_Qxx("i", "j") += polvec[1] * ttensors.QED_Dy("i", "j"))
(QED_Qxx("i", "j") += polvec[2] * ttensors.QED_Dz("i", "j"))
(QED_Qyy("i", "j") = QED_Qxx("i", "k") * ehf_tmp("k", "j"))
(ehf_tmp("i", "j") = QED_Qyy("i", "k") * QED_Qxx("k", "j"))
.execute();
// clang-format on

const double coupl_strength = pow(scf_options.qed_lambdas[i], 2);
energy_qed_et +=
0.5 * coupl_strength * tt_trace(ec, ttensors.D_last_alpha, ttensors.ehf_tmp);
}
sch.deallocate(ehf_tmp, QED_Qxx, QED_Qyy).execute();

#if defined(USE_SCALAPACK)
Tensor<double>::deallocate(X_a);
#endif
}
}
if(is_uhf) {
nelectrons = tt_trace(ec, ttensors.D_last_alpha, ttensors.S1);
Expand All @@ -72,6 +131,13 @@ void exachem::scf::SCFIO::print_energies(ExecutionContext& ec, ChemEnv& chem_env
chem_env.sys_data.results["output"]["SCF"]["kinetic_1e"] = kinetic_1e;
chem_env.sys_data.results["output"]["SCF"]["energy_1e"] = energy_1e;
chem_env.sys_data.results["output"]["SCF"]["energy_2e"] = energy_2e;

if(is_qed) {
std::cout << "QED energy = " << energy_qed << std::endl;
std::cout << "QED eT energy = " << energy_qed_et << std::endl;
chem_env.sys_data.results["output"]["SCF"]["energy_qed"] = energy_qed;
chem_env.sys_data.results["output"]["SCF"]["energy_qed_et"] = energy_qed_et;
}
}
}

Expand Down
Loading

0 comments on commit 6a5a6fc

Please sign in to comment.