add analytical HF gradient using torch tensors

tjira · Oct 16, 2024 · 1f1848d · 1f1848d
1 parent 519f787
commit 1f1848d
Show file tree

Hide file tree

Showing 15 changed files with 186 additions and 71 deletions.
diff --git a/example/input/ints.json b/example/input/ints.json
@@ -5,10 +5,12 @@
     },
     "integral" : {
         "data_export" : {
-            "hamiltonian" : true,
+            "kinetic" : true,
+            "nuclear" : true,
             "coulomb" : true,
             "overlap" : true,
-            "hamiltonian_d1" : true,
+            "kinetic_d1" : true,
+            "nuclear_d1" : true,
             "coulomb_d1" : true,
             "overlap_d1" : true
         }

diff --git a/include/classicaldynamics.h b/include/classicaldynamics.h
@@ -9,7 +9,7 @@
 class ClassicalDynamics {
 public:
     struct TrajectoryData {
-        Eigen::VectorXi state; Eigen::MatrixXd position, velocity; std::vector<Eigen::MatrixXd> diabatic_potential, adiabatic_potential; std::vector<Eigen::VectorXd> hopping_geometry;
+        Eigen::VectorXi state; Eigen::MatrixXd position, velocity; std::vector<Eigen::MatrixXd> diabatic_potential, adiabatic_potential; std::vector<Eigen::VectorXd> hopping_geometry; std::vector<double> hopping_time;
     };
 
     ClassicalDynamics(const Input::ClassicalDynamics& input) : input(input) {}

diff --git a/include/constant.h b/include/constant.h
@@ -1,8 +1,8 @@
 #pragma once
 
-#define PERIODIC_TABLE "H He Li Be B C N O F Ne"
+#define PERIODIC_TABLE "H He Li Be B C N O F Ne Na Mg Al Si P S Cl Ar K Ca Sc Ti V Cr Mn Fe Co Ni Cu Zn Ga Ge As Se Br Kr Rb Sr Y Zr Nb Mo Tc Ru Rh Pd Ag Cd In Sn Sb Te I Xe Cs Ba La Ce Pr Nd Pm Sm Eu Gd Tb Dy Ho Er Tm Yb Lu Hf Ta W Re Os Ir Pt Au Hg Tl Pb Bi Po At Rn Fr Ra Ac Th Pa U Np Pu Am Cm Bk Cf Es Fm Md No Lr Rf Db Sg Bh Hs Mt Ds Rg Cn Nh Fl Mc Lv Ts Og"
 
 #define OCCUPIED_ORBITALS "abcdefghijklmnopqrstuvwxyz"
-#define VIRTUAL_ORBITALS "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+#define VIRTUAL_ORBITALS  "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 
 #define ANGSTROM_TO_BOHR 1.889726124626
diff --git a/include/hartreefock.h b/include/hartreefock.h
@@ -13,6 +13,7 @@ class HartreeFock {
     torch::Tensor get_fock(const torch::Tensor& H_AO, const torch::Tensor& J_AO, const torch::Tensor& D_MO) const;
     std::tuple<torch::Tensor, torch::Tensor, double> run(const System& system, const torch::Tensor& H_AO, const torch::Tensor& S_AO, const torch::Tensor& J_AO, torch::Tensor D_MO) const;
     std::tuple<torch::Tensor, torch::Tensor, double> scf(const System& system, const torch::Tensor& H, const torch::Tensor& S, const torch::Tensor& J, torch::Tensor D) const;
+    torch::Tensor gradient(const System& system, const torch::Tensor& dT_AO, const torch::Tensor& dV_AO, const torch::Tensor& dS_AO, const torch::Tensor& dJ_AO, const torch::Tensor& F_AO, const torch::Tensor& C_MO) const;
 
 private:
     Input::HartreeFock input;

diff --git a/include/input.h b/include/input.h
@@ -14,7 +14,7 @@ struct Input {
     struct Integral {
         double precision;
         struct DataExport {
-            bool hamiltonian, coulomb, overlap, hamiltonian_d1, coulomb_d1, overlap_d1;
+            bool kinetic, nuclear, coulomb, overlap, kinetic_d1, nuclear_d1, coulomb_d1, overlap_d1;
         } data_export;
     } integral;
     struct HartreeFock {
@@ -44,7 +44,7 @@ struct Input {
     } quantum_dynamics;
     struct ClassicalDynamics {
         struct DataExport {
-            bool diabatic_population, adiabatic_population, total_energy, total_energy_mean, position, position_mean, momentum, momentum_mean, potential_energy, potential_energy_mean, kinetic_energy, kinetic_energy_mean, hopping_geometry;
+            bool diabatic_population, adiabatic_population, total_energy, total_energy_mean, position, position_mean, momentum, momentum_mean, potential_energy, potential_energy_mean, kinetic_energy, kinetic_energy_mean, hopping_geometry, hopping_time;
         } data_export;
         struct SurfaceHopping {
             int quantum_step_factor;
@@ -61,15 +61,15 @@ 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, hamiltonian_d1, coulomb_d1, overlap_d1);
+NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::Integral::DataExport, kinetic, nuclear, coulomb, overlap, kinetic_d1, nuclear_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, 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, hopping_geometry);
+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, hopping_geometry, hopping_time);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::ClassicalDynamics::SurfaceHopping, type, quantum_step_factor);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input::ClassicalDynamics, potential, iterations, trajectories, time_step, data_export, seed, adiabatic, surface_hopping, log_interval_step, log_interval_traj);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(Input, system, wavefunction, integral, hartree_fock, quantum_dynamics, classical_dynamics);
diff --git a/include/integral.h b/include/integral.h
@@ -13,14 +13,14 @@ class Integral {
     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);
+    static torch::Tensor double_electron_d1(libint2::Engine& engine, const libint2::BasisSet& shells, const std::vector<libint2::Atom>& atoms);
+    static torch::Tensor single_electron_d1(libint2::Engine& engine, const libint2::BasisSet& shells, const std::vector<libint2::Atom>& atoms);
 
     // calculate the basic integrals
-    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> calculate(const std::vector<libint2::Atom>& atoms, const libint2::BasisSet& shells) const;
+    std::tuple<torch::Tensor, 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;
+    std::tuple<torch::Tensor, 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/include/system.h b/include/system.h
@@ -18,6 +18,7 @@ class System {
     int get_multi() const;
     libint2::BasisSet get_shells() const;
     double nuclear_repulsion() const;
+    torch::Tensor nuclear_repulsion_d1() const;
     int virtual_spinorbitals() const;
 
 private:

diff --git a/research/fssh_vs_lzsh.sh b/research/fssh_vs_lzsh.sh
@@ -4,9 +4,10 @@
 # VARIABLES
 # ======================================================================================================================================================================================================
 
-CORES=1; PSTART=10.0; PSTEP=10.0; PEND=50.0; TRAJS=1000; LOG_INTERVAL_STEP=1000; LOG_INTERVAL_TRAJ=100; CLEAN=1
+CORES=12; PSTART=10.0; PSTEP=00.0; PEND=50.0; TRAJS=10000; LOG_INTERVAL_STEP=1000; LOG_INTERVAL_TRAJ=100; CLEAN=1
 
 MODELS=("TULLY_1" "TULLY_2" "DS_1" "DS_2" "TS_1" "TS_2" "TS_3" "TS_4")
+MODELS=("TULLY_1")
 
 # ======================================================================================================================================================================================================
 # START TEMPLATES
@@ -48,7 +49,8 @@ read -r -d '' TEMPLATE_SH_DYN <<- EOM
         "diabatic_population" : true, "adiabatic_population" : true,
         "position_mean"       : true, "momentum_mean"        : true,
         "total_energy_mean"   : true, "potential_energy_mean": true,
-        "kinetic_energy_mean" : true, "hopping_geometry"     : true
+        "kinetic_energy_mean" : true, "hopping_geometry"     : true,
+        "hopping_time"        : true
     }
 }
 EOM
@@ -225,11 +227,12 @@ for MODEL in ${MODELS[@]}; do
 
         # plot the hopping geometries
         plot-1d.py HOPPING-GEOMETRIES_FS-ADIABATIC.mat --bins 100 --legend "FSSH" --title "HOPPING GEOMETRIES: ${MODEL}" --xlabel "Coordinate (a.u.)" --ylabel "Relative Count" --output "HOPPING-GEOMETRIES_${MODEL}" --histogram --png
+        plot-1d.py HOPPING-TIMES_FS-ADIABATIC.mat      --bins 100 --legend "FSSH" --title "HOPPING TIMES: ${MODEL}"      --xlabel "Time (a.u.)"       --ylabel "Relative Count" --output "HOPPING-TIMES_${MODEL}"      --histogram --png
 
         # make the trajectory analysis image
         montage "POTENTIAL-ADIABATIC_${MODEL}.png" "POPULATION-ADIABATIC_${MODEL}_P=${MOMENTUM}.png" "POSITION_${MODEL}_P=${MOMENTUM}.png"         "MOMENTUM_${MODEL}_P=${MOMENTUM}.png"     -mode concatenate -tile x1 "TRAJECTORIES-GENERAL_${MODEL}_P=${MOMENTUM}.png"
         montage "POTENTIAL-ADIABATIC_${MODEL}.png" "KINETIC-ENERGY_${MODEL}_P=${MOMENTUM}.png"       "POTENTIAL-ENERGY_${MODEL}_P=${MOMENTUM}.png" "TOTAL-ENERGY_${MODEL}_P=${MOMENTUM}.png" -mode concatenate -tile x1 "TRAJECTORIES-ENERGETICS_${MODEL}_P=${MOMENTUM}.png"
-        montage "POTENTIAL-ADIABATIC_${MODEL}.png" "HOPPING-GEOMETRIES_${MODEL}.png"                                                                                                         -mode concatenate -tile x1 "TRAJECTORIES-TRANSITIONS_${MODEL}_P=${MOMENTUM}.png"
+        montage "POTENTIAL-ADIABATIC_${MODEL}.png" "HOPPING-GEOMETRIES_${MODEL}.png"                 "HOPPING-TIMES_${MODEL}.png"                                                            -mode concatenate -tile x1 "TRAJECTORIES-TRANSITIONS_${MODEL}_P=${MOMENTUM}.png"
     done
 
     # plot the population dependence on momentum

diff --git a/src/classicaldynamics.cpp b/src/classicaldynamics.cpp
@@ -110,8 +110,8 @@ void ClassicalDynamics::run(const Input::Wavefunction& initial_diabatic_wavefunc
         // define and initialie the surface hopping algorithms
         FewestSwitches fewestswitches(input.surface_hopping, input.adiabatic, seed); LandauZener landauzener(input.surface_hopping, input.adiabatic, seed);
 
-        // define the vector of hopping geometries
-        std::vector<Eigen::VectorXd> hopping_geometry_vector;
+        // define the vector of hopping geometries and times
+        std::vector<Eigen::VectorXd> hopping_geometry_vector; std::vector<double> hopping_time_vector;
 
         // define the initial conditions
         Eigen::MatrixXd position(input.iterations + 1, grid.cols()), velocity(input.iterations + 1, grid.cols()), acceleration(input.iterations + 1, grid.cols());
@@ -159,8 +159,8 @@ void ClassicalDynamics::run(const Input::Wavefunction& initial_diabatic_wavefunc
                 velocity.row(j).array() = std::sqrt(velocity.row(j).squaredNorm() - 2 * (potential(new_state, new_state) - potential(state(j), state(j))) / mass); state(j) = new_state;
             }
 
-            // append the hopping geometry
-            if (j && state(j) != state(j - 1)) hopping_geometry_vector.push_back(position.row(j));
+            // append the hopping geometry and time
+            if (j && state(j) != state(j - 1)) hopping_geometry_vector.push_back(position.row(j)), hopping_time_vector.push_back(j * input.time_step);
 
             // print the iteration info
             if ((j % input.log_interval_step == 0 || j && state(j) != state(j - 1)) && (i ? i + 1 : i) % input.log_interval_traj == 0) {
@@ -169,7 +169,7 @@ void ClassicalDynamics::run(const Input::Wavefunction& initial_diabatic_wavefunc
         }
 
         // set the trajectory data
-        trajectory_data_vector.at(i) = {state, position, velocity, diabatic_potential_vector, adiabatic_potential_vector, hopping_geometry_vector};
+        trajectory_data_vector.at(i) = {state, position, velocity, diabatic_potential_vector, adiabatic_potential_vector, hopping_geometry_vector, hopping_time_vector};
     }
 
     // create the vector of populations

diff --git a/src/export.cpp b/src/export.cpp
@@ -244,6 +244,24 @@ void Export::ClassicalTrajectories(const Input::ClassicalDynamics& input, const
         // export the energy mean
         Export::EigenMatrixDouble(std::string("HOPPING-GEOMETRIES_") + algorithm + (input.adiabatic ? "-ADIABATIC" : "-DIABATIC") + ".mat", data_matrix);
     }
+
+    // export the hopping times
+    if (input.data_export.hopping_time) {
+
+        // create the data of the hopping times
+        std::vector<double> hopping_times;
+
+        // fill the hopping geometry data
+        for (int i = 0; i < input.trajectories; i++) for (size_t j = 0; j < trajectory_data_vector.at(i).hopping_time.size(); j++) {
+            hopping_times.push_back(trajectory_data_vector.at(i).hopping_time.at(j));
+        }
+
+        // transform the hopping geometries to a matrix
+        data_matrix = Eigen::MatrixXd::Zero(hopping_times.size(), 1); for (size_t i = 0; i < hopping_times.size(); i++) data_matrix(i, 0) = hopping_times.at(i);
+
+        // export the energy mean
+        Export::EigenMatrixDouble(std::string("HOPPING-TIMES_") + algorithm + (input.adiabatic ? "-ADIABATIC" : "-DIABATIC") + ".mat", data_matrix);
+    }
 }
 
 void Export::EigenMatrixDouble(const std::string& path, const Eigen::MatrixXd& matrix, const Eigen::MatrixXd& independent_variable) {

diff --git a/src/hartreefock.cpp b/src/hartreefock.cpp
@@ -13,9 +13,59 @@ double HartreeFock::get_energy(const torch::Tensor& H_AO, const torch::Tensor& F
 }
 
 torch::Tensor HartreeFock::get_fock(const torch::Tensor& H_AO, const torch::Tensor& J_AO, const torch::Tensor& D_MO) const {
-    return H_AO + torch::einsum("ijkl,ij->kl", {J_AO - (input.generalized ? 1.0 : 0.5) * J_AO.swapaxes(0, 3), D_MO});
+    return H_AO + torch::einsum("ijkl,ij->kl", {J_AO - (input.generalized ? 1.0 : 0.5) * J_AO.swapaxes(1, 3), D_MO});
 }
 
+torch::Tensor HartreeFock::gradient(const System& system, const torch::Tensor& dT_AO, const torch::Tensor& dV_AO, const torch::Tensor& dS_AO, const torch::Tensor& dJ_AO, const torch::Tensor& F_AO, const torch::Tensor& C_MO) const {
+    // throw an error for generalized calculations
+    if (input.generalized) throw std::runtime_error("GRADIENTS NOT SUPPORTED FOR GENERALIZED HARTREE-FOCK");
+
+    // extract the one-integral integral derivative slices
+    torch::Tensor dT_AO_1 = dT_AO.index({Slice(), Slice(), Slice(None, 3)}), dV_AO_1 = dV_AO.index({Slice(), Slice(), Slice(None, 3)}), dS_AO_1 = dS_AO.index({Slice(), Slice(), Slice(None, 3)});
+
+    // extract the two-integral integral derivative slices and the density matrix
+    torch::Tensor dJ_AO_1 = dJ_AO.index({Slice(), Slice(), Slice(), Slice(), Slice(None, 3)}), D_MO = get_density(system, C_MO);
+
+    // extract the orbital energies, number of occupied orbitals and shell map
+    torch::Tensor eps = Transform::SingleSpatial(F_AO, C_MO).diagonal(); int nocc = system.electrons() / 2; auto atom2shell = system.get_shells().atom2shell(system.get_atoms());
+
+    // initialize the gradient matrix and weighted density matrix
+    torch::Tensor G = torch::zeros({(int)system.get_atoms().size(), 3}, torch::kDouble), W = torch::zeros_like(C_MO);
+
+    // fill the weighted density matrix
+    for (int i = 0; i < W.size(0); i++) for (int j = 0; j < W.size(1); j++) {
+        W.index_put_({i, j}, 2 * (C_MO.index({i, Slice(None, nocc)}) * C_MO.index({j, Slice(None, nocc)}) * eps.index({Slice(None, nocc)})).sum());
+    }
+
+    // contract the ERI derivative with the second density matrix
+    torch::Tensor dERI =  torch::einsum("kl,ijklx->ijx", {D_MO, dJ_AO_1 - 0.5 * dJ_AO_1.swapaxes(1, 3)});
+
+    // calculate the gradient matrix
+    for (int i = 0, si = 0, ss = 0; i < G.size(0); i++, si += ss, ss = 0) {
+
+        // calculate number of shells for current atom
+        for (long shell : atom2shell.at(i)) ss += system.get_shells().at(shell).size();
+
+        // define the hamiltonian derivative as the nuclear derivative terms
+        torch::Tensor dH_AO = dV_AO.index({Slice(), Slice(), Slice(6 + i * 3, 9 + i * 3)});
+
+        // add the orbital derivative terms
+        dH_AO.index({Slice(), Slice(si, si + ss), Slice(None, 3)}) += (dT_AO_1 + dV_AO_1).index({Slice(si, si + ss), Slice(), Slice(None, 3)}).swapaxes(0, 1);
+        dH_AO.index({Slice(si, si + ss), Slice(), Slice(None, 3)}) += (dT_AO_1 + dV_AO_1).index({Slice(si, si + ss), Slice(), Slice(None, 3)}).swapaxes(0, 0);
+
+        // add the first term of the gradient matrix
+        G.index({i, Slice()}) += torch::einsum("ijx,ij->x", {dH_AO, D_MO});
+
+        // add the second and third term of the gradient matrix
+        G.index({i, Slice()}) += 2 * torch::einsum("ijx,ij->x", {dERI.index({Slice(si, si + ss), Slice(), Slice()}), D_MO.index({Slice(si, si + ss), Slice()})});
+        G.index({i, Slice()}) -= 2 * torch::einsum("ijx,ij->x", {dS_AO_1.index({Slice(si, si + ss), Slice(), Slice()}), W.index({Slice(si, si + ss), Slice()})});
+    }
+
+    // return the gradient matrix
+    return G;
+}
+
+
 std::tuple<torch::Tensor, torch::Tensor, double> HartreeFock::run(const System& system, const torch::Tensor& H_AO, const torch::Tensor& S_AO, const torch::Tensor& J_AO, torch::Tensor D_MO) const {
     // throw an error if the system is open shell and not specified as generalized
     if (system.get_multi() != 1 && !input.generalized) throw std::runtime_error("OPEN SHELL SYSTEMS REQUIRE GENERALIZED HARTREE-FOCK");
@@ -101,9 +151,6 @@ std::tuple<torch::Tensor, torch::Tensor, double> HartreeFock::scf(const System&
         if (i == input.max_iter - 1) throw std::runtime_error("MAXIMUM NUMBER OF ITERATIONS IN THE SCF REACHED");
     }
 
-    // transform the results to the spin basis
-    torch::Tensor C_MS = input.generalized ? C : Transform::CoefficientSpin(C); torch::Tensor F_MS = input.generalized ? Transform::SingleSpatial(F, C_MS) : Transform::SingleSpin(F, C_MS);
-
     // return the converged results
-    return {F_MS, C_MS, energy_hf};
+    return {F, C, energy_hf};
 }
diff --git a/src/input.cpp b/src/input.cpp
@@ -19,10 +19,12 @@ nlohmann::json default_input = R"(
     "integral" : {
         "precision" : 1e-12,
         "data_export" : {
-            "hamiltonian" : false,
+            "kinetic" : false,
+            "nuclear" : false,
             "coulomb" : false,
             "overlap" : false,
-            "hamiltonian_d1" : false,
+            "kinetic_d1" : false,
+            "nuclear_d1" : false,
             "coulomb_d1" : false,
             "overlap_d1" : false
         }
@@ -96,7 +98,8 @@ nlohmann::json default_input = R"(
             "potential_energy_mean" : false,
             "kinetic_energy" : false,
             "kinetic_energy_mean" : false,
-            "hopping_geometry" : false
+            "hopping_geometry" : false,
+            "hopping_time" : false
         }
     }
 }