"""
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 .tdi import to_tdi_combination, to_tdi_generation
from .window import get_window
[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!).
window
Time-domain window, as a jax.Array of shape (n,), str or tuple following on
scipy.signal.get_window convention. Default is scipy.signal.windows.boxcar.
"""
def __init__(
self,
orbits: Orbits,
*,
t_obs: float = 6.2914560e7,
t0: float = 0.0,
n: int = 128,
window: str | tuple | jax.Array = "boxcar",
) -> 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)
r = jnp.zeros((3, 3, self.n)) # r convention: 12,23,31
r = r.at[0].set(self.position[0] - self.position[1]) # (3, n)
r = r.at[1].set(self.position[1] - self.position[2]) # (3, n)
r = r.at[2].set(self.position[2] - self.position[0]) # (3, n)
r /= self.arm_length
self.r = r
# time domain window function (if any)
self.window = get_window(window, size=n)
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,
apac_ssb: tuple[jax.Array, jax.Array],
epec_ssb: 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 (Ap, Ac) and polarisation 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.
apac_ssb
tuple with hplus and hcross strain amplitudes (complex-valued).
epec_ssb
tuple with Eplus, Ecross (polarisation tensors).
Returns
-------
jax.Array of shape (nsrc, 6, n)
Time-domain slow heterodyned response.
"""
# unpack f0, fdot, phi0
f0, fdot, phi0 = f0_fdot_phi0[:, 0], f0_fdot_phi0[:, 1], f0_fdot_phi0[:, 2]
# kdotp and xi
kdotp = jnp.dot(k.T, self.position)
kdotp /= c
xi = self.tm - kdotp
# kdotr convention: 12, 23, 31, 13, 32, 21, assuming 13=-31, 32=-23, 21=-12
nsrc = len(f0)
kdotr = jnp.zeros((nsrc, 6, self.n))
kdotr = kdotr.at[:, :3].set(jnp.dot(k.T, self.r))
kdotr = kdotr.at[:, 3:].set(-kdotr[:, 2::-1])
# frequency-related quantities
fi = f0[:, None, None] + fdot[:, None, None] * xi
fonfs = fi[:, [0, 1, 2, 0, 2, 1]] / self.fstar # true / transfer frequency
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
# A_ij: compute A_ij from r (12, 23, 31), assuming 21=12, 32=23, 13=31
# A_ij = (eplus @ r * r * dp + ecross @ r * r * dc).sum(axis=2)
a_ij = (
jnp.swapaxes(jnp.dot(epec_ssb[0], self.r), 1, 2)
* self.r
* apac_ssb[0][:, None, None, None]
+ jnp.swapaxes(jnp.dot(epec_ssb[1], self.r), 1, 2)
* self.r
* apac_ssb[1][:, None, None, None]
).sum(axis=2)
# compute second exponential argument
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
# G_ij: slow response calculation
alpha_ij = 0.5 * fonfs * (1 - kdotr)
gs = (
-0.5j
* a_ij[:, [0, 1, 2, 2, 1, 0]]
* fonfs
* jnp.sinc(alpha_ij / np.pi)
* jnp.exp(-1j * (alpha_ij + kdotp[:, [0, 1, 2, 0, 2, 1]]))
)
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.
"""
# sines and cosines of angles
sinra = jnp.sin(params[:, 3])
cosra = jnp.cos(params[:, 3])
sindec = jnp.sin(params[:, 4])
cosdec = jnp.cos(params[:, 4])
cosiota = jnp.cos(params[:, 6])
cos2psi = jnp.cos(2 * params[:, 5])
sin2psi = jnp.sin(2 * params[:, 5])
# aplus, across in source frame (real) and SSB frame (complex)
aplus = -params[:, 2] * (1.0 + cosiota * cosiota)
across = -2.0 * params[:, 2] * cosiota
apac_ssb = (
aplus * cos2psi + 1.0j * across * sin2psi,
aplus * sin2psi - 1.0j * across * cos2psi,
)
# (u, v, k): orthonormal triad
u = jnp.array([sinra, -cosra, jnp.zeros_like(sinra)], ndmin=2)
v = jnp.array([-sindec * cosra, -sindec * sinra, cosdec], ndmin=2)
k = jnp.array([-cosdec * cosra, -cosdec * sinra, -sindec], ndmin=2)
# polarisation tensors associated with SSB frame
eplus = (u[None, :, :] * u[:, None, :] - v[None, :, :] * v[:, None, :]).T
ecross = (u[None, :, :] * v[:, None, :] + v[None, :, :] * u[:, None, :]).T
epec_ssb = (eplus, ecross)
# g_ij calculation and windowing
gs = self._construct_slow_part(params[:, [0, 1, 7]], k, apac_ssb, epec_ssb)
gs = gs * self.window[None, None, :]
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_link_responses(
self,
params: jax.Array,
kmin: int | None = None,
kmax: int | None = None,
) -> jax.Array:
"""Sum single-link response channels for a set of source parameters.
Each source is simulated independently, and the link 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.
Returns
-------
jax.Array of shape (6, nfreq)
Single-link frequency-domain response.
"""
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(response_one, response_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 * response_sum[:, idx],
idx + _kmin,
axis=1,
)
return jax.lax.fori_loop(0, self.n, _update_fpoint, response_one)
def _add_one(idx, link_responses):
_res = self.get_link_responses(
jax.lax.dynamic_slice_in_dim(params, idx, 1, axis=0),
)
return _add_to(link_responses, _res[0], kmins[idx] - kmin)
_response_fd = jnp.zeros((6, kmax - kmin), dtype=np.complex128)
_response_fd = jax.lax.fori_loop(0, params.shape[0], _add_one, _response_fd)
return _response_fd
[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
----------
extension
Dimensionless extension factor for slow response calculation,
for windowing purposes. The slow response will be calculated
on the domain [t0 - extension * t_obs, t0 + t_obs +
extension * t_obs], i.e., the extension is symmetric
around the time interval of interest [t0, t0 + t_obs].
extension needs to have a strictly positive integer or
half-integer value. Note that using extension also
modifies the kmin values: df = 1/(t_obs + 2 * extension *
t_obs).
window
Time-domain window, as a jax.Array of shape (n,), or None for
a rectangular window or str or tuple following on
scipy.signal.get_window convention. Default is tukey(size, alpha=0.45) ** 4
"""
def __init__(
self,
extension: float = 0.5,
window: str | tuple | jax.Array = ("tukey4", 0.45),
**kwargs,
) -> None:
super().__init__(window=window, **kwargs)
# extension value (strictly positive integer or half-integer)
self.extension = extension
# length of extended time grid
self.t_ext = self.t_obs + 2.0 * self.extension * self.t_obs
# time grid with windowing [-extension*t_obs, t_obs + extension*t_obs]
self.tm = (
np.linspace(
0,
self.t_ext,
num=self.n,
endpoint=False,
)
- self.extension * self.t_obs
)
# time step slow response on extended time grid
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
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]) # (3, n)
r = r.at[1].set(self.position[1] - self.positionemi[1]) # (3, n)
r = r.at[2].set(self.position[2] - self.positionemi[2]) # (3, n)
r = r.at[3].set(self.position[0] - self.positionemi[3]) # (3, n)
r = r.at[4].set(self.position[2] - self.positionemi[4]) # (3, n)
r = r.at[5].set(self.position[1] - self.positionemi[5]) # (3, n)
r /= self.arm_length[:, None, :]
self.r = r
[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,
apac_ssb: tuple[jax.Array, jax.Array],
epec_ssb: 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 (Ap, Ac) and polarisation tensors
(eplus, ecross). The output is an array of shape (nsrc, 6, n) [other equations
need to be included here].
Compared to JaxGB, this implementation drops the assumption that link vectors
on one arm (e.g., r12 vs. r21) have identical magnitudes.
Parameters
----------
f0_fdot_phi0
array (nrsc, 3) listing f0, fdot, and phi0 values of all sources.
k
array (nsrc, 3) with source k-vectors.
apac_ssb
tuple with hplus and hcross strain amplitudes (complex-valued).
epec_ssb
tuple with Eplus, Ecross (polarisation tensors).
Returns
-------
jax.Array of shape (nsrc, 6, n)
Time-domain slow heterodyned response.
"""
# unpack f0, fdot, phi0
f0, fdot, phi0 = f0_fdot_phi0[:, 0], f0_fdot_phi0[:, 1], f0_fdot_phi0[:, 2]
# kdotp and xi
kdotp = jnp.dot(k.T, self.position)
kdotp /= c
xi = self.tm - kdotp
# kdotr convention: 12, 23, 31, 13, 32, 21
kdotr = jnp.dot(k.T, self.r)
# frequency-related quantities
fi = f0[:, None, None] + fdot[:, None, None] * xi
fonfs = fi[:, [0, 1, 2, 0, 2, 1]] / self.fstar # true / transfer frequency
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
# A_ij: compute A_ij from r (12, 23, 31, 13, 32, 21)
# A_ij = (eplus @ r * r * dp + ecross @ r * r * dc).sum(axis=2)
a_ij = (
jnp.swapaxes(jnp.dot(epec_ssb[0], self.r), 1, 2)
* self.r
* apac_ssb[0][:, None, None, None]
+ jnp.swapaxes(jnp.dot(epec_ssb[1], self.r), 1, 2)
* self.r
* apac_ssb[1][:, None, None, None]
).sum(axis=2)
# compute second exponential argument
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
# G_ij: slow response calculation
alpha_ij = 0.5 * fonfs * (1 - kdotr)
factor = self.extension / (1 + 2 * self.extension)
gs = (
-0.5j
* a_ij
* fonfs
* jnp.sinc(alpha_ij / np.pi)
* jnp.exp(-1j * (alpha_ij + kdotp[:, [0, 1, 2, 0, 2, 1]]))
* jnp.exp(-2.0j * np.pi * q[:, None, None] * factor)
)
return gs # (nsrc, 6, n)