From 3bbd7e47c4125ac4f0d3d5334cca93cf89c44c82 Mon Sep 17 00:00:00 2001
From: Andrew Reisner <areisner@lanl.gov>
Date: Fri, 30 Aug 2024 11:36:59 -0600
Subject: [PATCH] Add local_air CSR test, fix communication bug it found.

---
 raptor/ruge_stuben/par_interpolation.cpp |  7 +--
 raptor/ruge_stuben/tests/test_air.cpp    | 57 +++++++++++++-----------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/raptor/ruge_stuben/par_interpolation.cpp b/raptor/ruge_stuben/par_interpolation.cpp
index 8778b293..babadc23 100644
--- a/raptor/ruge_stuben/par_interpolation.cpp
+++ b/raptor/ruge_stuben/par_interpolation.cpp
@@ -2686,9 +2686,10 @@ CSRMatrix * communicate_neighborhood(const T & A, const ParCSRMatrix & S,
 			}
 		};
 
-		add_neighborhood(*A.on_proc,  A.on_proc_column_map,  *S.on_proc,  split.on_proc);
-		add_neighborhood(*A.off_proc, A.off_proc_column_map, *S.off_proc, split.off_proc);
-
+		if (split.on_proc[i] == Unselected) { // Only communicate f-point rows
+			add_neighborhood(*A.on_proc,  A.on_proc_column_map,  *S.on_proc,  split.on_proc);
+			add_neighborhood(*A.off_proc, A.off_proc_column_map, *S.off_proc, split.off_proc);
+		}
 		rowptr[i+1] = colind.size();
 	}
 
diff --git a/raptor/ruge_stuben/tests/test_air.cpp b/raptor/ruge_stuben/tests/test_air.cpp
index ad4331e6..498e1bf1 100644
--- a/raptor/ruge_stuben/tests/test_air.cpp
+++ b/raptor/ruge_stuben/tests/test_air.cpp
@@ -14,7 +14,6 @@ int main(int argc, char** argv)
     return ret;
 }
 
-#if 0
 TEST(TestOnePointInterp, TestsInRuge_Stuben) {
 	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 	constexpr std::size_t n{16};
@@ -87,11 +86,9 @@ TEST(TestOnePointInterp, TestsInRuge_Stuben) {
 	}
 }
 
-
 TEST(TestLocalAIR, TestsInRuge_Stuben) {
 	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-	constexpr std::size_t n{16};
-	// constexpr std::size_t n{5};
+	constexpr std::size_t n{17};
 	std::vector<int> grid;
 
 	grid.resize(1, n);
@@ -99,19 +96,6 @@ TEST(TestLocalAIR, TestsInRuge_Stuben) {
 	std::vector<double> stencil{{-1., 2, -1}};
 
 	auto A = par_stencil_grid(stencil.data(), grid.data(), 1);
-	// {
-	// 	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-	// 	if (rank == 1) {
-	// 		auto & offd = *A->off_proc;
-	// 		auto [rowptr, colind, values] = offd.vecs();
-	// 		for (int r = 0; r < A->local_num_rows; ++r) {
-	// 			for (int off = rowptr[r]; off < rowptr[r+1]; ++off) {
-	// 				auto & gcol = A->off_proc_column_map[colind[off]];
-	// 				std::cout << "here: " << r << " " << A->off_proc_column_map[colind[off]] << std::endl;
-	// 			}
-	// 		}
-	// 	}
-	// }
 	auto get_split = [](const Matrix &, const std::vector<int> & colmap) {
 		std::vector<int> split(colmap.size(), 0);
 
@@ -122,20 +106,42 @@ TEST(TestLocalAIR, TestsInRuge_Stuben) {
 	};
 	splitting_t split{get_split(*A->on_proc, A->on_proc_column_map),
 	                  get_split(*A->off_proc, A->off_proc_column_map)};
-
-	{
-		int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-		if (rank == 1)
-			split.on_proc[2] = Unselected;
-	}
-
 	auto S = A->copy();
 
 	auto R = local_air(*A, *S, split);
 
+	using expect_t = std::map<int, double>;
+	auto get_expected = [](int row) -> expect_t {
+		if (row == 0)
+			return {{0, 1}, {1, 0.5}};
+		else if (row == 16)
+			return {{16, 1}, {15, 0.5}};
+		else
+			return {
+				{row - 1, 0.5},
+				{row, 1},
+				{row + 1, 0.5}};
+	};
+	auto rowmap = [first_row = A->partition->first_local_col](int i) {
+		auto first_selected = (first_row % 2 == 0) ? first_row : first_row + 1;
+		return (first_selected / 2 + i) * 2;
+	};
+	for (int i = 0; i < R->on_proc->n_rows; ++i) {
+		auto row = rowmap(i);
+		auto expected = get_expected(row);
+		expect_t cols;
+		auto add_cols = [&](const Matrix & mat, const auto & colmap) {
+			for (std::size_t off = mat.idx1[i]; off < mat.idx1[i + 1]; ++off) {
+				cols[colmap[mat.idx2[off]]] = mat.vals[off];
+			}
+		};
+		add_cols(*R->on_proc, R->on_proc_column_map);
+		add_cols(*R->off_proc, R->off_proc_column_map);
+		ASSERT_EQ(expected, cols);
+	}
 }
-#endif
 
+#if 0
 TEST(TestLocalAIRBSR, TestsInRuge_Stuben) {
 	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 	constexpr std::size_t n{16};
@@ -191,3 +197,4 @@ TEST(TestLocalAIRBSR, TestsInRuge_Stuben) {
 
 	auto R = local_air(*Absr, *S, split);
 }
+#endif