import numpy as np

# Parameters
c = 1.0
epsilon = 5e-4
x_start, x_end = -5, 5
t_start, t_end = 0, 10
nx = 100  # number of spatial points
nt = 500  # number of time steps
dx = (x_end - x_start) / (nx - 1)
dt = (t_end - t_start) / nt

# Discretization
x = np.linspace(x_start, x_end, nx)
u = np.exp(-x**2)  # initial condition

# Time-stepping loop
for n in range(nt):
    u_old = u.copy()
    for i in range(1, nx-1):
        u[i] = (u_old[i] - c * dt / (2 * dx) * (u_old[i+1] - u_old[i-1]) +
                epsilon * dt / dx**2 * (u_old[i+1] - 2 * u_old[i] + u_old[i-1]))
    # Periodic boundary conditions
    u[0] = (u_old[0] - c * dt / (2 * dx) * (u_old[1] - u_old[-2]) +
            epsilon * dt / dx**2 * (u_old[1] - 2 * u_old[0] + u_old[-2]))
    u[-1] = u[0]

# Save the final solution
np.save('/home/weichao/Downloads/Code_Generation_Benchmark/PDE_Benchmark/results/prediction/gpt-4o/prompts/u_1D_Linear_Convection.npy', u)