# forward Euler method

from numpy import *

def forward_euler(f, y0, n, h):
    """approximate the solution of y'(x) = f(x, y), y(0) = y0

    Uses the forward (or explicit) Euler method, with n steps and
    stepsize h.
    """
    y = [ y0 ]
    x = [ 0.0 ]
    for i in range(n):
        y.append(y[-1] + h * f(x[-1], y[-1]))
        x.append(i*h)
    return x, y

def forward_euler_alt(f, y0, n, h):
    x = [ i*h for i in range(n+1) ]
    y = [ 0 for i in range(n+1) ]
    y[0] = y0
    for i in range(1, n+1):
        y[i] = y[i-1] + h * f(x[i], y[i-1])
    return x, y

def forward_euler_array(f, y0, n, h):
    x = h * arange(n+1)
    y = zeros((n+1,) + shape(y0))
    y[0] = y0
    for i in range(1, n+1):
        y[i] = y[i-1] + h * f(x[i-1], y[i-1])
    return x, y

#forward_euler = forward_euler_array

def exp_example():
    """solve y' = mu * y, y(0) = 1"""
    y0 = 1.0
    mu = -1.0
    def f_exp(x, y):
        return mu * y
    maxx = 2.0
    n = 100
    h = maxx/n
    x, y = forward_euler(f_exp, y0, n, h)
    
    #x = arange(n+1) * h

    return x, y

def sincos_example():
    """solve y'(x) = A y, y(0) = [1 0]^T

    A = [ 0 -1 ]
        [ 1  0 ]
    """
    y0 = array([1.0, 0.0])
    A = array([[0.0, -1.0], 
               [1.0, 0.0]])
    def f_sincos(x, y):
        return dot(A, y)
    maxx = 3*pi
    n = 200
    h = maxx/n
    x, y = forward_euler(f_sincos, y0, n, h)
    x = array(x)
    y = array(y)
    
    #x = arange(n+1) * h

    return x, y

if __name__ == "__main__":
    x, y = exp_example()
    print x
    print y
    x, y = sincos_example()
    print x
    print y