Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+dynamics.png filter=lfs diff=lfs merge=lfs -text
+spectra.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,37 @@
+# Large data and model files (exceed GitHub's 100MB limit or kept out of repo)
+# See README for how to obtain these.
+spectraV1_225_50_50_250000_5e_06_0_005_1_0_1_0.npz
+Boltzmann3_macro_values_Adapt_FullExp.pkl
+DSMC3ModelsExp/DSMC3LearnModelFull6.pt
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+*.egg-info/
+.eggs/
+dist/
+build/
+
+# Virtual envs
+.env
+.venv
+venv/
+env/
+
+# IDE / OS
+.idea/
+.vscode/
+.DS_Store
+*.swp
+*~
+
+# Jupyter
+.ipynb_checkpoints/
+*.ipynb_checkpoints
+
+# Optional: uncomment to also ignore generated outputs
+# spectra.png
+# dynamics.png

README.md

@@ -1,3 +1,42 @@
----
-license: mit
----
+# Learning the Optimal Linear Hydrodynamic Closure
+
+Code for generating spectral and time-evolution comparisons used in the paper *Learning the Optimal Linear Hydrodynamic Closure*. The main entry point is `Run.py`.
+
+## Setup
+
+Create a fresh environment and install the dependencies:
+
+```bash
+conda create -n nc-code python=3.11
+conda activate nc-code
+pip install -r requirements.txt
+```
+
+If you prefer, you can use `venv` instead of Conda.
+
+## Required Files
+
+Before running `Run.py`, make sure these files are present:
+
+- `spectraV1_225_50_50_250000_5e_06_0_005_1_0_1_0.npz`
+- `Boltzmann3_macro_values_Adapt_FullExp.pkl`
+- `DSMC3ModelsExp/DSMC3LearnModelFull6.pt`
+
+## Run
+
+Run from the repository root:
+
+```bash
+python Run.py
+```
+
+The script loads the precomputed data and trained model, then writes:
+
+- `spectra.png`
+- `dynamics.png`
+
+## Notes
+
+- `Run.py` uses relative paths, so run it from the repository root.
+- JAX is set to CPU mode in the script.
+- `jax` and `torch` may need platform-specific installation steps on some systems.

Run.py

@@ -0,0 +1,1612 @@
+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

requirements.txt

@@ -0,0 +1,9 @@
+jax
+jaxlib
+numpy
+scipy
+matplotlib
+ipython
+jaxopt
+tqdm
+torch