evolve_elec.h

Go to the documentation of this file.

#ifndef W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_MODULE_HAMILT_LCAO_MODULE_TDDFT_EVOLVE_ELEC_H
#define W_ABACUS_DEVELOP_ABACUS_DEVELOP_SOURCE_MODULE_HAMILT_LCAO_MODULE_TDDFT_EVOLVE_ELEC_H
 
#include "source_base/global_function.h"
#include "source_base/global_variable.h"
#include "source_base/module_container/ATen/core/tensor.h"     // ct::Tensor
#include "source_base/module_container/ATen/core/tensor_map.h" // TensorMap
#include "source_base/module_device/device.h"                  // base_device
#include "source_base/module_device/memory_op.h"               // memory operations
#include "source_base/module_external/scalapack_connector.h"   // Cpxgemr2d
#include "source_esolver/esolver_ks_lcao.h"
#include "source_esolver/esolver_ks_lcao_tddft.h"
#include "source_lcao/hamilt_lcao.h"
#include "source_lcao/module_rt/gather_mat.h" // MPI gathering and distributing functions
#include "source_psi/psi.h"
 
//-----------------------------------------------------------
// mohan add 2021-02-09
// This class is used to evolve the electronic wave functions
// in TDDFT in terms of the multiple k points
// k is the index for the points in the first Brillouin zone
//-----------------------------------------------------------
 
//------------------------ Debugging utility function ------------------------//
 
// Print the shape of a Tensor
inline void print_tensor_shape(const ct::Tensor& tensor, const std::string& name)
{
    std::cout << "Shape of " << name << ": [";
    for (int i = 0; i < tensor.shape().ndim(); ++i)
    {
        std::cout << tensor.shape().dim_size(i);
        if (i < tensor.shape().ndim() - 1)
        {
            std::cout << ", ";
        }
    }
    std::cout << "]" << std::endl;
}
 
// Recursive print function
template <typename T>
inline void print_tensor_data_recursive(const T* data,
                                        const std::vector<int64_t>& shape,
                                        const std::vector<int64_t>& strides,
                                        int dim,
                                        std::vector<int64_t>& indices,
                                        const std::string& name)
{
    if (dim == shape.size())
    {
        // Recursion base case: print data when reaching the innermost dimension
        std::cout << name;
        for (size_t i = 0; i < indices.size(); ++i)
        {
            std::cout << "[" << indices[i] << "]";
        }
        std::cout << " = " << *data << std::endl;
        return;
    }
    // Recursively process the current dimension
    for (int64_t i = 0; i < shape[dim]; ++i)
    {
        indices[dim] = i;
        print_tensor_data_recursive(data + i * strides[dim], shape, strides, dim + 1, indices, name);
    }
}
 
// Generic print function
template <typename T>
inline void print_tensor_data(const ct::Tensor& tensor, const std::string& name)
{
    const std::vector<int64_t>& shape = tensor.shape().dims();
    const std::vector<int64_t>& strides = tensor.shape().strides();
    const T* data = tensor.data<T>();
    std::vector<int64_t> indices(shape.size(), 0);
    print_tensor_data_recursive(data, shape, strides, 0, indices, name);
}
 
// Specialization for std::complex<double>
template <>
inline void print_tensor_data<std::complex<double>>(const ct::Tensor& tensor, const std::string& name)
{
    const std::vector<int64_t>& shape = tensor.shape().dims();
    const std::vector<int64_t>& strides = tensor.shape().strides();
    const std::complex<double>* data = tensor.data<std::complex<double>>();
    std::vector<int64_t> indices(shape.size(), 0);
    print_tensor_data_recursive(data, shape, strides, 0, indices, name);
}
 
//------------------------ Debugging utility function ------------------------//
 
namespace module_rt
{
template <typename Device = base_device::DEVICE_CPU>
class Evolve_elec
{
    friend class ModuleESolver::ESolver_KS_LCAO<std::complex<double>, double>;
 
    // Template parameter is needed for the friend class declaration
    friend class ModuleESolver::ESolver_KS_LCAO_TDDFT<double, Device>;
    friend class ModuleESolver::ESolver_KS_LCAO_TDDFT<std::complex<double>, Device>;
 
  public:
    Evolve_elec();
    ~Evolve_elec();
 
  private:
    static void solve_psi(const int& istep,
                          const int nband,
                          const int nlocal,
                          const int& nks,
                          hamilt::Hamilt<std::complex<double>>* phm,
                          Parallel_Orbitals& para_orb,
                          psi::Psi<std::complex<double>>* psi,
                          psi::Psi<std::complex<double>>* psi_laststep,
                          ct::Tensor& Hk_laststep,
                          ct::Tensor& Sk_laststep,
                          ModuleBase::matrix& ekb,
                          std::ofstream& ofs_running,
                          const int propagator,
                          const bool use_tensor,
                          const bool use_lapack);
 
    // ct_device_type = ct::DeviceType::CpuDevice or ct::DeviceType::GpuDevice
    static ct::DeviceType ct_device_type;
    // ct_Device = ct::DEVICE_CPU or ct::DEVICE_GPU
    using ct_Device = typename ct::PsiToContainer<Device>::type;
 
    // Memory operations
    using syncmem_double_h2d_op = base_device::memory::synchronize_memory_op<double, Device, base_device::DEVICE_CPU>;
    using syncmem_double_d2h_op = base_device::memory::synchronize_memory_op<double, base_device::DEVICE_CPU, Device>;
    using syncmem_complex_h2d_op
        = base_device::memory::synchronize_memory_op<std::complex<double>, Device, base_device::DEVICE_CPU>;
    using syncmem_complex_d2h_op
        = base_device::memory::synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, Device>;
};
} // namespace module_rt
#endif