Awesome

ExoJAX

Differentiable spectral modelling of exoplanets/brown dwarfs/M dwarfs using JAX! Read the docs 🐕. In short, ExoJAX allows you to do gradient based optimizations and HMC-NUTS samplings using the latest database.

ExoJAX is at least compatible with

PPLs: NumPyro, blackjax
Optimizers: JAXopt, optax, bayeux

<img src="https://github.com/HajimeKawahara/exojax/assets/15956904/8aa9673b-b64b-4b65-a76c-2966ef1edbc7" Titie="exojax" Width=850px> <details><summary>ExoJAX Classes</summary>

Databases: *db (mdb: molecular, adb: atomic, cdb:continuum, pdb: particulates)
Opacity Calculators: opa (i.e. Voigt profile)
Atmospheric Radiative Transfer: art (emission w, w/o scattering, refelction, transmission)
Atompsheric Microphysics: amp (clouds etc)

</details>

Get Started

See this page for the first step!

Functions

<details open><summary>Voigt Profile :heavy_check_mark: </summary>

from exojax.spec import voigt
nu=numpy.linspace(-10,10,100)
voigt(nu,1.0,2.0) #sigma_D=1.0, gamma_L=2.0

</details> <details><summary>Cross Section using HITRAN/HITEMP/ExoMol :heavy_check_mark: </summary>

from exojax.utils.grids import wavenumber_grid
from exojax.spec.api import MdbExomol
from exojax.spec.opacalc import OpaPremodit
from jax import config
config.update("jax_enable_x64", True)

nu_grid,wav,res=wavenumber_grid(1900.0,2300.0,200000,xsmode="premodit",unit="cm-1",)
mdb = MdbExomol(".database/CO/12C-16O/Li2015",nu_grid)
opa = OpaPremodit(mdb,nu_grid,auto_trange=[900.0,1100.0])
xsv = opa.xsvector(1000.0, 1.0) # cross section for 1000K, 1 bar

<img src="https://user-images.githubusercontent.com/15956904/111430765-2eedf180-873e-11eb-9740-9e1a313d590c.png" Titie="exojax auto cross section" Width=850px> </details> <details><summary>Do you just want to plot the line strength at T=1000K? </summary>

mdb.change_reference_temperature(1000.) # at 1000K
plt.plot(mdb.nu_lines,mdb.line_strength_ref,".")

</details> <details><summary>Emission Spectrum :heavy_check_mark: </summary>

art = ArtEmisPure(nu_grid=nu_grid, pressure_btm=1.e2, pressure_top=1.e-8, nlayer=100)
F = art.run(dtau, Tarr)

<img src="https://user-images.githubusercontent.com/15956904/116488770-286ea000-a8ce-11eb-982d-7884b423592c.png" Titie="exojax auto \emission spectrum" Width=850px> </details> <details><summary>Transmission Spectrum :heavy_check_mark: </summary></details> <details><summary>Reflection Spectrum :heavy_check_mark: </summary></details>

Installation

pip install exojax

python setup.py install

<details><summary>Note on installation w/ GPU support</summary>

:books: You need to install CUDA, JAX w/ NVIDIA GPU support.

Visit here for the installation of GPU supported JAX.