add option for mac interpolation only in z

AMReX-Codes · Jan 5, 2024 · 7c103f1 · 7c103f1
1 parent fe52b98
commit 7c103f1
Showing 1 changed file with 78 additions and 2 deletions.
diff --git a/Src/Particle/AMReX_TracerParticle_mod_K.H b/Src/Particle/AMReX_TracerParticle_mod_K.H
@@ -110,7 +110,7 @@ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
 void mac_interpolate (const P& p,
                       amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
                       amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi,
-                      amrex::GpuArray<amrex::Array4<amrex::Real const>,AMREX_SPACEDIM> const& p_uccarr,
+                      amrex::GpuArray<amrex::Array4<amrex::Real const>,AMREX_SPACEDIM> const& umac_arr,
                       amrex::ParticleReal * val)
 {
     AMREX_ASSERT(val != nullptr);
@@ -145,10 +145,86 @@ void mac_interpolate (const P& p,
           for (int jj = 0; jj <= 1; ++jj) {
 #endif
               for (int ii = 0; ii <= 1; ++ii) {
-                  val[d] += static_cast<ParticleReal>((p_uccarr[d])(IntVect(AMREX_D_DECL(i+ii, j+jj, k+kk)), 0)*AMREX_D_TERM(sx[ii],*sy[jj],*sz[kk]));
+                  val[d] += static_cast<ParticleReal>((umac_arr[d])(IntVect(AMREX_D_DECL(i+ii, j+jj, k+kk)), 0)*AMREX_D_TERM(sx[ii],*sy[jj],*sz[kk]));
     AMREX_D_TERM(},},});
     }
 }
 
+template <typename P>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+void mac_interpolate_mapped_z (const P& p,
+                               amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& plo,
+                               amrex::GpuArray<amrex::Real,AMREX_SPACEDIM> const& dxi,
+                               amrex::GpuArray<amrex::Array4<amrex::Real const>,AMREX_SPACEDIM> const& umac_arr,
+                               const amrex::Array4<amrex::Real const>&  zloc_arr,
+                               bool use_terrain, amrex::ParticleReal * val)
+{
+    AMREX_ASSERT(val != nullptr);
+
+    for (int d=0; d < AMREX_SPACEDIM; ++d) {
+
+        if (use_terrain && d == 2) {continue;}
+
+        AMREX_D_TERM(amrex::Real lx = (Real(p.pos(0))-plo[0])*dxi[0] - static_cast<Real>(d != 0)*Real(0.5);,
+                     amrex::Real ly = (Real(p.pos(1))-plo[1])*dxi[1] - static_cast<Real>(d != 1)*Real(0.5);,
+                     amrex::Real lz = (Real(p.pos(2))-plo[2])*dxi[2] - static_cast<Real>(d != 2)*Real(0.5));
+
+        AMREX_D_TERM(int const i = static_cast<int>(amrex::Math::floor(lx));,
+                     int const j = static_cast<int>(amrex::Math::floor(ly));,
+                     int const k = static_cast<int>(amrex::Math::floor(lz)));
+
+        AMREX_D_TERM(amrex::Real const xint = lx - static_cast<Real>(i);,
+                     amrex::Real const yint = ly - static_cast<Real>(j);,
+                     amrex::Real const zint = lz - static_cast<Real>(k));
+
+#if AMREX_SPACEDIM > 2
+        amrex::Real sz[] = {Real(1.0) - zint, zint};
+#endif
+#if AMREX_SPACEDIM > 1
+      amrex::Real sy[] = {Real(1.0) - yint, yint};
+#endif
+      amrex::Real sx[] = {Real(1.0) - xint, xint};
+
+      val[d] = ParticleReal(0.0);
+#if AMREX_SPACEDIM > 2
+      for (int kk = 0; kk <=1; ++kk) {
+#endif
+#if AMREX_SPACEDIM > 1
+          for (int jj = 0; jj <= 1; ++jj) {
+#endif
+              for (int ii = 0; ii <= 1; ++ii) {
+                  val[d] += static_cast<ParticleReal>((umac_arr[d])(IntVect(AMREX_D_DECL(i+ii, j+jj, k+kk)), 0)*AMREX_D_TERM(sx[ii],*sy[jj],*sz[kk]));
+    AMREX_D_TERM(},},});
+    }
+
+#if AMREX_SPACEDIM == 3
+    if (use_terrain) {
+        amrex::Real lx = (Real(p.pos(0))-plo[0])*dxi[0] - Real(0.5);
+        amrex::Real ly = (Real(p.pos(1))-plo[1])*dxi[1] - Real(0.5);
+
+        int const i = static_cast<int>(amrex::Math::floor(lx));
+        int const j = static_cast<int>(amrex::Math::floor(ly));
+        int const k = p.idata(0);
+
+        amrex::Real const xint = lx - static_cast<Real>(i);
+        amrex::Real const yint = ly - static_cast<Real>(j);
+        amrex::Real const zint = (Real(p.pos(2)) - zloc_arr(IntVect(i,j,k))) / (zloc_arr(i,j,k+1) - zloc_arr(i,j,k));
+
+        amrex::Real sz[] = {Real(1.0) - zint, zint};
+        amrex::Real sy[] = {Real(1.0) - yint, yint};
+        amrex::Real sx[] = {Real(1.0) - xint, xint};
+
+      val[2] = ParticleReal(0.0);
+      for (int kk = 0; kk <=1; ++kk) {
+          for (int jj = 0; jj <= 1; ++jj) {
+              for (int ii = 0; ii <= 1; ++ii) {
+                  val[2] += static_cast<ParticleReal>((umac_arr[2])(IntVect(i+ii, j+jj, k+kk)), 0)*sx[ii],*sy[jj],*sz[kk];
+              }
+          }
+      }
+    }
+#endif
+}
+
 }  // namespace amrex
 #endif  // include guard