HealPixPrior (Resloves Issue #419)

Hi all, This MR introduces a HealPixMapPriorDist object to be used to create a 2D or 3D joint prior in ra/dec (ra/dec/distance) according to a passed in HealPix map or .fits file to address and resolve issue #419 (closed).

This PriorObject inherits from BaseJointPriorDist and needs to overwrite three methods (_rescale(), _sample(), and _lnprob()) To illustrate this is working I sample 5000 samples from the prior and show below that these samples are distribvuted according to whats in the HealPix Map .fits file.

Here I show that the _rescale() method is also working by generating 5000 samples uniformly on [0,1] then rescale to the HealPixPrior which also matches with what is in the map file.

bilby/gw/prior.py

@@ -5,24 +5,39 @@ import numpy as np
 from scipy.interpolate import InterpolatedUnivariateSpline, interp1d
 from scipy.stats import norm
 from scipy.integrate import cumtrapz
-
-from ..core.prior import (PriorDict, Uniform, Prior, DeltaFunction, Gaussian,
-                          Interped, Constraint, conditional_prior_factory, BaseJointPriorDist, JointPrior,
-                          JointPriorDistError)
+from scipy.special import erfi
+
+from ..core.prior import (
+    PriorDict,
+    Uniform,
+    Prior,
+    DeltaFunction,
+    Gaussian,
+    Interped,
+    Constraint,
+    conditional_prior_factory,
+    BaseJointPriorDist,
+    JointPrior,
+    JointPriorDistError,
+)
 from ..core.utils import infer_args_from_method, logger
 from .conversion import (
     convert_to_lal_binary_black_hole_parameters,
-    convert_to_lal_binary_neutron_star_parameters, generate_mass_parameters,
-    generate_tidal_parameters, fill_from_fixed_priors,
+    convert_to_lal_binary_neutron_star_parameters,
+    generate_mass_parameters,
+    generate_tidal_parameters,
+    fill_from_fixed_priors,
     chirp_mass_and_mass_ratio_to_total_mass,
-    total_mass_and_mass_ratio_to_component_masses)
+    total_mass_and_mass_ratio_to_component_masses,
+)
 from .cosmology import get_cosmology
 
 try:
     from astropy import cosmology as cosmo, units
 except ImportError:
-    logger.debug("You do not have astropy installed currently. You will"
-                 " not be able to use some of the prebuilt functions.")
+    logger.debug(
+        "You do not have astropy installed currently. You will" " not be able to use some of the prebuilt functions."
+    )
 
 
 class BilbyPriorConversionError(Exception):
@@ -45,12 +60,12 @@ def convert_to_flat_in_component_mass_prior(result, fraction=0.25):
 
     """
     if getattr(result, "priors") is not None:
-        for key in ['chirp_mass', 'mass_ratio']:
+        for key in ["chirp_mass", "mass_ratio"]:
             if key not in result.priors.keys():
                 BilbyPriorConversionError("{} Prior not found in result object".format(key))
             if isinstance(result.priors[key], Constraint):
                 BilbyPriorConversionError("{} Prior should not be a Constraint".format(key))
-        for key in ['mass_1', 'mass_2']:
+        for key in ["mass_1", "mass_2"]:
             if not isinstance(result.priors[key], Constraint):
                 BilbyPriorConversionError("{} Prior should be a Constraint Prior".format(key))
     else:
@@ -61,27 +76,30 @@ def convert_to_flat_in_component_mass_prior(result, fraction=0.25):
     old_priors = copy.copy(result.priors)
     posterior = result.posterior
 
-    for key in ['chirp_mass', 'mass_ratio']:
+    for key in ["chirp_mass", "mass_ratio"]:
         priors[key] = Constraint(priors[key].minimum, priors[key].maximum, key, latex_label=priors[key].latex_label)
-    for key in ['mass_1', 'mass_2']:
-        priors[key] = Uniform(priors[key].minimum, priors[key].maximum, key, latex_label=priors[key].latex_label,
-                              unit="$M_{\odot}$")
+    for key in ["mass_1", "mass_2"]:
+        priors[key] = Uniform(
+            priors[key].minimum, priors[key].maximum, key, latex_label=priors[key].latex_label, unit="$M_{\odot}$"
+        )
 
-    weights = np.array(result.get_weights_by_new_prior(old_priors, priors,
-                                                       prior_names=['chirp_mass', 'mass_ratio', 'mass_1', 'mass_2']))
+    weights = np.array(
+        result.get_weights_by_new_prior(
+            old_priors, priors, prior_names=["chirp_mass", "mass_ratio", "mass_1", "mass_2"]
+        )
+    )
     jacobian = posterior["mass_1"] ** 2 / posterior["chirp_mass"]
     weights = jacobian * weights
     result.posterior = posterior.sample(frac=fraction, weights=weights)
 
     logger.info("Resampling posterior to flat-in-component mass")
-    effective_sample_size = sum(weights)**2 / sum(weights**2)
+    effective_sample_size = sum(weights) ** 2 / sum(weights ** 2)
     n_posterior = len(posterior)
     if fraction > effective_sample_size / n_posterior:
         logger.warning(
             "Sampling posterior of length {} with fraction {}, but "
             "effective_sample_size / len(posterior) = {}. This may produce "
-            "biased results"
-            .format(n_posterior, fraction, effective_sample_size / n_posterior)
+            "biased results".format(n_posterior, fraction, effective_sample_size / n_posterior)
         )
     result.posterior = posterior.sample(frac=fraction, weights=weights, replace=True)
     result.meta_data["reweighted_to_flat_in_component_mass"] = True
@@ -89,46 +107,45 @@ def convert_to_flat_in_component_mass_prior(result, fraction=0.25):
 
 
 class Cosmological(Interped):
-
     @property
     def _default_args_dict(self):
         return dict(
-            redshift=dict(name='redshift', latex_label='$z$', unit=None),
-            luminosity_distance=dict(
-                name='luminosity_distance', latex_label='$d_L$', unit=units.Mpc),
-            comoving_distance=dict(
-                name='comoving_distance', latex_label='$d_C$', unit=units.Mpc))
-
-    def __init__(self, minimum, maximum, cosmology=None, name=None,
-                 latex_label=None, unit=None, boundary=None):
+            redshift=dict(name="redshift", latex_label="$z$", unit=None),
+            luminosity_distance=dict(name="luminosity_distance", latex_label="$d_L$", unit=units.Mpc),
+            comoving_distance=dict(name="comoving_distance", latex_label="$d_C$", unit=units.Mpc),
+        )
+
+    def __init__(self, minimum, maximum, cosmology=None, name=None, latex_label=None, unit=None, boundary=None):
         self.cosmology = get_cosmology(cosmology)
         if name not in self._default_args_dict:
             raise ValueError(
                 "Name {} not recognised. Must be one of luminosity_distance, "
-                "comoving_distance, redshift".format(name))
+                "comoving_distance, redshift".format(name)
+            )
         self.name = name
         label_args = self._default_args_dict[self.name]
         if latex_label is not None:
-            label_args['latex_label'] = latex_label
+            label_args["latex_label"] = latex_label
         if unit is not None:
             if not isinstance(unit, units.Unit):
                 unit = units.Unit(unit)
-            label_args['unit'] = unit
-        self.unit = label_args['unit']
+            label_args["unit"] = unit
+        self.unit = label_args["unit"]
         self._minimum = dict()
         self._maximum = dict()
         self.minimum = minimum
         self.maximum = maximum
-        if name == 'redshift':
+        if name == "redshift":
             xx, yy = self._get_redshift_arrays()
-        elif name == 'comoving_distance':
+        elif name == "comoving_distance":
             xx, yy = self._get_comoving_distance_arrays()
-        elif name == 'luminosity_distance':
+        elif name == "luminosity_distance":
             xx, yy = self._get_luminosity_distance_arrays()
         else:
-            raise ValueError('Name {} not recognized.'.format(name))
-        super(Cosmological, self).__init__(xx=xx, yy=yy, minimum=minimum, maximum=maximum,
-                                           boundary=boundary, **label_args)
+            raise ValueError("Name {} not recognized.".format(name))
+        super(Cosmological, self).__init__(
+            xx=xx, yy=yy, minimum=minimum, maximum=maximum, boundary=boundary, **label_args
+        )
 
     @property
     def minimum(self):
@@ -138,25 +155,21 @@ class Cosmological(Interped):
     def minimum(self, minimum):
         cosmology = get_cosmology(self.cosmology)
         self._minimum[self.name] = minimum
-        if self.name == 'redshift':
-            self._minimum['luminosity_distance'] =\
-                cosmology.luminosity_distance(minimum).value
-            self._minimum['comoving_distance'] =\
-                cosmology.comoving_distance(minimum).value
-        elif self.name == 'luminosity_distance':
+        if self.name == "redshift":
+            self._minimum["luminosity_distance"] = cosmology.luminosity_distance(minimum).value
+            self._minimum["comoving_distance"] = cosmology.comoving_distance(minimum).value
+        elif self.name == "luminosity_distance":
             if minimum == 0:
-                self._minimum['redshift'] = 0
+                self._minimum["redshift"] = 0
             else:
-                self._minimum['redshift'] = cosmo.z_at_value(
-                    cosmology.luminosity_distance, minimum * self.unit)
-            self._minimum['comoving_distance'] = self._minimum['redshift']
-        elif self.name == 'comoving_distance':
+                self._minimum["redshift"] = cosmo.z_at_value(cosmology.luminosity_distance, minimum * self.unit)
+            self._minimum["comoving_distance"] = self._minimum["redshift"]
+        elif self.name == "comoving_distance":
             if minimum == 0:
-                self._minimum['redshift'] = 0
+                self._minimum["redshift"] = 0
             else:
-                self._minimum['redshift'] = cosmo.z_at_value(
-                    cosmology.comoving_distance, minimum * self.unit)
-            self._minimum['luminosity_distance'] = self._minimum['redshift']
+                self._minimum["redshift"] = cosmo.z_at_value(cosmology.comoving_distance, minimum * self.unit)
+            self._minimum["luminosity_distance"] = self._minimum["redshift"]
         try:
             self._update_instance()
         except (AttributeError, KeyError):
@@ -170,19 +183,15 @@ class Cosmological(Interped):
     def maximum(self, maximum):
         cosmology = get_cosmology(self.cosmology)
         self._maximum[self.name] = maximum
-        if self.name == 'redshift':
-            self._maximum['luminosity_distance'] = \
-                cosmology.luminosity_distance(maximum).value
-            self._maximum['comoving_distance'] = \
-                cosmology.comoving_distance(maximum).value
-        elif self.name == 'luminosity_distance':
-            self._maximum['redshift'] = cosmo.z_at_value(
-                cosmology.luminosity_distance, maximum * self.unit)
-            self._maximum['comoving_distance'] = self._maximum['redshift']
-        elif self.name == 'comoving_distance':
-            self._maximum['redshift'] = cosmo.z_at_value(
-                cosmology.comoving_distance, maximum * self.unit)
-            self._maximum['luminosity_distance'] = self._maximum['redshift']
+        if self.name == "redshift":
+            self._maximum["luminosity_distance"] = cosmology.luminosity_distance(maximum).value
+            self._maximum["comoving_distance"] = cosmology.comoving_distance(maximum).value
+        elif self.name == "luminosity_distance":
+            self._maximum["redshift"] = cosmo.z_at_value(cosmology.luminosity_distance, maximum * self.unit)
+            self._maximum["comoving_distance"] = self._maximum["redshift"]
+        elif self.name == "comoving_distance":
+            self._maximum["redshift"] = cosmo.z_at_value(cosmology.comoving_distance, maximum * self.unit)
+            self._maximum["luminosity_distance"] = self._maximum["redshift"]
         try:
             self._update_instance()
         except (AttributeError, KeyError):
@@ -193,13 +202,13 @@ class Cosmological(Interped):
         args_dict = {key: getattr(self, key) for key in subclass_args}
         self._convert_to(new=name, args_dict=args_dict)
         if unit is not None:
-            args_dict['unit'] = unit
+            args_dict["unit"] = unit
         return self.__class__(**args_dict)
 
     def _convert_to(self, new, args_dict):
         args_dict.update(self._default_args_dict[new])
-        args_dict['minimum'] = self._minimum[args_dict['name']]
-        args_dict['maximum'] = self._maximum[args_dict['name']]
+        args_dict["minimum"] = self._minimum[args_dict["name"]]
+        args_dict["maximum"] = self._maximum[args_dict["name"]]
 
     def _get_comoving_distance_arrays(self):
         zs, p_dz = self._get_redshift_arrays()
@@ -222,8 +231,7 @@ class Cosmological(Interped):
     def from_repr(cls, string):
         if "FlatLambdaCDM" in string:
             logger.warning(
-                "Cosmological priors cannot be loaded from a string. "
-                "If the prior has a name, use that instead."
+                "Cosmological priors cannot be loaded from a string. " "If the prior has a name, use that instead."
             )
             return string
         else:
@@ -235,19 +243,17 @@ class Cosmological(Interped):
         Get a dictionary containing the arguments needed to reproduce this object.
         """
         dict_with_properties = super(Cosmological, self)._repr_dict
-        if isinstance(dict_with_properties['cosmology'], cosmo.core.Cosmology):
-            if dict_with_properties['cosmology'].name is not None:
-                dict_with_properties['cosmology'] = dict_with_properties['cosmology'].name
-        if isinstance(dict_with_properties['unit'], units.Unit):
-            dict_with_properties['unit'] = dict_with_properties['unit'].to_string()
+        if isinstance(dict_with_properties["cosmology"], cosmo.core.Cosmology):
+            if dict_with_properties["cosmology"].name is not None:
+                dict_with_properties["cosmology"] = dict_with_properties["cosmology"].name
+        if isinstance(dict_with_properties["unit"], units.Unit):
+            dict_with_properties["unit"] = dict_with_properties["unit"].to_string()
         return dict_with_properties
 
 
 class UniformComovingVolume(Cosmological):
-
     def _get_redshift_arrays(self):
-        zs = np.linspace(self._minimum['redshift'] * 0.99,
-                         self._maximum['redshift'] * 1.01, 1000)
+        zs = np.linspace(self._minimum["redshift"] * 0.99, self._maximum["redshift"] * 1.01, 1000)
         p_dz = self.cosmology.differential_comoving_volume(zs).value
         return zs, p_dz
 
@@ -262,16 +268,15 @@ class UniformSourceFrame(Cosmological):
     """
 
     def _get_redshift_arrays(self):
-        zs = np.linspace(self._minimum['redshift'] * 0.99,
-                         self._maximum['redshift'] * 1.01, 1000)
+        zs = np.linspace(self._minimum["redshift"] * 0.99, self._maximum["redshift"] * 1.01, 1000)
         p_dz = self.cosmology.differential_comoving_volume(zs).value / (1 + zs)
         return zs, p_dz
 
 
 class AlignedSpin(Interped):
-
-    def __init__(self, a_prior=Uniform(0, 1), z_prior=Uniform(-1, 1),
-                 name=None, latex_label=None, unit=None, boundary=None):
+    def __init__(
+        self, a_prior=Uniform(0, 1), z_prior=Uniform(-1, 1), name=None, latex_label=None, unit=None, boundary=None
+    ):
         """
         Prior distribution for the aligned (z) component of the spin.
 
@@ -294,19 +299,17 @@ class AlignedSpin(Interped):
         """
         self.a_prior = a_prior
         self.z_prior = z_prior
-        chi_min = min(a_prior.maximum * z_prior.minimum,
-                      a_prior.minimum * z_prior.maximum)
+        chi_min = min(a_prior.maximum * z_prior.minimum, a_prior.minimum * z_prior.maximum)
         chi_max = a_prior.maximum * z_prior.maximum
         xx = np.linspace(chi_min, chi_max, 800)
         a_prior_minimum = a_prior.minimum
         if a_prior_minimum == 0:
             a_prior_minimum += 1e-32
         aas = np.linspace(a_prior_minimum, a_prior.maximum, 1000)
-        yy = [np.trapz(np.nan_to_num(a_prior.prob(aas) / aas *
-                                     z_prior.prob(x / aas)), aas) for x in xx]
-        super(AlignedSpin, self).__init__(xx=xx, yy=yy, name=name,
-                                          latex_label=latex_label, unit=unit,
-                                          boundary=boundary)
+        yy = [np.trapz(np.nan_to_num(a_prior.prob(aas) / aas * z_prior.prob(x / aas)), aas) for x in xx]
+        super(AlignedSpin, self).__init__(
+            xx=xx, yy=yy, name=name, latex_label=latex_label, unit=unit, boundary=boundary
+        )
 
 
 class CBCPriorDict(PriorDict):
@@ -317,9 +320,9 @@ class CBCPriorDict(PriorDict):
         if "chirp_mass" in self:
             return self["chirp_mass"].minimum
         elif "mass_1" in self:
-            mass_1 = self['mass_1'].minimum
+            mass_1 = self["mass_1"].minimum
             if "mass_2" in self:
-                mass_2 = self['mass_2'].minimum
+                mass_2 = self["mass_2"].minimum
             elif "mass_ratio" in self:
                 mass_2 = mass_1 * self["mass_ratio"].minimum
         if mass_1 is not None and mass_2 is not None:
@@ -336,9 +339,9 @@ class CBCPriorDict(PriorDict):
         if "chirp_mass" in self:
             return self["chirp_mass"].maximum
         elif "mass_1" in self:
-            mass_1 = self['mass_1'].maximum
+            mass_1 = self["mass_1"].maximum
             if "mass_2" in self:
-                mass_2 = self['mass_2'].maximum
+                mass_2 = self["mass_2"].maximum
             elif "mass_ratio" in self:
                 mass_2 = mass_1 * self["mass_ratio"].maximum
         if mass_1 is not None and mass_2 is not None:
@@ -353,10 +356,8 @@ class CBCPriorDict(PriorDict):
         if "mass_2" in self:
             return self["mass_2"].minimum
         if "chirp_mass" in self and "mass_ratio" in self:
-            total_mass = chirp_mass_and_mass_ratio_to_total_mass(
-                self["chirp_mass"].minimum, self["mass_ratio"].minimum)
-            _, mass_2 = total_mass_and_mass_ratio_to_component_masses(
-                self["mass_ratio"].minimum, total_mass)
+            total_mass = chirp_mass_and_mass_ratio_to_total_mass(self["chirp_mass"].minimum, self["mass_ratio"].minimum)
+            _, mass_2 = total_mass_and_mass_ratio_to_component_masses(self["mass_ratio"].minimum, total_mass)
             return mass_2
         else:
             logger.warning("Unable to determine minimum component mass")
@@ -364,8 +365,7 @@ class CBCPriorDict(PriorDict):
 
 
 class BBHPriorDict(CBCPriorDict):
-    def __init__(self, dictionary=None, filename=None, aligned_spin=False,
-                 conversion_function=None):
+    def __init__(self, dictionary=None, filename=None, aligned_spin=False, conversion_function=None):
         """ Initialises a Prior set for Binary Black holes
 
         Parameters
@@ -379,19 +379,20 @@ class BBHPriorDict(CBCPriorDict):
             By default this generates many additional parameters, see
             BBHPriorDict.default_conversion_function
         """
-        basedir = os.path.join(os.path.dirname(__file__), 'prior_files')
+        basedir = os.path.join(os.path.dirname(__file__), "prior_files")
         if dictionary is None and filename is None:
-            fname = 'binary_black_holes.prior'
+            fname = "binary_black_holes.prior"
             if aligned_spin:
-                fname = 'aligned_spin_' + fname
-                logger.info('Using aligned spin prior')
+                fname = "aligned_spin_" + fname
+                logger.info("Using aligned spin prior")
             filename = os.path.join(basedir, fname)
-            logger.info('No prior given, using default BBH priors in {}.'.format(filename))
+            logger.info("No prior given, using default BBH priors in {}.".format(filename))
         elif filename is not None:
             if not os.path.isfile(filename):
-                filename = os.path.join(os.path.dirname(__file__), 'prior_files', filename)
-        super(BBHPriorDict, self).__init__(dictionary=dictionary, filename=filename,
-                                           conversion_function=conversion_function)
+                filename = os.path.join(os.path.dirname(__file__), "prior_files", filename)
+        super(BBHPriorDict, self).__init__(
+            dictionary=dictionary, filename=filename, conversion_function=conversion_function
+        )
 
     def default_conversion_function(self, sample):
         """
@@ -439,40 +440,43 @@ class BBHPriorDict(CBCPriorDict):
             Whether the key is redundant or not
         """
         if key in self:
-            logger.debug('{} already in prior'.format(key))
+            logger.debug("{} already in prior".format(key))
             return True
 
-        sampling_parameters = {key for key in self if not isinstance(
-            self[key], (DeltaFunction, Constraint))}
-
-        mass_parameters = {'mass_1', 'mass_2', 'chirp_mass', 'total_mass', 'mass_ratio', 'symmetric_mass_ratio'}
-        spin_tilt_1_parameters = {'tilt_1', 'cos_tilt_1'}
-        spin_tilt_2_parameters = {'tilt_2', 'cos_tilt_2'}
-        spin_azimuth_parameters = {'phi_1', 'phi_2', 'phi_12', 'phi_jl'}
-        inclination_parameters = {'theta_jn', 'cos_theta_jn'}
-        distance_parameters = {'luminosity_distance', 'comoving_distance', 'redshift'}
-
-        for independent_parameters, parameter_set in \
-                zip([2, 2, 1, 1, 1, 1],
-                    [mass_parameters, spin_azimuth_parameters,
-                     spin_tilt_1_parameters, spin_tilt_2_parameters,
-                     inclination_parameters, distance_parameters]):
+        sampling_parameters = {key for key in self if not isinstance(self[key], (DeltaFunction, Constraint))}
+
+        mass_parameters = {"mass_1", "mass_2", "chirp_mass", "total_mass", "mass_ratio", "symmetric_mass_ratio"}
+        spin_tilt_1_parameters = {"tilt_1", "cos_tilt_1"}
+        spin_tilt_2_parameters = {"tilt_2", "cos_tilt_2"}
+        spin_azimuth_parameters = {"phi_1", "phi_2", "phi_12", "phi_jl"}
+        inclination_parameters = {"theta_jn", "cos_theta_jn"}
+        distance_parameters = {"luminosity_distance", "comoving_distance", "redshift"}
+
+        for independent_parameters, parameter_set in zip(
+            [2, 2, 1, 1, 1, 1],
+            [
+                mass_parameters,
+                spin_azimuth_parameters,
+                spin_tilt_1_parameters,
+                spin_tilt_2_parameters,
+                inclination_parameters,
+                distance_parameters,
+            ],
+        ):
             if key in parameter_set:
-                if len(parameter_set.intersection(
-                        sampling_parameters)) >= independent_parameters:
+                if len(parameter_set.intersection(sampling_parameters)) >= independent_parameters:
                     logger.disabled = disable_logging
-                    logger.warning('{} already in prior. '
-                                   'This may lead to unexpected behaviour.'
-                                   .format(parameter_set.intersection(self)))
+                    logger.warning(
+                        "{} already in prior. "
+                        "This may lead to unexpected behaviour.".format(parameter_set.intersection(self))
+                    )
                     logger.disabled = False
                     return True
         return False
 
 
 class BNSPriorDict(CBCPriorDict):
-
-    def __init__(self, dictionary=None, filename=None, aligned_spin=True,
-                 conversion_function=None):
+    def __init__(self, dictionary=None, filename=None, aligned_spin=True, conversion_function=None):
         """ Initialises a Prior set for Binary Neutron Stars
 
         Parameters
@@ -487,17 +491,18 @@ class BNSPriorDict(CBCPriorDict):
             BNSPriorDict.default_conversion_function
         """
         if aligned_spin:
-            default_file = 'binary_neutron_stars.prior'
+            default_file = "binary_neutron_stars.prior"
         else:
-            default_file = 'precessing_binary_neutron_stars.prior'
+            default_file = "precessing_binary_neutron_stars.prior"
         if dictionary is None and filename is None:
-            filename = os.path.join(os.path.dirname(__file__), 'prior_files', default_file)
-            logger.info('No prior given, using default BNS priors in {}.'.format(filename))
+            filename = os.path.join(os.path.dirname(__file__), "prior_files", default_file)
+            logger.info("No prior given, using default BNS priors in {}.".format(filename))
         elif filename is not None:
             if not os.path.isfile(filename):
-                filename = os.path.join(os.path.dirname(__file__), 'prior_files', filename)
-        super(BNSPriorDict, self).__init__(dictionary=dictionary, filename=filename,
-                                           conversion_function=conversion_function)
+                filename = os.path.join(os.path.dirname(__file__), "prior_files", filename)
+        super(BNSPriorDict, self).__init__(
+            dictionary=dictionary, filename=filename, conversion_function=conversion_function
+        )
 
     def default_conversion_function(self, sample):
         """
@@ -538,19 +543,18 @@ class BNSPriorDict(CBCPriorDict):
             return True
         redundant = False
 
-        sampling_parameters = {key for key in self if not isinstance(
-            self[key], (DeltaFunction, Constraint))}
+        sampling_parameters = {key for key in self if not isinstance(self[key], (DeltaFunction, Constraint))}
 
-        tidal_parameters = \
-            {'lambda_1', 'lambda_2', 'lambda_tilde', 'delta_lambda_tilde'}
+        tidal_parameters = {"lambda_1", "lambda_2", "lambda_tilde", "delta_lambda_tilde"}
 
         if key in tidal_parameters:
             if len(tidal_parameters.intersection(sampling_parameters)) > 2:
                 redundant = True
                 logger.disabled = disable_logging
-                logger.warning('{} already in prior. '
-                               'This may lead to unexpected behaviour.'
-                               .format(tidal_parameters.intersection(self)))
+                logger.warning(
+                    "{} already in prior. "
+                    "This may lead to unexpected behaviour.".format(tidal_parameters.intersection(self))
+                )
                 logger.disabled = False
             elif len(tidal_parameters.intersection(sampling_parameters)) == 2:
                 redundant = True
@@ -558,38 +562,38 @@ class BNSPriorDict(CBCPriorDict):
 
 
 Prior._default_latex_labels = {
-    'mass_1': '$m_1$',
-    'mass_2': '$m_2$',
-    'total_mass': '$M$',
-    'chirp_mass': '$\mathcal{M}$',
-    'mass_ratio': '$q$',
-    'symmetric_mass_ratio': '$\eta$',
-    'a_1': '$a_1$',
-    'a_2': '$a_2$',
-    'tilt_1': '$\\theta_1$',
-    'tilt_2': '$\\theta_2$',
-    'cos_tilt_1': '$\cos\\theta_1$',
-    'cos_tilt_2': '$\cos\\theta_2$',
-    'phi_12': '$\Delta\phi$',
-    'phi_jl': '$\phi_{JL}$',
-    'luminosity_distance': '$d_L$',
-    'dec': '$\mathrm{DEC}$',
-    'ra': '$\mathrm{RA}$',
-    'iota': '$\iota$',
-    'cos_iota': '$\cos\iota$',
-    'theta_jn': '$\\theta_{JN}$',
-    'cos_theta_jn': '$\cos\\theta_{JN}$',
-    'psi': '$\psi$',
-    'phase': '$\phi$',
-    'geocent_time': '$t_c$',
-    'lambda_1': '$\\Lambda_1$',
-    'lambda_2': '$\\Lambda_2$',
-    'lambda_tilde': '$\\tilde{\\Lambda}$',
-    'delta_lambda_tilde': '$\\delta\\tilde{\\Lambda}$'}
+    "mass_1": "$m_1$",
+    "mass_2": "$m_2$",
+    "total_mass": "$M$",
+    "chirp_mass": "$\mathcal{M}$",
+    "mass_ratio": "$q$",
+    "symmetric_mass_ratio": "$\eta$",
+    "a_1": "$a_1$",
+    "a_2": "$a_2$",
+    "tilt_1": "$\\theta_1$",
+    "tilt_2": "$\\theta_2$",
+    "cos_tilt_1": "$\cos\\theta_1$",
+    "cos_tilt_2": "$\cos\\theta_2$",
+    "phi_12": "$\Delta\phi$",
+    "phi_jl": "$\phi_{JL}$",
+    "luminosity_distance": "$d_L$",
+    "dec": "$\mathrm{DEC}$",
+    "ra": "$\mathrm{RA}$",
+    "iota": "$\iota$",
+    "cos_iota": "$\cos\iota$",
+    "theta_jn": "$\\theta_{JN}$",
+    "cos_theta_jn": "$\cos\\theta_{JN}$",
+    "psi": "$\psi$",
+    "phase": "$\phi$",
+    "geocent_time": "$t_c$",
+    "lambda_1": "$\\Lambda_1$",
+    "lambda_2": "$\\Lambda_2$",
+    "lambda_tilde": "$\\tilde{\\Lambda}$",
+    "delta_lambda_tilde": "$\\delta\\tilde{\\Lambda}$",
+}
 
 
 class CalibrationPriorDict(PriorDict):
-
     def __init__(self, dictionary=None, filename=None):
         """ Initialises a Prior set for Binary Black holes
 
@@ -601,8 +605,7 @@ class CalibrationPriorDict(PriorDict):
             See superclass
         """
         if dictionary is None and filename is not None:
-            filename = os.path.join(os.path.dirname(__file__),
-                                    'prior_files', filename)
+            filename = os.path.join(os.path.dirname(__file__), "prior_files", filename)
         super(CalibrationPriorDict, self).__init__(dictionary=dictionary, filename=filename)
         self.source = None
 
@@ -625,8 +628,7 @@ class CalibrationPriorDict(PriorDict):
                 outfile.write("# prior source file is {}".format(self.source))
 
     @staticmethod
-    def from_envelope_file(envelope_file, minimum_frequency,
-                           maximum_frequency, n_nodes, label):
+    def from_envelope_file(envelope_file, minimum_frequency, maximum_frequency, n_nodes, label):
         """
         Load in the calibration envelope.
 
@@ -660,45 +662,43 @@ class CalibrationPriorDict(PriorDict):
         amplitude_sigma = (calibration_data[5] - calibration_data[3]) / 2
         phase_sigma = (calibration_data[6] - calibration_data[4]) / 2
 
-        log_nodes = np.linspace(np.log(minimum_frequency),
-                                np.log(maximum_frequency), n_nodes)
+        log_nodes = np.linspace(np.log(minimum_frequency), np.log(maximum_frequency), n_nodes)
 
-        amplitude_mean_nodes = \
-            InterpolatedUnivariateSpline(log_frequency_array, amplitude_median)(log_nodes)
-        amplitude_sigma_nodes = \
-            InterpolatedUnivariateSpline(log_frequency_array, amplitude_sigma)(log_nodes)
-        phase_mean_nodes = \
-            InterpolatedUnivariateSpline(log_frequency_array, phase_median)(log_nodes)
-        phase_sigma_nodes = \
-            InterpolatedUnivariateSpline(log_frequency_array, phase_sigma)(log_nodes)
+        amplitude_mean_nodes = InterpolatedUnivariateSpline(log_frequency_array, amplitude_median)(log_nodes)
+        amplitude_sigma_nodes = InterpolatedUnivariateSpline(log_frequency_array, amplitude_sigma)(log_nodes)
+        phase_mean_nodes = InterpolatedUnivariateSpline(log_frequency_array, phase_median)(log_nodes)
+        phase_sigma_nodes = InterpolatedUnivariateSpline(log_frequency_array, phase_sigma)(log_nodes)
 
         prior = CalibrationPriorDict()
         for ii in range(n_nodes):
             name = "recalib_{}_amplitude_{}".format(label, ii)
             latex_label = "$A^{}_{}$".format(label, ii)
-            prior[name] = Gaussian(mu=amplitude_mean_nodes[ii],
-                                   sigma=amplitude_sigma_nodes[ii],
-                                   name=name, latex_label=latex_label,
-                                   boundary='reflective')
+            prior[name] = Gaussian(
+                mu=amplitude_mean_nodes[ii],
+                sigma=amplitude_sigma_nodes[ii],
+                name=name,
+                latex_label=latex_label,
+                boundary="reflective",
+            )
         for ii in range(n_nodes):
             name = "recalib_{}_phase_{}".format(label, ii)
             latex_label = "$\\phi^{}_{}$".format(label, ii)
-            prior[name] = Gaussian(mu=phase_mean_nodes[ii],
-                                   sigma=phase_sigma_nodes[ii],
-                                   name=name, latex_label=latex_label,
-                                   boundary='reflective')
+            prior[name] = Gaussian(
+                mu=phase_mean_nodes[ii],
+                sigma=phase_sigma_nodes[ii],
+                name=name,
+                latex_label=latex_label,
+                boundary="reflective",
+            )
         for ii in range(n_nodes):
             name = "recalib_{}_frequency_{}".format(label, ii)
             latex_label = "$f^{}_{}$".format(label, ii)
-            prior[name] = DeltaFunction(peak=np.exp(log_nodes[ii]), name=name,
-                                        latex_label=latex_label)
+            prior[name] = DeltaFunction(peak=np.exp(log_nodes[ii]), name=name, latex_label=latex_label)
         prior.source = os.path.abspath(envelope_file)
         return prior
 
     @staticmethod
-    def constant_uncertainty_spline(
-            amplitude_sigma, phase_sigma, minimum_frequency, maximum_frequency,
-            n_nodes, label):
+    def constant_uncertainty_spline(amplitude_sigma, phase_sigma, minimum_frequency, maximum_frequency, n_nodes, label):
         """
         Make prior assuming constant in frequency calibration uncertainty.
 
@@ -725,8 +725,7 @@ class CalibrationPriorDict(PriorDict):
             Priors for the relevant parameters.
             This includes the frequencies of the nodes which are _not_ sampled.
         """
-        nodes = np.logspace(np.log10(minimum_frequency),
-                            np.log10(maximum_frequency), n_nodes)
+        nodes = np.logspace(np.log10(minimum_frequency), np.log10(maximum_frequency), n_nodes)
 
         amplitude_mean_nodes = [0] * n_nodes
         amplitude_sigma_nodes = [amplitude_sigma] * n_nodes
@@ -737,28 +736,33 @@ class CalibrationPriorDict(PriorDict):
         for ii in range(n_nodes):
             name = "recalib_{}_amplitude_{}".format(label, ii)
             latex_label = "$A^{}_{}$".format(label, ii)
-            prior[name] = Gaussian(mu=amplitude_mean_nodes[ii],
-                                   sigma=amplitude_sigma_nodes[ii],
-                                   name=name, latex_label=latex_label,
-                                   boundary='reflective')
+            prior[name] = Gaussian(
+                mu=amplitude_mean_nodes[ii],
+                sigma=amplitude_sigma_nodes[ii],
+                name=name,
+                latex_label=latex_label,
+                boundary="reflective",
+            )
         for ii in range(n_nodes):
             name = "recalib_{}_phase_{}".format(label, ii)
             latex_label = "$\\phi^{}_{}$".format(label, ii)
-            prior[name] = Gaussian(mu=phase_mean_nodes[ii],
-                                   sigma=phase_sigma_nodes[ii],
-                                   name=name, latex_label=latex_label,
-                                   boundary='reflective')
+            prior[name] = Gaussian(
+                mu=phase_mean_nodes[ii],
+                sigma=phase_sigma_nodes[ii],
+                name=name,
+                latex_label=latex_label,
+                boundary="reflective",
+            )
         for ii in range(n_nodes):
             name = "recalib_{}_frequency_{}".format(label, ii)
             latex_label = "$f^{}_{}$".format(label, ii)
-            prior[name] = DeltaFunction(peak=nodes[ii], name=name,
-                                        latex_label=latex_label)
+            prior[name] = DeltaFunction(peak=nodes[ii], name=name, latex_label=latex_label)
 
         return prior
 
 
 def secondary_mass_condition_function(reference_params, mass_1):
-    return dict(minimum=reference_params['minimum'], maximum=mass_1)
+    return dict(minimum=reference_params["minimum"], maximum=mass_1)
 
 
 ConditionalCosmological = conditional_prior_factory(Cosmological)
@@ -767,7 +771,6 @@ ConditionalUniformSourceFrame = conditional_prior_factory(UniformSourceFrame)
 
 
 class HealPixMapPriorDist(BaseJointPriorDist):
-
     def __init__(self, hp_file, names=None, bounds=None, distance=False):
         """
         Class defining prior according to given HealPix Map, defaults to 2D in ra and dec but can be set to include
@@ -785,9 +788,9 @@ class HealPixMapPriorDist(BaseJointPriorDist):
         self.hp = self._check_imports()
         self.hp_file = hp_file
         if names is None:
-            names = ['ra', 'dec']
+            names = ["ra", "dec"]
         if bounds is None:
-            bounds = [[0, 2 * np.pi], [-np.pi / 2., np.pi / 2.]]
+            bounds = [[0, 2 * np.pi], [-np.pi / 2.0, np.pi / 2.0]]
         elif isinstance(bounds, dict):
             bs = [[] for _ in bounds.keys()]
             for i, key in enumerate(bounds.keys()):
@@ -795,22 +798,23 @@ class HealPixMapPriorDist(BaseJointPriorDist):
             bounds = bs
         if distance:
             if len(names) == 2:
-                names.append('distance')
+                names.append("distance")
             if len(bounds) == 2:
                 bounds.append([0, np.inf])
             self.distance = True
-            self.prob, self.distmu, self.distsigma, self.distnorm = self.hp.read_map(hp_file, verbose=False,
-                                                                                     field=range(4))
+            self.prob, self.distmu, self.distsigma, self.distnorm = self.hp.read_map(
+                hp_file, verbose=False, field=range(4)
+            )
         else:
             self.distance = False
             self.prob = self.hp.read_map(hp_file, verbose=False)
 
         super(HealPixMapPriorDist, self).__init__(names=names, bounds=bounds)
-        self.distname = 'hpmap'
+        self.distname = "hpmap"
         self.npix = len(self.prob)
         self.nside = self.hp.npix2nside(self.npix)
         self.pixel_area = self.hp.nside2pixarea(self.nside)
-        self.pixel_length = self.pixel_area**(1 / 2.)
+        self.pixel_length = self.pixel_area ** (1 / 2.0)
         self.pix_xx = np.arange(self.npix)
         self._all_interped = interp1d(x=self.pix_xx, y=self.prob, bounds_error=False, fill_value=0)
         self.inverse_cdf = None
@@ -887,17 +891,21 @@ class HealPixMapPriorDist(BaseJointPriorDist):
         ----------
         None - just updates these functions at new pixel values
         """
-        self.distance_pdf = lambda r: self.distnorm[pix_idx] * norm(loc=self.distmu[pix_idx],
-                                                                    scale=self.distsigma[pix_idx]).pdf(r)
-        self.distance_icdf = lambda d: self.distnorm[pix_idx] * norm(loc=self.distmu[pix_idx],
-                                                                     scale=self.distsigma[pix_idx]).ppf(d)
+        self.distance_pdf = lambda r: self.distnorm[pix_idx] * norm(
+            loc=self.distmu[pix_idx], scale=self.distsigma[pix_idx]
+        ).pdf(r)
+
+        rs = np.linspace(self.bounds["distance"][0], self.bounds["distance"][1], 1000)
+        pdfs = rs ** 2 * norm(loc=self.distmu[pix_idx], scale=self.distsigma[pix_idx]).pdf(rs)
+        cdfs = np.cumsum(pdfs) / np.sum(pdfs)
 
         def sample_distance(n):
             gaussian = norm(loc=self.distmu[pix_idx], scale=self.distsigma[pix_idx]).rvs(size=100 * n)
-            probs = self._check_norm(gaussian**2)
+            probs = self._check_norm(gaussian ** 2)
             return np.random.choice(gaussian, size=n, p=probs, replace=True)
 
         self.distance_dist = sample_distance
+        self.distance_icdf = interp1d(cdfs, rs)
 
     @staticmethod
     def _check_norm(array):
@@ -989,9 +997,11 @@ class HealPixMapPriorDist(BaseJointPriorDist):
         """
         if not self.check_in_pixel(ra, dec, pix):
             self.draw_from_pixel(ra, dec, pix)
-        return np.array([
-            np.random.uniform(ra - self.pixel_length, ra + self.pixel_length),
-            np.random.uniform(dec - self.pixel_length, dec + self.pixel_length)]
+        return np.array(
+            [
+                np.random.uniform(ra - self.pixel_length, ra + self.pixel_length),
+                np.random.uniform(dec - self.pixel_length, dec + self.pixel_length),
+            ]
         )
 
     def check_in_pixel(self, ra, dec, pix):
@@ -1053,7 +1063,7 @@ class HealPixMapPriorDist(BaseJointPriorDist):
         return lnprob
 
     def __eq__(self, other):
-        skip_keys = ['_all_interped', 'inverse_cdf', 'distance_pdf', 'distance_dist', 'distance_icdf']
+        skip_keys = ["_all_interped", "inverse_cdf", "distance_pdf", "distance_dist", "distance_icdf"]
         if self.__class__ != other.__class__:
             return False
         if sorted(self.__dict__.keys()) != sorted(other.__dict__.keys()):
@@ -1061,7 +1071,7 @@ class HealPixMapPriorDist(BaseJointPriorDist):
         for key in self.__dict__:
             if key in skip_keys:
                 continue
-            if key == 'hp_file':
+            if key == "hp_file":
                 if self.__dict__[key] != other.__dict__[key]:
                     return False
             elif isinstance(self.__dict__[key], (np.ndarray, list)):