{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Waveform formatting" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We will see how to format generated waveforms into the types of TLT, such as {class}`~.types.HarmonicWaveform` and {class}`~.types.PlusCrossWaveform`." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Let's get started!\n", "\n", "We will show how to format the waveform of a MBHB.\n", "\n", "" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "import os\n", "\n", "import lovelyplots\n", "import matplotlib.pyplot as plt\n", "\n", "plt.style.use([\"paper\", \"use_mathtext\", \"colors5\"])\n", "del lovelyplots\n", "\n", "MOJITO_USERNAME = os.environ.get(\"MOJITO_USERNAME\", \"\")\n", "MOJITO_TOKEN = os.environ.get(\"MOJITO_TOKEN\", \"\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "First, we prepare a MBHB parameter, fetched from the Mojito-light dataset using the [mojito](https://mojito-e66317.io.esa.int/index.html) package." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", "\n", "jax.config.update(\"jax_enable_x64\", val=True)\n", "import mojito.download\n", "\n", "_mbhb = mojito.download.get_source_params(\n", " \"mbhb\", 18, username=MOJITO_USERNAME, token=MOJITO_TOKEN\n", ")\n", "mbhb = {k: jnp.asarray(v) for k, v in _mbhb.items()}\n", "\n", "# We will use a PhenomTHM approximant and the following parameters are required\n", "parameter = {\n", " \"m1\": mbhb[\"PrimaryMassSSBFrame\"],\n", " \"m2\": mbhb[\"SecondaryMassSSBFrame\"],\n", " \"chi1z\": mbhb[\"PrimarySpinCompZ\"],\n", " \"chi2z\": mbhb[\"SecondarySpinCompZ\"],\n", " \"distance\": mbhb[\"LuminosityDistance\"],\n", " \"phi_ref\": mbhb[\"PhaseReferenceSourceFrame\"],\n", " \"inclination\": mbhb[\"InclinationAngle\"],\n", " \"psi\": mbhb[\"PolarisationAngle\"],\n", "}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Then, we prepare the waveform generation wrapper from [phentax](https://asantini29.github.io/phentax/)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import phentax\n", "\n", "wavegen = phentax.waveform.IMRPhenomTHM()\n", "\n", "# The order of the harmonic modes\n", "modes = (\n", " (2, 2),\n", " (2, 1),\n", " (3, 3),\n", " (4, 4),\n", " (5, 5),\n", " (2, -2),\n", " (2, -1),\n", " (3, -3),\n", " (4, -4),\n", " (5, -5),\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Plus and cross polarizations\n", "\n", "We can format the plus and cross polarizations." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "times, mask, h_plus, h_cross = wavegen.compute_polarizations_at_once(**parameter)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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. \n", "\n", "We also [conventionalize the shape](https://l2d-interface-c43116.pages.in2p3.fr/overview.html#shape-convention) of the waveform arrays." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "times = times[0][mask[0]]\n", "h_plus = h_plus[0][mask[0]][None, None, None, None, ...] # Conventionalization\n", "h_cross = h_cross[0][mask[0]][None, None, None, None, ...]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "These components allow us to construct the TLT waveform objects." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import typed_lisa_toolkit as tlt\n", "\n", "axis = tlt.axis(times)\n", "\n", "plus = tlt.time_series(axis, h_plus)\n", "cross = tlt.time_series(axis, h_cross)\n", "polarizations = tlt.pcw({\"plus\": plus, \"cross\": cross})" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "A note on the static type checking: without running the code, a static type checker can already infer that `axis` is of type {class}`~.types.Axis`[{class}`~.types.misc.Array`], `plus` and `cross` are of type {class}`~.types.TimeSeries`[{class}`~.types.Axis`[{class}`~.types.misc.Array`]], and `polarizations` is of type {class}`~.types.PlusCrossWaveform`[{class}`~.types.TimeSeries`[{class}`~.types.Axis`[{class}`~.types.misc.Array`]]].\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = tlt.plot(polarizations, set_legend=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Per harmonic strain\n", "\n", "We can also format complex strain of each harmonic mode." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "times, mask, h_lms = wavegen.compute_hlms(**parameter)\n", "times = times[0][mask[0]]\n", "times = tlt.linspace_from_array(times)\n", "axis = tlt.axis(times)\n", "h_lms = h_lms[0][:, mask[0]][None, None, :, None, ...] # Conventionalization\n", "\n", "lms = tlt.hhw(\n", " {\n", " mode: tlt.time_series(axis, h_lms[:, :, slice(i_mode, i_mode + 1)])\n", " for i_mode, mode in enumerate(modes)\n", " }\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This time, `times` is a {class}`~.types.Linspace`, more memory efficient than arrays, `times` is {class}`~.types.Axis`[{class}`~.types.Linspace`], and `lms` is {class}`~.types.HomogeneousHarmonicWaveform`[{class}`~.types.Harmonic`, {class}`~.types.UniformTimeSeries`]." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = tlt.plot(lms.abs())\n", "_ = axs[0][0].set_ylabel(\"Strain Magnitude\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note that we plotted the modulus of the complex strain." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Per harmonic amplitude and phase\n", "\n", "TLT also supports waveforms represented by amplitude and phase separately." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "_, _, mask, amplitudes, phases = wavegen.compute_amp_phase(**parameter)\n", "amplitudes = amplitudes[0][:, mask[0]][None, None, :, None, ...] # Conventionalization\n", "phases = phases[0][:, mask[0]][None, None, :, None, ...]\n", "\n", "phasor = tlt.hhw(\n", " {\n", " mode: tlt.time_phasor(\n", " axis,\n", " amplitudes[:, :, slice(i_mode, i_mode + 1)],\n", " phases[:, :, slice(i_mode, i_mode + 1)],\n", " )\n", " for i_mode, mode in enumerate(modes[:5])\n", " }\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The type of `phasor` is {class}`~.types.HomogeneousHarmonicWaveform`[{class}`~.types.Harmonic`, {class}`~.types.TimePhasor`[{class}`~.types.Axis`[{class}`~.types.Linspace`]]]. The function {func}`~.plot` traces the phase of the object." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axs = tlt.plot(phasor)\n", "_ = axs[0][0].set_ylabel(\"Phase\")" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.11" } }, "nbformat": 4, "nbformat_minor": 2 }