From 8199a6fe6fb33efffea6a988119901166200b4be Mon Sep 17 00:00:00 2001
From: Ann Almgren <asalmgren@lbl.gov>
Date: Mon, 18 Nov 2024 18:42:32 -0800
Subject: [PATCH] fix hybrid FFT when we used stretched grids

---
 Source/Utils/ERF_TerrainMetrics.cpp | 16 +++++++++---
 Source/Utils/ERF_solve_with_fft.cpp | 38 +++++++++++++++++------------
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/Source/Utils/ERF_TerrainMetrics.cpp b/Source/Utils/ERF_TerrainMetrics.cpp
index 03ffe8245..c024cdf1c 100644
--- a/Source/Utils/ERF_TerrainMetrics.cpp
+++ b/Source/Utils/ERF_TerrainMetrics.cpp
@@ -63,9 +63,6 @@ init_zlevels (Vector<Vector<Real>>& zlevels_stag,
                 stretched_dz_h[lev][k] = (zlevels_stag[lev][k+1] - zlevels_stag[lev][k]);
             }
         }
-
-        stretched_dz_d[lev].resize(domain.length(2));
-        Gpu::copy(Gpu::hostToDevice, stretched_dz_h[lev].begin(), stretched_dz_h[lev].end(), stretched_dz_d[lev].begin());
     }
 
     // Try reading in terrain_z_levels, which allows arbitrarily spaced grid
@@ -97,6 +94,19 @@ init_zlevels (Vector<Vector<Real>>& zlevels_stag,
             int rr = ref_ratio[lev-1][2];
             expand_and_interpolate_1d(zlevels_stag[lev], zlevels_stag[lev-1], rr, false);
         }
+
+        for (int lev = 0; lev <= max_level; lev++) {
+            int nz = zlevels_stag[lev].size();
+            for (int k = 0; k < nz-1; k++)
+            {
+                stretched_dz_h[lev][k] = (zlevels_stag[lev][k+1] - zlevels_stag[lev][k]);
+            }
+        }
+    }
+
+    for (int lev = 0; lev <= max_level; lev++) {
+        stretched_dz_d[lev].resize(stretched_dz_h[lev].size());
+        Gpu::copy(Gpu::hostToDevice, stretched_dz_h[lev].begin(), stretched_dz_h[lev].end(), stretched_dz_d[lev].begin());
     }
 }
 
diff --git a/Source/Utils/ERF_solve_with_fft.cpp b/Source/Utils/ERF_solve_with_fft.cpp
index dfb0559b3..d18d3ee67 100644
--- a/Source/Utils/ERF_solve_with_fft.cpp
+++ b/Source/Utils/ERF_solve_with_fft.cpp
@@ -53,13 +53,10 @@ void ERF::solve_with_fft (int lev, MultiFab& rhs, MultiFab& phi, Array<MultiFab,
     AMREX_ALWAYS_ASSERT(use_fft);
 
     bool l_use_terrain = SolverChoice::terrain_type != TerrainType::None;
-    bool use_gmres     = (l_use_terrain && !SolverChoice::terrain_is_flat);
 
     auto const dom_lo = lbound(geom[lev].Domain());
     auto const dom_hi = ubound(geom[lev].Domain());
 
-    MultiFab r_hse(base_state[lev], make_alias, BaseState::r0_comp, 1);
-
     auto bclo = get_projection_bc(Orientation::low);
     auto bchi = get_projection_bc(Orientation::high);
 
@@ -107,10 +104,7 @@ void ERF::solve_with_fft (int lev, MultiFab& rhs, MultiFab& phi, Array<MultiFab,
         if (!m_2D_poisson) {
             m_2D_poisson = std::make_unique<FFT::PoissonHybrid<MultiFab>>(Geom(0));
         }
-        // m_2D_poisson->solve(phi, rhs[lev], stretched_dz_d[lev]);
-
-        Gpu::DeviceVector<Real> stretched_dz(dom_hi.z+1, geom[lev].CellSize(2));
-        m_2D_poisson->solve(phi, rhs, stretched_dz);
+        m_2D_poisson->solve(phi, rhs, stretched_dz_d[lev]);
 
     } else {
         amrex::Abort("FFT isn't appropriate for spatially varying terrain");
@@ -227,16 +221,28 @@ void ERF::solve_with_fft (int lev, MultiFab& rhs, MultiFab& phi, Array<MultiFab,
         });
 
         Box const& zbx = mfi.nodaltilebox(2);
-        const Real dz_inv = dxInv[2];
         Array4<Real> const& fz_arr  = fluxes[2].array(mfi);
-        ParallelFor(zbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
-        {
-            if (k == dom_lo.z || k == dom_hi.z+1) {
-                fz_arr(i,j,k) = 0.0;
-            } else {
-                fz_arr(i,j,k) = -(p_arr(i,j,k) - p_arr(i,j,k-1)) * dz_inv;
-            }
-        });
+        if (l_use_terrain && SolverChoice::terrain_is_flat) {
+            ParallelFor(zbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+                if (k == dom_lo.z || k == dom_hi.z+1) {
+                    fz_arr(i,j,k) = 0.0;
+                } else {
+                    Real dz = 0.5 * (stretched_dz_d[lev][k] + stretched_dz_d[lev][k-1]);
+                    fz_arr(i,j,k) = -(p_arr(i,j,k) - p_arr(i,j,k-1)) / dz;
+                }
+            });
+        } else {
+            const Real dz_inv = dxInv[2];
+            ParallelFor(zbx, [=] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
+            {
+                if (k == dom_lo.z || k == dom_hi.z+1) {
+                    fz_arr(i,j,k) = 0.0;
+                } else {
+                    fz_arr(i,j,k) = -(p_arr(i,j,k) - p_arr(i,j,k-1)) * dz_inv;
+                }
+            });
+        }
     } // mfi
 }
 #endif