Integrating a diffusion equation by convolution

diffusion_convolution.py

"""
In this script we integrate the diffusion equation by applying the stencil matrix through
the convolve function from scipy.ndimage. This function applies a convolution, which means 
that for each element in matrix phi it will multiply the weights in the intencil to the 
element itself and its 8 nearest neighbors, summing the result subsequently. 
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import convolve


dx = 1.0 # space interval between points in the space discretization
stencil = (1/(dx*dx)) * np.array([[0, +1, 0],
                                  [1, -4, 1],
                                  [0, +1, 0]])

# initial condition
condition = 'random'

if condition == 'dirac':
    phi = np.zeros((50, 50))
    phi[25, 25] = 10.
if condition == 'random':
    phi = np.random.rand(50, 50)

tsteps = 100 # total time iterations
dt = 0.01 # time step
diffusion_constant = 1
nshots = 4 # number of screeshots
print_nstep = tsteps/nshots 

fig = plt.figure()
ax = []
ishot = 1

for nstep in range(tsteps):
    
    if nstep % print_nstep == 0:
        print(nstep)
        
        ax.append(fig.add_subplot(1, nshots, ishot))
        im = ax[ishot-1].imshow(phi, interpolation='bicubic', cmap=plt.get_cmap('plasma'), origin='lower', vmin=0, vmax=1)
        ax[ishot-1].set_title(r'$t = '+str(nstep)+'$')
        
        
        ishot+=1
        
    # time integration
    phi = phi + dt * diffusion_constant * convolve(phi, stencil, mode='wrap')

for figure in ax:
    figure.set_yticklabels([])
fig.colorbar(im, ax=[ax[0], ax[1], ax[2], ax[3]], location='bottom')

plt.savefig('diffusion.png', dpi=100)
plt.show