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_vec_partition():
N = 100
v = np.random.rand(N)
v_sub = partition_vec_by_num(v, 10)
v_reconstructed = reconstruct_vec_from_sub_vectors(v_sub)
assert np.sum(v - v_reconstructed) == 0
