import numpy as np

import matplotlib.pyplot as plt

def heat_equation_solver(L, T, alpha, Nx, Nt):

"""
Solves the one-dimensional heat equation using the finite difference method.

Parameters:
L (float): Length of the rod.
T (float): Total simulation time.
alpha (float): Thermal diffusivity.
Nx (int): Number of spatial grid points.
Nt (int): Number of time steps.

Returns:
u (2D array): Temperature distribution over space and time.
x (1D array): Spatial coordinates.
t (1D array): Time values.
"""
dx = L / (Nx - 1)
dt = T / Nt

r = alpha * dt / dx**2

if r > 0.5:
print("Warning: Unstable solution. Consider using a smaller time step.")

x = np.linspace(0, L, Nx)
t = np.linspace(0, T, Nt)

u = np.zeros((Nx, Nt))
u[:, 0] = np.sin(np.pi * x) # Initial condition

for n in range(0, Nt - 1):

for i in range(1, Nx - 1):
u[i, n + 1] = u[i, n] + r * (u[i + 1, n] - 2 * u[i, n] + u[i - 1, n])

return u, x, t

# Parameters
L = 1.0 # Length of the rod
T = 0.1 # Total simulation time
alpha = 0.01 # Thermal diffusivity
Nx = 50 # Number of spatial grid points
Nt = 100 # Number of time steps

# Solve the heat equation

u, x, t = heat_equation_solver(L, T, alpha, Nx, Nt)

# Plot the results

plt.imshow(u, aspect='auto', extent=[0, T, 0, L], origin='lower', cmap='hot')
plt.colorbar(label='Temperature')
plt.title('Heat Equation Solution')
plt.xlabel('Time')
plt.ylabel('Position')
plt.show()