Skip to content

Added chi_eff functions

Jacob Golomb requested to merge jacob.golomb/gwpopulation_pipe:chi_eff into master

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 for gaussian_kde, which is called by the chi_p prior. So these functions are done on the CPU by explicitly calling to_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 by data_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 import xp now that the backend is set by set_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 to chi_eff. If both "chi_eff" and "chi_p" are in args.parameters do this using the joint chi_eff-chi-p jacobian. This is for both the PE samples and the injections.

main.py

  • if "spin" is a vt-parameter then check if chi-eff and chi_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

Merge request reports