heat_diffusion_main.py

import os
import numpy as np
import matplotlib.pyplot as plt
import scienceplots
from tqdm import tqdm
from schemes import *

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

# File path
data_path = 'data/test'
figure_path = 'figures/test'

# Time
t_max = 100
steps = int(t_max / dt)

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

# Circle initialization
ini_heat = np.zeros((500,500))
r1 = 50
r2 = 100
iC1 = 120; jC1 = 120
iC2 = 300; jC2 = 300
for i in range(500):
    for j in range(500):
        if (i - iC1)**2 + (j - jC1)**2 <= r1**2:
            ini_heat[i,j] = 1
        if (i - iC2)**2 + (j - jC2)**2 <= r2**2:
            ini_heat[i, j] = 0.8

# Initial file
heat = np.copy(ini_heat)
np.save(os.path.join(data_path,'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 = diamond_5point(ini_heat)
    if t % 10 == 0:
        np.save(os.path.join(data_path,f'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(os.path.join(figure_path,f'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!")