1 files changed, 83 insertions, 2 deletions

diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index c2cf97f30..65079428b 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -449,6 +449,39 @@ WarpX::ReadParameters ()
             fine_tag_hi = RealVect{hi};
         }
 
+        // select which particle comps to write
+        {
+            pp.queryarr("particle_plot_vars", particle_plot_vars);
+
+            if (particle_plot_vars.size() == 0)
+            {
+                if (WarpX::do_boosted_frame_diagnostic && WarpX::do_boosted_frame_particles)
+                {
+                    particle_plot_flags.resize(PIdx::nattribs + 6, 1);
+                }
+                else
+                {
+                    particle_plot_flags.resize(PIdx::nattribs, 1);
+                }
+            }
+            else
+            {
+                if (WarpX::do_boosted_frame_diagnostic && WarpX::do_boosted_frame_particles)
+                {
+                    particle_plot_flags.resize(PIdx::nattribs + 6, 0);
+                }
+                else
+                {
+                    particle_plot_flags.resize(PIdx::nattribs, 0);
+                }
+
+                for (const auto& var : particle_plot_vars)
+                {
+                    particle_plot_flags[ParticleStringNames::to_index.at(var)] = 1;
+                }
+            }
+        }
+
         pp.query("load_balance_int", load_balance_int);
         pp.query("load_balance_with_sfc", load_balance_with_sfc);
         pp.query("load_balance_knapsack_factor", load_balance_knapsack_factor);
@@ -555,8 +588,10 @@ WarpX::MakeNewLevelFromScratch (int lev, Real time, const BoxArray& new_grids,
     InitLevelData(lev, time);
 
 #ifdef WARPX_USE_PSATD
-    AllocLevelDataFFT(lev);
-    InitLevelDataFFT(lev, time);
+    if (fft_hybrid_mpi_decomposition){
+        AllocLevelDataFFT(lev);
+        InitLevelDataFFT(lev, time);
+    }
 #endif
 }
 
@@ -692,6 +727,37 @@ WarpX::AllocLevelData (int lev, const BoxArray& ba, const DistributionMapping& d
     // CKC solver requires one additional guard cell
     if (maxwell_fdtd_solver_id == 1) ngF = std::max( ngF, 1 );
 
+#ifdef WARPX_USE_PSATD
+    if (fft_hybrid_mpi_decomposition == false){
+        // All boxes should have the same number of guard cells
+        // (to avoid temporary parallel copies)
+        // Thus take the max of the required number of guards for each field
+        // Also: the number of guard cell should be enough to contain
+        // the stencil of the FFT solver. Here, this number (`ngFFT`)
+        // is determined *empirically* to be the order of the solver
+        // for nodal, and half the order of the solver for staggered.
+        IntVect ngFFT;
+        if (do_nodal) {
+            ngFFT = IntVect(AMREX_D_DECL(nox_fft, noy_fft, noz_fft));
+        } else {
+            ngFFT = IntVect(AMREX_D_DECL(nox_fft/2, noy_fft/2, noz_fft/2));
+        }
+        for (int i_dim=0; i_dim<AMREX_SPACEDIM; i_dim++ ){
+            int ng_required = ngFFT[i_dim];
+            // Get the max
+            ng_required = std::max( ng_required, ngE[i_dim] );
+            ng_required = std::max( ng_required, ngJ[i_dim] );
+            ng_required = std::max( ng_required, ngRho[i_dim] );
+            ng_required = std::max( ng_required, ngF );
+            // Set the guard cells to this max
+            ngE[i_dim] = ng_required;
+            ngJ[i_dim] = ng_required;
+            ngRho[i_dim] = ng_required;
+            ngF = ng_required;
+        }
+    }
+#endif
+
     AllocLevelMFs(lev, ba, dm, ngE, ngJ, ngRho, ngF);
 }
 
@@ -742,6 +808,21 @@ WarpX::AllocLevelMFs (int lev, const BoxArray& ba, const DistributionMapping& dm
         rho_fp[lev].reset(new MultiFab(amrex::convert(ba,IntVect::TheUnitVector()),dm,2,ngRho));
         rho_fp_owner_masks[lev] = std::move(rho_fp[lev]->OwnerMask(period));
     }
+    if (fft_hybrid_mpi_decomposition == false){
+        // Allocate and initialize the spectral solver
+        std::array<Real,3> dx = CellSize(lev);
+#if (AMREX_SPACEDIM == 3)
+        RealVect dx_vect(dx[0], dx[1], dx[2]);
+#elif (AMREX_SPACEDIM == 2)
+        RealVect dx_vect(dx[0], dx[2]);
+#endif
+        // Get the cell-centered box, with guard cells
+        BoxArray realspace_ba = ba;  // Copy box
+        realspace_ba.enclosedCells().grow(ngE); // cell-centered + guard cells
+        // Define spectral solver
+        spectral_solver_fp[lev].reset( new SpectralSolver( realspace_ba, dm,
+            nox_fft, noy_fft, noz_fft, do_nodal, dx_vect, dt[lev] ) );
+    }
 #endif
 
     //