Finite differences ode

finite-differences.py

# -*- coding: utf-8 -*-
"""
Created on Fri Dec  5 11:48:36 2014

@author: isra
"""
from __future__ import division
from scipy import sparse
from scipy.sparse.linalg import spsolve
import numpy as np
import matplotlib.pyplot as plt

N = 24
t = np.linspace(0, np.pi, N + 1)

u = t - 2 * t * np.cos(t)

v = np.empty(t.shape, dtype=np.float64)
v.fill(-1)

w = t

alpha, beta = 0, np.pi

h = np.pi / N

a = -(1 + h / 2 * w)
b = -h**2 * u
c = -(1 - h / 2 * w)
d = 2 + h**2 * v

#We can drop the first and last items as they aren't useful
a = a[1:-1]
b = b[1:-1]
c = c[1:-1]
d = d[1:-1]

Coefs = sparse.spdiags([a, d, c], [-1, 0, 1], N - 1, N - 1, 'csc')
rhs = b.copy()
rhs[0] -= a[0] * alpha
rhs[-1] -= c[-1] * beta

x = np.zeros(t.shape)
x[0] = alpha
x[-1] = beta
x[1:-1] = spsolve(Coefs, rhs)

plt.plot(t, x, 'r', t, t + 2 * np.sin(t), 'b')