Waveform formatting

We will see how to format generated waveforms into the types of TLT, such as HarmonicWaveform and PlusCrossWaveform.

Let’s get started!

We will show how to format the waveform of a MBHB.

Show code cell source

Hide code cell source

import os

import lovelyplots
import matplotlib.pyplot as plt

plt.style.use(["paper", "use_mathtext", "colors5"])
del lovelyplots

MOJITO_USERNAME = os.environ.get("MOJITO_USERNAME", "")
MOJITO_TOKEN = os.environ.get("MOJITO_TOKEN", "")

First, we prepare a MBHB parameter, fetched from the Mojito-light dataset using the mojito package.

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", val=True)
import mojito.download

_mbhb = mojito.download.get_source_params(
    "mbhb", 18, username=MOJITO_USERNAME, token=MOJITO_TOKEN
)
mbhb = {k: jnp.asarray(v) for k, v in _mbhb.items()}

# We will use a PhenomTHM approximant and the following parameters are required
parameter = {
    "m1": mbhb["PrimaryMassSSBFrame"],
    "m2": mbhb["SecondaryMassSSBFrame"],
    "chi1z": mbhb["PrimarySpinCompZ"],
    "chi2z": mbhb["SecondarySpinCompZ"],
    "distance": mbhb["LuminosityDistance"],
    "phi_ref": mbhb["PhaseReferenceSourceFrame"],
    "inclination": mbhb["InclinationAngle"],
    "psi": mbhb["PolarisationAngle"],
}

Downloading file 'brickmarket/mojito_light_v1_0_0/catalogues/mbhb_cat_mojito_lite_processed_MT_rounding_fixed.hdf5' from 'https://nextcloud-dcc-fi-csc-okd-globalstorage1.2.rahtiapp.fi/remote.php/dav/files/senwen-deng/brickmarket/mojito_light_v1_0_0/catalogues/mbhb_cat_mojito_lite_processed_MT_rounding_fixed.hdf5' to '/root/.cache/mojito'.

WARNING! Publications using Mojito data are currently not allowed! Please keep in touch, as publication policies will soon be published.

Then, we prepare the waveform generation wrapper from phentax.

import phentax

wavegen = phentax.waveform.IMRPhenomTHM()

# The order of the harmonic modes
modes = (
    (2, 2),
    (2, 1),
    (3, 3),
    (4, 4),
    (5, 5),
    (2, -2),
    (2, -1),
    (3, -3),
    (4, -4),
    (5, -5),
)

Plus and cross polarizations

We can format the plus and cross polarizations.

times, mask, h_plus, h_cross = wavegen.compute_polarizations_at_once(**parameter)

The above arrays have a batch dimension of size 1. In general, the batches obtained from phentax do not necessarily share the same time axis. In TLT, we require batches to share the same grid (but different modes are allowed to have different grids). Here we have only one waveform, so we can directly use the 0-component of the batch dimension.

We also conventionalize the shape of the waveform arrays.

times = times[0][mask[0]]
h_plus = h_plus[0][mask[0]][None, None, None, None, ...]  # Conventionalization
h_cross = h_cross[0][mask[0]][None, None, None, None, ...]

These components allow us to construct the TLT waveform objects.

import typed_lisa_toolkit as tlt

axis = tlt.axis(times)

plus = tlt.time_series(axis, h_plus)
cross = tlt.time_series(axis, h_cross)
polarizations = tlt.pcw({"plus": plus, "cross": cross})

A note on the static type checking: without running the code, a static type checker can already infer that axis is of type Axis[Array], plus and cross are of type TimeSeries[Axis[Array]], and polarizations is of type PlusCrossWaveform[TimeSeries[Axis[Array]]].

We can plot the polarizations using the TLT plotting utilities. TLT automatically handles the labeling. One can also customize the plot using the returned figure and axes objects.

fig, axs = tlt.plot(polarizations, set_legend=False)

Per harmonic strain

We can also format complex strain of each harmonic mode.

times, mask, h_lms = wavegen.compute_hlms(**parameter)
times = times[0][mask[0]]
times = tlt.linspace_from_array(times)
axis = tlt.axis(times)
h_lms = h_lms[0][:, mask[0]][None, None, :, None, ...]  # Conventionalization

lms = tlt.hhw(
    {
        mode: tlt.time_series(axis, h_lms[:, :, slice(i_mode, i_mode + 1)])
        for i_mode, mode in enumerate(modes)
    }
)

This time, times is a Linspace, more memory efficient than arrays, times is Axis[Linspace], and lms is HomogeneousHarmonicWaveform[Harmonic, UniformTimeSeries].

fig, axs = tlt.plot(lms.abs())
_ = axs[0][0].set_ylabel("Strain Magnitude")

Note that we plotted the modulus of the complex strain.

Per harmonic amplitude and phase

TLT also supports waveforms represented by amplitude and phase separately.

_, _, mask, amplitudes, phases = wavegen.compute_amp_phase(**parameter)
amplitudes = amplitudes[0][:, mask[0]][None, None, :, None, ...]  # Conventionalization
phases = phases[0][:, mask[0]][None, None, :, None, ...]

phasor = tlt.hhw(
    {
        mode: tlt.time_phasor(
            axis,
            amplitudes[:, :, slice(i_mode, i_mode + 1)],
            phases[:, :, slice(i_mode, i_mode + 1)],
        )
        for i_mode, mode in enumerate(modes[:5])
    }
)

The type of phasor is HomogeneousHarmonicWaveform[Harmonic, TimePhasor[Axis[Linspace]]]. The function plot() traces the phase of the object.

fig, axs = tlt.plot(phasor)
_ = axs[0][0].set_ylabel("Phase")