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
56 stars 14 forks source link

Fix opachord.py to avoid nan in the gradient of ArtTransPure model #464

Closed sh-tada closed 9 months ago

sh-tada commented 9 months ago

This pull request resolves the issue (#463) of NaN values in the gradient of the ArtTransPure transmission spectra model with respect to temperature, gravity_btm, radius_btm, and mean molecular weight. The NaN values originated from the jnp.sqrt function in exojax/spec/opachord.py.

To address this, I implemented jnp.tril before the jnp.sqrt operation. This modification zeros out irrelevant values, ensuring proper gradient calculations. Setting k=-1 in jnp.tril also zeroes the diagonal.

Fixes #463

sh-tada commented 9 months ago

I added transmission_grad_test.py in integration/unittest_long/transmission/. If necessary, I will change the file name.

HajimeKawahara commented 9 months ago

The newly added file is not pytest file. maybe remove _test from the filename?

sh-tada commented 9 months ago

I modified the filename and Nx.