1 files changed, 77 insertions, 0 deletions
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];
+        });
+    }
+}