Homec4science
Diffusion SP4E_Homework_Ashtari_Sieber (master) Browse

R9482/Homework1/conjugate_gradient.py
b3377f2bf369master

Last Change
Raw File

conjugate_gradient.py
View Options

import numpy as np
import argparse
import math
# Arguments parser
my_parser = argparse.ArgumentParser()
my_parser.add_argument('-A', action='store', type=float, nargs='+', required=True, help='The matrix A in rowmajor format: numbers separated by space')
my_parser.add_argument('-b', action='store', type=float, nargs='+', required=True, help='The vector b in rowmajor format: numbers separated by space')
my_parser.add_argument('-x0', action='store', type=float, nargs='+', required=True, help='The initial guess x0 in rowmajor format: numbers separated by space')
my_parser.add_argument('-Nmax', action='store', type=int, default=20, help='Maximum number of iterations. Default: 20')
my_parser.add_argument('-criteria', action='store', type=float, default=1e-6, help='Convergence criteria. Default: 1e-6')
args = my_parser.parse_args()
# Extracting matrix A, vector b, and the initial guess x0 from the input arguments
A = np.array(args.A)
b = np.array(args.b)
x0 = np.array(args.x0)
# Checking the compatibility of size of A with b (n: dimension of the problem)
n=np.size(b)
if np.size(A) != n*n:
print('\n=====ERROR=====\n'
'Matrix A and vector b are not compatible in size.\n'
'===============\n')
exit()
# Checking the compatibility of size of x0 with A and b
if np.size(x0) != n:
print('\n=====ERROR=====\n'
'Initial guess does not match to matrix A and vector b in size.\n'
'===============\n')
exit()
# Converting matrix A from rowmajor to square form
A=A.reshape((n,n))
# Converting optional settings from the arguments
N_max=args.Nmax
criteria=args.criteria
# Iteration procedure
step=0
x=x0
r=b-np.einsum('ij,j',A,x0)
norm_residual=math.sqrt(np.einsum('j,j',r,r))
p=r
print('==========RESULT==========')
while norm_residual>criteria and step<N_max:
step = step+1
alpha=np.einsum('j,j',r,r)/np.einsum('j,j',p,np.einsum('ij,j',A,p))
x=x+alpha*p
r_backup=r
p_backup=p
r=r-alpha*np.einsum('ij,j',A,p)
beta=np.einsum('j,j',r,r)/np.einsum('j,j',r_backup,r_backup)
p=r+beta*p
norm_residual=math.sqrt(np.einsum('j,j',r,r))
print('step: '+
str(step)+'\n'+
'residual: '+
str(norm_residual)+'\n'+
'solution: '+
str(x)+
'\n--------------------------')
c4science · Help