gstlal_fake_frames

#!/usr/bin/env python3
#
# 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 lal
 
from gstlal import kernels
from gstlal import pipeparts
from gstlal import datasource
from gstlal.psd import read_psd
from gstlal.stream import Stream
 

 
 

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 = "Whiten the data by tracking PSD changes, can be combined with --whiten-reference-psd to seed the process.")
    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("--output-frame-type", metavar = "name", help = "Frame type, required")
    parser.add_option("--output-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("--output-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.output_frame_type is None:
        raise ValueError("--frame-type is required")
 
 
    return options, filenames
 
 
#
# Main
#
 
options, filenames = parse_command_line()
 

os.makedirs(options.output_path, exist_ok=True)
 

gw_data_source = datasource.DataSourceInfo.from_optparse(options)
 

instrument = list(gw_data_source.channel_dict.keys())[0]
 
# 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, we will do them later
injections, gw_data_source.injection_filename = options.injection_file, None
 

gw_data_source.block_size = 8 * options.sample_rate
stream = Stream.from_datasource(
    gw_data_source,
    instrument,
    name=os.path.split(sys.argv[0])[1],
    verbose=options.verbose
)
 

if options.verbose:
    stream = stream.progressreport("frames")
 
if options.shift is not None:
    
    stream = stream.shift(shift=options.shift)
 
    
    stream = stream.progressreport("frames_shifted")
 
if options.whiten_reference_psd:
    
    wpsd = read_psd(options.whiten_reference_psd, verbose = options.verbose)[instrument]
else:
    
    wpsd = None
 
if options.whiten_reference_psd or options.whiten_track_psd:
    
    stream = stream.condition(options.sample_rate, instrument, psd=wpsd, track_psd=options.whiten_track_psd)
else:
    
    stream = stream.resample(quality=9).capsfilter(f"audio/x-raw, rate={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":
        renormalize_factor = 1.0 / (pipeparts.audioresample_variance_gain(9, 16384, options.sample_rate))**.5
        stream = stream.audioamplify(renormalize_factor)
 
# Apply a coloring filter
if options.color_psd:
 
    
    rpsd = read_psd(options.color_psd, verbose = options.verbose)[instrument]
    
    
    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])
    if len(rpsd_data) < max_sample: # FIXME this is off by one sample, but shouldn't be.
        rpsd_data = numpy.append(rpsd_data, numpy.zeros(max_sample - len(rpsd_data)))
 
    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 = kernels.PSDFirKernel()
    fir_matrix, latency, measured_sample_rate = FIRKernel.psd_to_linear_phase_whitening_fir_kernel(rpsd, invert = False)
 
    
    stream = stream.firbank(latency=latency, fir_matrix=[fir_matrix], block_stride=(1 * options.sample_rate))
 
 

tagstr = f"units=strain,channel-name={options.output_channel_name},instrument={instrument}"
 

stream = stream.taginject(tagstr)
 

if injections is not None:
    stream = stream.injections(injections)
 

stream = stream.framecppchannelmux(
    channels=f"{instrument}:{options.output_channel_name}",
    frame_duration=options.output_frame_duration,
    frames_per_file=options.output_frames_per_file
)
 

framesink = stream.framecppfilesink(frame_type=options.output_frame_type)
 
# Put O(100000 s) frames in each directory
framesink.connect("notify::timestamp", pipeparts.framecpp_filesink_ldas_path_handler, (options.output_path, 5))
 
# Run it
stream.start()