interactive heat simulation

heat_simulation2d.py

import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Ellipse

r=2
dt=0.1
dx=0.1
dy=0.1

L=50                    #długość
H=50                    #szerokość
maxT=200                #max czas
k=0.5                   #rozchodzenie ciepła

M=np.zeros([L,H])       #mapa ciepła
POM=np.zeros([L,H]) 
M[20:30,10:20]=2000     #grzejniki

T=np.arange(0,maxT,dt)
MM=[]
MM.append(M.copy())


fig = plt.figure()
pcm = plt.pcolormesh(MM[0])
plt.colorbar()

def computeCiepełko(e):
    POM=MM[e] if e==0 else MM[e-1]
    #for t in (nexMtx+number for number in range(len(T)-nexMtx)):
    for i in range(1,L-1):
        for j in range(1,H-1):
            POM[i,j] = (POM[i-1,j] + POM[i+1,j] + POM[i,j-1] + POM[i,j+1])/4
    POM[0:L-1,0]=0
    POM[0:L-1,H-1]=0
    POM[0,0:H-1]=0
    POM[L-1,0:H-1]=0
    if e>0:
        MM.append(POM.copy())
    

def onclick(event):
    print(int(event.xdata), int(event.ydata))
    xpos=int(event.xdata)
    ypos=int(event.ydata)
    r_xpos1= 0 if xpos-r < 0 else xpos-r
    r_xpos2= L if xpos+r > L else xpos+r
    r_ypos1= 0 if ypos-r < 0 else ypos-r
    r_ypos2= H if ypos+r > H else ypos+r

    MM[nextMtx-1][r_ypos1:r_ypos2,r_xpos1:r_xpos2]=2000

cid = fig.canvas.mpl_connect('button_press_event', onclick)

# Function called to update the graphic
def step(i):
    global nextMtx
    nextMtx=i+1
    if i >= len(T): return
    computeCiepełko(i)
    pcm.set_array(MM[i].ravel())
    plt.draw()

anim = animation.FuncAnimation(fig, step, interval=50)
plt.show()