Added chi_eff functions
Resolves #10.
This adopts several of the functions from Tom Callister's effective-spin-priors repository, as described in arXiv:2104.09508. I edited the functions significantly for efficiency and compatibility with gwpopulation. Here is a description of some of the key points:
spin_conversions.py:
- This is a new file that has the utils for doing the conversions. This could reasonably live in gwpopulation instead.
- There is some dependence on the
spence
function from scipy.special that does not yet have a cupy implementation (but does have a jax implementation). Similar story forgaussian_kde
, which is called by thechi_p
prior. So these functions are done on the CPU by explicitly callingto_numpy
on the arrays they will be acting on. I think this is generally being done on the CPU anyway since the functions in this script are likely being called bydata_collection
which already sets the backend to numpy. I still think this catch should be in here in case it is called by a user at some stage in which the arrays are on a GPU. - Is
from gwpopulation.hyperpe import xp
still the preferred way to importxp
now that the backend is set byset_backend
? If not I can change this line. - There are lots of functions that need to be applied conditionally. The best way I could think of doing this is by storing these functions in dicts whose keys are the same keys as the mask. I'm choosing to apply the functions conditionally rather than apply them to all the cases and mask afterward, as applying them conditionally may avoid any issues with
jax
encountering an inf or nan before a value is masked out. This is because several of the cases may return invalid values for the cases they are not intended for. - These are permanent lambda expressions so we can store them outside the functions that they are called by
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 "chi_eff" is in
args.parameters
, apply the correct jacobian to convert from isotropic spin prior tochi_eff
. If both "chi_eff" and "chi_p" are inargs.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. 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