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
|
#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 proc (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 ("mfi.isValid")
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;
}
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
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;
}
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
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;
}
/* TODO: Documentation: point to Fonberg paper ; explain recurrence relation
*/
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 );
// 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;
}
|
