Merge branch 'deepmodeling:develop' into dsp1

Author: A-006

cmake/FindCereal.cmake

@@ -15,7 +15,7 @@ if(NOT CEREAL_INCLUDE_DIR)
     include(FetchContent)
     FetchContent_Declare(
         cereal
-        URL https://github.com/USCiLab/cereal/archive/refs/tags/v1.3.2.tar.gz
+        URL https://codeload.github.com/USCiLab/cereal/tar.gz/master
     )
     FetchContent_Populate(cereal)
     set(CEREAL_INCLUDE_DIR ${cereal_SOURCE_DIR}/include)

python/pyabacus/src/ModuleNAO/py_m_nao.cpp

@@ -152,12 +152,9 @@ void bind_m_nao(py::module& m)
                 double* cvR = static_cast<double*>(pvR_info.ptr);
                 ModuleBase::Vector3<double> vR(cvR[0], cvR[1], cvR[2]);
                 double out[1] = {0.0};
-                double* grad_out = nullptr;
-                if (cal_grad)
-                {
-                    grad_out = new double[3];
-                }
-                self.calculate(itype1, l1, izeta1, m1, itype2, l2, izeta2, m2, vR, out, grad_out);
+                double grad_out[3] = {0.0, 0.0, 0.0};
+                double* grad_ptr = cal_grad ? grad_out : nullptr;
+                self.calculate(itype1, l1, izeta1, m1, itype2, l2, izeta2, m2, vR, out, grad_ptr);
                 py::array_t<double> out_array(1, out);
                 if (cal_grad)
                 {

python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp

@@ -132,6 +132,11 @@ class PyDiagoDavSubspace
             std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out);
         };
 
+        auto spsi_func = [this](const std::complex<double>* psi_in,
+                                std::complex<double>* spsi_out,
+                                const int ld_psi,
+                                const int nvec) { syncmem_op()(spsi_out, psi_in, static_cast<size_t>(ld_psi * nvec)); };
+
         obj = std::make_unique<hsolver::Diago_DavSubspace<std::complex<double>, base_device::DEVICE_CPU>>(
             precond_vec, 
             nband, 
@@ -145,7 +150,7 @@ class PyDiagoDavSubspace
             nb2d
         );
 
-        return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type);
+        return obj->diag(hpsi_func, spsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type);
     }
 
 private:
@@ -156,6 +161,10 @@ class PyDiagoDavSubspace
     int nband;
 
     std::unique_ptr<hsolver::Diago_DavSubspace<std::complex<double>, base_device::DEVICE_CPU>> obj;
+
+    base_device::DEVICE_CPU* ctx = {};
+    using syncmem_op = base_device::memory::
+        synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, base_device::DEVICE_CPU>;
 };
 
 } // namespace py_hsolver

python/pyabacus/tests/test_m_nao.py

@@ -73,7 +73,36 @@ def test_rc():
     assert orb(0,1,1).l == 1
     assert orb(0,2,0).l == 2
 
+def test_twocenterintegrator():
+    orb_dir = '../../../tests/PP_ORB/'
+    file_list = ["C_gga_8au_100Ry_2s2p1d.orb", "O_gga_10au_100Ry_2s2p1d.orb"]
+    file_list = [orb_dir + orbfile for orbfile in file_list]
+
+    orb = nao.RadialCollection()
+    orb.build(2, file_list, 'o')
+
+    alpha = nao.RadialCollection()
+    alpha.build(1, [file_list[0]])
+
+    dr = 0.01  # R spacing
+    rmax = max(orb.rcut_max, alpha.rcut_max)
+    cutoff = 2.0 * rmax
+    nr = int(rmax / dr) + 1
+
+    orb.set_uniform_grid(True, nr, cutoff, 'i', True)
+    alpha.set_uniform_grid(True, nr, cutoff, 'i', True)
+
+    sbt = base.SphericalBesselTransformer()
+    orb.set_transformer(sbt)
+    alpha.set_transformer(sbt)
+
+    integrator = nao.TwoCenterIntegrator()
+    integrator.tabulate(orb, alpha, 'S', nr, cutoff)
 
+    overlap = integrator.snap(0, 0, 0, 0, 0, np.array([0.0, 0.0, 0.0]), False)
+    assert 1 - overlap[0][0] < 1e-10
+    overlap = integrator.snap(1, 0, 0, 0, 0, np.array([3.0, 3.0, 3.0]), False)
+    assert abs(overlap[0][0] - 0.031136758774787342) < 1e-10
 
 
 

source/CMakeLists.txt

@@ -54,6 +54,10 @@ list(APPEND device_srcs
   source_pw/module_pwdft/kernels/vnl_op.cpp
   source_base/kernels/math_ylm_op.cpp
   source_hamilt/module_xc/kernels/xc_functional_op.cpp
+  source_pw/module_pwdft/kernels/cal_density_real_op.cpp
+  source_pw/module_pwdft/kernels/mul_potential_op.cpp
+  source_pw/module_pwdft/kernels/vec_mul_vec_complex_op.cpp
+  source_pw/module_pwdft/kernels/exx_cal_energy_op.cpp
 )
 
 if(USE_CUDA)
@@ -80,6 +84,10 @@ if(USE_CUDA)
     source_base/kernels/cuda/math_kernel_op.cu
     source_base/kernels/cuda/math_kernel_op_vec.cu
     source_hamilt/module_xc/kernels/cuda/xc_functional_op.cu
+    source_pw/module_pwdft/kernels/cuda/cal_density_real_op.cu
+    source_pw/module_pwdft/kernels/cuda/mul_potential_op.cu
+    source_pw/module_pwdft/kernels/cuda/vec_mul_vec_complex.cu
+    source_pw/module_pwdft/kernels/cuda/exx_cal_energy_op.cu
   )
 endif()
 

source/Makefile.Objects

@@ -342,6 +342,10 @@ OBJS_HAMILT=hamilt_pw.o\
     velocity_pw.o\
     radial_proj.o\
     exx_helper.o\
+    vec_mul_vec_complex_op.o\
+    exx_cal_energy_op.o\
+    cal_density_real_op.o\
+    mul_potential_op.o\
 
 OBJS_HAMILT_OF=kedf_tf.o\
     kedf_vw.o\

source/source_base/kernels/math_kernel_op.cpp

@@ -1,4 +1,5 @@
 #include "source_base/kernels/math_kernel_op.h"
+#include "source_base/module_external/blas_connector.h"
 
 #include <iomanip>
 #include <iostream>

source/source_base/kernels/math_kernel_op.h

@@ -5,7 +5,6 @@
 
 #include "source_base/macros.h"
 
-#include "source_base/module_external/blas_connector.h"
 #include "source_base/parallel_reduce.h"
 
 #include "source_base/module_device/memory_op.h"

source/source_base/kernels/math_kernel_op_vec.cpp

@@ -1,4 +1,6 @@
 #include "source_base/kernels/math_kernel_op.h"
+#include "source_base/module_external/blas_connector.h"
+
 
 namespace ModuleBase
 {

source/source_base/kernels/test/math_kernel_test.cpp

@@ -1,7 +1,7 @@
-#include "source_base/module_external/blas_connector.h"
 #include "source_base/constants.h"
 #include "source_base/module_device/memory_op.h"
 #include "source_base/kernels/math_kernel_op.h"
+#include "source_base/module_external/blas_connector.h"
 
 #include <complex>
 #include <gtest/gtest.h>