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

AmClouds is not differentiable yet #494

Closed HajimeKawahara closed 2 weeks ago

HajimeKawahara commented 1 month ago

searchsorted strategy in amclouds.compute_cloud_base_pressure_index does not work. I need to think about the implementation of clouds.

    rg_layer, MMRc = amp_nh3.calc_ammodel(
  File "/home/kawahara/exojax/src/exojax/atm/atmphys.py", line 111, in calc_ammodel
    pressure_base = pressures[ibase]
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape int32[].
The error occurred while tracing the function update at /home/kawahara/anaconda3/lib/python3.10/site-packages/jaxopt/_src/optax_wrapper.py:121 for jit. This value became a tracer due to JAX operations on these lines:
HajimeKawahara commented 1 month ago

searchsorted

It gives null derivative for x.

tests/unittests/atm/searchsorted_test.py

HajimeKawahara commented 1 month ago

bug: miegrid was defined as complex64, should be float.

HajimeKawahara commented 2 weeks ago

it looks OK now.

HajimeKawahara commented 1 week ago

518