From d8af9bbc69ba1ce8ecbc982171bd7a72459c907b Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 15 May 2023 13:30:13 -0400 Subject: [PATCH 1/3] DEV: add option for jax backend --- .github/workflows/unit_test.yml | 1 + gpucbc/__init__.py | 35 ++++++---------- gpucbc/backend.py | 16 ++++++++ gpucbc/likelihood.py | 72 +++++++++++++-------------------- gpucbc/pn.py | 4 +- gpucbc/waveforms.py | 24 +++++------ pyproject.toml | 5 ++- test/test_backends.py | 23 +++++++++++ test/test_waveform.py | 28 ++++++------- 9 files changed, 112 insertions(+), 96 deletions(-) create mode 100644 gpucbc/backend.py create mode 100644 test/test_backends.py diff --git a/.github/workflows/unit_test.yml b/.github/workflows/unit_test.yml index e9f57fb..2f74279 100644 --- a/.github/workflows/unit_test.yml +++ b/.github/workflows/unit_test.yml @@ -25,6 +25,7 @@ jobs: run: | conda update pip setuptools conda install numpy astropy bilby python-lal python-lalsimulation + conda install jax conda install pytest-cov pip install . - name: Test with pytest diff --git a/gpucbc/__init__.py b/gpucbc/__init__.py index 5cca709..2757093 100644 --- a/gpucbc/__init__.py +++ b/gpucbc/__init__.py @@ -1,25 +1,16 @@ -from . import likelihood, waveforms +from . import likelihood, pn, waveforms +from .backend import Backend -def disable_cupy(): - import numpy - from scipy.special import i0e - likelihood.xp = numpy - likelihood.i0e = i0e - waveforms.xp = numpy - - -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(numpy=numpy, scipy=scipy) + likelihood.BACKEND = BACKEND + pn.BACKEND = BACKEND + waveforms.BACKEND = BACKEND diff --git a/gpucbc/backend.py b/gpucbc/backend.py new file mode 100644 index 0000000..abf7f5a --- /dev/null +++ b/gpucbc/backend.py @@ -0,0 +1,16 @@ +from importlib import import_module + + +class Backend: + + def __init__(self, numpy, scipy=None): + 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}") + +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/waveforms.py b/gpucbc/waveforms.py index 278e6a0..faf13ba 100644 --- a/gpucbc/waveforms.py +++ b/gpucbc/waveforms.py @@ -5,11 +5,7 @@ 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 @@ -102,7 +98,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 +130,7 @@ def _lal_phasing_coefficients(self): from lalsimulation import ( SimInspiralWaveformParamsInsertTidalLambda1, SimInspiralWaveformParamsInsertTidalLambda2, + SimInspiralWaveformParamsInsertPNTidalOrder, SimInspiralSetQuadMonParamsFromLambdas, SimInspiralTaylorF2AlignedPhasing, ) @@ -141,6 +138,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 @@ -173,9 +171,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 += ( @@ -227,9 +225,9 @@ def call_cupy_tf2( in_band = frequency_array >= minimum_frequency - frequency_array = xp.asarray(frequency_array) + frequency_array = B.np.asarray(frequency_array) - h_out_of_band = xp.zeros(int(xp.sum(~in_band))) + h_out_of_band = B.np.zeros(int(B.np.sum(~in_band))) wf = TF2( mass_1, @@ -241,7 +239,7 @@ def call_cupy_tf2( luminosity_distance=luminosity_distance, ) strain = wf(frequency_array[in_band], phi_c=phase) - strain = xp.hstack([h_out_of_band, 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 +259,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 +280,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..be75cbc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 diff --git a/test/test_backends.py b/test/test_backends.py new file mode 100644 index 0000000..b2fad63 --- /dev/null +++ b/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.BACKEND.np, jnp) + self.assertEqual(gpucbc.pn.BACKEND.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.BACKEND.np, bilby) + self.assertEqual(gpucbc.pn.BACKEND.special, None) diff --git a/test/test_waveform.py b/test/test_waveform.py index 3811110..9070713 100644 --- a/test/test_waveform.py +++ b/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,7 @@ from bilby.gw.utils import noise_weighted_inner_product from bilby.gw.waveform_generator import WaveformGenerator +from gpucbc import set_backend from gpucbc.waveforms import TF2WFG, TF2 @@ -51,7 +54,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 +72,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): + @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 +110,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)) @@ -164,12 +169,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) From 7b734159688531894dab1fb3c6c25bfbc11b4784 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 15 May 2023 14:00:03 -0400 Subject: [PATCH 2/3] DEV: fix backend testing --- .github/workflows/unit_test.yml | 2 +- gpucbc/__init__.py | 12 ++++---- gpucbc/backend.py | 13 +++++++++ gpucbc/waveforms.py | 51 +++++++++++++-------------------- pyproject.toml | 2 +- test/test_backends.py | 8 +++--- test/test_waveform.py | 8 ++++-- 7 files changed, 50 insertions(+), 46 deletions(-) diff --git a/.github/workflows/unit_test.yml b/.github/workflows/unit_test.yml index 2f74279..62f5963 100644 --- a/.github/workflows/unit_test.yml +++ b/.github/workflows/unit_test.yml @@ -26,7 +26,7 @@ jobs: conda update pip setuptools conda install numpy astropy bilby python-lal python-lalsimulation conda install jax - conda install pytest-cov + conda install pytest-cov parameterized pip install . - name: Test with pytest run: | diff --git a/gpucbc/__init__.py b/gpucbc/__init__.py index 2757093..f4b8641 100644 --- a/gpucbc/__init__.py +++ b/gpucbc/__init__.py @@ -1,5 +1,4 @@ -from . import likelihood, pn, waveforms -from .backend import Backend +from . import backend, likelihood, pn, waveforms def set_backend(numpy): @@ -10,7 +9,8 @@ def set_backend(numpy): jax="jax.scipy", ).get(numpy, None) numpy = dict(jax="jax.numpy").get(numpy, numpy) - BACKEND = Backend(numpy=numpy, scipy=scipy) - likelihood.BACKEND = BACKEND - pn.BACKEND = BACKEND - waveforms.BACKEND = BACKEND + 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 index abf7f5a..aa504c1 100644 --- a/gpucbc/backend.py +++ b/gpucbc/backend.py @@ -4,6 +4,7 @@ class Backend: def __init__(self, numpy, scipy=None): + self.module = numpy try: self.np = import_module(numpy) if scipy is not None: @@ -13,4 +14,16 @@ def __init__(self, numpy, scipy=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/waveforms.py b/gpucbc/waveforms.py index faf13ba..fe0c615 100644 --- a/gpucbc/waveforms.py +++ b/gpucbc/waveforms.py @@ -2,6 +2,7 @@ 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 @@ -13,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. @@ -47,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 ): @@ -84,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, @@ -146,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) @@ -227,7 +216,7 @@ def call_cupy_tf2( frequency_array = B.np.asarray(frequency_array) - h_out_of_band = B.np.zeros(int(B.np.sum(~in_band))) + h_out_of_band = B.np.zeros(B.np.sum(~in_band)) wf = TF2( mass_1, diff --git a/pyproject.toml b/pyproject.toml index be75cbc..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"] diff --git a/test/test_backends.py b/test/test_backends.py index b2fad63..0c25766 100644 --- a/test/test_backends.py +++ b/test/test_backends.py @@ -10,8 +10,8 @@ class TestBackend(unittest.TestCase): def test_setting_jax(self): gpucbc.set_backend("jax") - self.assertEqual(gpucbc.pn.BACKEND.np, jnp) - self.assertEqual(gpucbc.pn.BACKEND.special, jsp) + 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): @@ -19,5 +19,5 @@ def test_unknown_backend_raises_error(self): def test_no_scipy_backend(self): gpucbc.set_backend("bilby") - self.assertEqual(gpucbc.pn.BACKEND.np, bilby) - self.assertEqual(gpucbc.pn.BACKEND.special, None) + self.assertEqual(gpucbc.pn.B.np, bilby) + self.assertEqual(gpucbc.pn.B.special, None) diff --git a/test/test_waveform.py b/test/test_waveform.py index 9070713..cdfbfcf 100644 --- a/test/test_waveform.py +++ b/test/test_waveform.py @@ -17,6 +17,7 @@ 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 @@ -75,9 +76,9 @@ def test_native_phasing(self, backend): 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) + 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): @@ -149,6 +150,7 @@ def test_absolute_overlap(self, backend): ) 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 From 453bd1a114dba15fa4911ead314f9c95efefc3e9 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Mon, 15 May 2023 14:31:21 -0400 Subject: [PATCH 3/3] COMPATIBIILTY: backend waveform compatibility --- .coveragerc | 5 +++++ gpucbc/__init__.py | 3 +++ gpucbc/test/__init__.py | 0 {test => gpucbc/test}/test_backends.py | 0 {test => gpucbc/test}/test_waveform.py | 0 gpucbc/waveforms.py | 5 ++--- 6 files changed, 10 insertions(+), 3 deletions(-) create mode 100644 .coveragerc create mode 100644 gpucbc/test/__init__.py rename {test => gpucbc/test}/test_backends.py (100%) rename {test => gpucbc/test}/test_waveform.py (100%) 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/gpucbc/__init__.py b/gpucbc/__init__.py index f4b8641..92e3984 100644 --- a/gpucbc/__init__.py +++ b/gpucbc/__init__.py @@ -1,6 +1,9 @@ from . import backend, likelihood, pn, waveforms +from ._version import __version__ + + def set_backend(numpy): scipy = dict( diff --git a/gpucbc/test/__init__.py b/gpucbc/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/test_backends.py b/gpucbc/test/test_backends.py similarity index 100% rename from test/test_backends.py rename to gpucbc/test/test_backends.py diff --git a/test/test_waveform.py b/gpucbc/test/test_waveform.py similarity index 100% rename from test/test_waveform.py rename to gpucbc/test/test_waveform.py diff --git a/gpucbc/waveforms.py b/gpucbc/waveforms.py index fe0c615..0de3383 100644 --- a/gpucbc/waveforms.py +++ b/gpucbc/waveforms.py @@ -213,11 +213,8 @@ def call_cupy_tf2( minimum_frequency = waveform_kwargs["minimum_frequency"] in_band = frequency_array >= minimum_frequency - frequency_array = B.np.asarray(frequency_array) - h_out_of_band = B.np.zeros(B.np.sum(~in_band)) - wf = TF2( mass_1, mass_2, @@ -227,7 +224,9 @@ def call_cupy_tf2( lambda_2=lambda_2, luminosity_distance=luminosity_distance, ) + strain = wf(frequency_array[in_band], phi_c=phase) + 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)