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):
