wip

WeiqunZhang · Dec 5, 2024 · 68a6062 · 68a6062
1 parent beb320a
commit 68a6062
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 19 deletions.
diff --git a/Src/LinearSolvers/AMReX_AlgPartition.cpp b/Src/LinearSolvers/AMReX_AlgPartition.cpp
@@ -26,13 +26,13 @@ void AlgPartition::define (Long global_size)
 void AlgPartition::define (Vector<Long> const& rows)
 {
     m_ref->m_row = rows;
-    AMREX_ASSERT(m_row.size() == ParallelDescriptor::NProcs());
+    AMREX_ASSERT(m_ref->m_row.size() == ParallelDescriptor::NProcs());
 }
 
 void AlgPartition::define (Vector<Long>&& rows)
 {
     m_ref->m_row = std::move(rows);
-    AMREX_ASSERT(m_row.size() == ParallelDescriptor::NProcs());
+    AMREX_ASSERT(m_ref->m_row.size() == ParallelDescriptor::NProcs());
 }
 
 AlgPartition::Ref::Ref (Long global_size)

diff --git a/Tests/Algebra/GMRES/main.cpp b/Tests/Algebra/GMRES/main.cpp
@@ -1,6 +1,7 @@
 #include <AMReX_GMRES_MV.H>
 
 #include <AMReX.H>
+#include <AMReX_Random.H>
 
 using namespace amrex;
 
@@ -14,20 +15,7 @@ int main (int argc, char* argv[])
         AlgVector<Real> bvec(xvec.partition());
 
         // Initialzie bvec
-        {
-            auto nrows = bvec.numLocalRows();
-            auto* p = bvec.data();
-
-            amrex::ParallelForRNG(nrows,
-            [=] AMREX_GPU_DEVICE (Long i, RandomEngine const& engine) noexcept
-            {
-                p[i] = amrex::RandomNormal(0.,1.,engine);
-            });
-
-            auto bsum = bvec.sum();
-            auto offset = bsum / Real(bvec.numGlobalRows());
-            ForEach(bvec, [=] AMREX_GPU_DEVICE (auto& b) { b -= bsum; });
-        }
+        amrex::FillRandomNormal(bvec.data(), bvec.numLocalRows(), Real(0), Real(1));
 
         // Initial guess
         xvec.setVal(0);
@@ -36,14 +24,35 @@ int main (int argc, char* argv[])
 
         auto num_non_zeros = [=] AMREX_GPU_DEVICE (Long row) -> int
         {
-            int nnz = 2*AMREX_SPACEDIM+1; // cross stencil
-            IntVect iv = box_indexer.intVect(row);
-            return nnz;
+            return 2*AMREX_SPACEDIM+1; // cross stencil w/ periodic boundaries
         };
 
+        // a * phi - del dot grad phi. For simplicity, let a=1 and dx=1.
         auto set_stencil = [=] AMREX_GPU_DEVICE (Long row, Long* col, Real* value)
         {
+            IntVect cell = box_indexer.intVect(row);
+            for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+                IntVect cell2 = cell;
+                if (cell[idim] == domain.smallEnd(idim)) {
+                    cell2[idim] = domain.bigEnd(idim);
+                } else {
+                    cell2[idim] = cell[idim] - 1;
+                }
+                Long row2 = domain.index(cell2);
+                *col++ = row2;
+                *value++ = Real(-1.0);
 
+                if (cell[idim] == domain.bigEnd(idim)) {
+                    cell2[idim] = domain.smallEnd(idim);
+                } else {
+                    cell2[idim] = cell[idim] + 1;
+                }
+                row2 = domain.index(cell2);
+                *col++ = row2;
+                *value++ = Real(-1.0);
+            }
+            *col++ = row;
+            *value++ = Real(2*AMREX_SPACEDIM+1);
         };
 
         SpMatrix<Real> mat(xvec.partition(), num_non_zeros);