Heat diffusion model using ndimage.generic_filter

[rougier]

Author: Nicolas Rougier (@rougier) Submitted by: Nicolas Rougier (@rougier)

Simple heat diffusion model using the ndimage.generic_filter.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
''' Simple model of heat diffusion '''
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import generic_filter

n  = 40
k  = 4
t  = 10.0
dt = 0.005

S = np.ones((n,n))

def diffuse(Z):
    return k*(Z[0] + Z[1] + Z[3] + Z[4] - 4*Z[2])

for i in xrange(int(t/dt)):
    # Set border temperature
    S[:,n-1] = S[n-1,:] = S[:,0] = 1
    S[0,:] = 0

    # Diffusion equation
    S += dt*4*generic_filter(S, diffuse, footprint= [[0, 1, 0],
                                                     [1, 1, 1],
                                                     [0, 1, 0]])

fig = plt.figure(figsize=(10,7.5))
plt.imshow(S, cmap=plt.cm.hot, origin='lower', extent=[0,n-1,0,n-1],
           interpolation='bicubic', vmin=0, vmax=1)
plt.colorbar()
CS = plt.contour(S, 10, colors='k')
plt.clabel(CS, inline=1, fontsize=16)
plt.grid()
plt.show()

[jni]

@rougier these are all awesome, thank you! And with this many options, we can do some as examples and use others as exercises, which is brilliant.