<reponame>MECH3750/2021-Tutorials
import numpy as np
import matplotlib.pyplot as plt

u0 = 1
L = 10
T = 10

# 10 / 2 -> 5
# 11 // 2 -> 5

# 1D initial conditions (space)

x = np.arange(0, L+1)  # 0 -- 1 -- ... -- (L-1) -- (L)

f = np.zeros(L+1)  # 0 -- 1 -- ... -- (L-1) -- (L)
f[:L//2] = 2*u0 / L * x[:L//2]
f[L//2:] = 2*u0 / L * (L - x[L//2:])

fig, ax = plt.subplots()
ax.plot(x, f)
plt.show()

# 2D solution (time + space)

X_array, T_array = np.meshgrid(
    np.arange(0, L+1), np.arange(0, T),
    indexing='ij'  # NOTE: Ensures that first index is `x`, second is `t`
)

f = np.zeros_like(X_array, dtype=float)
# NOTE: `zeros_like` with copy the DATATYPE as well as the shape
#       We need to force `f` to have `float`, so it can contain decimals
f[:L//2, 0] = 2*u0 / L * X_array[:L//2, 0]
f[L//2:, 0] = 2*u0 / L * (L - X_array[L//2:, 0])

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(X_array, T_array, f)
plt.show()

# Analytical solution

D = 1

u = u0/2 * np.ones_like(X_array, dtype=float)
for n in range(10):  # NOTE: The number of terms to add
    k = 2*np.pi * (2*n + 1)
    u -= 16*u0/k**2 * np.exp(-D*(k/L)**2*T_array) * np.cos(k*X_array/L)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.plot_surface(X_array, T_array, u)
plt.show()
