Use FFT to calculate derivative in FFT solvers

gyselax · Jan 22, 2024 · 1fad537 · 1fad537
1 parent 31cac54
commit 1fad537
Show file tree

Hide file tree

Showing 12 changed files with 198 additions and 55 deletions.
diff --git a/simulations/geometryXVx/bump_on_tail/bumpontail_fft.cpp b/simulations/geometryXVx/bump_on_tail/bumpontail_fft.cpp
@@ -225,7 +225,7 @@ int main(int argc, char** argv)
     DFieldVx const quadrature_coeffs = neumann_spline_quadrature_coefficients(gridvx, builder_vx);
     Quadrature<IDimVx> const integrate_v(quadrature_coeffs);
     ChargeDensityCalculator rhs(integrate_v);
-    FftPoissonSolver const poisson(builder_x, spline_x_evaluator, rhs);
+    FftPoissonSolver const poisson(rhs);
 
     PredCorr const predcorr(vlasov, poisson);
 

diff --git a/simulations/geometryXVx/landau/landau_fft.cpp b/simulations/geometryXVx/landau/landau_fft.cpp
@@ -229,7 +229,7 @@ int main(int argc, char** argv)
     DFieldVx const quadrature_coeffs = neumann_spline_quadrature_coefficients(gridvx, builder_vx);
     Quadrature<IDimVx> const integrate_v(quadrature_coeffs);
     ChargeDensityCalculator rhs(integrate_v);
-    FftPoissonSolver const poisson(builder_x, spline_x_evaluator, rhs);
+    FftPoissonSolver const poisson(rhs);
 
     PredCorr const predcorr(vlasov, poisson);
 

diff --git a/simulations/geometryXVx/sheath/sheath.cpp b/simulations/geometryXVx/sheath/sheath.cpp
@@ -319,7 +319,7 @@ int main(int argc, char** argv)
     ChargeDensityCalculator rhs(integrate_v);
 #ifdef PERIODIC_RDIMX
     ddc::init_fourier_space<RDimX>(ddc::select<IDimX>(meshSpXVx));
-    FftPoissonSolver const poisson(builder_x, spline_x_evaluator, rhs);
+    FftPoissonSolver const poisson(rhs);
 #else
     FemNonPeriodicPoissonSolver const poisson(builder_x, spline_x_evaluator, rhs);
 #endif

diff --git a/simulations/geometryXYVxVy/landau/landau4d_fft.cpp b/simulations/geometryXYVxVy/landau/landau4d_fft.cpp
@@ -251,7 +251,7 @@ int main(int argc, char** argv)
     ddc::init_fourier_space<RDimX, RDimY>(ddc::select<IDimX, IDimY>(meshSpXYVxVy));
 
     ChargeDensityCalculator const rhs(builder_vxvy, spline_vxvy_evaluator);
-    FftPoissonSolver const poisson(builder_xy, spline_xy_evaluator, rhs);
+    FftPoissonSolver const poisson(rhs);
 
     // Create predcorr operator
     PredCorr const predcorr(vlasov, poisson);

diff --git a/src/geometryXVx/poisson/README.md b/src/geometryXVx/poisson/README.md
@@ -27,5 +27,3 @@ The Poisson equation can be solved with a variety of different methods. Here we
 -   FemPeriodicPoissonSolver
 
 These classes return the electric potential $\phi$ and the electric field $\frac{d \phi}{dx}$.
-
-The FftPoissonSolver does not calculate the electric field using the Fourier modes. Rather it uses a spline interpolation to approximate this value. This interpolation is calculated by the operator ElectricField.
diff --git a/src/geometryXVx/poisson/fftpoissonsolver.cpp b/src/geometryXVx/poisson/fftpoissonsolver.cpp
@@ -17,12 +17,8 @@
 
 #include "fftpoissonsolver.hpp"
 
-FftPoissonSolver::FftPoissonSolver(
-        SplineXBuilder const& spline_x_builder,
-        SplineEvaluator<BSplinesX> const& spline_x_evaluator,
-        IChargeDensityCalculator const& compute_rho)
+FftPoissonSolver::FftPoissonSolver(IChargeDensityCalculator const& compute_rho)
     : m_compute_rho(compute_rho)
-    , m_electric_field(spline_x_builder, spline_x_evaluator)
 {
 }
 
@@ -57,7 +53,7 @@ void FftPoissonSolver::operator()(
                 ddc::kwArgs_fft {norm});
 
     // Solve Poisson's equation -d2Phi/dx2 = rho
-    // in Fourier space as -kx*kx*FFT(Phi)=FFT(rho))
+    //   in Fourier space as kx*kx*FFT(Phi)=FFT(rho)
 
     // First, 0 mode of Phi is set to 0 to avoid divergency. Note: to allow writing in the GPU memory, starting a device kernel is necessary which is performed by iterating on a 0D domain (single element).
     ddc::for_each(
@@ -74,18 +70,34 @@ void FftPoissonSolver::operator()(
                 intermediate_chunk(ikx)
                         = intermediate_chunk(ikx) / (coordinate(ikx) * coordinate(ikx));
             });
+
+    // Find the electric field in Fourier space
+    // FFT(efield) = -kx*i*FFT(Phi)
+    Kokkos::Profiling::pushRegion("ElectricField");
+    Kokkos::complex<double> imaginary_unit(0.0, 1.0);
+    device_t<ddc::Chunk<Kokkos::complex<double>, IDomainFx>> fourier_efield_alloc(k_mesh);
+    ddc::ChunkSpan fourier_efield = fourier_efield_alloc.span_view();
+
+    ddc::for_each(
+            ddc::policies::parallel_device,
+            k_mesh,
+            KOKKOS_LAMBDA(IndexFx const ikx) {
+                fourier_efield(ikx) = -imaginary_unit * coordinate(ikx) * intermediate_chunk(ikx);
+            });
+
+    // Perform the inverse 1D FFT of the solution to deduce the electric field
+    ddc::
+            ifft(Kokkos::DefaultExecutionSpace(),
+                 electric_field,
+                 fourier_efield,
+                 ddc::kwArgs_fft {norm});
+    Kokkos::Profiling::popRegion();
+
     // Perform the inverse 1D FFT of the solution to deduce the electrostatic potential
     ddc::
             ifft(Kokkos::DefaultExecutionSpace(),
                  electrostatic_potential.span_view(),
                  intermediate_chunk,
                  ddc::kwArgs_fft {norm});
-
-    // Compute efield = -dPhi/dx where Phi is the electrostatic potential
-    DFieldX electrostatic_potential_host(x_dom);
-    DFieldX electric_field_host(x_dom);
-    ddc::deepcopy(electrostatic_potential_host, electrostatic_potential);
-    m_electric_field(electric_field_host, electrostatic_potential_host);
-    ddc::deepcopy(electric_field, electric_field_host);
     Kokkos::Profiling::popRegion();
 }
diff --git a/src/geometryXVx/poisson/fftpoissonsolver.hpp b/src/geometryXVx/poisson/fftpoissonsolver.hpp
@@ -25,20 +25,13 @@ class FftPoissonSolver : public IPoissonSolver
 {
     IChargeDensityCalculator const& m_compute_rho;
 
-    ElectricField m_electric_field;
-
 public:
     /**
      * Construct the FftPoissonSolver operator.
      *
      * @param compute_rho The operator which calculates the charge density, the right hand side of the equation.
-     * @param spline_x_builder A spline builder which calculates the coefficients of a spline representation.
-     * @param spline_x_evaluator A spline evaluator which provides the value of a spline representation from its coefficients.
      */
-    FftPoissonSolver(
-            SplineXBuilder const& spline_x_builder,
-            SplineEvaluator<BSplinesX> const& spline_x_evaluator,
-            IChargeDensityCalculator const& compute_rho);
+    FftPoissonSolver(IChargeDensityCalculator const& compute_rho);
 
     ~FftPoissonSolver() override = default;
 

diff --git a/src/geometryXYVxVy/poisson/fftpoissonsolver.cpp b/src/geometryXYVxVy/poisson/fftpoissonsolver.cpp
@@ -16,12 +16,8 @@
 
 #include "fftpoissonsolver.hpp"
 
-FftPoissonSolver::FftPoissonSolver(
-        SplineXYBuilder const& spline_xy_builder,
-        SplineXYEvaluator const& spline_xy_evaluator,
-        IChargeDensityCalculator const& compute_rho)
+FftPoissonSolver::FftPoissonSolver(IChargeDensityCalculator const& compute_rho)
     : m_compute_rho(compute_rho)
-    , m_electric_field(spline_xy_builder, spline_xy_evaluator)
 {
 }
 
@@ -70,14 +66,46 @@ void FftPoissonSolver::operator()(
         }
     });
 
+    // Find the electric field in Fourier space
+    // FFT(efield) = [-kx*i*FFT(Phi), -ky*i*FFT(Phi)]
+    Kokkos::Profiling::pushRegion("ElectricField");
+    Kokkos::complex<double> imaginary_unit(0.0, 1.0);
+    ddc::Chunk<Kokkos::complex<double>, IDomainFxFy> fourier_efield
+            = ddc::Chunk(k_mesh, ddc::HostAllocator<Kokkos::complex<double>>());
+
+    // Calculate x component
+    ddc::for_each(k_mesh, [&](IndexFxFy const ikxky) {
+        IndexFx const ikx = ddc::select<IDimFx>(ikxky);
+        fourier_efield(ikxky) = -imaginary_unit * coordinate(ikx) * intermediate_chunk(ikxky);
+    });
+
+    // Perform the inverse 1D FFT of the solution to deduce the electric field
+    ddc::
+            ifft(Kokkos::DefaultHostExecutionSpace(),
+                 electric_field_x.span_view(),
+                 fourier_efield.span_view(),
+                 ddc::kwArgs_fft {.normalization = norm});
+
+    // Calculate y component
+    ddc::for_each(k_mesh, [&](IndexFxFy const ikxky) {
+        IndexFy const iky = ddc::select<IDimFy>(ikxky);
+        fourier_efield(ikxky) = -imaginary_unit * coordinate(iky) * intermediate_chunk(ikxky);
+    });
+
+    // Perform the inverse 1D FFT of the solution to deduce the electric field
+    ddc::
+            ifft(Kokkos::DefaultHostExecutionSpace(),
+                 electric_field_y.span_view(),
+                 fourier_efield.span_view(),
+                 ddc::kwArgs_fft {.normalization = norm});
+    Kokkos::Profiling::popRegion();
+
     // Perform the inverse 1D FFT of the solution to deduce the electrostatic potential
     ddc::
             ifft(Kokkos::DefaultHostExecutionSpace(),
                  electrostatic_potential.span_view(),
                  intermediate_chunk.span_view(),
                  (ddc::kwArgs_fft) {.normalization = norm});
 
-    // Compute efield = -dPhi/dx where Phi is the electrostatic potential
-    m_electric_field(electric_field_x, electric_field_y, electrostatic_potential);
     Kokkos::Profiling::popRegion();
 }
diff --git a/src/geometryXYVxVy/poisson/fftpoissonsolver.hpp b/src/geometryXYVxVy/poisson/fftpoissonsolver.hpp
@@ -25,20 +25,13 @@ class FftPoissonSolver : public IPoissonSolver
 {
     IChargeDensityCalculator const& m_compute_rho;
 
-    ElectricField m_electric_field;
-
 public:
     /**
      * Construct the FftPoissonSolver operator.
      *
-     * @param spline_xy_builder A spline builder which calculates the coefficients of a spline representation.
-     * @param spline_xy_evaluator A spline evaluator which provides the value of a spline representation from its coefficients.
      * @param compute_rho The operator which calculates the charge density, the right hand side of the equation.
      */
-    FftPoissonSolver(
-            SplineXYBuilder const& spline_xy_builder,
-            SplineXYEvaluator const& spline_xy_evaluator,
-            IChargeDensityCalculator const& compute_rho);
+    FftPoissonSolver(IChargeDensityCalculator const& compute_rho);
 
     ~FftPoissonSolver() override = default;
 

diff --git a/tests/geometryXVx/fftpoissonsolver.cpp b/tests/geometryXVx/fftpoissonsolver.cpp
@@ -45,13 +45,10 @@ TEST(FftPoissonSolver, CosineSource)
     ddc::DiscreteDomain<IDimX> interpolation_domain_x(SplineInterpPointsX::get_domain());
     ddc::DiscreteDomain<IDimVx> interpolation_domain_vx(SplineInterpPointsVx::get_domain());
 
-    SplineBuilder<BSplinesX, IDimX, BoundCond::PERIODIC, BoundCond::PERIODIC> const builder_x(
-            interpolation_domain_x);
-
     SplineVxBuilder const builder_vx(interpolation_domain_vx);
 
-    IDomainX const gridx = builder_x.interpolation_domain();
-    IDomainVx const gridvx = builder_vx.interpolation_domain();
+    IDomainX const gridx = interpolation_domain_x;
+    IDomainVx const gridvx = interpolation_domain_vx;
     IDomainSp const gridsp = IDomainSp(my_iion, IVectSp(1));
 
     IDomainSpXVx const mesh(gridsp, gridx, gridvx);
@@ -60,13 +57,6 @@ TEST(FftPoissonSolver, CosineSource)
 
     // Creating operators
 
-    SplineEvaluator<BSplinesX> const
-            spline_x_evaluator(g_null_boundary<BSplinesX>, g_null_boundary<BSplinesX>);
-
-
-    SplineEvaluator<BSplinesVx> const
-            spline_vx_evaluator(g_null_boundary<BSplinesVx>, g_null_boundary<BSplinesVx>);
-
     FieldSp<int> charges(dom_sp);
     charges(my_ielec) = -1;
     charges(my_iion) = 1;
@@ -87,7 +77,7 @@ TEST(FftPoissonSolver, CosineSource)
     DFieldVx const quadrature_coeffs = neumann_spline_quadrature_coefficients(gridvx, builder_vx);
     Quadrature<IDimVx> const integrate_v(quadrature_coeffs);
     ChargeDensityCalculator rhs(integrate_v);
-    FftPoissonSolver poisson(builder_x, spline_x_evaluator, rhs);
+    FftPoissonSolver poisson(rhs);
 
     DFieldX electrostatic_potential(gridx);
     DFieldX electric_field(gridx);

diff --git a/tests/geometryXYVxVy/CMakeLists.txt b/tests/geometryXYVxVy/CMakeLists.txt
@@ -5,6 +5,7 @@ cmake_minimum_required(VERSION 3.15)
 add_executable(unit_tests_xy_vxvy
     ../main.cpp
     quadrature.cpp
+    fft_poisson.cpp
 )
 target_compile_features(unit_tests_xy_vxvy PUBLIC cxx_std_17)
 target_link_libraries(unit_tests_xy_vxvy
@@ -16,6 +17,7 @@ target_link_libraries(unit_tests_xy_vxvy
         gslx::geometry_xyvxvy
         sll::splines
         gslx::advection
+        gslx::poisson_xy
         gslx::quadrature
 )
 

diff --git a/tests/geometryXYVxVy/fft_poisson.cpp b/tests/geometryXYVxVy/fft_poisson.cpp
@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: MIT
+
+#include <sll/null_boundary_value.hpp>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "chargedensitycalculator.hpp"
+#include "fftpoissonsolver.hpp"
+#include "geometry.hpp"
+#include "species_info.hpp"
+
+TEST(FftPoissonSolver, CosineSource)
+{
+    CoordX const x_min(0.0);
+    CoordX const x_max(2.0 * M_PI);
+    IVectX const x_size(10);
+
+    CoordY const y_min(0.0);
+    CoordY const y_max(2.0 * M_PI);
+    IVectY const y_size(10);
+
+    CoordVx const vx_min(-0.5);
+    CoordVx const vx_max(0.5);
+    IVectVx const vx_size(11);
+
+    CoordVy const vy_min(-0.5);
+    CoordVy const vy_max(0.5);
+    IVectVy const vy_size(11);
+
+    IVectSp const nb_species(2);
+    IDomainSp const dom_sp(IndexSp(0), nb_species);
+    IndexSp const my_ielec = dom_sp.front();
+    IndexSp const my_iion = dom_sp.back();
+
+    // Creating mesh & supports
+    ddc::init_discrete_space<BSplinesVx>(vx_min, vx_max, vx_size);
+    ddc::init_discrete_space<BSplinesVy>(vy_min, vy_max, vy_size);
+
+    ddc::init_discrete_space<IDimVx>(SplineInterpPointsVx::get_sampling());
+    ddc::init_discrete_space<IDimVy>(SplineInterpPointsVy::get_sampling());
+
+    ddc::init_discrete_space(IDimX::init(x_min, x_max, x_size + 1));
+    ddc::init_discrete_space(IDimY::init(y_min, y_max, y_size + 1));
+
+    IDomainSp const gridsp = IDomainSp(my_iion, IVectSp(1));
+
+    IDomainX gridx = IDomainX(IndexX(0), x_size);
+    IDomainY gridy = IDomainY(IndexY(0), y_size);
+
+    IDomainXY gridxy(gridx, gridy);
+
+    IDomainVx gridvx = SplineInterpPointsVx::get_domain();
+    IDomainVy gridvy = SplineInterpPointsVy::get_domain();
+
+    IDomainVxVy gridvxvy(gridvx, gridvy);
+
+    ddc::init_fourier_space<RDimX>(gridx);
+    ddc::init_fourier_space<RDimY>(gridy);
+
+    SplineVxVyBuilder const builder_vx_vy(gridvxvy);
+    SplineVxVyEvaluator const evaluator_vx_vy(
+            g_null_boundary_2d<BSplinesVx, BSplinesVy>,
+            g_null_boundary_2d<BSplinesVx, BSplinesVy>,
+            g_null_boundary_2d<BSplinesVx, BSplinesVy>,
+            g_null_boundary_2d<BSplinesVx, BSplinesVy>);
+
+    IDomainSpXYVxVy const mesh(gridsp, gridxy, gridvxvy);
+
+    // Initialise infomation about species
+    FieldSp<int> charges(dom_sp);
+    charges(my_ielec) = -1;
+    charges(my_iion) = 1;
+    DFieldSp masses(dom_sp);
+    ddc::fill(masses, 1);
+    FieldSp<int> init_perturb_mode(dom_sp);
+    ddc::fill(init_perturb_mode, 0);
+    DFieldSp init_perturb_amplitude(dom_sp);
+    ddc::fill(init_perturb_amplitude, 0);
+
+    ddc::init_discrete_space<IDimSp>(
+            std::move(charges),
+            std::move(masses),
+            std::move(init_perturb_amplitude),
+            std::move(init_perturb_mode));
+
+    ChargeDensityCalculator rhs(builder_vx_vy, evaluator_vx_vy);
+    FftPoissonSolver poisson(rhs);
+
+    DFieldXY electrostatic_potential(gridxy);
+    DFieldXY electric_field_x(gridxy);
+    DFieldXY electric_field_y(gridxy);
+    DFieldSpXYVxVy allfdistribu(mesh);
+
+    // Initialization of the distribution function --> fill values
+    auto c_dom = ddc::remove_dims_of(mesh, gridxy);
+    ddc::for_each(gridxy, [&](IndexXY const ixy) {
+        IndexX ix = ddc::select<IDimX>(ixy);
+        IndexY iy = ddc::select<IDimY>(ixy);
+        double x = ddc::coordinate(ix);
+        double y = ddc::coordinate(iy);
+        double fdistribu_val = cos(x) + cos(y);
+        ddc::for_each(c_dom, [&](auto const ispvxvy) {
+            allfdistribu(ixy, ispvxvy) = fdistribu_val;
+        });
+    });
+
+    poisson(electrostatic_potential, electric_field_x, electric_field_y, allfdistribu);
+
+    double error_pot = 0.0;
+    double error_field_x = 0.0;
+    double error_field_y = 0.0;
+
+    ddc::for_each(gridxy, [&](IndexXY const ixy) {
+        IndexX ix = ddc::select<IDimX>(ixy);
+        IndexY iy = ddc::select<IDimY>(ixy);
+        double const exact_pot = cos(ddc::coordinate(ix)) + cos(ddc::coordinate(iy));
+        error_pot = fmax(fabs(electrostatic_potential(ixy) - exact_pot), error_pot);
+        double const exact_field_x = sin(ddc::coordinate(ix));
+        double const exact_field_y = sin(ddc::coordinate(iy));
+        error_field_x = fmax(fabs(electric_field_x(ixy) - exact_field_x), error_field_x);
+        error_field_y = fmax(fabs(electric_field_y(ixy) - exact_field_y), error_field_y);
+    });
+    EXPECT_LE(error_pot, 1e-12);
+    EXPECT_LE(error_field_x, 1e-10);
+    EXPECT_LE(error_field_y, 1e-10);
+}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -27,5 +27,3 @@ The Poisson equation can be solved with a variety of different methods. Here we
		- FemPeriodicPoissonSolver

		These classes return the electric potential $\phi$ and the electric field $\frac{d \phi}{dx}$.

		The FftPoissonSolver does not calculate the electric field using the Fourier modes. Rather it uses a spline interpolation to approximate this value. This interpolation is calculated by the operator ElectricField.