Malha quadrada, pulso senoidal em repouso no centro. Contornos fixos.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D


#pulso em repouso no centro do grid, contorno fixo
#parâmetros do pulso:
kx = 3          #número de picos no eixo x
ky = 3          #número de picos no eixo y
L = 1           #largura do grid (LxL)

#parâmetros do programa:
dt, tf = 1e-2, 1e+1
dx, xf = 1e-2, L
dy, yf = dx, xf
Nt, Nj = int(tf/dt), int(xf/dx)
x = np.linspace(0, L, Nj)
y = np.linspace(0, L, Nj)
v = 0.4**0.5*dx/dt          #velocidade no meio
a = v*v*dt*dt/(dx*dx)

#condições iniciais:
f = np.zeros((Nt, Nj, Nj))
psi = np.zeros((Nj, Nj))

for j in range(int((Nj-1)*1/3), int((Nj-1)*2/3)):
    for i in range(int((Nj-1)*1/3), int((Nj-1)*2/3)):
        f[0, j, i] = (np.sin(kx*np.pi*i*dx*Nj/(Nj-1))*
                      np.sin(ky*np.pi*j*dy*Nj/(Nj-1)))**2

for j in range(1, Nj-1):
    for i in range(1, Nj-1):
        f[1, j, i] = ((1-2*a)*f[0, j, i] + dt*psi[j, i] +
                      0.5*a*(f[0, j-1, i] + f[0, j+1, i] +
                             f[0, j, i-1] + f[0, j, i+1]))


#laço:
for n in range(2, Nt):
    for j in range(1, Nj-1):
        for i in range(1, Nj-1):
            f[n, j, i] = (2 - 4*a)*f[n-1, j, i] - f[n-2, j, i] + a*(
                f[n-1, j-1, i] + f[n-1, j+1, i] + f[n-1, j, i-1] + f[n-1, j, i+1]
                )


#plot:
X, Y = np.meshgrid(x, y)
vmin, vmax = -0.2, 0.2

fig, ab = plt.subplots(subplot_kw={"projection": "3d"})
surf = ab.plot_surface(X, Y, f[0], cmap=cm.plasma, linewidth=45, vmin=vmin, vmax=vmax)
ab.set_zlim(-1.01, 1.01)
ab.set_xlabel(r'$x$')
ab.set_ylabel(r'$y$')
ab.set_zlabel(r'$f$')
cbar = fig.colorbar(surf, ax=ab, orientation='vertical', pad=0.2, shrink=0.7, ticks=np.linspace(vmin, vmax, 5))
title_text = f'$t = {0:.2f}$, $a = {a:.2f}$'
ab.set_title(title_text)

# Function to update the plot for each frame
def update(frame):
    ab.cla()                                        # Clear the previous frame
    surf = ab.plot_surface(X, Y, f[frame], cmap=cm.plasma, linewidth=45, vmin=vmin, vmax=vmax)
    ab.set_zlim(-1.01, 1.01)
    ab.set_xlabel(r'$x$')
    ab.set_ylabel(r'$y$')
    ab.set_zlabel(r'$f$')
    title_text = f'$t = {frame*dt:.2f}$, $a = {a:.2f}$'
    ab.set_title(title_text)

# Create the animation and assign it to a variable
animation = FuncAnimation(fig, update, frames=np.arange(0, int(0.5*Nt), 5), interval=100)

# Keep a reference to the animation object
# This line is crucial to prevent the animation from being deleted prematurely
anim_variable = animation

plt.show()

animation.save('onda2d.gif')