AMS 209: Project: Python Code

'''------------------------------------------------------------------
 Final Project Type A – Numerical Linear Algebra:
 Gaussian Elimination with and without Pivoting for Solving Linear
 Systems

 This code is written by Geetanjali Rakshit for AMS 209.
---------------------------------------------------------------------
'''

'''

Directory structure:

   project/
       |-- LinAlg/setup_mode.f90
       |          read_data.f90
       |          write_data.f90
       |          linear_solve.f90
       |          LU_decomp.f90
       |          backward_solve.f90
       |          forward_solve.f90
       |          write_to_screen.f90          
       |	  makefile
       |
       |-- pyRun/pyRun.py
       

All .dat files containing A, b and x values as well as the plots produced in .png files are stored in the LinAlg directory, as is setup.init.
'''

import os
import sys
import glob
import subprocess
import numpy as np
import matplotlib.pyplot as plt


def createFiles(ifCustomValues):
	# This function either takes A and b as input from user, or loads a set of predefined values for A and b using setValues(eqSet) and 		# then writes them to .dat files. It also returns the matrix A and the vector b.

	global eqSet


	if(ifCustomValues):
		if(glob.glob('A_*.dat')):
			cmd = 'ls A_*.dat | wc -l'
			i = int(subprocess.check_output(cmd, shell=True))
		else:
			i = 0

		eqSet = i+1

		A = []
		b = []

		f1 = open('A_'+str(eqSet)+'.dat', 'w')
		f2 = open('b_'+str(eqSet)+'.dat', 'w')
		
		n = int(input("Enter the number of variables in the system of equations: "))
		
		for i in range(0, n):
			A_matrix_row_str = ''
			print("Enter the values for row "+(str(i+1))+" of A matrix ")
			A.append([])
			
			for j in range(0, n):
				A[i].append(int(input()))
				A_matrix_row_str = A_matrix_row_str+str(A[i][j])+","
				
			f1.write(str(A_matrix_row_str)+"\n")
			
			print("Enter the value for index "+(str(i+1))+" of B vector ")
			b.append(int(input()))
			f2.write(str(b[i])+"\n")
			
		f1.close()
		f2.close()
	
		A = np.asarray(A)
		b = np.asarray(b)

	else:
		A, b = setValues(eqSet)
		
		f1 = open('A_'+str(eqSet)+'.dat', 'w')
		f2 = open('b_'+str(eqSet)+'.dat', 'w')
		for i in range(0, len(b)):
			A_matrix_row_str = ''
			for j in range(0, len(b)):
				A_matrix_row_str = A_matrix_row_str+str(A[i][j])+","
				
			f1.write(str(A_matrix_row_str)+"\n")
			f2.write(str(b[i])+"\n")
		f1.close()
		f2.close()

	#Generating a new "setup.init" input file.
	filename = "setup.init"
	f = open(filename, 'w')
	f.write(str(eqSet)+"\n")
	f.write(str(len(A))+"\n")
	f.close()

	print("A:")
	print(A)

	print("b:")
	print(b)
	
	return (A, b)

	
def setValues(eqSet):
	#This function sets A and b to predefined values and returns them.

	# Initialize A and b with predefined values

	if(eqSet == 1):
		A = np.array([[1, 1, -1], [1, 2, -2], [-2, 1, 1]])

		b = np.array([1, 0, 1])

	if(eqSet == 2):
		A = np.array([[4, 3, 2, 1], [3, 4, 3, 2], [2, 3, 4, 3], [1, 2, 3, 4]])

		b = np.array([1, 1, -1, -1])

	if(eqSet == 3):
		A = np.array([[1, -1, 1, -1], [-1, 3, -3, 3], [2, -4, 7, -7], [-3, 7, -10, 14]])

		b = np.array([0, 2, -2, -8])

	return (A, b)


def make_make():
	# Compile the Fortran code.
	
	# Change directory from "pyRun/" to "LinAlg/" to write A and b to files.
	os.chdir('../LinAlg')	
	
	if os.path.isfile('linear_solve.exe'):
		os.system('make clean')
		os.system('make')
	else:
		os.system('make')
	

def runFortranSolver():

	# This function executes the Fortran executable, "linear_solve.exe" and returns the solution in an array x.
	x = []
	cmd = './linear_solve.exe'
	os.system(cmd)

	# Read solutions for Fortran from "x_i" file
	x = readFortranSolution()
	return(x)


def plotData(A, b, x):
	# This function plots A, b, and x.

	global eqSet

	#Plot A matrix
	plt.imshow(A, interpolation='nearest', cmap=plt.cm.ocean)
	plt.colorbar()
	plt.title('Matrix A')
	plt.savefig('A_matrix'+str(eqSet)+'.png')
	plt.show()

	#Plot x and B vectors
	fig = plt.figure()
	ax = fig.add_subplot(1,2,1)
	plt.imshow(np.matrix(x).transpose(), interpolation='nearest', cmap=plt.cm.ocean)
	plt.colorbar()
	plt.title('Vector x')
	
	ax = fig.add_subplot(1,2,2)
	plt.imshow(np.matrix(b).transpose(), interpolation='nearest', cmap=plt.cm.ocean)
	plt.colorbar()
	plt.title('Vector b')
	plt.savefig('x_and_b_vectors'+str(eqSet)+'.png')
	plt.show()
	

def readFortranSolution():
	# Reads the solution saved in file by the fortran code and into an array and returns x.

	filename = "x_"+str(eqSet)+".dat"
	f = open(filename,'r')
	x = []
	for line in f:
		x.append(float(line))
	x = np.asarray(x)
	f.close()
	return(x)


def checkSolution(A, b, x_fortran, threshold):
	# This function outputs Pass/Fail by comparing the Fortran and Python solutions, based on a threshold value. If FAIL, it outputs the 		# error together with the Fortran and Python solutions to screen, else prints PASS.

	x_python = np.dot(np.linalg.inv(A), b)
	
	#Compare solutions and print the error
	if(np.linalg.norm(x_fortran - x_python) > threshold):
		print("FAIL")
		print("Fortran solution: "+str(x_fortran))
		print("Python solution: "+str(x_python))
		print("Error: "+str(abs(x_python-x_fortran).transpose()))
	else:
		print("PASS")
        

def runScheduler(ifCustomValues):
	# The runScheduler function performs all the steps (initialize, run Fortran code, compare with numpy solution and plotting) 		# sequentially for a system of equations.

	# Create A matrix and B vector and write them to .dat files
	A, b = createFiles(ifCustomValues)
	
	# Run Fortran code
	x_fortran = runFortranSolver()

	# Compare Fortran and Python solutions
	checkSolution(A, b, x_fortran, threshold)

	# Visualize the data
	plotData(A, b, x_fortran)



if __name__ == '__main__':

	# Compile the fortran code.
	make_make()

	# Set a threshold value to compare the Fortran solution with the Python numpy library solution
	threshold = 0.00001

	# The user can choose to enter his own system of equations to solve or choose from a set of 3 predefined equations (mentioned in the 		# project handout). The runScheduler function performs all the steps (initialize, run Fortran code, compare with numpy solution and 		# plotting) sequentially for a system of equations.

	eqSet = 1	# This variable identifies a set of equations.
	ifCustomValues = input("Do you want to solve for your own set of equations? (y/n): ")

	if(ifCustomValues == 'y'):
		runScheduler(True)
	else:
		customValues = False
		while(eqSet < 4):
			print("\nRunning for "+str(eqSet)+" system of equations...")
			runScheduler(False)
			eqSet = eqSet+1