from scitools.std import *

def nonlinearSystem(u0,v0,T,dt):
    """Numerical solution with implicit scheme of 
       u'(t)=-v(t)^3, u(0)=u0,
       v'(t)=u(t)^3, v(0)=v0,
       for 0<t<T and with time step dt. 
    """ 

    times=arange(0,T+dt,dt)
    N=len(times)
    
    # The first row of uv contains the approximation of u(t).
    # The second row of uv contains the approximation of v(t).
    uv=matrix(zeros([2,N]))    
    uv[:,0]=matrix([[u0],[v0]])
    for n in range(N-1): 
        uv[:,n+1]=solveNonLinSystem(uv[0,n],uv[1,n],uv[0,n], uv[1,n],dt,0.000001) 
    plot(times,array(uv[0,:])) # plot approximation of u
    hold('on')
    plot(times,array(uv[1,:])) # plot approximation of v
    xlabel('t')
    legend('u','v') 
    hold('off')


def solveNonLinSystem(x0,y0,a,b,dt,eps):
    """Numerical solution with Newton's method of 
    x+dt*y^3-a=0,
    y-dt*x^3-b=0,
    using (x0,y0) as intial guess. 
    """

    f=x0+dt*pow(y0,3)-a
    g=y0-dt*pow(x0,3)-b
    
    xy=matrix([[x0],[y0]])

    while fabs(f)+fabs(g)>eps:
       A=matrix([[1, 3*dt*pow(xy[1,0],2)],[-3*dt*pow(xy[0,0],2), 1]])
       xy=xy-A.I*matrix([[f],[g]])
       f=xy[0,0]+dt*pow(xy[1,0],3)-a
       g=xy[1,0]-dt*pow(xy[0,0],3)-b
 
    return xy