2_1

import numpy
from matplotlib import pyplot
from matplotlib import animation as animation


N = 100
x = numpy.linspace(-1, 1, N)
y = numpy.linspace(-1, 1, N)
t = 0
a = 1

level=[0,10,20,30,40,50,60,70,80,90,100,110,120]

dt = 0.00001
dx = x[1]-x[0]
dy = y[1]-y[0]

kx = dt * (a**2) / (dx**2)
ky = dt * (a**2) / (dy**2)

X, Y = numpy.meshgrid(x,y)

T = numpy.ones(X.shape)+ numpy.where(x > 0, numpy.exp(- ((X * X + Y * Y)**0.5 - 0.5)**2/ (2 * 0.03 * 0.03) ) * 1000, 0)

def f(x, y):
    return numpy.exp(- ((x + 0.4)**2 + y**2)/ (2*0.03*0.03) )*100000

u = T.copy()
un = u.copy()

pyplot.rcParams["animation.html"] = "jshtml"

fig, ax = pyplot.subplots(1,1,figsize=(8,7))

cont = pyplot.contourf(X,Y,u,levels=level)
cb = pyplot.colorbar()
pyplot.axis('equal')

def steps(i = 1):
    global u
    global un
    global t
    for _ in range (i):
        un[1:-1,1:-1] = u[1:-1,1:-1] + kx * ( u[1:-1,2:] - 2 * u[1:-1,1:-1] \
                                             + u[1:-1,:-2]) + dt * f(X[1:-1,1:-1], Y[1:-1,1:-1])\
                                        + ky * ( u[2:,1:-1] - 2 * u[1:-1,1:-1] + u[:-2,1:-1])
        un, u = u, un
        t += dt

def update(frame):
    global cont, u
    global t
    for c in cont.collections:
        c.remove()
    cont = pyplot.contourf(X,Y,u,levels=level)
    ax.set_title('time={:.4f}'.format(t))
    steps(100)
    return cont

anim = animation.FuncAnimation(fig, update, frames=150, init_func=init)

anim