potential solution that reduces divergences in fbased rt

[HajimeKawahara]

I found that the change of sum to nansum can avoid nan values of the fbased computation in some cases (See "File changed"). Can you check if this change reduces the number of divergences in your case? @ykawashima

[github-actions[bot]]

Unit Test Results

139 tests  ±0   139 :heavy_check_mark: ±0   33s :stopwatch: +2s     1 suites ±0       0 :zzz: ±0      1 files   ±0       0 :x: ±0 

Results for commit 36e7b479. ± Comparison against base commit f4c4d954.

:recycle: This comment has been updated with latest results.

[ykawashima]

Thank you! Just for your information, the computational time for the spectrum (in my setting) with the current develop and fbased_nansum branches (so just the difference of sum and nansum ) are 0.0271 $\pm$ 0.0015 and 0.0270 $\pm$ 0.0020 [s] (10 times calculations).

[ykawashima]

It seems there exists a ~19% difference in the computational time for the first 57 samplings in the warm-up phase (nansum calculation is faster). (Previously, I ran exactly the same hmc with different GPUs (both A100) and confirmed that the difference in the computation time was less than 0.1%.)

[ykawashima]

Just for your information, HMC (500 warm-up steps + 1000 samples, H2O + CIA) for 2 Subaru/IRD orders dominated by H2O absorption took 3.0 hr (33,571 total steps) 3.6 hr (40,921 total steps) and 3.5 hr (40,921 total steps) for without and with nansum calculations, respectively (cf. 2.7 hr and 29,855 total steps for 8-stream intensity-base case). On the other hand, HMC (500 warm-up steps + 1000 samples, H2O, CH4, NH3 + CIA) for 2 Subaru/IRD orders dominated by H2O/CH4/CIA absorption took 118 hr (398,410 total steps) 66 hr (227,615 total steps) and 66 hr (227,615 total steps) for without and with nansum calculations, respectively (cf. 28 hr and 93,022 total steps for 8-stream intensity-base case).

• comments with strikethrough are for the calculations with the older jax version

[HajimeKawahara]

Thanks! So, the replacement of sum to nansum looks useful. Oh, the intensity-based 8 stream is much faster than the flux-based one for the latter case!

[ykawashima]

Yes, it really does for the CIA region!

[ykawashima]

I'm sorry. I realized that I used different accounts for the calculations of the cases of sum and nansum, so the computational time differences might come from the different versions of packages. The difference for exojax is just between sum and nansum, and the numpyro version is the same, but for example, the jax version is different.

I re-ran HMC for 2 Subaru/IRD orders dominated by H2O absorption using the same account used for the case with nansum (I have updated the above post). The calculation for the CIA region will take another 4 days or so. I will let you know once it finishes.

[ykawashima]

Ah, sorry. The difference in computational time indeed came from the different package versions. Given that there is no difference between nansum and sum, it can be just kept. Sorry for the confusion.

[ykawashima]

The faster calculation uses jax 0.4.16, while the slower calculation uses jax 0.4.8. It might be better to upgrade jax.

[HajimeKawahara]

oh, good to know!