from scitools.std import *

def explicit(T,n,m):
    """Numerical solution of  
       u_t=(ku_x)_x for 0<x<1 and 0<t<T 
       u(0,t)=g0(t) for 0<t
       k(1,t)*u_x(1,t)=g1(t) 0<t
       u(0,x)=f(x).
       We use an explicit scheme with n spatial intervals and m time steps. 
    """ 
    dx=1.0/n
    dt=float(T)/m

    x=arange(0,1+dx,dx)
    t=arange(0,T+dt,dt)

    u=array(zeros(n+1))    
    uOld=array(zeros(n+1)) 

    # Initial condition
    for i in range(n+1):
       uOld[i]=f(x[i])

    # Compute numerical approximations
    alpha=dt/(dx*dx)
    print "alpha=", alpha
    for l in range(m):
      for i in range(1,n):
        kl=k(x[i]-dx/2.0,t[l])
        kr=k(x[i]+dx/2.0,t[l])
        u[i]=alpha*(kr*uOld[i+1]-(kr+kl)*uOld[i]+kl*uOld[i-1])+uOld[i]
      # Dirichlet condition
      u[0]=g0(t[l])
      # Neumann condition
      u[n]=dx/k(x[n],t[l])*g1(t[l])+u[n-1]
      # New values -> old values
      uOld=u.copy()
  
    # Plot the solution at time T
    plot(x,u,"r-")   
    xlabel('x')
    ylabel('u')
    legend('Explicit')

def f(x):
    return cos(2.0*pi*x)

def k(x,t): 
    return 1.0

def g0(t): 
    return 1.0

def g1(t):
    return 0.0    


def implicit(T,n,m):
    """Numerical solution of  
       u_t=u_xx for 0<x<1 and 0<t<T 
       u(0,t)=g0(t) for 0<t
       k(1,t) u_x(1,t)=g1(t) 0<t
       u(0,x)=f(x).
       We use an implicit scheme.
    """ 
    dx=1.0/n
    dt=float(T)/m
    alpha=dt/(dx*dx)

    x=arange(0,1+dx,dx)
    t=arange(0,T+dt,dt)

    A=matrix(zeros([n+1,n+1]))    
    u=matrix(zeros([n+1,1]))    
    b=matrix(zeros([n+1,1]))

    # Make matrix A
    
    A[0,0]=1.0
    for i in range(1,n):
        A[i,i-1]=-alpha    
        A[i,i]=1.0+2.0*alpha
        A[i,i+1]=-alpha
    A[n,n-1]=-1        
    A[n,n]=1        

    # Initial condition
    for i in range(n+1):
       u[i]=f(x[i])

    # Compute numerical approximations
    print "alpha=", alpha
    invA=A.I # inverse of A
    for l in range(m):
      b[0]=g0(t[l+1])
      for i in range(1,n):
         b[i]=u[i]
      b[n]=dx*g1(t[l+1])
      u=invA*b
  
    # Plot the solution at time T
    plot(x,array(u),"b-")   
    xlabel('x')
    ylabel('u')
    legend('Implicit')
    #axis('equal')


def solve(T,n,m):
    """Numerical solution of  
       u_t=u_xx for 0<x<1 and 0<t<T 
       u(0,t)=g0(t) for 0<t
       k(1,t) u_x(1,t)=g1(t) 0<t
       u(0,x)=f(x).
       We use both implicit and explicit schemes. 
       Make sure that the conductivity k used in the explicit scheme is equal 
       to 1.0 for all x and t. 
    """ 
    explicit(T,n,m)
    hold('on')
    implicit(T,n,m)
    hold('off')