Benchmark.py

import os 
import sys
import numpy as np
import matplotlib.pyplot as plt

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from FFT_Solver_2D import Solver2D_FFT as Solver2D_FFT_torch
from FFT_Solver_2D_numpy import Solver2D_FFT as Solver2D_FFT_np
from Initial_conditions.TaylorGreen_IC import setup_Taylor_Green2D

# %% Setup domain

n_iter = 200 # Total iterations
dt = 1 # Timestep

std = 200000
n_iter_inflow = 1 # Forcing duration = 1 to just "pulse" the simulation at the start.

Nx = np.array([64,256,512,1024])
Ny = np.array([64,256,512,1024])

t_np = np.zeros((len(Nx)))
t_torch_cpu = np.zeros((len(Nx)))
t_torch_gpu = np.zeros((len(Nx)))

n_repeats = 10

for i in range(len(Nx)):
    
    t_np_tmp = 0
    t_torch_cpu_tmp = 0
    t_torch_gpu_tmp = 0
    
    for _ in range(n_repeats):
        X, Y = np.meshgrid(range(0,Nx[i]),range(0,Ny[i]), indexing='xy')
    
        texture_field = np.exp( -1/(std) * ((X - Nx[i]/2)**2 + (Y-Ny[i]/2)**2))
        
        U_TG, V_TG = setup_Taylor_Green2D(Nx[i],Ny[i],no_vortices = 2)
        
        U_TG *= 10
        V_TG *= 10
        
        solver_torch_np = Solver2D_FFT_np(Nx[i], Ny[i], dt, n_iter)                  
        solver_torch_cpu = Solver2D_FFT_torch(Nx[i], Ny[i], dt, n_iter, gpu = 0)   
        solver_torch_gpu = Solver2D_FFT_torch(Nx[i], Ny[i], dt, n_iter, gpu = 1)               
    
        _, _, _ = solver_torch_np.simulate_2D(U_TG,V_TG,texture_field,n_iter_inflow)
        _, _, _ = solver_torch_cpu.simulate_2D(U_TG,V_TG,texture_field,n_iter_inflow)
        _, _, _ = solver_torch_gpu.simulate_2D(U_TG,V_TG,texture_field,n_iter_inflow)

        t_np_tmp += solver_torch_np.elapsed_time
        t_torch_cpu_tmp += solver_torch_cpu.elapsed_time
        t_torch_gpu_tmp += solver_torch_gpu.elapsed_time

    t_np[i] = t_np_tmp / n_repeats
    t_torch_cpu[i] = t_torch_cpu_tmp / n_repeats
    t_torch_gpu[i] = t_torch_gpu_tmp / n_repeats

# %% Bar chart
categories = ['64x64', '256x256', '512x512', '1024x1024']           # 4 categories

x = np.arange(len(categories))             # Category positions (e.g., [0, 1, 2, 3])
width = 0.2                                # Width of each bar

plt.figure(figsize=(8,4))
plt.bar(x - width, t_np/n_iter, width,  label='Numpy')
plt.bar(x,         t_torch_cpu/n_iter, width,  label='PyTorch')
plt.bar(x + width, t_torch_gpu/n_iter, width,  label='PyTorch w/ GPU')
plt.xlabel('Domain size')
plt.ylabel('Average time per iteration [s]')
plt.xticks(x, categories)
plt.legend()
plt.tight_layout()
plt.show()