diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 0000000..7a52897 --- /dev/null +++ b/.coveragerc @@ -0,0 +1,5 @@ +[report] +show_missing = true + +[run] +omit = gpucbc/_version.py,gpucbc/test/* diff --git a/.github/workflows/unit_test.yml b/.github/workflows/unit_test.yml index e9f57fb..62f5963 100644 --- a/.github/workflows/unit_test.yml +++ b/.github/workflows/unit_test.yml @@ -25,7 +25,8 @@ jobs: run: | conda update pip setuptools conda install numpy astropy bilby python-lal python-lalsimulation - conda install pytest-cov + conda install jax + conda install pytest-cov parameterized pip install . - name: Test with pytest run: | diff --git a/gpucbc/__init__.py b/gpucbc/__init__.py index 5cca709..92e3984 100644 --- a/gpucbc/__init__.py +++ b/gpucbc/__init__.py @@ -1,25 +1,19 @@ -from . import likelihood, waveforms +from . import backend, likelihood, pn, waveforms -def disable_cupy(): - import numpy - from scipy.special import i0e - likelihood.xp = numpy - likelihood.i0e = i0e - waveforms.xp = numpy +from ._version import __version__ -def enable_cupy(): - try: - import cupy - from .cupy_utils import i0e - likelihood.xp = cupy - likelihood.i0e = i0e - waveforms.xp = cupy - except ImportError: - print("Cannot import cupy") - disable_cupy() - - -enable_cupy() +def set_backend(numpy): + scipy = dict( + numpy="scipy", + cupy="cupyx.scipy", + jax="jax.scipy", + ).get(numpy, None) + numpy = dict(jax="jax.numpy").get(numpy, numpy) + BACKEND = backend.Backend(numpy=numpy, scipy=scipy) + backend.BACKEND = BACKEND + likelihood.B = BACKEND + pn.B = BACKEND + waveforms.B = BACKEND diff --git a/gpucbc/backend.py b/gpucbc/backend.py new file mode 100644 index 0000000..aa504c1 --- /dev/null +++ b/gpucbc/backend.py @@ -0,0 +1,29 @@ +from importlib import import_module + + +class Backend: + + def __init__(self, numpy, scipy=None): + self.module = numpy + try: + self.np = import_module(numpy) + if scipy is not None: + self.special = import_module(f"{scipy}.special") + else: + self.special = None + except ImportError as e: + raise ImportError(f"Cannot initialize backend for {numpy} {scipy}.\n{e}") + + def to_numpy(self, array): + if self.module == "numpy": + return array + elif self.module == "jax.numpy": + import numpy as np + return np.asarray(array) + elif self.module == "cupy": + return self.np.asnumpy(array) + else: + return array + + +BACKEND = Backend("numpy", "scipy") diff --git a/gpucbc/likelihood.py b/gpucbc/likelihood.py index 5ca118b..0362a96 100644 --- a/gpucbc/likelihood.py +++ b/gpucbc/likelihood.py @@ -1,12 +1,7 @@ import numpy as np from bilby.core.likelihood import Likelihood -try: - import cupy as xp - from cupyx.special import i0e, logsumexp -except ImportError: - xp = np - from scipy.special import i0e, logsumexp +from .backend import BACKEND as B class CUPYGravitationalWaveTransient(Likelihood): @@ -61,25 +56,17 @@ def __init__( def _data_to_gpu(self): for ifo in self.interferometers: - self.psds[ifo.name] = xp.asarray( + self.psds[ifo.name] = B.np.asarray( ifo.power_spectral_density_array[ifo.frequency_mask] ) - self.strain[ifo.name] = xp.asarray( + self.strain[ifo.name] = B.np.asarray( ifo.frequency_domain_strain[ifo.frequency_mask] ) - self.frequency_array[ifo.name] = xp.asarray( + self.frequency_array[ifo.name] = B.np.asarray( ifo.frequency_array[ifo.frequency_mask] ) self.duration = ifo.strain_data.duration - def __repr__(self): - return ( - self.__class__.__name__ - + "(interferometers={},\n\twaveform_generator={})".format( - self.interferometers, self.waveform_generator - ) - ) - def noise_log_likelihood(self): """Calculates the real part of noise log-likelihood @@ -95,7 +82,7 @@ def noise_log_likelihood(self): log_l -= ( 2.0 / self.duration - * xp.sum(xp.abs(self.strain[name]) ** 2 / self.psds[name]) + * (B.np.abs(self.strain[name]) ** 2 / self.psds[name]).sum() ) self._noise_log_l = float(log_l) return self._noise_log_l @@ -134,29 +121,26 @@ def log_likelihood_ratio(self): d_inner_h=d_inner_h, h_inner_h=h_inner_h ) else: - log_l = - h_inner_h / 2 + xp.real(d_inner_h) + log_l = - h_inner_h / 2 + d_inner_h return float(log_l.real) def calculate_snrs(self, interferometer, waveform_polarizations): name = interferometer.name - signal_ifo = xp.sum( - xp.vstack( - [ - waveform_polarizations[mode] - * float( - interferometer.antenna_response( - self.parameters["ra"], - self.parameters["dec"], - self.parameters["geocent_time"], - self.parameters["psi"], - mode, - ) + signal_ifo = B.np.vstack( + [ + waveform_polarizations[mode] + * float( + interferometer.antenna_response( + self.parameters["ra"], + self.parameters["dec"], + self.parameters["geocent_time"], + self.parameters["psi"], + mode, ) - for mode in waveform_polarizations - ] - ), - axis=0, - )[interferometer.frequency_mask] + ) + for mode in waveform_polarizations + ] + ).sum(axis=0)[interferometer.frequency_mask] time_delay = ( self.parameters["geocent_time"] @@ -168,10 +152,10 @@ def calculate_snrs(self, interferometer, waveform_polarizations): ) ) - signal_ifo *= xp.exp(-2j * np.pi * time_delay * self.frequency_array[name]) + signal_ifo *= B.np.exp(-2j * np.pi * time_delay * self.frequency_array[name]) - d_inner_h = xp.sum(xp.conj(signal_ifo) * self.strain[name] / self.psds[name]) - h_inner_h = xp.sum(xp.abs(signal_ifo) ** 2 / self.psds[name]) + d_inner_h = (signal_ifo.conj() * self.strain[name] / self.psds[name]).sum() + h_inner_h = (B.np.abs(signal_ifo) ** 2 / self.psds[name]).sum() d_inner_h *= 4 / self.duration h_inner_h *= 4 / self.duration return d_inner_h, h_inner_h @@ -192,13 +176,13 @@ def distance_marglinalized_likelihood(self, d_inner_h, h_inner_h): d_inner_h=d_inner_h_array, h_inner_h=h_inner_h_array ) else: - log_l_array = - h_inner_h_array / 2 + xp.real(d_inner_h_array) + log_l_array = - h_inner_h_array / 2 + d_inner_h_array.real log_l = logsumexp(log_l_array, b=self.distance_prior_array) return log_l def phase_marginalized_likelihood(self, d_inner_h, h_inner_h): - d_inner_h = xp.abs(d_inner_h) - d_inner_h = xp.log(i0e(d_inner_h)) + d_inner_h + d_inner_h = B.np.abs(d_inner_h) + d_inner_h = B.np.log(B.special.i0e(d_inner_h)) + d_inner_h log_l = - h_inner_h / 2 + d_inner_h return log_l @@ -208,10 +192,10 @@ def _setup_distance_marginalization(self): self.priors["luminosity_distance"].maximum, 10000, ) - self.distance_prior_array = xp.asarray( + self.distance_prior_array = B.np.asarray( self.priors["luminosity_distance"].prob(self.distance_array) ) * (self.distance_array[1] - self.distance_array[0]) - self.distance_array = xp.asarray(self.distance_array) + self.distance_array = B.np.asarray(self.distance_array) def generate_posterior_sample_from_marginalized_likelihood(self): return self.parameters.copy() diff --git a/gpucbc/pn.py b/gpucbc/pn.py index e14980b..a42630d 100644 --- a/gpucbc/pn.py +++ b/gpucbc/pn.py @@ -1,5 +1,7 @@ import numpy as np +from .backend import BACKEND as B + PI = np.pi @@ -135,7 +137,7 @@ def taylor_f2_phase_6(args): + args.eta * (-15737765635 / 3048192 + 2255 / 12 * PI ** 2) + args.eta ** 2 * 76055 / 1728 - args.eta ** 3 * 127825 / 1296 - + taylor_f2_phase_6l(args) * np.log(4) + + taylor_f2_phase_6l(args) * B.np.log(4) ) phase += (32675 / 112 + 5575 / 18 * args.eta) * args.eta * args.chi_1 * args.chi_2 for m_on_m, chi, qm_def in zip( diff --git a/gpucbc/test/__init__.py b/gpucbc/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/gpucbc/test/test_backends.py b/gpucbc/test/test_backends.py new file mode 100644 index 0000000..0c25766 --- /dev/null +++ b/gpucbc/test/test_backends.py @@ -0,0 +1,23 @@ +import unittest + +import bilby +import gpucbc +import jax.numpy as jnp +from jax.scipy import special as jsp + + +class TestBackend(unittest.TestCase): + + def test_setting_jax(self): + gpucbc.set_backend("jax") + self.assertEqual(gpucbc.pn.B.np, jnp) + self.assertEqual(gpucbc.pn.B.special, jsp) + + def test_unknown_backend_raises_error(self): + with self.assertRaises(ImportError): + gpucbc.set_backend("unknown") + + def test_no_scipy_backend(self): + gpucbc.set_backend("bilby") + self.assertEqual(gpucbc.pn.B.np, bilby) + self.assertEqual(gpucbc.pn.B.special, None) diff --git a/test/test_waveform.py b/gpucbc/test/test_waveform.py similarity index 86% rename from test/test_waveform.py rename to gpucbc/test/test_waveform.py index 3811110..cdfbfcf 100644 --- a/test/test_waveform.py +++ b/gpucbc/test/test_waveform.py @@ -1,6 +1,8 @@ #!/usr/bin/env python3 import unittest +import pytest +from parameterized import parameterized import numpy as np import pandas as pd @@ -15,6 +17,8 @@ from bilby.gw.utils import noise_weighted_inner_product from bilby.gw.waveform_generator import WaveformGenerator +import gpucbc +from gpucbc import set_backend from gpucbc.waveforms import TF2WFG, TF2 @@ -51,7 +55,10 @@ def setUp(self) -> None: parameter_conversion=convert_to_lal_binary_neutron_star_parameters, ) - def test_native_phasing(self): + @parameterized.expand(["numpy", "jax", "cupy"]) + def test_native_phasing(self, backend): + pytest.importorskip(backend) + set_backend(backend) priors = PriorDict() priors["mass_1"] = Uniform(1, 100) priors["mass_2"] = Uniform(1, 100) @@ -66,21 +73,21 @@ def test_native_phasing(self): priors["luminosity_distance"] = Uniform(10, 200) wf = TF2(**priors.sample()) + TF2.pn_tidal_order = 15 lal_phasing = wf._lal_phasing_coefficients() my_phasing = wf.phasing_coefficients() - self.assertLess(max(abs(lal_phasing.v - my_phasing.v)), 1e-8) - self.assertLess(max(abs(lal_phasing.vlogv - my_phasing.vlogv)), 1e-8) - self.assertLess(max(abs(lal_phasing.vlogvsq - my_phasing.vlogvsq)), 1e-8) - - def test_absolute_overlap(self): + self.assertLess(max(abs(lal_phasing.v - my_phasing.v)), 1e-5) + self.assertLess(max(abs(lal_phasing.vlogv - my_phasing.vlogv)), 1e-5) + self.assertLess(max(abs(lal_phasing.vlogvsq - my_phasing.vlogvsq)), 1e-5) + + @parameterized.expand(["numpy", "jax", "cupy"]) + def test_absolute_overlap(self, backend): + pytest.importorskip(backend) + set_backend(backend) np.random.seed(42) priors = BNSPriorDict(aligned_spin=True) priors["mass_1"] = Uniform(1, 100) priors["mass_2"] = Uniform(1, 100) - # priors["total_mass"] = Uniform(2, 100) - # priors["mass_ratio"] = Uniform(name="mass_ratio", minimum=0.5, maximum=1) - # priors["mass_1"] = Constraint(name="mass_1", minimum=1, maximum=50) - # priors["mass_2"] = Constraint(name="mass_2", minimum=1, maximum=50) del priors["mass_ratio"], priors["chirp_mass"] priors["luminosity_distance"] = UniformSourceFrame( name="luminosity_distance", minimum=1e2, maximum=5e3 @@ -104,8 +111,7 @@ def test_absolute_overlap(self): ) priors["lambda_1"] = Uniform(name="lambda_1", minimum=0, maximum=5000) priors["lambda_2"] = Uniform(name="lambda_2", minimum=0, maximum=5000) - priors["geocent_time"] = 0 - # priors["geocent_time"] = Uniform(-10, 10) + priors["geocent_time"] = Uniform(-10, 10) n_samples = 100 all_parameters = pd.DataFrame(priors.sample(n_samples)) @@ -144,6 +150,7 @@ def test_absolute_overlap(self): ) bilby_pols = dict(plus=lal_strain[0].data.data, cross=lal_strain[1].data.data) gpu_pols = self.gpu_wfg.frequency_domain_strain(parameters) + gpu_pols = {key: gpucbc.backend.BACKEND.to_numpy(value) for key, value in gpu_pols.items()} bilby_strain = self.ifo.get_detector_response( waveform_polarizations=bilby_pols, parameters=parameters @@ -164,12 +171,5 @@ def test_absolute_overlap(self): / self.ifo.optimal_snr_squared(signal=bilby_strain) ** 0.5 / self.ifo.optimal_snr_squared(signal=gpu_strain) ** 0.5 ) - # print(self.ifo.optimal_snr_squared(signal=bilby_strain) ** 0.5, self.ifo.optimal_snr_squared(signal=gpu_strain) ** 0.5) - # print(np.max(abs( - # np.fft.fft(bilby_strain[:-1] * gpu_strain.conj().real[:-1]) / self.ifo.power_spectral_density_array[:-1] - # / self.ifo.optimal_snr_squared(signal=bilby_strain) ** 0.5 - # / self.ifo.optimal_snr_squared(signal=gpu_strain) ** 0.5 - # ))) overlaps.append(overlap) - print(overlaps) self.assertTrue(min(np.abs(overlaps)) > 0.99) diff --git a/gpucbc/waveforms.py b/gpucbc/waveforms.py index 278e6a0..0de3383 100644 --- a/gpucbc/waveforms.py +++ b/gpucbc/waveforms.py @@ -2,14 +2,11 @@ import numpy as np from astropy import constants +from attrs import define from bilby.gw.conversion import convert_to_lal_binary_neutron_star_parameters from bilby.core.utils import create_frequency_series -try: - import cupy as xp -except ImportError: - xp = np - +from .backend import BACKEND as B from . import pn SOLAR_RADIUS_IN_M = constants.GM_sun.si.value / constants.c.si.value ** 2 @@ -17,14 +14,27 @@ MEGA_PARSEC_SI = constants.pc.si.value * 1e6 -class Phasing(object): +class Phasing: def __init__(self, order=16): self.v = np.zeros(order, dtype=float) self.vlogv = np.zeros(order, dtype=float) self.vlogvsq = np.zeros(order, dtype=float) -class TF2(object): +@define +class PhaseParameters: + eta: float + chi_1: float + chi_2: float + m1_on_m: float + m2_on_m: float + qm_def_1: float + qm_def_2: float + lambda_1: float + lambda_2: float + + +class TF2: """ A copy of the TaylorF2 waveform. @@ -51,21 +61,6 @@ class TF2(object): Dimensionless tidal deformability of the less massive object. """ - phase_parameters = namedtuple( - "PhaseParameters", - [ - "eta", - "chi_1", - "chi_2", - "m1_on_m", - "m2_on_m", - "qm_def_1", - "qm_def_2", - "lambda_1", - "lambda_2", - ], - ) - def __init__( self, mass_1, mass_2, chi_1, chi_2, luminosity_distance, lambda_1=0, lambda_2=0 ): @@ -88,7 +83,7 @@ def __init__( self.pn_spin_order = -1 self.pn_tidal_order = -1 - self.args = self.phase_parameters( + self.args = PhaseParameters( self.eta, self.chi_1, self.chi_2, @@ -102,7 +97,7 @@ def __init__( def __call__(self, frequency_array, tc=0, phi_c=0): orbital_speed = self.orbital_speed(frequency_array=frequency_array) - hoff = self.amplitude(frequency_array, orbital_speed=orbital_speed) * xp.exp( + hoff = self.amplitude(frequency_array, orbital_speed=orbital_speed) * B.np.exp( -1j * self.phase(frequency_array, phi_c=phi_c, orbital_speed=orbital_speed) ) return hoff @@ -134,6 +129,7 @@ def _lal_phasing_coefficients(self): from lalsimulation import ( SimInspiralWaveformParamsInsertTidalLambda1, SimInspiralWaveformParamsInsertTidalLambda2, + SimInspiralWaveformParamsInsertPNTidalOrder, SimInspiralSetQuadMonParamsFromLambdas, SimInspiralTaylorF2AlignedPhasing, ) @@ -141,6 +137,7 @@ def _lal_phasing_coefficients(self): param_dict = CreateDict() SimInspiralWaveformParamsInsertTidalLambda1(param_dict, self.lambda_1) SimInspiralWaveformParamsInsertTidalLambda2(param_dict, self.lambda_2) + SimInspiralWaveformParamsInsertPNTidalOrder(param_dict, self.pn_tidal_order) SimInspiralSetQuadMonParamsFromLambdas(param_dict) return SimInspiralTaylorF2AlignedPhasing( self.mass_1, self.mass_2, self.chi_1, self.chi_2, param_dict @@ -148,18 +145,8 @@ def _lal_phasing_coefficients(self): def phasing_coefficients(self): scale = 3 / (128 * self.eta) - self._phasing.v[0] = pn.taylor_f2_phase_0(self.args) - self._phasing.v[1] = pn.taylor_f2_phase_1(self.args) - self._phasing.v[2] = pn.taylor_f2_phase_2(self.args) - self._phasing.v[3] = pn.taylor_f2_phase_3(self.args) - self._phasing.v[4] = pn.taylor_f2_phase_4(self.args) - self._phasing.v[5] = pn.taylor_f2_phase_5(self.args) - self._phasing.v[6] = pn.taylor_f2_phase_6(self.args) - self._phasing.v[7] = pn.taylor_f2_phase_7(self.args) - self._phasing.v[10] = pn.taylor_f2_phase_10(self.args) - self._phasing.v[12] = pn.taylor_f2_phase_12(self.args) - self._phasing.v[13] = pn.taylor_f2_phase_13(self.args) - self._phasing.v[14] = pn.taylor_f2_phase_14(self.args) + for ii in range(15): + self._phasing.v[ii] = getattr(pn, f"taylor_f2_phase_{ii}")(self.args) if self.pn_tidal_order > 14: self._phasing.v[15] = pn.taylor_f2_phase_15(self.args) self._phasing.vlogv[5] = pn.taylor_f2_phase_5l(self.args) @@ -173,9 +160,9 @@ def phase(self, frequency_array, phi_c=0, orbital_speed=None): if orbital_speed is None: orbital_speed = self.orbital_speed(frequency_array=frequency_array) phase_coefficients = self.phasing_coefficients() - phasing = xp.zeros(orbital_speed.shape) + phasing = B.np.zeros(orbital_speed.shape) cumulative_power_frequency = orbital_speed ** -5 - log_orbital_speed = xp.log(orbital_speed) + log_orbital_speed = B.np.log(orbital_speed) for ii in range(len(phase_coefficients.v)): phasing += phase_coefficients.v[ii] * cumulative_power_frequency phasing += ( @@ -226,10 +213,7 @@ def call_cupy_tf2( minimum_frequency = waveform_kwargs["minimum_frequency"] in_band = frequency_array >= minimum_frequency - - frequency_array = xp.asarray(frequency_array) - - h_out_of_band = xp.zeros(int(xp.sum(~in_band))) + frequency_array = B.np.asarray(frequency_array) wf = TF2( mass_1, @@ -240,8 +224,10 @@ def call_cupy_tf2( lambda_2=lambda_2, luminosity_distance=luminosity_distance, ) + strain = wf(frequency_array[in_band], phi_c=phase) - strain = xp.hstack([h_out_of_band, strain]) + h_out_of_band = B.np.zeros(len(frequency_array) - len(strain)) + strain = B.np.hstack([h_out_of_band, strain]) h_plus = strain * (1 + np.cos(theta_jn) ** 2) / 2 h_cross = -1j * strain * np.cos(theta_jn) @@ -261,7 +247,7 @@ def __init__( waveform_arguments = dict(minimum_frequency=10) self.fdsm = frequency_domain_source_model self.waveform_arguments = waveform_arguments - self.frequency_array = xp.asarray( + self.frequency_array = B.np.asarray( create_frequency_series( duration=duration, sampling_frequency=sampling_frequency ) @@ -282,7 +268,7 @@ def eos_q_from_lambda(lamb, tolerance=0.5): """ def worker(log_lambda): - return np.exp( + return B.np.exp( 0.194 + 0.0936 * log_lambda + 0.0474 * log_lambda ** 2 diff --git a/pyproject.toml b/pyproject.toml index 05b60ae..d4ee42e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ classifiers = [ "Topic :: Scientific/Engineering", "Intended Audience :: Science/Research", ] -dependencies = ["numpy >= 1.16.5", "scipy", "astropy", "bilby >= 2.0"] +dependencies = ["attrs", "numpy >= 1.16.5", "scipy", "astropy", "bilby >= 2.0"] readme = {file = "README.md", content-type = "text/markdown"} dynamic = ["version"] @@ -38,4 +38,7 @@ homepage = "https://github.com/ColmTalbot/GPUCBC" namespaces = false [tool.setuptools_scm] -write_to = "gpucbc/_version.py" \ No newline at end of file +write_to = "gpucbc/_version.py" + +[tool.pytest.ini_options] +addopts = "--cov gpucbc --cov-report term-missing" \ No newline at end of file