Jaxgb Reference

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

# 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")

class jaxgb.jaxgb.JaxGB(orbits: Orbits, *, t_obs: float = 62914560.0, t0: float = 0.0, n: int = 128)[source]

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!).

get_frequency_grid(kmin: Array) → Array[source]

Return the frequency grid associated to the returned signal.

Parameters:: kmin – Index of the first frequency bin.
Returns:: Array with frequency values (in Hz).
Return type:: jax.Array of shape (nsrc, n)

get_kmin(f0: Array) → Array[source]

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 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:: Array with kmin values (integer).
Return type:: jax.Array of shape (nsrc,)

get_link_responses(params: Array) → Array[source]

Calculate the 6 single-link responses in frequency domain.

Parameters:: params – Parameter array of shape (nsrc, 8) of source parameters, as defined in jaxgb.params.
Returns:: Single-link response vector.
Return type:: jax.Array of shape (nsrc, 6, n)

get_tdi(params: Array, tdi_generation: float = 1.5, tdi_combination: str = 'XYZ') → tuple[Array, Array, Array][source]

Calculate TDI XYZ in frequency domain.

Parameters:

params – Parameter array of shape (nsrc, 8) of source parameters, as defined in jaxgb.params.
tdi_generation – TDI generation (either 1.5 or 2).
tdi_combination – TDI combinations (either “XYZ” or “AET”).

Returns:

TDI frequency-domain response channels (X, Y, Z) or (A, E, T).

Return type:

tuple of arrays, each of shape (nsrc, n)

sum_tdi(params: Array, kmin: int | None = None, kmax: int | None = None, **kwargs) → tuple[Array, Array, Array][source]

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 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 get_tdi().

Returns:

TDI frequency-domain response channels (X, Y, Z) or (A, E, T) [strain].

Return type:

tuple (jax.Array, jax.Array, jax.Array, ), each of shape (nfreq)

class jaxgb.jaxgb.JaxGBaccurate(window: float = 0.5, **kwargs)[source]

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).

get_frequency_grid(kmin: Array) → Array[source]

Return the frequency grid associated to the returned signal.

Parameters:: kmin – Index of the first frequency bin.
Returns:: Array with frequency values (in Hz).
Return type:: jax.Array of shape (nsrc, n)

get_kmin(f0: Array) → Array[source]

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:: Array with kmin values (integer)
Return type:: jax.Array of shape (nsrc,)