gw_animate.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# The MIT License (MIT)
# Copyright (c) 2017 Duncan Macleod
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
# OR OTHER DEALINGS IN THE SOFTWARE.

"""
Animate a gravitational-wave signal, and with detectior noise, and
with filtered detector noise

This is a modified version that allows to create animation for any GW detections
It also fixes a couple of imports due to deprecated api used in the original Duncan script
Huge kudos to Duncan for all of the hard work
"""

from __future__ import division

import h5py
import numpy
from astropy.utils.data import download_file
from matplotlib import animation
from matplotlib import pyplot
from gwpy.timeseries import TimeSeries
from gwpy.plot import Plot
from gwpy.plot.rc import rcParams
from gwpy.signal import filter_design
from pycbc.catalog import Merger
from pycbc.waveform import get_td_waveform
from pycbc.detector import Detector
from statistics import mean


pyplot.style.use('dark_background')
rcParams.update({
    'figure.dpi': 600,
    'figure.subplot.left': 0.,
    'figure.subplot.right': 1.,
    'figure.subplot.bottom': 0.,
    'figure.subplot.top': 1.,
    'text.usetex': False,
    'font.family': 'sans-serif',
    'font.sans-serif': ['Helvetica Neue', 'Helvetica'],
})


# -- animation helpers --------------------------------------------------------

def create_figure(xlim=None, ylim=None):
    plot = Plot(figsize=[12, 4])
    ax = plot.gca()
    ax.set_axis_off()
    ax.arrow(0.02, 0.05, 0.02, 0, color='white',
             head_width=0.02, head_length=0.01, transform=ax.transAxes)
    ax.text(0.06, 0.05, 'Time',
            transform=ax.transAxes, color='white', va='center')
    return plot, ax


def animate(target, fig, animate_func, init_func, fps=30., duration=10.):
    nf = fps * duration
    delay = int(1 / fps * 1000)
    anim = animation.FuncAnimation(fig, animate_func, init_func=init_func,
                                   frames=int(nf), interval=delay, blit=True)
    anim.save(target, fps=fps, extra_args=['-vcodec', 'libxvid'])
    print('{} written'.format(target))
    return anim


def legend(ax, artists, labels):
    leg = ax.legend(artists, labels, loc='upper left', frameon=False,
                    fontsize=16)
    for text in leg.get_texts():  # change font colour
        text.set_color('white')
    return leg


# -- get data and simulate waveform -------------------------------------------

event = 'GW170818'
# create waveform
m = Merger(event)
hp, hx = get_td_waveform(approximant='IMRPhenomD', delta_t=1 / 4096.,
                         f_lower=30, **m.__dict__)

posteriors = download_file('https://dcc.ligo.org/public/0157/P1800370/005/%s_GWTC-1.hdf5' % event, cache=True)
# project waveform onto LHO at the right time (and sky location)
f = h5py.File(posteriors, mode='r')

ra = mean(f['IMRPhenomPv2_posterior']['right_ascension'])
dec = mean(f['IMRPhenomPv2_posterior']['declination'])
gps = m.time

available_detectors = {}
for x in ['H1', 'L1', 'V1', 'G1', 'K1', 'I1']:
    try:
        m.strain(x)
        available_detectors[x] = Detector(x)
    except:
        pass

print('Generated waveform')
Fp, Fx = available_detectors['H1'].antenna_pattern(ra, dec, 0, gps)
sig = TimeSeries.from_pycbc(Fp * hp + Fx * hx).pad((0, 512))

print('Downloading data')
time_series = {}
for x in available_detectors.keys():
    time_series[x] = TimeSeries.fetch_open_data(x, gps - 16, gps + 16, cache=True)

# shift data to account for time-delay and orientation
for x in available_detectors.keys():
    if x == 'H1':
        continue
    time_delta = available_detectors['H1'].time_delay_from_detector(available_detectors[x], ra, dec, gps)
    time_series[x].t0 += time_delta * time_series[x].t0.unit
    time_series[x] *= -1 * time_series[x].unit

# shift detector data to merge at t=0
t0 = gps * time_series['H1'].t0.unit
for x in available_detectors.keys():
    time_series[x].t0 -= t0

# rescale to sensible amplitude (makes plotting easier)
for x in available_detectors.keys():
    time_series[x] *= 1e21
sig *= 1e21

# -- animate waveform ---------------------------------------------------------

print('Animating waveform...')

# crop signal to interval we want to animate
sigc = sig.crop(-.4, .2)

# set animation parameters
fps = 30.
duration = 10.
times = numpy.arange(sigc.size)

# create figure and animation methods
plot, ax = create_figure()
line, = ax.plot([], [], color='white')
ax.set_xlim(0, times.size)
ax.set_ylim(-2, 2)
npf = sig.size / (fps * duration)


def init():
    line.set_data([], [])
    return line,


def update(n):
    n = int(n * npf)
    line.set_data(times[:n], sigc[:n])
    return line,


# write animation
animate('%s-signal.mp4' % event, plot, update, init, fps=fps, duration=duration)
plot.close()

# -- animate filtered data ----------------------------------------------------

filter1 = filter_design.highpass(10, time_series['H1'].sample_rate)
bp = filter_design.bandpass(50, 250, time_series['H1'].sample_rate)
notches = [filter_design.notch(f, time_series['H1'].sample_rate) for f in (60, 120, 180)]
filter2 = filter_design.concatenate_zpks(bp, *notches)

for zpk, ylim, tag in [
    (filter1, (-300, 300), 'highpass'),
    (filter2, (-2, 2), 'filtered'),
]:
    print('Animating {} data...'.format(tag))

    filtered_lines = {}
    for x in time_series.keys():
        filtered_lines[x] = time_series[x].filter(zpk, filtfilt=True).crop(-.4, .2)
    sigf = sig.filter(zpk, filtfilt=True).crop(-.4, .2)[:filtered_lines['H1'].size]
    times = numpy.arange(sigf.size)

    # create figure and animation methods
    plot, ax = create_figure()
    lines = {}
    for x in time_series.keys():
        if x == 'H1':
            color = 'gwpy:ligo-hanford'
        elif x == 'L1':
            color = 'gwpy:ligo-livingston'
        elif x == 'V1':
            color = 'gwpy:virgo'
        elif x == 'G1':
            color = 'gwpy:geo600'
        elif x == 'K1':
            color = 'gwpy:kagra'
        elif x == 'I1':
            color = 'gwpy:ligo-india'
        else:
            raise RuntimeError('wtf?')
        lines[x] = ax.plot([], [], color=color, linewidth=1)[0]
    sigline, = ax.plot([], [], color='white')
    ax.set_xlim(0, times.size)
    ax.set_ylim(*ylim)

    # add a legend
    legend_list = []
    for x in lines.keys():
        if x == 'H1':
            legend_list.append('LIGO-Hanford')
        elif x == 'L1':
            legend_list.append('LIGO-Livingston')
        elif x == 'V1':
            legend_list.append('VIRGO')
        elif x == 'G1':
            legend_list.append('GEO600')
        elif x == 'K1':
            legend_list.append('KAGRA')
        elif x == 'I1':
            legend_list.append('LIGO-India')

    legend(ax, [sigline] + list(lines.values()),
           ['%s signal' % event] + legend_list)

    def init():
        sigline.set_data([], [])
        ret = (sigline, )
        for x in lines.keys():
            lines[x].set_data([], [])
            ret += (lines[x], )
        return ret


    def update(n):
        n = int(n * npf)
        sigline.set_data(times[:n], sigf.value[:n])
        ret = (sigline, )
        for x in lines.keys():
            lines[x].set_data(times[:n], filtered_lines[x].value[:n])
            ret += (lines[x], )
        return ret


    # write animation
    animate(event + ('-{}.mp4'.format(tag)), plot, update, init,
            fps=fps, duration=duration)
    plot.close()