gstlal_fake_frames

#!/usr/bin/env python
#
# Copyright (C) 2011  Kipp Cannon, Chad Hanna, Drew Keppel
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import numpy
from optparse import OptionParser
import os
import sys

import gi
gi.require_version('Gst', '1.0')
from gi.repository import GObject, Gst
GObject.threads_init()
Gst.init(None)

import lal

from gstlal import pipeparts
from gstlal import reference_psd
from gstlal import simplehandler
from gstlal import datasource
from gstlal import multirate_datasource
from glue.ligolw import utils as ligolw_utils

### This program will make fake data in a variety of ways.  Its input is anything
### supported by datasource.py.  You can additionally whiten the data or apply a
### frequency domain coloring filter.  It writes its output to frame files.
###
### Overview graph of the pipeline
### ------------------------------
###
### Gray boxes are optional and depend on the command line given
###
### .. graphviz::
###
###    digraph G {
###
###        // graph properties
###
###        rankdir=LR;
###        compound=true;
###        node [shape=record fontsize=10 fontname="Verdana"];
###        edge [fontsize=8 fontname="Verdana"];
###
###        // nodes
###
###        "mkbasicsrc()" [URL="\ref datasource.mkbasicsrc()"];
###        "whitened_multirate_src()" [label="whitened_multirate_src()", URL="\ref multirate_datasource.mkwhitened_multirate_src()", style=filled, color=lightgrey];
###        capsfilter [style=filled, color=lightgrey, URL="\ref pipeparts.mkcapsfilter()"];
###        resample [style=filled, color=lightgrey, URL="\ref pipeparts.mkresample()"];
###        audioamplify [style=filled, color=lightgrey, URL="\ref pipeparts.mkaudioamplify()", label="audioamplify \n if --data-source=white \n in order to have unit variance \n after resampling"];
###        lal_shift[style=filled, color=lightgrey, URL="\ref pipeparts.mkshift()", label="lal_shift \n [iff --shift provided]"];
###        progressreport1 [URL="\ref pipeparts.mkprogressreport()", label="progressreport \n [iff --verbose provided]"];
###        progressreport2 [style=filled, color=lightgrey, URL="\ref pipeparts.mkprogressreport()", label="progressreport \n [iff --verbose provided]"];
###        lal_firbank [style=filled, color=lightgrey, URL="\ref pipeparts.mkfirbank()", label="lal_firbank \n [iff --color-psd provided]"];
###        taginject [URL="\ref pipeparts.mktaginject()"];
###        lal_simulation [style=filled, color=lightgrey, URL="\ref pipeparts.mkinjections()", label="lal_simulation \n [iff --injections provided]"];
###        framecppchannelmux [URL="\ref pipeparts.mkframecppchannelmux()"];
###        framecppfilesink [URL="\ref pipeparts.mkframecppfilesink()"];
###
###        // connect it up
###
###        "mkbasicsrc()" -> progressreport1-> lal_shift  -> progressreport2;
###        progressreport2 -> "whitened_multirate_src()" [label="if --whiten given"];
###        "whitened_multirate_src()" -> lal_firbank;
###        progressreport2 -> resample [label="if --whiten not given"];
###        resample -> capsfilter;
###        capsfilter -> audioamplify
###        audioamplify -> lal_firbank;
###        lal_firbank -> taginject -> lal_simulation -> framecppchannelmux -> framecppfilesink;
###    }
###
###
### Usage cases
### -----------
###
### 1. Making initial LIGO colored noise, Advanced LIGO noise, etc for different
### instruments.   See datasource.append_options() for other choices::
###
###    $ gstlal_fake_frames --data-source=LIGO --channel-name=H1=FAKE-STRAIN \
###    --frame-type=H1_FAKE --gps-start-time=900000000 --gps-end-time=900005000 \
###    --output-path=testframes  --verbose
###
###    $ gstlal_fake_frames --data-source=AdvLIGO --channel-name=L1=FAKE-STRAIN \
###    --frame-type=L1_FAKE --gps-start-time=900000000 --gps-end-time=900005000 \
###    --output-path=testframes  --verbose
###
###
### 2. Custom colored noise.
###
### Obviously it is not practical to code up every possible noise curve to
### simulate as a custom data source.  However, it is possible to provide your
### own custom noise curve as an ASCII file with frequency in one column and
### strain/Hz in the second. e.g., early advanced ligo noise curve
### https://git.ligo.org/lscsoft/gstlal/raw/master/gstlal/share/early_aligo_asd.txt
### You will need to convert ASD text files to PSD xml files using
### gstlal_psd_xml_from_asd_txt first.::
###
###    $ gstlal_fake_frames --data-source=white --color-psd v1psd.xml.gz \
###    --channel-name=V1=FAKE-STRAIN --frame-type=V1_FAKE --gps-start-time=900000000 \
###    --gps-end-time=900005000 --output-path=testframes  --verbose
###
###
### 3. Recoloring existing frame data
###
### This procedure is very similar to the above except that instead of
### using white noise to drive the fake frame generation, we start with
### real frame data and whiten it.  Recoloring can be thought of as first
### whitening data and then applying a coloring filter.  You must first
### make a reference PSD of the data with e.g. gstlal_reference_psd.  You
### will need to make a frame cache of the frames to recolor.::
###
###    $ gstlal_fake_frames --whiten-reference-psd reference_psd.xml.gz \
###    --color-psd recolor_psd.xml.gz --data-source frames \
###    --output-path /path/to/output --output-channel-name h_16384Hz \
###    --gps-start-time 966384031 --frame-type T1300121_V1_EARLY_RECOLORED_V2 \
###    --gps-end-time 966389031 --frame-duration 16 --frames-per-file 256 \
###    --frame-cache frame.cache --channel-name=V1=h_16384Hz --verbose
###
###
### 4. Creating injections into silence (assumes you have an injection xml file from e.g. lalapps_inspinj)::
###
###    $ gstlal_fake_frames --data-source silence --output-path /path/to/output \
###    --gps-start-time 966384031 --frame-type V1_INJECTIONS \
###    --gps-end-time 966389031 --frame-duration 16 --frames-per-file 256 \
###    --verbose --channel-name=V1=INJECTIONS --injections injections.xml
###
###
### 5. Other things are certainly possible, please add some!
###
### Debug
### -----
###
### It is possible to check the pipeline graph by interupting the running process
### with ctrl+c if you have the GST_DEBUG_DUMP_DOT_DIR evironment set.  A dot
### graph will be written to gstlal_fake_frames.  Here is an example::
###
###    $ GST_DEBUG_DUMP_DOT_DIR="." gstlal_fake_frames --data-source=LIGO \
###    --channel-name=H1=FAKE-STRAIN --frame-type=H1_FAKE \
###    --gps-start-time=900000000 --gps-end-time=900005000 \
###    --output-path=testframes  --verbose``
###
### You should see::
###
###    Wrote pipeline to ./gstlal_fake_frames_PLAYING.dot
###
### After Ctrl+c you should see::
###
###    ^C*** SIG 2 attempting graceful shutdown (this might take several minutes) ... ***
###        Wrote pipeline to ./gstlal_fake_frames.dot
###
### You can then turn the pipeline graph into an image with dot, e.g.,::
###
###    $ dot -Tpng gstlal_fake_frames_PLAYING.dot > test.png
###
### Review
### ------
###


##
# Extract and validate the command line options
#
def parse_command_line():
   parser = OptionParser(description = __doc__)

   #
   # Append data source options
   #

   datasource.append_options(parser)

   #
   # Append program specific options
   #

   parser.add_option("--shift", metavar = "ns", help = "Number of nanoseconds to delay (negative) or advance (positive) the time stream", type = "int")
   parser.add_option("--sample-rate", metavar = "Hz", default = 16384, type = "int", help = "Sample rate at which to generate the data, should be less than or equal to the sample rate of the measured psds provided, default = 16384 Hz, max 16384 Hz")
   parser.add_option("--whiten-reference-psd", metavar = "name", help = "Set the name of psd xml file to whiten the data with")
   parser.add_option("--whiten-track-psd", action = "store_true", help = "Calculate PSD from input data and track with time.  Always enabled if --whiten-reference-psd is not given.")
   parser.add_option("--color-psd", metavar = "name", help = "Set the name of psd xml file to color the data with")
   parser.add_option("--output-path", metavar = "name", default = ".", help = "Path to output frame files (default = \".\").")
   parser.add_option("--output-channel-name", metavar = "name", help = "The name of the channel in the output frames. The default is the same as the channel name")
   parser.add_option("--frame-type", metavar = "name", help = "Frame type, required")
   parser.add_option("--frame-duration", metavar = "s", default = 16, type = "int", help = "Set the duration of the output frames.  The duration of the frame file will be multiplied by --frames-per-file.  Default: 16s")
   parser.add_option("--frames-per-file", metavar = "n", default = 256, type = "int", help = "Set the number of frames per file.  Default: 256")
   parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose (optional).")

   #
   # Parse options
   #

   options, filenames = parser.parse_args()

   if options.sample_rate > 16384:
       raise ValueError("--sample-rate must be <= 16384")

   if options.frame_type is None:
       raise ValueError("--frame-type is required")

   if options.whiten_reference_psd is None:
       options.whiten_track_psd = True

   return options, filenames


#
# Main
#

options, filenames = parse_command_line()

## Parse the command line options into a python.datasource.GWDataSourceInfo class instance
gw_data_source = datasource.GWDataSourceInfo(options)

## Assume instrument is the first and only key of the python.datasource.GWDataSourceInfo.channel_dict
instrument = gw_data_source.channel_dict.keys()[0]

# disable progress reports if not verbose
if not options.verbose:
   pipeparts.mkprogressreport = lambda pipeline, src, *args: src

# set default output channel if not set by user
if options.output_channel_name is None:
   options.output_channel_name = gw_data_source.channel_dict[instrument]

# do not do injections in datasource.mkbasicsrc(), we will do them later
injections, gw_data_source.injection_filename = options.injections, None

## Setup the pipeline
pipeline = Gst.Pipeline(name=os.path.split(sys.argv[0])[1])

## Main loop 
mainloop = GObject.MainLoop()

## An instance of the python.simplehandler.Handler class
handler = simplehandler.Handler(mainloop, pipeline)

## 1) Set the pipeline head to basic input from datasource.mkbasicsrc()
# FIXME: fake source causes problems when making large buffers, so block_size needs to be overwritten
gw_data_source.block_size = 8 * options.sample_rate
head, _, _ = datasource.mkbasicsrc(pipeline, gw_data_source, instrument, verbose = options.verbose)

## 2) Set the pipeline head to be verbose with pipeparts.mkprogressreport()
head = pipeparts.mkprogressreport(pipeline, head, "frames")

if options.shift is not None:
   ## 3) Set the pipeline head to apply a time shift if requested with a pipeparts.mkshift() element
   head = pipeparts.mkshift(pipeline, head, shift = options.shift)

   ## 4) Set the pipeline head to be verbose with a pipeparts.mkprogressreport() element
   head = pipeparts.mkprogressreport(pipeline, head, "frames_shifted")

if options.whiten_reference_psd:
   ## If whitening read the PSD
   wpsd = lal.series.read_psd_xmldoc(ligolw_utils.load_filename(options.whiten_reference_psd, verbose = options.verbose, contenthandler = lal.series.PSDContentHandler))[instrument]
else:
   ## else set wpsd to None
   wpsd = None

if options.whiten_reference_psd or options.whiten_track_psd:
   ## 5) Set the pipeline head to a whitened data stream if requested using a multirate_datasource.mkwhitened_multirate_src()
   head = multirate_datasource.mkwhitened_multirate_src(pipeline, head, [options.sample_rate], instrument, psd = wpsd, track_psd = options.whiten_track_psd)[options.sample_rate]
else:
   ## 6) Otherwise simply add a pipeparts.mkcapsfilter() and pipeparts.mkresample()
   head = pipeparts.mkcapsfilter(pipeline, pipeparts.mkresample(pipeline, head, quality = 9), "audio/x-raw, rate=%d" % options.sample_rate)
   # FIXME this is a bit hacky, should datasource.mkbasicsrc be patched to change the sample rate?
   # FIXME don't hardcode sample rate (though this is what datasource generates for all fake data, period 
   # Renormalize if datasource is "white"
   if options.data_source == "white":
       head = pipeparts.mkaudioamplify(pipeline, head, 1.0 / (pipeparts.audioresample_variance_gain(9, 16384, options.sample_rate))**.5)

# Apply a coloring filter
if options.color_psd:

   ## read coloring psd file and convert to an FIR filter
   rpsd = lal.series.read_psd_xmldoc(ligolw_utils.load_filename(options.color_psd, verbose = options.verbose, contenthandler = lal.series.PSDContentHandler))[instrument]
   
   ## Calculate the maximum sample rate
   max_sample = int(round(1.0 / rpsd.deltaF * options.sample_rate / 2.0)) + 1

   # Truncate to requested output sample rate, if it is higher than the psd provides an assert will fail later
   rpsd_data = numpy.array(rpsd.data.data[:max_sample])
   rpsd = lal.CreateCOMPLEX16FrequencySeries(name = rpsd.name, epoch = rpsd.epoch, f0 = rpsd.f0, deltaF = rpsd.deltaF, length = max_sample, sampleUnits = rpsd.sampleUnits)
   rpsd.data.data = rpsd_data
   
   # create the coloring FIR kernel from reference_psd.psd_to_fir_kernel()
   FIRKernel = reference_psd.PSDFirKernel()
   fir_matrix, latency, measured_sample_rate = FIRKernel.psd_to_linear_phase_whitening_fir_kernel(rpsd, invert = False)

   ## 7) Set the pipeline head to a pipeparts.mkfirbank() recoloring filter
   head = pipeparts.mkfirbank(pipeline, head, latency = latency, fir_matrix = [fir_matrix], block_stride = 1 * options.sample_rate)


## Set the tags for the output data
tagstr = "units=strain,channel-name=%s,instrument=%s" % (options.output_channel_name, instrument)

## 8) Put the units back to strain before writing to frames.  Additionally, override the output channel name if provided from the command line
head = pipeparts.mktaginject(pipeline, head, tagstr)


if injections is not None:
   ## 9) Do injections into the final fake data
   head = pipeparts.mkinjections(pipeline, head, injections)

if not os.path.isdir(options.output_path):
   try:
       os.makedirs(options.output_path)
   except:
       print >> sys.stderr, "Unable to make directory ", options.output_path
       pass
else:
   print "Target output directory already exists."

## 10) create frames
head = pipeparts.mkframecppchannelmux(pipeline, {"%s:%s" % (instrument, options.output_channel_name): head}, frame_duration = options.frame_duration, frames_per_file = options.frames_per_file)

## 11) Write the frames to disk
head = pipeparts.mkframecppfilesink(pipeline, head, frame_type = options.frame_type)

# Put O(100000 s) frames in each directory
head.connect("notify::timestamp", pipeparts.framecpp_filesink_ldas_path_handler, (options.output_path, 5))

# Run it
if pipeline.set_state(Gst.State.READY) == Gst.StateChangeReturn.FAILURE:
   raise RuntimeError("pipeline failed to enter READY state")
datasource.pipeline_seek_for_gps(pipeline, gw_data_source.seg[0], gw_data_source.seg[1])
if pipeline.set_state(Gst.State.PLAYING) == Gst.StateChangeReturn.FAILURE:
   raise RuntimeError("pipeline failed to enter PLAYING state")

## Debugging output
if "GST_DEBUG_DUMP_DOT_DIR" in os.environ:
   pipeparts.write_dump_dot(pipeline, "%s_PLAYING" % pipeline.get_name(), verbose = True)

   ## Setup a signal handler to intercept SIGINT in order to write the pipeline graph at ctrl+C before cleanly shutting down
   class SigHandler(simplehandler.OneTimeSignalHandler):
       def do_on_call(self, signum, frame):
           pipeparts.write_dump_dot(pipeline, "%s_SIGINT" % pipeline.get_name(), verbose = True)

   sighandler = SigHandler(pipeline)

mainloop.run()