FFT_Solver_2D_numpy.py

import numpy as np
from tqdm import tqdm
from numpy.fft import rfft2, irfft2
from timeit import default_timer as timer

class Solver2D_FFT:
        
    def advect_BFECC2D_Periodic(self, input_matrix, U, V):
        
        # Forward warp
        x = self.X - self.dt * U
        y = self.Y - self.dt * V
        x_floor = np.floor(x).astype(int)
        y_floor = np.floor(y).astype(int)
        s = x - x_floor
        t = y - y_floor
        i0 = (self.Nx + (x_floor % self.Nx)) % self.Nx
        i1 = (i0 + 1) % self.Nx
        j0 = (self.Ny + (y_floor %self. Ny)) % self.Ny
        j1 = (j0 + 1) % self.Ny
        U00 = input_matrix[j0, i0]
        U01 = input_matrix[j0, i1]
        U10 = input_matrix[j1, i0]
        U11 = input_matrix[j1, i1]
        forward_warp = (1-s) * ((1-t) * U00 + t * U10) + s * ((1-t) * U01 + t * U11)
        # Backward warp
        backward_warp = self.advect_Linear2D_Periodic(forward_warp, -U, -V)
        # Warp error
        corrected = 0.5  * (3 * input_matrix - backward_warp)
        # Final warp
        advected_result = self.advect_Linear2D_Periodic(corrected, U, V)
        
        # Limiter
        sampling_stack = np.stack((U00,U01,U10,U11), axis = 2)
        min_sampling = np.min(sampling_stack, axis = 2)
        max_sampling = np.max(sampling_stack, axis = 2)
        
        advected_result = np.clip(advected_result,min_sampling,max_sampling)
        
        return advected_result
        
    def advect_MacCormack2D_Periodic(self, input_matrix, U, V):
        
        # Forward warp with min max of sampling coordinates extracted for limiter    
        x = self.X - self.dt * U
        y = self.Y - self.dt * V
        x_floor = np.floor(x).astype(int)
        y_floor = np.floor(y).astype(int)
        s = x - x_floor
        t = y - y_floor
        i0 = (self.Nx + (x_floor % self.Nx)) % self.Nx
        i1 = (i0 + 1) % self.Nx
        j0 = (self.Ny + (y_floor % self.Ny)) % self.Ny
        j1 = (j0 + 1) % self.Ny
        U00 = input_matrix[j0, i0]
        U01 = input_matrix[j0, i1]
        U10 = input_matrix[j1, i0]
        U11 = input_matrix[j1, i1]
        forward_warp = (1-s) * ((1-t) * U00 + t * U10) + s * ((1-t) * U01 + t * U11)
        # Backward warp
        backward_warp = self.advect_Linear2D_Periodic(forward_warp, -U, -V)
        # Correct
        advected_result = forward_warp + 0.5 * (input_matrix - backward_warp)
        
        # Limiter
        sampling_stack = np.stack((U00,U01,U10,U11), axis = 2)
        min_sampling = np.min(sampling_stack, axis = 2)
        max_sampling = np.max(sampling_stack, axis = 2)
        
        advected_result = np.clip(advected_result,min_sampling,max_sampling)
        
        return advected_result
    
    def advect_Fedkiw2D_Periodic(self, input_matrix, U, V):
        
        x = self.X - self.dt * U
        y = self.Y - self.dt * V
        
        x_floor = np.floor(x).astype(int)
        y_floor = np.floor(y).astype(int)
        
        s = x - x_floor
        t = y - y_floor
        
        # Extract sampling coordinate
        i0 = (self.Nx + (x_floor % self.Nx)) % self.Nx
        i_n1 = (i0 - 1) % self.Nx
        i1 = (i0 + 1) % self.Nx
        i2 = (i1 + 1) % self.Nx
        
        j0 = (self.Ny + (y_floor % self.Ny)) % self.Ny
        j_n1 = (j0 - 1) % self.Ny
        j1 = (j0 + 1) % self.Ny
        j2 = (j1 + 1) % self.Ny
    
        # Sample coordinates
        # i - 1
        q_jn1_in1 = input_matrix[j_n1, i_n1]    # j - 1   
        q_j0_in1  = input_matrix[j0,   i_n1]    # j       
        q_j1_in1  = input_matrix[j1,   i_n1]    # j + 1
        q_j2_in1  = input_matrix[j2,   i_n1]    # j + 2  
        # i
        q_jn1_i0 = input_matrix[j_n1, i0]    # j - 1    
        q_j0_i0  = input_matrix[j0,   i0]    # j, j      
        q_j1_i0  = input_matrix[j1,   i0]    # j + 1    
        q_j2_i0  = input_matrix[j2,   i0]    # j + 2     
        # i + 1
        q_jn1_i1 = input_matrix[j_n1, i1]  # j - 1  
        q_j0_i1  = input_matrix[j0, i1]    # j      
        q_j1_i1  = input_matrix[j1, i1]    # j + 1 
        q_j2_i1  = input_matrix[j2, i1]    # j + 2  
        # i + 2
        q_jn1_i2 = input_matrix[j_n1, i2]  # j - 1  
        q_j0_i2  = input_matrix[j0, i2]    # j      
        q_j1_i2  = input_matrix[j1, i2]    # j + 1  
        q_j2_i2  = input_matrix[j2, i2]    # j + 2 
        
        # y-axis interpolation
        # in1         
        delQ = q_j1_in1 - q_j0_in1
        di = (q_j1_in1 - q_jn1_in1)/2
        di_plus1 = (q_j2_in1 - q_j0_in1)/2     
        # Monotonic correction        
        sign_delQ = np.sign(delQ)
        sign_di = np.sign(di)
        sign_di_plus1 = np.sign(di_plus1)
        mask_di = (sign_di != sign_delQ) 
        mask_di_plus1 = (sign_di_plus1 != sign_delQ)
        di[mask_di] = 0
        di_plus1[mask_di_plus1] = 0
        # Interp
        q_in1 = q_j0_in1 + di*t + (3 * delQ - 2 * di - di_plus1)*t**2 + (-2*delQ + di + di_plus1)*t**3

        # i0
        delQ = q_j1_i0 - q_j0_i0
        di = (q_j1_i0 - q_jn1_i0)/2
        di_plus1 = (q_j2_i0 - q_j0_i0)/2  
        # Monotonic correction        
        sign_delQ = np.sign(delQ)
        sign_di = np.sign(di)
        sign_di_plus1 = np.sign(di_plus1)
        mask_di = (sign_di != sign_delQ) 
        mask_di_plus1 = (sign_di_plus1 != sign_delQ)
        di[mask_di] = 0
        di_plus1[mask_di_plus1] = 0
        # Interp
        q_i0 = q_j0_i0 + di*t + (3 * delQ - 2 * di - di_plus1)*t**2 + (-2*delQ + di + di_plus1)*t**3

        # i1
        delQ = q_j1_i1 - q_j0_i1
        di = (q_j1_i1 - q_jn1_i1)/2
        di_plus1 = (q_j2_i1 - q_j0_i1)/2  
        # Monotonic correction        
        sign_delQ = np.sign(delQ)
        sign_di = np.sign(di)
        sign_di_plus1 = np.sign(di_plus1)
        mask_di = (sign_di != sign_delQ) 
        mask_di_plus1 = (sign_di_plus1 != sign_delQ)
        di[mask_di] = 0
        di_plus1[mask_di_plus1] = 0
        # Interp
        q_i1 = q_j0_i1 + di*t + (3 * delQ - 2 * di - di_plus1)*t**2 + (-2*delQ + di + di_plus1)*t**3
        
        # i2
        delQ = q_j1_i2 - q_j0_i2
        di = (q_j1_i2 - q_jn1_i2)/2
        di_plus1 = (q_j2_i2 - q_j0_i2)/2    
        # Monotonic correction        
        sign_delQ = np.sign(delQ)
        sign_di = np.sign(di)
        sign_di_plus1 = np.sign(di_plus1)
        mask_di = (sign_di != sign_delQ) 
        mask_di_plus1 = (sign_di_plus1 != sign_delQ)
        di[mask_di] = 0
        di_plus1[mask_di_plus1] = 0
        # Interp
        q_i2 = q_j0_i2 + di*t + (3 * delQ - 2 * di - di_plus1)*t**2 + (-2*delQ + di + di_plus1)*t**3

        # x-axis interpolation        
        delQ = q_i1 - q_i0
        di = (q_i1 - q_in1)/2
        di_plus1 = (q_i2 - q_i0)/2  
        # Monotonic correction        
        sign_delQ = np.sign(delQ)
        sign_di = np.sign(di)
        sign_di_plus1 = np.sign(di_plus1)
        mask_di = (sign_di != sign_delQ) 
        mask_di_plus1 = (sign_di_plus1 != sign_delQ)
        di[mask_di] = 0
        di_plus1[mask_di_plus1] = 0
        # Interp
        advected_result = q_i0 + di*s + (3 * delQ - 2 * di - di_plus1)*s**2 + (-2*delQ + di + di_plus1)*s**3
        
        return advected_result
        
        
    def advect_Linear2D_Periodic(self, input_matrix, U, V):
            
        x = self.X - self.dt * U
        y = self.Y - self.dt * V
        
        x_floor = np.floor(x).astype(int)
        y_floor = np.floor(y).astype(int)
        
        s = x - x_floor
        t = y - y_floor
        
        i0 = (self.Nx + (x_floor % self.Nx)) % self.Nx
        i1 = (i0 + 1) % self.Nx
        j0 = (self.Ny + (y_floor % self.Ny)) % self.Ny
        j1 = (j0 + 1) % self.Ny
            
        U00 = input_matrix[j0, i0]
        U01 = input_matrix[j0, i1]
        U10 = input_matrix[j1, i0]
        U11 = input_matrix[j1, i1]
        
        advected_result = (1-s) * ((1-t) * U00 + t * U10) + s * ((1-t) * U01 + t * U11)

        return advected_result
    
    def advect_Cubic2D_Periodic(self, input_matrix, U, V):
            
        x = self.X - self.dt * U
        y = self.Y - self.dt * V
        
        x_floor = np.floor(x).astype(np.int16)
        y_floor = np.floor(y).astype(np.int16)
        
        s = x - x_floor
        t = y - y_floor
        
        # Interpolation weights x
        s_sq = s**2
        s_cub = s_sq*s
        s_Wn1 = -1/3 * s + 1/2 * s_sq - 1/6 * s_cub
        s_W0 = 1 - s_sq + 1/2 * (s_cub - s)
        s_W1 = s + 1/2 * (s_sq - s_cub)
        s_W2 = 1/6 * (s_cub - s)
        
        # Interpolation weights y
        t_sq = t**2
        t_cub = t_sq*t
        t_Wn1 = -1/3 * t + 1/2 * t_sq - 1/6 * t_cub
        t_W0 = 1 - t_sq + 1/2 * (t_cub - t)
        t_W1 = t + 1/2 * (t_sq - t_cub)
        t_W2 = 1/6 * (t_cub - t)
        
        # Extract sampling coordinate
        i0 = (self.Nx + (x_floor % self.Nx)) % self.Nx
        i_n1 = (i0 - 1) % self.Nx
        i1 = (i0 + 1) % self.Nx
        i2 = (i1 + 1) % self.Nx
        
        j0 = (self.Ny + (y_floor % self.Ny)) % self.Ny
        j_n1 = (j0 - 1) % self.Ny
        j1 = (j0 + 1) % self.Ny
        j2 = (j1 + 1) % self.Ny
        
        # Sample coordinates
        # j - 1
        q_jn1_in1 = input_matrix[j_n1, i_n1]    # i - 1   for s_Wn1
        q_j0_in1  = input_matrix[j0,   i_n1]    # i       for s_W0
        q_j1_in1  = input_matrix[j1,   i_n1]    # i + 1   for s_W1
        q_j2_in1  = input_matrix[j2,   i_n1]    # i + 2   for s_W2
        # j
        q_jn1_i0 = input_matrix[j_n1, i0]    # i - 1      for s_Wn1
        q_j0_i0  = input_matrix[j0,   i0]    # i, j       for s_W0
        q_j1_i0  = input_matrix[j1,   i0]    # i + 1      for s_W1
        q_j2_i0  = input_matrix[j2,   i0]    # i + 2      for s_W2
        # j + 1
        q_jn1_i1 = input_matrix[j_n1, i1]  # i - 1   for s_Wn1
        q_j0_i1  = input_matrix[j0, i1]    # i       for s_W0
        q_j1_i1  = input_matrix[j1, i1]    # i + 1   for s_W1
        q_j2_i1  = input_matrix[j2, i1]    # i + 2   for s_W2
        # j + 2
        q_jn1_i2 = input_matrix[j_n1, i2]  # i - 1   for s_Wn1
        q_j0_i2  = input_matrix[j0, i2]    # i       for s_W0
        q_j1_i2  = input_matrix[j1, i2]    # i + 1   for s_W1
        q_j2_i2  = input_matrix[j2, i2]    # i + 2   for s_W2
        
        # y-axis interpolation
        q_in1 = t_Wn1 * q_jn1_in1 + t_W0 * q_j0_in1 + t_W1 * q_j1_in1 + t_W2 * q_j2_in1
        q_i0  = t_Wn1 * q_jn1_i0  + t_W0 * q_j0_i0  + t_W1 * q_j1_i0  + t_W2 * q_j2_i0
        q_i1  = t_Wn1 * q_jn1_i1  + t_W0 * q_j0_i1  + t_W1 * q_j1_i1  + t_W2 * q_j2_i1
        q_i2  = t_Wn1 * q_jn1_i2  + t_W0 * q_j0_i2  + t_W1 * q_j1_i2  + t_W2 * q_j2_i2

        # x-axis interpolation
        advected_result = s_Wn1 * q_in1 + s_W0 * q_i0 + s_W1 * q_i1 + s_W2 * q_i2

        return advected_result
    
    def __init__(self, 
                 Nx, Ny, 
                 dt, n_iter, 
                 Nu = 0, Nu_scalar = 0, C_dissip = 0, 
                 VC_weight = 0,
                 advection_type = 'cubic'):
        
        valid_advection_type = {'linear','cubic','bfecc','maccormack','fedkiw'}
        assert advection_type in valid_advection_type, f'Invalid input: {advection_type}. Must be one of: {valid_advection_type}.'
       
        # Interpolation method for advection
        if advection_type == 'cubic':
            self.advect = self.advect_Cubic2D_Periodic
        elif advection_type == 'linear':
            self.advect = self.advect_Linear2D_Periodic
        elif advection_type == 'bfecc':
            self.advect = self.advect_BFECC2D_Periodic
        elif advection_type == 'maccormack':
            self.advect = self.advect_MacCormack2D_Periodic
        elif advection_type == 'fedkiw':
                self.advect = self.advect_Fedkiw2D_Periodic
        else:
            raise ValueError('Unknown advection type')  # Program stops with ValueError
            
        # Dimensions and grid
        self.Nx = Nx
        self.Ny = Ny
        self.X, self.Y = np.meshgrid(np.arange(0,Nx),np.arange(0,Ny), indexing='xy')
        # Iterations and timestep
        self.n_iter = n_iter
        self.dt = dt
        # Viscosities and dissipation constants
        self.Nu = Nu
        self.Nu_scalar = Nu_scalar
        self.C_dissip = C_dissip
        # Vorticity confinement parameter
        self.VC_weight = VC_weight
        self.advection_type = advection_type
        self.elapsed_time = None

    def print_parameters(self):
        
        print('----- Parameter summary -----')
        print(f'Grid resolution: X={self.Nx}, Y={self.Ny}')
        print(f'Iterations: {self.n_iter}')
        print(f'dt: {self.dt}')
        print(f'Fluid viscosity: {self.Nu}')
        print(f'Scalar viscosity: {self.Nu_scalar}')
        print(f'Scalar dissipation: {self.C_dissip}')
        print(f'Advection: {self.advection_type}')
        print(f'Vorticity confinement: {self.VC_weight}')
        
        
    def calc_vorticity_2D(self,U,V): # curl = vorticty = 2X the angular velocity
           
        dU_dy, dU_dx = np.gradient(U)
        dV_dy, dV_dx = np.gradient(V)
        
        result = dV_dx - dU_dy 
        
        return result

    def vorticity_confinement_2D(self,U,V):
        
        omega_z = self.calc_vorticity_2D(U,V) # 2D Vorticity
        
        omega = np.abs(omega_z) # |vorticity|
        
        eta_y, eta_x = np.gradient(omega) # nabla |vorticity|
        
        norm = np.sqrt(eta_x**2 + eta_y**2) + 0.0000001

        Nx = eta_x / norm
        Ny = eta_y / norm

        fx = self.VC_weight *  Ny * omega_z
        fy = self.VC_weight * -Nx * omega_z
        
        return fx, fy
        
    def simulate_2D(self,U_force,V_force,S_force,n_iter_force):
        
        # Forcing function
        assert U_force.shape == V_force.shape, 'Mismatch in U and V forcing function dimensions'
        assert U_force.shape[0:1] == S_force.shape[0:1], 'Mismatch in velocity and density forcing function dimensions'
        assert (U_force.shape[0], U_force.shape[1]) == (self.Ny, self.Nx), \
        'Forcing function dimensions do not match the simulation domain'
        assert n_iter_force <= self.n_iter, 'Forcing exceeds max iteration'
        
        # time_vector = np.arange(0,self.T,self.dt)
        
        U_full = np.zeros((self.Ny,self.Nx,self.n_iter))
        V_full = np.zeros((self.Ny,self.Nx,self.n_iter))
        S_full = np.zeros((self.Ny,self.Nx,self.n_iter))
        
        # Setup wavenumbers
        kx = np.fft.rfftfreq(self.Nx) * self.Nx
        ky = np.fft.fftfreq(self.Ny) * self.Ny
        
        kx, ky = np.meshgrid(kx, ky, indexing='xy')
        k_squared = kx**2 + ky**2
        k_squared[0, 0] = 1  # Avoid division by zero at the zero frequency
        
        kx_norm = kx / k_squared
        ky_norm = ky / k_squared

        U = np.zeros((self.Ny,self.Nx))
        V = np.zeros((self.Ny,self.Nx))
        S = np.zeros((self.Ny,self.Nx))
        
        # Precompute constant denominators
        vel_diffu_denom = (1 + self.Nu * self.dt * k_squared)
        scal_diffu_denom = (1 + self.Nu_scalar * self.dt * k_squared)
        diss_denom = (1 + self.dt * self.C_dissip)
        
        start_time = timer()
        for i in tqdm(range(self.n_iter)): # [!] dt as a step + enumerate?
                        
            #            *** Velocity  step ***
            # -------------- Force velocity --------------
            if i < n_iter_force:  
                U += self.dt * U_force
                V += self.dt * V_force
                S += self.dt * S_force
                
            # -------------- Vorticity confinement --------------
            if self.VC_weight > 0:
                f_vort_x, f_vort_y = self.vorticity_confinement_2D(U,V)
                U += self.dt * f_vort_x
                V += self.dt * f_vort_y
        
            # -------------- Advect --------------    
            U_0 = U
            V_0 = V
            U = self.advect(U, U_0, V_0) # Cubic polynomial
            V = self.advect(V, U_0, V_0)

            U_fft = rfft2(U)
            V_fft = rfft2(V)
            
            # -------------- Diffuse velocity --------------    
            U_fft /= vel_diffu_denom
            V_fft /= vel_diffu_denom
                    
            # -------------- Project --------------    
            div_fft = (U_fft * kx + V_fft * ky)  # (* 1j) - Compute divergence
        
            U = irfft2(U_fft - div_fft * kx_norm)
            V = irfft2(V_fft - div_fft * ky_norm)
                
            #            *** Density  step ***
            # -------------- Advect --------------    
            S = self.advect(S, U, V) # Cubic polynomial

            # -------------- Diffuse density --------------  
            S_fft = rfft2(S)
            S_fft /= scal_diffu_denom
            S = irfft2(S_fft)
            
            # -------------- Dissipate density --------------    
            S /= diss_denom
                
            # -------------- Accumulate timestep result --------------
            U_full[:,:,i] = U
            V_full[:,:,i] = V
            S_full[:,:,i] = S
        
        end_time = timer()
        self.elapsed_time = end_time - start_time
        print(f'Simulation complete in: {self.elapsed_time} seconds')
        
        return U_full, V_full, S_full