A

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


#equação da onda 2d, anel nas extremidades x,
#extremidades y = 0 e y = Ly fixas,
#anteparo em y = alph*Ly, duas fendas.

#parâmetros
#região de contorno:
Ly, Lx = 7, 7       #comprimento x e y
fp, fc = 0.5, 0.2   #posição fendas em relação a Ly/2 e comprimento
alph = 0.7          #posição anteparo

#laços:
dt, tf = 1e-2, 1e+1
dx = 1e-2
dy = dx
Nt, Ny, Nx = int(tf/dt), int(Ly/dy), int(Lx/dx)
x = np.linspace(0, Lx, Nx)
y = np.linspace(0, Ly, Ny)
v = 0.4**0.5*dx/dt          #velocidade no meio
a = v*v*dt*dt/(dx*dx)

#onda inicial:
ky = 4              #número de picos após anteparo
A = 0.25            #amplitude inicial
K = np.pi*ky/(Ly*(1-alph))


#condições iniciais:
f = np.zeros((Nt, Ny, Nx))  #f[tempo, coordenada y, coordenada x]
psi = np.zeros((Ny, Nx))
for j in range(int(alph*Ny), Ny-1):
    for i in range(Nx):
        f[0, j, i] = A*np.sin(ky*np.pi*(j*dy-Ly)/(Ly*(1-alph)))
for j in range(1, Ny-1):
    for i in range(Nx):
        if i == 0:
            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, Nx-1]))
        elif i == Nx-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, 0]))
        else:
            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, Ny-1):
        for i in range(Nx):
            if j == int(alph*(Ny-1)): #anteparo
                if i in range(int(Nx*(0.5*Lx-fp-fc/2)/Lx), int(Nx*(0.5*Lx-fp+fc/2)/Lx)): #fendas
                    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]
                )
                elif i in range(int(Nx*(0.5*Lx+fp-fc/2)/Lx), int(Nx*(0.5*Lx+fp+fc/2)/Lx)):
                    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]
                )
                else:
                    f[n, j, i] = 0
            elif i == 0:
                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, Nx-1] + f[n-1, j, i+1])
            elif i == Nx-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, 0])
            else:
                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.01, 0.01

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(-0.5, 0.5)
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(-0.5, 0.5)
    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.7*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('duplafenda.gif')