import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib import rc
import numpy as np

rc('animation', html='html5')
plt.rcParams['figure.figsize'] = 8, 6

def updateStateExplicit(state, const):
    nextState = np.zeros_like(state)
    for j, u in enumerate(state):
        if j==0:
            nextState[j] = u + const*(state[j+1]-2*u + 0)
        elif j == len(state)-1:
            nextState[j] = u + const*(0-2*u + state[j-1])
        else:
            nextState[j] = u + const*(state[j+1]-2*u + state[j-1])
    return nextState

initialState = np.zeros(501)
initialState[251] = 1.
D = 1.
dx = 1.
dts = [1, .5, .1]
states = [[] for _ in dts]
for i, dt in enumerate(dts):
    states[i].append(initialState)
    for s in range(300):
        states[i].append(updateStateExplicit(states[i][-1], D*dt/(dx**2)))


fig, ax = plt.subplots()
ax.set_ylim(-1, 1)
ax.set_xlim(0, 501)
line, = ax.plot([], [], lw=2)
ax.grid()

def run1(i):
    # update the data
    line.setData(np.arange(0,501), states[0][i])
    return line,
def run2(i):
    line.setData(np.arange(0,501), states[1][i])
    return line,
def run3(i):
    line.setData(np.arange(0,501), states[2][i])
    return line,

rc('animation', html='html5')
#ani1 = animation.FuncAnimation(fig, run1, list(range(299)), blit=False, interval=50,repeat=True)
#ani1.save('diffusion_anim1.gif', writer='imagemagick')
#ani2 = animation.FuncAnimation(fig, run2, list(range(299)), blit=False, interval=50,repeat=True)
#ani2.save('diffusion_anim2.gif', writer='imagemagick')
#ani3 = animation.FuncAnimation(fig, run3, list(range(299)), blit=False, interval=50,repeat=True)
#ani3.save('diffusion_anim3.gif', writer='imagemagick')

For dt = 1.0

For dt = 0.5

For dt = 0.1

Compressed the Gifs a bit too much = weird color scheme (REDO)

Stopped here due to remote working madness - will come back to it since this is important to some projects of mine