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, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (13, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (22, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (23, 'numpy.meshgrid', 'np.meshgrid', True, 'import numpy as np\n'), (25, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (35, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (17, 'numpy.copy', 'np.copy', True, 'import numpy as np\n')]