Fix(deltaspin,psi): disable subspace for nspin=4 and fix PAGED_GPU memory risks

- lambda_loop.cpp: add nspin_==2 guard to subspace acceleration,
  since nspin=4 P_I_sub loses spin-channel separation from I_spin pre_hr
- psi.cpp: fix copy constructor crash when source has PAGED_GPU flag
  but psi_cpu_ not allocated (set_storage_mode edge case)
- psi.cpp: fix_k() auto-routes to psi_cpu_ when GPU lacks k-point data,
  covering psiToRho/cal_becsum/cal_occupations without per-caller changes
- psi.cpp: operator= supports PAGED_GPU (copy psi_cpu_, skip GPU buffer)
- psi.cpp: zero_out() handles PAGED_GPU (memset CPU + zero GPU buffer)
- psi.cpp: operator()(ik,ib,ibasis) reads from psi_cpu_ in PAGED_GPU
- psi.cpp: fix_kb/fix_b add PAGED_GPU guard
- psi.h/psi_paging.cpp: add get_cpu_pointer_safe() for null-safe access
- setup_psi_pw.cpp: copy_d2h_impl reads from psi_cpu_ in PAGED_GPU

GPU memory invariant preserved: psi->psi holds only single k-point.

Author: dyzheng

source/source_lcao/module_deltaspin/lambda_loop.cpp

@@ -691,49 +691,21 @@ void spinconstrain::SpinConstrain<std::complex<double>>::run_lambda_loop(int out
         mean_error = sum_2d(temp_1) / nat;
         rms_error = std::sqrt(mean_error);
 
+        // =============================================================
         // =============================================================
         // ACCELERATION ACTIVATION CHECK (LCAO nspin=2 only)
         // =============================================================
         //
-        // PURPOSE: When the BFGS optimizer is near convergence (RMS error
-        // drops below sc_acceleration_rms_thr), lambda changes become small.
-        // At this point, subspace methods can approximate the response much
-        // faster than full diagonalization while maintaining accuracy.
-        //
-        // ACTIVATION CONDITIONS (all must be true):
-        //   1. PARAM.inp.basis_type == "lcao" (only LCAO basis supported)
-        //   2. this->nspin_ == 2 (only collinear spin supported;
-        //      nspin=4 requires complex-type operator template instantiation)
-        //   3. sc_acceleration_mode != "off" (user must explicitly enable;
-        //      valid values: "first_order", "subspace")
-        //   4. sc_acceleration_rms_thr > 0 (threshold must be set)
-        //   5. rms_error < sc_acceleration_rms_thr (current error is small enough)
-        //
-        // ONCE-ONLY ACTIVATION:
-        //   The check uses `!this->acceleration_active_` to ensure subspace
-        //   is built only ONCE per SCF iteration. After activation, all
-        //   subsequent cal_mw_from_lambda calls use the accelerated path
-        //   until the next SCF iteration resets the flag.
-        //
-        // WHY build subspace at current lambda?
-        //   The subspace approximation is only valid for small perturbations
-        //   around the reference lambda. Building it at the current (near-converged)
-        //   lambda ensures the reference point is close to the optimal solution.
-        //   If we built it at lambda=0, the approximation would be poor for
-        //   large lambda values.
-        //
-        // WHY use i_step=-2?
-        //   The main BFGS loop uses i_step = -1, 0, 1, ..., nsc-1.
-        //   Using -2 ensures this call is handled by a special branch in
-        //   cal_mw_from_lambda that does full diagonalization + cache build
-        //   WITHOUT affecting the main loop state (psi, delta_lambda, etc.).
-        //   This is a clean separation of concerns.
-        //
-        // WHY free_lcao_subspace_cache() first?
-        //   To prevent memory leaks from previous SCF iterations. Each SCF
-        //   iteration should build a fresh subspace at its own converged lambda.
+        // nspin=4: subspace acceleration is disabled. The pre_hr matrix
+        // stores I_spin structure (same value in all 4 spin blocks),
+        // causing pre-computed P_I_sub to lose spin-channel separation.
+        // Full diagonalization is always used for nspin=4.
         // =============================================================
+
+        const bool nspin_ok = (this->nspin_ == 2);
+
         const bool accel_enabled = (PARAM.inp.basis_type == "lcao") &&
+                                    nspin_ok &&
                                     (this->sc_acceleration_mode_ != "off") &&
                                     (this->sc_acceleration_rms_thr_ > 0.0) &&
                                     (rms_error < this->sc_acceleration_rms_thr_);

source/source_psi/psi.cpp

@@ -206,9 +206,21 @@ Psi<T, Device>::Psi(const Psi& psi_in)
         const size_t total_size = static_cast<size_t>(this->nk) * this->nbands * this->nbasis;
         const T* src_ptr = psi_in.get_pointer() - psi_in.get_psi_bias();
 
-        if (psi_in.get_storage_mode() == PsiStorageMode::PAGED_GPU)
+         if (psi_in.get_storage_mode() == PsiStorageMode::PAGED_GPU)
         {
-            std::memcpy(this->psi_cpu_, psi_in.get_cpu_pointer(), sizeof(T) * total_size);
+            // Source may have PAGED_GPU flag set but psi_cpu_ not yet allocated
+            // (e.g. after set_storage_mode() in setup_psi_pw). In that case,
+            // copy from the source's main buffer (psi) into our psi_cpu_.
+            const T* cp_ptr = psi_in.get_cpu_pointer_safe();
+            if (cp_ptr != nullptr)
+            {
+                std::memcpy(this->psi_cpu_, cp_ptr, sizeof(T) * total_size);
+            }
+            else
+            {
+                // Fallback: source data is in psi buffer (ALL_GPU layout)
+                std::memcpy(this->psi_cpu_, src_ptr, sizeof(T) * total_size);
+            }
         }
         else if (std::is_same<Device, base_device::DEVICE_CPU>::value)
         {
@@ -255,10 +267,23 @@ Psi<T, Device>::Psi(const Psi<T_in, Device_in>& psi_in)
     {
         const size_t total_size = static_cast<size_t>(this->nk) * this->nbands * this->nbasis;
 
-        if (psi_in.get_storage_mode() == PsiStorageMode::PAGED_GPU)
+         if (psi_in.get_storage_mode() == PsiStorageMode::PAGED_GPU)
         {
-            const T_in* src_ptr = psi_in.get_cpu_pointer();
-            detail::TypeCopy<T, T_in, std::is_same<T, T_in>::value>::copy(this->psi_cpu_, src_ptr, total_size);
+            // Source may have PAGED_GPU flag set but psi_cpu_ not yet allocated.
+            const T_in* cp_ptr = psi_in.get_cpu_pointer_safe();
+            if (cp_ptr != nullptr)
+            {
+                detail::TypeCopy<T, T_in, std::is_same<T, T_in>::value>::copy(this->psi_cpu_, cp_ptr, total_size);
+            }
+            else
+            {
+                // Fallback: source data is in main psi buffer (ALL_GPU layout)
+                auto* arr = (T*)malloc(sizeof(T) * total_size);
+                base_device::memory::cast_memory_op<T, T_in, Device_in, Device_in>()(
+                    arr, psi_in.get_pointer() - psi_in.get_psi_bias(), total_size);
+                std::memcpy(this->psi_cpu_, arr, sizeof(T) * total_size);
+                free(arr);
+            }
         }
         else
         {
@@ -317,13 +342,30 @@ Psi<T, Device>& Psi<T, Device>::operator=(const Psi<T, Device>& psi_in)
     this->k_first = psi_in.get_k_first();
     // this function will copy psi_in.psi to this->psi no matter the device types of each other.
 
+    this->storage_mode_ = psi_in.get_storage_mode();
     this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis());
-    base_device::memory::synchronize_memory_op<T, Device, Device>()(this->psi,
-                                                                    psi_in.psi,
-                                                                    psi_in.size());
-    this->psi_bias = psi_in.get_psi_bias();
+
+    if (this->storage_mode_ == PsiStorageMode::PAGED_GPU)
+    {
+        // PAGED_GPU: copy the full CPU buffer, GPU buffer stays single-k
+        if (psi_cpu_ != nullptr && psi_in.psi_cpu_ != nullptr)
+        {
+            const size_t total_size = static_cast<size_t>(this->nk) * this->nbands * this->nbasis;
+            std::memcpy(this->psi_cpu_, psi_in.psi_cpu_, sizeof(T) * total_size);
+        }
+        this->current_k_gpu_ = -1;
+        this->psi_bias = 0;
+        this->psi_current = this->psi;
+    }
+    else
+    {
+        base_device::memory::synchronize_memory_op<T, Device, Device>()(this->psi,
+                                                                        psi_in.psi,
+                                                                        psi_in.size());
+        this->psi_bias = psi_in.get_psi_bias();
+        this->psi_current = this->psi + psi_in.get_psi_bias();
+    }
     this->current_nbasis = psi_in.get_current_nbas();
-    this->psi_current = this->psi + psi_in.get_psi_bias();
 
     return *this;
 }
@@ -485,7 +527,17 @@ void Psi<T, Device>::fix_k(const int ik) const
     if (storage_mode_ == PsiStorageMode::PAGED_GPU)
     {
         this->psi_bias = 0;
-        this->psi_current = const_cast<T*>(this->psi);
+        if (this->current_k_gpu_ == ik)
+        {
+            // This k-point is loaded on GPU by load_k_to_gpu
+            this->psi_current = const_cast<T*>(this->psi);
+        }
+        else
+        {
+            // Read from CPU full buffer (psi_cpu_ has all k-points)
+            this->psi_current = const_cast<T*>(psi_cpu_
+                + static_cast<size_t>(ik) * this->nbands * this->nbasis);
+        }
         return;
     }
 
@@ -508,6 +560,13 @@ void Psi<T, Device>::fix_b(const int ib) const
     assert(ib >= 0);
     this->current_b = ib;
 
+    if (storage_mode_ == PsiStorageMode::PAGED_GPU)
+    {
+        this->psi_bias = 0;
+        this->psi_current = const_cast<T*>(this->psi);
+        return;
+    }
+
     if (!this->k_first)
     {
         this->current_k = 0;
@@ -532,6 +591,14 @@ void Psi<T, Device>::fix_kb(const int ik, const int ib) const
     assert(ik >= 0 && ib >= 0);
     this->current_k = ik;
     this->current_b = ib;
+
+    if (storage_mode_ == PsiStorageMode::PAGED_GPU)
+    {
+        this->psi_bias = 0;
+        this->psi_current = const_cast<T*>(this->psi);
+        return;
+    }
+
     if (ik >= this->nk || ib >= this->nbands)
     { // fix to 0
         this->psi_bias = 0;
@@ -549,6 +616,14 @@ T& Psi<T, Device>::operator()(const int ikb1, const int ikb2, const int ibasis)
 {
     assert(ikb1 >= 0 && ikb2 >= 0 && ibasis >= 0);
     assert(this->k_first ? ikb1 < this->nk && ikb2 < this->nbands : ikb1 < this->nbands && ikb2 < this->nk);
+
+    if (storage_mode_ == PsiStorageMode::PAGED_GPU)
+    {
+        // PAGED_GPU: GPU buffer is single-k, read from CPU full buffer (k_first layout)
+        assert(psi_cpu_ != nullptr);
+        return const_cast<T*>(psi_cpu_)[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis];
+    }
+
     return this->k_first ? this->psi[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis]
                          : this->psi[(ikb1 * this->nk + ikb2) * this->nbasis + ibasis];
 }
@@ -597,8 +672,21 @@ const int& Psi<T, Device>::get_ngk(const int ik_in) const
 template <typename T, typename Device>
 void Psi<T, Device>::zero_out()
 {
-    // this->psi.assign(this->psi.size(), T(0));
-    set_memory_op()(this->psi, 0, this->size());
+    if (storage_mode_ == PsiStorageMode::PAGED_GPU)
+    {
+        // Zero full CPU buffer and single-k GPU buffer
+        if (psi_cpu_ != nullptr)
+        {
+            const size_t total_size = static_cast<size_t>(this->nk) * this->nbands * this->nbasis;
+            std::memset(psi_cpu_, 0, sizeof(T) * total_size);
+        }
+        const size_t k_size = static_cast<size_t>(this->nbands) * this->nbasis;
+        set_memory_op()(this->psi, 0, k_size);
+    }
+    else
+    {
+        set_memory_op()(this->psi, 0, this->size());
+    }
 }
 
 template <typename T, typename Device>

source/source_psi/psi.h

@@ -158,6 +158,7 @@ class Psi
     int get_current_k_gpu() const { return current_k_gpu_; }
     T* get_cpu_pointer(int ik = 0);
     const T* get_cpu_pointer(int ik = 0) const;
+    const T* get_cpu_pointer_safe(int ik = 0) const;
 
   private:
     T* psi = nullptr; // avoid using C++ STL

source/source_psi/psi_paging.cpp

@@ -114,6 +114,17 @@ const T* Psi<T, Device>::get_cpu_pointer(int ik) const
     }
 }
 
+template <typename T, typename Device>
+const T* Psi<T, Device>::get_cpu_pointer_safe(int ik) const
+{
+    if (storage_mode_ == PsiStorageMode::PAGED_GPU && psi_cpu_ != nullptr)
+    {
+        if (ik < 0 || ik >= this->nk) return nullptr;
+        return psi_cpu_ + static_cast<size_t>(ik) * this->nbands * this->nbasis;
+    }
+    return nullptr;
+}
+
 template <typename T, typename Device>
 void Psi<T, Device>::load_k_to_gpu(int ik)
 {
@@ -199,6 +210,7 @@ void Psi<T, Device>::set_psi_cpu_external(T* ext_cpu_buf)
     template void Psi<T, Device>::set_storage_mode(PsiStorageMode);                                                    \
     template T* Psi<T, Device>::get_cpu_pointer(int);                                                                  \
     template const T* Psi<T, Device>::get_cpu_pointer(int) const;                                                      \
+    template const T* Psi<T, Device>::get_cpu_pointer_safe(int) const;                                                 \
     template void Psi<T, Device>::load_k_to_gpu(int);                                                                  \
     template void Psi<T, Device>::store_k_from_gpu(int);                                                               \
     template void Psi<T, Device>::ensure_k_on_gpu(int);                                                                \

source/source_psi/setup_psi_pw.cpp

@@ -200,9 +200,18 @@ template <typename T, typename Device>
 void Setup_Psi_pw::copy_d2h_impl()
 {
     auto* psi_t = this->get_psi_t<T, Device>();
-    this->castmem_d2h_impl<T, Device>(this->psi_cpu[0].get_pointer() - this->psi_cpu[0].get_psi_bias(),
-                                      psi_t->get_pointer() - psi_t->get_psi_bias(),
-                                      this->psi_cpu[0].size());
+    if (psi_t->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU)
+    {
+        // PAGED_GPU: full data is on CPU in psi_cpu_, just memcpy
+        const size_t total_size = sizeof(T) * psi_t->get_nk() * psi_t->get_nbands() * psi_t->get_nbasis();
+        std::memcpy(this->psi_cpu[0].get_pointer(), psi_t->get_cpu_pointer(), total_size);
+    }
+    else
+    {
+        this->castmem_d2h_impl<T, Device>(this->psi_cpu[0].get_pointer() - this->psi_cpu[0].get_psi_bias(),
+                                          psi_t->get_pointer() - psi_t->get_psi_bias(),
+                                          this->psi_cpu[0].size());
+    }
 }
 
 void Setup_Psi_pw::copy_d2h()