8
\$\begingroup\$

I wrote this equation parser and solver, I feel like it is well documented and tested:

import doctest

def swap_sign(num):
    """
    Swaps the sign of a number-representing string.

    >>> swap_sign('9')
    '-9'
    >>> swap_sign('-4x')
    '4x'
    """
    if num.startswith('-'):
        return ''.join(num[1:])
    elif num.startswith('+'):
        return '-'+''.join(num[1:])
    else:
        return '-'+num

def add_one_before_x(expr):
    """
    It is a common usage to omit the one before
    the x saying 'x' when you really mean '1x'.

    >>> add_one_before_x('x +2x = 9 -x^2')
    '1x +2x = 9 -1x^2'
    """
    if expr.startswith('x'):
        expr = '1'+expr
    return expr.replace(' x',' 1x').replace(' -x',' -1x').replace(' +x',' +1x')

def move_operators(expr):
    """
    Sticks the operators the next term for easier
    further parsing.

    >>> move_operators('3 + 9x = 3 - 3x^2 + 4')
    '3 +9x = 3 -3x^2 +4'
    """
    return expr.replace('+ ','+').replace('- ','-').replace('--','-')

def move_to_the_same_side(expr):
    """
    Moves all the numbers and x-s to one side,
    changing sign when needed.

    >>> move_to_the_same_side("3x -4x^2 +5 -2 = 2x")
    ['3x', '-4x^2', '+5', '-2', '-2x']
    >>> move_to_the_same_side("9x = 1x^2 -4 +2 -1 +7x")
    ['9x', '-1x^2', '4', '-2', '1', '-7x']
    """
    if ' = ' not in expr:
        raise ValueError("There is no second term, remember spaces around the equal sign.")
    left,right = expr.split(' = ')
    return left.split(' ') + [swap_sign(token) for token in right.split(' ')]

def to_normal_form(expr):
    """
    Performs further parsing on all the coefficients on one side
    and return the coefficients (a,b,c).

    >>> to_normal_form(['3x', '-4x^2', '+5', '-2', '-2x'])
    (-4, 1, 3)
    """
    bare_nums = [i for i in expr if 'x' not in i]
    xes = [i for i in expr if 'x' in i and '^2' not in i]
    two_pow_xes = [i for i in expr if 'x' in i and '^2' in i]

    final_num = sum([int(n) for n in bare_nums])
    final_x = sum([int(x.replace('x','')) for x in xes])
    final_pow_x = sum([int(x.replace('x^2','')) for x in two_pow_xes])

    return final_pow_x, final_x, final_num


def first_grade_solve(coeffs):
    """
    Solves the first grade equations using the trivial
    first grade formula. Also solves equations of grade 0.

    >>> first_grade_solve((0,0,3))  # 0x + 3 = 0
    'Never'
    >>> first_grade_solve((0,2,3))  # 2x + 3 = 0
    (-1.5,)
    """

    _,a,b = coeffs
    if a == 0:
        return "Every time" if b == 0 else "Never"
    return ((-b) / a,)

def second_grade_solve(coeffs):
    """
    Solves second grade equations using the well known
    second grade formula.

    >>> second_grade_solve((1,2,1)) # 1x^2 + 2x + 1 = 0
    (-1.0, -1.0)
    >>> second_grade_solve((1,5,6)) # 1x^2 + 5x + 6 = 0
    (-2.0, -3.0)
    """
    a,b,c = coeffs
    delta = b**2 - 4*a*c
    return (((-b)+delta**.5)/(2*a),
            ((-b)-delta**.5)/(2*a))

def solve(coefficients):
    """
    Dispatches solving to the correct method or aborts
    if the equation grade is too high.

    >>> solve((1,5,6))
    (-2.0, -3.0)
    >>> solve((0,2,4))
    (-2.0,)
    """
    if coefficients[0] == 0:
        return first_grade_solve(coefficients)
    elif len(coefficients) == 3:
        return second_grade_solve(coefficients)
    raise NotImplementedError("Only 0th, 1st and 2nd grade equations can be solved")

def parse_and_solve(expr):
    """
    Connects the other functions to provide full equation
    solving.

    >>> parse_and_solve("2x - 4 = 0")
    (2.0,)
    >>> parse_and_solve("1 = 2")
    'Never'
    >>> parse_and_solve("x^2 + 2x = -1")
    (-1.0, -1.0)
    """
    simpler = add_one_before_x(move_operators(expr))
    same_sided = move_to_the_same_side(simpler)
    normal_form = to_normal_form(same_sided)
    return solve(normal_form)

def interface():
    print("Welcome to the symbolic equation solver.")
    print("Please enter equations in the form:")
    print("3x - 4x^2 + 5 - 2 = 2x + x^2\n")
    while True:
        expr = input("> ")
        try:
            result = parse_and_solve(expr)
            if result in ("Every time","Never"):
                print(result)
            else:
                print('\n'.join(['x = '+str(x) for x in result]))
        except Exception as E:
            print("Invalid Expression:", E)

if __name__ == "__main__":
    doctest.testmod()
    interface()
\$\endgroup\$
  • 1
    \$\begingroup\$ Please add more context to your question, such as explain what your code does. This makes the life of reviewers easier, and will most likely attract more reviewers. \$\endgroup\$ – TheCoffeeCup Apr 21 '15 at 23:16
  • 2
    \$\begingroup\$ @MannyMeng The purpose of the code seems obvious to me. \$\endgroup\$ – 200_success Apr 21 '15 at 23:53
6
\$\begingroup\$
  1. This is well documented code, clearly written and sensibly factored into functions.

  2. The docstrings need a bit of work. It's best to imagine reading a docstring from the point of view of a user who wants to know how to call a function and so has just typed help(function) in the Python interpreter. But if they see something like this:

    >>> help(solve)
    Help on function solve in module __main__:
    
    solve(coefficients)
        Dispatches solving to the correct method or aborts
        if the equation grade is too high.
    

    then they are none the wiser. What do I pass as coefficients? What does it return? What does it mean by "correct method"? The user was hoping for something more like this:

    solve(coefficients)
        Given the sequence of coefficients of x^0, x^1, ..., in a
        polynomial, returns a tuple of roots of that polynomial.
        Raises NotImplementedError for polynomials of degree
        greater than 2.
    
  3. I think degree is the more usual term, not grade.

  4. first_grade_solve sometimes returns a length-1 tuple of solutions, and sometimes returns a string. It's not a good idea to write functions that return different kinds of result in exceptional situations, because it's too easy for the caller to forget to handle the exceptional case and write something like:

    x = first_grade_solve(coeffs)[0]
    

    which will mistakenly set x to the string 'N' if the equation has no solution. It's better to deal with an exceptional situation by raising an exception.

  5. Running all the doctests each time the program is executed is a mistake. If the user just wants to run the program then they probably don't want to wait for the tests to finish running first. On the other hand, if they just want to test the program they probably don't want to have to deal with the interactive interface. So just run the program; a user who wants to run the doctests can easily do so using the -m command-line option:

    python -m doctest program.py
    
  6. The parsing relies on repeated string substitution, which is hard to make reliable:

    >>> parse_and_solve('x  = 1')
    Traceback (most recent call last):
        ...
        final_num = sum([int(n) for n in bare_nums])
    ValueError: invalid literal for int() with base 10: ''
    

    and is prone to false positives:

    >>> parse_and_solve('xx = 1')
    (1.0,)
    >>> parse_and_solve('x2x = 12')
    (1.0,)
    

    It's usually more reliable and robust to use a real parser.

Here's how you might implement an equation parser using PyParsing:

from collections import Counter
from pyparsing import Group, Literal, Optional, Regex, Suppress, ZeroOrMore

def equation_grammar():
    """Return a PyParsing grammar for a polynomial equation."""
    plus = Literal('+').setParseAction(lambda:1)
    minus = Literal('-').setParseAction(lambda:-1)
    sign = plus | minus
    number = Regex('[0-9]+').setParseAction(lambda t:int(t[0]))
    exponent = Optional(Suppress('^') + number, default=1)('exponent')
    power = 'x' + exponent.setParseAction(lambda t:t[0])
    atom = number('coeff') + Optional(power) | power
    monomial = Optional(sign('sign')) + atom
    polynomial = Group(monomial) + ZeroOrMore(Group(sign('joiner') + monomial))
    return polynomial('left') + '=' + polynomial('right')

_grammar = equation_grammar()

def parse(s):
    """Parse the string s as a polynomial equation in x; gather terms to
    express it as a power series in x equated to zero; and return a
    list of the coefficients.

    >>> parse('x = 1')
    [1, -1]
    >>> parse('7x^7 + -5x^5 + 3x^3 + -1 = 3x^6 - x^4 + 7')
    [7, -3, -5, 1, 3, 0, 0, -8]

    """
    equation = _grammar.parseString(s)
    series = Counter()
    for side, terms in ((+1, equation.left), (-1, equation.right)):
        for t in terms:
            n = side * t.get('joiner', 1) * t.get('sign', 1) * t.get('coeff', 1)
            series[t.get('exponent', 0)] += n
    degree = max(e for e, c in series.items() if c)
    return [series[e] for e in range(degree, -1, -1)]

I've chosen the output format of parse so that it can be fed directly into numpy.roots:

>>> from numpy import roots
>>> roots(parse('7x = 21'))
array([ 3.])
>>> roots(parse('x^2 + 3x = -2'))
array([-2., -1.])
>>> roots(parse('2x^3 + 22x = 12x^2 + 12'))
array([ 3.,  2.,  1.])
>>> roots(parse('x^5 = 1'))
array([-0.80901699+0.58778525j, -0.80901699-0.58778525j,
        0.30901699+0.95105652j,  0.30901699-0.95105652j,
        1.00000000+0.j        ])
\$\endgroup\$

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.