import jax import jax.numpy as np jax.config.update("jax_platform_name", "cpu") jax.config.update("jax_enable_x64", True) import scipy import scipy.special as special import scipy.signal as signal import matplotlib.pyplot as plt import copy import numpy.f2py import os import copy from IPython.display import clear_output with np.load(r"./spectraV1_225_50_50_250000_5e_06_0_005_1_0_1_0.npz", allow_pickle = True) as load: densitiesF = load["densitiesF"] parameters = load["parameters"].item() try: velocitiesF = load["velocitiesF"] has_velocitiesF = True except: has_velocitiesF = False for key in parameters.keys(): exec( key + "=parameters['"+key+"']" ) Boltz = 1.380622e-23 print("total_Time", total_Time) for key in parameters.keys(): exec( key + "=parameters['"+key+"']" ) import scipy.special as special import scipy.signal as signal time_ratio = 0.05 domain_T = total_Time*(1-time_ratio) cutoff_time_step = int(densitiesF.shape[0]*time_ratio) spectra = np.sqrt(2*np.pi)**3/domain_LX/domain_LY/domain_LZ*np.sqrt(2*np.pi)/domain_T*np.abs( delta_x*delta_t*np.fft.fftn(densitiesF[cutoff_time_step:], axes=(-2, -1))/np.sqrt(2*np.pi)/np.sqrt(2*np.pi) )**2 kx_freq = 2*np.pi*np.fft.fftfreq(spectra.shape[1], d=delta_x) omega_freq = 2*np.pi*np.fft.fftfreq(spectra.shape[0], d=delta_t) tau_BGK = viscosity_coef/(density_0/molecule_mass*Boltz*Temp_0) time_scale = 10*tau_BGK# non-dim tau_BGK = 0.1 kx_freq = 2*np.pi*np.fft.fftfreq(spectra.shape[1], d=delta_x) omega_freq = 2*np.pi*np.fft.fftfreq(spectra.shape[0], d=delta_t) space_scale = time_scale*np.sqrt(Boltz*Temp_0/molecule_mass) # Non-dim k_titla_freq = kx_freq*space_scale omega_titla_freq = omega_freq*time_scale ''' k_number=11 omega_plot = np.fft.fftshift(omega_freq)/(kx_freq[k_number]*speed_of_sound) omega_spacing = np.mean( omega_plot[1:]-omega_plot[:-1] ) plot_spectra = np.fft.fftshift( spectra[:,k_number])/space_scale**3/time_scale b, a = signal.butter(1, 20, fs = 1/omega_spacing) filtered_spectra = signal.filtfilt(b, a, plot_spectra) def CE_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) tau1 = -kabs*(1-kabs**2*mu1) tau2 = kabs**2*mu2 tau3 = -kabs*(1-kabs**2*mu3) tau4 = 2/3*kabs**2*kap1 tau5 = -2/3*kabs*(1+kap2) tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec #m = molecule_mass*num_effect/(4*np.pi**2*wavelength**3) #spec = -((complex(0,1)*m*(-4*k**2-20*k**4*Kn**2-23*complex(0,1)*k**2*Kn*omega_freq+6*omega_freq**2))/(density_0*(15*complex(0,1)*k**4*Kn-10*k**2*omega_freq-20*k**4*Kn**2*omega_freq-23*complex(0,1)*k**2*Kn*omega_freq**2+6*omega_freq**3))) #return 2*np.real(spec) def TH_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k, mask = -1): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) print(kabs) tau1 = -0.910858854965221635639183580413607*kabs-0.0681031747662745635208073664825305*kabs**2 tau2 = -0.137647159761080650249789096758410*kabs**2+0.00258989586650062128792005410176610*kabs**3 tau3 = -0.980736834542179305331651160476348*kabs+0.002878781203443508188808986885745442*kabs**3 tau4 = 0.00270043719488066581208735095210*kabs**2-0.000956430940350492874510704526398*kabs**3+0.0000398354558076273539332410866928*kabs**4 #tau4 = 2/3*kabs**2*kap1 tau5 = -0.677006693600986515897927777975082*kabs-0.02720695551116714037510321445687466*kabs**2 tau6 = -0.248523533056964259160237397019776*kabs**2+0.00978711916205478871774188307120469*kabs**3 if mask == 0: tau1 = -kabs*(1-kabs**2*mu1) elif mask == 1: tau2 = kabs**2*mu2 elif mask == 2: tau3 = -kabs*(1-kabs**2*mu3) elif mask == 3: tau4 = 2/3*kabs**2*kap1 elif mask == 4: tau5 = -2/3*kabs*(1+kap2) elif mask == 5: tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec def THBGK_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k, mask = -1): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) print(kabs) #tau1 = -0.910858854965221635639183580413607*kabs-0.0681031747662745635208073664825305*kabs**2 tau1 = -0.9858867248408030*kabs-0.04701941111476196*kabs**2 #tau2 = -0.137647159761080650249789096758410*kabs**2+0.00258989586650062128792005410176610*kabs**3 tau2 = -0.129198118429724*kabs**2+0.00400848955369925*kabs**3 #tau3 = -0.980736834542179305331651160476348*kabs+0.002878781203443508188808986885745442*kabs**3 tau3 = -1.002964233770220*kabs-0.0000911571759799050*kabs**3 #tau4 = 0.00270043719488066581208735095210*kabs**2-0.000956430940350492874510704526398*kabs**3+0.0000398354558076273539332410866928*kabs**4 tau4 = -0.0084451424529939*kabs**2-0.00060032464997635*kabs**3+0.000017125859875981*kabs**4 #tau4 = 2/3*kabs**2*kap1 #tau5 = -0.677006693600986515897927777975082*kabs-0.02720695551116714037510321445687466*kabs**2 tau5 = -0.677965329349480*kabs-0.002348308957845470*kabs**2 #tau6 = -0.248523533056964259160237397019776*kabs**2+0.00978711916205478871774188307120469*kabs**3 tau6 = -0.185259190536085*kabs**2+0.0068634125525231*kabs**3 if mask == 0: tau1 = -kabs*(1-kabs**2*mu1) elif mask == 1: tau2 = kabs**2*mu2 elif mask == 2: tau3 = -kabs*(1-kabs**2*mu3) elif mask == 3: tau4 = 2/3*kabs**2*kap1 elif mask == 4: tau5 = -2/3*kabs*(1+kap2) elif mask == 5: tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec def MW_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) tau1 = -kabs*(1-kabs**2*mu1) tau2 = kabs**2*mu2 tau3 = -kabs*(1-kabs**2*mu3) tau4 = 2/3*kabs**2*kap1 tau5 = -2/3*kabs*(1+kap2) tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec #m = molecule_mass*num_effect/(4*np.pi**2*wavelength**3) #spec = -((complex(0,1)*m*(-4*k**2-20*k**4*Kn**2-23*complex(0,1)*k**2*Kn*omega_freq+6*omega_freq**2))/(density_0*(15*complex(0,1)*k**4*Kn-10*k**2*omega_freq-20*k**4*Kn**2*omega_freq-23*complex(0,1)*k**2*Kn*omega_freq**2+6*omega_freq**3))) #return 2*np.real(spec) plt.title("k = "+ str(k_titla_freq[k_number])) plt.plot(omega_plot , filtered_spectra, label = "DSMC") plt.plot( np.fft.fftshift(omega_freq)/(kx_freq[k_number]*speed_of_sound) , np.fft.fftshift( CE_spectra(space_scale,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_titla_freq,num_effect,k_titla_freq[k_number])) , label = "NS Equation") plt.legend() plt.xlim(-2,2) plt.show() plt.title("k = "+ str(k_titla_freq[k_number])) plt.plot(omega_plot , filtered_spectra, label = "DSMC") plt.plot( np.fft.fftshift(omega_freq)/(kx_freq[k_number]*speed_of_sound) , np.fft.fftshift( TH_spectra(space_scale,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_titla_freq,num_effect,k_titla_freq[k_number])) , label = "TheorysBGK") plt.legend() plt.xlim(-2,2) plt.show() plt.title("k = "+ str(k_titla_freq[k_number])) plt.plot(omega_plot , filtered_spectra, label = "DSMC") plt.plot( np.fft.fftshift(omega_freq)/(kx_freq[k_number]*speed_of_sound) , np.fft.fftshift( THBGK_spectra(space_scale,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_titla_freq,num_effect,k_titla_freq[k_number])) , label = "TheoryBGK") plt.legend() plt.xlim(-2,2) plt.show() ''' k_traval = np.logspace(-2, 2.79817986094985, 2048, endpoint=True) from jax import grad, jit, vmap from numpy.polynomial.hermite import hermgauss from scipy import optimize import jaxopt from scipy.sparse.linalg import LinearOperator from jax.scipy.optimize import minimize from functools import partial # Define the function \tilde{h}^{(+)} @jit def h_plus(u, k, omega): # Dummy implementation of \tilde{h}^{(+)} for demonstration purposes #out = np.sin(np.dot(u, k)[:, np.newaxis] + omega) out = np.sin(np.dot(u, k)[:, np.newaxis] + omega)*0 return out @jit def compute_n_plus(k, h_values, u_mesh, u_weights): u_squared = np.sum(u_mesh**2, axis=-1) pre_factor = 1 # Vectorize the computation over omega integrand = h_values[:, :, np.newaxis] integral = np.sum(integrand * u_weights[:, np.newaxis, np.newaxis], axis=0) q_plus = pre_factor * integral return q_plus @jit def compute_v_plus(k, h_values, u_mesh, u_weights): u_squared = np.sum(u_mesh**2, axis=-1) pre_factor = 1 # Vectorize the computation over omega integrand = u_mesh[:, np.newaxis, :] * h_values[:, :, np.newaxis] integral = np.sum(integrand * u_weights[:, np.newaxis, np.newaxis], axis=0) q_plus = pre_factor * integral return q_plus @jit def compute_T_plus(k, h_values, u_mesh, u_weights): u_squared = np.sum(u_mesh**2, axis=-1) pre_factor = 1 # Vectorize the computation over omega integrand = 1/3 * (u_squared[:, np.newaxis, np.newaxis] * h_values[:, :, np.newaxis]) integral = np.sum(integrand * u_weights[:, np.newaxis, np.newaxis], axis=0) q_plus = pre_factor * integral return q_plus @jit def compute_sigma_plus(k, h_values, u_mesh, u_weights): u_squared = np.sum(u_mesh**2, axis=-1) pre_factor = 1 # Vectorize the computation over omega integrand = ( ((u_mesh[:,np.newaxis,:]*u_mesh[:,:,np.newaxis]).reshape(u_mesh.shape[0],1,-1)- 1/3 * (np.sum(u_mesh**2, axis=-1)[:,np.newaxis,np.newaxis]*np.eye(3).reshape(-1))) * h_values[:, :, np.newaxis]) integral = np.sum(integrand * u_weights[:, np.newaxis, np.newaxis], axis=0) q_plus = pre_factor * integral return q_plus # Function to compute \tilde{q}^{(+)} _i @jit def compute_q_plus(k, h_values, u_mesh, u_weights): u_squared = np.sum(u_mesh**2, axis=-1) pre_factor = 1 # Vectorize the computation over omega integrand = 1/2*u_mesh[:, np.newaxis, :] * (u_squared[:, np.newaxis, np.newaxis] * h_values[:, :, np.newaxis]) integral = np.sum(integrand * u_weights[:, np.newaxis, np.newaxis], axis=0) q_plus = pre_factor * integral return q_plus @jit def taus(k): m = molecule_mass mu1 = 0 Kn = (np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/space_scale)[0] mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) tau1 = -kabs*(1-kabs**2*mu1) tau2 = kabs**2*mu2 tau3 = -kabs*(1-kabs**2*mu3) tau4 = 2/3*kabs**2*kap1 tau5 = -2/3*kabs*(1+kap2) tau6 = 2/3*kabs**2*kap3 return tau1, tau2, tau3, tau4, tau5, tau6 @jit def denominator_reg(denominator): sign = 1-2*(denominator < 0) return sign*np.maximum(np.abs(denominator), 1e-10) @jit def get_rhoRe(k,omega): tau1, tau2, tau3, tau4, tau5, tau6 = taus(k) # Define the numerator numerator = k * omega**2 * (tau1 * tau2 + tau3 * tau4) - k * (tau3 * tau4 - tau1 * tau6) * (tau2 * tau6 + tau3 * tau5) # Define the denominator denominator = (omega**2 * (k**2 * tau1**2 - 2 * k * (tau1 * tau3 * tau5 - tau1 * tau6**2 + tau2 * tau3 * tau4 + tau3 * tau4 * tau6) + (tau2 * tau6 + tau3 * tau5)**2) + k**2 * (tau3 * tau4 - tau1 * tau6)**2 + omega**4 * (2 * k * tau1 + tau2**2 - 2 * tau3 * tau5 + tau6**2) + omega**6) # Define the expression expression = numerator / denominator_reg(denominator) return expression @jit def get_rhoIm(k,omega): tau1, tau2, tau3, tau4, tau5, tau6 = taus(k) # Define the numerator numerator_2 = omega * (-(tau3 * tau5 + tau2 * tau6)**2 - (tau2**2 - 2 * tau3 * tau5 + tau6**2) * omega**2 - omega**4 + k * (tau2 * tau3 * tau4 + tau1 * tau3 * tau5 + tau3 * tau4 * tau6 - tau1 * tau6**2 - tau1 * omega**2)) # Define the denominator denominator_2 = (k**2 * (tau3 * tau4 - tau1 * tau6)**2 + (k**2 * tau1**2 + (tau3 * tau5 + tau2 * tau6)**2 - 2 * k * (tau2 * tau3 * tau4 + tau1 * tau3 * tau5 + tau3 * tau4 * tau6 - tau1 * tau6**2)) * omega**2 + (2 * k * tau1 + tau2**2 - 2 * tau3 * tau5 + tau6**2) * omega**4 + omega**6) # Define the expression expression_2 = numerator_2 / denominator_reg(denominator_2) return expression_2 @jit def get_vRe(k,omega): tau1, tau2, tau3, tau4, tau5, tau6 = taus(k) # Define the numerator numerator_3 = omega * (tau3 * tau4 - tau1 * tau6) * (tau2 * tau6 + tau3 * tau5) - omega**3 * (tau1 * tau2 + tau3 * tau4) # Define the denominator denominator_3 = (omega**2 * (k**2 * tau1**2 - 2 * k * (tau1 * tau3 * tau5 - tau1 * tau6**2 + tau2 * tau3 * tau4 + tau3 * tau4 * tau6) + (tau2 * tau6 + tau3 * tau5)**2) + k**2 * (tau3 * tau4 - tau1 * tau6)**2 + omega**4 * (2 * k * tau1 + tau2**2 - 2 * tau3 * tau5 + tau6**2) + omega**6) # Define the expression expression_3 = numerator_3 / denominator_reg(denominator_3) return expression_3 @jit def get_vIm(k,omega): tau1, tau2, tau3, tau4, tau5, tau6 = taus(k) # Define the numerator numerator_4 = (-k * (tau3 * tau4 - tau1 * tau6)**2 + (-k * tau1**2 + tau2 * tau3 * tau4 + tau1 * tau3 * tau5 + tau3 * tau4 * tau6 - tau1 * tau6**2) * omega**2 - tau1 * omega**4) # Define the denominator denominator_4 = (k**2 * (tau3 * tau4 - tau1 * tau6)**2 + (k**2 * tau1**2 + (tau3 * tau5 + tau2 * tau6)**2 - 2 * k * (tau2 * tau3 * tau4 + tau1 * tau3 * tau5 + tau3 * tau4 * tau6 - tau1 * tau6**2)) * omega**2 + (2 * k * tau1 + tau2**2 - 2 * tau3 * tau5 + tau6**2) * omega**4 + omega**6) # Define the expression expression_4 = numerator_4 / denominator_reg(denominator_4) return expression_4 @jit def get_TRe(k,omega): tau1, tau2, tau3, tau4, tau5, tau6 = taus(k) # Define the numerator numerator_5 = (-omega**2 * (k * tau1 * tau4 + tau1 * tau2 * tau5 + tau1 * tau5 * tau6 + tau2**2 * tau4 - tau3 * tau4 * tau5) + k * (tau1 * tau5 + tau2 * tau4) * (tau3 * tau4 - tau1 * tau6) - tau4 * omega**4) # Define the denominator denominator_5 = (omega**2 * (k**2 * tau1**2 - 2 * k * (tau1 * tau3 * tau5 - tau1 * tau6**2 + tau2 * tau3 * tau4 + tau3 * tau4 * tau6) + (tau2 * tau6 + tau3 * tau5)**2) + k**2 * (tau3 * tau4 - tau1 * tau6)**2 + omega**4 * (2 * k * tau1 + tau2**2 - 2 * tau3 * tau5 + tau6**2) + omega**6) # Define the expression expression_5 = numerator_5 / denominator_reg(denominator_5) return expression_5 @jit def get_TIm(k,omega): tau1, tau2, tau3, tau4, tau5, tau6 = taus(k) # Define the numerator numerator_6 = omega * (-k * (tau3 * tau4**2 + tau1**2 * tau5 + tau1 * tau4 * (tau2 - tau6)) + (tau2 * tau4 + tau1 * tau5) * (tau3 * tau5 + tau2 * tau6) + (-tau1 * tau5 + tau4 * tau6) * omega**2) # Define the denominator denominator_6 = (k**2 * (tau3 * tau4 - tau1 * tau6)**2 + (k**2 * tau1**2 + (tau3 * tau5 + tau2 * tau6)**2 - 2 * k * (tau2 * tau3 * tau4 + tau1 * tau3 * tau5 + tau3 * tau4 * tau6 - tau1 * tau6**2)) * omega**2 + (2 * k * tau1 + tau2**2 - 2 * tau3 * tau5 + tau6**2) * omega**4 + omega**6) # Define the expression expression_6 = numerator_6 / denominator_reg(denominator_6) return expression_6 @jit def get_rhoRe_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Re = Kn/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_n_plus(k, h_values_Re, u_mesh, u_weights) @jit def get_rhoIm_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Im = -Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_n_plus(k, h_values_Im, u_mesh, u_weights) @jit def get_vRe_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Re = Kn/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_v_plus(k, h_values_Re, u_mesh, u_weights) @jit def get_vIm_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Im = -Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_v_plus(k, h_values_Im, u_mesh, u_weights) @jit def get_TRe_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Re = Kn/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_T_plus(k, h_values_Re, u_mesh, u_weights) - compute_n_plus(k, h_values_Re, u_mesh, u_weights) @jit def get_TIm_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Im = -Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_T_plus(k, h_values_Im, u_mesh, u_weights) - compute_n_plus(k, h_values_Im, u_mesh, u_weights) @jit def get_sigmaRe_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Re = Kn/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_sigma_plus(k, h_values_Re, u_mesh, u_weights) @jit def get_sigmaIm_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Im = -Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_sigma_plus(k, h_values_Im, u_mesh, u_weights) @jit def get_qRe_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Re = Kn/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_q_plus(k, h_values_Re, u_mesh, u_weights) - 5/2*compute_v_plus(k, h_values_Re, u_mesh, u_weights) @jit def get_qIm_Asymp(k, Kn, omega_mesh, u_mesh, u_weights): h_values_Im = -Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) return compute_q_plus(k, h_values_Im, u_mesh, u_weights) - 5/2*compute_v_plus(k, h_values_Im, u_mesh, u_weights) @jit def func(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh): lamb = 1 h_values_Re, h_values_Im = x[:len(x)//2], x[len(x)//2:] n_plus_result_Re = compute_n_plus(k, h_values_Re, u_mesh, u_weights) n_plus_result_Im = compute_n_plus(k, h_values_Im, u_mesh, u_weights) v_plus_result_Re = compute_v_plus(k, h_values_Re, u_mesh, u_weights) v_plus_result_Im = compute_v_plus(k, h_values_Im, u_mesh, u_weights) T_plus_result_Re = compute_T_plus(k, h_values_Re, u_mesh, u_weights) - n_plus_result_Re T_plus_result_Im = compute_T_plus(k, h_values_Im, u_mesh, u_weights) - n_plus_result_Im sigma_plus_result_Re = compute_sigma_plus(k, h_values_Re, u_mesh, u_weights) sigma_plus_result_Im = compute_sigma_plus(k, h_values_Im, u_mesh, u_weights) q_plus_result_Re = compute_q_plus(k, h_values_Re, u_mesh, u_weights) - 5/2*v_plus_result_Re q_plus_result_Im = compute_q_plus(k, h_values_Im, u_mesh, u_weights) - 5/2*v_plus_result_Im # Spectra rho0k = 1 # Density impulse h_values_Re_new = Kn*(h_values_Im*omega_mesh + u_mesh.dot(k)[:,np.newaxis]*h_values_Im + rho0k) + lamb*(-\ (1-Pr)*(1 - np.sum( u_mesh**2, axis = -1 )[:,np.newaxis]/5 )*np.sum( u_mesh[:, np.newaxis,:]*q_plus_result_Re, axis = -1 ) +\ n_plus_result_Re[:,0] + np.sum( u_mesh[:, np.newaxis,:]*v_plus_result_Re, axis = -1 ) +\ (np.sum( u_mesh**2, axis = -1 )[:,np.newaxis]/2 - 3/2 )*T_plus_result_Re[:,0] ) h_values_Im_new = -Kn*(h_values_Re*omega_mesh + u_mesh.dot(k)[:,np.newaxis]*h_values_Re) + lamb*( -\ (1-Pr)*(1 - np.sum( u_mesh**2, axis = -1 )[:,np.newaxis]/5 )*np.sum( u_mesh[:, np.newaxis,:]*q_plus_result_Im, axis = -1 )+\ n_plus_result_Im[:,0] + np.sum( u_mesh[:, np.newaxis,:]*v_plus_result_Im, axis = -1 ) +\ (np.sum( u_mesh**2, axis = -1 )[:,np.newaxis]/2 - 3/2 )*T_plus_result_Im[:,0] ) resRe = (h_values_Re_new - h_values_Re)#*(u_weights/np.max(u_weights))[:, np.newaxis] resIm = (h_values_Im_new - h_values_Im)#*(u_weights/np.max(u_weights))[:, np.newaxis] return np.concatenate( (resRe , resIm) , axis = 0) @jit def mv(h_flat, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re): shape = h_values_Re.shape h_values_Re = h_flat[:len(h_flat)//2].reshape(shape) h_values_Im = h_flat[len(h_flat)//2:].reshape(shape) x = np.concatenate( (h_values_Re , h_values_Im) , axis = 0) return func(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh).reshape(-1) @jit def rescompute_mv(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re): return np.mean( mv(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re)**2 ) @jit def grad_rescompute_mv(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re): return grad(rescompute_mv)(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re) @partial(jit, static_argnames=["method"]) def compute_h(k, omega_mesh, u_mesh, u_weights, molecule_mass, Boltz, Temp_0, viscosity_coef, density_0, space_scale, tol=1e-8, verbose = False, method = "SBGK"): #N = 2000 #u_mesh = np.random.randn(N, 3) #u_weights = np.zeros(N) + 1/N #print("u_mesh", u_mesh.shape) #print("u_weights", u_weights.shape) #print("omega_mesh", omega_mesh.shape) khat = (k/np.sqrt(denominator_reg(np.sum(k**2)))) h_values_Re = h_plus(u_mesh, k, omega_mesh) h_values_Im = h_plus(u_mesh, k, omega_mesh) Kn = (np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/space_scale)[0] rhoRe = get_rhoRe(np.abs(k[0]), omega_mesh) rhoIm = get_rhoIm(np.abs(k[0]), omega_mesh) vRe = get_vRe(np.abs(k[0]), omega_mesh) vIm = get_vIm(np.abs(k[0]), omega_mesh) TRe = get_TRe(np.abs(k[0]), omega_mesh) TIm = get_TIm(np.abs(k[0]), omega_mesh) h_values_Re_ini_NS = (u_mesh[:,0:1]*0+rhoRe + np.sum( ( u_mesh[:,np.newaxis,:]*(vRe[:,np.newaxis]*khat) ) ,axis=-1) ) + (np.sum( u_mesh[:,np.newaxis,:]**2 ,axis=-1)/2 - 3/2)*TRe h_values_Im_ini_NS = (u_mesh[:,0:1]*0+rhoIm + np.sum( ( u_mesh[:,np.newaxis,:]*(vIm[:,np.newaxis]*khat) ) ,axis=-1) ) + (np.sum( u_mesh[:,np.newaxis,:]**2 ,axis=-1)/2 - 3/2)*TIm h_values_Re_ini_Asmp = Kn/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) h_values_Im_ini_Asmp = -Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)/(1+Kn**2*( u_mesh.dot(k)[...,None] + omega_mesh)**2) less20 = k[0] < 20 h_values_Re = h_values_Re_ini_NS * less20 + (1 - less20) * h_values_Re_ini_Asmp h_values_Im = h_values_Im_ini_NS * less20 + (1 - less20) * h_values_Im_ini_Asmp ''' h_values_Re = np.zeros_like(h_values_Re_ini) h_values_Im = np.zeros_like(h_values_Im_ini) for i, omega in enumerate(omega_mesh): nearest_index = np.argmin(np.abs(omega_mesh_ini - omega)) #print(len(h_values_Re_ini)) h_values_Re = h_values_Re.at[:,i].set(h_values_Re_ini[:,nearest_index]) h_values_Im = h_values_Im.at[:,i].set(h_values_Im_ini[:,nearest_index]) ''' # sBGK model lamb = 1.0 stepI = 0.005 nite = 10000 reshist = [] h_flat = np.concatenate( (h_values_Re.reshape(-1) , h_values_Im.reshape(-1)) , axis = 0) #print("h_flat", h_flat.shape) ''' A = LinearOperator((len(h_flat),len(h_flat)), matvec=lambda x: mv(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re.shape)) from scipy.sparse.linalg import bicgstab, lgmres, tfqmr x, exit_code = tfqmr(A, np.zeros_like(h_flat), atol=1e-15, x0 = h_flat, maxiter = 1000) print("exit_code",exit_code) print("residual", np.linalg.norm(A@x - np.zeros_like(h_flat))) #x = spsolve(A, np.zeros_like(h_flat), maxiter = 1000) h_values_Re = x[:len(x)//2].reshape(h_values_Re.shape) h_values_Im = x[len(x)//2:].reshape(h_values_Im.shape) ''' #h_ini = np.concatenate( (h_values_Re , h_values_Im) , axis = 0) if method == "SBGK": rescompute = lambda x: rescompute_mv(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re) jac_rescompute = lambda x: grad_rescompute_mv(x, k, Kn, Pr, u_mesh, u_weights, omega_mesh, h_values_Re) solver = jaxopt.LBFGS(fun=rescompute, maxiter=1000, tol=tol) #solver = jaxopt.ScipyMinimize(fun=rescompute, method="L-BFGS-B", maxiter=1000, tol=tol) x = solver.run(h_flat)[0] #x = jax.scipy.optimize.minimize(rescompute,h_flat, tol=1e-3, options = {'maxiter': 1000}, method="BFGS").x #if k[0] > 1 and k[0] < 200: # x = jaxopt.ScipyMinimize(rescompute, h_flat, tol=tol, options = {"disp":False, 'maxiter': 1000}, method="L-BFGS-B").x h_values_Re = x[:len(x)//2].reshape(h_values_Re.shape) h_values_Im = x[len(x)//2:].reshape(h_values_Im.shape) rho_spec_S = np.sum( 2*h_values_Re * u_weights[:,np.newaxis] ,axis = 0) return h_values_Re, h_values_Im, rho_spec_S def hermGauss(N): gh_nodes, gh_weights = hermgauss(N) # integrate exp(- x^2) gh_nodes, gh_weights = 2**0.5*gh_nodes, 2**0.5*gh_weights #https://en.wikipedia.org/wiki/Gauss%E2%80%93Hermite_quadrature # integrate exp(- x^2/2) # Create the 3D mesh for \tilde{\mathbf{u}} u_mesh_x, u_mesh_y, u_mesh_z = np.meshgrid(gh_nodes, gh_nodes, gh_nodes, indexing='ij') u_weights_x, u_weights_y, u_weights_z = np.meshgrid(gh_weights, gh_weights, gh_weights, indexing='ij') # Flatten the mesh and weights for vectorized operations u_mesh = np.stack([u_mesh_x.ravel(), u_mesh_y.ravel(), u_mesh_z.ravel()], axis=-1) u_weights = (1 / (2 * np.pi)**(3/2)) * np.prod(np.stack([u_weights_x.ravel(), u_weights_y.ravel(), u_weights_z.ravel()], axis=-1), axis=-1) return u_mesh, u_weights def monteCarlo(N): u_mesh = np.random.randn(N, 3) u_weights = np.zeros(N) + 1/N return u_mesh, u_weights from scipy.stats import qmc, norm def quasiMonteCarlo(N): # Initialize a Sobol sequence sampler for 3 dimensions sampler = qmc.Sobol(d=3, scramble=True) # Generate N quasi-random samples in the unit cube [0, 1]^3 sample = sampler.random(N) # Transform the samples to follow a standard normal distribution # by applying the inverse cumulative distribution function (CDF) u_mesh = norm.ppf(sample) # Assign equal weights to each sample u_weights = np.full(N, 1/N) return u_mesh, u_weights def quasiMonteCarloGaussian(N): sampler = qmc.MultivariateNormalQMC(mean=np.zeros(3), cov=np.eye(3)) sample = sampler.random(N) u_weights = np.full(N, 1/N) return sample, u_weights # Example usage Pr = 2/3 N = 2**14 #h_values_Res = [] #h_values_Ims = [] omegas = [] k_numbers = [] macro_values = [] amplitude = molecule_mass*num_effect/density_0 u_mesh, u_weights = quasiMonteCarloGaussian(N) impulse = 1/(2*np.pi)**2/space_scale**3*amplitude import tqdm for k_val in tqdm.tqdm(k_traval[1900:1901]): #k_val = k_traval[k_number] k = np.array([k_val, 0, 0]) Kn = (np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/space_scale)[0] eps = 1e-1 peakwidth_NS = (15*(k_val + eps)**2*Kn)/(4*(1+5*(k_val + eps)**2*Kn**2)) peakwidth_Asympt = np.exp(-(1/(2*k_val**2*Kn**2)))*k_val*np.sqrt(2/np.pi) peakwidth = np.maximum(peakwidth_NS, peakwidth_Asympt) domainwidth = 10*abs( (k_val + eps) ) number_of_points = int(20*domainwidth/peakwidth)//2 + 1 #print("number_of_points", number_of_points) omega_mesh = np.linspace(-domainwidth/2, domainwidth/2, number_of_points) rhoRe = get_rhoRe(np.abs(k[0]), omega_mesh) omega_plt = omega_mesh/(denominator_reg(k[0])*(5/3)**0.5) #for N in [9,10, 11, 12]: #for N in [9,10, 11, 12]: #h_values_Re, h_values_Im, rho_spec_S = compute_h(k,omega_mesh, *hermGauss(N), verbose = True) if k_val < 1 or k_val > 200: optmethod = "Theory" else: optmethod = "SBGK" h_values_Re, h_values_Im, rho_spec_S = compute_h(k,omega_mesh, u_mesh, u_weights, molecule_mass, Boltz, Temp_0, viscosity_coef, density_0, space_scale, tol = 1e-8, verbose = True, method = optmethod) n_plus_result_Re = compute_n_plus(k, h_values_Re, u_mesh, u_weights) n_plus_result_Im = compute_n_plus(k, h_values_Im, u_mesh, u_weights) v_plus_result_Re = compute_v_plus(k, h_values_Re, u_mesh, u_weights) v_plus_result_Im = compute_v_plus(k, h_values_Im, u_mesh, u_weights) T_plus_result_Re = compute_T_plus(k, h_values_Re, u_mesh, u_weights) - n_plus_result_Re T_plus_result_Im = compute_T_plus(k, h_values_Im, u_mesh, u_weights) - n_plus_result_Im sigma_plus_result_Re = compute_sigma_plus(k, h_values_Re, u_mesh, u_weights) sigma_plus_result_Im = compute_sigma_plus(k, h_values_Im, u_mesh, u_weights) q_plus_result_Re = compute_q_plus(k, h_values_Re, u_mesh, u_weights) - 5/2*v_plus_result_Re q_plus_result_Im = compute_q_plus(k, h_values_Im, u_mesh, u_weights) - 5/2*v_plus_result_Im current_result = (omega_mesh, n_plus_result_Re, n_plus_result_Im, v_plus_result_Re, v_plus_result_Im, T_plus_result_Re, T_plus_result_Im ) macro_values.append(current_result) # append the current_result to a file jax.clear_caches() #h_values_Res.append(h_values_Re) #h_values_Ims.append(h_values_Im) #plt.plot( omega_plt, impulse*2*rhoRe, label = "NS") #plt.plot( omega_plt, impulse*rho_spec_S, label = "N = "+str(N) ) #plt.xlim(-2*peakwidth /(denominator_reg(k[0])*(5/3)**0.5), 2*peakwidth /(denominator_reg(k[0])*(5/3)**0.5) ) #plt.legend() #plt.show() # save the macro values results import pickle #save_data = { "values": macro_values , "k": k_traval, "omega": omega_mesh } #with open("Boltzmann3_macro_values_Adapt_FullExp.pkl", "wb") as f: # pickle.dump(save_data, f) # load the macro values results #"Boltzmann3_macro_values_500.pkl" def denominator_reg(denominator): sign = 1-2*(denominator < 0) return sign*np.maximum(np.abs(denominator), 1e-10) with open("Boltzmann3_macro_values_Adapt_FullExp.pkl", "rb") as f: save_data = pickle.load(f) macro_values = save_data["values"] k_traval = save_data["k"] omega_mesh = save_data["omega"] def bump_fourier(a, k, b): result1 = 1/(2*np.sqrt(2)*np.pi**(3/2))-(a**2*k**2)/(20*(np.sqrt(2)*np.pi**(3/2))) k = denominator_reg(k) coefficient = 3/(2*np.sqrt(2)*np.pi**(3/2)) numerator = -a * k * np.cos(a * k) + np.sin(a * k) denominator = a**3 * k**3 exponential = np.exp(-k**2 * b**2 / 2) result2 = (coefficient * numerator / denominator) * exponential mask = np.abs(k) < 5e-1 result = result1*exponential*mask + result2*(1-mask) return result #!/usr/bin/env python import numpy as nnp from scipy.special import gamma, spherical_jn from scipy.fft import ifft, rfft, irfft # usefull constants PI = nnp.pi TPI = 2.0 * nnp.pi II = 1.0j class pyNumSBT(object): ''' Numerically perform spherical Bessel transform (SBT) in O(Nln(N)) time based on the algorithm proposed by J. Talman. Talman, J. Computer Physics Communications 2009, 180, 332-338. For a function "f(r)" defined numerically on a LOGARITHMIC radial grid "r", the SBT for the function gives g(k) = \sqrt{2\over\pi} \int_0^\infty j_l(kr) f(r) r^2 dr (c1) and the inverse SBT (iSBT) gives f(r) = \sqrt{2\over\pi} \int_0^\infty j_l(kr) g(k) k^2 dk (c2) ''' def __init__(self, rr, kmax: float=500, lmax: int=10): ''' Init ''' self.rr = nnp.asarray(rr, dtype=float) # the real-space logarithmic radial grid self.nr = self.rr.size # No. of grid points self.nr2 = 2 * self.nr self.sbt_init(rr, kmax) self.lmax = lmax self.sbt_mltb(self.lmax) def sbt_init(self, rr, kmax: float): ''' Initialize the real-space grid (rr) and momentum-space grid (kk). \rho = \ln(rr) \kappa = \ln(kk) The algorithm by Talman requries \Delta\kappa = \Delta\rho ''' self.r_min = rr[0] self.r_max = rr[-1] self.rho_min = nnp.log(self.r_min) self.rho_max = nnp.log(self.r_max) # \Delta\rho = \ln(rr[1] - rr[0]) self.drho = (self.rho_max - self.rho_min) / (self.nr - 1) self.kappa_max = nnp.log(kmax) self.kappa_min = self.kappa_max - (self.rho_max - self.rho_min) # the reciprocal-space (momentum-space) logarithmic grid self.kk = nnp.exp(self.kappa_min) * nnp.exp(nnp.arange(self.nr)*self.drho) self.k_min = self.kk[0] self.k_max = self.kk[-1] self.rr3 = self.rr**3 self.kk3 = self.kk**3 # r values for the extended mesh as discussed in Sec.4 of Talman paper. self.rr_ext = self.r_min * nnp.exp(nnp.arange(-self.nr, 0) * self.drho) # the multiply factor as in Talman paper after Eq.(32) self.rr15 = nnp.zeros((2, self.nr), dtype=float) self.rr15[0] = self.rr_ext**1.5 / self.r_min**1.5 # the extended r-mesh self.rr15[1] = self.rr**1.5 / self.r_min**1.5 # the normal r-mesh # \Delta\rho \times \Delta t = \frac{2\pi}{N} # as in Talman paper after Eq.(30), where N = 2 * nr as discussed in # Sec. 4 self.dt = TPI / (self.nr2 * self.drho) # the post-division factor (1 / k^1.5) as in Eq. (32) of Talman paper. # Note that the factor is (1 / r^1.5) for inverse-SBT, as a result, we # omit the r_min^1.5 / k_min^1.5 here. self.post_div_fac = nnp.exp(-nnp.arange(self.nr)*self.drho*1.5) # Simpson integration of a function on the logarithmic radial grid. # # Setup weights for simpson integration on radial grid any radial integral # can then be evaluated by just summing all radial grid points with the # weights SI # # \int dr = \sum_i w(i) * f(i) self.simp_wht_rr = nnp.zeros_like(self.rr) self.simp_wht_kk = nnp.zeros_like(self.kk) for ii in range(self.nr-1, 1, -2): self.simp_wht_rr[ii] = self.drho * self.rr[ii] / 3. \ + self.simp_wht_rr[ii] self.simp_wht_rr[ii-1] = self.drho * self.rr[ii-1] * 4. / 3. self.simp_wht_rr[ii-2] = self.drho * self.rr[ii-2] / 3. self.simp_wht_kk[ii] = self.drho * self.kk[ii] / 3. \ + self.simp_wht_kk[ii] self.simp_wht_kk[ii-1] = self.drho * self.kk[ii-1] * 4. / 3. self.simp_wht_kk[ii-2] = self.drho * self.kk[ii-2] / 3. def sbt_mltb(self, lmax_in: int): ''' construct the M_l(t) table according to Eq. (15), (16) and (24) of Talman paper. Note that Talman paper use Stirling's approximaton to evaluate the Gamma function, whereas I just call scipy.special.gamma here. ''' if lmax_in > self.lmax: lmax = lmax_in self.lmax = lmax else: lmax = self.lmax tt = nnp.arange(self.nr) * self.dt # M_lt1 is quantity defined by Eq. (12) self.M_lt1 = nnp.zeros((lmax+1, self.nr), dtype=complex) # Eq. (15) in Talman paper # self.M_lt1[0] = gamma(0.5 - II*tt) * nnp.sin(0.5*PI*(0.5 - II*tt)) / self.nr # Eq. (19) and Eq. (15) in Talman paper are equivalent, while the former # is more stable for larger tt self.M_lt1[0] = nnp.sqrt(nnp.pi / 2) * nnp.exp( II * ( nnp.log(gamma(0.5 - II*tt)).imag # + nnp.log(nnp.sin(0.5*PI*(0.5 - II*tt))).imag - nnp.arctan(nnp.tanh(PI*tt/2)) ) ) / self.nr self.M_lt1[0,0] /= 2.0 self.M_lt1[0] *= nnp.exp(II*tt*(self.kappa_min + self.rho_min)) # Eq. (16) in Talman paper phi = nnp.arctan(nnp.tanh(PI*tt/2)) - nnp.arctan(2*tt) self.M_lt1[1] = self.M_lt1[0] * nnp.exp(2*II*phi) # Eq. (24) in Talman paper for ll in range(1, lmax): phi_l = nnp.arctan(2*tt / (2*ll + 1)) self.M_lt1[ll+1] = nnp.exp(-2*II*phi_l) * self.M_lt1[ll-1] ll = nnp.arange(lmax+1) # Eq. (35) in Talman paper xx = nnp.exp(self.rho_min + self.kappa_min + nnp.arange(self.nr2)*self.drho) # M_lt2 is just the Fourier transform of spherical Bessel function j_l self.M_lt2 = ifft( spherical_jn(ll[:,None], xx[None,:]), axis=1 ).conj()[:,:self.nr+1] def run(self, ff, l: int=0, direction: int = 1, norm: bool=False, np_in: int =0, return_rr: bool=False, include_zero: bool=False, ): ''' Perform SBT or inverse-SBT. Input parapeters: ff: the function defined on the logarithmic radial grid l: the "l" as in the underscript of "j_l(kr)" of Eq. (c1) and (c2) direction: 1 for forward SBT and -1 for inverse SBT norm: whether to multiply the prefactor \sqrt{2\over\pi} in Eq. (c1) and (c2). If False, then subsequent applicaton of SBT and iSBT will yield the original data scaled by a factor of 2/pi. np_in: the asymptotic bahavior of ff when r -> 0 ff(r\to 0) \approx r^{np_in + l} include_zero: the SBT does not include the k = 0, i.e. self.kk.min() != 0, term by default. ''' assert l <= self.lmax, \ "lmax = {} smaller than l = {}! Increase lmax!".format(self.lmax, l) ff = nnp.asarray(ff, dtype=float) gg = nnp.zeros_like(ff) r2c_in = nnp.zeros(self.nr2, dtype=float) r2c_out = nnp.zeros(self.nr + 1, dtype=complex) c2r_in = nnp.zeros(self.nr + 1, dtype=complex) c2r_out = nnp.zeros(self.nr2, dtype=float) # The prefactor as in Eq. (c1) and (c2) of the Class docstring. sqrt_2_over_pi = nnp.sqrt(2 / PI) if norm else 1.0 if direction == 1: rmin = self.r_min kmin = self.k_min C = ff[0] / self.r_min**(np_in + l) elif direction == -1: rmin = self.k_min kmin = self.r_min C = ff[0] / self.k_min**(np_in + l) else: raise ValueError("Use direction=1/-1 for forward- and inverse-SBT!") # SBT for LARGE k values extend the input to the doubled mesh, # extrapolating the input as C r^(np_in + l) # Step 1 in the procedure after Eq. (32) of Talman paper r2c_in[:self.nr] = C * self.rr_ext**(np_in + l) * self.rr15[0] r2c_in[self.nr:] = ff * self.rr15[1] # Step 2 in the procedure after Eq. (32) of Talman paper r2c_out = rfft(r2c_in) # Step 3 in the procedure after Eq. (32) of Talman paper tmp1 = nnp.zeros(self.nr2, dtype=complex) tmp1[:self.nr] = r2c_out[:self.nr].conj() * self.M_lt1[l] # Step 4 and 5 in the procedure after Eq. (32) of Talman paper tmp2 = ifft(tmp1) * self.nr2 gg = (rmin / kmin)**1.5 * tmp2[self.nr:].real * self.post_div_fac * sqrt_2_over_pi # obtain the SMALL k results in the array c2r_out if direction == 1: r2c_in[:self.nr] = self.rr3 * ff else: r2c_in[:self.nr] = self.kk3 * ff r2c_in[self.nr:] = 0.0 r2c_out = rfft(r2c_in) c2r_in = r2c_out.conj() * self.M_lt2[l] * sqrt_2_over_pi c2r_out = irfft(c2r_in) * self.nr2 c2r_out[:self.nr] *= self.drho # compare the minimum difference between large and small k as described # in the paragraph above Eq. (39) of Talman paper. gdiff = nnp.abs(gg - c2r_out[:self.nr]) minloc = nnp.argmin(gdiff) gg[:minloc+1] = c2r_out[:minloc+1] # include k = 0 in the SBT and r = 0 in the i-SBT. if include_zero: if direction == 1: gg0 = sqrt_2_over_pi * spherical_jn(l, 0) * nnp.sum( self.simp_wht_rr * self.rr**2 * ff ) else: gg0 = sqrt_2_over_pi * spherical_jn(l, 0) * nnp.sum( self.simp_wht_kk * self.kk**2 * ff ) if return_rr: if direction == 1: return (nnp.r_[0, self.kk], nnp.r_[gg0, gg]) if include_zero else (self.kk, gg) else: return (nnp.r_[0, self.rr], nnp.r_[gg0, gg]) if include_zero else (self.rr, gg) else: return nnp.r_[gg0, gg] if include_zero else gg def run_int(self, ff, l: int=0, direction: int = 1, norm: bool=False, return_rr: bool=False, include_zero: bool=False, ): ''' Perform SBT or inverse-SBT by numerically integrating Eq. (c1) and (c2). Input parapeters: ff: the function defined on the logarithmic radial grid l: the "l" as in the underscript of "j_l(kr)" of Eq. (c1) and (c2) direction: 1 for forward SBT and -1 for inverse SBT norm: whether to multiply the prefactor \sqrt{2\over\pi} in Eq. (c1) and (c2). If False, then subsequent applicaton of SBT and iSBT will yield the original data scaled by a factor of 2/pi. include_zero: the SBT does not include the k = 0, i.e. self.kk.min() != 0, term by default. ''' kr = self.rr[:,None] * self.kk[None,:] jl_kr = spherical_jn(l, kr) ff = nnp.asarray(ff, dtype=float) # The prefactor as in Eq. (c1) and (c2) of the Class docstring. sqrt_2_over_pi = nnp.sqrt(2 / PI) if norm else 1.0 if direction == 1: gg = sqrt_2_over_pi * nnp.sum( jl_kr * (self.rr**2 * ff * self.simp_wht_rr)[:,None], axis=0 ) # include k = 0 in the SBT and r = 0 in the i-SBT. if include_zero: gg0 = sqrt_2_over_pi * spherical_jn(l, 0) * nnp.sum( self.simp_wht_rr * self.rr**2 * ff ) elif direction == -1: gg = sqrt_2_over_pi * nnp.sum( jl_kr * (self.kk**2 * ff * self.simp_wht_kk)[None,:], axis=1 ) # include k = 0 in the SBT and r = 0 in the i-SBT. if include_zero: gg0 = sqrt_2_over_pi * spherical_jn(l, 0) * nnp.sum( self.simp_wht_kk * self.kk**2 * ff ) else: raise ValueError("Use direction=1/-1 for forward- and inverse-SBT!") if return_rr: if direction == 1: return (nnp.r_[0, self.kk], nnp.r_[gg0, gg]) if include_zero else (self.kk, gg) else: return (nnp.r_[0, self.rr], nnp.r_[gg0, gg]) if include_zero else (self.rr, gg) else: return nnp.r_[gg0, gg] if include_zero else gg N = 256 rmin = 2.0 / 1024 / 32 rmax = 30 rr = nnp.logspace(nnp.log10(rmin), nnp.log10(rmax), N, endpoint=True) f1 = 2*nnp.exp(-rr) xx = pyNumSBT(rr) g1 = xx.run(f1, direction=1, norm=False) f2 = xx.run(g1 * g1, direction=-1, norm=False) import matplotlib.pyplot as plt # fig = plt.figure( # dpi=100, # ) # ax = plt.subplot() # # ax.plot(rr, f1) # ax.plot(rr, f2 * 2 / nnp.pi, ls='-') # ax.plot(rr, nnp.exp(-rr) * (1 + rr + rr**2 / 3), ls='--') # ax.set_xlabel('X-label', labelpad=5) # ax.set_ylabel('Y-label', labelpad=5) # plt.tight_layout() # plt.show() from scipy.interpolate import krogh_interpolate n_plus_resoult_Re_fines = [] v_plus_result_Re_fines = [] T_plus_result_Re_fines = [] omega_mesh = np.linspace(-5* 50 - 1e-6, 5*50 + 1e-6, 10419) for i in range(len(macro_values)): omega_mesh_data = macro_values[i][0] n_plus_resoult_Re = macro_values[i][1][...,0] v_plus_result_Re = macro_values[i][3][...,0] T_plus_result_Re = macro_values[i][5][...,0] # interpolate the data n_plus_resoult_Re_fine = np.interp(omega_mesh, omega_mesh_data, n_plus_resoult_Re, left = 0, right = 0) v_plus_result_Re_fine = np.interp(omega_mesh, omega_mesh_data, v_plus_result_Re, left = 0, right = 0) T_plus_result_Re_fine = np.interp(omega_mesh, omega_mesh_data, T_plus_result_Re, left = 0, right = 0) # interpolate the data using the n_plus_resoult_Re_fines.append(n_plus_resoult_Re_fine) v_plus_result_Re_fines.append(v_plus_result_Re_fine) T_plus_result_Re_fines.append(T_plus_result_Re_fine) n_plus_resoult_Re_fines = np.array(n_plus_resoult_Re_fines) v_plus_result_Re_fines = np.array(v_plus_result_Re_fines) T_plus_result_Re_fines = np.array(T_plus_result_Re_fines) kk = np.logspace(-2, 2.79817986094985, 2048, endpoint=True) rho0k = np.array( [ bump_fourier(1, kvalue, 1) for kvalue in kk] ) SBT = pyNumSBT(kk.__array__()) # plt.plot( kk,rho0k, label = "NS") # plt.show() # g1 = SBT.run(rho0k.__array__(), direction=1, norm=True) # plt.plot( SBT.kk, g1) # plt.xlim(0, 10) # plt.show() # f1 = SBT.run(g1.__array__(), direction=-1, norm=True) # plt.plot( kk, f1 ) # plt.plot( kk, rho0k, linestyle = "dotted") # plt.xlim(0, 10) # plt.show() # Compute delta_rho(k) from delta_rho(r) (Forward Transform) def forward_transform(delta_rho_r, r, k): # Compute kr matrix kr = np.outer(k, r) # Shape: (N_k, N_r) # Compute the sine of kr sin_kr = np.sinc(kr / np.pi) # Compute the integrand integrand = delta_rho_r * r**2 * sin_kr # Broadcasting over r # Compute the integral over r for each k integral = np.trapezoid(integrand, x=r, axis=1) # Shape: (N_k,) # Compute the prefactor prefactor = (4 * np.pi) / ((2 * np.pi) ** 1.5 * np.ones_like(k)) # Compute delta_rho(k) delta_rho_k = prefactor * integral return delta_rho_k # Compute delta_rho(r) from delta_rho(k) (Inverse Transform) def inverse_transform(delta_rho_k, k, r): # Compute kr matrix kr = np.outer(r, k) # Shape: (N_r, N_k) # Compute the sine of kr sin_kr = np.sinc(kr / np.pi) # Compute the integrand integrand = delta_rho_k * k**2 * sin_kr # Broadcasting over k # Compute the integral over k for each r integral = np.trapezoid(integrand, x=k, axis=1) # Shape: (N_r,) # Compute the prefactor prefactor = (4 * np.pi) / ((2 * np.pi) ** 1.5 * np.ones_like(r) ) # Compute delta_rho(r) delta_rho_r_rec = prefactor * integral return delta_rho_r_rec rho0k = np.array( [ bump_fourier(2, kvalue, 100) for kvalue in k_traval] ) import torch import torch.nn as nn import numpy as nnp device = "cpu" dtype = torch.float sNet= torch.load("./DSMC3ModelsExp/DSMC3LearnModelFull6.pt", weights_only=False, map_location=torch.device('cpu')) #sNet.load_state_dict(torch.load("./DSMC3ModelsV/DSMC3LearnModelN250", map_location=torch.device('cpu'))) sNet.eval() kTest = 5 with torch.no_grad(): net_spectra, net_spectra_omega, _, _, _, _, _ = sNet(torch.tensor( nnp.array( kTest ), device = device, dtype = dtype), torch.tensor( nnp.array(omega_mesh[None,:]/(kTest *(5/3)**0.5) ), device = device, dtype = dtype)) net_spectra = net_spectra.cpu().numpy() net_spectra_omega = net_spectra_omega.cpu().numpy() # plt.plot( omega_mesh, 2*get_rhoRe(kTest , omega_mesh), label = "NS", linewidth = 2) # plt.plot(omega_mesh, net_spectra[0]/impulse, label = "NN") # plt.legend() # plt.show() #k_traval = np.linspace( 0, 30, 200 ) # suppress the outbound omega spectra by a Gaussian filter in omega space #omega_fft_range = np.max( np.abs(omega_mesh) ) #fft_omega_filter = np.exp( - (omega_mesh)**2 / (2 * (omega_fft_range/2)**2) ) rho_k_omega_Bol3_re = 2*n_plus_resoult_Re_fines#*fft_omega_filter rho_k_omega_NS_re = np.array([2*get_rhoRe(k_traval[i], omega_mesh) for i in range(len(k_traval))])#*fft_omega_filter #rho_k_omega_Bol3_re = rho_k_omega_Bol3_re.at[k_traval < 1].set(rho_k_omega_NS_re[k_traval < 1]) rho_k_omega_Learn_re = [] #rho_k_omega_Learn_re[0] = 2*get_rhoRe(k_traval[0], omega_mesh) v_plus_k_omega_Bol3_re = 2*v_plus_result_Re_fines#*fft_omega_filter v_plus_k_omega_NS_re = np.array([2*get_vRe(k_traval[i], omega_mesh) for i in range(len(k_traval))])#*fft_omega_filter #v_plus_k_omega_Bol3_re = v_plus_k_omega_Bol3_re.at[k_traval < 1].set(v_plus_k_omega_NS_re[k_traval < 1]) T_plus_k_omega_Bol3_re = 2*T_plus_result_Re_fines#*fft_omega_filter T_plus_k_omega_NS_re = np.array([2*get_TRe(k_traval[i], omega_mesh) for i in range(len(k_traval))])#*fft_omega_filter T_plus_k_omega_Learn_re = [] for kTest in k_traval: with torch.no_grad(): net_spectra, net_spectra_omega, _, rhoSpec, vSpec, TSpec, _ = sNet(torch.tensor( nnp.array( kTest ), device = device, dtype = dtype), torch.tensor( nnp.array(omega_mesh[None,:]/(kTest *(5/3)**0.5) ), device = device, dtype = dtype)) net_spectra = net_spectra.cpu().numpy() net_spectra_omega = net_spectra_omega.cpu().numpy() rho_k_omega_Learn_re.append( rhoSpec[0].cpu().numpy()/impulse ) T_plus_k_omega_Learn_re.append( TSpec[0].cpu().numpy()/impulse ) rho_k_omega_Learn_re = np.array(rho_k_omega_Learn_re)#*fft_omega_filter T_plus_k_omega_Learn_re = np.array(T_plus_k_omega_Learn_re)#*fft_omega_filter #T_plus_k_omega_Bol3_re = T_plus_k_omega_Bol3_re.at[k_traval < 1].set(T_plus_k_omega_NS_re[k_traval < 1]) #for i in range(20,25): #plt.plot( rho_k_omega_Bol3_re[i], label = "Bol3") #plt.plot( rho_k_omega_NS_re[i], label = "NS") #plt.legend() #plt.show() # rho_k_omega_Bol3_re = rho_k_omega_Bol3_re.at[i].set(rho_k_omega_NS_re[i]) # # plt.plot( omega_mesh, rho_k_omega_Bol3_re[100], linewidth = 3, label = "Bol3") # plt.plot( omega_mesh, rho_k_omega_NS_re[100], label = "NS", linewidth = 1) # plt.plot( omega_mesh, rho_k_omega_Learn_re[100], label = "Learn", linewidth = 0.5) # plt.xlim(-30,30) # plt.legend() # plt.show() rho0k = np.array( [ bump_fourier(0.5, kvalue, 0.02) for kvalue in k_traval] ) v0k = np.zeros_like(rho0k) T0k = np.zeros_like(rho0k) # plt.plot( k_traval,rho0k, label = "NS") # plt.show() time_travel = 2*np.pi*np.fft.fftfreq(len(omega_mesh), d = omega_mesh[1] - omega_mesh[0]) SBT = pyNumSBT(k_traval.__array__()) r_travel = SBT.kk rho0x = SBT.run(rho0k.__array__(), direction=1, norm=True) # plt.plot( r_travel, rho0x) # plt.xlim(0,2) # plt.show() ''' r_travel = np.linspace( 0, np.pi/(k_traval[1] - k_traval[0]), len(k_traval) ) rho0x = inverse_transform(rho0k, k_traval, r_travel) plt.plot( r_travel, rho0x) plt.xlim(0,2) plt.show() ''' wavelength = 2*np.pi/k_traval*space_scale Kn_number = viscosity_coef/(density_0*wavelength)/np.sqrt(Boltz*Temp_0/molecule_mass) #rho_k_omega_re = np.array([compute_n_plus( None, 2*h_values_Res[i], u_mesh, u_weights)[...,0] for i in range(len(k_traval))]) #rho_k_omega_re = rho_k_omega_re.at[0].set(2*get_rhoRe(5e-2, omega_mesh)) #plt.plot( omega_plt, rho_k_omega_re[0], label = "NS") #rhoRe = get_rhoRe(k_titla_freq[0], omega_mesh) #plt.plot(omega_plt,2*rhoRe) #plt.show() #plt.show() def time_evolution(rho_k_omega_re, rho0k, k_traval, omega_mesh): #print( rho0k ) rho_k_omega_re = rho_k_omega_re * rho0k[:,np.newaxis] /(2*np.pi)**0.5 rho_k_t = np.real( np.fft.ifft( np.fft.fftshift( rho_k_omega_re[:,:-1] , axes=-1) , axis = -1) ) / (2*np.pi)**0.5 * ( (omega_mesh[1] - omega_mesh[0]) * len(omega_mesh) ) rho_k_values = k_traval SBT = pyNumSBT(k_traval.__array__()) x_travel = SBT.kk #x_travel = np.linspace(0, np.pi/(k_traval[1] - k_traval[0]), len(k_traval) ) rho_x_t = np.array( [SBT.run(rho_k_t[:,i].__array__(), direction=1, norm=True) for i in range( rho_k_t.shape[1] ) ] ).T #rho_x_t = np.array( [ inverse_transform(rho_k_t[:,i], k_traval, x_travel) for i in range( rho_k_t.shape[1] ) ] ).T time_travel = 2*np.pi*np.fft.fftfreq(len(omega_mesh), d = omega_mesh[1] - omega_mesh[0]) return x_travel, time_travel, rho_x_t x_travel, time_travel, rho_x_t_NS = time_evolution(rho_k_omega_NS_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, rho_x_t_Bol3 = time_evolution(rho_k_omega_Bol3_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, rho_x_t_Learn = time_evolution(rho_k_omega_Learn_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, v_x_t_NS = time_evolution(v_plus_k_omega_NS_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, v_x_t_Bol3 = time_evolution(v_plus_k_omega_Bol3_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, T_x_t_NS = time_evolution(T_plus_k_omega_NS_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, T_x_t_Bol3 = time_evolution(T_plus_k_omega_Bol3_re, rho0k, k_traval, omega_mesh) x_travel, time_travel, T_x_t_Learn = time_evolution(T_plus_k_omega_Learn_re, rho0k, k_traval, omega_mesh) # plot the time evolution in 2 x 2 grid tplots = [0.0, 0.2, 0.4] fig, axs = plt.subplots(len(tplots)+1, 2, figsize=(8, 2.5*(len(tplots)+1)), dpi=200) for i,time in enumerate([0.0, 0.2, 0.4]): time_index = np.argmin(np.abs(time_travel - time)) # +1 here is to add reference rho and T back to the perturbation value axs[i, 0].plot(x_travel, rho_x_t_Bol3[:, time_index]+1, label = "SBGK t=" + "{:.2f}".format( time_travel[ time_index ] ) , linewidth = 2) axs[i, 0].plot(x_travel, rho_x_t_NS[:, time_index]+1, linestyle = 'dotted', label = "NS t=" + "{:.2f}".format( time_travel[ time_index ] ) , linewidth = 2) axs[i, 0].plot(x_travel, rho_x_t_Learn[:, time_index]+1, linestyle = 'dashed', label = "Learn t=" + "{:.2f}".format( time_travel[ time_index ] ) , linewidth = 2) axs[i, 0].set_title("Density") axs[i, 0].legend() axs[i, 0].set_xlim(0, 1) axs[i, 0].set_ylim(1-0.1, 1+2.1) axs[i, 0].set_xlabel("r") axs[i, 0].set_ylabel("ρ") axs[i, 1].plot(x_travel, T_x_t_Bol3[:, time_index]+1, label = "SBGK t=" + "{:.2f}".format( time_travel[ time_index ] ) , linewidth = 2) axs[i, 1].plot(x_travel, T_x_t_NS[:, time_index]+1, linestyle = 'dotted', label = "NS t=" + "{:.2f}".format( time_travel[ time_index ] ) , linewidth = 2) axs[i, 1].plot(x_travel, T_x_t_Learn[:, time_index]+1, linestyle = 'dashed', label = "Learn t=" + "{:.2f}".format( time_travel[ time_index ] ) , linewidth = 2) axs[i, 1].set_title("Temperature") axs[i, 1].legend() axs[i, 1].set_xlim(0, 1) axs[i, 1].set_ylim(1-1, 1+0.2) axs[i, 1].set_xlabel("r") axs[i, 1].set_ylabel("T") #axs[0, 1].set_title("Velocity") axs[i+1, 0].plot( Kn_number,rho0k, label = "t=0.00") axs[i+1, 0].set_xlim(0,0.5) axs[i+1, 0].set_xlabel("Kn") axs[i+1, 0].set_ylabel("ρ") axs[i+1, 0].set_title("Initial Density (Fourier Space)") axs[i+1, 0].legend() axs[i+1, 1].plot( Kn_number,rho0k, label = "t=0.00") axs[i+1, 1].set_xlim(0.5,2) axs[i+1, 1].set_ylim(-0.001,0.001) axs[i+1, 1].set_xlabel("Kn") axs[i+1, 1].set_ylabel("ρ") from matplotlib.ticker import FormatStrFormatter axs[i+1, 1].xaxis.set_major_formatter(FormatStrFormatter('%.2f')) axs[i+1, 1].yaxis.set_major_formatter(FormatStrFormatter('%.3f')) axs[i+1, 1].set_title("Initial Density (Fourier Space)") axs[i+1, 1].legend() fig.tight_layout() fig.savefig("dynamics.png", dpi=200) def CE_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) tau1 = -kabs*(1-kabs**2*mu1) tau2 = kabs**2*mu2 tau3 = -kabs*(1-kabs**2*mu3) tau4 = 2/3*kabs**2*kap1 tau5 = -2/3*kabs*(1+kap2) tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec #m = molecule_mass*num_effect/(4*np.pi**2*wavelength**3) #spec = -((complex(0,1)*m*(-4*k**2-20*k**4*Kn**2-23*complex(0,1)*k**2*Kn*omega_freq+6*omega_freq**2))/(density_0*(15*complex(0,1)*k**4*Kn-10*k**2*omega_freq-20*k**4*Kn**2*omega_freq-23*complex(0,1)*k**2*Kn*omega_freq**2+6*omega_freq**3))) #return 2*np.real(spec) def R13_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) tau1 = -kabs*(1-kabs**2*mu1) tau2 = kabs**2*mu2 tau3 = -kabs*(1-kabs**2*mu3) tau4 = 2/3*kabs**2*kap1 tau5 = -2/3*kabs*(1+kap2) tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = 2*(15*k**4*Kn*m*Neff*(5*k**2*(5+6*k**2*Kn**2)*(25+210*k**2*Kn**2+162*k**4*Kn**4)+(500+5775*k**2*Kn**2+29295*k**4*Kn**4+17496*k**6*Kn**6)*omega**2+225*Kn**2*(5+6*k**2*Kn**2)*omega**4))/(L**3*np.pi**2*density_0*(5625*k**8*Kn**2*(5+6*k**2*Kn**2)**2+25*k**4*(2500+3*k**2*Kn**2*(3500+35175*k**2*Kn**2+60480*k**4*Kn**4+34992*k**6*Kn**6))*omega**2-30*k**2*(2500+3*k**2*Kn**2*(625+20850*k**2*Kn**2-6030*k**4*Kn**4+34992*k**6*Kn**6))*omega**4+18*(1250-7750*k**2*Kn**2+50775*k**4*Kn**4-125820*k**6*Kn**6+52488*k**8*Kn**8)*omega**6+1125*Kn**2*(65+72*k**2*Kn**2*(-2+9*k**2*Kn**2))*omega**8+50625*Kn**4*omega**10)) return spec def TH_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k, mask = -1): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) print(kabs) tau1 = -0.910858854965221635639183580413607*kabs-0.0681031747662745635208073664825305*kabs**2 tau2 = -0.137647159761080650249789096758410*kabs**2+0.00258989586650062128792005410176610*kabs**3 tau3 = -0.980736834542179305331651160476348*kabs+0.002878781203443508188808986885745442*kabs**3 tau4 = 0.00270043719488066581208735095210*kabs**2-0.000956430940350492874510704526398*kabs**3+0.0000398354558076273539332410866928*kabs**4 #tau4 = 2/3*kabs**2*kap1 tau5 = -0.677006693600986515897927777975082*kabs-0.02720695551116714037510321445687466*kabs**2 tau6 = -0.248523533056964259160237397019776*kabs**2+0.00978711916205478871774188307120469*kabs**3 if mask == 0: tau1 = -kabs*(1-kabs**2*mu1) elif mask == 1: tau2 = kabs**2*mu2 elif mask == 2: tau3 = -kabs*(1-kabs**2*mu3) elif mask == 3: tau4 = 2/3*kabs**2*kap1 elif mask == 4: tau5 = -2/3*kabs*(1+kap2) elif mask == 5: tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec def THBGK_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k, mask = -1): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) print(kabs) #tau1 = -0.910858854965221635639183580413607*kabs-0.0681031747662745635208073664825305*kabs**2 #tau1 = -0.9858867248408030*kabs-0.04701941111476196*kabs**2 tau1 = -0.991013044792487766036126122635640*kabs-0.0403239332346867445520412638803096*kabs**2-0.002166184469076816223924007295632155*kabs**3 #tau2 = -0.137647159761080650249789096758410*kabs**2+0.00258989586650062128792005410176610*kabs**3 #tau2 = -0.129198118429724*kabs**2+0.00400848955369925*kabs**3 tau2 = -0.149161289391781081302603295336010*kabs**2 + 0.00594195008578441066220857129869587*kabs**3 - 0.000232348591684703766188347436129974*kabs**4 #tau3 = -0.980736834542179305331651160476348*kabs+0.002878781203443508188808986885745442*kabs**3 #tau3 = -1.002964233770220*kabs-0.0000911571759799050*kabs**3 tau3 = -1.064641200298265511954132389066412*kabs + 0.02791589569574715853735528833289102*kabs**2 + 0.000762400543447067776385421347639320*kabs**3 #tau4 = 0.00270043719488066581208735095210*kabs**2-0.000956430940350492874510704526398*kabs**3+0.0000398354558076273539332410866928*kabs**4 #tau4 = -0.0084451424529939*kabs**2-0.00060032464997635*kabs**3+0.000017125859875981*kabs**4 tau4 = 0.00270043719488066581208735095210*kabs**2 - 0.000956430940350492874510704526398*kabs**3 + 0.0000398354558076273539332410866928*kabs**4 #tau4 = 2/3*kabs**2*kap1 #tau5 = -0.677006693600986515897927777975082*kabs-0.02720695551116714037510321445687466*kabs**2 #tau5 = -0.677965329349480*kabs-0.002348308957845470*kabs**2 tau5 = -0.664300463755129273501867308536487*kabs - 0.0316105859215623339424006394153918*kabs^2 + 0.000343388633980720341529015301083001*kabs^3 #tau6 = -0.248523533056964259160237397019776*kabs**2+0.00978711916205478871774188307120469*kabs**3 #tau6 = -0.185259190536085*kabs**2+0.0068634125525231*kabs**3 tau6 = -0.280753259910037449250341065089074*kabs^2 + 0.0191700074349251339913376848978025*kabs^3 - 0.000650377569549622175919978516795108*kabs^4 if mask == 0: tau1 = -kabs*(1-kabs**2*mu1) elif mask == 1: tau2 = kabs**2*mu2 elif mask == 2: tau3 = -kabs*(1-kabs**2*mu3) elif mask == 3: tau4 = 2/3*kabs**2*kap1 elif mask == 4: tau5 = -2/3*kabs*(1+kap2) elif mask == 5: tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec def MW_spectra(delta_x,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,omega_freq,num_effect,k): m = molecule_mass mu1 = 0 Kn = np.sqrt(molecule_mass/(Boltz*Temp_0))*viscosity_coef/density_0/delta_x mu2 = -4/3*Kn mu3 = 0 kap1 = 0 kap2 = 0 kap3 = -15/4*Kn kabs = np.abs(k) tau1 = -kabs*(1-kabs**2*mu1) tau2 = kabs**2*mu2 tau3 = -kabs*(1-kabs**2*mu3) tau4 = 2/3*kabs**2*kap1 tau5 = -2/3*kabs*(1+kap2) tau6 = 2/3*kabs**2*kap3 omega = omega_freq Neff = num_effect rho0 = density_0 L = delta_x spec = (k*m*Neff*(-(tau3*tau4-tau1*tau6)*(tau3*tau5+tau2*tau6)+(tau1*tau2+tau3*tau4)*omega**2))/(2*L**3*np.pi**2*rho0*(k**2*(tau3*tau4-tau1*tau6)**2+(k**2*tau1**2+(tau3*tau5+tau2*tau6)**2-2*k*(tau2*tau3*tau4+tau1*tau3*tau5+tau3*tau4*tau6-tau1*tau6**2))*omega**2+(2*k*tau1+tau2**2-2*tau3*tau5+tau6**2)*omega**4+omega**6)) return spec #m = molecule_mass*num_effect/(4*np.pi**2*wavelength**3) #spec = -((complex(0,1)*m*(-4*k**2-20*k**4*Kn**2-23*complex(0,1)*k**2*Kn*omega_freq+6*omega_freq**2))/(density_0*(15*complex(0,1)*k**4*Kn-10*k**2*omega_freq-20*k**4*Kn**2*omega_freq-23*complex(0,1)*k**2*Kn*omega_freq**2+6*omega_freq**3))) #return 2*np.real(spec) wavelength = 2*np.pi/k_traval*space_scale Kn_number = viscosity_coef/(density_0*wavelength)/np.sqrt(Boltz*Temp_0/molecule_mass) Kn_number plt_grid = (2, 3) fig, axs = plt.subplots(2, 3, figsize=(14, 8), dpi = 200) for i, kn_plot in enumerate( [0.01511645, 0.04534935, 0.10116449, 0.63489085, 2.51188643, 10. ] ): ind = np.argmin( np.abs(Kn_number - kn_plot) ) k_plot = k_traval[ind] axs[i//plt_grid[1], i%plt_grid[1]].plot(macro_values[ind][0]/(k_plot*(5/3)**0.5), impulse*2*macro_values[ind][1][...,0], label = "SBGK (ref)") axs[i//plt_grid[1], i%plt_grid[1]].plot(macro_values[ind][0]/(k_plot*(5/3)**0.5), impulse*2*get_rhoRe(k_traval[ind], macro_values[ind][0]), label = "NS", linestyle = 'dotted') axs[i//plt_grid[1], i%plt_grid[1]].plot(macro_values[ind][0]/(k_plot*(5/3)**0.5) , R13_spectra(space_scale,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,macro_values[ind][0],num_effect,k_plot) , label = "R13 Equation", linestyle = 'dashed') with torch.no_grad(): net_spectra, net_spectra_omega, _, rhoSpec, vSpec, TSpec,_ = sNet(torch.tensor( nnp.array( k_plot ), device = device, dtype = dtype), torch.tensor( nnp.array(macro_values[ind][0][None,:]/(k_plot *(5/3)**0.5) ), device = device, dtype = dtype)) net_spectra = net_spectra.cpu().numpy() net_spectra_omega = net_spectra_omega.cpu().numpy() axs[i//plt_grid[1], i%plt_grid[1]].plot(net_spectra_omega[0] , net_spectra[0], label = "Learned") if k_plot < 50: k_number = np.argmin( np.abs(k_titla_freq - k_plot) ) omega_plot = np.fft.fftshift(omega_freq)/(kx_freq[k_number]*speed_of_sound) omega_spacing = np.mean( omega_plot[1:]-omega_plot[:-1] ) plot_spectra = np.fft.fftshift( spectra[:,k_number])/space_scale**3/time_scale b, a = signal.butter(1, 20, fs = 1/omega_spacing) filtered_spectra = signal.filtfilt(b, a, plot_spectra) axs[i//plt_grid[1], i%plt_grid[1]].plot(omega_plot , filtered_spectra, label = "DSMC (ref)") if k_plot < 13: #axs[i//plt_grid[1], i%plt_grid[1]].plot(macro_values[ind][0]/(k_plot*(5/3)**0.5) , TH_spectra(space_scale,viscosity_coef,Boltz,Temp_0,molecule_mass,density_0,macro_values[ind][0],num_effect,k_plot) , label = "Hydro Manifold") pass axs[i//plt_grid[1], i%plt_grid[1]].set_xlim(0,2) axs[i//plt_grid[1], i%plt_grid[1]].set_title("Kn = " + "{:.2f}".format(Kn_number[ind])) axs[i//plt_grid[1], i%plt_grid[1]].set_xlabel("ω") axs[i//plt_grid[1], i%plt_grid[1]].set_ylabel("$<ρ^2>$") axs[i//plt_grid[1], i%plt_grid[1]].legend() # adjust the space between the plots plt.tight_layout() fig.savefig("spectra.png", dpi=200) #rho_k_omega_re = ( np.flip(rho_k_omega_re,axis=-1) + rho_k_omega_re )/2