# -*- coding: utf-8 -*-
"""
Created on Mon Feb  8 08:06:19 2016

@author: nemenman
"""
import numpy as np

def Euler(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1):
    """
    Solves 1-d ODE using the Euler method.
            
    Euler(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1):
    """
    t = np.arange(0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        x[i] = x[i-1] + dt * xPrime(x[i-1])
    
    return (t,x)        
    
def RK2(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1):
    """
    Solves 1-d ODE using the Runge-Kutta 2 method.
            
    RK2(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1):
    """    
    
    t = np.arange(0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        fx = xPrime(x[i-1])
        guess = x[i-1] + dt * fx
        fxdx = xPrime(guess)
        x[i] = x[i-1]+0.5*(fx+fxdx)*dt
  
    return (t,x)     
    
    
def EulerArg(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1,arg=()):
    """
    Solves 1-d ODE using the Euler method.
            
    EulerArg(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1),arg=():
    """
    t = np.arange(0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        x[i] = x[i-1] + dt * xPrime(x[i-1],t[i-1],arg)
    
    return (t,x)        
        
    
def RK2Arg(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1,arg=()):
    """
    Solves 1-d ODE using the Runge-Kutta 2 method.
            
    RK2Arg(xPrime,t0=0.0,x0=0.0,T=1.0,dt=0.1,arg=()):
    """    
    
    t = np.arange(0, T+dt, dt)
    x = np.zeros(t.size)
    x[0] = x0
    for i in range(1, t.size):
        fx = xPrime(x[i-1],t[i-1],arg)
        guess = x[i-1] + dt * fx
        fxdx = xPrime(guess,t[i-1],arg)
        x[i] = x[i-1]+0.5*(fx+fxdx)*dt
  
    return (t,x)