diff --git a/pycbc/distributions/__init__.py b/pycbc/distributions/__init__.py index d6acccd1599..1ab45c41f24 100644 --- a/pycbc/distributions/__init__.py +++ b/pycbc/distributions/__init__.py @@ -29,7 +29,7 @@ from pycbc.distributions.arbitrary import Arbitrary, FromFile from pycbc.distributions.gaussian import Gaussian from pycbc.distributions.power_law import UniformPowerLaw, UniformRadius -from pycbc.distributions.sky_location import UniformSky, FisherSky +from pycbc.distributions.sky_location import UniformSky, FisherSky, HealpixSky from pycbc.distributions.uniform import Uniform from pycbc.distributions.uniform_log import UniformLog10 from pycbc.distributions.spins import IndependentChiPChiEff @@ -61,7 +61,8 @@ FixedSamples.name: FixedSamples, MchirpfromUniformMass1Mass2.name: MchirpfromUniformMass1Mass2, QfromUniformMass1Mass2.name: QfromUniformMass1Mass2, - FisherSky.name: FisherSky + FisherSky.name: FisherSky, + HealpixSky.name: HealpixSky } def read_distributions_from_config(cp, section="prior"): diff --git a/pycbc/distributions/sky_location.py b/pycbc/distributions/sky_location.py index 76cb26c508c..64ebaa3d599 100644 --- a/pycbc/distributions/sky_location.py +++ b/pycbc/distributions/sky_location.py @@ -20,6 +20,7 @@ import logging import numpy + from scipy.spatial.transform import Rotation from pycbc.distributions import angular @@ -36,13 +37,14 @@ class UniformSky(angular.UniformSolidAngle): names are "dec" (declination) for the polar angle and "ra" (right ascension) for the azimuthal angle, instead of "theta" and "phi". """ + name = 'uniform_sky' _polardistcls = angular.CosAngle _default_polar_angle = 'dec' _default_azimuthal_angle = 'ra' -class FisherSky(): +class FisherSky: """A distribution that returns a random angle drawn from an approximate `Von_Mises-Fisher distribution`_. Assumes that the Fisher concentration parameter is large, so that we can draw the samples from a simple @@ -76,6 +78,7 @@ class FisherSky(): angle_unit: str Unit for the angle parameters: either "deg" or "rad". """ + name = 'fisher_sky' _params = ['ra', 'dec'] @@ -93,7 +96,7 @@ def __init__(self, **params): raise ValueError( f'The mean RA must be between 0 and 2π, {mean_ra} rad given' ) - if mean_dec < -numpy.pi/2 or mean_dec > numpy.pi/2: + if mean_dec < -numpy.pi / 2 or mean_dec > numpy.pi / 2: raise ValueError( 'The mean declination must be between ' f'-π/2 and π/2, {mean_dec} rad given' @@ -106,13 +109,13 @@ def __init__(self, **params): if sigma > 0.35: logger.warning( 'Warning: sigma = %s rad is probably too large for the ' - 'Fisher approximation to be valid', sigma + 'Fisher approximation to be valid', + sigma, ) self.rayleigh_scale = 0.66 * sigma # Prepare a rotation that puts the North Pole at the mean position self.rotation = Rotation.from_euler( - 'yz', - [numpy.pi / 2 - mean_dec, mean_ra] + 'yz', [numpy.pi / 2 - mean_dec, mean_ra] ) @property @@ -124,8 +127,10 @@ def from_config(cls, cp, section, variable_args): tag = variable_args variable_args = variable_args.split(VARARGS_DELIM) if set(variable_args) != set(cls._params): - raise ValueError("Not all parameters used by this distribution " - "included in tag portion of section name") + raise ValueError( + "Not all parameters used by this distribution " + "included in tag portion of section name" + ) mean_ra = float(cp.get_opt_tag(section, 'mean_ra', tag)) mean_dec = float(cp.get_opt_tag(section, 'mean_dec', tag)) sigma = float(cp.get_opt_tag(section, 'sigma', tag)) @@ -134,20 +139,13 @@ def from_config(cls, cp, section, variable_args): mean_ra=mean_ra, mean_dec=mean_dec, sigma=sigma, - angle_unit=angle_unit + angle_unit=angle_unit, ) def rvs(self, size): # Draw samples from a distribution centered on the North pole - np_ra = numpy.random.uniform( - low=0, - high=(2*numpy.pi), - size=size - ) - np_dec = numpy.random.rayleigh( - scale=self.rayleigh_scale, - size=size - ) + np_ra = numpy.random.uniform(low=0, high=(2 * numpy.pi), size=size) + np_dec = numpy.random.rayleigh(scale=self.rayleigh_scale, size=size) # Convert the samples to intermediate cartesian representation np_cart = numpy.empty(shape=(size, 3)) @@ -161,13 +159,7 @@ def rvs(self, size): # Convert the samples back to spherical coordinates. # Some unpleasant conditional operations are needed # to get the correct angle convention. - rot_radec = FieldArray( - size, - dtype=[ - ('ra', ' coverage: + break + + # A safety margin is added to ensure that the pixels at the edges + # of the distribution are fully accounted for. + # width of one pixel : < π/4nside-1 + + margin = 2 * numpy.pi / (4 * nside - 1) + + delta_max = min(delta_max + margin, numpy.pi / 2) + delta_min = max(delta_min - margin, -numpy.pi / 2) + alpha_max = min(alpha_max + margin, 2 * numpy.pi) + alpha_min = max(alpha_min - margin, 0) + + return delta_min, delta_max, alpha_min, alpha_max + + file_name = params['healpix_file'] + + coverage = params['coverage'] + + if coverage > 1 or coverage < 0: + raise ValueError( + f'Coverage must be between 0 and 1, {coverage} is not correct' + ) + + rasterization_nside = params['rasterization_nside'] + + if bin(rasterization_nside).count('1') != 1: + raise ValueError( + f'Rasterization_nside must be a power of 2,' + f'{rasterization_nside} is not correct' + ) + self.healpix_map = mhealpy.HealpixMap.read_map(file_name) + self.boundaries = boundaries( + self.healpix_map, rasterization_nside, coverage + ) + logging.info('HealpixSky boundary is %s', self.boundaries) + + @property + def params(self): + return self._params + + @classmethod + def from_config(cls, cp, section, variable_args): + tag = variable_args + variable_args = variable_args.split(VARARGS_DELIM) + if set(variable_args) != set(cls._params): + raise ValueError( + "Not all parameters used by this distribution " + "included in tag portion of section name" + ) + healpix_file = str(cp.get_opt_tag(section, 'healpix_file', tag)) + coverage = 0.9999 + if cp.has_option_tag(section, 'coverage', tag): + coverage = float(cp.get_opt_tag(section, 'coverage', tag)) + + rasterization_nside = 64 + if cp.has_option_tag(section, 'rasterization_nside', tag): + rasterization_nside = int( + cp.get_opt_tag(section, 'rasterization_nside', tag) + ) + return cls( + healpix_file=healpix_file, + coverage=coverage, + rasterization_nside=rasterization_nside, + ) + + def rvs(self, size): + def simple_rejection_sampling(healpix_map, size, boundaries): + """Start from a uniform distribution of points, and accepts those + whose values on the map are greater than a random value + following a uniform law + + Parameters + ---------- + healpix_map : HealpixMap instance + size : int + number of points tested by the method + boundaries : list -> tuple of 4 floats + delta_min,delta_max,alpha_min,alpha_max = boundaries + delta is the declination in radians [-pi/2, pi/2] + alpha is the right ascention in radians [0,2pi] + + Returns + ------- + coordinates of the accepted points, + following the mhealpy conventions + + """ + # The angles are in radians and follow the radec convention + delta_min, delta_max, alpha_min, alpha_max = boundaries + + # draw points uniformly distributed inside the region delimited by + # boundaries + u = numpy.random.uniform(0, 1, size) + delta = numpy.arcsin( + (numpy.sin(delta_max) - numpy.sin(delta_min)) * u + + numpy.sin(delta_min) + ) + alpha = numpy.random.uniform(alpha_min, alpha_max, size) + # a conversion is required to use get_interp_val + theta, phi = numpy.pi / 2 - delta, alpha + + data = healpix_map.data + random_data = numpy.random.uniform(0, data.max(), size) + + # the version of mhealpy 0.3.4 or later is needed to run + # get_interp_val + # get_interp_val might return an astropy object depending on the + # units of the column of the fits file, + # hence the need to transform into an array + d_data = numpy.array( + healpix_map.get_interp_val(theta, phi, lonlat=False) + ) + + dist_theta = theta[d_data > random_data] + dist_phi = phi[d_data > random_data] + + return (dist_theta, dist_phi) + + # Sampling method to generate the desired number of points + theta, phi = numpy.array([]), numpy.array([]) + while len(theta) < size: + new_theta, new_phi = simple_rejection_sampling( + self.healpix_map, size, self.boundaries + ) + theta = numpy.concatenate((theta, new_theta), axis=0) + phi = numpy.concatenate((phi, new_phi), axis=0) + + if len(theta) > size: + theta = theta[:size] + phi = phi[:size] + + # convert back to the radec convention + radec = FieldArray(size, dtype=[('ra', '