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.

from optparse import OptionParser
import os
import numpy
import sys

import pygtk
pygtk.require("2.0")
import gobject
gobject.threads_init()
import pygst
pygst.require("0.10")
import gst

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 ligolw
from glue.ligolw import array as ligolw_array
from glue.ligolw import param as ligolw_param
from glue.ligolw import utils as ligolw_utils
from pylal.datatypes import LIGOTimeGPS
from pylal import series as lalseries

class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
    pass
ligolw_array.use_in(LIGOLWContentHandler)
ligolw_param.use_in(LIGOLWContentHandler)

## @file gstlal_fake_frames
# This program will make fake data in a variety of ways; see gstlal_fake_frames for help and usage
#
# 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
#
# @dot
# 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;
# }
# @enddot
#
#
# ### Usage cases ###
#
#
# -# 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
#
#
# -# 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 found <a
#   href=http://www.lsc-group.phys.uwm.edu/cgit/gstlal/plain/gstlal/share/v1_early_asd.txt>here</a>.
#   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
#
#
# -# 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
#
#
# -# 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
#
#
# -# Other things are certainly possible, please add some!
#
# ### Command line options ###
#
# See datasource.append_options() for common command line options shared among different programs
#
#   + `--shift` [int] (ns): Number of nanoseconds, \f$\tau\f$ to delay (negative) or advance (positive) the input time series \f$x\f$ relative to the output time series \f$y\f$. \f$ y(t) = x(t+\tau)\f$
#   + `--sample-rate` [int] (Hz): 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.
#   + `--whiten-reference-psd` [file name]: Set the name of psd xml file to whiten the data with.
#   + `--whiten-track-psd` []: Calculate PSD from input data and track with time.
#   + `--color-psd` [file name]: Set the name of psd xml file to color the data with
#   + `--output-path` [file path]: Path to output frame files. Default = "."
#   + `--output-channel-name` [string]: The name of the channel in the output frames. The default is the same as the channel name.
#   + `--frame-type` [string]: Frame type, required.
#   + `--frame-duration` [int] (s): Set the duration of the output frames.  The duration of the frame file will be multiplied by --frames-per-file.  Default is 16s.
#   + `--frames-per-file` [int]: Set the number of frames per file.  Default is 256.
#   + `--verbose` []: Be verbose.
#
# ### 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
# 
# 
# @cond

##
# Return the number of digits in a number
#
def digits(num):
    return int(numpy.ceil(numpy.log10(float(num))))


##
# 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.")
    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")

    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(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()
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 = lalseries.read_psd_xmldoc(ligolw_utils.load_filename(options.whiten_reference_psd, verbose = options.verbose, contenthandler = LIGOLWContentHandler))[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, seekevent = gw_data_source.seekevent, 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-float, 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 = lalseries.read_psd_xmldoc(ligolw_utils.load_filename(options.color_psd, verbose = options.verbose, contenthandler = LIGOLWContentHandler))[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 = rpsd.data[:max_sample]
    
    # create the coloring FIR kernel from reference_psd.psd_to_fir_kernel()
    fir_matrix, latency, measured_sample_rate = reference_psd.psd_to_fir_kernel(rpsd)

    ## 7) Set the pipeline head to a pipeparts.mkfirbank() recoloring filter
    head = pipeparts.mkfirbank(pipeline, head, latency = latency, fir_matrix = [fir_matrix], block_stride = 32 * 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(100) frames in each directory
head.connect("notify::timestamp", pipeparts.framecpp_filesink_ldas_path_handler, (options.output_path, digits(options.gps_start_time) - digits(options.frame_duration * options.frames_per_file * 100)))

# Run it
if pipeline.set_state(gst.STATE_PLAYING) == gst.STATE_CHANGE_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()

## @endcond