gwsnr.jax
Submodules
Package Contents
Functions
|
Function to calculate the chirp time from minimum frequency to last stable orbit (JAX implementation). |
|
Function to calculate the plus polarization antenna response for gravitational wave detection (JAX implementation). |
|
Function to calculate the cross polarization antenna response for gravitational wave detection (JAX implementation). |
|
Function to calculate the antenna response for multiple detectors and sources (JAX implementation). |
|
Calculate interpolated signal-to-noise ratio (SNR) for aligned spin gravitational wave signals using JAX. |
|
Calculate interpolated signal-to-noise ratio (SNR) for aligned spin gravitational wave signals using JAX. |
- gwsnr.jax.findchirp_chirptime_jax(m1, m2, fmin)[source]
Function to calculate the chirp time from minimum frequency to last stable orbit (JAX implementation).
Time taken from f_min to f_lso (last stable orbit). 3.5PN in fourier phase considered.
- Parameters:
- m1float
Mass of the first body in solar masses.
- m2float
Mass of the second body in solar masses.
- fminfloat
Lower frequency cutoff in Hz.
- Returns:
- chirp_timefloat
Time taken from f_min to f_lso (last stable orbit frequency) in seconds.
Notes
Calculates chirp time using 3.5PN approximation for gravitational wave Fourier phase. The time represents frequency evolution from fmin to last stable orbit frequency. Uses post-Newtonian expansion coefficients optimized for efficient JAX computation. JAX implementation supports automatic differentiation and GPU acceleration.
- gwsnr.jax.antenna_response_plus(ra, dec, time, psi, detector_tensor)[source]
Function to calculate the plus polarization antenna response for gravitational wave detection (JAX implementation).
- Parameters:
- rafloat
Right ascension of the source in radians.
- decfloat
Declination of the source in radians.
- timefloat
GPS time of the source in seconds.
- psifloat
Polarization angle of the source in radians.
- detector_tensorjax.numpy.ndarray
Detector tensor for the detector (3x3 matrix).
- Returns:
- antenna_response_plusfloat
Plus polarization antenna response of the detector.
Notes
Computes the plus polarization antenna response by calculating the Frobenius inner product between the detector tensor and the plus polarization tensor. The polarization tensor is determined by the source location (ra, dec), observation time, and polarization angle (psi). JAX implementation provides automatic differentiation for parameter estimation workflows.
- gwsnr.jax.antenna_response_cross(ra, dec, time, psi, detector_tensor)[source]
Function to calculate the cross polarization antenna response for gravitational wave detection (JAX implementation).
- Parameters:
- rafloat
Right ascension of the source in radians.
- decfloat
Declination of the source in radians.
- timefloat
GPS time of the source in seconds.
- psifloat
Polarization angle of the source in radians.
- detector_tensorjax.numpy.ndarray
Detector tensor for the detector (3x3 matrix).
- Returns:
- antenna_response_crossfloat
Cross polarization antenna response of the detector.
Notes
Computes the cross polarization antenna response by calculating the Frobenius inner product between the detector tensor and the cross polarization tensor. The polarization tensor is determined by the source location (ra, dec), observation time, and polarization angle (psi). JAX implementation provides automatic differentiation for parameter estimation workflows.
- gwsnr.jax.antenna_response_array(ra, dec, time, psi, detector_tensor)[source]
Function to calculate the antenna response for multiple detectors and sources (JAX implementation).
- Parameters:
- rajax.numpy.ndarray
Array of right ascension values for sources in radians.
- decjax.numpy.ndarray
Array of declination values for sources in radians.
- timejax.numpy.ndarray
Array of GPS times for sources in seconds.
- psijax.numpy.ndarray
Array of polarization angles for sources in radians.
- detector_tensorjax.numpy.ndarray
Detector tensor array for multiple detectors (n×3×3 matrix), where n is the number of detectors.
- Returns:
- Fpjax.numpy.ndarray
Plus polarization antenna response array with shape (n_detectors, n_sources).
- Fcjax.numpy.ndarray
Cross polarization antenna response array with shape (n_detectors, n_sources).
Notes
Computes antenna responses for both plus and cross polarizations across multiple detectors and source parameters simultaneously. Uses JAX’s vmap for efficient vectorized computation with automatic differentiation support. Each antenna response is calculated using the Frobenius inner product between detector tensors and polarization tensors derived from source sky location and polarization angle. Optimized for GPU acceleration and gradient-based optimization.
- gwsnr.jax.get_interpolated_snr_aligned_spins_jax(mass_1, mass_2, luminosity_distance, theta_jn, psi, geocent_time, ra, dec, a_1, a_2, detector_tensor, snr_partialscaled, ratio_arr, mtot_arr, a1_arr, a_2_arr, batch_size=100000)[source]
Calculate interpolated signal-to-noise ratio (SNR) for aligned spin gravitational wave signals using JAX. This function computes the SNR for gravitational wave signals with aligned spins across multiple detectors using 4D cubic spline interpolation. It calculates the effective distance, partial SNR, and combines results from multiple detectors to produce the effective SNR.
- Parameters:
- mass_1jax.numpy.ndarray
Primary mass of the binary system in solar masses.
- mass_2jax.numpy.ndarray
Secondary mass of the binary system in solar masses.
- luminosity_distancejax.numpy.ndarray
Luminosity distance to the source in Mpc.
- theta_jnjax.numpy.ndarray
Inclination angle between the orbital angular momentum and line of sight in radians.
- psijax.numpy.ndarray
Polarization angle in radians.
- geocent_timejax.numpy.ndarray
GPS time of coalescence at the geocenter in seconds.
- rajax.numpy.ndarray
Right ascension of the source in radians.
- decjax.numpy.ndarray
Declination of the source in radians.
- a_1jax.numpy.ndarray
Dimensionless spin magnitude of the primary black hole.
- a_2jax.numpy.ndarray
Dimensionless spin magnitude of the secondary black hole.
- detector_tensorjax.numpy.ndarray
Detector tensor array containing detector response information. Shape: (n_detectors, …)
- snr_partialscaledjax.numpy.ndarray
Pre-computed scaled partial SNR values for interpolation. Shape: (n_detectors, …)
- ratio_arrjax.numpy.ndarray
Mass ratio grid points for interpolation (q = m2/m1).
- mtot_arrjax.numpy.ndarray
Total mass grid points for interpolation.
- a1_arrjax.numpy.ndarray
Primary spin grid points for interpolation.
- a_2_arrjax.numpy.ndarray
Secondary spin grid points for interpolation.
- Returns:
- snrjax.numpy.ndarray
SNR values for each detector. Shape: (n_detectors, n_samples)
- snr_effectivejax.numpy.ndarray
Effective SNR combining all detectors. Shape: (n_samples,)
- snr_partial_jax.numpy.ndarray
Interpolated partial SNR values for each detector. Shape: (n_detectors, n_samples)
- d_effjax.numpy.ndarray
Effective distance for each detector accounting for antenna response. Shape: (n_detectors, n_samples)
Notes
Uses 4D cubic spline interpolation for efficient SNR calculation
Assumes aligned spins (no precession)
Effective SNR is calculated as sqrt(sum(SNR_i^2)) across detectors
Chirp mass and inclination-dependent factors are computed analytically
- gwsnr.jax.get_interpolated_snr_no_spins_jax(mass_1, mass_2, luminosity_distance, theta_jn, psi, geocent_time, ra, dec, a_1, a_2, detector_tensor, snr_partialscaled, ratio_arr, mtot_arr, a1_arr, a_2_arr, batch_size=100000)[source]
Calculate interpolated signal-to-noise ratio (SNR) for aligned spin gravitational wave signals using JAX. This function computes the SNR for gravitational wave signals with aligned spins across multiple detectors using 4D cubic spline interpolation. It calculates the effective distance, partial SNR, and combines results from multiple detectors to produce the effective SNR.
- Parameters:
- mass_1jax.numpy.ndarray
Primary mass of the binary system in solar masses.
- mass_2jax.numpy.ndarray
Secondary mass of the binary system in solar masses.
- luminosity_distancejax.numpy.ndarray
Luminosity distance to the source in Mpc.
- theta_jnjax.numpy.ndarray
Inclination angle between the orbital angular momentum and line of sight in radians.
- psijax.numpy.ndarray
Polarization angle in radians.
- geocent_timejax.numpy.ndarray
GPS time of coalescence at the geocenter in seconds.
- rajax.numpy.ndarray
Right ascension of the source in radians.
- decjax.numpy.ndarray
Declination of the source in radians.
- a_1jax.numpy.ndarray
Dimensionless spin magnitude of the primary black hole.
- a_2jax.numpy.ndarray
Dimensionless spin magnitude of the secondary black hole.
- detector_tensorjax.numpy.ndarray
Detector tensor array containing detector response information. Shape: (n_detectors, …)
- snr_partialscaledjax.numpy.ndarray
Pre-computed scaled partial SNR values for interpolation. Shape: (n_detectors, …)
- ratio_arrjax.numpy.ndarray
Mass ratio grid points for interpolation (q = m2/m1).
- mtot_arrjax.numpy.ndarray
Total mass grid points for interpolation.
- a1_arrjax.numpy.ndarray
Primary spin grid points for interpolation.
- a_2_arrjax.numpy.ndarray
Secondary spin grid points for interpolation.
- Returns:
- snrjax.numpy.ndarray
SNR values for each detector. Shape: (n_detectors, n_samples)
- snr_effectivejax.numpy.ndarray
Effective SNR combining all detectors. Shape: (n_samples,)
- snr_partial_jax.numpy.ndarray
Interpolated partial SNR values for each detector. Shape: (n_detectors, n_samples)
- d_effjax.numpy.ndarray
Effective distance for each detector accounting for antenna response. Shape: (n_detectors, n_samples)
Notes
Uses 4D cubic spline interpolation for efficient SNR calculation
Assumes aligned spins (no precession)
Effective SNR is calculated as sqrt(sum(SNR_i^2)) across detectors
Chirp mass and inclination-dependent factors are computed analytically