drop: adding rowmap

raptor/multilevel/par_level.hpp

@@ -34,6 +34,7 @@ namespace raptor
 
             ParCSRMatrix* A;
             ParCSRMatrix* P;
+		    ParCSRMatrix* R;
             ParVector x;
             ParVector b;
             ParVector tmp;

raptor/ruge_stuben/par_air_solver.hpp

@@ -34,7 +34,7 @@ struct ParAIRSolver : ParMultilevel {
 		ParCSRMatrix *A = levels[level_ctr]->A;
 		ParBSRMatrix *A_bsr = dynamic_cast<ParBSRMatrix*>(A);
 		if (A_bsr) extend_hier(*A_bsr);
-		// else extend_hier(*A);
+		else extend_hier(*A);
 	}
 
 private:
@@ -50,7 +50,46 @@ struct ParAIRSolver : ParMultilevel {
 		auto P = one_point_interpolation(A, *S, split);
 		auto R = local_air(A, *S, split, fpoint_distance::two);
 
+		levels.back()->P = P;
+		levels.back()->R = R;
 		levels.emplace_back(new ParLevel());
+		auto & level = *levels.back();
+		auto AP = A.mult(P, tap_level);
+		auto coarse_A = R->mult(AP);
+		// auto coarse_A = AP->mult_T(P);
+
+		level.A = coarse_A;
+
+		finalize_coarse_op(tap_amg >= 0 && tap_amg <= level_ctr + 1);
+		init_coarse_vectors();
+
+		level.P = nullptr;
+
+		delete AP;
+		delete S;
+	}
+
+	void finalize_coarse_op(bool tap_init) {
+		auto & coarse_op = *levels.back()->A;
+		auto & fine_op = *(**(levels.end() - 2)).A;
+
+		coarse_op.sort();
+		coarse_op.on_proc->move_diag();
+		coarse_op.comm = new ParComm(coarse_op.partition, coarse_op.off_proc_column_map,
+		                             coarse_op.on_proc_column_map,
+		                             fine_op.comm->key, fine_op.comm->mpi_comm);
+
+		if (tap_init) {
+			coarse_op.init_tap_communicators(RAPtor_MPI_COMM_WORLD);
+		}
+	}
+
+	void init_coarse_vectors() {
+		auto & clevel = *levels.back();
+		[global_rows = clevel.A->global_num_rows,
+		 local_rows = clevel.A->local_num_rows](auto & ... vecs) {
+			(vecs.resize(global_rows, local_rows),...);
+		}(clevel.x, clevel.b, clevel.tmp);
 	}
 
 	coarsen_t coarsen_type;

raptor/ruge_stuben/par_interpolation.cpp

@@ -2143,6 +2143,19 @@ T * create_R(const T & A,
 		return global;
 	}(local_rows);
 
+
+	auto local_rowmap = [local_rows]() {
+		auto row_start = [](int local) {
+			MPI_Exscan(MPI_IN_PLACE, &local, 1, MPI_INT, MPI_SUM, RAPtor_MPI_COMM_WORLD);
+			int rank; MPI_Comm_rank(RAPtor_MPI_COMM_WORLD, &rank);
+			return (rank == 0) ? 0 : local;
+		}(local_rows);
+
+		std::vector<int> rowmap(local_rows);
+		std::iota(std::begin(rowmap), std::end(rowmap), row_start);
+		return rowmap;
+	};
+
 	auto off_proc_num_cols = rowptr_and_colmap.offd_colmap.size();
 	auto move_data = [&](offd_map_rowptr && rac, ParMatrix & mat) {
 		mat.on_proc->idx1 = std::move(rac.diag_rowptr);
@@ -2163,11 +2176,13 @@ T * create_R(const T & A,
 		                            local_rows, S.on_proc_num_cols, off_proc_num_cols,
 		                            A.on_proc->b_rows, A.on_proc->b_cols);
 		move_data(std::move(rowptr_and_colmap), *R);
+		R->local_row_map = local_rowmap();
 		return R;
 	} else {
 		auto * R = new ParCSRMatrix(S.partition, global_rows, S.global_num_cols,
 		                            local_rows, S.on_proc_num_cols, off_proc_num_cols);
 		move_data(std::move(rowptr_and_colmap), *R);
+		R->local_row_map = local_rowmap();
 		return R;
 	}
 }

raptor/ruge_stuben/tests/test_air.cpp

@@ -1,6 +1,7 @@
 #include "gtest/gtest.h"
 
 #include "raptor/raptor.hpp"
+#include "raptor/ruge_stuben/par_air_solver.hpp"
 #include "raptor/tests/compare.hpp"
 
 using namespace raptor;
@@ -14,6 +15,7 @@ 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};
@@ -141,7 +143,6 @@ TEST(TestLocalAIR, TestsInRuge_Stuben) {
 	}
 }
 
-#if 0
 TEST(TestLocalAIRBSR, TestsInRuge_Stuben) {
 	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
 	constexpr std::size_t n{16};
@@ -198,3 +199,77 @@ TEST(TestLocalAIRBSR, TestsInRuge_Stuben) {
 	auto R = local_air(*Absr, *S, split);
 }
 #endif
+
+void dump(const ParCSRMatrix & A) {
+	std::cout << A.global_num_cols << "  " << A.global_num_rows << std::endl;
+	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+	std::ofstream ofile("coarse-" + std::to_string(rank));
+	auto write_row = [&](int i, const Matrix & mat, const std::vector<int> & colmap) {
+		const auto & rowmap = A.local_row_map;
+		for (int j = mat.idx1[i]; j < mat.idx1[i+1]; ++j) {
+			ofile << rowmap[i] << " " << colmap[mat.idx2[j]] << " " << std::scientific << mat.vals[j] << '\n';
+
+		}
+	};
+	for (int i = 0; i < A.local_num_rows; ++i) {
+		write_row(i, *A.on_proc, A.on_proc_column_map);
+		write_row(i, *A.off_proc, A.off_proc_column_map);
+	}
+}
+
+ParCSRMatrix * gen(std::size_t n) {
+	auto A = new ParCSRMatrix(n, n);
+
+	int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+	int size; MPI_Comm_size(MPI_COMM_WORLD, &size);
+
+	auto num_local = A->partition->local_num_rows;
+	auto row_start = A->partition->first_local_row;
+	auto init = [&](Matrix & mat, int ncols) {
+		mat.n_rows = num_local;
+		mat.n_cols = ncols;
+		mat.nnz = 0;
+		mat.idx1.resize(num_local+1);
+		mat.idx2.reserve(num_local*2);
+		mat.vals.reserve(.3*num_local*2);
+	};
+	init(*A->on_proc, num_local);
+	init(*A->off_proc, n);
+
+	A->on_proc->idx1[0] = 0;
+	A->off_proc->idx1[0] = 0;
+	for (int i = 0; i < num_local; ++i) {
+		A->add_value(i, i + row_start, 1.);
+		if (!(rank == 0 && i == 0))
+			A->add_value(i, i + row_start - 1, -1.);
+
+		auto set_rowptr = [i](Matrix & mat) {
+			mat.idx1[i+1] = mat.idx2.size();
+		};
+
+		set_rowptr(*A->on_proc);
+		set_rowptr(*A->off_proc);
+	}
+
+	A->on_proc->nnz = A->on_proc->idx2.size();
+	A->off_proc->nnz = A->off_proc->idx2.size();
+
+	A->finalize();
+
+	return A;
+}
+TEST(UpwindAdvection, TestsInRuge_Stuben) {
+	// constexpr std::size_t n{100};
+	// std::vector<int> grid;
+	// grid.resize(1, n);
+	// std::vector<double> stencil{{-1, 1, 0}};
+	// auto A = par_stencil_grid(stencil.data(), grid.data(), 1);
+	auto A = gen(100);
+
+	// dump(*A);
+	ParAIRSolver ml;
+	// ParRugeStubenSolver ml;
+	ml.max_levels = 2;
+	ml.setup(A);
+	dump(*ml.levels[1]->A);
+}

raptor/ruge_stuben/tests/test_air.py

@@ -0,0 +1,28 @@
+import numpy as np
+
+from numpy.testing import TestCase, assert_array_almost_equal
+from scipy.sparse import bsr_matrix, diags, coo_matrix
+from pyamg.classical.air import air_solver
+from pyamg.gallery import poisson
+
+n = 100
+A = diags([np.ones((n,)), -1*np.ones((n-1,))], [0, -1]).tocsr()
+f_relax = ('fc_jacobi', {'iterations': 1, 'f_iterations': 1,
+                         'c_iterations': 0})
+ml = air_solver(A, postsmoother=f_relax, max_levels=2)
+# print(ml.levels[1].A)
+
+ra_coords = np.loadtxt('/home/andrew/src/raptor/build/raptor/ruge_stuben/tests/coarse')
+I = np.array(ra_coords[:,0], dtype=int)
+J = np.array(ra_coords[:,1], dtype=int)
+# Jinds = np.unique(J)
+# jmap = {}
+# for i in range(Jinds.shape[0]):
+#     jmap[Jinds[i]] = i
+
+# for i in range(J.shape[0]):
+#     J[i] = jmap[J[i]]
+
+ra = coo_matrix((ra_coords[:,2], (I, J)))
+
+# print(np.max(ra - ml.levels[1].A))