Feature: Implement k-point Continuity Wavefunction Initialization Strategy

[jieli-matrix]

Background

We will introduce a new wavefunction initialization method that leverages k-point continuity in reciprocal space to accelerate SCF convergence. The method may reduce diagonalization iterations by 20% for dense k-meshes through FFT-based wavefunction propagation between adjacent k-points.

Describe the solution you'd like

Implementation Roadmap

1. Core Propagation Logic

source/module_hsolver/hsolver_pw.cpp

• Add k-point adjacency handling in HSolverPW::solve loop
• Implement phase-shift propagation between nearest k-points
• Maintain initialized k-point cache using octree spatial partitioning

2. Initializer Extension

source/module_psi/psi_initializer_atomic.h

• Add new psi initializer consistent with psi_initializer_atomic e.g.

class psi_initializer_kcontinuity : public psi_initializer<T> {
public:
    void propagate(const ModulePW::PW_Basis_K& pw_wfc,
                 const K_Vectors& kv,
                 psi::Psi<T, Device>& psi,
                 int ik_current);
private:
    std::map<int, int> k_adjacency; // k_current -> k_prev mapping
    FFT_Phase_Shift phase_shift;    // Handles (k_n - k_m)·r phase factors
};

3. k-point Sorting Enhancement

• Implement octree-based k-point sorting in K_Vectors.

Key Integration Points

1.Initialization Flow Control

// Modified initialization flow
if(init_wfc == "kcontinuity" && ik > 0) {
    auto& psi_prev = get_psi_from_nearest_initialized_k();
    psi_initer->propagate(pw_wfc, kv, psi, ik);
} else {
    // Fallback to traditional methods
    psi_initer->init_psig(...);
}

2. FFT Phase Handling

void apply_kphase_shift(psi::Psi<T,Device>& psi,
                       const ModuleBase::Vector3<double> delta_k,
                       const ModulePW::PW_Basis_K& pw_wfc) {
    // Real-space grid iteration
    for(int ir=0; ir<pw_wfc.nrxx; ir++) {
        const double phase = delta_k * pw_wfc.get_gr(ir);
        psi(ir) *= std::exp(std::complex<double>(0, phase));
    }
}

Validation Strategy

1. Unit Tests

2. Integration Tests

• Compare total energy convergence vs atomic/random initialization
• Monitor iteration counts
• Check charge density continuity between k-points

Task list only for developers

• [ ] Notice possible changes of behavior
• [ ] Explain the changes of codes in core modules of ESolver, HSolver, ElecState, Hamilt, Operator or Psi

Notice Possible Changes of Behavior (Reminder only for developers)

No response

Notice any changes of core modules (Reminder only for developers)

No response

Notice Possible Changes of Core Modules (Reminder only for developers)

No response

Additional Context

No response

Task list for Issue attackers (only for developers)

• [ ] Review and understand the proposed feature and its importance.
• [ ] Research on the existing solutions and relevant research articles/resources.
• [ ] Discuss with the team to evaluate the feasibility of implementing the feature.
• [ ] Create a design document outlining the proposed solution and implementation details.
• [ ] Get feedback from the team on the design document.
• [ ] Develop the feature following the agreed design.
• [ ] Write unit tests and integration tests for the feature.
• [ ] Update the documentation to include the new feature.
• [ ] Perform code review and address any issues.
• [ ] Merge the feature into the main branch.
• [ ] Monitor for any issues or bugs reported by users after the feature is released.
• [ ] Address any issues or bugs reported by users and continuously improve the feature.