QuickStart
==========

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

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

.. code-block:: python

    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.