implement integral derivatives

tjira · Oct 12, 2024 · e027a1e · e027a1e
1 parent 565ffd2
commit e027a1e
Show file tree

Hide file tree

Showing 7 changed files with 180 additions and 12 deletions.
diff --git a/example/input/ints.json b/example/input/ints.json
@@ -0,0 +1,16 @@
+{
+    "system" : {
+        "path" : "example/molecule/water.xyz",
+        "basis" : "sto-3g"
+    },
+    "integral" : {
+        "data_export" : {
+            "hamiltonian" : true,
+            "coulomb" : true,
+            "overlap" : true,
+            "hamiltonian_d1" : true,
+            "coulomb_d1" : true,
+            "overlap_d1" : true
+        }
+    }
+}
diff --git a/include/input.h b/include/input.h
@@ -14,11 +14,11 @@ struct Input {
     struct Integral {
         double precision;
         struct DataExport {
-            bool hamiltonian, coulomb, overlap;
+            bool hamiltonian, coulomb, overlap, hamiltonian_d1, coulomb_d1, overlap_d1;
         } data_export;
     } integral;
     struct HartreeFock {
-        bool generalized;
+        bool generalized, gradient;
         int diis_size, max_iter;
         double threshold;
         struct ConfigurationInteraction {
@@ -61,12 +61,12 @@ extern nlohmann::json default_input; extern int nthread;
 
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::System, basis, path, charge, multiplicity);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::Wavefunction, dimension, mass, momentum, grid_limits, grid_points, guess);
-NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::Integral::DataExport, hamiltonian, coulomb, overlap);
+NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::Integral::DataExport, hamiltonian, coulomb, overlap, hamiltonian_d1, coulomb_d1, overlap_d1);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::Integral, precision, data_export);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::HartreeFock::ConfigurationInteraction, cas, triplet);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::HartreeFock::CoupledCluster, excitation, perturbation, max_iter, threshold);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::HartreeFock::MollerPlesset, order);
-NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::HartreeFock, diis_size, max_iter, threshold, configuration_interaction, coupled_cluster, moller_plesset, generalized);
+NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::HartreeFock, diis_size, max_iter, threshold, configuration_interaction, coupled_cluster, moller_plesset, generalized, gradient);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::QuantumDynamics::DataExport, diabatic_wavefunction, adiabatic_wavefunction, diabatic_density, adiabatic_density, diabatic_population, adiabatic_population, diabatic_potential, adiabatic_potential, total_energy, position, momentum, acf, potential_energy, kinetic_energy);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::QuantumDynamics, potential, imaginary, real, iterations, time_step, data_export);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::ClassicalDynamics::DataExport, total_energy, position, momentum, diabatic_population, adiabatic_population, total_energy_mean, position_mean, momentum_mean, potential_energy, potential_energy_mean, kinetic_energy, kinetic_energy_mean);

diff --git a/include/integral.h b/include/integral.h
@@ -8,11 +8,20 @@ class Integral {
 public:
     Integral(const Input::Integral& input) : precision(input.precision) {}
 
+    // basic integrals
     static torch::Tensor double_electron(libint2::Engine& engine, const libint2::BasisSet& shells);
     static torch::Tensor single_electron(libint2::Engine& engine, const libint2::BasisSet& shells);
 
+    // first derivatives
+    static torch::Tensor double_electron_d1(libint2::Engine& engine, const libint2::BasisSet& shells);
+    static torch::Tensor single_electron_d1(libint2::Engine& engine, const libint2::BasisSet& shells);
+
+    // calculate the basic integrals
     std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> calculate(const std::vector<libint2::Atom>& atoms, const libint2::BasisSet& shells) const;
 
+    // calculate the first derivatives
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> calculate_d1(const std::vector<libint2::Atom>& atoms, const libint2::BasisSet& shells) const;
+
 private:
     double precision;
 };
diff --git a/src/export.cpp b/src/export.cpp
@@ -270,6 +270,17 @@ void Export::TorchTensorDouble(const std::string& path, const torch::Tensor& ten
         }
     }
 
+    // write 3rd order tensor
+    if (tensor.sizes().size() == 3) {
+        for (int i = 0; i < tensor.sizes().at(0); i++) {
+            for (int j = 0; j < tensor.sizes().at(1); j++) {
+                for (int k = 0; k < tensor.sizes().at(2); k++) {
+                    file << std::setw(22) << tensor.index({i, j, k}).item().toDouble() << (j < tensor.sizes().at(1) - 1 ? " " : "");
+                }
+            } file << "\n";
+        }
+    }
+
     // write 4th order tensor
     if (tensor.sizes().size() == 4) {
         for (int i = 0; i < tensor.sizes().at(0); i++) {
@@ -282,6 +293,21 @@ void Export::TorchTensorDouble(const std::string& path, const torch::Tensor& ten
             }
         }
     }
+
+    // write 5th order tensor
+    if (tensor.sizes().size() == 5) {
+        for (int i = 0; i < tensor.sizes().at(0); i++) {
+            for (int j = 0; j < tensor.sizes().at(1); j++) {
+                for (int k = 0; k < tensor.sizes().at(2); k++) {
+                    for (int l = 0; l < tensor.sizes().at(3); l++) {
+                        for (int m = 0; m < tensor.sizes().at(4); m++) {
+                            file << std::setw(22) << tensor.index({i, j, k, l, m}).item().toDouble() << (k < tensor.sizes().at(2) - 1 ? " " : "");
+                        }
+                    }
+                } file << "\n";
+            }
+        }
+    }
 }
 
 void Export::WavefunctionTrajectory(const Input::QuantumDynamics& input, const std::vector<QuantumDynamics::IterationData>& iteration_data, const Eigen::MatrixXd& grid, int state, bool imaginary) {

diff --git a/src/input.cpp b/src/input.cpp
@@ -21,11 +21,15 @@ nlohmann::json default_input = R"(
         "data_export" : {
             "hamiltonian" : false,
             "coulomb" : false,
-            "overlap" : false
+            "overlap" : false,
+            "hamiltonian_d1" : false,
+            "coulomb_d1" : false,
+            "overlap_d1" : false
         }
     },
     "hartree_fock" : {
         "generalized" : false,
+        "gradient" : false,
         "diis_size" : 5,
         "max_iter" : 100,
         "threshold" : 1e-8,

diff --git a/src/integral.cpp b/src/integral.cpp
@@ -64,6 +64,84 @@ torch::Tensor Integral::single_electron(libint2::Engine& engine, const libint2::
     return torch::from_blob(ints.data(), {nbf, nbf}, torch::kDouble).clone();
 }
 
+torch::Tensor Integral::double_electron_d1(libint2::Engine& engine, const libint2::BasisSet& shells) {
+    // define the number of basis functions, matrix of integrals and shell to basis function map
+    int nbf = shells.nbf(), dim = 3; std::vector<double> ints(nbf * nbf * nbf * nbf * dim, 0); std::vector<size_t> sh2bf = shells.shell2bf();
+
+    // generate engine for every thread
+    std::vector<libint2::Engine> engines(nthread, engine);
+
+    // loop over all unique elements
+    #pragma omp parallel for num_threads(nthread)
+    for (size_t i = 0; i < shells.size(); i++) for (size_t j = i; j < shells.size(); j++) for (size_t k = i; k < shells.size(); k++) for (size_t l = (i == k ? j : k); l < shells.size(); l++) {
+
+        // calculate the integral and skip if it is zero
+        int integral_index = 0; engines.at(omp_get_thread_num()).compute(shells.at(i), shells.at(j), shells.at(k), shells.at(l)); if (engines.at(omp_get_thread_num()).results().at(0) == nullptr) continue;
+
+        // for every shell combination
+        for (size_t m = 0; m < shells.at(i).size(); m++) {
+            for (size_t n = 0; n < shells.at(j).size(); n++) {
+                for (size_t o = 0; o < shells.at(k).size(); o++) {
+                    for (size_t p = 0; p < shells.at(l).size(); p++) {
+
+                        // for every dimension
+                        for (int q = 0; q < dim; q++) {
+                            ints.at((m + sh2bf.at(i)) * nbf * nbf * nbf * dim + (n + sh2bf.at(j)) * nbf * nbf * dim + (o + sh2bf.at(k)) * nbf * dim + (p + sh2bf.at(l)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((m + sh2bf.at(i)) * nbf * nbf * nbf * dim + (n + sh2bf.at(j)) * nbf * nbf * dim + (p + sh2bf.at(l)) * nbf * dim + (o + sh2bf.at(k)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((n + sh2bf.at(j)) * nbf * nbf * nbf * dim + (m + sh2bf.at(i)) * nbf * nbf * dim + (o + sh2bf.at(k)) * nbf * dim + (p + sh2bf.at(l)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((n + sh2bf.at(j)) * nbf * nbf * nbf * dim + (m + sh2bf.at(i)) * nbf * nbf * dim + (p + sh2bf.at(l)) * nbf * dim + (o + sh2bf.at(k)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((o + sh2bf.at(k)) * nbf * nbf * nbf * dim + (p + sh2bf.at(l)) * nbf * nbf * dim + (m + sh2bf.at(i)) * nbf * dim + (n + sh2bf.at(j)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((o + sh2bf.at(k)) * nbf * nbf * nbf * dim + (p + sh2bf.at(l)) * nbf * nbf * dim + (n + sh2bf.at(j)) * nbf * dim + (m + sh2bf.at(i)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((p + sh2bf.at(l)) * nbf * nbf * nbf * dim + (o + sh2bf.at(k)) * nbf * nbf * dim + (m + sh2bf.at(i)) * nbf * dim + (n + sh2bf.at(j)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                            ints.at((p + sh2bf.at(l)) * nbf * nbf * nbf * dim + (o + sh2bf.at(k)) * nbf * nbf * dim + (n + sh2bf.at(j)) * nbf * dim + (m + sh2bf.at(i)) * dim + q) = engines.at(omp_get_thread_num()).results().at(q)[integral_index];
+                        }
+
+                        // increment index
+                        integral_index++;
+                    }
+                }
+            }
+        }
+    }
+
+    // return the integrals
+    return torch::from_blob(ints.data(), {nbf, nbf, nbf, nbf, dim}, torch::kDouble).clone();
+}
+
+torch::Tensor Integral::single_electron_d1(libint2::Engine& engine, const libint2::BasisSet& shells) {
+    // define the number of basis functions and dimensions, matrix of integrals and shell to basis function map
+    int nbf = shells.nbf(), dim = 3; std::vector<double> ints(nbf * nbf * dim, 0); std::vector<size_t> sh2bf = shells.shell2bf();
+
+    // generate engine for every thread
+    std::vector<libint2::Engine> engines(nthread, engine);
+
+    // loop over all unique elements
+    #pragma omp parallel for num_threads(nthread)
+    for (size_t i = 0; i < shells.size(); i++) for (size_t j = i; j < shells.size(); j++) {
+
+        // calculate the integral and skip if it is zero
+        int integral_index = 0; engines.at(omp_get_thread_num()).compute(shells.at(i), shells.at(j)); if (engines.at(omp_get_thread_num()).results().at(0) == nullptr) continue;
+
+        // for every shell pair
+        for (size_t k = 0; k < shells.at(i).size(); k++) {
+            for (size_t l = 0; l < shells.at(j).size(); l++) {
+
+                // for every dimension
+                for (int m = 0; m < dim; m++) {
+                    ints.at((k + sh2bf.at(i)) * nbf * dim + (l + sh2bf.at(j)) * dim + m) = engines.at(omp_get_thread_num()).results().at(m)[integral_index];
+                    ints.at((l + sh2bf.at(j)) * nbf * dim + (k + sh2bf.at(i)) * dim + m) = engines.at(omp_get_thread_num()).results().at(m)[integral_index];
+                }
+
+                // increment index
+                integral_index++;
+            }
+        }
+    }
+
+    // return the integrals
+    return torch::from_blob(ints.data(), {nbf, nbf, dim}, torch::kDouble).clone();
+}
+
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> Integral::calculate(const std::vector<libint2::Atom>& atoms, const libint2::BasisSet& shells) const {
     // initialize libint
     libint2::initialize();
@@ -89,3 +167,29 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> Integral::calculate(cons
     // return the integrals
     return std::make_tuple(T + V, S, J);
 }
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> Integral::calculate_d1(const std::vector<libint2::Atom>& atoms, const libint2::BasisSet& shells) const {
+    // initialize libint
+    libint2::initialize();
+
+    // define the engines
+    libint2::Engine kinetic_engine_d1(libint2::Operator::kinetic, shells.max_nprim(), shells.max_l(), 1, precision);
+    libint2::Engine nuclear_engine_d1(libint2::Operator::nuclear, shells.max_nprim(), shells.max_l(), 1, precision);
+    libint2::Engine overlap_engine_d1(libint2::Operator::overlap, shells.max_nprim(), shells.max_l(), 1, precision);
+    libint2::Engine coulomb_engine_d1(libint2::Operator::coulomb, shells.max_nprim(), shells.max_l(), 1, precision);
+
+    // set charges to the nuclear engine
+    nuclear_engine_d1.set_params(libint2::make_point_charges(atoms));
+
+    // calculate the integrals
+    torch::Tensor dT = single_electron_d1(kinetic_engine_d1, shells);
+    torch::Tensor dV = single_electron_d1(nuclear_engine_d1, shells);
+    torch::Tensor dS = single_electron_d1(overlap_engine_d1, shells);
+    torch::Tensor dJ = double_electron_d1(coulomb_engine_d1, shells);
+
+    // finalize libint
+    libint2::finalize();
+
+    // return the integrals
+    return std::make_tuple(dT + dV, dS, dJ);
+}
diff --git a/src/main.cpp b/src/main.cpp
@@ -79,11 +79,12 @@ int main(int argc, char** argv) {
         auto [input_json, input] = parse_input(program_input); System system; if (input_json.contains("system")) system = System(input.system);
 
         // define all the variables to store the results
-        double energy_hf = 0, energy_mp = 0, energy_cc = 0, energy_ci = 0; torch::Tensor H_AO, S_AO, J_AO, H_MS, J_MS_AP, F_MS, C_MS, E_CI, C_CI;
+        double energy_hf = 0, energy_mp = 0, energy_cc = 0, energy_ci = 0; torch::Tensor H_AO, dH_AO, S_AO, dS_AO, J_AO, dJ_AO, H_MS, J_MS_AP, F_MS, C_MS, E_CI, C_CI;
 
         // extract all the method booleans
         bool do_hartree_fock = input_json.contains("hartree_fock");
         bool do_integral = input_json.contains("integral") || do_hartree_fock;
+        bool do_integral_d1 = do_integral || (do_hartree_fock && input_json.at("hartree_fock").contains("gradient"));
         bool do_configuration_interaction = do_hartree_fock && input_json.at("hartree_fock").contains("configuration_interaction");
         bool do_coupled_cluster = do_hartree_fock && input_json.at("hartree_fock").contains("coupled_cluster");
         bool do_moller_plesset = do_hartree_fock && input_json.at("hartree_fock").contains("moller_plesset");
@@ -98,30 +99,38 @@ int main(int argc, char** argv) {
         if (input_json.contains("system")) printf("%s BASIS, %lu ATOMS, %d ELECTRONS, %d BASIS FUNCTIONS\n\n", system.get_basis().c_str(), system.get_atoms().size(), system.electrons(), system.basis_functions());
 
         // integral calculation
-        if (do_integral) {
+        if (do_integral || do_integral_d1) {
 
             // create the timer and print the calculation header
             Timepoint integral_calculation_timer = Timer::Now(); std::printf("INTEGRAL CALCULATION: "); std::flush(std::cout);
 
-            // calculate the one- and two-electron integrals
-            std::tie(H_AO, S_AO, J_AO) = Integral(input.integral).calculate(system.get_atoms(), system.get_shells());
+            // calculate the base integrals
+            if (do_integral) std::tie(H_AO, S_AO, J_AO) = Integral(input.integral).calculate(system.get_atoms(), system.get_shells());
+
+            // calculate the first derivatives
+            if (do_integral_d1) std::tie(dH_AO, dS_AO, dJ_AO) = Integral(input.integral).calculate_d1(system.get_atoms(), system.get_shells());
 
             // print the time taken to calculate the integrals
             std::printf("%s\n", Timer::Format(Timer::Elapsed(integral_calculation_timer)).c_str());
 
             // wrte the integrals to the disk
-            if (input.integral.data_export.hamiltonian || input.integral.data_export.coulomb || input.integral.data_export.overlap) {
+            if (input.integral.data_export.hamiltonian || input.integral.data_export.coulomb || input.integral.data_export.overlap || input.integral.data_export.hamiltonian_d1 || input.integral.data_export.coulomb_d1 || input.integral.data_export.overlap_d1) {
 
                 // create the export timer and print the header
                 Timepoint integral_export_timer = Timer::Now(); std::printf("WRITING INTS TO DISK: "); std::flush(std::cout);
 
-                // export the integrals
+                // export the base integrals
                 if (input.integral.data_export.hamiltonian) Export::TorchTensorDouble("H_AO.mat", H_AO);
                 if (input.integral.data_export.overlap    ) Export::TorchTensorDouble("S_AO.mat", S_AO);
                 if (input.integral.data_export.coulomb    ) Export::TorchTensorDouble("J_AO.mat", J_AO);
 
+                // export the first derivatives
+                if (input.integral.data_export.hamiltonian_d1) Export::TorchTensorDouble("dH_AO.mat", dH_AO);
+                if (input.integral.data_export.overlap_d1    ) Export::TorchTensorDouble("dS_AO.mat", dS_AO);
+                if (input.integral.data_export.coulomb_d1    ) Export::TorchTensorDouble("dJ_AO.mat", dJ_AO);
+
                 // print the time taken to export the integrals
-                std::printf("%s\n", Timer::Format(Timer::Elapsed(integral_export_timer)).c_str());
+                std::printf("%s%s", Timer::Format(Timer::Elapsed(integral_export_timer)).c_str(), do_hartree_fock ? "\n" : "");
             }
 
             // print the new line