I have coded a 2D heat equation animation with heat sources inside the plane. Now I am having trouble implementing that in a 3D array. I have tried with Axes3d library in Matplotlib but getting confused how to properly functionalize it and even in the animation output. Could someone help me pointing out the tweaks in the following code please?

我已经编写了一个平面内有热源的2D热方程动画。现在我在3D数组中实现这一点时遇到了困难。我试过在Matplotlib中使用Axes3d库，但我搞不清如何正确地将其函数化，甚至在动画输出中也是如此。有人能帮我指出以下代码中的调整吗？

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

print("2D heat equation solver")

plate_length = 50
max_iter_time = 1500

alpha = 2
delta_x = 1

delta_t = (delta_x ** 2)/(4 * alpha)
gamma = (alpha * delta_t) / (delta_x ** 2)

# Initialize solution: the grid of u(k, i, j)
u = np.empty((max_iter_time, plate_length, plate_length))

# Initial condition everywhere inside the grid
u_initial = 0.0
u.fill(u_initial)

# Heat source 1
u_top_size = (6, 6) # (height, width)

# Heat source 2
u_right_size = (9, 9) # (height, width)

# Heat source 3
u_bottom_size = (8, 8) # (height, width)
u_new_size = (7, 7)
# Calculate the constant for each heat source based on the desired proportionality
total_area =u_top_size[0]*u_top_size[1]+ u_right_size[0]*u_right_size[1] + u_bottom_size[0]*u_bottom_size[1]+u_new_size[0]*u_new_size[1]
total_heat = 280.0 # sum of desired temperature of all sources
heat_constant = total_heat / total_area

# Set the boundary conditions
u[:, 36:41, 10:15] = heat_constant*u_top_size[0]*u_top_size[1]
u[:, 34:43, 34:43] = heat_constant*u_right_size[0]*u_right_size[1]
u[:, 9:16, 10:17] = heat_constant*u_bottom_size[0]*u_bottom_size[1]

# Heat source 4
 # (height, width)
u[:, 9:15, 37:43] = heat_constant*u_new_size[0]*u_new_size[1]
def calculate(u):
    for k in range(0, max_iter_time-1, 1):
        # Calculate u at time k+1
        for i in range(1, plate_length-1, 1):
            for j in range(1, plate_length-1, 1):
                # check if cell is a heat source
                if (36<=i<=40 and 10<=j<=14) or (34<=i<=42 and 34<=j<=42) or (9<=i<=15 and 10<=j<=16) or (9<=i<=14 and 37<=j<=42):
                    # check if a neighboring cell has a higher temperature
                    if (u[k][i+1][j] > u[k][i][j]) or (u[k][i-1][j] > u[k][i][j]) or (u[k][i][j+1] > u[k][i][j]) or (u[k][i][j-1] > u[k][i][j]):
                        # update temperature of heat source to the maximum temperature of neighboring cells
                        u[k + 1, i, j] = max(u[k][i+1][j], u[k][i-1][j], u[k][i][j+1], u[k][i][j-1])
                    else:
                        u[k + 1, i, j] = u[k][i][j]
                # for all other cells
                else:
                    u[k + 1, i, j] = gamma * (u[k][i+1][j] + u[k][i-1][j] + u[k][i][j+1] + u[k][i][j-1] - 4*u[k][i][j]) + u[k][i][j]
        # Add Neumann boundary condition at the top and bottom borders
        for i in range(plate_length):
            u[k + 1, 0, i] = u[k + 1, 1, i]
            u[k + 1, plate_length-1, i] = u[k + 1, plate_length-2, i]

        # Add Neumann boundary condition at the left and right borders
        for i in range(plate_length):
            u[k + 1, i, 0] = u[k + 1, i, 1]
            u[k + 1, i, plate_length-1] = u[k + 1, i, plate_length-2]

        # Check if every point in the plane has reached the desired temperature
        if np.all(u[k+1] >= heat_constant*u_right_size[0]*u_right_size[1]):
            break

    return u




def plotheatmap(u_k, k):
    # Clear the current plot figure
    plt.clf()

    plt.title(f"Temperature at t = {k*delta_t:.3f} unit time")
    plt.xlabel("x")
    plt.ylabel("y")

    # This is to plot u_k (u at time-step k)
    plt.pcolormesh(u_k, cmap=plt.cm.jet, vmin=0, vmax=100)
    plt.colorbar()

    return plt

# Do the calculation here
u = calculate(u)

def animate(k):
    plotheatmap(u[k], k)
    return plt


anim = animation.FuncAnimation(plt.figure(), animate, interval=1, frames=max_iter_time, repeat=False)
anim
anim.save("heat_equation1.gif")

I tried the Axes3d in Matplotlib. I am looking for the output same as the 2D one but in 3D.

我在Matplotlib中尝试了Axes3d。我正在寻找的输出与2D一样，但在3D。