Line data Source code
1 : //
2 : // Class FFTOpenPoissonSolver
3 : // FFT-based Poisson Solver for open boundaries.
4 : // Solves laplace(phi) = -rho, and E = -grad(phi).
5 : //
6 : //
7 :
8 : #ifndef IPPL_FFT_OPEN_POISSON_SOLVER_H_
9 : #define IPPL_FFT_OPEN_POISSON_SOLVER_H_
10 :
11 : #include <Kokkos_MathematicalConstants.hpp>
12 : #include <Kokkos_MathematicalFunctions.hpp>
13 :
14 : #include "Types/Vector.h"
15 :
16 : #include "Utility/IpplException.h"
17 : #include "Utility/IpplTimings.h"
18 :
19 : #include "Field/Field.h"
20 :
21 : #include "Communicate/Archive.h"
22 : #include "FFT/FFT.h"
23 : #include "Field/HaloCells.h"
24 : #include "FieldLayout/FieldLayout.h"
25 : #include "Meshes/UniformCartesian.h"
26 : #include "Poisson.h"
27 :
28 : namespace ippl {
29 : namespace detail {
30 :
31 : /*!
32 : * Access a view that either contains a vector field or a scalar field
33 : * in such a way that the correct element access is determined at compile
34 : * time, reducing the number of functions needed to achieve the same
35 : * behavior for both kinds of fields
36 : * @tparam tensorRank indicates whether scalar, vector, or matrix field
37 : * @tparam - the view type
38 : */
39 : template <int tensorRank, typename>
40 : struct ViewAccess;
41 :
42 : template <typename View>
43 : struct ViewAccess<2, View> {
44 0 : KOKKOS_INLINE_FUNCTION constexpr static auto& get(View&& view, unsigned dim1,
45 : unsigned dim2, size_t i, size_t j,
46 : size_t k) {
47 0 : return view(i, j, k)[dim1][dim2];
48 : }
49 : };
50 :
51 : template <typename View>
52 : struct ViewAccess<1, View> {
53 0 : KOKKOS_INLINE_FUNCTION constexpr static auto& get(View&& view, unsigned dim1,
54 : [[maybe_unused]] unsigned dim2,
55 : size_t i, size_t j, size_t k) {
56 0 : return view(i, j, k)[dim1];
57 : }
58 : };
59 :
60 : template <typename View>
61 : struct ViewAccess<0, View> {
62 0 : KOKKOS_INLINE_FUNCTION constexpr static auto& get(View&& view,
63 : [[maybe_unused]] unsigned dim1,
64 : [[maybe_unused]] unsigned dim2,
65 : size_t i, size_t j, size_t k) {
66 0 : return view(i, j, k);
67 : }
68 : };
69 : } // namespace detail
70 :
71 : template <typename FieldLHS, typename FieldRHS>
72 : class FFTOpenPoissonSolver : public Poisson<FieldLHS, FieldRHS> {
73 : constexpr static unsigned Dim = FieldLHS::dim;
74 : using Trhs = typename FieldRHS::value_type;
75 : using mesh_type = typename FieldLHS::Mesh_t;
76 : using Tg = typename FieldLHS::value_type::value_type;
77 :
78 : public:
79 : // type of output
80 : using Base = Poisson<FieldLHS, FieldRHS>;
81 :
82 : // types for LHS and RHS
83 : using typename Base::lhs_type, typename Base::rhs_type;
84 :
85 : // define a type for the 3 dimensional real to complex Fourier transform
86 : typedef FFT<RCTransform, FieldRHS> FFT_t;
87 :
88 : // enum type for the algorithm
89 : enum Algorithm {
90 : HOCKNEY = 0b01,
91 : VICO = 0b10,
92 : BIHARMONIC = 0b11,
93 : DCT_VICO = 0b100
94 : };
95 :
96 : // define a type for a 3 dimensional field (e.g. charge density field)
97 : // define a type of Field with integers to be used for the helper Green's function
98 : // also define a type for the Fourier transformed complex valued fields
99 : // define matrix and matrix field types for the Hessian
100 : typedef FieldRHS Field_t;
101 : typedef typename FieldLHS::Centering_t Centering;
102 : typedef Field<int, Dim, mesh_type, Centering> IField_t;
103 : typedef Field<Tg, Dim, mesh_type, Centering> Field_gt;
104 : typedef Field<Kokkos::complex<Tg>, Dim, mesh_type, Centering> CxField_gt;
105 : typedef typename FFT_t::ComplexField CxField_t;
106 : typedef Vector<Trhs, Dim> Vector_t;
107 : typedef typename mesh_type::matrix_type Matrix_t;
108 : typedef Field<Matrix_t, Dim, mesh_type, Centering> MField_t;
109 :
110 : // define type for field layout
111 : typedef FieldLayout<Dim> FieldLayout_t;
112 :
113 : // type for communication buffers
114 : using memory_space = typename FieldLHS::memory_space;
115 : using buffer_type = mpi::Communicator::buffer_type<memory_space>;
116 :
117 : // types of mesh and mesh spacing
118 : using vector_type = typename mesh_type::vector_type;
119 : using scalar_type = typename mesh_type::value_type;
120 :
121 : // constructor and destructor
122 : FFTOpenPoissonSolver();
123 : FFTOpenPoissonSolver(rhs_type& rhs, ParameterList& params);
124 : FFTOpenPoissonSolver(lhs_type& lhs, rhs_type& rhs, ParameterList& params);
125 0 : ~FFTOpenPoissonSolver() = default;
126 :
127 : // override the setRhs function of the Solver class
128 : // since we need to call initializeFields()
129 : void setRhs(rhs_type& rhs) override;
130 :
131 : // allows user to set gradient of phi = Efield instead of spectral
132 : // calculation of Efield (which uses FFTs)
133 : void setGradFD();
134 :
135 : // solve the Poisson equation using FFT;
136 : // more specifically, compute the scalar potential given a density field rho using
137 : void solve() override;
138 :
139 : // override getHessian to return Hessian field if flag is on
140 0 : MField_t* getHessian() override {
141 0 : bool hessian = this->params_m.template get<bool>("hessian");
142 0 : if (!hessian) {
143 0 : throw IpplException(
144 : "FFTOpenPoissonSolver::getHessian()",
145 : "Cannot call getHessian() if 'hessian' flag in ParameterList is false");
146 : }
147 0 : return &hess_m;
148 : }
149 :
150 : // compute standard Green's function
151 : void greensFunction();
152 :
153 : // function called in the constructor to initialize the fields
154 : void initializeFields();
155 :
156 : // communication used for multi-rank Vico-Greengard's Green's function
157 : void communicateVico(Vector<int, Dim> size, typename CxField_gt::view_type view_g,
158 : const ippl::NDIndex<Dim> ldom_g, const int nghost_g,
159 : typename Field_t::view_type view, const ippl::NDIndex<Dim> ldom,
160 : const int nghost);
161 :
162 : // needed for improved Vico-Greengard (DCT VICO)
163 : void communicateVico(Vector<int, Dim> size, typename Field_t::view_type view_g,
164 : const ippl::NDIndex<Dim> ldom_g, const int nghost_g,
165 : typename Field_t::view_type view, const ippl::NDIndex<Dim> ldom,
166 : const int nghost);
167 :
168 : private:
169 : // create a field to use as temporary storage
170 : // references to it can be created to make the code where it is used readable
171 : Field_t storage_field;
172 :
173 : Field_t& rho2_mr =
174 : storage_field; // the charge-density field with mesh doubled in each dimension
175 : Field_t& grn_mr = storage_field; // the Green's function
176 :
177 : // rho2tr_m is the Fourier transformed charge-density field
178 : // domain3_m and mesh3_m are used
179 : CxField_t rho2tr_m;
180 :
181 : // grntr_m is the Fourier transformed Green's function
182 : // domain3_m and mesh3_m are used
183 : CxField_t grntr_m;
184 :
185 : // temp_m field for the E-field computation
186 : CxField_t temp_m;
187 :
188 : // fields that facilitate the calculation in greensFunction()
189 : IField_t grnIField_m[Dim];
190 :
191 : // hessian matrix field (only if hessian parameter is set)
192 : MField_t hess_m;
193 :
194 : // the FFT object
195 : std::unique_ptr<FFT_t> fft_m;
196 :
197 : // mesh and layout objects for rho_m (RHS)
198 : mesh_type* mesh_mp;
199 : FieldLayout_t* layout_mp;
200 :
201 : // mesh and layout objects for rho2_m
202 : std::unique_ptr<mesh_type> mesh2_m;
203 : std::unique_ptr<FieldLayout_t> layout2_m;
204 :
205 : // mesh and layout objects for the Fourier transformed Complex fields
206 : std::unique_ptr<mesh_type> meshComplex_m;
207 : std::unique_ptr<FieldLayout_t> layoutComplex_m;
208 :
209 : // domains for the various fields
210 : NDIndex<Dim> domain_m; // original domain, gridsize
211 : NDIndex<Dim> domain2_m; // doubled gridsize (2*Nx,2*Ny,2*Nz)
212 : NDIndex<Dim> domainComplex_m; // field for the complex values of the RC transformation
213 :
214 : // mesh spacing and mesh size
215 : vector_type hr_m;
216 : Vector<int, Dim> nr_m;
217 :
218 : // string specifying algorithm: Hockney or Vico-Greengard
219 : std::string alg_m;
220 :
221 : // members for Vico-Greengard
222 : CxField_gt grnL_m;
223 :
224 : std::unique_ptr<FFT<CCTransform, CxField_gt>> fft4n_m;
225 :
226 : std::unique_ptr<mesh_type> mesh4_m;
227 : std::unique_ptr<FieldLayout_t> layout4_m;
228 :
229 : NDIndex<Dim> domain4_m;
230 :
231 : // members for improved Vico-Greengard (DCT VICO)
232 : Field_t grn2n1_m;
233 :
234 : std::unique_ptr<FFT<Cos1Transform, Field_t>> fft2n1_m;
235 :
236 : std::unique_ptr<mesh_type> mesh2n1_m;
237 : std::unique_ptr<FieldLayout_t> layout2n1_m;
238 :
239 : NDIndex<Dim> domain2n1_m;
240 :
241 : // bool indicating whether we want gradient of solution to calculate E field
242 : bool isGradFD_m;
243 :
244 : // buffer for communication
245 : detail::FieldBufferData<Trhs> fd_m;
246 :
247 : protected:
248 0 : virtual void setDefaultParameters() override {
249 : using heffteBackend = typename FFT_t::heffteBackend;
250 0 : heffte::plan_options opts = heffte::default_options<heffteBackend>();
251 0 : this->params_m.add("use_pencils", opts.use_pencils);
252 0 : this->params_m.add("use_reorder", opts.use_reorder);
253 0 : this->params_m.add("use_gpu_aware", opts.use_gpu_aware);
254 0 : this->params_m.add("r2c_direction", 0);
255 :
256 0 : switch (opts.algorithm) {
257 0 : case heffte::reshape_algorithm::alltoall:
258 0 : this->params_m.add("comm", a2a);
259 0 : break;
260 0 : case heffte::reshape_algorithm::alltoallv:
261 0 : this->params_m.add("comm", a2av);
262 0 : break;
263 0 : case heffte::reshape_algorithm::p2p:
264 0 : this->params_m.add("comm", p2p);
265 0 : break;
266 0 : case heffte::reshape_algorithm::p2p_plined:
267 0 : this->params_m.add("comm", p2p_pl);
268 0 : break;
269 0 : default:
270 0 : throw IpplException("FFTOpenPoissonSolver::setDefaultParameters",
271 : "Unrecognized heffte communication type");
272 : }
273 :
274 0 : this->params_m.add("algorithm", HOCKNEY);
275 0 : this->params_m.add("hessian", false);
276 0 : }
277 : };
278 : } // namespace ippl
279 :
280 : #include "PoissonSolvers/FFTOpenPoissonSolver.hpp"
281 : #endif
|
