Added k-space filter for RZ spectral solver (#1006)

* Added k-space filter for RZ spectral solver

* Added SpectralBinomialFilter files for RZ spectral solver

* Added RZspectral binomial filter to CMakeLists.txt

* Update Docs/source/running_cpp/parameters.rst

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Update Docs/source/running_cpp/parameters.rst

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Update Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.H

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Update Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.cpp

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Update Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* Fixed literals in SpectralBinomialFilter.cpp

* For RZ spectral, apply filter to rho old and new

* Added SpectralBinomialFilter::InitFilterArray

* For SpectralBinomialFilter, combine R and Z into one routine

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* For SpectralBinomialFilter, combine R and Z into one routine, part 2

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

* For SpectralBinomialFilter, combine R and Z into one routine, part 3

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

11 files changed, 217 insertions, 5 deletions

diff --git a/Source/FieldSolver/SpectralSolver/CMakeLists.txt b/Source/FieldSolver/SpectralSolver/CMakeLists.txt
index 34012aabd..3380ec7d7 100644
--- a/Source/FieldSolver/SpectralSolver/CMakeLists.txt
+++ b/Source/FieldSolver/SpectralSolver/CMakeLists.txt
@@ -23,6 +23,7 @@ if(WarpX_DIMS STREQUAL RZ)
         SpectralSolverRZ.cpp
         SpectralFieldDataRZ.cpp
         SpectralKSpaceRZ.cpp
+        SpectralBinomialFilter.cpp
     )
     add_subdirectory(SpectralHankelTransform)
 endif()
diff --git a/Source/FieldSolver/SpectralSolver/Make.package b/Source/FieldSolver/SpectralSolver/Make.package
index ba41619a1..04ae352be 100644
--- a/Source/FieldSolver/SpectralSolver/Make.package
+++ b/Source/FieldSolver/SpectralSolver/Make.package
@@ -11,6 +11,7 @@ ifeq ($(USE_RZ),TRUE)
   CEXE_sources += SpectralSolverRZ.cpp
   CEXE_sources += SpectralFieldDataRZ.cpp
   CEXE_sources += SpectralKSpaceRZ.cpp
+  CEXE_sources += SpectralBinomialFilter.cpp
   include $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/Make.package
 endif
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.H b/Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.H
new file mode 100644
index 000000000..9d3ab70f5
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.H
@@ -0,0 +1,48 @@
+/* Copyright 2020 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef WARPX_SPECTRAL_BINOMIAL_FILTER_H_
+#define WARPX_SPECTRAL_BINOMIAL_FILTER_H_
+
+#include "SpectralHankelTransform/HankelTransform.H"
+#include "SpectralKSpace.H"
+
+/**
+ * \brief Class that sets up binomial filtering in k space
+ *
+ * Contains the filter arrays for r and z.
+ */
+class SpectralBinomialFilter
+{
+    public:
+
+        // `KFilterArray` holds a one 1D array ("ManagedVector") that
+        // implements the filter.
+        using KFilterArray = amrex::Gpu::ManagedVector<amrex::Real>;
+
+        SpectralBinomialFilter () {};
+        void InitFilterArray (RealKVector const & kvec,
+                              amrex::Real const dels,
+                              int const npasses,
+                              bool const compensation,
+                              KFilterArray & filter);
+        void InitFilterArray (RealKVector const & kr,
+                              RealKVector const & kz,
+                              amrex::RealVect const dx,
+                              amrex::IntVect const filter_npass_each_dir,
+                              bool const compensation);
+
+        KFilterArray const & getFilterArrayR () {return filter_r;};
+        KFilterArray const & getFilterArrayZ () {return filter_z;};
+
+    protected:
+
+        KFilterArray filter_r;
+        KFilterArray filter_z;
+
+};
+
+#endif
diff --git a/Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.cpp b/Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.cpp
new file mode 100644
index 000000000..c4ecd0969
--- /dev/null
+++ b/Source/FieldSolver/SpectralSolver/SpectralBinomialFilter.cpp
@@ -0,0 +1,49 @@
+/* Copyright 2020 David Grote
+ *
+ * This file is part of WarpX.
+ *
+ * License: BSD-3-Clause-LBNL
+ */
+
+#include "SpectralBinomialFilter.H"
+#include <AMReX_REAL.H>
+
+#include <cmath>
+
+using namespace amrex::literals;
+
+/* \brief Initialize the general filter array */
+void
+SpectralBinomialFilter::InitFilterArray (HankelTransform::RealVector const & kvec,
+                                         amrex::Real const dels,
+                                         int const npasses,
+                                         bool const compensation,
+                                         KFilterArray & filter)
+{
+
+    filter.resize(kvec.size());
+
+    for (int i=0 ; i < kvec.size() ; i++) {
+        amrex::Real const ss = std::sin(0.5_rt*kvec[i]*dels);
+        amrex::Real const ss2 = ss*ss;
+        amrex::Real filt = std::pow(1._rt - ss2, npasses);
+        if (compensation) {
+            filt *= (1._rt + npasses*ss2);
+        }
+        filter[i] = filt;
+    }
+
+}
+
+/* \brief Initialize the radial and longitudinal filter arrays */
+void
+SpectralBinomialFilter::InitFilterArray (RealKVector const & kr,
+                                         RealKVector const & kz,
+                                         amrex::RealVect const dx,
+                                         amrex::IntVect const filter_npass_each_dir,
+                                         bool const compensation)
+{
+    // Note that this includes the kr values for all modes
+    InitFilterArray(kr, dx[0], filter_npass_each_dir[0], compensation, filter_r);
+    InitFilterArray(kz, dx[1], filter_npass_each_dir[1], compensation, filter_z);
+}
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
index d0e29d070..a2ab2b9cb 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.H
@@ -10,6 +10,7 @@
 #include "SpectralKSpaceRZ.H"
 #include "SpectralFieldData.H"
 #include "SpectralHankelTransform/SpectralHankelTransformer.H"
+#include "SpectralBinomialFilter.H"
 #include <AMReX_MultiFab.H>
 
 /* \brief Class that stores the fields in spectral space, and performs the
@@ -28,9 +29,11 @@ class SpectralFieldDataRZ
 #else
         using FFTplans = amrex::LayoutData<fftw_plan>;
 #endif
-        // Similarly, define the Hankel transformers for each box.
+        // Similarly, define the Hankel transformers and filter for each box.
         using MultiSpectralHankelTransformer = amrex::LayoutData<SpectralHankelTransformer>;
 
+        using BinomialFilter = amrex::LayoutData<SpectralBinomialFilter>;
+
         SpectralFieldDataRZ(const amrex::BoxArray& realspace_ba,
                             const SpectralKSpaceRZ& k_space,
                             const amrex::DistributionMapping& dm,
@@ -57,6 +60,12 @@ class SpectralFieldDataRZ
                                    amrex::MultiFab & tempHTransformedSplit,
                                    const bool is_nodal_z);
 
+        void InitFilter (amrex::IntVect const & filter_npass_each_dir, bool const compensation,
+                         SpectralKSpaceRZ const & k_space);
+
+        void ApplyFilter (int const field_index);
+        void ApplyFilter (int const field_index1, int const field_index2, int const field_index3);
+
         // Returns an array that holds the kr for all of the modes
         HankelTransform::RealVector const & getKrArray(amrex::MFIter const & mfi) const {
             return multi_spectral_hankel_transformer[mfi].getKrArray();
@@ -78,6 +87,7 @@ class SpectralFieldDataRZ
         // a cell-centered grid in real space, instead of a nodal grid
         SpectralShiftFactor zshift_FFTfromCell, zshift_FFTtoCell;
         MultiSpectralHankelTransformer multi_spectral_hankel_transformer;
+        BinomialFilter binomialfilter;
 
 };
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp
index 65d121cd3..21f84fce7 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldDataRZ.cpp
@@ -447,3 +447,80 @@ SpectralFieldDataRZ::BackwardTransform (amrex::MultiFab& field_mf_r, int const f
     }
 
 }
+
+/* \brief Initialize arrays used for filtering */
+void
+SpectralFieldDataRZ::InitFilter (amrex::IntVect const & filter_npass_each_dir, bool const compensation,
+                                 SpectralKSpaceRZ const & k_space)
+{
+    binomialfilter = BinomialFilter(multi_spectral_hankel_transformer.boxArray(),
+                                    multi_spectral_hankel_transformer.DistributionMap());
+
+    auto const & dx = k_space.getCellSize();
+    auto const & kz = k_space.getKzArray();
+
+    for (amrex::MFIter mfi(binomialfilter); mfi.isValid(); ++mfi){
+        binomialfilter[mfi].InitFilterArray(multi_spectral_hankel_transformer[mfi].getKrArray(),
+                                            kz[mfi], dx, filter_npass_each_dir, compensation);
+    }
+}
+
+/* \brief Apply K-space filtering on a scalar */
+void
+SpectralFieldDataRZ::ApplyFilter (int const field_index)
+{
+
+    for (amrex::MFIter mfi(binomialfilter); mfi.isValid(); ++mfi){
+        auto const & filter_r = binomialfilter[mfi].getFilterArrayR();
+        auto const & filter_z = binomialfilter[mfi].getFilterArrayZ();
+        auto const & filter_r_arr = filter_r.dataPtr();
+        auto const & filter_z_arr = filter_z.dataPtr();
+
+        amrex::Array4<Complex> const& fields_arr = fields[mfi].array();
+
+        int const modes = n_rz_azimuthal_modes;
+        constexpr int n_fields = SpectralFieldIndex::n_fields;
+
+        amrex::Box const& spectralspace_bx = fields[mfi].box();
+        int const nr = spectralspace_bx.length(0);
+
+        ParallelFor(spectralspace_bx, modes,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept {
+            int const ic = field_index + mode*n_fields;
+            int const ir = i + nr*mode;
+            fields_arr(i,j,k,ic) *= filter_r_arr[ir]*filter_z_arr[j];
+        });
+    }
+}
+
+/* \brief Apply K-space filtering on a vector */
+void
+SpectralFieldDataRZ::ApplyFilter (int const field_index1, int const field_index2, int const field_index3)
+{
+
+    for (amrex::MFIter mfi(binomialfilter); mfi.isValid(); ++mfi){
+        auto const & filter_r = binomialfilter[mfi].getFilterArrayR();
+        auto const & filter_z = binomialfilter[mfi].getFilterArrayZ();
+        auto const & filter_r_arr = filter_r.dataPtr();
+        auto const & filter_z_arr = filter_z.dataPtr();
+
+        amrex::Array4<Complex> const& fields_arr = fields[mfi].array();
+
+        int const modes = n_rz_azimuthal_modes;
+        constexpr int n_fields = SpectralFieldIndex::n_fields;
+
+        amrex::Box const& spectralspace_bx = fields[mfi].box();
+        int const nr = spectralspace_bx.length(0);
+
+        ParallelFor(spectralspace_bx, modes,
+        [=] AMREX_GPU_DEVICE(int i, int j, int k, int mode) noexcept {
+            int const ic1 = field_index1 + mode*n_fields;
+            int const ic2 = field_index2 + mode*n_fields;
+            int const ic3 = field_index3 + mode*n_fields;
+            int const ir = i + nr*mode;
+            fields_arr(i,j,k,ic1) *= filter_r_arr[ir]*filter_z_arr[j];
+            fields_arr(i,j,k,ic2) *= filter_r_arr[ir]*filter_z_arr[j];
+            fields_arr(i,j,k,ic3) *= filter_r_arr[ir]*filter_z_arr[j];
+        });
+    }
+}
diff --git a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H
index 0640b181a..46fb28447 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralHankelTransform/HankelTransform.H
@@ -21,7 +21,7 @@ class HankelTransform
 {
     public:
 
-        using RealVector = amrex::Gpu::ManagedDeviceVector<amrex::Real>;
+        using RealVector = amrex::Gpu::ManagedVector<amrex::Real>;
 
         // Constructor
         HankelTransform(const int hankel_order,
diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
index f26677700..d41e05a73 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.H
@@ -17,8 +17,8 @@
 // `KVectorComponent` and `SpectralShiftFactor` hold one 1D array
 // ("ManagedVector") for each box ("LayoutData"). The arrays are
 // only allocated if the corresponding box is owned by the local MPI rank.
-using KVectorComponent = amrex::LayoutData<
-                           amrex::Gpu::ManagedVector<amrex::Real> >;
+using RealKVector = amrex::Gpu::ManagedVector<amrex::Real>;
+using KVectorComponent = amrex::LayoutData< RealKVector >;
 using SpectralShiftFactor = amrex::LayoutData<
                            amrex::Gpu::ManagedVector<Complex> >;
 
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
index e1179d9cf..7a6f80278 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.H
@@ -80,6 +80,24 @@ class SpectralSolverRZ
             algorithm->pushSpectralFields(field_data);
         };
 
+        /* \brief Initialize K space filtering arrays */
+        void InitFilter (amrex::IntVect const & filter_npass_each_dir,
+                         bool const compensation)
+        {
+            field_data.InitFilter(filter_npass_each_dir, compensation, k_space);
+        };
+
+        /* \brief Apply K space filtering for a scalar */
+        void ApplyFilter (int const field_index)
+        {
+            field_data.ApplyFilter(field_index);
+        };
+
+        /* \brief Apply K space filtering for a vector */
+        void ApplyFilter (int const field_index1, int const field_index2, int const field_index3)
+        {
+            field_data.ApplyFilter(field_index1, field_index2, field_index3);
+        };
         /**
           * \brief Public interface to call the member function ComputeSpectralDivE
           * of the base class SpectralBaseAlgorithmRZ from objects of class SpectralSolverRZ
@@ -120,6 +138,7 @@ class SpectralSolverRZ
 
     private:
 
+        SpectralKSpaceRZ k_space; // Save the instance to initialize filtering
         SpectralFieldDataRZ field_data; // Store field in spectral space
                                         // and perform the Fourier transforms
         std::unique_ptr<SpectralBaseAlgorithmRZ> algorithm;
diff --git a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
index 3d7aa77b4..60e981181 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralSolverRZ.cpp
@@ -29,6 +29,7 @@ SpectralSolverRZ::SpectralSolverRZ(amrex::BoxArray const & realspace_ba,
                                    amrex::RealVect const dx, amrex::Real const dt,
                                    int const lev,
                                    bool const pml )
+    : k_space(realspace_ba, dm, dx)
 {
 
     // Initialize all structures using the same distribution mapping dm
@@ -36,7 +37,6 @@ SpectralSolverRZ::SpectralSolverRZ(amrex::BoxArray const & realspace_ba,
     // - The k space object contains info about the size of
     //   the spectral space corresponding to each box in `realspace_ba`,
     //   as well as the value of the corresponding k coordinates.
-    SpectralKSpaceRZ k_space(realspace_ba, dm, dx);
 
     // - Select the algorithm depending on the input parameters
     //   Initialize the corresponding coefficients over k space
diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index 8bce9f3d2..6c14757ea 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -70,6 +70,13 @@ namespace {
         solver.ForwardTransform(*current[2], Idx::Jz);
         solver.ForwardTransform(*rho, Idx::rho_old, 0);
         solver.ForwardTransform(*rho, Idx::rho_new, 1);
+#ifdef WARPX_DIM_RZ
+        if (WarpX::use_kspace_filter) {
+            solver.ApplyFilter(Idx::rho_old);
+            solver.ApplyFilter(Idx::rho_new);
+            solver.ApplyFilter(Idx::Jx, Idx::Jy, Idx::Jz);
+        }
+#endif
         // Advance fields in spectral space
         solver.pushSpectralFields();
         // Perform backward Fourier Transform