<jupyter_start><jupyter_text>More molecular species still single-layer Earth atmosphere, now for $H_2 O$, $CO_2$, and $O_2$ Using HITRAN, `jax`, and NVIDIA GPUs<jupyter_code>import astropy.units as u
import matplotlib.pyplot as plt

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

import hapi
hapi.db_begin('../../hapi/data/')

from jax.config import config
config.update("jax_enable_x64", False)

import jax.numpy as np
from jax import vmap, jit, grad, random
from jax.lax import stop_gradient
from jax import jacfwd
from functools import partial

import pandas as pd
#hapi.fetch('../../hapi/data/CO2',2,1,3900,12700)
@jit
def gamma_of_p_and_T(p, T, p_self, n_air, gamma_air_ref, gamma_self_ref):
    '''Compute the Lorentz HWHM with pressure and temp'''
    return (296.0/T)**n_air *(gamma_air_ref*(p-p_self) + gamma_self_ref*(p_self))
@jit
def lorentz_profile(nu, p, nu_ij, gamma, dp_ref, S_ij):
    '''Return the Lorentz line profile given vectors and parameters'''
    return S_ij/np.pi * gamma / ( gamma**2 + (nu - (nu_ij + dp_ref*p))**2)
@jit
def tips_Q_of_T(T, g_k, E_k):
    '''Total Internal Partition Sum'''
    c_2 = 1.4387770 #cm K
    return np.sum( g_k * np.exp(- c_2 * E_k / T) )
@jit
def S_ij_of_T(T, S_ij_296, nu_ij, g_lower, E_lower):
    '''The Spectral Line Intensity'''
    c_2 = 1.4387770 #cm K
    return (S_ij_296 * 
            stop_gradient(tips_Q_of_T(296.0, g_lower, E_lower)) / 
            tips_Q_of_T(T, g_lower, E_lower) *
            np.exp(-c_2 * E_lower/T) /np.exp(-c_2 * E_lower/296.0) *
           (1-np.exp(-c_2*nu_ij/T) )/(1-np.exp(-c_2*nu_ij/296.0)))
@jit
def transmission_of_T_p(T, p, nus, vol_mix_ratio=0.001, path_length_km=3.0, atomic_data=None):
    '''Return the absorption coefficient as a function of T'''
    
    (n_air, gamma_air, gamma_self, E_lower, g_lower,\
     nu_lines, delta_air, S_ij_296) = atomic_data
    
    gammas = vmap(gamma_of_p_and_T, in_axes=(None, None, None, 0, 0, 0)\
                                )(p, T,vol_mix_ratio, n_air, gamma_air, gamma_self)
    
    S_ij = S_ij_of_T(T, S_ij_296, nu_lines, g_lower, E_lower)
    
    abs_coeff = vmap(lorentz_profile, in_axes=(None, None, 0, 0,0, 0))(nus, p,
                                     nu_lines, gammas, delta_air, S_ij).sum(axis=0)
    
    tau = abs_coeff * (vol_mix_ratio*2.688392857142857e+19) * (path_length_km*100000.0)
    return np.exp(-tau)<jupyter_output><empty_output><jupyter_text>Compute the transmission spectrumRequires a path length.  Let's assume constant density of the Earth's atmosphere, and 3 km path length.The HITRAN absorption coefficient $\alpha$ units are: $cm^{−1}/(molecule⋅cm^{−2})$So we have the optical depth:$$ \tau_{ij}(\nu, T, p) = \alpha \cdot n \cdot \ell$$where $n$ is the number density of particles, and $\ell$ is the path length.<jupyter_code>def get_hapi_molec_data(species):
    '''return the order columns in device arrays
    TODO: a bit fragile, consider returning a dict instead!
    '''
    ordered_cols = ['n_air','gamma_air','gamma_self','elower','gpp','nu','delta_air','sw']
    return [np.array(hapi.getColumn(species, col)) for col in ordered_cols]
#(n_air, gamma_air, gamma_self, E_lower, g_lower, nu_lines, delta_air, S_ij_296) = get_hapi_molec_data('H2O')
atomic_data = get_hapi_molec_data('O2')
df = pd.read_csv('../../ucdwhpf/data/lblrtm_hpf_50humidity.csv')
nus = np.array((df.wl.values*u.Angstrom).to(1/u.cm, equivalencies=u.spectral()).value)
nus.shape
trans = transmission_of_T_p(291.0, 0.97, nus, vol_mix_ratio=0.21,atomic_data=atomic_data)<jupyter_output><empty_output><jupyter_text>Read in a demo LBLRTM spectrum from TelFit<jupyter_code>plt.figure(figsize=(18,4))
#plt.plot(df.wl, df.trans, label='TelFit')
plt.plot(df.wl, transmission_of_T_p(291.0, 0.97, nus, vol_mix_ratio=0.01,
                                    atomic_data=get_hapi_molec_data('H2O')),label='H$_2$O')
plt.plot(df.wl, transmission_of_T_p(291.0, 0.97, nus, vol_mix_ratio=0.21,
                                    atomic_data=get_hapi_molec_data('O2')),label='O$_2$')
plt.plot(df.wl, transmission_of_T_p(291.0, 0.97, nus, vol_mix_ratio=365e-6,
                                    atomic_data=get_hapi_molec_data('CO2')),label='CO$_2$')

plt.legend(loc='lower left')
plt.ylabel('Transmission')
plt.xlabel('$\lambda \;(\AA)$')
plt.title('Single 3 km layer');
plt.figure(figsize=(18,4))
plt.plot(df.wl, df.trans, label='TelFit', color='k');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.147931654676259, 'fraction_numerical': 0.039568345323741004, 'mean_word_length': 4.010135135135135, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '479102', 'license': 'MIT', 'n_tokens_mistral': 1896, 'n_tokens_neox': 1764, 'n_words': 367, 'path': 'notebooks/05_more_molecular_species.ipynb', 'repo_name': 'BrownDwarf/fiatlux', 'size': '4448'}