Poisson equation solution using TDMA

gistfile1.py

import math
import matplotlib.pyplot as plt

def tdma(a, b, c, d):
    '''
    a[i] * x[i-1] + b[i] * x[i] + c[i] * x[i+1] = d[i]
    '''
    length = len(a)
    for i in xrange(1, length):
        tmp = a[i] / b[i-1]
        b[i] = b[i] - tmp * c[i-1]
        d[i] = d[i] - tmp * d[i-1]

    x = a
    x[-1] = d[-1] / b[-1]

    for i in xrange(length - 2, -1, -1):
        x[i] = (d[i] - c[i] * x[i+1]) / b[i]

    return x

def poisson_equation(h, x0, xn, a0, b0, c0, an, bn, cn, function):
    '''
    Problem definition:
    Laplace(u) = f(x), x0 <= x <= xn
    a0 * du/dx + b0 * u(x) = c0, x = x0
    an * du/dx + db * u(x) = cn, x = xn

    Linear equations system:
    (h * b0 - a0) * u0 + a0 * u1 = c0 * h;
    u[i-1] - 2u[i] + u[i+1] = h^2 * f[i], i = 1..n-1;
    -a0 * un + (an + h * bn) * u[n+1] = cn * h;
    '''
    n = int((xn - x0) / h)
    a, b, c, d = [], [], [], []
    for i in xrange(n + 1):
        a.append(1)
        b.append(-2)
        c.append(1)
        d.append(h * h * function(x0 + i * h))

    a[n] = -an
    b[0] = h * b0 - a0
    b[n] = an + h * bn
    c[0] = a0
    d[0] = c0 * h
    d[n] = cn * h

    u = tdma(a, b, c, d)
    return u

def test():
    def example_function(x):
        return math.sin(x)

    def exact_solution(i, h, x0):
        x = x0 + i * h
        return -math.sin(x) - 0.0116 * x + 1.7082

    def plot(u, exact_u):
        plt.plot(range(len(u)), u)
        plt.plot(range(len(exact_u)), exact_u)
        plt.show()

    h = 0.1 # step
    x0 = 0.0 # left border
    xn = 20.0 # right border
    a0, b0, c0, an, db, cn = 0.7, 1.0, 1.0, 0.3, 2.0, 1.0 # coefficients of boundary values
    u = poisson_equation(h, x0, xn, a0, b0, c0, an, db, cn, example_function)
    exact_u = map(lambda i: exact_solution(i, h, x0), range(len(u)))
    plot(u, exact_u)

test()