Adds jaxified hyp2f1 function and epl profile with tests #10
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #10 +/- ##
==========================================
Coverage 100.00% 100.00%
==========================================
Files 4 7 +3
Lines 184 388 +204
==========================================
+ Hits 184 388 +204
|
ahuang314
changed the title
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
|
Thanks for the work behind this implementation @ahuang314. Out of curiosity, how does it impact the runtime (compilation time and compiled runtime) for the deflection angle calculation, compared to the iterative version? |
for more information, see https://pre-commit.ci
|
@aymgal Because this hyp2f1 function is more general than the ones implemented in herculens, the compile time is slightly longer. In the specific case of the EPL profile, since -f exp(2 i phi) is always located in the unit disk, both implementations end up using almost the same iterative calculation, so the runtime is mostly unchanged. The herculens implementation using real numbers is still faster by a small margin (it's roughly 1 second faster when comparing one million executions, where they each take around 10 - 11 seconds on my GPU) |
|
Thanks for the details @ahuang314. Given your nice implementation (which seems to be much more general than for the sole use in the EPL profile), it may perhaps make more sense to propose it as a PR to |
|
Unfortunately there are still a few edge cases which need to be addressed. The analytic continuation formula used for 2F1(a,b,c,z) does not work when b - a is an integer, due to the presence of gamma function prefactors gamma(b - a) and gamma(a - b), one of which would have a pole. I have not found a way around this issue. There is a similar problem when c - a - b is an integer. I have found a way around this issue but it is not yet implemented. So there is a bit more work to be done before this can be submitted as a PR to JAX. I may or may not get around to doing this in the future, but for now I have some other tasks I've been neglecting and need to get to. (Note that neither of these issues are relevant in the specific case of the EPL profile, but may potentially be an issue for other profiles) |
There was a problem hiding this comment.
Looks great @ahuang314, I had one question whether you think the limitation of the hyperbolic functions might have impact on real problems when using the EPL profile
| from jax import jit, lax | ||
| from jax.scipy.special import gamma | ||
|
|
||
| # TODO: The analytic continuation formula used in hyp2f1_continuation only works |
There was a problem hiding this comment.
will this lead to some anticipated failures in the use for the EPL function? Just checking
There was a problem hiding this comment.
The EPL profile only uses the hyp2f1_series function, which will not lead to any issues at all. The EPL profile should be 100% free of issues.
Other profiles which involve z outside of the hyp2f1 series' radius of convergence may have some issues since they have to rely on analytic continuation and transformation formulas which may not always work. For example, the gnfw, pseudo double power law, splcore, and uldm profiles may run into issues.
There was a problem hiding this comment.
I see, so these then need to be tested when implemented. So it's ready to merge :)
One goal of this PR is to implement a jaxified version the hyp2f1 function in scipy. One implementation computes the value iteratively using the standard series representation of hyp2f1 which converges on the unit disk in the complex plane (done in herculens for the EPL profile). However, this method is not too reliable when z is near the boundary where the series converges slowly, and also cannot be used when z is outside of the unit disk.
Making use of transformation formulas and analytic continuation, a jaxified version of hyp2f1 (along with a test file) has been implemented in hyp2f1_util in order to allow for the calculation of hyp2f1 anywhere on the complex plane. This can now be used to jaxify any lens profiles which rely on hyp2f1, including the EPL profile, which is also added in this PR along with a test file.
Since the EPL profile relies on the
rotatefunction in util, I've copy-pasted just that part of the util file from lenstronomy into jaxtronomy. The rest of the util file remains to be jaxified (when applicable). Additionally, there is one function from the SPP profile that is used in EPL's mass_3d_lens function. This SPP function has not been jaxified.