import numpy as np
import matplotlib.pyplot as plt
L = 10
a = 0.4
type = "COS"
fig = plt.figure(figsize=(12, 10))
ax = plt.axes(projection='3d')
ax.grid(False)
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
x = np.arange(0, L, 0.05)
y = np.arange(0, 20, 0.05)
X, Y = np.meshgrid(x, y)
def f(x, t, n):
    return np.cos(n * np.pi * x / L) * np.exp(-a * n * n * np.pi * np.pi * a * t / (L * L))
def s(n):
    if type == "PAR":
        if n > 0:
            return L * L * (4 * np.pi * n + 4 * np.pi * n * ((-1) ** n)) / (2 * (np.pi ** 3) * (n ** 3))
        else:
            return L * L / 12
    elif type == "COS":
        if n > 0:
            return 2 * (L * np.sin(L) * ((-1) ** n)) / (L * L - np.pi * np.pi * n * n)
        else:
            return np.sin(L) / L
def g(x, t):
    ans = s(0)
    for n in range(1, 200):
        ans += s(n) * f(x, t, n)
    return ans
Z = g(X, Y)
surf = ax.plot_surface(X, Y, Z, cmap=plt.cm.inferno)
ax.set_xlabel('x', labelpad=20)
ax.set_ylabel('t', labelpad=20)
ax.set_zlabel('T', labelpad=20)
plt.show()

surf = ax.plot_surface(X, Y, Z, cmap=plt.cm.inferno, antialiased=False)