propagator.h

Go to the documentation of this file.

 
#ifndef PROPAGATOR_H
#define PROPAGATOR_H
 
#include "source_base/constants.h"
#include "source_base/module_container/ATen/core/tensor.h" // ct::Tensor
#include "source_basis/module_ao/parallel_orbitals.h"
 
#include <complex>
 
namespace module_rt
{
//--------------------------------- Utility function ---------------------------------//
#ifdef __MPI
inline int globalIndex(int localindex, int nblk, int nprocs, int myproc)
{
    int iblock, gIndex;
    iblock = localindex / nblk;
    gIndex = (iblock * nprocs + myproc) * nblk + localindex % nblk;
    return gIndex;
}
#endif // __MPI
 
// Auxiliary function: process non-complex types, return value 1.0
template <typename T>
inline T init_value(typename std::enable_if<!std::is_same<T, std::complex<float>>::value
                                            && !std::is_same<T, std::complex<double>>::value>::type* = nullptr)
{
    return T(1.0);
}
 
// Auxiliary function: process complex types, return value 1.0 + 0.0i
template <typename T>
inline T init_value(typename std::enable_if<std::is_same<T, std::complex<float>>::value
                                            || std::is_same<T, std::complex<double>>::value>::type* = nullptr)
{
    return T(1.0, 0.0);
}
 
// Create an identity matrix of size n×n
template <typename T>
ct::Tensor create_identity_matrix(const int n, ct::DeviceType device = ct::DeviceType::CpuDevice)
{
    // Choose the data type of the Tensor
    ct::DataType data_type;
    if (std::is_same<T, float>::value)
    {
        data_type = ct::DataType::DT_FLOAT;
    }
    else if (std::is_same<T, double>::value)
    {
        data_type = ct::DataType::DT_DOUBLE;
    }
    else if (std::is_same<T, std::complex<float>>::value)
    {
        data_type = ct::DataType::DT_COMPLEX;
    }
    else if (std::is_same<T, std::complex<double>>::value)
    {
        data_type = ct::DataType::DT_COMPLEX_DOUBLE;
    }
    else
    {
        static_assert(std::is_same<T, float>::value || std::is_same<T, double>::value
                          || std::is_same<T, std::complex<float>>::value
                          || std::is_same<T, std::complex<double>>::value,
                      "Unsupported data type!");
    }
 
    ct::Tensor tensor(data_type, device, ct::TensorShape({n, n}));
    tensor.zero();
 
    // Set the diagonal elements to 1
    if (device == ct::DeviceType::CpuDevice)
    {
        // For CPU, we can directly access the data
        T* data_ptr = tensor.data<T>();
        for (int i = 0; i < n; ++i)
        {
            data_ptr[i * n + i] = init_value<T>();
        }
    }
#if ((defined __CUDA))
    else if (device == ct::DeviceType::GpuDevice)
    {
        // For GPU, we need to use a kernel to set the diagonal elements
        T* data_ptr = tensor.data<T>();
        for (int i = 0; i < n; ++i)
        {
            T value = init_value<T>();
            ct::kernels::set_memory<T, ct::DEVICE_GPU>()(data_ptr + i * n + i, value, 1);
        }
    }
#endif
 
    return tensor;
}
//--------------------------------- Utility function ---------------------------------//
 
class Propagator
{
  public:
    Propagator(const int ptype, const Parallel_Orbitals* pv, const double& dt)
    {
        this->ptype = ptype;
        this->ParaV = pv;
        this->dt = dt / ModuleBase::AU_to_FS;
    }
    ~Propagator();
 
#ifdef __MPI
    void compute_propagator(const int nlocal,
                            const std::complex<double>* Stmp,
                            const std::complex<double>* Htmp,
                            const std::complex<double>* H_laststep,
                            std::complex<double>* U_operator,
                            std::ofstream& ofs_running,
                            const int print_matrix) const;
 
    template <typename Device>
    void compute_propagator_tensor(const int nlocal,
                                   const ct::Tensor& Stmp,
                                   const ct::Tensor& Htmp,
                                   const ct::Tensor& H_laststep,
                                   ct::Tensor& U_operator,
                                   std::ofstream& ofs_running,
                                   const int print_matrix,
                                   const bool use_lapack) const;
#endif // __MPI
 
  private:
    int ptype; // type of propagator
    const Parallel_Orbitals* ParaV;
    double dt; // time step
 
#ifdef __MPI
 
    void compute_propagator_cn2(const int nlocal,
                                const std::complex<double>* Stmp,
                                const std::complex<double>* Htmp,
                                std::complex<double>* U_operator,
                                std::ofstream& ofs_running,
                                const int print_matrix) const;
 
    void compute_propagator_cn2_tensor(const int nlocal,
                                       const ct::Tensor& Stmp,
                                       const ct::Tensor& Htmp,
                                       ct::Tensor& U_operator,
                                       std::ofstream& ofs_running,
                                       const int print_matrix) const;
 
    template <typename Device>
    void compute_propagator_cn2_tensor_lapack(const int nlocal,
                                              const ct::Tensor& Stmp,
                                              const ct::Tensor& Htmp,
                                              ct::Tensor& U_operator,
                                              std::ofstream& ofs_running,
                                              const int print_matrix) const;
 
    void compute_propagator_taylor(const int nlocal,
                                   const std::complex<double>* Stmp,
                                   const std::complex<double>* Htmp,
                                   std::complex<double>* U_operator,
                                   std::ofstream& ofs_running,
                                   const int print_matrix,
                                   const int tag) const;
 
    void compute_propagator_etrs(const int nlocal,
                                 const std::complex<double>* Stmp,
                                 const std::complex<double>* Htmp,
                                 const std::complex<double>* H_laststep,
                                 std::complex<double>* U_operator,
                                 std::ofstream& ofs_running,
                                 const int print_matrix) const;
#endif // __MPI
};
} // namespace module_rt
 
#endif