MLLinOp::postSolve (AMReX-Codes#2981)

Add a virtual function MLLinOp::postSolve. This allows WarpX to set EB covered nodes to prescribed values in the solver's output for visualization purpose.
Thierry992 · Oct 11, 2022 · 0019b3a · 0019b3a
1 parent 2d87a4c
commit 0019b3a
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 0 deletions.
diff --git a/Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLaplacian.H b/Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLaplacian.H
@@ -119,6 +119,8 @@ public:
                           Array4<Real const> const& bfab) const override;
 #endif
 
+    virtual void postSolve (Vector<Any>& sol) const override;
+
 private:
     GpuArray<Real,AMREX_SPACEDIM> m_sigma{{AMREX_D_DECL(1_rt,1_rt,1_rt)}};
     Real m_s_phi_eb = std::numeric_limits<Real>::lowest();

diff --git a/Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLaplacian.cpp b/Src/LinearSolvers/MLMG/AMReX_MLEBNodeFDLaplacian.cpp
@@ -655,4 +655,39 @@ MLEBNodeFDLaplacian::fillRHS (MFIter const& /*mfi*/, Array4<int const> const& /*
 }
 #endif
 
+void
+MLEBNodeFDLaplacian::postSolve (Vector<Any>& sol) const
+{
+#ifdef AMREX_USE_EB
+    for (int amrlev = 0; amrlev < m_num_amr_levels; ++amrlev) {
+        const auto phieb = m_s_phi_eb;
+        auto factory = dynamic_cast<EBFArrayBoxFactory const*>(m_factory[amrlev][0].get());
+        auto const& levset_mf = factory->getLevelSet();
+        auto const& levset_ar = levset_mf.const_arrays();
+        MultiFab& mf = sol[amrlev].get<MultiFab>();
+        auto const& sol_ar = mf.arrays();
+        if (phieb == std::numeric_limits<Real>::lowest()) {
+            auto const& phieb_ar = m_phi_eb[amrlev].const_arrays();
+            amrex::ParallelFor(mf, IntVect(1),
+            [=] AMREX_GPU_DEVICE (int bi, int i, int j, int k) noexcept
+            {
+                if (levset_ar[bi](i,j,k) >= Real(0.0)) {
+                    sol_ar[bi](i,j,k) = phieb_ar[bi](i,j,k);
+                }
+            });
+        } else {
+            amrex::ParallelFor(mf, IntVect(1),
+            [=] AMREX_GPU_DEVICE (int bi, int i, int j, int k) noexcept
+            {
+                if (levset_ar[bi](i,j,k) >= Real(0.0)) {
+                    sol_ar[bi](i,j,k) = phieb;
+                }
+            });
+        }
+    }
+#else
+    amrex::ignore_unused(sol);
+#endif
+}
+
 }
diff --git a/Src/LinearSolvers/MLMG/AMReX_MLLinOp.H b/Src/LinearSolvers/MLMG/AMReX_MLLinOp.H
@@ -330,6 +330,8 @@ public:
 
     virtual void AnyAverageDownAndSync (Vector<Any>& sol) const = 0;
 
+    virtual void postSolve (Vector<Any>& sol) const;
+
     Real MFNormInf (MultiFab const& mf, iMultiFab const* fine_mask, bool local) const;
 
     bool isMFIterSafe (int amrlev, int mglev1, int mglev2) const;

diff --git a/Src/LinearSolvers/MLMG/AMReX_MLLinOp.cpp b/Src/LinearSolvers/MLMG/AMReX_MLLinOp.cpp
@@ -1159,6 +1159,9 @@ MLLinOp::AnyInterpAssignMG (int amrlev, int fmglev, Any& fine, Any& crse) const
     interpAssign(amrlev, fmglev, fine.get<MultiFab>(), crse.get<MultiFab>());
 }
 
+void
+MLLinOp::postSolve (Vector<Any>& /* sol */) const {}
+
 bool
 MLLinOp::isMFIterSafe (int amrlev, int mglev1, int mglev2) const
 {

diff --git a/Src/LinearSolvers/MLMG/AMReX_MLMG.cpp b/Src/LinearSolvers/MLMG/AMReX_MLMG.cpp
@@ -218,6 +218,8 @@ MLMG::solve (Vector<Any>& a_sol, const Vector<Any>& a_rhs,
         timer[iter_time] = amrex::second() - iter_start_time;
     }
 
+    linop.postSolve(sol);
+
     IntVect ng_back = final_fill_bc ? IntVect(1) : IntVect(0);
     if (linop.hasHiddenDimension()) {
         ng_back[linop.hiddenDirection()] = 0;