Independent k/q meshes in RTA case, improve storage of velocities in memory

src/apps/electron_wannier_transport_app.cpp

@@ -58,18 +58,45 @@ void ElectronWannierTransportApp::run(Context &context) {
     std::cout << "Done computing electronic band structure.\n" << std::endl;
   }
 
-  // Old code for using all the band structure
-  //  bool withVelocities = true;
-  //  bool withEigenvectors = true;
-  //  FullBandStructure bandStructure = electronH0.populate(
-  //      fullPoints, withVelocities, withEigenvectors);
-  //  // set the chemical potentials to zero, load temperatures
-  //  StatisticsSweep statisticsSweep(context, &bandStructure);
-
   // build/initialize the scattering matrix and the smearing
   Kokkos::Profiling::pushRegion("ETapp.setupScatteringMatrix");
-  ElScatteringMatrix scatteringMatrix(context, statisticsSweep, bandStructure,
-                                      bandStructure, phononH0, &couplingElPh);
+
+  // This is a less dense K mesh of final states for RTA, the mesh 
+  // corresponds to what we will use for the phonon sampling
+  // In the RTA only case () 
+  std::shared_ptr<ActiveBandStructure> innerBandStructure;  
+
+  if(!context.getScatteringMatrixInMemory()) {
+
+    Eigen::Vector3i qMesh = context.getQMesh();
+    Eigen::Vector3i kMesh = context.getKMesh();
+
+    if(qMesh.prod() == 0 || qMesh == kMesh ) {
+      Warning("When running RTA, if qMesh is not set, we default to using the kMesh,\n"
+              "This may be more expensive than needed.");
+
+      // use the same bandstructure
+      innerBandStructure = std::make_shared<ActiveBandStructure>(bandStructure);
+
+    } else { // check that the k and q meshes are commensurate 
+
+      if( kMesh(0)%qMesh(0) != 0 || kMesh(1)%qMesh(1) != 0 || kMesh(2)%qMesh(2) != 0 ) {
+        Error("k and q meshes must be commensurate.");
+      }
+      // if everything is ok, we setup the second different band structure 
+      Points fullPointsRTA(crystal, context.getQMesh());
+      auto tuple = ActiveBandStructure::builder(context, electronH0, fullPointsRTA);
+      innerBandStructure = std::make_shared<ActiveBandStructure>(std::get<0>(tuple));
+    }
+  } else {
+    // non-RTA case always needs matching bandstructures
+    innerBandStructure = std::make_shared<ActiveBandStructure>(bandStructure);  
+  }
+
+  ElScatteringMatrix scatteringMatrix(context, statisticsSweep, 
+                                      *innerBandStructure.get(), bandStructure,
+                                      phononH0, &couplingElPh);
+
   scatteringMatrix.setup();
   scatteringMatrix.outputToJSON("rta_el_relaxation_times.json");
   Kokkos::Profiling::popRegion();

src/bands/active_bandstructure.cpp

@@ -64,9 +64,26 @@ ActiveBandStructure::ActiveBandStructure(const Points &points_,
   numStates = numFullBands * numPoints;
   hasEigenvectors = withEigenvectors;
 
-  energies.resize(numPoints * numFullBands, 0.);
-  if(withVelocities) velocities.resize(numPoints * numFullBands * numFullBands * 3, complexZero);
-  if(withEigenvectors) eigenvectors.resize(numPoints * numFullBands * numFullBands, complexZero);
+  energies.resize(size_t(numPoints) * size_t(numFullBands), 0.);
+
+  if (withVelocities) {
+    try {
+      size_t size = size_t(numPoints) * size_t(numFullBands) * size_t(numFullBands) * size_t(3);
+      velocities.resize(size);
+    } catch(std::bad_alloc &) {
+      Error("Failed to allocate band structure velocities.\n"
+        "You are likely running out of memory.");
+    }
+  }
+  if (withEigenvectors) {
+    try {
+      size_t size = size_t(numPoints) * size_t(numFullBands) * size_t(numFullBands);
+      eigenvectors.resize(size);
+    } catch(std::bad_alloc &) {
+      Error("Failed to allocate band structure eignevectors.\n"
+        "You are likely running out of memory.");
+    }
+  }
 
   windowMethod = Window::nothing;
   buildIndices();
@@ -652,10 +669,10 @@ void ActiveBandStructure::buildOnTheFly(Window &window, Points points_,
   // this isn't a constant number.
   // Also, we look for the size of the arrays containing band structure.
   numBands = Eigen::VectorXi::Zero(numPoints);
-  int numEnStates = 0;
-  int numVelStates = 0;
-  int numEigStates = 0;
-  for (int ik = 0; ik < numPoints; ik++) {
+  size_t numEnStates = 0;
+  size_t numVelStates = 0;
+  size_t numEigStates = 0;
+  for (size_t ik = 0; ik < numPoints; ik++) {
     numBands(ik) = filteredBands(ik, 1) - filteredBands(ik, 0) + 1;
     numEnStates += numBands(ik);
     numVelStates += 3 * numBands(ik) * numBands(ik);
@@ -945,9 +962,9 @@ StatisticsSweep ActiveBandStructure::buildAsPostprocessing(
   // this isn't a constant number.
   // Also, we look for the size of the arrays containing band structure.
   numBands = Eigen::VectorXi::Zero(numPoints);
-  int numEnStates = 0;
-  int numVelStates = 0;
-  int numEigStates = 0;
+  size_t numEnStates = 0;
+  size_t numVelStates = 0;
+  size_t numEigStates = 0;
   for (int ik = 0; ik < numPoints; ik++) {
     numBands(ik) = filteredBands(ik, 1) - filteredBands(ik, 0) + 1;
     numEnStates += numBands(ik);

src/bands/bandstructure.cpp

@@ -213,6 +213,11 @@ Eigen::VectorXd FullBandStructure::getEnergies(WavevectorIndex &ik) {
   return x;
 }
 
+double FullBandStructure::getMaxEnergy() {
+  Error("Developer error: getMaxEnergy not implemented for fullbandstructure.");
+  return 0;
+}
+
 Eigen::Vector3d FullBandStructure::getGroupVelocity(StateIndex &is) {
   int stateIndex = is.get();
   auto tup = decompress2Indices(stateIndex, numPoints, numBands);

src/bands/bandstructure.h

@@ -14,10 +14,10 @@
  */
 class BaseBandStructure {
  public:
- 
+
   /** Base destructor for bandstructure class, silences warnings */
   virtual ~BaseBandStructure() = default;
- 
+
   /** Get the Particle object associated with this class
    * @return particle: a Particle object, describing e.g. whether this
    * is a phonon or electron bandStructure
@@ -111,6 +111,12 @@ class BaseBandStructure {
   virtual const double &getEnergy(StateIndex &is) = 0;
   virtual Eigen::VectorXd getEnergies(WavevectorIndex &ik) = 0;
 
+  /** Return the maximum energy of a bandstructure.
+  * Used when ph energy maximum is a cutoff for phel scattering.
+  * @return maxEnergy: maximum energy value of bandstructure in Ry
+  */
+  virtual double getMaxEnergy() = 0;
+
   /** Returns the energy of a quasiparticle from its Bloch index
    * Used for accessing the band structure in the BTE.
    * @param stateIndex: an integer index in range [0,numStates[
@@ -307,7 +313,7 @@ class FullBandStructure : public BaseBandStructure {
                     bool withEigenvectors, Points &points_,
                     bool isDistributed_ = false);
 
-  FullBandStructure();
+  //FullBandStructure();
 
   /** Copy constructor
    */
@@ -447,6 +453,12 @@ class FullBandStructure : public BaseBandStructure {
    */
   Eigen::VectorXd getEnergies(WavevectorIndex &ik) override;
 
+  /** Return the maximum energy of a bandstructure.
+  * Used when ph energy maximum is a cutoff for phel scattering.
+  * @return maxEnergy: maximum energy value of bandstructure in Ry
+  */
+  double getMaxEnergy() override;
+
   /** Returns the group velocity of a quasiparticle from its Bloch index.
    * Used for accessing the band structure in the BTE.
    * @param stateIndex: a StateIndex(is) object where 'is' is an integer index
@@ -675,7 +687,7 @@ class FullBandStructure : public BaseBandStructure {
   bool hasVelocities = false;
 
   // matrices storing the raw data
-  Matrix<double> energies;       // size(bands,points)
+  Matrix<double> energies;                    // size(bands,points)
   Matrix<std::complex<double>> velocities;    // size(3*bands^2,points)
   Matrix<std::complex<double>> eigenvectors;  // size(bands^2,points)
 

src/bte/el_scattering.h

@@ -1,7 +1,6 @@
 #ifndef EL_SCATTERING_H
 #define EL_SCATTERING_H
 
-#include "electron_h0_wannier.h"
 #include "interaction_elph.h"
 #include "phonon_h0.h"
 #include "scattering.h"

src/bte/helper_el_scattering.cpp

@@ -8,31 +8,33 @@
 HelperElScattering::HelperElScattering(BaseBandStructure &innerBandStructure_,
                                BaseBandStructure &outerBandStructure_,
                                StatisticsSweep &statisticsSweep_,
-                               const int &smearingType_, PhononH0 &h0_, InteractionElPhWan *coupling)
-    : innerBandStructure(innerBandStructure_),
-      outerBandStructure(outerBandStructure_),
-      statisticsSweep(statisticsSweep_),
-      smearingType(smearingType_),
-      h0(h0_), couplingElPhWan(coupling) {
+                               const int &smearingType_, PhononH0 &h0_,
+                               InteractionElPhWan *coupling)
+          : innerBandStructure(innerBandStructure_),
+            outerBandStructure(outerBandStructure_),
+            statisticsSweep(statisticsSweep_),
+            smearingType(smearingType_), h0(h0_),
+            couplingElPhWan(coupling) {
+
   // three conditions must be met to avoid recomputing q3
-  // 1 - q1 and q2 mesh must be the same
+  // 1 - k1 and k2 mesh must be the same
   // 2 - the mesh is gamma-centered
   // 3 - the mesh is complete (if q1 and q2 are only around 0, q3 might be
   //     at the border)
 
   Kokkos::Profiling::pushRegion("HelperElScattering");
 
   auto t1 = outerBandStructure.getPoints().getMesh();
-//  auto mesh = std::get<0>(t1);
   auto offset = std::get<1>(t1);
   storedAllQ3 = false;
 
   if (mpi->mpiHead()) {
     std::cout << "Computing phonon band structure." << std::endl;
   }
 
-  if ((&innerBandStructure == &outerBandStructure) && (offset.norm() == 0.) &&
-      innerBandStructure.hasWindow() == 0) {
+  // bandstructure are the same and are unfiltered
+  if ((outerBandStructure.getPoints() == innerBandStructure.getPoints())
+      && (offset.norm() == 0.) && innerBandStructure.hasWindow() == 0) {
 
     storedAllQ3 = true;
     storedAllQ3Case = storedAllQ3Case1;
@@ -43,19 +45,21 @@ HelperElScattering::HelperElScattering(BaseBandStructure &innerBandStructure_,
 
     fullPoints3 = std::make_unique<Points>(
         innerBandStructure.getPoints().getCrystal(), mesh2, offset2);
+
     if (mpi->mpiHead()) { // print info on memory
       double x = h0.getNumBands() * fullPoints3->getNumPoints();
       x *= sizeof(x) / pow(1024,3);
       std::cout << std::setprecision(4);
       std::cout << "Allocating " << x << " GB (per MPI process)." << std::endl;
     }
-    bool withVelocities = true;
+
+    bool withVelocities = false; // these are never needed except sym adaptive smearing
     bool withEigenvectors = true;
-    FullBandStructure bs = h0.populate(*fullPoints3, withVelocities,
-                                               withEigenvectors);
+    FullBandStructure bs = h0.populate(*fullPoints3, withVelocities, withEigenvectors);
     bandStructure3 = std::make_unique<FullBandStructure>(bs);
 
-  } else if ((&innerBandStructure == &outerBandStructure) &&
+  // bandstructures the same, no offset, filtered bands
+  } else if ((outerBandStructure.getPoints() == innerBandStructure.getPoints()) &&
              (offset.norm() == 0.) && innerBandStructure.hasWindow() != 0) {
 
     storedAllQ3 = true;
@@ -155,7 +159,7 @@ HelperElScattering::HelperElScattering(BaseBandStructure &innerBandStructure_,
 
     // build band structure
     bool withEigenvectors = true;
-    bool withVelocities = true;
+    bool withVelocities = false;
     // note: bandStructure3 stores a copy of ap3: can't pass unique_ptr
     bandStructure3 = std::make_unique<ActiveBandStructure>(
         ap3, &h0, withEigenvectors, withVelocities);
@@ -278,7 +282,10 @@ std::tuple<Eigen::Vector3d, Eigen::VectorXd, int, Eigen::MatrixXcd,
     int ik2Counter = ik2 - cacheOffset;
     Eigen::VectorXd energies3 = cacheEnergies[ik2Counter];
     Eigen::MatrixXcd eigenVectors3 = cacheEigenVectors[ik2Counter];
-    Eigen::MatrixXd v3s = cacheVelocity[ik2Counter];
+    Eigen::MatrixXd v3s;
+    if (smearingType == DeltaFunction::adaptiveGaussian) {
+      v3s = cacheVelocity[ik2Counter];
+    }
     Eigen::MatrixXd bose3Data = cacheBose[ik2Counter];
     Eigen::VectorXcd polarData = cachePolarData[ik2Counter];
     int nb3 = int(energies3.size());
@@ -295,7 +302,9 @@ void HelperElScattering::prepare(const Eigen::Vector3d &k1,
     cacheEigenVectors.resize(numPoints);
     cacheBose.resize(numPoints);
     cachePolarData.resize(numPoints);
-    cacheVelocity.resize(numPoints);
+    if (smearingType == DeltaFunction::adaptiveGaussian) {
+      cacheVelocity.resize(numPoints);
+    }
     cacheOffset = k2Indexes[0];
 
     Particle particle = h0.getParticle();
@@ -308,6 +317,8 @@ void HelperElScattering::prepare(const Eigen::Vector3d &k1,
 
       Eigen::Vector3d q3 = k2 - k1;
 
+      // TODO this whole thing should be converted to kokkos
+      // batched diagonalize from coordinates
       auto t1 = h0.diagonalizeFromCoordinates(q3);
       auto energies3 = std::get<0>(t1);
       auto eigenVectors3 = std::get<1>(t1);
@@ -319,6 +330,7 @@ void HelperElScattering::prepare(const Eigen::Vector3d &k1,
 
       for (int iCalc = 0; iCalc < statisticsSweep.getNumCalculations(); iCalc++) {
         double temp = statisticsSweep.getCalcStatistics(iCalc).temperature;
+        #pragma omp parallel for default(none) shared(bose3Data,iCalc,nb3,particle,energies3,temp)
         for (int ib3 = 0; ib3 < nb3; ib3++) {
           bose3Data(iCalc, ib3) = particle.getPopulation(energies3(ib3), temp);
         }
@@ -337,6 +349,7 @@ void HelperElScattering::prepare(const Eigen::Vector3d &k1,
             v3s(ib3, i) = v3sTmp(ib3, ib3, i).real();
           }
         }
+        cacheVelocity[ik2Counter] = v3s;
       }
 
       Eigen::VectorXcd polarData = couplingElPhWan->polarCorrectionPart1(q3, eigenVectors3);
@@ -345,7 +358,7 @@ void HelperElScattering::prepare(const Eigen::Vector3d &k1,
       cacheEigenVectors[ik2Counter] = eigenVectors3;
       cacheBose[ik2Counter] = bose3Data;
       cachePolarData[ik2Counter] = polarData;
-      cacheVelocity[ik2Counter] = v3s;
+      //cacheVelocity[ik2Counter] = v3s;
 
     }
   }

src/context.cpp

@@ -861,6 +861,8 @@ void Context::printInputSummary(const std::string &fileName) {
     if (appName.find("elPh") != std::string::npos || appName == "electronLifetimes" ||
         appName == "electronWannierTransport") {
 
+      if(qMesh.prod() != 0) std::cout << "qMesh = [" << qMesh(0) << ", " << qMesh(1) << ", " << qMesh(2) << " ]" << std::endl;
+
       std::cout << "elphFileName = " << elphFileName << std::endl;
       if (!elPhInterpolation.empty())
         std::cout << "elPhInterpolation = " << elPhInterpolation << std::endl;
@@ -928,11 +930,11 @@ void Context::printInputSummary(const std::string &fileName) {
         }
       } else {
         if(g2PlotStyle == "qFixed") {
-          std::cout << "kMesh = " << kMesh(0) << " " << kMesh(1) << " " << kMesh(2)
+          std::cout << "kMesh = [" << kMesh(0) << ", " << kMesh(1) << ", " << kMesh(2) << " ]"
                 << std::endl;
         }
         else if(g2PlotStyle == "kFixed") {
-          std::cout << "qMesh = " << qMesh(0) << " " << qMesh(1) << " " << qMesh(2)
+          std::cout << "qMesh = [" << qMesh(0) << ", " << qMesh(1) << ", " << qMesh(2) << " ]"
                 << std::endl;
         }
       }
@@ -949,7 +951,7 @@ void Context::printInputSummary(const std::string &fileName) {
       appName.find("Lifetimes") != std::string::npos) {
 
     if (appName.find("honon") != std::string::npos || appName.find("elPh") != std::string::npos) {
-      std::cout << "qMesh = " << qMesh(0) << " " << qMesh(1) << " " << qMesh(2) << std::endl;
+          std::cout << "qMesh = [" << qMesh(0) << ", " << qMesh(1) << ", " << qMesh(2) << " ]" << std::endl;
       if(!getElphFileName().empty()) {
         std::cout << "electronH0Name = " << electronH0Name << std::endl;
         std::cout << "hasSpinOrbit = " << hasSpinOrbit << std::endl;
@@ -958,7 +960,7 @@ void Context::printInputSummary(const std::string &fileName) {
     }
     if (appName.find("lectron") != std::string::npos || appName.find("elPh") != std::string::npos
                         || (appName.find("honon") != std::string::npos && !getElphFileName().empty())) {
-         std::cout << "kMesh = " << kMesh(0) << " " << kMesh(1) << " " << kMesh(2) << std::endl;
+         std::cout << "kMesh = [" << kMesh(0) << ", " << kMesh(1) << ", " << kMesh(2) << " ]" << std::endl;
     }
 
     if (!std::isnan(constantRelaxationTime))
@@ -1101,7 +1103,7 @@ void Context::printInputSummary(const std::string &fileName) {
     std::cout << "epaFileName = " << epaFileName << std::endl;
     std::cout << "electronH0Name = " << electronH0Name << std::endl;
     std::cout << "hasSpinOrbit = " << hasSpinOrbit << std::endl;
-    std::cout << "kMesh = " << kMesh(0) << " " << kMesh(1) << " " << kMesh(2)
+    std::cout << "kMesh = [" << kMesh(0) << ", " << kMesh(1) << ", " << kMesh(2) << " ]"
               << std::endl;
     if (!std::isnan(epaEnergyRange))
       std::cout << "epaEnergyRange = " << epaEnergyRange * energyRyToEv << " eV"