HajimeKawahara / exojax

🐈 Automatic differentiable spectrum modeling of exoplanets/brown dwarfs using JAX, compatible with NumPyro and JAXopt
http://secondearths.sakura.ne.jp/exojax/
MIT License
45 stars 14 forks source link

ExoJAX

License Docs arxiv paper

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

ExoJAX Classes - 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)

Get Started

See this page for the first step!

Functions

Voigt Profile :heavy_check_mark: ```python3 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 ```
Cross Section using HITRAN/HITEMP/ExoMol :heavy_check_mark: ```python 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 ```
Do you just want to plot the line strength at T=1000K? ```python mdb.change_reference_temperature(1000.) # at 1000K plt.plot(mdb.nu_lines,mdb.line_strength_ref,".") ```
Emission Spectrum :heavy_check_mark: ```python art = ArtEmisPure(nu_grid=nu_grid, pressure_btm=1.e2, pressure_top=1.e-8, nlayer=100) F = art.run(dtau, Tarr) ```
Transmission Spectrum :heavy_check_mark:
Reflection Spectrum :heavy_check_mark:

Installation

pip install exojax

or

python setup.py install
Note on installation w/ GPU support :books: You need to install CUDA, JAX w/ NVIDIA GPU support. Visit [here](https://github.com/google/jax) for the installation of GPU supported JAX.

References

paper

License

🐈 Copyright 2020-2024 ExoJAX contributors. ExoJAX is publicly available under the MIT license.