import numpy as np

def cg(x0, A, b):
 r0 = np.dot(A, x0) - b
 p0 = -r0
 rk = r0
 pk = p0
 xk = x0
 t = 0 #記錄迭代次數
 while np.linalg.norm(rk) >= 1e-6:
  rr = np.dot(rk.T, rk)
  ak = rr / np.dot(np.dot(pk.T, A), pk)
  xk = xk + ak * pk
  rk = rk + ak * np.dot(A, pk)
  bk = np.dot(rk.T, rk) / rr
  pk = -rk + bk * pk
  t += 1
 return xk, t

#輸入列表，生成以列表為對角元素的對角矩陣
def Diagonal_matrix(D):
 n = len(D)
 diag = np.zeros((n,n))
 for i in range(n):
  diag[i][i] = D[i]
 return diag
#矩陣對角線元素
D_1 = [1, 1, 1, 1, 1, 6, 7, 8, 9, 10]
D_2 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
D_3 = [0.8, 0.9, 1, 1.1, 1.2, 6, 7, 8, 9, 10]
D_4 = [1 - 2*1e-7, 1 - 1e-7, 1, 1 + 1e-7, 1 + 2*1e-7, 6, 7, 8, 9, 10]
D_5 = [1, 1, 1, 2, 2, 2, 3, 3, 3, 10]
#初始值
x0 = np.zeros((10,1))
b = np.ones((10,1))  
#生成對角矩陣
diag1 = Diagonal_matrix(D_1)
diag2 = Diagonal_matrix(D_2)
diag3 = Diagonal_matrix(D_3)
diag4 = Diagonal_matrix(D_4)
diag5 = Diagonal_matrix(D_5)
#共軛梯度法迭代
x_1, n_1 = cg(x0, diag1, b)
x_2, n_2 = cg(x0, diag2, b)
x_3, n_3 = cg(x0, diag3, b)
x_4, n_4 = cg(x0, diag4, b)
x_5, n_5 = cg(x0, diag5, b)
n = [n_1, n_2, n_3, n_4, n_5]
#輸出
for i in range(5):
  print('矩陣',i + 1 ,'的迭代次數為: ', n[i])