heat_diffusion_main.py.backup

import numpy as np
import matplotlib.pyplot as plt
import scienceplots
from tqdm import tqdm

plt.style.use(['science'])

dx = 0.1
dt = 0.1
k = 0.05

t_max = 1000
steps = int(t_max / dt)

ini_heat = np.zeros((500, 500))
ini_heat[100:180, 300:380] = 1
ini_heat[250:400, 100:250] = 0.8

heat = np.copy(ini_heat)
np.save(f'data/square/heat_diffuse_0000', ini_heat)

# Create a progress bar
progress_bar = tqdm(total=steps, unit=' Steps', ncols=100, ascii=True)

for t in range(steps):
    ini_heat = np.copy(heat)
    heat = np.zeros((500, 500))
    for j in range(1, 499):
        for i in range(1, 499):
            heat[i,j] = ini_heat[i,j] + dt*k*(ini_heat[i+1,j+1] + ini_heat[i-1,j-1] + ini_heat[i-1,j+1] + ini_heat[i+1,j-1] - 4*ini_heat[i,j])/(2*dx**2)
    if t % 10 == 0:
        np.save(f'data/square/heat_diffuse_{t:04d}', heat)
        plt.figure(figsize=(8, 7))
        plt.imshow(ini_heat, origin='lower', cmap='jet', vmin=0, vmax=1)

        # Colorbar with label
        cbar = plt.colorbar(fraction=0.046, pad=0.04)
        cbar.set_label('Temperature')

        plt.minorticks_on()
        plt.tick_params(axis='both', which='both', direction='out', top=False, right=False)
        # Labels for x and y axes
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.title(f"Step = {t}")

        plt.savefig(f'figures/square/heat_diffuse_{t:04d}')
        plt.close()
    else:
        pass
    
    # Update the progress bar
    progress_bar.update(1)

# Close the progress bar
progress_bar.close()

print("Process complete!")