from numpy import array, zeros, float, dot

from copy import copy

def GaussJordan(A,b):
"""
Returns the vector x such that Ax=b.
A is assumed to be an n by n matrix and b an n-element vector.
"""
n,m = A.shape
# should put in some checks of our assumptions, e.g. n == m
C = zeros((n,m+1),float)
C[:,0:n],C[:,n] = A, b
for j in range(n):
# First, do partial pivoting.
p = j # the current diagonal element is pivot by default
# look for alternate pivot by searching for largest element in column
for i in range(j+1,n):
if abs(C[i,j]) > abs(C[p,j]): p = i
if abs(C[p,j]) < 1.0e-16:
print "matrix is (likely) singular"
return b
# swap rows to get largest magnitude element on the diagonal
C[p,:],C[j,:] = copy(C[j,:]),copy(C[p,:])
# Now, do scaling and elimination.
pivot = C[j,j]
C[j,:] = C[j,:] / pivot
for i in range(n):
if i == j: continue
C[i,:] = C[i,:] - C[i,j]*C[j,:]
I,x = C[:,0:n],C[:,n]
return x