Branch data 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
|
