phonon_h0.h

#ifndef PHONON_H0_H
#define PHONON_H0_H
 
#include <cmath>
 
#include "bandstructure.h"
#include "constants.h"
#include "crystal.h"
#include "eigen.h"
#include "harmonic.h"
#include "particle.h"
#include "points.h"
#include "common_kokkos.h"
 
class PhononH0 : public HarmonicHamiltonian {
 
 public:
  PhononH0(Crystal &crystal, 
           Eigen::Tensor<double, 5> &forceConstants_,
           Eigen::Vector3i& qCoarseGrid,
           const Eigen::MatrixXd& bravaisVectors_,
           const Eigen::VectorXd& weights_, const int fcRangeType);
 
  PhononH0(const PhononH0 &that);
 
  PhononH0 &operator=(const PhononH0 &that);
 
  ~PhononH0();
 
  int getNumBands() override;
 
  Particle getParticle() override;
 
  std::tuple<Eigen::VectorXd, Eigen::MatrixXcd> diagonalize(Point &point) override;
 
  std::tuple<Eigen::VectorXd, Eigen::MatrixXcd> diagonalizeFromCoordinates(
      Eigen::Vector3d &q, const bool &withMassScaling);
  std::tuple<Eigen::VectorXd, Eigen::MatrixXcd> diagonalizeFromCoordinates(
      Eigen::Vector3d &q) override;
 
  Eigen::Tensor<std::complex<double>, 3> diagonalizeVelocity(Point &point) override;
  Eigen::Tensor<std::complex<double>, 3> diagonalizeVelocityFromCoordinates(
      Eigen::Vector3d &coordinates) override;
 
  FullBandStructure populate(Points &points, const bool &withVelocities,
                             const bool &withEigenvectors,
                             const bool isDistributed = false) override;
  FullBandStructure cpuPopulate(Points &points, const bool &withVelocities,
                                const bool &withEigenvectors, bool isDistributed = false);
  FullBandStructure kokkosPopulate(Points &points, const bool &withVelocities,
                                   const bool &withEigenvectors, const bool isDistributed = false);
  StridedComplexView3D kokkosBatchedBuildBlochHamiltonian(
      const DoubleView2D &cartesianCoordinates) override;
  std::tuple<DoubleView2D, StridedComplexView3D> kokkosBatchedDiagonalizeFromCoordinates(
      const DoubleView2D &cartesianCoordinates, const bool withMassScaling = true) override;
  std::tuple<DoubleView2D, StridedComplexView3D, ComplexView4D>
  kokkosBatchedDiagonalizeWithVelocities(
      const DoubleView2D &cartesianCoordinates) override;
  void kokkosBatchedScaleEigenvectors(StridedComplexView3D& eigenvectors);
 
  Eigen::Vector3i getCoarseGrid();
 
  int getIndexEigenvector(const int &iAt, const int &iPol) const;
 
  static int getIndexEigenvector(const int &iAt, const int &iPol, const int &nAtoms);
 
  Eigen::Matrix3d getDielectricMatrix();
 
  Eigen::Tensor<double, 3> getBornCharges();
 
  int estimateBatchSize(const bool& withVelocity) override;
 
  void printDynToHDF5(Eigen::Vector3d& qCrys);
 
  // we want to designate these as medium or short range 
  // "medium" range are fcs that come from finite difference codes
  // "short" range are fcs coming from q2r
  static const int mediumRange = 0;
  static const int shortRange = 1;
 
protected:
 
  Particle particle;
  Crystal &crystal;
 
  bool hasDielectric = false;
  int numAtoms;
  int numBands;
  double volumeUnitCell;
  Eigen::MatrixXi atomicSpecies;
  Eigen::VectorXd speciesMasses;
  Eigen::MatrixXd atomicPositions;
  Eigen::Matrix3d dielectricMatrix;
  Eigen::Tensor<double, 3> bornCharges;
  Eigen::Vector3i qCoarseGrid;
  Eigen::Matrix3d directUnitCell;
  int dimensionality;
  int fcRangeType;  // medium or short 
  
    // TODO -- when long range correction for dim=2 has been really well checked,
  // uncomment this to activate it
  bool longRange2d = false;
  double alpha;
  double alat;
 
  // the R vectors on the WS cell used in the Fourier transform
  int numBravaisVectors = 0;
  Eigen::MatrixXd bravaisVectors;
  Eigen::VectorXd weights;
 
  // container to store the D(R) matrix/the harmonic force constants
  Eigen::Tensor<double,5> mat2R;
 
  Eigen::MatrixXd gVectors;
  Eigen::Tensor<double,3> longRangeCorrection1;
  const double gMax = 14.; // cutoff for ewald summation
  const double e2 = 4; 
 
  // kokkos members:
  DoubleView1D atomicMasses_d;
  DoubleView3D longRangeCorrection1_d;
  DoubleView2D gVectors_d;
  DoubleView2D dielectricMatrix_d;
  DoubleView3D bornCharges_d;
  DoubleView2D atomicPositions_d;
  DoubleView2D bravaisVectors_d;
  DoubleView1D weights_d;
  DoubleView3D mat2R_d;
 
  // private methods, used to diagonalize the Dyn matrix
 
  Eigen::Tensor<double, 5> wsInit(const Eigen::MatrixXd &unitCell,
                                  const Eigen::Matrix3d &directUnitCell,
                                  const int& nr1Big,
                                  const int& nr2Big,
                                  const int& nr3Big);
 
 
  void addLongRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
                        const Eigen::VectorXd &q);
 
  void shortRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
                      const Eigen::VectorXd &q);
 
  void addMediumRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
                            const Eigen::VectorXd &q); 
 
  std::tuple<Eigen::VectorXd, Eigen::MatrixXcd> dynDiagonalize(
      Eigen::Tensor<std::complex<double>, 4> &dyn);
 
  void correctMediumRangeIFCs(Eigen::Tensor<double, 4> &fq, const Eigen::VectorXd &q); 
 
  double getDeviceMemoryUsage();
 
};
 
// utility functions which are related but not necessarily connected to the class
 
void setAcousticSumRule(const std::string &sumRule, Crystal& crystal,
                        const Eigen::Vector3i& qCoarseGrid,
                        Eigen::Tensor<double, 7>& forceConstants);
 
std::tuple<Eigen::Tensor<double, 5>,Eigen::MatrixXd,Eigen::VectorXd>
          reorderHarmonicForceConstants(Crystal& crystal,
                                const Eigen::Tensor<double, 7>& forceConstants,
                                Eigen::Vector3i& qCoarseGrid);
 
double wsWeight(const Eigen::VectorXd &r, const Eigen::MatrixXd &rws);
 
#endif