# -*- coding: utf-8 -*-
"""
Created on Wed Nov 26 21:36:52 2014

@author: susanhen
"""
from sys import exit
from numpy import *

def varmeledning (ustart, u0, u1, N, r):
    '''l?s varmeledningslikninga med startverdi ustart
    og med randkrav u0 til venstre og u1 til h?yre,
    utf?r N tidsskritt og la r = Delta t / (Delta x)**2 v?re gitt'''
    
    # sjekk f?rst om det ble brukt vektor eller matrise
    if ustart.shape == ustart.size:
        M, MM = ustart.shape
        if ((MM > 1) or (M < 3)):
            exit('Temperaturvektoren har ikke riktig st?rrelse, m? v?re Mx1 hvor M>2')
        u = ustart[:,0]    
    else:
        M = ustart.size
        if M < 3 :
            exit('Temperaturvektoren har ikke riktig st?rrelse, m? v?re Mx1 hvor M>2')
        u = ustart
    u[0]=u0
    u[M-1]=u1
    for n in range(1,N):        
        u = r*roll(u,1) + (1 - 2*r)*u + r*roll(u,-1)
        u[0] = u0
        u[M-1] = u1    
    
    return u