Python script to plot gravitational wave waveforms.

waveform_plotting.py

"""Module to plot Gravitational Wave Templates

Can be used to test if a waveform approximant works or not. 

"""

from __future__ import division, print_function

import bilby 
import matplotlib.pyplot as plt
import plotly
import plotly.tools as tls
import plotly.graph_objects as go
from matplotlib.ticker import (AutoMinorLocator)
import logging
logging.getLogger().setLevel(logging.INFO)

approximants = [
    "IMRPhenomXPHM",
    "IMRPhenomPv3HM",
    "SEOBNRv4PHM",
    "IMRPhenomPv2_NRTidal", 
    "IMRPhenomPv2", 
    "SEOBNRv4",
    "NRSur7dq4"
]


injection_parameters = dict(
    mass_1=141.0741,
    mass_2=113.0013,
    a_1=0.9434,
    a_2=0.2173,
    tilt_1=0,
    tilt_2=0,
    phi_jl=0,
    phi_12=0,
    luminosity_distance=1782.1610,
    theta_jn=0.9614,
    psi=1.6831,
    phase=5.2220,
    geocent_time=1242424473.5880,
    ra=0.9978,
    dec=-0.4476
)


def get_data(approximants, injection_parameters):
    """Passes injection_param to the approximants and calls the approximant to generate the waveforms."""
    generator_args = dict(
        duration=4,
        sampling_frequency=2048.,
        frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
        parameters=injection_parameters
    )

    waveform_args = dict(reference_frequency=50.)

    approximant_data = dict()
    for approximant in approximants:
        logging.info(f"Creating strain from {approximant}")
        waveform_args.update(dict(waveform_approximant=approximant))
        generator_args.update(waveform_arguments=waveform_args)
        generator = bilby.gw.WaveformGenerator(**generator_args)
        h_signal = generator.time_domain_strain(injection_parameters)
        h_time = generator.time_array
        approximant_data.update({approximant: dict(x=h_time, y=h_signal['plus'])})
    return approximant_data


def matplot_plot_data(approximant_data):
    """Plots the approximant data as a matplotlib png"""
    logging.info("plotting waveforms")
    fig, ax = plt.subplots()
    for approximant, data in approximant_data.items():
        ax.plot(data['x'], data['y'], label=approximant)
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.set_title('Comparing Time-Domain Data')
    ax.set_ylabel('Strain')
    ax.set_xlabel('Time (s)')
    ax.legend()
    return fig

def plotly_plot_data(fig):
    """Converts the matplotlib plot to a plotly plot."""
    ax_list = fig.axes
    for ax in ax_list:
        ax.get_legend().remove()

    plotly_fig = tls.mpl_to_plotly(fig)
    legend = go.layout.Legend(
        x=0.05,
        y=0.95
    )
    plotly_fig.update_layout(showlegend=True, legend=legend)
    plotly.offline.plot(plotly_fig, filename="waveform_test.html")


def main():
    data = get_data(approximants, injection_parameters)
    fig = matplot_plot_data(data)
    fig.savefig('waveform_test.png')
    plotly_plot_data(fig)

if __name__ == '__main__':
    main()