psd.py

# Copyright (C) 2020  Patrick Godwin (patrick.godwin@ligo.org)
#
# 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 gstlal import plugins
from gstlal.datafind import DataType, DataCache
from gstlal.dags import Argument, Option
from gstlal.dags import util as dagutil
from gstlal.dags.layers import Layer, Node
 
 
def reference_psd_layer(config, dag):
    layer = Layer(
        "gstlal_reference_psd",
        requirements={
            "request_cpus": 2, 
            "request_memory": 2000, 
            "request_disk": "2GB",
            **config.condor.submit
        },
        transfer_files=config.condor.transfer_files,
    )
 
    psd_cache = DataCache.generate(DataType.REFERENCE_PSD, config.ifo_combos, config.time_bins)
 
    frame_opts = [
        Option("data-source", "frames"),
        Option("psd-fft-length", config.psd.fft_length),
        Option("frame-segments-name", config.source.frame_segments_name),
    ]
 
    for (ifo_combo, span), psds in psd_cache.groupby("ifo", "time").items():
        ifos = config.to_ifo_list(ifo_combo)
        start, end = span
 
        arguments = [
            Option("gps-start-time", int(start)),
            Option("gps-end-time", int(end)),
            Option("channel-name", dagutil.format_ifo_args(ifos, config.source.channel_name)),
            *frame_opts,
        ]
        inputs = [Option("frame-segments-file", config.source.frame_segments_file, track=False)]
 
        if config.source.frame_cache:
            inputs.append(Option("frame-cache", config.source.frame_cache, track=False))
        else:
            arguments.extend([
                Option("frame-type", dagutil.format_ifo_args(ifos, config.source.frame_type)),
                Option("data-find-server", config.source.data_find_server),
            ])
 
        layer += Node(
            arguments = arguments,
            inputs = inputs,
            outputs = Option("write-psd", psds.files)
        )
 
    dag.attach(layer)
    return psd_cache
 
 
def median_psd_layer(config, dag, psd_cache):
    layer = Layer(
        "gstlal_median_of_psds",
        requirements={
            "request_cpus": 2, 
            "request_memory": 2000, 
            "request_disk": "2GB",
            **config.condor.submit
        },
        transfer_files=config.condor.transfer_files,
    )
 
    median_psd_cache = DataCache.generate(DataType.MEDIAN_PSD, config.all_ifos, config.span)
 
    layer += Node(
        inputs = Argument("psds", psd_cache.files),
        outputs = Option("output-name", median_psd_cache.files)
    )
 
    dag.attach(layer)
    return median_psd_cache
 
 
def smoothen_psd_layer(config, dag, psd_cache):
    layer = Layer(
        "gstlal_psd_polyfit",
        requirements={
            "request_cpus": 2, 
            "request_memory": 2000, 
            "request_disk": "2GB",
            **config.condor.submit
        },
        transfer_files=config.condor.transfer_files,
    )
 
    smooth_psd_cache = DataCache.generate(DataType.SMOOTH_PSD, config.ifo_combos, config.time_bins)
 
    smooth_psds = smooth_psd_cache.groupby("ifo", "time")
    for (ifo_combo, span), psds in psd_cache.groupby("ifo", "time").items():
        layer += Node(
            arguments = Option("low-fit-freq", 10),
            inputs = Argument("psds", psd_cache.files),
            outputs = Option("output-name", smooth_psds[span].files)
        )
 
    dag.attach(layer)
    return smooth_psd_cache
 
 
@plugins.register
def layers():
    return {
        "reference_psd": reference_psd_layer,
        "median_psd": median_psd_layer,
        "smoothen_psd": smoothen_psd_layer,
    }