"""
JaxGB parameter module
======================
This is a module clarifying the parameters (names, units, and LaTeX notation) and
their ordering in JaxGB. It also enables conversions to/from the parameterisation
used in other tools, by defining a `GBObject` containing all necessary parameters.
.. autoclass:: GBObject
:members:
"""
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from astropy.coordinates import SkyCoord
PARAM_NAMES = np.array(
[
"Frequency",
"Frequency derivative",
"Amplitude",
"Right ascension",
"Declination",
"Polarisation angle",
"Inclination",
"Initial phase",
]
)
"""Array of ordered parameter names."""
PARAM_LATEX = np.array(
[
"$f_0$",
r"$\dot{f}$",
"$A$",
r"$\mathrm{RA}$",
r"$\mathrm{dec}$",
r"$\psi$",
r"$\iota$",
r"$\phi_0$",
]
)
"""Array of ordered Latex-formatted parameter symbols."""
PARAM_UNITS = np.array(
[
r"$\mathrm{Hz}$",
r"$\mathrm{Hz}/\mathrm{s}$",
"",
r"$\mathrm{rad}$",
r"$\mathrm{rad}$",
r"$\mathrm{rad}$",
r"$\mathrm{rad}$",
r"$\mathrm{rad}$",
]
)
"""Array of ordered parameter expected units."""
[docs]
class GBObject:
"""Circular Galactic binary parameter set.
GBObject brings together a canonical parameterisation of a set of circular
Galactic Binaries. It enables conversion into the parameterisation used by
different (circular) Galactic Binary response codes, like LISA GW Response, LDC
tools, fastgb, and gbgpu. At the moment, we ignore the inclusion of a second
derivative of the frequency.
Parameters
----------
f0
Gravitational-wave signal frequency at epoch [Hz].
fdot
Gravitational-wave signal frequency derivative [Hz/s].
A
Gravitational-wave signal amplitude [strain].
ra
Right ascension (equatorial longitude) [rad].
dec
Declination, equatorial latitude [rad].
psi
Polarisation angle [rad].
iota
Inclination angle [rad].
phi0
Initial phase at epoch [rad].
t_init
Epoch used to define the initial phase phi0 and the frequency f0, in TCB [s].
"""
def __init__(
self,
*,
f0: np.ndarray,
fdot: np.ndarray,
A: np.ndarray,
ra: np.ndarray,
dec: np.ndarray,
psi: np.ndarray,
iota: np.ndarray,
phi0: np.ndarray,
t_init: float,
) -> None:
self.f0 = f0
"""Gravitational-wave signal frequency at epoch [Hz]."""
self.fdot = fdot
"""Gravitational-wave signal frequency derivative [Hz/s]."""
self.A = A
"""Gravitational-wave signal amplitude [strain]."""
self.ra = ra
"""Right ascension (equatorial longitude) [rad]."""
self.dec = dec
"""Declination, equatorial latitude [rad]."""
self.psi = psi
"""Polarisation angle [rad]."""
self.iota = iota
"""Inclination angle [rad]."""
self.phi0 = phi0
"""Initial phase at epoch [rad]."""
self.t_init = t_init
"""Epoch used to define the initial phase phi0 and the frequency f0, in TCB
[s]."""
[docs]
@classmethod
def from_jaxgb_params(cls, jaxgb_params: np.ndarray, t_init: float) -> "GBObject":
"""Instantiate GB from a JaxGB params array.
Parameters
----------
jaxgb_params : array-like of shape (nsrc, 8)
Parameters, as defined in :attr:`PARAM_NAMES` and units
as defined in :attr:`PARAM_UNITS`.
t_init : float
Epoch used to define the initial phase phi0 and the frequency f0, in TCB
[s].
Returns
-------
A new GB instance.
"""
gb_object = cls(
f0=jaxgb_params[:, 0],
fdot=jaxgb_params[:, 1],
A=jaxgb_params[:, 2],
ra=jaxgb_params[:, 3],
dec=jaxgb_params[:, 4],
psi=jaxgb_params[:, 5],
iota=jaxgb_params[:, 6],
phi0=jaxgb_params[:, 7],
t_init=t_init,
)
return gb_object
[docs]
@classmethod
def from_pandas_dataframe(
cls, pandas_df: pd.DataFrame, t_init: float
) -> "GBObject":
"""Instantiate GB from a Pandas dataframe.
Parameters
----------
pandas_df
Pandas dataframe of parameters.
t_init : float
Epoch used to define the initial phase phi0 and the frequency f0, in TCB
[s].
Returns
-------
A new GB instance.
"""
gb_object = cls(
f0=pandas_df["Frequency"].to_numpy(),
fdot=pandas_df["FrequencyDerivative"].to_numpy(),
A=pandas_df["Amplitude"].to_numpy(),
ra=pandas_df["RightAscension"].to_numpy(),
dec=pandas_df["Declination"].to_numpy(),
psi=pandas_df["Polarization"].to_numpy(),
iota=pandas_df["Inclination"].to_numpy(),
phi0=pandas_df["InitialPhase"].to_numpy(),
t_init=t_init,
)
return gb_object
[docs]
@classmethod
def from_records(cls, rec_array: np.recarray | dict, t_init: float) -> "GBObject":
"""Instantiate GB from a numpy record array.
Parameters
----------
rec_array
Numpy record array or dictionnary of parameters.
t_init : float
Epoch used to define the initial phase phi0 and the frequency f0, in TCB
[s].
Returns
-------
A new GB instance.
"""
gb_object = cls(
f0=rec_array["Frequency"],
fdot=rec_array["FrequencyDerivative"],
A=rec_array["Amplitude"],
ra=rec_array["RightAscension"],
dec=rec_array["Declination"],
psi=rec_array["Polarization"],
iota=rec_array["Inclination"],
phi0=rec_array["InitialPhase"],
t_init=t_init,
)
return gb_object
@staticmethod
def _equatorial_to_ecliptic(ra: float, dec: float) -> tuple[float, float]:
"""Convert (ra, dec) into ecliptic longitude and lattitude (lambda, beta).
Parameters
----------
ra
Right ascension [rad].
dec
Declination [rad].
Returns
-------
tuple (lambda, beta)
Ecliptic longitude and latitude [rad, rad].
"""
sky_icrs = SkyCoord(ra=ra, dec=dec, frame="icrs", unit="rad")
ecl_coords = sky_icrs.barycentrictrueecliptic
return ecl_coords.lon.rad, ecl_coords.lat.rad
def _to_array_at_reference_time(self, t0: float) -> np.ndarray:
"""Create a numpy array with parameter at reference time `t0`.
Returns
-------
np.array of shape (nsrc, 8)
Array of parameters in the same order as :attr:`PARAM_NAMES`
and with units :attr:`PARAM_UNITS`.
"""
# shift frequency `f0`, which evolves at the linear rate `fdot`
f0_t0 = self.f0 - self.fdot * (self.t_init - t0)
# shift `phi0`, to account for difference `t_init` (catalogue epoch) and `t0`
_shift = np.pi * (
2 * self.f0 * (self.t_init - t0) - self.fdot * (self.t_init - t0) ** 2
)
phi0_t0 = (self.phi0 + _shift) % (2 * np.pi)
return np.array(
[
f0_t0, # shifted f0
self.fdot,
self.A,
self.ra,
self.dec,
self.psi,
self.iota,
phi0_t0, # shifted phi0
]
).T
[docs]
def to_jaxgb_array(self, t0: float) -> jax.Array:
"""Create JaxGB array, respecting format and ordering.
Parameters
----------
t0
Starting time of `JAXGB` response calculation in TCB (s). This time is used
to shift the GW frequency `f0` and the phase `phi0`, defined at `t_init`, to
the simulation starting time.
Returns
-------
jax.Array of shape (nsrc, 8)
Array of parameters in the same order as :attr:`PARAM_NAMES`
and with units :attr:`PARAM_UNITS`.
"""
return jnp.array(self._to_array_at_reference_time(t0=t0))
[docs]
def to_ldc_dict(self) -> dict[str, np.ndarray]:
"""Create LDC dictionary, respecting format and ordering.
Sky positions are converted to ecliptic coordinates. Parameters are specified
at the epoch `t_init`.
Returns
-------
Dictionary of parameters.
"""
# LDC assumes no difference between `t0` and `t_init`
params = self._to_array_at_reference_time(t0=self.t_init)
# convert equatorial to ecliptic coordinates, change order
params[4], params[3] = GBObject._equatorial_to_ecliptic(params[3], params[4])
return dict(
zip(
[
"Frequency",
"FrequencyDerivative",
"Amplitude",
"EclipticLatitude",
"EclipticLongitude",
"Polarization",
"Inclination",
"InitialPhase",
],
params,
)
)
[docs]
def to_fastgb_array(self) -> np.ndarray:
"""Create fastgb array, respecting format and ordering.
The parameter array assumes the epoch `t_init`, and fastgb is doing the
conversion between `tinit` and `t0` internally.
Additional comments:
fastgb expects that `params[:, 3]` contain ecliptic latitudes (beta) and
`params[:, 4]` contain ecliptic longitudes (lambda). However, from LISA
Orbits v3 onwards, orbit information is expressed in equatorial coordinates,
so implictly, the sky location params in fastgb need to be expressed in
equatorial coordinates. In principle, we can therefore pass the equatorial
coordinates, as in JAXGB, provided we swap the longitude and latitude order:
- In JAXGB, `params[:, 3]` is longitude (ra) while `params[:, 4]` includes
latitude (dec),
- In fastgb, `params[:, 3]` (what used to be beta) includes latitude, while
`params[:, 4]` includes longitude (what used to be beta)
Returns
-------
np.ndarray of shape (nsrc, 8)
Array of parameters.
"""
# fastgb applies internal conversion between `t0` and `t_init`, so we assume
# them to be identical here
params = self._to_array_at_reference_time(t0=self.t_init)
return params[[0, 1, 2, 4, 3, 5, 6, 7]]
[docs]
def to_gwresponse_dict(self) -> dict[str, np.ndarray]:
"""Create GW Response input dictionary, respecting format and ordering.
This can be used with :class:`lisagwresponse.GalacticBinary`. The parameter
dictionary assumes the epoch `t_init`.
Returns
-------
Dictionary of parameters, compatible with
:class:`lisagwresponse.GalacticBinary`.
"""
# GW response applies own internal time shift for `t0` and `t_init`, so we
# assume them to be identical here
params = self._to_array_at_reference_time(t0=self.t_init)
return dict(
zip(
[
"f",
"df",
"A",
"ra",
"dec",
"psi",
"iota",
"phi0",
],
params,
)
)
[docs]
def to_gbgpu_array(self) -> np.ndarray:
"""Create gbgpu array, respecting format and ordering.
Sky positions are converted to ecliptic coordinates. ``fddot`` is
assumed to be zero. The returned parameters have the epoch `t_init`.
Returns
-------
np.ndarray of shape (nsrc, 9)
Array of parameters.
"""
# gbgbpu assumes no difference between `t0` and `t_init`
params = self._to_array_at_reference_time(t0=self.t_init)
params = params[[2, 0, 1, 1, 7, 6, 5, 3, 4]]
params[3] = 0 # f2dot
# convert equatorial to corresponding ecliptic coordinates
params[7], params[8] = GBObject._equatorial_to_ecliptic(params[7], params[8])
return np.array([np.full(1, p) for p in params])