5
\$\begingroup\$

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 <= 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
\$\endgroup\$
4
\$\begingroup\$

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 <= 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)
\$\endgroup\$
  • \$\begingroup\$ Is it a good idea to use nested list comprehension? \$\endgroup\$ – Hungry Blue Dev Dec 25 '16 at 18:34
  • 1
    \$\begingroup\$ If it doesn't hurt readability, then yes. Otherwise no. No need to force the list comprehensions, but in the examples I gave it's easy to see to agree that it's an improvement. When it becomes questionable, then it's probably better to use a good old-fashioned loop. \$\endgroup\$ – janos Dec 25 '16 at 18:37
2
\$\begingroup\$

Just looking at the solve function:

  1. There's no docstring. What does this function do? What are the constraints on the inputs? What does it return?

  2. B.transpose() can be written B.T.

  3. The function relies on its inputs being np.matrix objects, so that the * operation is matrix multiplication (rather than element-wise multiplication). But this makes the function awkward to use — normally when working with NumPy you have array objects, not matrices, and so you have to remember to convert them to matrices whenever you called the function.

    Although this isn't a problem here, it's the kind of thing that becomes a pain later on, and so I recommend that even in simple programs like this, you stick to np.array (avoiding np.matrix), and when you need to compute a matrix multiplication, you use the np.dot function. (Or, in Python 3.5 or later, the new @ operator.)

  4. Instead of:

    X = np.linalg.inv(A) * B.transpose()
    

    use np.linalg.solve, like this:

    X = np.linalg.solve(A, B.T)
    

    (Or np.linalg.solve(A, B) if you are using np.array as recommended above.)

    The reason for doing this is that computing the inverse of a matrix requires more computation than solving a system of linear equations, and so can result in larger errors due to making more floating-point approximations.

  5. The vars argument has a default value of ['x', 'y'] but this only works if A is a 2×2 matrix. It would be better if the default value for vars supplied as many variables as needed. This could be implemented using an generator that yielded an unbounded sequence of variable names (see the function variable_names below).

  6. The limit_denominator method takes an argument max_denominator which the user of solve might want to change. It would make sense to add it as a keyword argument to the solve function.

Revised code (using np.array and not np.matrix as recommended in point 3 above):

from fractions import Fraction
from itertools import count, product
import numpy as np
from string import ascii_lowercase

def variable_names():
    """Generate variable names a, b, ..., z, a', b', ..., z', a'', ..."""
    for primes in count():
        for letter in ascii_lowercase:
            yield letter + "'" * primes

def solve(A, B, vars=None, max_denominator=1000000):
    """Solve the system of linear equations A @ vars = B.

    Arguments:
    A -- array with shape (m, m)
    B -- array with shape (m)
    vars -- iterable of variable names. If omitted,
        variable_names() is used.
    max_denominator -- values in solution are converted to fractions
        with denominator at most this.

    Returns: dictionary mapping variable names to their values.

    """
    X = np.linalg.solve(A, B)
    values = [Fraction(i).limit_denominator(max_denominator) for i in X]
    if vars is None:
        vars = variable_names()
    return dict(zip(vars, values))
\$\endgroup\$
  • \$\begingroup\$ I really like this version better. :-) BTW... What if I want to use x, y and z if I need 3 variables? How do you suggest one should do it? \$\endgroup\$ – Hungry Blue Dev Dec 31 '16 at 12:06
  • \$\begingroup\$ If you want the variable names to start with x, then write 'xyzabcdefghijklmnopqrstuvw' (or whatever order you prefer) instead of ascii_lowercase. \$\endgroup\$ – Gareth Rees Dec 31 '16 at 17:17
  • \$\begingroup\$ Well... I want to start with x only when I need 3 or less. The rest of the time... I'm happy to start from a. \$\endgroup\$ – Hungry Blue Dev Jan 1 '17 at 6:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.