QuickStart

In short, JAXGB aims at bringing together several separate evolutions.

Below, a minimal example of how to use JAXGB is shown.

import jax
import lisaorbits
import numpy as np
import pandas as pd
import jax.numpy as jnp

from jaxgb.jaxgb import JaxGB
from jaxgb.params import GBObject

# always configure jax to use 64-bit precision!
jax.config.update("jax_enable_x64", True)

# Catalogue parameter definition and definition of epoch `t_init`.
ncol = 10  # number of objects
catalogue_pd = pd.DataFrame(
    {
        "Frequency": 1e-3 * np.abs(np.random.randn(ncol)),
        "FrequencyDerivative": 1e-13 * np.abs(np.random.randn(ncol)),
        "Amplitude": 1e-22 * np.abs(np.random.randn(ncol)),
        "Inclination": np.arccos(np.random.uniform(size=ncol, low=-1.0, high=1.0)),
        "RightAscension": np.random.uniform(size=ncol, low=0, high=2 * np.pi),
        "Declination": np.random.uniform(size=ncol, low=-np.pi / 2, high=np.pi / 2),
        "Polarization": np.random.uniform(size=ncol, low=0, high=2 * np.pi),
        "InitialPhase": np.random.uniform(size=ncol, low=0, high=2 * np.pi),
    }
)
t_init = 1.0e3  # catalogue epoch
t_obs = 62914560.0  # approximately 2 years

# orbits definition
orbits = lisaorbits.EqualArmlengthOrbits()

# Definition of `JaxGB` simulator, with simulation `t0` (simulation epoch).
t0 = 1.0e4  # simulation t0
myjaxgb = JaxGB(orbits, t_obs=t_obs, t0=t0, n=2000)

# 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. The result is a tuple (X, Y, Z),
# each of shape (nsrc, n)
tdi_fd = myjaxgb.get_tdi(jnp.array(params), tdi_generation=1.5, tdi_combination="XYZ")

This will generate the TDI response in frequency domain, for 10 objects.