Curl Curl Solver: Variable beta coefficient

Add support for variable beta coefficient.
AMReX-Codes · Mar 16, 2024 · dc17bd1 · dc17bd1
1 parent 3fe7aad
commit dc17bd1
Show file tree

Hide file tree

Showing 7 changed files with 520 additions and 138 deletions.
diff --git a/Src/Base/AMReX_LUSolver.H b/Src/Base/AMReX_LUSolver.H
@@ -2,7 +2,7 @@
 #define AMREX_LU_SOLVER_H_
 #include <AMReX_Config.H>
 
-#include <AMReX_Arena.H>
+#include <AMReX_Algorithm.H>
 #include <AMReX_Array.H>
 #include <cmath>
 #include <limits>
@@ -18,6 +18,7 @@ public:
 
     LUSolver () = default;
 
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     LUSolver (Array2D<T, 0, N-1, 0, N-1, Order::C> const& a_mat);
 
     void define (Array2D<T, 0, N-1, 0, N-1, Order::C> const& a_mat);
@@ -74,6 +75,7 @@ public:
 
 private:
 
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
     void define_innard ();
 
     Array2D<T, 0, N-1, 0, N-1, Order::C> m_mat;
@@ -82,6 +84,7 @@ private:
 };
 
 template <int N, typename T>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
 LUSolver<N,T>::LUSolver (Array2D<T, 0, N-1, 0, N-1, Order::C> const& a_mat)
     : m_mat(a_mat)
 {
@@ -96,6 +99,7 @@ void LUSolver<N,T>::define (Array2D<T, 0, N-1, 0, N-1, Order::C> const& a_mat)
 }
 
 template <int N, typename T>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
 void LUSolver<N,T>::define_innard ()
 {
     static_assert(N > 1);
@@ -121,9 +125,9 @@ void LUSolver<N,T>::define_innard ()
         }
 
         if (imax != i) {
-            std::swap(m_piv(i), m_piv(imax));
+            amrex::Swap(m_piv(i), m_piv(imax));
             for (int j = 0; j < N; ++j) {
-                std::swap(m_mat(i,j), m_mat(imax,j));
+                amrex::Swap(m_mat(i,j), m_mat(imax,j));
             }
             ++m_npivs;
         }

diff --git a/Src/Base/AMReX_TagParallelFor.H b/Src/Base/AMReX_TagParallelFor.H
@@ -72,6 +72,26 @@ struct Array4BoxValTag {
     Box const& box () const noexcept { return dbox; }
 };
 
+template <class T>
+struct Array4BoxOrientationTag {
+    Array4<T> fab;
+    Box bx;
+    Orientation face;
+
+    [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    Box const& box() const noexcept { return bx; }
+};
+
+template <class T>
+struct Array4BoxOffsetTag {
+    Array4<T> fab;
+    Box bx;
+    Dim3 offset;
+
+    [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    Box const& box() const noexcept { return bx; }
+};
+
 template <class T>
 struct VectorTag {
     T* p;

diff --git a/Src/LinearSolvers/MLMG/AMReX_MLCurlCurl.H b/Src/LinearSolvers/MLMG/AMReX_MLCurlCurl.H
@@ -11,7 +11,7 @@ namespace amrex {
  * \brief curl (alpha curl E) + beta E = rhs
  *
  * Here E is an Array of 3 MultiFabs on staggered grid, alpha is a positive
- * scalar, and beta is a non-negative scalar.
+ * scalar, and beta is either a non-negative scalar or a MultiFab.
  *
  * It's the caller's responsibility to make sure rhs has consistent nodal
  * data. If needed, one could call prepareRHS for this.
@@ -45,6 +45,9 @@ public:
 
     void setScalars (RT a_alpha, RT a_beta) noexcept;
 
+    //! This is needed only if there is variable beta coefficient.
+    void setBeta (const Vector<Array<MultiFab const*,3>>& a_bcoefs);
+
     //! Synchronize RHS on nodal points. If the user can guarantee it, this
     //! function does not need to be called.
     void prepareRHS (Vector<MF*> const& rhs) const;
@@ -133,6 +136,7 @@ private:
     static constexpr int m_ncomp = 1;
     Vector<Vector<std::unique_ptr<Gpu::DeviceScalar
                                   <LUSolver<AMREX_SPACEDIM*2,RT>>>>> m_lusolver;
+    Vector<Vector<Array<std::unique_ptr<MultiFab>,3>>> m_bcoefs;
 };
 
 }