Newton's method to solve cubic equations

Question

I have used the Newton-Raphson method to solve Cubic equations of the form $$ax^3+bx^2+cx+d=0$$ by first iteratively finding one solution, and then reducing the polynomial to a quadratic $$a1*x^2+b1*x+c=0$$ and solving for it using the quadratic formula. It also gives the imaginary roots.

def deg3(a,b,c,d,g):  #using newton raphson method, I designed this to solve degree3 equations
    y=a*g**3+b*g**2+c*g+d
    return y
def solvedeg3equation():
    '''solves for cubic equations of the form ax^3+bx^2+cx+d=0 with the rough estimate of error within e'''
    import math
    e=float(input("e= "))
    a=float(input("a= "))
    b=float(input("b= "))
    c=float(input("c= "))
    d=float(input("d= "))
    count=1
    g=0.01
    while abs(deg3(a,b,c,d,g))>e and count<=100 and not d==0:
        count=count+1
        if 3*a*g**2+2*b*g+c==0:
            g=g+0.001
        g=g-deg3(a,b,c,d,g)/(3*a*g**2+2*b*g+c)


   if count<=100:
        if d==0:
            a1=a
            b1=b
            c1=c
            g=0
        else:
            c1=-d/g                                 #This is generation 3, which provides all three solutions, including imaginary solutions, if any.
            a1=a
            b1=(c1-c)/g                             #imagg=imaginary part of root, realg=real part
        if b1**2-4*a1*c1<0:
            realg=-b1/(2*a1)
            imagg=math.sqrt(4*a1*c1-(b1**2))/(2*a1)
            if a1>0:
                g2=str(realg)+'+'+str(imagg)+'i'
                g3=str(realg)+'-'+str(imagg)+'i'
            if a1<0:
                g2=str(realg)+'+'+str(-imagg)+'i'
                g3=str(realg)+'-'+str(-imagg)+'i'                
        if abs(b1**2-4*a1*c1)<e:
            g2=-b1/(2*a1)
            g3=None
        if b1**2-4*a1*c1>e:
            g2=(-b1+math.sqrt(b1**2-4*a1*c1))/2*a1
            g3=(-b1-math.sqrt(b1**2-4*a1*c1))/2*a1
        print("Solved. The best guess is:",g,'and',g2,'and',g3)
        print("Iterations required: ",count)
    else:
        print("Maximum iterations exceeded ")
        print("Iterations: ",count,"Current guess: ",g)

I would like to know how I can make this code more efficient, both mathematically and structurally. Some places which could do with a more efficient piece of code, according to me, are:

1. To avoid crashing of the program when the derivative equals zero, I have used the provision of g+0.001, to avoid the case. Is there a better way to solve this problem? Both mathematical optimizations or structural (code-related) are welcome.

2. Reducing the number of conditional statements and elementary operations (like comparisons, assignments and arithmetic operations) to reduce the computation time

3. Improve the Newton-Raphson algorithm mathematically to reduce the number of iterations required (since 3-4 digit numbers require more than 100 iterations currently). I'm incorporating the provision of better guesses.

I am a beginner to Python (this is my second program) and programming in general, so please try to suggest as simple modifications as possible. Also, I would appreciate any tips on the general nature of my writing style.

Answer 1

Handling of complex roots

When there are complex roots, you have two cases, for a1 > 0 and a1 < 0. You could just collapse the two cases by using abs(a1) when defining imagg.

Better yet, Python has built-in support for a complex type. Why not use it instead of building a string? Then you don't even have to worry about the sign of the imaginary part for the sake of a pretty printout.

Separation of concerns

Your solvedeg3equation() function performs input, calculation, and output. That may do the job for you, but it ensures that your code will never be reusable, other than by copying and pasting — or by piping the input to you and attempting to parse the output! Furthermore, it makes your code just as difficult to unit test. What you want is a function that accepts a polynomial (more on that shortly) and returns a 3-tuple of solutions. The caller would be responsible for input and output.

Notation

Representing a third-degree polynomial as variables a, b, c, and d is unwieldy. You can't pass it around easily as one entity. The proliferation of variables imposes a mental load. Most importantly, I think that it leads to poor notation in your code.

I propose the following notation:

• x0, x1, x2: The first, second, and third roots (better than g, g2, g3)
• Polynomial(a, b, c, d): An object representing the polynomial a x³ + b x² + c x + d
• f: The polynomial to be solved (better than a, b, c, d)
• df_dx: The derivative of f (better than 3*a*g**2+2*b*g+c)
• q: The quadratic polynomial remaining after x0 has been found (better than c1=-d/g; a1=a; b1=(c1-c)/g)
• q.discriminant(): Shorthand for b1**2-4*a1*c1
• tolerance: better than e (which unfortunately looks like it's related to a, b, c, d)

Furthermore, you shouldn't hard-code a 100-iteration limit, especially twice as a magic number in the code.

Polynomial class

As noted above, having a class for polynomials is extremely useful for your problem. In addition to the polynomial to be solved, you need its derivative, and you also have a quadratic equation to solve. A polynomial class lets you

• pass the polynomial into the solver function conveniently
• collapse several variables into one
• evaluate the value of the polynomial easily
• take its derivative

Here is an implementation:

class Polynomial(object):
    def __init__(self, *coeffs):
        """
        Polynomial(3, 5, 0, -2) represents f(x) = 3x^3 + 5x^2 - 2.
        """
        self.coeffs = list(coeffs)[::-1]
        while self.coeffs[-1] == 0:
            self.coeffs.pop()

    def __call__(self, x):
        """
        >>> f = Polynomial(3, 5, 0, -2)
        >>> f(2)
        42
        """
        return sum([self[n] * x ** n for n in range(len(self.coeffs))])

    def __getitem__(self, n):
        """
        Gets the coefficient for x^n
        >>> f = Polynomial(3, 5, 0, -2)
        >>> f[2]
        5
        """
        return self.coeffs[n] if n < len(self.coeffs) else 0

    def __str__(self):
        """
        >>> f = Polynomial(3, 5, 0, -2)
        >>> str(f)
        3 x^3 + 5 x^2 + 0 x^1 + -2
        """
        return ' + '.join(["%d x^%d" % (self[n], n) for n in range(len(self.coeffs))][::-1])

    def degree(self):
        """
        >>> f = Polynomial(3, 5, 0, -2)
        >>> f.degree()
        3
        """
        return len(self.coeffs) - 1

    def derivative(self):
        """
        >>> f = Polynomial(3, 5, 0, -2)
        >>> f.derivative()
        9 x^2 + 5 x^1 + 0 x^0
        """
        deriv = [n * self[n] for n in range(len(self.coeffs) - 1, 0, -1)]
        return Polynomial(*deriv)

    def discriminant(self):
        """
        For quadratic polynomial a x^2 + b x + c, returns b^2 - 4 * a * c

        >>> f = Polynomial(4, -1, 2)
        >>> f.discriminant()
        -31
        """
        if self.degree() != 2:
            raise ValueError, "Discriminant of polynomial of degree %d not supported" % (self.degree())
        return self[1] ** 2 - 4 * self[2] * self[0]

Error handling

If something goes wrong, don't just print something. Raise an exception, which ensures that the caller will notice the problem!

class SolutionNotFound(ValueError):
    """
    Raised when a root of a polynomial cannot be found.
    """

Iteration limit

Python has a language feature for executing an else clause when a loop terminates naturally through exhaustion (i.e., when its condition becomes false, rather than by an early break). You can use it like this:

for _ in range(max_iterations):
    if abs(f(x0)) <= tolerance:
        break
    # TODO: refine x0 here
else:
    # Loop exhaustion, i.e. max_iterations reached
    # TODO: raise an error
# TODO: Calculate x1 and x2 here

Pretty!

All of that setup lets your code look like mathematics, not like a C program. Now you can clear your head of the minutiae and focus on things that matter, like your mathematical technique.

I've made a few remarks in the code comments as well.

from math import sqrt

def solve_degree_3_polynomial(f, tolerance=0.00000001, initial_guess=0.01, max_iterations=100):
    if f.degree() != 3:
        raise ValueError, "Input must be a polynomial of degree 3"

    # If 0 is a root, make x0 exactly 0 to trigger a special case for
    # the quadratic equation below.
    x0 = 0 if f(0) == 0 else initial_guess
    df_dx = f.derivative()
    for _ in range(max_iterations):
        if abs(f(x0)) <= tolerance:
            break
        if df_dx(x0) == 0:
            x0 += 0.001
        x0 = x0 - f(x0) / df_dx(x0)
    else:
        raise SolutionNotFound, "Exceeded %d iterations.  Current guess: f(%d) = %d" % (max_iterations, x0, f(x0))

    # q = Quadratic
    q = Polynomial(f[3], f[2], f[1]) if x0 == 0 else \
        Polynomial(f[3], (-f[0] / x0 - f[1]) / x0, -f[0] / x0)

    # These three cases are mutually exclusive, right?  Then write them that way.
    if abs(q.discriminant()) < tolerance:
        # I think that returning a double root is better than returning
        # None for one of the roots.  It's mathematically more correct,
        # and less likely to cause bugs involving NoneType.
        x1 = x2 = -q[1] / (2 * q[2])

    elif q.discriminant() < 0:
        # You could just let cmath.sqrt() take care of this for you instead,
        # in which case all three cases collapse down to one!
        x1 = complex(-q[1] / (2 * q[2]), +sqrt(-q.discriminant()) / (2 * q[2]))
        x2 = complex(-q[1] / (2 * q[2]), -sqrt(-q.discriminant()) / (2 * q[2]))

    else:
        # If you change
        #     from math import sqrt
        # to
        #     from cmath import sqrt
        # then the two lines below handle all three cases, regardless of
        # whether the discriminant is negative, 0, or positive.
        x1 = (-q[1] + sqrt(q.discriminant())) / (2 * q[2])
        x2 = (-q[1] - sqrt(q.discriminant())) / (2 * q[2])

    return x0, x1, x2

Example usage

f = Polynomial(3, -1, 6, -2)
(x0, x1, x2) = solve_degree_3_polynomial(f)

print "f(x) = %s" % (f)
print "f(%s) = %s" % (x0, f(x0))
print "f(%s) = %s" % (x1, f(x1))
print "f(%s) = %s" % (x2, f(x2))

… prints

f(x) = 3 x^3 + -1 x^2 + 6 x^1 + -2 x^0
f(0.333333333333) = 2.423339307e-13
f((3.44613226844e-13+1.41421356237j)) = (-4.36495284361e-12-1.7763568394e-15j)
f((3.44613226844e-13-1.41421356237j)) = (-4.36495284361e-12+1.7763568394e-15j)

Next steps

As noted in the comments, I would recommend using cmath.sqrt() to collapse the three cases for the discriminant into one.

It would also be a good idea to decompose the cubic equation solver into a generic Newton's method solver for any polynomial, followed by a quadratic equation solver.

def solve_degree_3_polynomial(f, ...):
   if f.degree() != 3:
       raise ValueError
   x0 = solve_polynomial_by_newton(f, ...)
   q = Polynomial(f[3], (-f[0] / x0 - f[1]) / x0, -f[0] / x0)
   x1, x2 = solve_degree_2_polynomial(q)
   return x0, x1, x2

Answer 2

In addition, I would suggest that you can win some speed by evaluating the coefficients of the quadratic equation as:

a1= a
b1 = a1*g + b
c1 = b1*g + c

where, g is the root found out using Newton Raphson.

In this way you avoid using the division operator (like in your method, c1 = -d/g , ...) - small but some gain at least! Besides, no fears if the denominator becomes 0.