SciPy ODE Solver

I originally computed this ODE's solution with Wolfram's Mathematica and quite a bit of manual labour. Since I've been managing to finish homework ahead of time for the past weeks, I've decided to come back to this problem and solve it with brains instead of brawn. SciPy's ODE integrator is a little trickier to learn than Mathematica's, but does the job competently and I enjoy Python.
$$-\frac{1}{2}\ddot{\psi}(u)+(u^4-e)\psi(u)=0$$

Ground State Wave Function ($e_0$)
The program repeatedly shrinks the interval containing "e" (eigen energies of the solutions for Schrodinger's equation with potential $V(u)=u^4$). Estimates for "e" that deviate from the exact value eventually make the solution blow up due to numerical errors. Closer is your estimate to the true values of "e", longer it takes for the solution to blow up. I was able to find estimates for the ground state energy and for the excited states up to $e_5$ with 7-8 significant digits of accuracy. I called the eigen energy "a" in the program instead of "e" because I feel like "e" should be reserved to Euler's number... I ended up spending much more time than I wanted on this, but I got a very good grasp of SciPy's ODE solver.

import scipy.integrate as integrate
import matplotlib.pyplot as plt
import numpy as np

FIG, AX = plt.subplots()

def setup_plot(min_t, max_t, min_y, max_y):
    plt.axis([min_t, max_t, min_y, max_y])
    AX.set_autoscale_on(False)
    plt.grid(True)

def add_plot(t, y, label):
    AX.plot(t, y, label=label)

def show_plots():
    plt.legend()
    plt.show()

def run_solver(solver, initial_values, t_start, t_end, delta_t):
    """
    Runs solver for the given interval with step size delta
    """
    num_steps = np.floor((t_end-t_start)/delta_t)
    t = np.zeros((num_steps, 1))
    y = np.zeros((num_steps, 1))
    t[0] = t_start
    y[0] = initial_values[0]
    i = 1
    while solver.successful() and i < num_steps:
        solver.integrate(solver.t + delta_t)
        t[i], y[i] = solver.t, solver.y[0]
        i += 1
    return t, y

def get_derivatives(a):
    """
    Creates array of derivatives representing the ODE
    """
    # ODE
    # -(1/2)*f'' + (t**4-a)*f=0
    # f'' = 2*(t**4-a)*f
    # derivatives
    # d/dt(y[0]) = y[1]
    # d/dt(y[1]) = 2*(t**4-a)*y[0]
    derivatives=lambda t, y: np.array([y[1], 2.*(t**4.-a)*y[0]])
    return derivatives

def get_solution(a, initial_values, t_start, t_end, delta_t):
    """
    Gets numerical solution for the given parameters and interval
    """
    # creates LSODA solver
    solver = integrate.ode(get_derivatives(a)).set_integrator('lsoda', rtol=1e-12)
    solver.set_initial_value(initial_values, t_start)
    t, y = run_solver(solver, initial_values, t_start, t_end, delta_t)
    return t, y

def update_interval(a_start, a_end, delta_a, initial_values, t_start, t_end, delta_t):
    """
    If "a" is smaller than the exact value, the solution blows up to a positive value,
    while if "a" is greater than the exact value, solution blows up to a negative value.
    The sequence of choices of "a" that causes the solution to blow up positive and
    blow up negative defines an interval that contains the exact value of "a".
    If we repeatedly call this method with smaller and smaller delta a, we should get
    smaller and smaller intervals and therefore better estimates for "a".
    """
    a = a_start
    t, y = get_solution(a, initial_values, t_start, t_end, delta_t)
    previous_a, previous_y = a, y[-1]
    while a <= a_end:
        a += delta_a
        t, y = get_solution(a, initial_values, t_start, t_end, delta_t)
        if previous_y/y[-1] < 0:
            return previous_a, a, t, y
        previous_a, previous_y = a, y[-1]
    return a_start, a_end, t, y

def compute_interval(a_start, a_end, delta_a, initial_values, t_start, t_end, delta_t, num_digits):
    for _ in range(num_digits):
        a_start, a_end, t, y = update_interval(a_start, a_end, delta_a, initial_values, t_start, t_end, delta_t)
        label = "%.15f < a < %.15f" % (a_start, a_end)
        add_plot(t, y, label)
        print(label)
        delta_a = delta_a/10.

if __name__ == '__main__':
    # set parameters here
    t_start, t_end, delta_t = 0., 5, 0.01
    num_digits = 6

    # setup plot
    setup_plot(t_start, t_end, -1, 1)

    # experiment with values of "a" here
    a = 10
    t, y = get_solution(a, [1, 0], t_start, t_end, delta_t)
    #add_plot(t, y, "a = %.30f" % a)

    # compute interval for a with given accuracy
    compute_interval(0.6, 0.7, 0.01, [1, 0], t_start, t_end, delta_t, num_digits)   #e0
    #compute_interval(2.3, 2.4, 0.01, [0, 1], t_start, t_end, delta_t, num_digits)   #e1
    #compute_interval(4.6, 4.7, 0.01, [1, 0], t_start, t_end, delta_t, num_digits)   #e2
    #compute_interval(7.3, 7.4, 0.01, [0, 1], t_start, t_end, delta_t, num_digits)   #e3
    #compute_interval(10.2, 10.3, 0.01, [1, 0], t_start, t_end, delta_t, num_digits) #e4
    #compute_interval(13.3, 13.4, 0.01, [0, 1], t_start, t_end, delta_t, num_digits) #e5

    show_plots()

First Excited State Wave Function ($e_1$)

$e_2$

$e_3$

$e_4$

SymPy

I've been using Wolfram Alpha for my integration, derivation and algebraic manipulation needs, but I tried SymPy and I have to say that it's pretty neat. I used it to implement the Gram-Schmidt Process to obtain a set of orthonormal basis functions from a set of non-orthogonal, non-normalized and linearly independent polynomial functions and SymPy has done the job beautifully. If you already know a little Python, the learning curve is quite gentle.

from __future__ import division
from sympy import *
from prettytable import PrettyTable
import functools
import string

x=symbols('x')
init_printing(use_unicode=True)

def print_basis(basis):
    """
    pretty print basis functions
    """
    table = PrettyTable(["e%d" %i for i in range(len(basis))])
    table.padding_width = 1
    table.add_row(map(pretty, basis))
    print table

def print_projection_coefficients(coefficients, significant_digits):
    """
    pretty print projection coefficients
    """
    table = PrettyTable(string.ascii_uppercase[:len(coefficients)])
    table.padding_width = 1
    table.add_row(map(pretty, coefficients))
    table.add_row(map(functools.partial(N, n=significant_digits), coefficients))
    print table

def normalize(f, inner_product):
    """
    normalize input function using input inner product operation
    """
    return f/sqrt(inner_product(f, f))

def project(f, basis):
    """
    project input function in space generated by input basis
    """
    projection = 0
    coefficients = []
    for b in basis:
        coefficient = inner_product(f, b)
        coefficients.append(coefficient)
        projection = projection + coefficient * b
    return projection, coefficients

def gram_schmidt(v_basis, inner_product):
    """
    orthonormalize input basis vectors using gram-schmidt algorithm
    """
    e = normalize(v_basis[0], inner_product)
    e_basis = [e]
    for v in v_basis[1:]:
        f = v
        for e in e_basis:
            f = f - inner_product(e, v) * e
        e_basis.append(normalize(f, inner_product))
    return e_basis

# define non orthonormal basis functions here
v_basis = [1, x, x**2, x**3, x**4]

# define inner_product here
inner_product=lambda f1, f2: integrate(f1*f2,(x, -pi, pi))

# orthonormalize basis using gram-schmidt algorithm
e_basis = gram_schmidt(v_basis, inner_product)
print("orthonormalized basis functions:")
print_basis(e_basis)

# get projection of cos(x) in orthonormal basis
significant_digits=6
approximation, coefficients = project(cos(x), e_basis)
print("projection coefficients of cos(x):")
print_projection_coefficients(coefficients, significant_digits)
print("projection of cos(x):")
pprint(approximation.evalf(significant_digits))

SymPy formats the output nicely in unicode. I used PrettyTable to print the output nicely in table form. I'm very pleased with the result. Together with NumPy and Matplotlib, it makes the perfect combo for scientific programming.

orthonormalized basis functions:
+---------+---------+--------------------+------------------------+--------------------------------------+
|    e0   |    e1   |         e2         |           e3           |                  e4                  |
+---------+---------+--------------------+------------------------+--------------------------------------+
|    ___  |   ___   |          ⎛      2⎞ |          ⎛        2  ⎞ |           ⎛          ⎛      2⎞     ⎞ |
|  ╲╱ 2   | ╲╱ 6 ⋅x |     ____ ⎜ 2   π ⎟ |     ____ ⎜ 3   3⋅π ⋅x⎟ |           ⎜        2 ⎜ 2   π ⎟     ⎟ |
| ─────── | ─────── | 3⋅╲╱ 10 ⋅⎜x  - ──⎟ | 5⋅╲╱ 14 ⋅⎜x  - ──────⎟ |           ⎜     6⋅π ⋅⎜x  - ──⎟    4⎟ |
|     ___ |     3/2 |          ⎝     3 ⎠ |          ⎝       5   ⎠ |       ___ ⎜ 4        ⎝     3 ⎠   π ⎟ |
| 2⋅╲╱ π  |  2⋅π    | ────────────────── | ────────────────────── | 105⋅╲╱ 2 ⋅⎜x  - ────────────── - ──⎟ |
|         |         |          5/2       |            7/2         |           ⎝           7          5 ⎠ |
|         |         |       4⋅π          |         4⋅π            | ──────────────────────────────────── |
|         |         |                    |                        |                   9/2                |
|         |         |                    |                        |               16⋅π                   |
+---------+---------+--------------------+------------------------+--------------------------------------+
projection coefficients of cos(x):
+---+---+------------+---+----------------------------------------------------------+
| A | B |     C      | D |                            E                             |
+---+---+------------+---+----------------------------------------------------------+
| 0 | 0 |      ____  | 0 |             ⎛            3⎞             ⎛      3       ⎞ |
|   |   | -3⋅╲╱ 10   |   |         ___ ⎜        16⋅π ⎟         ___ ⎜  16⋅π        ⎟ |
|   |   | ────────── |   |   105⋅╲╱ 2 ⋅⎜-24⋅π + ─────⎟   105⋅╲╱ 2 ⋅⎜- ───── + 24⋅π⎟ |
|   |   |     3/2    |   |             ⎝          7  ⎠             ⎝    7         ⎠ |
|   |   |    π       |   | - ───────────────────────── + ────────────────────────── |
|   |   |            |   |                9/2                         9/2           |
|   |   |            |   |            16⋅π                        16⋅π              |
| 0 | 0 |  -1.70371  | 0 |                         0.486659                         |
+---+---+------------+---+----------------------------------------------------------+
projection of cos(x):
           4             2           
0.0261599⋅x  - 0.452289⋅x  + 0.978327