cuda/WaveList.h

#ifndef PRDC_CUDA_WAVE_LIST_H
#define PRDC_CUDA_WAVE_LIST_H
 
/*
* PSCF - Polymer Self-Consistent Field
*
* Copyright 2015 - 2025, The Regents of the University of Minnesota
* Distributed under the terms of the GNU General Public License.
*/
 
#include <prdc/cuda/RField.h>
#include <prdc/crystal/UnitCell.h>
 
#include <pscf/mesh/Mesh.h>
#include <pscf/math/IntVec.h>
#include <pscf/cuda/DeviceArray.h>
#include <pscf/cuda/HostDArray.h>
 
#include <util/containers/DArray.h>
 
namespace Pscf {
namespace Prdc {
namespace Cuda {
 
   using namespace Util;

   template <int D>
   class WaveList
   {
   public:

      WaveList(bool isRealField = true);

      ~WaveList();

      void allocate(Mesh<D> const & m, UnitCell<D> const & c);

      void clearUnitCellData();

      void computeMinimumImages();

      void computeKSq();

      void computedKSq();

      DeviceArray<int> const & minImages_d() const;

      HostDArray< IntVec<D> > const & minImages_h() const;

      RField<D> const & kSq() const;

      RField<D> const & dKSq(int i) const;

      DeviceArray<bool> const & implicitInverse() const;

      bool isAllocated() const
      {  return isAllocated_; }

      bool hasMinImages() const
      {  return hasMinImages_; }

      bool hasKSq() const
      {  return hasKSq_; }

      bool hasdKSq() const
      {  return hasdKSq_; }
 
      bool isRealField() const
      {  return isRealField_; }
 
   private:

      DeviceArray<int> minImages_;

      mutable
      HostDArray< IntVec<D> > minImages_h_;

      RField<D> kSq_;

      DeviceArray<cudaReal> dKSq_;

      DArray<RField<D> > dKSqSlices_;

      DeviceArray<bool> implicitInverse_;
 
      IntVec<D> kMeshDimensions_;
 
      int kSize_;

      bool isAllocated_;

      bool hasMinImages_;

      mutable
      bool hasMinImages_h_;

      bool hasKSq_;

      bool hasdKSq_;

      bool isRealField_;

      UnitCell<D> const * unitCellPtr_;

      Mesh<D> const * meshPtr_;

      UnitCell<D> const & unitCell() const
      {  return *unitCellPtr_; }

      Mesh<D> const & mesh() const
      {  return *meshPtr_; }
 
   };
 
   // Get the array of minimum images on the device by reference.
   template <int D>
   inline DeviceArray<int> const & WaveList<D>::minImages_d() const
   {
      UTIL_CHECK(hasMinImages_);
      return minImages_;
   }
 
   // Get the kSq array on the device by reference.
   template <int D>
   inline RField<D> const & WaveList<D>::kSq() const
   {
      UTIL_CHECK(hasKSq_);
      return kSq_;
   }
 
   // Get a slice of the dKSq array on the device by reference.
   template <int D>
   inline RField<D> const & WaveList<D>::dKSq(int i) const
   {
      UTIL_CHECK(hasdKSq_);
      return dKSqSlices_[i];
   }
 
   // Get the implicitInverse array by reference.
   template <int D>
   inline DeviceArray<bool> const & WaveList<D>::implicitInverse() const
   {
      UTIL_CHECK(isAllocated_);
      UTIL_CHECK(isRealField_);
      return implicitInverse_;
   }
 
   #ifndef PRDC_CUDA_WAVE_LIST_TPP
   // Suppress implicit instantiation
   extern template class WaveList<1>;
   extern template class WaveList<2>;
   extern template class WaveList<3>;
   #endif
 
} // Cuda
} // Prdc
} // Pscf
#endif