Run.py

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