"""
JaxGB module
============
Fast waveform and response generator for galactic binaries, as seen by the
LISA instrument. Galactic binaries are described by the following 8
parameters: frequency (f0), frequency evolution (fdot), amplitude
(ampl), sky location (ra, dec), polarisation (psi), inclination
(incl) and initial phase (phi0).
Usage
-----
.. code-block:: python
# Catalogue parameter definition and definition of epoch `t_init`.
catalogue_pd = pd.DataFrame(
{
"Frequency": #...,
# other parameters...
"InitialPhase": #...,
}
)
t_init = 1.0e3 # catalogue epoch
# Definition of `JaxGB` simulator, with simulation `t0` (simulation epoch).
t0 = 1.0e4 # simulation t0
myjaxgb = JaxGB(..., t0=t0, ...)
# Use GBObject to create the jaxgb input parameter array, which also
# accounts for the evalution of `f0` and `phi0` for the particular epoch
# of simulation `t0`.
# Step 1: define GBObject, holding all source parameters
catalogue_as_gbo = GBObject.from_pandas_dataframe(catalogue_pd, t_init=t_init)
# Step 2: get parameters in JaxGB input format, evaluated at `t0`
params = catalogue_as_gbo.to_jaxgb_array(t0=t0)
# Step 3: calculate TDI variables in frequency domain
myjaxgb.get_tdi(jnp.array(params), tdi_generation=1.5, tdi_combination="XYZ")
.. autoclass:: JaxGB
:members:
.. autoclass:: JaxGBaccurate
:members:
"""
import jax
import jax.numpy as jnp
import numpy as np
from lisaconstants import c # pylint: disable=no-name-in-module
from lisaorbits import Orbits
from scipy.signal.windows import tukey
from .tdi import to_tdi_combination, to_tdi_generation
[docs]
class JaxGB:
"""Galactic binary fast response generation.
Parameters
----------
orbits
Orbits instance, based on lisaorbits.
t_obs
Duration (s).
t0
Start time (s).
n
Number of points for slow response evaluation (even-valued integer!).
"""
def __init__(
self,
orbits: Orbits,
*,
t_obs: float = 6.2914560e7,
t0: float = 0.0,
n: int = 128,
) -> None:
# duration (s)
self.t_obs = t_obs
# t0 (s)
self.t0 = t0
# number of points for slow response evaluation
self.n = n
# time grid
self.tm = np.linspace(0.0, self.t_obs, num=self.n, endpoint=False)
# time step slow response
self.dtm = self.t_obs / self.n
# initialize orbits and calculate positions of M = 3 spacecraft at time
# of reception
self.orbits = orbits
ltt = self._get_ltt_simple()
self.arm_length = ltt * c
self.position = (
self.orbits.compute_position(self.t0 + self.tm, [1, 2, 3])
.swapaxes(0, 1)
.swapaxes(1, 2)
) # (3, 3, n)
self.fstar = c / (self.arm_length * 2 * np.pi)
def _get_ltt_simple(self) -> jax.Array:
"""Get light-travel times: either pick "L" from orbits file or calculate the an
ltt average using ``compute_ltt``.
Returns
-------
jax.Array of shape ()
LTT from orbits file, or averaged calculated LTT (both are scalar
quantities).
"""
if hasattr(self.orbits, "L"):
return jnp.array(self.orbits.L / c)
if hasattr(self.orbits, "original_attrs"):
if "L" in self.orbits.original_attrs:
return jnp.array(self.orbits.original_attrs["L"] / c)
ltt = self.orbits.compute_ltt(self.t0 + self.tm).T # (6, n)
return jnp.mean(ltt)
[docs]
def get_kmin(self, f0: jax.Array) -> jax.Array:
"""Calculates kmin, the index of the first heterodyned frequency bin in the
full frequency grid. Given the observation time baseline `t_obs`, the
corresponding Fourier frequencies are integer multiples of `1 / t_obs`. The
nearest frequency bin number to `f0` is therefore `rint(f0 * t_obs)`. This
bin is centred on the grid of heterodyned Fourier frequencies, of length `n`,
hence kmin (i.e., the first bin) is shifted by `n/2`.
Note: `f0` needs to be evaluated at simulation epoch `self.t0`, which in general
differs from the catalogue epoch `t_init`. Use the :mod:`params.GBObject` to
calculate source frequencies at simulation epoch `t0` from catalogue frequencies
(specified at `t_init`).
Assuming `params` is the `JAXGB`-compatible array of parameters, and `myjaxgb`
is a `JaxGB` instance, we have:
.. code-block:: python
myjaxgb.get_kmin(f0=params[:, 0])
# `params[:, 0]` needs to contain the catalogue frequencies at the
# simulation reference time `t0`, which is in general different from the
# catalogue epoch `t_init`.
Parameters
----------
f0
Frequency f0 for source(s), at epoch self.t0 (Hz).
Returns
-------
jax.Array of shape (nsrc,)
Array with kmin values (integer).
"""
q = jnp.rint(f0 * self.t_obs)
return jnp.round(q - self.n / 2).astype(np.int32)
[docs]
def get_frequency_grid(self, kmin: jax.Array) -> jax.Array:
"""Return the frequency grid associated to the returned signal.
Parameters
----------
kmin
Index of the first frequency bin.
Returns
-------
jax.Array of shape (nsrc, n)
Array with frequency values (in Hz).
"""
return (1 / self.t_obs) * (jnp.arange(self.n) + kmin[:, None])
def _construct_slow_part(
self,
f0_fdot_phi0: jax.Array,
k: jax.Array,
dpdc: tuple[jax.Array, jax.Array],
pol_tensors: tuple[jax.Array, jax.Array],
) -> jax.Array:
"""This method implements the slow, heterodyned response in time domain. It
accepts the f0, fdot, and phi0 of all the sources, their direction (k),
the hplus/hcross strain amplitudes (Dp, Dc) and polarization tensors
(eplus, ecross). The output is an array of shape (nsrc, 6, n) [other equations
need to be included here].
Parameters
----------
f0_fdot_phi0
array (nsrc, 3) listing f0, fdot, and phi0 values of all sources.
k
array (nsrc, 3) with source k-vectors.
dpdc
tuple with hplus and hcross strain amplitudes (complex-valued).
pol_tensors
tuple with Eplus, Ecross (polarization tensors).
Returns
-------
jax.Array of shape (nsrc, 6, n)
Time-domain slow heterodyned response.
"""
nsrc = len(f0_fdot_phi0)
r = jnp.zeros((4, 3, self.n)) # r convention: 12,13,23,31
r = r.at[0].set(self.position[1] - self.position[0]) # (3, n)
r = r.at[1].set(self.position[2] - self.position[0]) # (3, n)
r = r.at[2].set(self.position[2] - self.position[1]) # (3, n)
r = r.at[3].set(-r[1]) # (3, n)
r /= self.arm_length
# kdotr convention: 12, 21, 13, 31, 23, 32
kdotr = jnp.zeros((nsrc, 6, self.n))
kdotr = kdotr.at[:, ::2].set(jnp.dot(k.T, r[:3]))
kdotr = kdotr.at[:, 1::2].set(-kdotr[:, ::2])
kdotp = jnp.dot(k.T, self.position)
kdotp /= c
f0, fdot, phi0 = f0_fdot_phi0[:, 0], f0_fdot_phi0[:, 1], f0_fdot_phi0[:, 2]
xi = self.tm - kdotp
fi = f0[:, None, None] + fdot[:, None, None] * xi
fonfs = fi / self.fstar # Ratio of true frequency to transfer frequency
### compute transfer f-n
q = jnp.rint(f0 * self.t_obs) # index of nearest Fourier bin
df = 2.0 * np.pi * (q / self.t_obs)
om = 2.0 * np.pi * f0
# Reduce to (3, 3, n).
# A = (eplus @ r * r * dp + ecross @ r * r * dc).sum(axis=2)
idx = jnp.array([0, 2, 3])
eplus, ecross = pol_tensors
aij = (
jnp.swapaxes(jnp.dot(eplus, r[idx]), 1, 2)
* r[idx]
* dpdc[0][:, None, None, None]
+ jnp.swapaxes(jnp.dot(ecross, r[idx]), 1, 2)
* r[idx]
* dpdc[1][:, None, None, None]
)
asum = aij.sum(axis=2) # A convention: 12, 23, 31
arg_s = (phi0[:, None] + (om - df)[:, None] * self.tm[None, :])[:, None, :] + (
np.pi * fdot[:, None, None] * (xi**2)
)
kdotp = om[:, None, None] * kdotp - arg_s
# gs convention: 12, 23, 31, 21, 32, 13
idx = jnp.array([0, 4, 3, 1, 5, 2])
arg_ij = (
0.5 * fonfs[:, [0, 1, 2, 1, 2, 0]] * (1 + kdotr[:, idx])
) # note the plus-sign!
gs = (
0.5j
* (f0[:, None, None] + fdot[:, None, None] * self.tm[None, None, :])
/ self.fstar
* jnp.sinc(arg_ij / np.pi)
* jnp.exp(-1j * (arg_ij + kdotp[:, [0, 1, 2, 1, 2, 0]]))
* asum[:, [0, 1, 2, 0, 1, 2]]
) # last term returns a frequency-dependent sign, related to the windowing
# shift convention: 12, 23, 31, 21, 32, 13 to 12, 23, 31, 13, 32, 21
gs = gs[:, [0, 1, 2, 5, 4, 3]]
return gs # (nsrc, 6, n)
def _get_gs(self, params: jax.Array) -> jax.Array:
"""Calculate slow response in time domain. Generic for single-link and
TDI response calculation.
Parameters
----------
params
Parameter array of source parameters, as defined in :mod:`jaxgb.params`.
Returns
-------
jax.Array of shape (nsrc, 6, n)
Slow response in time domain.
"""
# sign flip phase
params = params.at[:, 7].set(params[:, 7] * -1)
# equatorial latitude to theta
params = params.at[:, 4].set(np.pi / 2 - params[:, 4])
cosiota = jnp.cos(params[:, 6])
cosps = jnp.cos(2 * params[:, 5])
sinps = jnp.sin(2 * params[:, 5])
aplus = params[:, 2] * (1.0 + cosiota * cosiota)
across = -2.0 * params[:, 2] * cosiota
dpdc = (
aplus * cosps - 1.0j * across * sinps,
-aplus * sinps - 1.0j * across * cosps,
)
sinph = jnp.sin(params[:, 3])
cosph = jnp.cos(params[:, 3])
sinth = jnp.sin(params[:, 4])
costh = jnp.cos(params[:, 4])
u = jnp.array([costh * cosph, costh * sinph, -sinth], ndmin=2)
v = jnp.array([sinph, -cosph, jnp.zeros_like(sinph)], ndmin=2)
k = jnp.array([-sinth * cosph, -sinth * sinph, -costh], ndmin=2)
eplus = (v[None, :, :] * v[:, None, :] - u[None, :, :] * u[:, None, :]).T
ecross = (u[None, :, :] * v[:, None, :] + v[None, :, :] * u[:, None, :]).T
pol_tensors = (eplus, ecross)
gs = self._construct_slow_part(params[:, [0, 1, 7]], k, dpdc, pol_tensors)
return gs # (nsrc, 6, n)
def _compute_response(self, gs: jax.Array) -> jax.Array:
"""Compute single-link response from gs.
Parameters
----------
gs
Slow (heterodyned) response in time domain.
Returns
-------
jax.Array of shape (nsrc, 6, n)
Frequency-domain response vector.
"""
yf_slow = jnp.fft.fft(gs, axis=2)
# yf convention: 12, 23, 31, 13, 32, 21
yf = jnp.fft.fftshift(yf_slow, axes=2) # (nsrc, 6, n)
fctr = 0.5 * self.dtm # normalization of (half-sided) fourier transform
return yf * fctr
def _compute_tdi_xyz(
self, gs: jax.Array, f0: jax.Array, fdot: jax.Array
) -> jax.Array:
"""Compute TDI X, Y, Z from gs.
Parameters
----------
gs
Slow (heterodyned) response in time domain.
Returns
-------
jax.Array of shape (nsrc, 3, n)
Frequency-domain TDI response vector.
"""
f = f0[:, None] + fdot[:, None] * self.tm[None, :]
oml = f / np.mean(self.fstar) # constant arm length assumed
soml = jnp.sin(oml)
fctr = jnp.exp(-1.0j * oml)
fctr2 = -2.0j * soml * fctr
# gs convention: 12, 23, 31, 13, 32, 21
xsl = gs[:, 5] - gs[:, 2] + (gs[:, 0] - gs[:, 3]) * fctr
ysl = gs[:, 4] - gs[:, 0] + (gs[:, 1] - gs[:, 5]) * fctr
zsl = gs[:, 3] - gs[:, 1] + (gs[:, 2] - gs[:, 4]) * fctr
xyzsl = fctr2[:, None, :] * jnp.concatenate(
[xsl[:, None, :], ysl[:, None, :], zsl[:, None, :]], axis=1
)
xyzf_slow = jnp.fft.fft(xyzsl, axis=2)
xyzf = jnp.fft.fftshift(xyzf_slow, axes=2)
fctr3 = 0.5 * self.dtm
xyzf *= fctr3
return xyzf
[docs]
def get_link_responses(self, params: jax.Array) -> jax.Array:
"""Calculate the 6 single-link responses in frequency domain.
Parameters
----------
params
Parameter array of shape (nsrc, 8) of source parameters, as
defined in :mod:`jaxgb.params`.
Returns
-------
jax.Array of shape (nsrc, 6, n)
Single-link response vector.
"""
params = params.copy()
gs = self._get_gs(params)
yf = self._compute_response(gs)
return yf
[docs]
def get_tdi(
self,
params: jax.Array,
tdi_generation: float = 1.5,
tdi_combination: str = "XYZ",
) -> tuple[jax.Array, jax.Array, jax.Array]:
"""Calculate TDI XYZ in frequency domain.
Parameters
----------
params
Parameter array of shape (nsrc, 8) of source parameters, as
defined in :mod:`jaxgb.params`.
tdi_generation
TDI generation (either 1.5 or 2).
tdi_combination
TDI combinations (either "XYZ" or "AET").
Returns
-------
tuple of arrays, each of shape (nsrc, n)
TDI frequency-domain response channels (X, Y, Z) or
(A, E, T).
"""
_params = params.copy()
if len(_params.shape) == 1:
_params = _params.reshape(1, -1)
gs = self._get_gs(_params)
tdi = self._compute_tdi_xyz(gs, _params[:, 0], _params[:, 1])
kmin = self.get_kmin(f0=_params[:, 0])
freqs = self.get_frequency_grid(kmin)
tdi = to_tdi_generation(
tdi, freqs, to=tdi_generation, arm_length=self.arm_length
)
tdi = to_tdi_combination(tdi, tdi_combination)
x = tdi[:, 0, :]
y = tdi[:, 1, :]
z = tdi[:, 2, :]
if len(params.shape) == 1:
return x[0], y[0], z[0]
return x, y, z
[docs]
def sum_tdi(
self,
params: jax.Array,
kmin: int | None = None,
kmax: int | None = None,
**kwargs,
) -> tuple[jax.Array, jax.Array, jax.Array]:
r"""Sum TDI response channels for a set of source parameters.
Each source is simulated independently, and the TDI responses
are then summed (after aligning the frequency grids) and
returned.
Parameters
----------
params
Parameter array with shape (nsrc, 8) of source parameters,
as defined in :mod:`jaxgb.params`.
kmin
First index of the frequency grid. If None, takes the
first index of the lowest frequency source.
kmax
Last index of the frequency grid. If None, takes the
last index of the highest frequency source.
kwargs
Extra keywords arguments passed to :meth:`get_tdi`.
Returns
-------
tuple (jax.Array, jax.Array, jax.Array, ), each of shape (nfreq)
TDI frequency-domain response channels (X, Y, Z) or
(A, E, T) [strain].
"""
kmins = self.get_kmin(f0=params[:, 0])
kmin = jnp.min(kmins).item() if kmin is None else kmin
kmax = jnp.max(kmins).item() + self.n if kmax is None else kmax
def _add_to(tdi_one, tdi_sum, _kmin):
def _update_fpoint(idx, val):
_mask = ((idx + _kmin) >= 0) & ((idx + _kmin) < kmax)
return jax.lax.dynamic_update_index_in_dim(
val,
val[:, idx + _kmin] + _mask * tdi_sum[:, idx],
idx + _kmin,
axis=1,
)
return jax.lax.fori_loop(0, self.n, _update_fpoint, tdi_one)
def _add_one(idx, tdi):
_x, _y, _z = self.get_tdi(
jax.lax.dynamic_slice_in_dim(params, idx, 1, axis=0),
**kwargs,
)
return _add_to(tdi, jnp.vstack([_x, _y, _z]), kmins[idx] - kmin)
_tdi = jnp.zeros((3, kmax - kmin), dtype=np.complex128)
_tdi = jax.lax.fori_loop(0, params.shape[0], _add_one, _tdi)
return tuple(_tdi)
[docs]
class JaxGBaccurate(JaxGB):
"""Galactic binary less fast but more accurate response generation. Considers
time-varying arm lengths and reduces edge artefacts (ringing).
Parameters
----------
window
Dimensionless extension factor for slow response calculation, for
windowing purposes. The slow response will be calculated on the domain
[t0 - window * t_obs, t0 + t_obs + window * t_obs], i.e., the extension is
symmetric around the time interval of interest [t0, t0 + t_obs].
window needs to have a strictly positive integer or half-integer value. Note
that using window also modifies the kmin values:
df = 1/(t_obs + 2 * window * t_obs).
"""
def __init__(self, window: float = 0.5, **kwargs) -> None:
super().__init__(**kwargs)
# window value (strictly positive integer or half-integer)
self.window = window
# length of extended time grid (with windowing)
self.t_ext = self.t_obs + 2.0 * self.window * self.t_obs
# time grid with windowing [-window*t_obs, t_obs + 2*window*t_obs]
self.tm = (
np.linspace(
0,
self.t_ext,
num=self.n,
endpoint=False,
)
- self.window * self.t_obs
)
# time step slow response
self.dtm = self.t_ext / self.n
# calculate positions of M = 3 spacecraft at time of emission
ltt = self.orbits.compute_ltt(self.t0 + self.tm).T # (6, n), recalculate
self.position = (
self.orbits.compute_position(self.t0 + self.tm, [1, 2, 3])
.swapaxes(0, 1)
.swapaxes(1, 2)
) # (M, 3, n)
self.positionemi = jnp.array(
[
self.orbits.compute_position(self.t0 + self.tm - ltt[0], [2])
.squeeze()
.T,
self.orbits.compute_position(self.t0 + self.tm - ltt[1], [3])
.squeeze()
.T,
self.orbits.compute_position(self.t0 + self.tm - ltt[2], [1])
.squeeze()
.T,
self.orbits.compute_position(self.t0 + self.tm - ltt[3], [3])
.squeeze()
.T,
self.orbits.compute_position(self.t0 + self.tm - ltt[4], [2])
.squeeze()
.T,
self.orbits.compute_position(self.t0 + self.tm - ltt[5], [1])
.squeeze()
.T,
]
) # (6, 3, n)
self.arm_length = ltt * c
fstar = c / (self.arm_length * 2 * np.pi)
self.fstar = fstar[None, :, :]
# convention: 12, 23, 31, 13, 32, 21
[docs]
def get_kmin(self, f0: jax.Array) -> jax.Array:
"""Calculates kmin, the index of the first heterodyned frequency bin in the
full frequency grid. For windowed response calculations, the observation time
baseline is extended from `t_obs` to `t_ext = t_obs + 2 * window * t_obs`, so
the corresponding Fourier frequencies are integer multiples of `1 / t_ext`.
The nearest frequency bin number to `f0` is therefore `rint(f0 * t_ext)`. This
bin is centred on the grid of heterodyned Fourier frequencies, of length `n`,
hence kmin (i.e., the first bin) is shifted by `n/2`.
If `t_init` is different from `t0`, `f0` needs to be corrected for the linear
drift proportional to `fdot`.
Parameters
----------
f0
Frequency f0 for source(s) (Hz)
Returns
-------
jax.Array of shape (nsrc,)
Array with kmin values (integer)
"""
q = jnp.rint(f0 * self.t_ext)
return jnp.round(q - self.n / 2).astype(np.int32)
[docs]
def get_frequency_grid(self, kmin: jax.Array) -> jax.Array:
"""Return the frequency grid associated to the returned signal.
Parameters
----------
kmin
Index of the first frequency bin.
Returns
-------
jax.Array of shape (nsrc, n)
Array with frequency values (in Hz).
"""
return (1 / self.t_ext) * (jnp.arange(self.n) + kmin[:, None])
def _construct_slow_part(
self,
f0_fdot_phi0: jax.Array,
k: jax.Array,
dpdc: tuple[jax.Array, jax.Array],
pol_tensors: tuple[jax.Array, jax.Array],
) -> jax.Array:
"""This method implements the slow, heterodyned response in time domain. It
accepts the f0, fdot, and phi0 of all the sources, their direction (k),
the hplus/hcross strain amplitudes (Dp, Dc) and polarization tensors
(eplus, ecross). The output is an array of shape (nsrc, 6, n) [other equations
need to be included here].
For this implementation, we include tukey windowing immediately here, which
tapers off the slow response at the edges of the extended domain.
Parameters
----------
f0_fdot_phi0
array (nrsc, 3) listing f0, fdot, and phi0 values of all sources.
k
array (nsrc, 3) with source k-vectors.
dpdc
tuple with hplus and hcross strain amplitudes (complex-valued).
pol_tensors
tuple with Eplus, Ecross (polarization tensors).
Returns
-------
jax.Array of shape (nsrc, 6, n)
Time-domain slow heterodyned response.
"""
r = jnp.zeros((6, 3, self.n)) # r convention: 12, 23, 31, 13, 32, 21
r = r.at[0].set(self.position[0] - self.positionemi[0]) ## [3xn]
r = r.at[1].set(self.position[1] - self.positionemi[1]) ## [3xn]
r = r.at[2].set(self.position[2] - self.positionemi[2]) ## [3xn]
r = r.at[3].set(self.position[0] - self.positionemi[3]) ## [3xn]
r = r.at[4].set(self.position[2] - self.positionemi[4]) ## [3xn]
r = r.at[5].set(self.position[1] - self.positionemi[5]) ## [3xn]
r /= self.arm_length[:, None, :]
# kdotr convention: 12, 23, 31, 13, 32, 21
kdotr = jnp.dot(k.T, r)
kdotp = jnp.dot(k.T, self.position)
kdotp /= c
f0, fdot, phi0 = f0_fdot_phi0[:, 0], f0_fdot_phi0[:, 1], f0_fdot_phi0[:, 2]
xi = self.tm - kdotp
fi = f0[:, None, None] + fdot[:, None, None] * xi
# Ratio of true frequency to transfer frequency
fonfs = fi[:, [0, 1, 2, 0, 2, 1]] / self.fstar
### compute transfer f-n
q = jnp.rint(f0 * self.t_ext) # index of nearest Fourier bin
df = 2.0 * np.pi * (q / self.t_ext)
om = 2.0 * np.pi * f0
# Reduce to (3, 3, n).
# A = (eplus @ r * r * dp + ecross @ r * r * dc).sum(axis=2)
eplus, ecross = pol_tensors
aij = (
jnp.swapaxes(jnp.dot(eplus, r), 1, 2) * r * dpdc[0][:, None, None, None]
+ jnp.swapaxes(jnp.dot(ecross, r), 1, 2) * r * dpdc[1][:, None, None, None]
)
asum = aij.sum(axis=2) # A convention: 12, 23, 31, 13, 32, 21
arg_s = (phi0[:, None] + (om - df)[:, None] * self.tm[None, :])[:, None, :] + (
np.pi * fdot[:, None, None] * (xi**2)
)
kdotp = om[:, None, None] * kdotp - arg_s
# gs convention: 12, 23, 31, 13, 32, 21
idx = jnp.array([0, 1, 2, 0, 2, 1])
arg_ij = 0.5 * fonfs * (1 - kdotr) # note the minus-sign!
gs = (
0.5j
* (f0[:, None, None] + fdot[:, None, None] * self.tm[None, None, :])
/ self.fstar
* jnp.sinc(arg_ij / np.pi)
* jnp.exp(-1j * (arg_ij + kdotp[:, idx]))
* asum
* jnp.exp(
-2.0j * np.pi * q[:, None, None] * self.window * self.t_obs / self.t_ext
)
) # last term returns a frequency-dependent sign, related to the windowing
# finally, apply windowing
w = tukey(len(self.tm), alpha=0.45) ** 4
gs *= w[None, None, :]
return gs # (nsrc, 6, n)