Commit f1cda596 authored by Leo Pound Singer's avatar Leo Pound Singer

Refactor and add unit test for get_max_z

parent 5c8a41c6
Pipeline #95822 failed with stages
in 43 minutes and 32 seconds
......@@ -23,7 +23,7 @@ not reject an excessive number of events. We divide the intrinsic parameter
space into a very coarse grid and we calculate the maximum horizon distance in
each grid cell."""
import functools
from functools import partial
from astropy import cosmology
from astropy.cosmology.core import vectorize_if_needed
......@@ -36,6 +36,7 @@ from scipy.interpolate import interp1d
from scipy.optimize import root_scalar
from scipy.ndimage import maximum_filter
from ..util import progress_map
from ..bayestar.filter import sngl_inspiral_psd
from . import (
ArgumentParser, FileType, random_parser, register_to_xmldoc, write_fileobj)
......@@ -73,7 +74,7 @@ def lo_hi_nonzero(x):
return nonzero[0], nonzero[-1]
def z_at_snr(cosmo, psds, waveform, f_low, snr, params):
def z_at_snr(cosmo, psds, waveform, f_low, snr, mass1, mass2, spin1z, spin2z):
"""
Get redshift at which a waveform attains a given SNR.
......@@ -98,7 +99,6 @@ def z_at_snr(cosmo, psds, waveform, f_low, snr, params):
Comoving distance in Mpc.
"""
# Construct waveform
mass1, mass2, spin1z, spin2z = params
series = sngl_inspiral_psd(waveform, f_low=f_low,
mass1=mass1, mass2=mass2,
spin1z=spin1z, spin2z=spin2z)
......@@ -134,6 +134,21 @@ def z_at_snr(cosmo, psds, waveform, f_low, snr, params):
return root_scalar(lambda z: snr_at_z(z) - snr, bracket=(0, 1e3)).root
def get_max_z(cosmo, psds, waveform, f_low, snr, mass1, mass2, spin1z, spin2z,
jobs=1):
# Calculate the maximum distance on the grid.
params = [mass1, mass2, spin1z, spin2z]
result = progress_map(
partial(z_at_snr, cosmo, psds, waveform, f_low, snr),
*(param.ravel() for param in np.meshgrid(*params, indexing='ij')),
jobs=jobs)
result = np.reshape(result, tuple(len(param) for param in params))
assert np.all(result >= 0), 'some redshifts are negative'
assert np.all(np.isfinite(result)), 'some redshifts are not finite'
return result
def _sensitive_volume_integral(cosmo, z):
dh3_sr = cosmo.hubble_distance**3 / units.sr
......@@ -226,8 +241,6 @@ def main(args=None):
import lal.series
from scipy import stats
from ..util import progress_map
p = parser()
args = p.parse_args(args)
......@@ -342,20 +355,9 @@ def main(args=None):
params = m1, m2, x1, x2
# Calculate the maximum distance on the grid.
max_z = np.reshape(
progress_map(
functools.partial(
z_at_snr, cosmo, psds,
args.waveform, args.f_low, args.min_snr),
np.column_stack([param.ravel() for param
in np.meshgrid(*params, indexing='ij')]),
jobs=args.jobs),
tuple(len(param) for param in params))
# Make sure that all redshifts are valid.
assert np.all(max_z >= 0), 'some redshifts are negative'
assert np.all(np.isfinite(max_z)), 'some redshifts are not finite'
max_z = get_max_z(
cosmo, psds, args.waveform, args.f_low, args.min_snr, m1, m2, x1, x2,
jobs=args.jobs)
max_distance = sensitive_distance(cosmo, max_z).to_value(units.Mpc)
# Find piecewise constant approximate upper bound on distance.
......
......@@ -9,7 +9,8 @@ import pytest
from scipy.misc import derivative
from ..bayestar_inject import (
get_decisive_snr, z_at_snr, cell_max, sensitive_distance, sensitive_volume)
get_decisive_snr, z_at_snr, get_max_z, cell_max, sensitive_distance,
sensitive_volume)
def test_get_decisive_snr():
......@@ -49,11 +50,41 @@ def test_z_at_snr(mtotal, z):
snr = get_snr_at_z_lalsimulation(
cosmo, z, mass1, mass2, f_low, f_high, psd)
z_solution = z_at_snr(
cosmo, [psd], 'IMRPhenomPv2', f_low, snr, (mass1, mass2, 0, 0))
cosmo, [psd], 'IMRPhenomPv2', f_low, snr, mass1, mass2, 0, 0)
assert z_solution == pytest.approx(z, rel=1e-2)
def test_get_max_z():
cosmo = default_cosmology.get_cosmology_from_string('Planck15')
f_low = 10
f_high = 4096
df = 0.1
waveform = 'IMRPhenomPv2'
snr = 8
m1 = np.asarray([50.0])
m2 = np.asarray([30.0, 50.0])
x1 = np.asarray([-1.0, 0.0, 1.0])
x2 = np.asarray([-1.0, -0.5, 0.5, 1.0])
psd = lal.CreateREAL8FrequencySeries(
'', 0, f_low, df, lal.DimensionlessUnit, int((f_high - f_low) // df))
lalsimulation.SimNoisePSDaLIGODesignSensitivityP1200087(psd, f_low)
result = get_max_z(
cosmo, [psd], waveform, f_low, snr, m1, m2, x1, x2)
# Check that shape matches
assert result.shape == (1, 2, 3, 4)
# Spot check some individual cells
for im1, m1_ in enumerate(m1):
for im2, m2_ in enumerate(m2):
for ix1, x1_ in enumerate(x1):
for ix2, x2_ in enumerate(x2):
expected = z_at_snr(
cosmo, [psd], waveform, f_low, snr, m1_, m2_, x1_, x2_)
assert result[im1, im2, ix1, ix2] == expected
def test_sensitive_volume_0():
cosmo = default_cosmology.get_cosmology_from_string('Planck15')
assert sensitive_volume(cosmo, 0) == 0
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment