rpg/solvers/Block.h

#ifndef RPG_BLOCK_H
#define RPG_BLOCK_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 <pscf/solvers/BlockTmpl.h>      // base class template
 
#include <prdc/cuda/RField.h>            // member
#include <prdc/cuda/RFieldDft.h>         // member
#include <prdc/cuda/FFTBatched.h>        // member
#include <pscf/cuda/DeviceArray.h>       // member
#include <util/containers/FSArray.h>     // member
 
#include <prdc/cuda/resources.h>
 
// Forward declarations
namespace Pscf {
   template <int D> class Mesh;
   namespace Prdc {
      template <int D> class UnitCell;
      namespace Cuda {
         template <int D> class WaveList;
         template <int D> class FFT;
      }
   }
   namespace Rpg {
      template <int D> class Propagator;
   }
}
 
namespace Pscf {
namespace Rpg {
 
   using namespace Util;
   using namespace Pscf::Prdc;
   using namespace Pscf::Prdc::Cuda;

   template <int D>
   class Block : public BlockTmpl< Propagator<D>, RField<D> >
   {
 
   public:
 
      // Public type name aliases

      using BlockTmplT = BlockTmpl< Propagator<D>, RField<D> > ;

      using typename BlockTmplT::PropagatorT;

      using typename BlockTmplT::FieldT;

      using FFTT = FFT<D>;

      using WaveListT = WaveList<D>;
 
      // Public member functions

      Block();

      ~Block();


      void associate(Mesh<D> const & mesh,
                     FFT<D> const & fft,
                     UnitCell<D> const & cell,
                     WaveList<D>& waveList);

      void allocate(double ds, bool useBatchedFFT = true);

      void clearUnitCellData();

      void setLength(double newLength);

      void setKuhn(double kuhn);

      void setupSolver(RField<D> const & w);


      void stepThread(RField<D> const & qin, RField<D>& qout) const;

      void stepBead(RField<D> const & qin, RField<D>& qout) const;

      void stepFieldBead(RField<D> & q) const;

      void stepBondBead(RField<D> const & qin, RField<D>& qout) const;

      void stepHalfBondBead(RField<D> const & qin, RField<D>& qout) const;


      void computeConcentrationThread(double prefactor);

      void computeConcentrationBead(double prefactor);


      void computeStressThread(double prefactor);

      void computeStressBead(double prefactor);

      double stress(int n) const;


      double ds() const;

      int ns() const;

 
      // Inherited functions with non-dependent names (selected)
      using BlockTmplT::propagator;
      using BlockTmplT::cField;
 
   private:

      FFTBatched<D> fftBatchedAll_;

      FFTBatched<D> fftBatchedPair_;

      FSArray<double, 6> stress_;

      RField<D> expKsq_;

      RField<D> expW_;

      RField<D> expKsq2_;

      RField<D> expW2_;

      RField<D> expWInv_;

      mutable DeviceArray<cudaReal> qrPair_;

      mutable DeviceArray<cudaComplex> qkPair_;
 
      // R-grid work space (used in productAverage)
      mutable RField<D> qr_;
 
      // K-grid work space (used for FFT of q in stepBondBead)
      mutable RFieldDft<D> qk_;

      mutable DeviceArray<cudaComplex> q0kBatched_;

      mutable DeviceArray<cudaComplex> q1kBatched_;
 
      // Slices of forward and reverse propagator on a k-grid (for stress)
      mutable RFieldDft<D> q0k_;
      mutable RFieldDft<D> q1k_;

      Mesh<D> const * meshPtr_;

      FFT<D> const * fftPtr_;

      UnitCell<D> const * unitCellPtr_;

      WaveList<D> * waveListPtr_;

      IntVec<D> kMeshDimensions_;

      int kSize_;

      double ds_;

      double dsTarget_;

      int ns_;

      int nParams_;

      bool isAllocated_;

      bool hasExpKsq_;

      bool useBatchedFFT_;

      Mesh<D> const & mesh() const;

      FFT<D> const & fft() const;

      UnitCell<D> const & unitCell() const;

      WaveList<D> & waveList();

      void computeExpKsq();
 
   };
 
   // Inline member functions
 
   // Get number of contour steps.
   template <int D>
   inline int Block<D>::ns() const
   {  return ns_; }
 
   // Get contour length step size.
   template <int D>
   inline double Block<D>::ds() const
   {  return ds_; }
 
   // Get derivative of free energy w/ respect to a unit cell parameter.
   template <int D>
   inline double Block<D>::stress(int n) const
   {  return stress_[n]; }
 
   // Get Mesh by reference.
   template <int D>
   inline Mesh<D> const & Block<D>::mesh() const
   {
      UTIL_ASSERT(meshPtr_);
      return *meshPtr_;
   }
 
   // Get FFT by reference.
   template <int D>
   inline FFT<D> const & Block<D>::fft() const
   {
      UTIL_ASSERT(fftPtr_);
      return *fftPtr_;
   }
 
   // Get associated UnitCell<D> by const reference (private).
   template <int D>
   UnitCell<D> const & Block<D>::unitCell() const
   {  return *unitCellPtr_; }
 
   // Get the WaveList by non-const reference (private).
   template <int D>
   WaveList<D> & Block<D>::waveList()
   {  return *waveListPtr_; }
 
} // namespace Rpg
} // namespace Pscf
 
// Explicit instantiation declarations
namespace Pscf {
 
   extern template
   class BlockTmpl< Rpg::Propagator<1>, Prdc::Cuda::RField<1> >;
   extern template
   class BlockTmpl< Rpg::Propagator<2>, Prdc::Cuda::RField<2> >;
   extern template
   class BlockTmpl< Rpg::Propagator<3>, Prdc::Cuda::RField<3> >;
 
   namespace Rpg {
      extern template class Block<1>;
      extern template class Block<2>;
      extern template class Block<3>;
   }
}
#endif