%pylab inline

from scipy import sparse  # to define sparse matrices
from scipy.sparse.linalg import spsolve   # to solve sparse systems

def solve_BVP_direct(f,x,u_left,u_right,k):
    """
    On input:
    f should be the function defining the ODE  u''(x) = f(x),
    x should be an array of equally spaced points where the 
        solution is to be computed,
    u_left should be the boundary condition at the left edge x[0],
    u_right should be the boundary condition at the right edge x[-1],
    
    Returns: 
    u, vector of approximate values to solution at the points x.
    """
    
    n = len(x) - 2  # number of interior points
    print "Solving tridiagonal system with n =%4i" % n
    dx = x[1]-x[0]
    
    # Form the matrix:
    d1 = ones(n)
    d0 = (k**2 * dx**2 - 2) * ones(n) 
    A = sparse.spdiags([d1,d0,d1], [-1,0,1],n,n,format='csc')
    
    # Form the right-hand side
    fx = f(x)
    rhs = -dx**2 * fx[1:-1] 
    rhs[0] = rhs[0] - u_left  
    rhs[-1] = rhs[-1] - u_right
    
    # Solve the system for the interior points:
    u_int = spsolve(A,rhs)   # vector of length n-2
    
    # Append the known boundary values:
    u = hstack((u_left, u_int, u_right)) 
    return u

def check_solution(x,u,u_true):
    # u should be a vector of approximate solution values at the points x
    # u_true should be a function defining the true solution.
    plot(x,u,'ro')
    plot(x,u_true(x),'b')
    error = u - u_true(x)
    error_max = abs(abs(error)).max()
    print "Maximum error is %10.3e" % error_max

k = 15.
u_left = 2.
u_right = -1.
c1 = u_left
c2 = (u_right - c1*cos(k)) / sin(k)
print "c1 = %g, c2 = %g" % (c1,c2)
f = lambda x: 0. * x
u_true = lambda x: c1 * cos(k*x) + c2 * sin(k*x)


n = 50  # number of interior points
x = linspace(0, 1, n+2)
u = solve_BVP_direct(f,x,u_left,u_right,k)

check_solution(x,u,u_true)

nvals = 10*array([10,20,40,80,160])
errors = zeros(nvals.shape)
for j,n in enumerate(nvals):
    x = linspace(0, 1, n+2)
    u = solve_BVP_direct(f,x,u_left,u_right,k)
    errors[j] = abs(u - u_true(x)).max()   # maximum abs error over interval

print "\n   n         error"
for j,n in enumerate(nvals):
    print "%4i  %15.9f" % (n,errors[j])
loglog(nvals,errors,'o-')

def solve_BVP_split(f,x,u_left,u_right,k):
    n2 = len(x)
    nmid = int(floor(n2/2.))
    xhalf1 = x[:nmid+1]
    xhalf2 = x[nmid:]
    u_mid = u_true(x[nmid])  # Assumes we know true solution!!
    print "u_mid = ",u_mid
    uhalf1 = solve_BVP_direct(f,xhalf1,u_left,u_mid,k)
    uhalf2 = solve_BVP_direct(f,xhalf2,u_mid,u_right,k)
    u = hstack((uhalf1, uhalf2[1:]))
    return u

x = linspace(0, 1, 51)
u = solve_BVP_split(f,x,u_left,u_right,k)
check_solution(x,u,u_true)

def solve_BVP_split_mismatch(f,x,u_left,u_right,k):
    n2 = len(x)
    dx = x[1]-x[0]
    nmid = int(floor(n2/2.))
    xhalf1 = x[:nmid+1]
    xhalf2 = x[nmid:]
    x_mid = x[nmid]
    u_mid_guesses = linspace(0,2,5)
    mismatch = zeros(u_mid_guesses.shape)
    for j,u_mid in enumerate(u_mid_guesses):
        uhalf1 = solve_BVP_direct(f,xhalf1,u_left,u_mid,k)
        uhalf2 = solve_BVP_direct(f,xhalf2,u_mid,u_right,k)
        u = hstack((uhalf1, uhalf2[1:]))
        plot(x,u)
        mismatch[j] = (uhalf1[-2] + (k**2*dx**2 - 2.)*u_mid + uhalf2[1]) + dx**2 * f(x_mid)
        print "With u_mid = %g, the mismatch is %g" % (u_mid, mismatch[j])
    figure()
    plot(u_mid_guesses, mismatch,'o-')
    
x = linspace(0, 1, 40)
solve_BVP_split_mismatch(f,x,u_left,u_right,k)


def solve_BVP_split(f,x,u_left,u_right,k):
    n2 = len(x)
    dx = x[1]-x[0]
    nmid = int(floor(n2/2.))
    xhalf1 = x[:nmid+1]
    xhalf2 = x[nmid:]
    x_mid = x[nmid]
    
    # solve the sub-problems twice with different values of u_mid
    # Note that any two distinct values can be used.
    
    u_mid = 0.
    uhalf1 = solve_BVP_direct(f,xhalf1,u_left,u_mid,k)
    uhalf2 = solve_BVP_direct(f,xhalf2,u_mid,u_right,k)
    G0 = (uhalf1[-2] + (k**2*dx**2 - 2.)*u_mid + uhalf2[1]) + dx**2 * f(x_mid)
    v0 = hstack((uhalf1, uhalf2[1:]))
    
    u_mid = 1.
    uhalf1 = solve_BVP_direct(f,xhalf1,u_left,u_mid,k)
    uhalf2 = solve_BVP_direct(f,xhalf2,u_mid,u_right,k)
    G1 = (uhalf1[-2] + (k**2*dx**2 - 2.)*u_mid + uhalf2[1]) + dx**2 * f(x_mid)
    v1 = hstack((uhalf1, uhalf2[1:]))

    z =  G1 / (G1 - G0)
    u = z*v0 + (1-z)*v1
    
    return u

x = linspace(0, 1, 80)
u = solve_BVP_split(f,x,u_left,u_right,k)
print x.shape, u.shape
check_solution(x,u,u_true)

def solve_BVP_split_recursive(f,x,u_left,u_right,k):
    n2 = len(x)
    print "Entering solve_BVP_split_recursive with n2 =%4i, for x from %5.3f to %5.3f" \
            % (n2, x[0], x[-1])
    if n2 < 20:
        # Stop recursing if the problem is small enough:
        u = solve_BVP_direct(f,x,u_left,u_right,k)
        return u
    else:
        # recursive
        nmid = int(floor(n2/2.))
        x_mid = x[nmid]
        xhalf1 = x[:nmid+1]
        xhalf2 = x[nmid:]
        dx = x[1]-x[0]
        u_mid = 0.
        uhalf1 = solve_BVP_split_recursive(f,xhalf1,u_left,u_mid,k)
        uhalf2 = solve_BVP_split_recursive(f,xhalf2,u_mid,u_right,k)
        G0 = (uhalf1[-2] + (k**2*dx**2 - 2.)*u_mid + uhalf2[1]) + dx**2 * f(x_mid)
        v0 = hstack((uhalf1, uhalf2[1:]))

        u_mid = 1.
        uhalf1 = solve_BVP_split_recursive(f,xhalf1,u_left,u_mid,k)
        uhalf2 = solve_BVP_split_recursive(f,xhalf2,u_mid,u_right,k)
        G1 = (uhalf1[-2] + (k**2*dx**2 - 2.)*u_mid + uhalf2[1]) + dx**2 * f(x_mid)
        v1 = hstack((uhalf1, uhalf2[1:]))

        z =  G1 / (G1 - G0)
        u = z*v0 + (1-z)*v1
    
    return u

x = linspace(0, 1, 40)
u = solve_BVP_split_recursive(f,x,u_left,u_right,k)
check_solution(x,u,u_true)