wip

WeiqunZhang · Oct 15, 2024 · fd95855 · fd95855
1 parent ff7e950
commit fd95855
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 11 deletions.
diff --git a/Src/FFT/AMReX_FFT.H b/Src/FFT/AMReX_FFT.H
@@ -131,10 +131,10 @@ template <typename F>
 void R2C::post_forward_doit (F const& post_forward)
 {
  auto& spectral_fab = m_cz[ParallelDescriptor::MyProc()];
- auto const& a = spectral_fab.array();
- ParallelFor(spectral_fab.box(), [=] AMREX_GPU_DEVICE (int i, int j, int k)
+ auto const& a = spectral_fab.array(); // m_cz's ordering is z,x,y
+ ParallelFor(spectral_fab.box(), [=] AMREX_GPU_DEVICE (int iz, int jx, int ky)
  {
- post_forward(j,k,i,a(i,j,k)); // m_cz's ordering is z,x,y
+ post_forward(jx,ky,iz,a(iz,jx,ky));
  });
 }
 

diff --git a/Src/FFT/AMReX_FFT.cpp b/Src/FFT/AMReX_FFT.cpp
@@ -72,7 +72,7 @@ R2C::R2C (Box const& domain)
  DistributionMapping cdmz = dmx; // xxxxx todo
  m_cz.define(cbaz, cdmz, 1, 0);
 
- std::tie(m_fftw_fwd_z, m_fftw_bwd_z) = make_c2c_plans(m_cy, m_spectral_domain_z);
+ std::tie(m_fftw_fwd_z, m_fftw_bwd_z) = make_c2c_plans(m_cz, m_spectral_domain_z);
 
  // comm meta-data between y and z phases
  m_cmd_y2z = std::make_unique<MultiBlockCommMetaData>
@@ -109,14 +109,13 @@ void R2C::backward_doit (MultiFab& outmf, Scaling scaling)
 {
  // xxxxx todo: scaling
 
- fftw_execute(m_fftw_bwd_z); //
-
+ fftw_execute(m_fftw_bwd_z);
  ParallelCopy(m_cy, m_cz, *m_cmd_z2y, 0, 0, 1, m_dtos_z2y);
- fftw_execute(m_fftw_bwd_y);
 
+ fftw_execute(m_fftw_bwd_y);
  ParallelCopy(m_cx, m_cy, *m_cmd_y2x, 0, 0, 1, m_dtos_y2x);
- fftw_execute(m_fftw_bwd_x);
 
+ fftw_execute(m_fftw_bwd_x);
  outmf.ParallelCopy(m_rx, 0, 0, 1);
 }
 

diff --git a/Tests/FFT/Poisson/main.cpp b/Tests/FFT/Poisson/main.cpp
@@ -52,14 +52,15 @@ int main (int argc, char* argv[])
  Real x = (i+0.5_rt) * dx[0];
  Real y = (AMREX_SPACEDIM>=2) ? (j+0.5_rt) * dx[1] : 0._rt;
  Real z = (AMREX_SPACEDIM==3) ? (k+0.5_rt) * dx[2] : 0._rt;
- rhsma[b](i,j,k) = std::exp(-10._rt*((x-0.5_rt)*(x-0.5_rt) +
- (y-0.5_rt)*(y-0.5_rt) +
+ rhsma[b](i,j,k) = std::exp(-10._rt*((x-0.5_rt)*(x-0.5_rt)*1.0_rt +
+ (y-0.5_rt)*(y-0.5_rt)*1.0_rt +
  (z-0.5_rt)*(z-0.5_rt)));
  });
 
  Real facx = 2._rt*Math::pi<Real>()/std::abs(prob_hi_x-prob_lo_x);
  Real facy = 2._rt*Math::pi<Real>()/std::abs(prob_hi_y-prob_lo_y);
  Real facz = 2._rt*Math::pi<Real>()/std::abs(prob_hi_z-prob_lo_z);
+ Real scale = 1._rt/(Real(n_cell_z)*Real(n_cell_y)*Real(n_cell_z));
 
  auto post_forward = [=] AMREX_GPU_DEVICE (int i, int j, int k,
  GpuComplex<Real>& spectral_data)
@@ -80,12 +81,31 @@ int main (int argc, char* argv[])
  // solution is zero
  spectral_data = 0._rt;
  }
+ spectral_data *= scale;
  };
 
  FFT::R2C fft(geom.Domain());
  fft.forwardThenBackward(rhs, soln, post_forward);
 
- VisMF::Write(soln, "soln");
+ {
+ MultiFab phi(soln.boxArray(), soln.DistributionMap(), 1, 1);
+ MultiFab res(soln.boxArray(), soln.DistributionMap(), 1, 0);
+ MultiFab::Copy(phi, soln, 0, 0, 1, 0);
+ phi.FillBoundary(geom.periodicity());
+ auto const& res_ma = res.arrays();
+ auto const& phi_ma = phi.const_arrays();
+ auto const& rhs_ma = rhs.const_arrays();
+ ParallelFor(res, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
+ {
+ auto const& phia = phi_ma[b];
+ auto lap = (phia(i-1,j,k)-2.*phia(i,j,k)+phia(i+1,j,k)) / (dx[0]*dx[0])
+ + (phia(i,j-1,k)-2.*phia(i,j,k)+phia(i,j+1,k)) / (dx[1]*dx[1])
+ + (phia(i,j,k-1)-2.*phia(i,j,k)+phia(i,j,k+1)) / (dx[2]*dx[2]);
+ res_ma[b](i,j,k) = rhs_ma[b](i,j,k) - lap;
+ });
+ amrex::Print() << " rhs.min & max: " << rhs.min(0) << " " << rhs.max(0) << "\n"
+ << " res.min & max: " << res.min(0) << " " << res.max(0) << "\n";
+ }
  }
  amrex::Finalize();
 }