Added chi_eff functions
Resolves #10 (closed).
This adopts several of the functions from Iwaya et al. in DCC P2400395. Coded up the equations in that paper in a way that can be handled with jax.
data_collection.py:
- We now read in "chi_eff" and "chi_p" from the PE files.
- We now include "chi_eff" and "chi_p" from the injections if these parameters are in
args.parameters - If both "chi_eff" and "chi_p" are in
args.parametersdo this using the jointchi_eff-chi-pjacobian. This is for both the PE samples and the injections.
main.py
- if "spin" is a
vt-parameterthen check ifchi-effandchi_pare the specified spin parameters. If not, use the magnitude and tilt models.
analytic_spin_prior.py
- The analytic integrals (I_1, I_2, J_3, I_4) are implemented in a jax-agnostic way under the function
prior_chieff_chip_isotropic, which is@maybe_jitted - These functions call the Dilogarithm/Spence function which is defined under
Li2. Ifjaxis not the active backend, then it uses the scipyspencefunction. Otherwise, I wrote ajaxspenceimplementation ofspencewhich is used, as thejax.scipyversion ofspencedoes not work with complex inputs
There are probably some improvements that can be made, but I think this is in a form that we can start testing. @colm.talbot
Edited by Jacob Golomb