Fix size mismatch for equivariant version DeePKS.

Author: ErjieWu

source/module_hamilt_lcao/module_deepks/deepks_force.cpp

@@ -49,16 +49,16 @@ void DeePKS_domain::cal_f_delta(const std::vector<std::vector<TK>>& dm,
             const int nw2_tot,
             ModuleBase::Vector3<int> dR2)
         {
-            double r0[3] = {0, 0, 0};
             double r1[3] = {0, 0, 0};
+            double r2[3] = {0, 0, 0};
             if (isstress)
             {
                 r1[0] = (tau1.x - tau0.x);
                 r1[1] = (tau1.y - tau0.y);
                 r1[2] = (tau1.z - tau0.z);
-                r0[0] = (tau2.x - tau0.x);
-                r0[1] = (tau2.y - tau0.y);
-                r0[2] = (tau2.z - tau0.z);
+                r2[0] = (tau2.x - tau0.x);
+                r2[1] = (tau2.y - tau0.y);
+                r2[2] = (tau2.z - tau0.z);
             }
 
             auto row_indexes = pv.get_indexes_row(ibt1);
@@ -255,7 +255,7 @@ void DeePKS_domain::cal_f_delta(const std::vector<std::vector<TK>>& dm,
                             for (int jpol = ipol; jpol < 3; jpol++)
                             {
                                 svnl_dalpha(ipol, jpol)
-                                    += *dm_current * (nlm[ipol] * r0[jpol] + nlm_t[ipol] * r1[jpol]);
+                                    += *dm_current * (nlm[ipol] * r2[jpol] + nlm_t[ipol] * r1[jpol]);
                             }
                         }
                     }

source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp

@@ -27,6 +27,7 @@
 void DeePKS_domain::read_pdm(bool read_pdm_file,
                              bool is_equiv,
                              bool& init_pdm,
+                             const int nat,
                              const int inlmax,
                              const int lmaxd,
                              const int* inl_l,
@@ -68,9 +69,9 @@ void DeePKS_domain::read_pdm(bool read_pdm_file,
                 nproj += (2 * il + 1) * alpha.getNchi(il);
             }
             pdm_size = nproj * nproj;
-            for (int inl = 0; inl < inlmax; inl++)
+            for (int iat = 0; iat < nat; iat++)
             {
-                auto accessor = pdm[inl].accessor<double, 1>();
+                auto accessor = pdm[iat].accessor<double, 1>();
                 for (int ind = 0; ind < pdm_size; ind++)
                 {
                     double c;
@@ -128,9 +129,9 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
             nproj += (2 * il + 1) * orb.Alpha[0].getNchi(il);
         }
         pdm_size = nproj * nproj;
-        for (int inl = 0; inl < inlmax; inl++)
+        for (int iat = 0; iat < ucell.nat; iat++)
         {
-            pdm[inl] = torch::zeros({pdm_size}, torch::kFloat64);
+            pdm[iat] = torch::zeros({pdm_size}, torch::kFloat64);
         }
     }
 
@@ -268,7 +269,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
                             s_2t[i * col_size + icol] = col_ptr->get_value(col_indexes[icol], trace_alpha_col[i]);
                         }
                     }
-                    // prepare DM_gamma from DMR
+                    // prepare DM from DMR
                     std::vector<double> dm_array(row_size * col_size, 0.0);
                     const double* dm_current = nullptr;
                     for (int is = 0; is < dm->get_DMR_vector().size(); is++)
@@ -286,6 +287,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
                             dRy = dR2.y - dR1.y;
                             dRz = dR2.z - dR1.z;
                         }
+                        // dm_R
                         auto* tmp = dm->get_DMR_vector()[is]->find_matrix(ibt1, ibt2, dRx, dRy, dRz);
                         if (tmp == nullptr)
                         {
@@ -306,7 +308,10 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
                     }
 
                     dm_current = dm_array.data();
-                    // dgemm for s_2t and dm_current to get g_1dmt
+                    // use s_2t and dm_current to get g_1dmt
+                    // dgemm_: C = alpha * A * B + beta * C
+                    // C = g_1dmt, A = dm_current, B = s_2t
+                    // all the input should be data pointer
                     constexpr char transa = 'T', transb = 'N';
                     const double gemm_alpha = 1.0, gemm_beta = 1.0;
                     dgemm_(&transa,
@@ -338,6 +343,8 @@ void DeePKS_domain::cal_pdm(bool& init_pdm,
                             {
                                 for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d
                                 {
+                                    // ddot_: dot product of two vectors
+                                    // inc means the increment of the index
                                     accessor[m1][m2] += ddot_(&row_size,
                                                               g_1dmt.data() + index * row_size,
                                                               &inc,

source/module_hamilt_lcao/module_deepks/deepks_pdm.h

@@ -32,6 +32,7 @@ namespace DeePKS_domain
 void read_pdm(bool read_pdm_file,
               bool is_equiv,
               bool& init_pdm,
+              const int nat,
               const int inlmax,
               const int lmaxd,
               const int* inl_l,

source/module_hamilt_lcao/module_deepks/deepks_spre.cpp

@@ -57,14 +57,14 @@ void DeePKS_domain::cal_gdmepsl(const int lmaxd,
             const int nw2_tot,
             ModuleBase::Vector3<int> dR2)
         {
-            double r0[3] = {0, 0, 0};
             double r1[3] = {0, 0, 0};
+            double r2[3] = {0, 0, 0};
             r1[0] = (tau1.x - tau0.x);
             r1[1] = (tau1.y - tau0.y);
             r1[2] = (tau1.z - tau0.z);
-            r0[0] = (tau2.x - tau0.x);
-            r0[1] = (tau2.y - tau0.y);
-            r0[2] = (tau2.z - tau0.z);
+            r2[0] = (tau2.x - tau0.x);
+            r2[1] = (tau2.y - tau0.y);
+            r2[2] = (tau2.z - tau0.z);
             auto row_indexes = pv.get_indexes_row(ibt1);
             auto col_indexes = pv.get_indexes_col(ibt2);
             if (row_indexes.size() * col_indexes.size() == 0)
@@ -152,7 +152,7 @@ void DeePKS_domain::cal_gdmepsl(const int lmaxd,
                                             accessor[mm][inl][m2][m1]
                                                 += ucell.lat0 * *dm_current
                                                    * (grad_overlap_2[jpol]->get_value(col_indexes[iw2], ib + m2)
-                                                      * overlap_1->get_value(row_indexes[iw1], ib + m1) * r0[ipol]);
+                                                      * overlap_1->get_value(row_indexes[iw1], ib + m1) * r2[ipol]);
                                             accessor[mm][inl][m2][m1]
                                                 += ucell.lat0 * *dm_current
                                                    * (overlap_2->get_value(col_indexes[iw2], ib + m1)