from pylab import *
from copy import deepcopy

def euler_step(t, h, f, y):
    """Evaluates one step of the Euler method.  Takes the current value of the
    independent variable (t), the step size (h), a system of first-order
    differential equations (f), and the current values of the equations (y).
    Returns new values for t and y."""
    y = deepcopy(y) # to avoid changing the inputs
    for i in xrange(len(f)): # for each equation in the system
        # find the slope at the beginning of the interval
        k = f[i](t, y)
        # estimate the next point with the slope
        y[i] = y[i] + k * h
    return t + h, y

def RK4_step(t, h, f, y):
    """Evaluates one step of the fourth-order Runge-Kutta method.  Same parameters
    as euler_step, above."""
    y = deepcopy(y) # to avoid changing the inputs
    
    o = len(f)
    for i in xrange(o):
        k1 = []; k2 = []; k3 = []; k4 = []
    
    # find k terms
    for i in xrange(o):
        k1.append(h * f[i](t, y))
    for i in xrange(o):
        k2.append(h * f[i](t + h/2.0, y + array(k1)/2.0))
    for i in xrange(o):
        k3.append(h * f[i](t + h/2.0, y + array(k2)/2.0))
    for i in xrange(o):
        k4.append(h * f[i](t + h, y + array(k3)))

    # compute new values of y
    for i in xrange(o):
        y[i] = y[i] + (k1[i]/6.0 + k2[i]/3.0 + k3[i]/3.0 + k4[i]/6.0)

    return t+h, y
    
def solver(step_fn, f, f0, interval, h):
    """Generic numeric solver. step_fn is a function that evaluates one step of the
    method to be used, given the current value of the independent variable, the step
    size, a system of first order differential equations, and current values for
    each of the equations.  f is a system of first-order differential equations; f0
    are initial values for the functions in f, h is the step size, and interval is a
    tuple.  Returned values are t and y, where y are the values of each equation in
    the system at each step."""
    a, b = interval
    t = [a]
    y = map(lambda x: [x], f0)
    y_ = array(f0); t_ = a
    err_local = [0.0]

    for n in xrange(int((float(b) - float(a)) / h)):
        # evaluate as two steps for local error estimate
        t2, y2 = step_fn(t_, h/2.0, f, y_)
        t2, y2 = step_fn(t2, h/2.0, f, y2)
        
        # then actually take the step
        t_, y_ = step_fn(t_, h, f, y_)

        err_local.append(abs(y2[0] - y_[0]))

        for i in xrange(len(f)):
            y[i].append(y_[i])
        t.append(t_) 
    return t, y, err_local

def solver_adaptive(step_fn, f, f0, interval, err_target, err_max, h_initial, kp):
    a, b = interval
    t = [a]
    y = map(lambda x: [x], f0)
    y_ = array(f0); t_ = a
    err_local = [0.0]
    err_prev = 0.0
    err_int = 0.0
    h = h_initial
    stepsize = [h]
    steps_discarded = 0
    
    while t_ < b:
        while True:
            # evaluate as two steps for local error estimate
            t2, y2 = step_fn(t_, h/2.0, f, y_)
            t2, y2 = step_fn(t2, h/2.0, f, y2)
            
            # then actually take the step
            t1, y1 = step_fn(t_, h, f, y_)
            
            # compute the error, update the step size with a very simple control loop
            err = abs(y2[0] - y1[0])
            h += (err_target - err) * kp
            if err < err_max:
                break
            else:
                steps_discarded += 1

        t_ = t1
        y_ = y1

        stepsize.append(h)
        print "%f" % err

        err_local.append(err)

        for i in xrange(len(f)):
            y[i].append(y_[i])
        t.append(t_) 
    print "Steps discarded: %d" % steps_discarded
    return t, y, err_local, stepsize

if __name__ == "__main__":

    # Our system of differential equations
    f = (
        lambda t, y: y[1],         # y1' = y2
        lambda t, y: -1.0 * y[0]   # y2' = -y1
    )

    # Initial conditions
    f0 = (1.0, 0.0)

    # Step size
    h = 0.1

    # Interval
    interval = (0.0, 100.0 * pi)

    # Solve with Euler method and RK4:
    t, y, err_local_euler = solver(euler_step, f, f0, interval, h)
    t2, y2, err_local_rk4 = solver(RK4_step, f, f0, interval, h)
    t3, y3, err_local_adaptive, h_adaptive = solver_adaptive(RK4_step, f, f0, interval,
        0.00001, 0.00002, h, 750.0)

    #mean = lambda x: sum(x) / len(x)
    print "Average local error:"
    print " Euler: ",  mean(err_local_euler)
    print " RK4:   ", mean(err_local_rk4), max(err_local_rk4), len(err_local_rk4)
    print ( " RK4a:  ", mean(err_local_adaptive), max(err_local_adaptive),
        len(err_local_adaptive) )

    print "Average stepsize: ", mean(h_adaptive)

    # Evaluate the analytic solution for reference:
    xa = frange(interval[0], interval[1], 0.01)
    ya = cos(xa)
    
    # Plot!
    plot(xa, ya, t, y[0], t2, y2[0], t3, y3[0])
    legend(("Analytic", "Euler", "RK4", "Adaptive"))
    #plot(t3, h_adaptive)
    show()