add change for the npsin_4

Author: A-006

source/source_pw/module_pwdft/kernels/veff_op.cpp

@@ -20,20 +20,17 @@ struct veff_pw_op<FPTYPE, base_device::DEVICE_CPU>
                     const int& size,
                     std::complex<FPTYPE>* out,
                     std::complex<FPTYPE>* out1,
-                    const FPTYPE** in)
+                    const std::complex<FPTYPE>* in)
     {
+
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
-        for (int ir = 0; ir < size; ir++) {
-            auto sup = out[ir] * (in[0][ir] + in[3][ir])
-                + out1[ir]
-                        * (in[1][ir]
-                        - std::complex<FPTYPE>(0.0, 1.0) * in[2][ir]);
-            auto sdown = out1[ir] * (in[0][ir] - in[3][ir])
-                    + out[ir]
-                        * (in[1][ir]
-                            + std::complex<FPTYPE>(0.0, 1.0) * in[2][ir]);
+        for (int ir = 0; ir < size; ir++) 
+        {
+            const int base = ir * 4;
+            auto sup = out[ir] * (in[base]) + out1[ir] * (in[base + 1]);
+            auto sdown = out1[ir] * (in[base + 2]) + out[ir] * (in[base + 3]);
             out[ir] = sup;
             out1[ir] = sdown;
         }

source/source_pw/module_pwdft/kernels/veff_op.h

@@ -48,7 +48,7 @@ struct veff_pw_op {
         const int& size,
         std::complex<FPTYPE>* out,
         std::complex<FPTYPE>* out1,
-        const FPTYPE** in);
+        const std::complex<FPTYPE>* in);
 };
 
 #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM

source/source_pw/module_pwdft/operator_pw/veff_pw.cpp

@@ -25,6 +25,8 @@ Veff<OperatorPW<T, Device>>::Veff(const int* isk_in,
     this->veff_row = veff_row;
     this->veff_col = veff_col;
     this->wfcpw = wfcpw_in;
+    resmem_complex_op()(nspin_4_veff, 4*veff_col, "Veff<PW>::porter");
+    // this->nspin_4_veff=new std::complex<double>[4*veff_row];
     resmem_complex_op()(this->porter, this->wfcpw->nmaxgr, "Veff<PW>::porter");
     resmem_complex_op()(this->porter1, this->wfcpw->nmaxgr, "Veff<PW>::porter1");
 
@@ -35,6 +37,7 @@ Veff<OperatorPW<T, Device>>::~Veff()
 {
     delmem_complex_op()(this->porter);
     delmem_complex_op()(this->porter1);
+    delmem_complex_op()(this->nspin_4_veff);
 }
 
 template<typename T, typename Device>
@@ -111,17 +114,26 @@ void Veff<OperatorPW<T, Device>>::act(
     }
     else if (npol == 2)
     {
-        const Real* current_veff[4]={nullptr};
-        for (int is = 0; is < 4; is++)
+        const Real* current_veff={nullptr};
+        const std::complex<Real> imag=std::complex<Real>(0.0, 1.0);
+        for (int ir=0; ir < veff_col; ir++)
         {
-            current_veff[is] = this->veff + is * this->veff_col;
+            const int base = 4 *ir;
+            Real part_1 = this->veff[ir];
+            Real part_2 = this->veff[ir + veff_col];
+            Real part_3 = this->veff[ir + 2*veff_col];
+            Real part_4 = this->veff[ir + 3*veff_col];
+            nspin_4_veff[base ] = part_1 + part_4;
+            nspin_4_veff[base + 1] = part_2 - imag * part_3;
+            nspin_4_veff[base + 2] = part_1 - part_4;
+            nspin_4_veff[base + 3] = part_2 + imag * part_3;
         }
         for (int ib = 0; ib < nbands; ib += npol)
         {
             // FFT to real space and do things.
             wfcpw->recip_to_real<T, Device>(tmpsi_in, this->porter, this->ik);
             wfcpw->recip_to_real<T, Device>(tmpsi_in + max_npw, this->porter1, this->ik);
-            veff_op()(this->ctx, this->veff_col, this->porter, this->porter1, current_veff);
+            veff_op()(this->ctx, this->veff_col, this->porter, this->porter1, nspin_4_veff);
             // FFT back to G space.
             wfcpw->real_to_recip<T, Device>(this->porter, tmhpsi, this->ik, true);
             wfcpw->real_to_recip<T, Device>(this->porter1, tmhpsi + max_npw, this->ik, true);

source/source_pw/module_pwdft/operator_pw/veff_pw.h

@@ -68,6 +68,7 @@ class Veff<OperatorPW<T, Device>> : public OperatorPW<T, Device>
     const Real *veff = nullptr, *h_veff = nullptr, *d_veff = nullptr;
     T *porter = nullptr;
     T *porter1 = nullptr;
+    mutable T* nspin_4_veff=nullptr;
     base_device::AbacusDevice_t device = {};
     using veff_op = veff_pw_op<Real, Device>;