galax
galax copied to clipboard
Add Zhao potential model
This adds the (alpha,beta,gamma) Zhao (1996) double power law model. This replaces #755. Fixes #726.
~~The tests are currently failing but will pass once #760 is merged.~~
From @jnibauer:
Sorry missed this earlier. On the Zhao double power-law model – are force evaluations super slow using this model? I had tried to implement it but gradients of the incomplete Beta-function were annoyingly slow...
I haven't run any benchmarks yet, but step 1 is to get a working model in and then we can see if there are any optimizations we can make!
Also I didn't export this to the public API in case we get #757 in.
It is slow! Some benchmarks:
Codecov Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 95.90%. Comparing base (f99ac8d) to head (578e535).
Additional details and impacted files
@@ Coverage Diff @@
## main #761 +/- ##
==========================================
+ Coverage 95.84% 95.90% +0.06%
==========================================
Files 158 159 +1
Lines 5963 6053 +90
==========================================
+ Hits 5715 5805 +90
Misses 248 248
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
:rocket: New features to boost your workflow:
- :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
The latest commits make a conceptual change to the potential model after I realized there was a slight bug in the previous math. The mass parameter can't be total mass because models with beta <= 3 do not have finite mass -- that's fine and was working before. But: The density normalization (rho0 in my notation, C in Zhao's notation) was failing for beta<=3, causing the model to return nan for those cases. After a bit of a rabbit hole, long story short, this is because scipy.special.betainc (and so jax.scipy.special.betainc) computes the regularized incomplete beta function, which becomes nan because of the evaluation of the complete beta function for the arguments set by the power law indices with beta<=3. I didn't come up with a straightforward way of setting the normalization such that it is continuous through the beta=3 pole such that mass can mean "total mass" on one side and "scale mass" on the other. So instead, I opted to change the meaning of the mass parameter: it now means "mass enclosed within the scale radius," which of course is well defined for any model with valid power-law indices. I also added an additional constructor to create a model from a total mass (when using beta > 3).
@adrn we should make a benchmark suite.