Branch data     Line data    Source code

       1                 :             : //
       2                 :             : // Class FFTPeriodicPoissonSolver
       3                 :             : //   Solves the periodic Poisson problem using Fourier transforms
       4                 :             : //   cf. https://math.mit.edu/~stevenj/fft-deriv.pdf Algorithm 5
       5                 :             : //
       6                 :             : //
       7                 :             : 
       8                 :             : #ifndef IPPL_FFT_PERIODIC_POISSON_SOLVER_H
       9                 :             : #define IPPL_FFT_PERIODIC_POISSON_SOLVER_H
      10                 :             : 
      11                 :             : #include <Kokkos_MathematicalConstants.hpp>
      12                 :             : 
      13                 :             : #include "Types/ViewTypes.h"
      14                 :             : 
      15                 :             : #include "FFT/FFT.h"
      16                 :             : #include "FieldLayout/FieldLayout.h"
      17                 :             : #include "Index/NDIndex.h"
      18                 :             : #include "Poisson.h"
      19                 :             : 
      20                 :             : namespace ippl {
      21                 :             : 
      22                 :             :     template <typename FieldLHS, typename FieldRHS>
      23                 :             :     class FFTPeriodicPoissonSolver : public Poisson<FieldLHS, FieldRHS> {
      24                 :             :         constexpr static unsigned Dim = FieldLHS::dim;
      25                 :             :         using Trhs                    = typename FieldRHS::value_type;
      26                 :             :         using mesh_type               = typename FieldRHS::Mesh_t;
      27                 :             : 
      28                 :             :     public:
      29                 :             :         using Field_t   = FieldRHS;
      30                 :             :         using FFT_t     = FFT<RCTransform, FieldRHS>;
      31                 :             :         using Complex_t = typename FFT_t::Complex_t;
      32                 :             :         using CxField_t = typename FFT_t::ComplexField;
      33                 :             :         using Layout_t  = FieldLayout<Dim>;
      34                 :             :         using Vector_t  = Vector<Trhs, Dim>;
      35                 :             : 
      36                 :             :         using Base = Poisson<FieldLHS, FieldRHS>;
      37                 :             :         using typename Base::lhs_type, typename Base::rhs_type;
      38                 :             :         using scalar_type = typename FieldLHS::Mesh_t::value_type;
      39                 :             :         using vector_type = typename FieldLHS::Mesh_t::vector_type;
      40                 :             : 
      41                 :           0 :         FFTPeriodicPoissonSolver()
      42   [ #  #  #  # ]:           0 :             : Base() {
      43                 :             :             using T = typename FieldLHS::value_type::value_type;
      44                 :             :             static_assert(std::is_floating_point<T>::value, "Not a floating point type");
      45                 :             : 
      46         [ #  # ]:           0 :             setDefaultParameters();
      47                 :           0 :         }
      48                 :             : 
      49                 :             :         FFTPeriodicPoissonSolver(lhs_type& lhs, rhs_type& rhs)
      50                 :             :             : Base(lhs, rhs) {
      51                 :             :             using T = typename FieldLHS::value_type::value_type;
      52                 :             :             static_assert(std::is_floating_point<T>::value, "Not a floating point type");
      53                 :             : 
      54                 :             :             setDefaultParameters();
      55                 :             :         }
      56                 :             : 
      57                 :             :         //~FFTPeriodicPoissonSolver() {}
      58                 :             : 
      59                 :             :         void setRhs(rhs_type& rhs) override;
      60                 :             : 
      61                 :             :         void solve() override;
      62                 :             : 
      63                 :             :     private:
      64                 :             :         void initialize();
      65                 :             : 
      66                 :             :         std::shared_ptr<FFT_t> fft_mp;
      67                 :             :         CxField_t fieldComplex_m;
      68                 :             :         CxField_t tempFieldComplex_m;
      69                 :             :         NDIndex<Dim> domain_m;
      70                 :             :         std::shared_ptr<Layout_t> layoutComplex_mp;
      71                 :             : 
      72                 :             :     protected:
      73                 :           0 :         virtual void setDefaultParameters() override {
      74                 :             :             using heffteBackend       = typename FFT_t::heffteBackend;
      75         [ #  # ]:           0 :             heffte::plan_options opts = heffte::default_options<heffteBackend>();
      76   [ #  #  #  # ]:           0 :             this->params_m.add("use_pencils", opts.use_pencils);
      77   [ #  #  #  # ]:           0 :             this->params_m.add("use_reorder", opts.use_reorder);
      78   [ #  #  #  # ]:           0 :             this->params_m.add("use_gpu_aware", opts.use_gpu_aware);
      79   [ #  #  #  # ]:           0 :             this->params_m.add("r2c_direction", 0);
      80                 :             : 
      81   [ #  #  #  #  :           0 :             switch (opts.algorithm) {
                      # ]
      82                 :           0 :                 case heffte::reshape_algorithm::alltoall:
      83   [ #  #  #  # ]:           0 :                     this->params_m.add("comm", a2a);
      84                 :           0 :                     break;
      85                 :           0 :                 case heffte::reshape_algorithm::alltoallv:
      86   [ #  #  #  # ]:           0 :                     this->params_m.add("comm", a2av);
      87                 :           0 :                     break;
      88                 :           0 :                 case heffte::reshape_algorithm::p2p:
      89   [ #  #  #  # ]:           0 :                     this->params_m.add("comm", p2p);
      90                 :           0 :                     break;
      91                 :           0 :                 case heffte::reshape_algorithm::p2p_plined:
      92   [ #  #  #  # ]:           0 :                     this->params_m.add("comm", p2p_pl);
      93                 :           0 :                     break;
      94                 :           0 :                 default:
      95   [ #  #  #  #  :           0 :                     throw IpplException("FFTPeriodicPoissonSolver::setDefaultParameters",
                   #  # ]
      96                 :             :                                         "Unrecognized heffte communication type");
      97                 :             :             }
      98                 :           0 :         }
      99                 :             :     };
     100                 :             : }  // namespace ippl
     101                 :             : 
     102                 :             : #include "PoissonSolvers/FFTPeriodicPoissonSolver.hpp"
     103                 :             : #endif