Added chi_eff functions
Resolves #10.
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.parameters
do this using the jointchi_eff-chi-p
jacobian. This is for both the PE samples and the injections.
main.py
- if "spin" is a
vt-parameter
then check ifchi-eff
andchi_p
are 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_jit
ted - These functions call the Dilogarithm/Spence function which is defined under
Li2
. Ifjax
is not the active backend, then it uses the scipyspence
function. Otherwise, I wrote ajaxspence
implementation ofspence
which is used, as thejax.scipy
version ofspence
does 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