from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
import os
import contorno
from constantes import INTERVALOS, PASSOS, TAMANHO_BARRA, DELTA_T, DELTA_X

z_temp = contorno.p_3
TAMANHO_BARRA = 2

x = np.linspace(0.0, TAMANHO_BARRA, INTERVALOS+1)
y = np.linspace(0.0, DELTA_T, PASSOS+1)

z = []
for k in range(PASSOS+1):
    z_k = np.copy(z_temp)
    z.append(z_k)
    for i in range(1, INTERVALOS):
        z_temp[i] = z_k[i] + (DELTA_T/(DELTA_X**2)) * (z_k[i+1]-2*z_k[i]+z_k[i-1])

z = np.asarray(z)
x, y = np.meshgrid(x, y)

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(x, y, z, cmap=cm.coolwarm, antialiased=False)

ax.set_xlabel('x')
ax.set_ylabel('t')
ax.set_zlabel('T(x,t)')

fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

[((12, 4, 12, 49), 'numpy.linspace', 'np.linspace', ({(12, 16, 12, 19): '0.0', (12, 21, 12, 34): 'TAMANHO_BARRA', (12, 36, 12, 48): 'INTERVALOS + 1'}, {}), '(0.0, TAMANHO_BARRA, INTERVALOS + 1)', True, 'import numpy as np\n'), ((13, 4, 13, 39), 'numpy.linspace', 'np.linspace', ({(13, 16, 13, 19): '0.0', (13, 21, 13, 28): 'DELTA_T', (13, 30, 13, 38): 'PASSOS + 1'}, {}), '(0.0, DELTA_T, PASSOS + 1)', True, 'import numpy as np\n'), ((22, 4, 22, 17), 'numpy.asarray', 'np.asarray', ({(22, 15, 22, 16): 'z'}, {}), '(z)', True, 'import numpy as np\n'), ((23, 7, 23, 24), 'numpy.meshgrid', 'np.meshgrid', ({(23, 19, 23, 20): 'x', (23, 22, 23, 23): 'y'}, {}), '(x, y)', True, 'import numpy as np\n'), ((25, 6, 25, 18), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((35, 0, 35, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((17, 10, 17, 25), 'numpy.copy', 'np.copy', ({(17, 18, 17, 24): 'z_temp'}, {}), '(z_temp)', True, 'import numpy as np\n')]