def solve_block_jacobi(A, b, n_partitions=2, x0=None, maxiter=100): assert A.shape[0] == A.shape[1], "Square matrix required" N = len(b) x = parse_initial_cond(x0, N) A_sub = partition_mat_by_num(A, n_partitions) x_sub_old = partition_vec_by_num(x, n_partitions) b_sub = partition_vec_by_num(b, n_partitions) rhs_sub = dict() x_sub = dict() residual = 0.0 for _ in xrange(maxiter): for i in xrange(n_partitions): # Go through row blocks sum_mat_vec_results = np.zeros(len(b_sub[i])) for j in xrange(n_partitions): if i == j: continue sum_mat_vec_results += A_sub[i, j] * x_sub_old[j] rhs_sub[i] = b_sub[i] - sum_mat_vec_results x_sub[i] = spsolve(A_sub[i, i], rhs_sub[i]) for i in xrange(n_partitions): x_sub_old[i][:] = x_sub[i][:] x = reconstruct_vec_from_sub_vectors(x_sub) residual = compute_normalized_linf_residual(A, x, b) return x, residual
def test_mat_partition(): N = 100 A = poisson_1d(N) A_sub = partition_mat_by_num(A, 10) A_block = [[A_sub[i, j] for j in xrange(10)] for i in xrange(10)] A_reconstructed_from_block = sparse.bmat(A_block).tocsr() assert (A - A_reconstructed_from_block).nnz == 0
if __name__ == "__main__": comm = MPI.COMM_WORLD my_id = comm.Get_rank() n_proc = comm.Get_size() if my_id == 0: n_partitions = 20 N = 1311 A, b = poisson_system_1d(N) x = np.zeros(N) A_sub = partition_mat_by_num(A, n_partitions) b_sub = partition_vec_by_num(b, n_partitions) x_sub = partition_vec_by_num(x, n_partitions) indices_ptrs = np.linspace(0, n_partitions, num=n_proc + 1, dtype=np.int) G = nx.Graph() for i in xrange(n_partitions): for j in xrange(n_partitions): if A_sub[i, j].nnz != 0 and i != j: G.add_edge(i, j) for n in G.nodes(): G.node[n]['size'] = len(b_sub[n]) for owner in xrange(n_proc):