A command-line utility to solve a system of linear equations I was re-checking my answers for a Math paper on my laptop. I needed to verify my solution for a system of linear equations. Unfortunately, I found no quick way to do that. That's why I wrote a (hopefully) simple Python program (using NumPy) to do this job quickly. General comments on the code are appreciated.

linearsolver.py

from sys import argv, exit
from fractions import Fraction
import numpy as np
import string


def solve(A, B, vars=['x', 'y']):
    X = np.linalg.inv(A) * B.transpose()
    values = []
    for item in X:
        values.append(Fraction(item[0, 0]).limit_denominator())
    return dict(zip(vars, values))


def vars(len):
    alph = list(string.ascii_lowercase)
    # x is the 24th alphabet
    start = 23 if len &lt;= 3 else 0
    return list(alph[start:start + len])


def interact():
    print "======================\nLinear Equation Solver\n======================\nHow many equations (and variables) ? ",

    len = int(raw_input())

    # Input the variable names
    vars = []
    print "Enter the variables involved :"
    for i in range(0, len):
        vars.append(raw_input())

    # Input the coefficients
    a = []
    print "Enter the coefficients of the variables in order, for each equation:"
    for i in range(1, len + 1):
        print "-----%d-----" % i
        temp = []
        for i in range(0, len):
            temp.append(float(raw_input()))
        a.append(temp)

    # Input the values they equal to
    print "Enter the constants for the RHS :"
    b = []
    for i in range(0, len):
        b.append(float(raw_input()))

    A = np.matrix(a)
    B = np.matrix(b)
    answer = solve(A, B, vars)
    print_answer(answer)


def print_answer(answer):
    for variable, value in answer.items():
        print "%4s = %4s" % (variable, value)


def main(argv):
    num = len(argv)
    varlist = None
    if num == 1:
        # Only the script name: enter interactive mode
        interact()
        exit()
    elif num == 3:
        script, a, b = argv
    elif num == 4:
        script, a, varlist, b = argv
    else:
        print "Too many arguments."
        sys.exit()

    A = np.matrix(a)
    B = np.matrix(b)

    rows = A.shape[0]
    cols = A.shape[1]

    if rows != cols:
        print "Number of rows must be equal to number of columns."
        exit()

    if varlist is None:
        answer = solve(A, B, vars(rows))
    else:
        answer = solve(A, B, list(varlist))
    print_answer(answer)


if __name__ == '__main__':
    main(argv)


Some example runs:



$ python linearsolver.py "2 5 -9 3; 5 6 -4 2; 3 -4 2 7; 11 7 4 -8" "151 103 16 -32"
   a =    3
   c =  -11
   b =    5
   d =    7
$ python linearsolver.py "5 -11 9; 13 7 -10; -6 5 -8" "gkm" "-125 2 60"
   k =   12
   m =    3
   g =   -4
$ python linearsolver.py
======================
Linear Equation Solver
======================
How many equations (and variables) ?  7
Enter the variables involved :
s
t
u
v
w
x
y
Enter the coefficients of the variables in order, for each equation:
-----1-----
2
-7
3
5
-7
11
8
-----2-----
5
7
-1
9
-6
3
-5
-----3-----
13
11
8
7
4
2
1
-----4-----
12
1
6
5
9
3
11
-----5-----
1
-4
7
-3
9
-2
6
-----6-----
4
1
-6
-3
8
-7
-4
-----7-----
-7
-2
5
1
9
3
-4
Enter the constants for the RHS :
114
-11
-25
-37
4
-91
-50
   s =    5
   u =    8
   t =   -7
   w =   -9
   v =   -6
   y =   -3
   x =    2
List comprehensions

List comprehensions are a cool feature that could shorten your code, at many places.
For example instead of this:


    values = []
    for item in X:
        values.append(Fraction(item[0, 0]).limit_denominator())



You can write like this:

    values = [Fraction(item[0, 0]).limit_denominator() for item in X]


Another example:


    vars = []
    for i in range(0, len):
        vars.append(raw_input())



Better this way:

    vars = [raw_input() for _ in range(len)]


There are some more tips in there:


range(len) is equivalent to range(0, len)
When you don't use the loop variable in the loop body, the convention is to name the loop variable _


There are other places in your code that you could shorten using list comprehensions, I suggest to rewrite them all.

Magic values

The number 23 is a magic value here:    


    # x is the 24th alphabet
    start = 23 if len &lt;= 3 else 0



Although the comment explains it, it would be better to put the value in a constant with a name that describes its purpose.
And instead of the hard-coded number, I suggest to express it in a programmatic way, for example:

POS_X = ord('x') - ord('a')


When written this way, we don't even have to know what the value really is (23), the important thing is its relative position from the letter 'a'.

Code duplication

The solve calls are duplicated here, except for the 3rd parameter:


    if varlist is None:
        answer = solve(A, B, vars(rows))
    else:
        answer = solve(A, B, list(varlist))
    print_answer(answer)



It would be better to reorganize the code to eliminate such duplication:

    if varlist is None:
        varlist = vars(rows)

    answer = solve(A, B, varlist)
    print_answer(answer)


Of course, in your current implementation varlist is not really a list.
It would be better to fix that, and make it a list when you assign it earlier, for example:

elif num == 4:
    script, a, symbols, b = argv
    varlist = list(symbols)