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
|
#include <WarpXConst.H>
#include <SpectralKSpace.H>
#include <cmath>
using namespace amrex;
using namespace Gpu;
SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
const DistributionMapping& dm,
const RealVect realspace_dx )
{
// 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, each box in spectral space starts at 0 in each direction
// and has the same number of points as the real space box (including guard cells)
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 );
}
}
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
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.");
// Fill positive values of k (FFT conventions: first half is positive)
for (int i=0; i<(N+1)/2; i++ ){
k[i] = i*dk;
}
// Fill negative values of k (FFT conventions: second half is negative)
for (int i=(N+1)/2; i<N; i++){
k[i] = (N-i)*dk;
}
// TODO: This should be quite different for the hybrid spectral code:
// In that case we should take into consideration the actual indices of the box
// and distinguish the size of the local box and that of the global FFT
// This will also be different for the real-to-complex transform
}
return k_comp;
}
KVectorComponent
SpectralKSpace::getModifiedKComponent(
const DistributionMapping& dm, const int i_dim, const int order ) const
{
// Initialize an empty ManagedVector in each box
KVectorComponent modified_k_comp = KVectorComponent( 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<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 now, this simply copies the infinite order k
// TODO: Use the formula for finite-order modified k vector
modified_k[i] = k[i];
}
}
return modified_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::CenteredToNodal: sign = -1.; break;
case ShiftType::NodalToCentered: 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;
}
|