import matplotlib.pyplot as plt  
import numpy as np  
from tqdm import tqdm  #进度条不要也罢

N_POINTS = 41  #网格点数
DOMAIN_SIZE = 1.0  #模拟区域大小
N_ITERATIONS = 500  #迭代次数  
TIME_STEP_LENGTH = 0.001  #时间步长
KINEMATIC_VISCOSITY = 0.1  #运动粘度 
DENSITY = 1.0  #密度
HORIZONTAL_VELOCITY_TOP = 1.0  #顶部水平速度 
N_PRESSURE_POISSON_ITERATIONS = 50  #Poisson迭代的次数  
  
def initialize_fields():  
    element_length = DOMAIN_SIZE / (N_POINTS - 1)  
    x = np.linspace(0.0, DOMAIN_SIZE, N_POINTS)  
    y = np.linspace(0.0, DOMAIN_SIZE, N_POINTS)  
    X, Y = np.meshgrid(x, y)  
    u = np.zeros((N_POINTS, N_POINTS))  
    v = np.zeros((N_POINTS, N_POINTS))  
    p = np.zeros((N_POINTS, N_POINTS))  
      
    return X, Y, u, v, p, element_length  #element网格间距
  
def central_difference(f, axis, element_length):    
    diff = np.zeros_like(f)    
    if axis == 'x':    
        diff[1:-1, 1:-1] = (f[1:-1, 2:] - f[1:-1, 0:-2]) / (2 * element_length)    
    elif axis == 'y':    
        diff[1:-1, 1:-1] = (f[2:, 1:-1] - f[0:-2, 1:-1]) / (2 * element_length)    
    return diff 

#计算拉普拉丝算子
def laplace(f, element_length):    
    diff = np.zeros_like(f)    
    diff[1:-1, 1:-1] = (f[1:-1, 0:-2] + f[0:-2, 1:-1] - 4 * f[1:-1, 1:-1] + f[1:-1, 2:] + f[2:, 1:-1]) / (element_length**2)    
    return diff  
  
def set_boundary_conditions(u, v, p):  
    #u=水平速度 v=垂直速度 p=压力
    u[0, :] = 0.0  #上部水平速度设为常数，其他全是0
    u[:, 0] = 0.0
    u[:, -1] = 0.0  
    u[-1, :] = HORIZONTAL_VELOCITY_TOP  
    v[0, :] = 0.0  
    v[:, 0] = 0.0  
    v[:, -1] = 0.0  
    v[-1, :] = 0.0  
    p[:, -1] = p[:, -2]  
    p[0, :] = p[1, :]  
    p[:, 0] = p[:, 1]  
    p[-1, :] = 0.0  
  
def main():     
    X, Y, u_prev, v_prev, p_prev, element_length = initialize_fields()    
    
    for _ in tqdm(range(N_ITERATIONS)):    
        d_u_prev__d_x = central_difference(u_prev, 'x', element_length)    
        d_u_prev__d_y = central_difference(u_prev, 'y', element_length)    
        d_v_prev__d_x = central_difference(v_prev, 'x', element_length)    
        d_v_prev__d_y = central_difference(v_prev, 'y', element_length)    
        laplace__u_prev = laplace(u_prev, element_length)    
        laplace__v_prev = laplace(v_prev, element_length)    
  
        u_tent = u_prev + TIME_STEP_LENGTH * (-(u_prev * d_u_prev__d_x + v_prev * d_u_prev__d_y) + KINEMATIC_VISCOSITY * laplace__u_prev)  
        v_tent = v_prev + TIME_STEP_LENGTH * (-(u_prev * d_v_prev__d_x + v_prev * d_v_prev__d_y) + KINEMATIC_VISCOSITY * laplace__v_prev)
        set_boundary_conditions(u_tent, v_tent, p_prev)  
  
        d_u_tent__d_x = central_difference(u_tent, 'x',  element_length)  
        d_v_tent__d_y = central_difference(v_tent, 'y',  element_length)  
  
        rhs = (DENSITY / TIME_STEP_LENGTH) * (d_u_tent__d_x + d_v_tent__d_y)  
  
        for _ in range(N_PRESSURE_POISSON_ITERATIONS):  
            p_next = np.zeros_like(p_prev)  
            p_next[1:-1, 1:-1] = 0.25 * (p_prev[1:-1, 0:-2] + p_prev[0:-2, 1:-1] + p_prev[1:-1, 2:] + p_prev[2:, 1:-1] - element_length**2 * rhs[1:-1, 1:-1])  
  
            p_next[:, -1] = p_next[:, -2]  
            p_next[0, :] = p_next[1, :]  
            p_next[:, 0] = p_next[:, 1]  
            p_next[-1, :] = 0.0  
  
            p_prev = p_next  
  
        d_p_next__d_x = central_difference(p_next, 'x', element_length)  
        d_p_next__d_y = central_difference(p_next, 'y',  element_length)  
  
        u_next = u_tent - TIME_STEP_LENGTH / DENSITY * d_p_next__d_x  
        v_next = v_tent - TIME_STEP_LENGTH / DENSITY * d_p_next__d_y  
  
        set_boundary_conditions(u_next, v_next, p_next)  
  
        u_prev = u_next  
        v_prev = v_next  
        p_prev = p_next  
    
    plt.figure()  
    plt.contour(X, Y, p_next)  
    plt.colorbar()  
    plt.quiver(X, Y, u_next, v_next, color="black")  
    save_path = 'images/my_plot.png'  
    plt.savefig(save_path) 
 
    plt.close() #显示图像换成plt.show()
  
if __name__ == "__main__":  
    main()