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>        // member
#include <pscf/cuda/DeviceArray.h>   // member
#include <pscf/cuda/HostDArray.h>    // member
#include <pscf/math/IntVec.h>        // member
#include <util/containers/DArray.h>  // member
#include <util/containers/GArray.h>  // member
#include <util/containers/Pair.h>    // member
 
// Forward declarations
namespace Pscf {
   template <int D> class Mesh;
   namespace Prdc {
      template <int D> class UnitCell;
   }
}
 
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();

      void sortWaves();


      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;

      DArray<int> const & sortedIds() const;

      GArray< Pair<int> > const & sortedBunches() const;

      DArray<int> const & bunchIds() const;

      IntVec<D> const & kMeshDimensions() const
      {  return kMeshDimensions_; }

      int kSize() const
      {  return kSize_; }

      int nBunch() const
      {
         UTIL_CHECK(isSorted_);  
         return nBunch_; 
      }


      bool isAllocated() const
      {  return isAllocated_; }

      bool hasMinImages() const
      {  return hasMinImages_; }

      bool hasKSq() const
      {  return hasKSq_; }

      bool hasdKSq() const
      {  return hasdKSq_; }

      bool isSorted() const
      {  return isSorted_; }

      bool isRealField() const
      {  return isRealField_; }

 
   private:
 
      // Private member variables

      DeviceArray<int> minImages_;

      mutable
      HostDArray< IntVec<D> > minImages_h_;

      RField<D> kSq_;

      DeviceArray<cudaReal> dKSq_;

      DArray< RField<D> > dKSqSlices_;

      DeviceArray<bool> implicitInverse_;

      DArray<int> sortedIds_;

      GArray< Pair<int> > sortedBunches_;

      DArray<int> bunchIds_;

      IntVec<D> kMeshDimensions_;

      int kSize_;

      int nBunch_;

      bool isAllocated_;

      bool hasMinImages_;

      mutable
      bool hasMinImages_h_;

      bool hasKSq_;

      bool hasdKSq_;

      bool isSorted_;

      bool isRealField_;

      UnitCell<D> const * unitCellPtr_;

      Mesh<D> const * meshPtr_;
 
      // Private member functions

      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_;
   }
 
   // Get the sortedIds array by const reference.
   template <int D>
   inline
   DArray<int> const & WaveList<D>::sortedIds() const
   {
      UTIL_CHECK(isSorted_);
      return sortedIds_;
   }
 
   // Get the sortedBunches array by const reference.
   template <int D> inline
   GArray< Pair<int> > const & WaveList<D>::sortedBunches() const
   {
      UTIL_CHECK(isSorted_);
      return sortedBunches_;
   }
 
   // Get the bunchIds array by const reference.
   template <int D> inline
   DArray<int> const & WaveList<D>::bunchIds() const
   {
      UTIL_CHECK(isSorted_);
      return bunchIds_;
   }
 
   // Explicit instantiation declarations
   extern template class WaveList<1>;
   extern template class WaveList<2>;
   extern template class WaveList<3>;
 
} // namespace Cuda
} // namespace Prdc
} // namespace Pscf
#endif