io.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 Option
from gstlal.dags.layers import Layer, Node
 
 
def create_frames_layer(config, dag, psd_cache):
    layer = Layer(
        "gstlal_fake_frames",
        requirements={
            "request_cpus": 2, 
            "request_memory": 2000, 
            "request_disk": "2GB",
            **config.condor.submit
        },
        transfer_files=config.condor.transfer_files,
    )
 
    frame_cache = DataCache.generate(DataType.FRAMES, config.ifo_combos, config.time_bins, create_dirs=False)
 
    # if given a single PSD, assign this to all times
    if len(psd_cache) == 1:
        psds = {key: psd_cache for key in frame_cache.groupby("ifo", "time").keys()}
    else:
        psds = psd_cache.groupby("ifo", "time")
 
    # set options
    common_opts = [
        Option("data-source", "frames"),
        Option("frame-segments-name", config.source.frame_segments_name),
        Option("data-find-server", config.source.data_find_server),
    ]
    if config.frames.injections:
        common_opts.append(Option("injections", config.frames.injections))
    if config.frames.recolor_psd:
        common_opts.append(Option("color-psd", config.frames.recolor_psd))
    if config.frames.track_psd:
        common_opts.append(Option("whiten-track-psd"))
    if config.frames.shift:
        common_opts.append(Option("shift", config.frames.shift))
    if config.frames.frame_duration:
        common_opts.append(Option("output-frame-duration", config.frames.frame_duration))
    if config.frames.frames_per_file:
        common_opts.append(Option("output-frames-per-file", config.frames.frames_per_file))
 
    for (ifo_combo, span), frames in frame_cache.groupby("ifo", "time").items():
        ifos = config.to_ifo_list(ifo_combo)
        start, end = span
 
        for ifo in ifos:
            frame_opts = [
                Option("gps-start-time", int(start)),
                Option("gps-end-time", int(end)),
                Option("channel-name", f"{ifo}={config.source.channel_name[ifo]}"),
                Option("frame-type", f"{ifo}={config.source.frame_type[ifo]}"),
                Option("output-channel-name", f"{ifo}={config.frames.output_channel_name[ifo]}"),
                Option("output-frame-type", f"{ifo}_{config.frames.output_frame_type[ifo]}"),
            ]
            frame_opts.extend(common_opts)
 
            layer += Node(
                arguments = frame_opts,
                inputs = [
                    Option("whiten-reference-psd", psds[(ifo_combo, span)].files),
                    Option("frame-segments-file", config.source.frame_segments_file),
                ],
                outputs = Option("output-path", config.frames.get("output_path", "frames"))
            )
 
    dag.attach(layer)
    return frame_cache
 
 
@plugins.register
def layers():
    return {
        "create_frames": create_frames_layer,
    }