1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
|
#include <WarpXConst.H>
#include <SpectralKSpace.H>
#include <cmath>
using namespace amrex;
using namespace Gpu;
/* \brief Initialize k space object.
*
* \param realspace_ba Box array that corresponds to the decomposition
* of the fields in real space (cell-centered ; includes guard cells)
* \param dm Indicates which MPI proc owns which box, in realspace_ba.
* \param realspace_dx Cell size of the grid in real space
*/
SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
const DistributionMapping& dm,
const RealVect realspace_dx )
{
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
realspace_ba.ixType()==IndexType::TheCellType(),
"SpectralKSpace expects a cell-centered box.");
// Store the cell size
dx = realspace_dx;
// Create the box array that corresponds to spectral space
BoxList spectral_bl; // Create empty box list
// Loop over boxes and fill the box list
for (int i=0; i < realspace_ba.size(); i++ ) {
// For local FFTs, boxes in spectral space start at 0 in
// each direction and have the same number of points as the
// (cell-centered) real space box
// TODO: this will be different for the real-to-complex FFT
// TODO: this will be different for the hybrid FFT scheme
Box realspace_bx = realspace_ba[i];
Box bx = Box( IntVect::TheZeroVector(),
realspace_bx.bigEnd() - realspace_bx.smallEnd() );
spectral_bl.push_back( bx );
}
spectralspace_ba.define( spectral_bl );
// Allocate the components of the k vector: kx, ky (only in 3D), kz
for (int i_dim=0; i_dim<AMREX_SPACEDIM; i_dim++) {
k_vec[i_dim] = getKComponent( dm, i_dim );
}
}
/* For each box, in `spectralspace_ba`, which is owned by the local MPI rank
* (as indicated by the argument `dm`), compute the values of the
* corresponding k coordinate along the dimension specified by `i_dim`
*/
KVectorComponent
SpectralKSpace::getKComponent( const DistributionMapping& dm,
const int i_dim ) const
{
// Initialize an empty ManagedVector in each box
KVectorComponent k_comp = KVectorComponent(spectralspace_ba, dm);
// Loop over boxes and allocate the corresponding ManagedVector
// for each box owned by the local MPI proc
for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
Box bx = spectralspace_ba[mfi];
ManagedVector<Real>& k = k_comp[mfi];
// Allocate k to the right size
int N = bx.length( i_dim );
k.resize( N );
// Fill the k vector
const Real dk = 2*MathConst::pi/(N*dx[i_dim]);
AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.smallEnd(i_dim) == 0,
"Expected box to start at 0, in spectral space.");
AMREX_ALWAYS_ASSERT_WITH_MESSAGE( bx.bigEnd(i_dim) == N-1,
"Expected different box end index in spectral space.");
const int mid_point = (N+1)/2;
// Fill positive values of k (FFT conventions: first half is positive)
for (int i=0; i<mid_point; i++ ){
k[i] = i*dk;
}
// Fill negative values of k (FFT conventions: second half is negative)
for (int i=mid_point; i<N; i++){
k[i] = (i-N)*dk;
}
// TODO: this will be different for the real-to-complex transform
// TODO: this will be different for the hybrid FFT scheme
}
return k_comp;
}
/* For each box, in `spectralspace_ba`, which is owned by the local MPI rank
* (as indicated by the argument `dm`), compute the values of the
* corresponding correcting "shift" factor, along the dimension
* specified by `i_dim`.
*
* (By default, we assume the FFT is done from/to a nodal grid in real space
* It the FFT is performed from/to a cell-centered grid in real space,
* a correcting "shift" factor must be applied in spectral space.)
*/
SpectralShiftFactor
SpectralKSpace::getSpectralShiftFactor( const DistributionMapping& dm,
const int i_dim,
const int shift_type ) const
{
// Initialize an empty ManagedVector in each box
SpectralShiftFactor shift_factor = SpectralShiftFactor( spectralspace_ba, dm );
// Loop over boxes and allocate the corresponding ManagedVector
// for each box owned by the local MPI proc
for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
const ManagedVector<Real>& k = k_vec[i_dim][mfi];
ManagedVector<Complex>& shift = shift_factor[mfi];
// Allocate shift coefficients
shift.resize( k.size() );
// Fill the shift coefficients
Real sign = 0;
switch (shift_type){
case ShiftType::TransformFromCellCentered: sign = -1.; break;
case ShiftType::TransformToCellCentered: sign = 1.;
}
constexpr Complex I{0,1};
for (int i=0; i<k.size(); i++ ){
shift[i] = std::exp( I*sign*k[i]*0.5*dx[i_dim] );
}
}
return shift_factor;
}
/* \brief For each box, in `spectralspace_ba`, which is owned by the local MPI
* rank (as indicated by the argument `dm`), compute the values of the
* corresponding finite-order modified k vector, along the
* dimension specified by `i_dim`
*
* The finite-order modified k vector is the spectral-space representation
* of a finite-order stencil in real space.
*
* \param n_order Order of accuracy of the stencil, in discretizing
* a spatial derivative
* \param nodal Whether the stencil is to be applied to a nodal or
staggered set of fields
*/
KVectorComponent
SpectralKSpace::getModifiedKComponent( const DistributionMapping& dm,
const int i_dim,
const int n_order,
const bool nodal ) const
{
// Initialize an empty ManagedVector in each box
KVectorComponent modified_k_comp = KVectorComponent(spectralspace_ba, dm);
// Compute real-space stencil coefficients
Vector<Real> stencil_coef = getFonbergStencilCoefficients(n_order, nodal);
// Loop over boxes and allocate the corresponding ManagedVector
// for each box owned by the local MPI proc
for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
Real delta_x = dx[i_dim];
const ManagedVector<Real>& k = k_vec[i_dim][mfi];
ManagedVector<Real>& modified_k = modified_k_comp[mfi];
// Allocate modified_k to the same size as k
modified_k.resize( k.size() );
// Fill the modified k vector
for (int i=0; i<k.size(); i++ ){
for (int n=1; n<stencil_coef.size(); n++){
if (nodal){
modified_k[i] = stencil_coef[n]* \
std::sin( k[i]*n*delta_x )/( n*delta_x );
} else {
modified_k[i] = stencil_coef[n]* \
std::sin( k[i]*(n-0.5)*delta_x )/( (n-0.5)*delta_x );
}
}
}
}
return modified_k_comp;
}
/* Returns an array of coefficients, corresponding to the weight
* of each point in a finite-difference approximation (to order `n_order`)
* of a derivative.
*
* `nodal` indicates whether this finite-difference approximation is
* taken on a nodal grid or a staggered grid.
*/
Vector<Real>
getFonbergStencilCoefficients( const int n_order, const bool nodal )
{
AMREX_ALWAYS_ASSERT_WITH_MESSAGE( n_order%2 == 0,
"n_order should be even.");
const int m = n_order/2;
Vector<Real> coefs;
coefs.resize( m+1 );
// Note: there are closed-form formula for these coefficients,
// but they result in an overflow when evaluated numerically.
// One way to avoid the overflow is to calculate the coefficients
// by recurrence.
// Coefficients for nodal (a.k.a. centered) finite-difference
if (nodal == true) {
coefs[0] = -2.; // First coefficient
for (int n=1; n<m+1; n++){ // Get the other coefficients by recurrence
coefs[n] = - (m+1-n)*1./(m+n)*coefs[n-1];
}
}
// Coefficients for staggered finite-difference
else {
Real prod = 1.;
for (int k=1; k<m+1; k++){
prod *= (m+k)*1./(4*k);
}
coefs[0] = 4*m*prod*prod; // First coefficient
for (int n=1; n<m+1; n++){ // Get the other coefficients by recurrence
coefs[n] = - ((2*n-3)*(m+1-n))*1./((2*n-1)*(m-1+n))*coefs[n-1];
}
}
return coefs;
}
|
