diff --git a/bilby_cython/geometry.pyx b/bilby_cython/geometry.pyx index 2eeead7..c097bb4 100644 --- a/bilby_cython/geometry.pyx +++ b/bilby_cython/geometry.pyx @@ -6,7 +6,13 @@ from .time import greenwich_mean_sidereal_time cdef double CC = 299792458.0 -cpdef time_delay_geocentric(np.ndarray detector1, np.ndarray detector2, double ra, double dec, double time): +cpdef time_delay_geocentric( + np.ndarray[np.float64_t, ndim=1] detector1, + np.ndarray[np.float64_t, ndim=1] detector2, + double ra, + double dec, + double time, +): """ Calculate time delay between two detectors in geocentric coordinates based on XLALArrivaTimeDiff in TimeDelay.c @@ -53,7 +59,12 @@ cpdef time_delay_geocentric(np.ndarray detector1, np.ndarray detector2, double r _GEOCENTER = np.zeros(3, dtype=float) -cpdef time_delay_from_geocenter(np.ndarray detector1, double ra, double dec, double time): +cpdef time_delay_from_geocenter( + np.ndarray[np.float64_t, ndim=1] detector1, + double ra, + double dec, + double time, +): """ Calculate time delay between a detectors and the geocenter based on XLALArrivalTimeDiff in TimeDelay.c @@ -78,7 +89,14 @@ cpdef time_delay_from_geocenter(np.ndarray detector1, double ra, double dec, dou return time_delay_geocentric(detector1, _GEOCENTER, ra, dec, time) -cdef _vectors_for_polarization_tensor(double phi, double theta, double psi): +cdef _vectors_for_polarization_tensor( + double phi, + double theta, + double psi, + double[:] omega_view, + double[:] m_view, + double[:] n_view, +): r""" Compute the three vectors that can be used to construct the different population modes. @@ -129,27 +147,25 @@ cdef _vectors_for_polarization_tensor(double phi, double theta, double psi): omega_view[2] = m_view[0] * n_view[1] - m_view[1] * n_view[0] -m = np.zeros(3) -n = np.zeros(3) -omega = np.zeros(3) -cdef double[:] m_view = m -cdef double[:] n_view = n -cdef double[:] omega_view = omega - - -cpdef _polarization_tensor(double[:, :] output_view, str mode): +cpdef _polarization_tensor( + double[:, :] output_view, + str mode, + double[:] omega_view, + double[:] m_view, + double[:] n_view, +): if mode == 'plus': - _plus(output_view) + _plus(output_view, m_view, n_view) elif mode == 'cross': - _cross(output_view) + _cross(output_view, m_view, n_view) elif mode == 'breathing': - _breathing(output_view) + _breathing(output_view, m_view, n_view) elif mode == 'longitudinal': - _longitudinal(output_view) + _longitudinal(output_view, omega_view) elif mode == 'x': - _x(output_view) + _x(output_view, omega_view, m_view) elif mode == 'y': - _y(output_view) + _y(output_view, omega_view, n_view) else: raise ValueError("{} not a polarization mode!".format(mode)) @@ -182,15 +198,21 @@ cpdef get_polarization_tensor(double ra, double dec, double time, double psi, st """ cdef double gmst, phi, theta - output = np.zeros((3, 3)) + output = np.empty((3, 3)) cdef double[:, :] output_view = output + omega = np.empty(3) + m = np.empty(3) + n = np.empty(3) + cdef double[:] omega_view = omega + cdef double[:] m_view = m + cdef double[:] n_view = n gmst = fmod(greenwich_mean_sidereal_time(time), 2 * pi) phi = ra - gmst theta = pi / 2 - dec - _vectors_for_polarization_tensor(phi, theta, psi) + _vectors_for_polarization_tensor(phi, theta, psi, omega_view, m_view, n_view) - _polarization_tensor(output_view, mode) + _polarization_tensor(output_view, mode, omega_view, m_view, n_view) return output @@ -225,22 +247,28 @@ cpdef get_polarization_tensor_multiple_modes(double ra, double dec, double time, """ cdef double gmst, phi, theta cdef double[:, :] output_view + omega = np.empty(3) + m = np.empty(3) + n = np.empty(3) + cdef double[:] omega_view + cdef double[:] m_view + cdef double[:] n_view output = list() gmst = fmod(greenwich_mean_sidereal_time(time), 2 * pi) phi = ra - gmst theta = pi / 2 - dec - _vectors_for_polarization_tensor(phi, theta, psi) + _vectors_for_polarization_tensor(phi, theta, psi, omega_view, m_view, n_view) for mode in modes: tensor = np.zeros((3, 3)) output_view = tensor - _polarization_tensor(output_view, mode) + _polarization_tensor(output_view, mode, omega_view, m_view, n_view) output.append(tensor) return output -cdef _plus(double[:, :] output): +cdef _plus(double[:, :] output, double[:] m_view, double[:] n_view): cdef int ii, jj for ii in range(3): @@ -250,7 +278,7 @@ cdef _plus(double[:, :] output): output[jj][ii] = output[ii][jj] -cdef _breathing(double[:, :] output): +cdef _breathing(double[:, :] output, double[:] m_view, double[:] n_view): cdef int ii, jj for ii in range(3): @@ -260,7 +288,7 @@ cdef _breathing(double[:, :] output): output[jj][ii] = output[ii][jj] -cdef _longitudinal(double[:, :] output): +cdef _longitudinal(double[:, :] output, double[:] omega_view): cdef int ii, jj for ii in range(3): @@ -280,19 +308,19 @@ cdef _symmetric_response(double[:, :] output, double[:] input_1, double[:] input output[jj][ii] = output[ii][jj] -cdef _cross(double[:, :] output): +cdef _cross(double[:, :] output, double[:] m_view, double[:] n_view): _symmetric_response(output, m_view, n_view) -cdef _x(double[:, :] output): +cdef _x(double[:, :] output, double[:] omega_view, double[:] m_view): _symmetric_response(output, m_view, omega_view) -cdef _y(double[:, :] output): +cdef _y(double[:, :] output, double[:] omega_view, double[:] n_view): _symmetric_response(output, n_view, omega_view) -cpdef three_by_three_matrix_contraction(np.ndarray x, np.ndarray y): +cpdef three_by_three_matrix_contraction(np.ndarray[np.float64_t, ndim=2] x, np.ndarray[np.float64_t, ndim=2] y): """ Doubly contract two 3x3 input matrices following Einstein summation. @@ -323,7 +351,7 @@ cpdef three_by_three_matrix_contraction(np.ndarray x, np.ndarray y): return output -cpdef detector_tensor(np.ndarray x, np.ndarray y): +cpdef detector_tensor(np.ndarray[np.float64_t, ndim=1] x, np.ndarray[np.float64_t, ndim=1] y): r""" Compute the detector tensor given the two unit arm vectors. @@ -469,7 +497,7 @@ cpdef rotation_matrix_from_delta(delta_x): return rotation -cpdef zenith_azimuth_to_theta_phi(double zenith, double azimuth, np.ndarray delta_x): +cpdef zenith_azimuth_to_theta_phi(double zenith, double azimuth, np.ndarray[np.float64_t, ndim=1] delta_x): """ Convert from the 'detector frame' to the Earth frame. diff --git a/bilby_cython/test/test_concurrent.py b/bilby_cython/test/test_concurrent.py new file mode 100644 index 0000000..ffb4252 --- /dev/null +++ b/bilby_cython/test/test_concurrent.py @@ -0,0 +1,36 @@ +import concurrent.futures +import sys + +import numpy as np +import pytest +from bilby_cython import geometry + + +if hasattr(sys, '_is_gil_enabled'): + GIL = sys._is_gil_enabled() +else: + GIL = True + + +@pytest.mark.skipif(GIL, reason="Thread test only works with freethreaded python") +def test_polarization_tensor_threadsafe(): + """ + A basic test of thread safety for the polarization tensor calculation. + Previously, this was not thread safe due to the use of global variables + to store intermediate results. + """ + + def dummy_func(val): + return geometry.get_polarization_tensor(*val, "plus") + + values = np.random.uniform(0, 1, (10000, 4)) + + truths = np.array([geometry.get_polarization_tensor(*val, "plus") for val in values]) + + results = truths.copy() + with concurrent.futures.ThreadPoolExecutor(max_workers=15) as executor: + jobs = {executor.submit(dummy_func, val): ii for ii, val in enumerate(values)} + for job in concurrent.futures.as_completed(jobs): + results[jobs[job]] = job.result() + + assert np.allclose(truths, results) diff --git a/setup.py b/setup.py index db02ebe..fe2469e 100644 --- a/setup.py +++ b/setup.py @@ -8,6 +8,8 @@ class LazyImportBuildExtCmd(build_ext): def finalize_options(self): from Cython.Build import cythonize + from Cython.Compiler.Version import version as cython_version + from packaging.version import Version compiler_directives = dict( language_level=3, @@ -22,6 +24,8 @@ def finalize_options(self): annotate = True else: annotate = False + if Version(cython_version) >= Version("3.1.0a1"): + compiler_directives["freethreading_compatible"] = True self.distribution.ext_modules = cythonize( self.distribution.ext_modules, compiler_directives=compiler_directives,