interaction_elph.h

#ifndef EL_PH_INTERACTION_H
#define EL_PH_INTERACTION_H
 
#include <complex>
 
#include "constants.h"
#include "crystal.h"
#include "eigen.h"
#include "phonon_h0.h"
#include "points.h"
#include "utilities.h"
#include "context.h"
#include "common_kokkos.h"
#include <Kokkos_Core.hpp>
 
// TODO: add flag to let user decide whether to use or not polar corrections
class InteractionElPhWan {
 
  Crystal crystal;
  PhononH0 phononH0;
  int numPhBands, numElBands, numWsR1Vectors, numWsR2Vectors;
  std::vector<Eigen::Tensor<double, 3>> cacheCoupling;
  bool usePolarCorrection = false;
  const int phaseConvention;
 
  // kokkos objects for GPU accelerated elph coupling
  ComplexView4D elPhCached;
  ComplexView5D couplingWannier_k;
  DoubleView2D wsR2Vectors_k;
  DoubleView1D wsR2VectorsDegeneracies_k;
  DoubleView2D wsR1Vectors_k;
  DoubleView1D wsR1VectorsDegeneracies_k;
  std::vector<ComplexView4D::HostMirror> elPhCached_hs;
 
  // TODO TEMPORARY, for old interpolation method
  Eigen::MatrixXd wsR1Vectors;
  Eigen::VectorXd wsR1VectorsDegeneracies;
  Eigen::MatrixXd wsR2Vectors;
  Eigen::VectorXd wsR2VectorsDegeneracies;
  Eigen::Tensor<std::complex<double>, 4> elPhCached_old;
  Eigen::Vector3d cachedK1;
  Eigen::Tensor<std::complex<double>, 5> couplingWannier;
 
#ifdef MPI_AVAIL
  std::vector<MPI_Request> mpi_requests;
#endif
 
public:
 
  // phase convention options -- TODO switch to ENUM
  // Giustino uses Re, Rp for R vectors
  // JDFTx uses Re, Re' R vectors 
  static const int GiustinoPhaseConvention = 0; 
  static const int JdftxPhaseConvention = 1; 
  
  // spin types -- TODO switch to ENUM
  static const int spinNonPolarized = 1; 
  static const int spinPolarizedOrSOC = 2; 
 
  InteractionElPhWan(
    const Crystal& crystal_,
    const Eigen::Tensor<std::complex<double>, 5> &couplingWannier_,
    const Eigen::MatrixXd &wsR1Vectors_,
    const Eigen::VectorXd &wsR1VectorsDegeneracies_,
    const Eigen::MatrixXd &wsR2Vectors_,
    const Eigen::VectorXd &wsR2VectorsDegeneracies_,
    const int& phaseConvention, const PhononH0& phononH0_);
 
    void resetK1(); // just resets the first transform for a new cycle
 
  //InteractionElPhWan(Crystal &crystal_);
 
  //InteractionElPhWan(const InteractionElPhWan &that);
 
  //InteractionElPhWan &operator=(const InteractionElPhWan &that);
 
  ~InteractionElPhWan();
 
  // Core elph interpolation functions --------------------------------
 
  void calcCouplingSquared(
      const Eigen::MatrixXcd &eigvec1,
      const std::vector<Eigen::MatrixXcd> &eigvecs2,
      const std::vector<Eigen::MatrixXcd> &eigvecs3,
      const std::vector<Eigen::Vector3d> &q3Cs,
      const Eigen::Vector3d &k1C,
      const std::vector<Eigen::VectorXcd> &polarData);
 
  void oldCalcCouplingSquared(
    const Eigen::MatrixXcd &eigvec1,
    const std::vector<Eigen::MatrixXcd> &eigvecs2,
    const std::vector<Eigen::MatrixXcd> &eigvecs3, const Eigen::Vector3d &k1C,
    const std::vector<Eigen::Vector3d> &k2Cs,
    const std::vector<Eigen::Vector3d> &q3Cs);
 
  void cacheElPh(const Eigen::MatrixXcd &eigvec1, const Eigen::Vector3d &k1C);
 
  Eigen::Tensor<double, 3>& getCouplingSquared(const int &ik2);
 
  static InteractionElPhWan parse(Context &context, Crystal &crystal,
                                  PhononH0& phononH0_);
 
  // Helper functions --------------------------------
 
  Eigen::Tensor<std::complex<double>, 3>
  getPolarCorrection(const Eigen::Vector3d &q3, const Eigen::MatrixXcd &ev1,
                     const Eigen::MatrixXcd &ev2, const Eigen::MatrixXcd &ev3);
 
  double getDeviceMemoryUsage();
 
  Eigen::VectorXi getCouplingDimensions();
 
  int estimateNumBatches(const int &nk2, const int &nb1);
 
  // Internal polar correction functions --------------------------------
 
  // TODO these functions need full documentation
 
  // functions to help with the calculation of the polar correction
  // as described in doi:10.1103/physRevLett.115.176401, Eq. 4
 
  // This calculates the long range V_L component of g_L, to be used
  // in the qe->Wannier transformation
  static Eigen::Tensor<std::complex<double>, 3> getPolarCorrectionStatic(
      const Eigen::Vector3d &q3, const Eigen::MatrixXcd &ev1,
      const Eigen::MatrixXcd &ev2, const Eigen::MatrixXcd &ev3,
      const double &volume, const Eigen::Matrix3d &reciprocalUnitCell,
      const Eigen::Matrix3d &epsilon,
      const Eigen::Tensor<double, 3> &bornCharges,
      const Eigen::MatrixXd &atomicPositions,
      const Eigen::Vector3i &qCoarseMesh, const int dimensionality);
 
  // this function calculates V_L for the "static" case used in the bloch->Wannier transform
  static Eigen::VectorXcd polarCorrectionPart1Static(
        const Eigen::Vector3d &q3, const Eigen::MatrixXcd &ev3,
        const double &volume, const Eigen::Matrix3d &reciprocalUnitCell,
        const Eigen::Matrix3d &epsilon, const Eigen::Tensor<double, 3> &bornCharges,
        const Eigen::MatrixXd &atomicPositions,
        const Eigen::Vector3i &qCoarseMesh, const int dimensionality);
 
  // Sets up a call to polarCorrectionPart1Static for the calculation of V_L
  // during wannier interpolation of the matrix elements
  // TODO this function is also public because of elph coupling app.
  // seems like it should be private
  Eigen::VectorXcd polarCorrectionPart1(
        const Eigen::Vector3d &q3, const Eigen::MatrixXcd &ev3);
 
  // adds the transformation of V_L to the wannier gauge
  // regardless of the use case, this calculates <psi|e^{i(G+q).r}|psi> part of polar correction
  static Eigen::Tensor<std::complex<double>, 3> polarCorrectionPart2(
        const Eigen::MatrixXcd &ev1, const Eigen::MatrixXcd &ev2, const Eigen::VectorXcd &x);
 
  Eigen::MatrixXcd precomputeQDependentPolar(BaseBandStructure &phBandStructure,
                         const bool useMinusQ = false);
 
};
 
#endif