heat_diffusion_2order.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/diamond-5point'
figure_path = 'figures/diamond-5point'

# Time
t_max = 500
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 - iC1)**2 + (j - jC2)**2 <= r2**2:
#            ini_heat[i, j] = 0.8

# Rerun
# ini_heat = np.load('data/diamond/heat_diffuse_1990.npy')

# 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)
    # Change the schemes
    heat = diamond_5point(ini_heat)
    if t % 10 == 0:
        np.save(os.path.join(data_path,f'heat_diffuse_{t:05d}'), 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"Time = {np.round(t*dt,2)}")

        plt.savefig(os.path.join(figure_path,f'heat_diffuse_{t:05d}'))
        plt.close()
    else:
        pass
    # Update the progress bar
    progress_bar.update(1)

# Close the progress bar
progress_bar.close()

print("Process complete!")