cimport numpy
import numpy

cimport cython

@cython.boundscheck(False)
def solve_heat_buf_nocheck(initial_conditions, double dx, double dt, int iter):
    cdef numpy.ndarray[double, ndim=2] cur = initial_conditions.copy()
    cdef numpy.ndarray[double, ndim=2] next = numpy.zeros_like(initial_conditions)
    cdef int count, i, j, M, N
    M, N = cur.shape[0], cur.shape[1]
    cdef double step
    for count in range(iter):
        for i in range(1, M-1):
            for j in range(1, N-1):
                step = cur[i-1,j] + cur[i+1,j] + cur[i,j-1] + cur[i,j+1] - 4*cur[i,j]
                next[i,j] = cur[i,j] + dt*step/dx**2
        cur, next = next, cur
    return cur
    



@cython.boundscheck(False)
def solve_heat_buf_noneg(initial_conditions, double dx, double dt, int iter):
    cdef numpy.ndarray[double, ndim=2] cur = initial_conditions.copy()
    cdef numpy.ndarray[double, ndim=2] next = numpy.zeros_like(initial_conditions)
    cdef unsigned int count, i, j, M, N
    M, N = cur.shape[0], cur.shape[1]
    cdef double step
    for count in range(iter):
        for i in range(1, M-1):
            for j in range(1, N-1):
                step = (  cur[<unsigned int>(i-1),j] + cur[<unsigned int>(i+1),j] 
                        + cur[i,<unsigned int>(j-1)] + cur[i,<unsigned int>(j+1)] 
                        - 4*cur[i,j] )
                next[i,j] = cur[i,j] + dt*step/dx**2
        cur, next = next, cur
    return cur
    
    
    
    
import numpy
N = 100


# the initial data        
L = numpy.ndarray(shape=(N, N))
for i in range(N):
    x = 2.0*i/N - 1.0
    a, b = .5-x**2 , .75 - 2*x**2
    if a < b:
        L[i, a*N:b*N] = .5
        

# solve it
arr = solve_heat_buf_noneg(L, 1, .001, 5000)

