Change function interface

1 files changed, 54 insertions, 52 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
index 76e8aef15..3c600221e 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp
@@ -11,7 +11,7 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
     // Store the cell size
     dx = realspace_dx;
 
-  // Create the box array that corresponds to spectral space
+    // 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++ ) {
@@ -23,66 +23,68 @@ SpectralKSpace::SpectralKSpace( const BoxArray& realspace_ba,
     }
     spectralspace_ba.define( spectral_bl );
 
-    // Allocate the 1D vectors
-    kx_vec = SpectralKVector( spectralspace_ba, dm );
-#if (AMREX_SPACEDIM == 3)
-    ky_vec = SpectralKVector( spectralspace_ba, dm );
-#endif
-    kz_vec = SpectralKVector( spectralspace_ba, dm );
-    // Initialize the values on each box
-    for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
-        Box bx = spectralspace_ba[mfi];
-        AllocateAndFillKvector( kx_vec[mfi], bx, dx, 0 );
-#if (AMREX_SPACEDIM == 3)
-        AllocateAndFillKvector( ky_vec[mfi], bx, dx, 1 );
-        AllocateAndFillKvector( kz_vec[mfi], bx, dx, 2 );
-#else
-        AllocateAndFillKvector( kz_vec[mfi], bx, dx, 1 );
-#endif
+    // 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] = AllocateAndFillKVector( dm, i_dim );
     }
 }
 
-void
-AllocateAndFillKvector( ManagedVector<Real>& k, const Box& bx,
-                        const RealVect dx, const int i_dim )
+SpectralKVector&
+SpectralKSpace::AllocateAndFillKVector( const DistributionMapping& dm, const int i_dim ) const
 {
-    // Alllocate k to the right size
-    int N = bx.length( i_dim );
-    k.resize( N );
+    // Initialize an empty vector in each box
+    SpectralKVector k_vec = SpectralKVector(spectralspace_ba, dm);
+    // Loop over boxes
+    for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
+        Box bx = spectralspace_ba[mfi];
+        ManagedVector<Real>& k = k_vec[mfi];
 
-    // 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
-    // TODO: For real-to-complex,
+        // 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_vec;
 }
 
-void
-ComputeModifiedKVector( ManagedVector<Real>& modified_k,
-                        const ManagedVector<Real>& k,
-                        const Box& bx, const Real dx, const int norder )
+SpectralKVector&
+SpectralKSpace::AllocateAndFillModifiedKVector(
+        const DistributionMapping& dm, const int i_dim, const int order ) const
 {
-    // Allocate modified_k to the right size
-    int N = k.size();
-    modified_k.resize( N );
+    // Initialize an empty vector in each box
+    SpectralKVector modified_k_vec = SpectralKVector( 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_vec[mfi];
 
-    // For now, this simply copies the infinite order k
-    for (int i=0; i<N; i++ ){
-        modified_k[i] = k[i];
-    }
+        // 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_vec;
 }