__init__.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.
 
 
import itertools
import getpass
import os
 
import yaml
 
from lal import LIGOTimeGPS
from ligo.lw import utils as ligolw_utils
from ligo.lw.utils import segments as ligolw_segments
from ligo.segments import segment, segmentlist, segmentlistdict
 
from gstlal import segments
from gstlal.dags import profiles
 
 
class Config:
    """
    Hold configuration used for analyzes.
    """
    def __init__(self, **kwargs):
        # normalize options
        kwargs = replace_keys(kwargs)
 
        # basic analysis options
        self.tag = kwargs.get("tag", "test")
        self.rootdir = os.getcwd()
 
        # instrument options
        if isinstance(kwargs["instruments"], list):
            self.ifos = kwargs["instruments"]
        else:
            self.ifos = self.to_ifo_list(kwargs["instruments"])
        self.min_ifos = kwargs.get("min_instruments", 1)
        self.all_ifos = frozenset(self.ifos)
 
        # get instrument combinations
        self.ifo_combos = []
        for n_ifos in range(self.min_ifos, len(self.ifos) + 1):
            for combo in itertools.combinations(self.ifos, n_ifos):
                self.ifo_combos.append(frozenset(combo))
 
        # source options
        self.source = dotdict(replace_keys(kwargs["source"]))
 
        # time options
        if not "start" in kwargs:
            self.span = segment(0, 0)
        else:
            self.start = LIGOTimeGPS(kwargs["start"])
            if "stop" in kwargs:
                self.stop = LIGOTimeGPS(kwargs["stop"])
                self.duration = self.stop - self.start
            else:
                self.duration = kwargs["duration"]
                self.stop = self.start + self.duration
            self.span = segment(self.start, self.stop)
 
        # section-specific options
        if "psd" in kwargs:
            self.psd = dotdict(replace_keys(kwargs["psd"]))
        if "data" in kwargs:
            self.data = dotdict(replace_keys(kwargs["data"]))
        if "frames" in kwargs:
            self.frames = dotdict(replace_keys(kwargs["frames"]))
        if "segments" in kwargs:
            self.segments = dotdict(replace_keys(kwargs["segments"]))
        if "injections" in kwargs:
            self.injections = dotdict(replace_keys(kwargs["injections"]))
 
        # condor options
        self.condor = dotdict(replace_keys(kwargs["condor"]))
 
        # file transfer installed by default
        if self.condor.transfer_files is None:
            self.condor.transfer_files = True
 
        self.condor.submit = self.create_condor_submit_options(
            self.condor,
            x509_proxy=self.source.x509_proxy,
        )
 
        # validate config
        self.validate()
 
    def create_time_bins(
        self,
        start_pad = 512,
        overlap = 512,
        min_instruments = 1,
        one_ifo_only = False,
        one_ifo_length = (3600 * 8)
    ):
        self.time_boundaries = segments.split_segments_by_lock(self.ifos, self.segments, self.span)
        self.time_bins = segmentlistdict()
        if not one_ifo_only:
            for span in self.time_boundaries:
                analysis_segs = segments.analysis_segments(
                    self.ifos,
                    self.segments,
                    span,
                    start_pad=start_pad,
                    overlap=overlap,
                    min_instruments=min_instruments,
                    one_ifo_length=one_ifo_length,
                )
                self.time_bins.extend(analysis_segs)
        else:
            for span in self.time_boundaries:
                time_bin = segmentlistdict()
                for ifo, segs in self.segments.items():
                    ifo_key = frozenset([ifo])
                    segs = segs & segmentlist([span])
                    time_bin[ifo_key] = segments.split_segments(segs, one_ifo_length, start_pad)
                self.time_bins.extend(time_bin)
 
    def create_condor_submit_options(self, condor_config, x509_proxy=False):
        if "accounting_group_user" in condor_config:
            accounting_group_user = condor_config["accounting_group_user"]
        else:
            accounting_group_user = getpass.getuser()
 
        submit_opts = {
            "want_graceful_removal": "True",
            "kill_sig": "15",
            "accounting_group": condor_config["accounting_group"],
            "accounting_group_user": accounting_group_user,
        }
        requirements = []
        environment = {}
 
        # load site profile
        profile = profiles.load_profile(condor_config["profile"])
        assert profile["scheduler"] == "condor", "only scheduler=condor is allowed currently"
 
        # profile options
        if "directives" in profile:
            submit_opts.update(profile["directives"])
        if "requirements" in profile:
            requirements.extend(profile["requirements"])
        if "environment" in profile:
            environment.update(profile["environment"])
 
        # singularity options
        if "singularity_image" in condor_config:
            singularity_image = condor_config["singularity_image"]
            submit_opts['+SingularityImage'] = f'"{singularity_image}"'
            submit_opts['transfer_executable'] = False
 
        # proxy options
        if x509_proxy:
            submit_opts['x509userproxy'] = x509_proxy
            submit_opts['use_x509userproxy'] = True
 
        # file transfer options
        if not self.condor.transfer_files:
            # set the job's working directory to be the current working directory
            submit_opts['initialdir'] = self.rootdir
 
        # config options
        if "directives" in condor_config:
            submit_opts.update(condor_config["directives"])
        if "requirements" in condor_config:
            requirements.extend(condor_config["requirements"])
        if "environment" in condor_config:
            environment.update(condor_config["environment"])
 
        # condor requirements
        submit_opts['requirements'] = " && ".join(requirements)
 
        # condor environment
        env_opts = [f"{key}={val}" for (key,val) in environment.items()]
        if "environment" in submit_opts:
            env_opts.append(submit_opts["environment"].strip('"'))
        submit_opts["environment"] = f'"{" ".join(env_opts)}"'
 
        return submit_opts
 
    def validate(self):
        """
        Validate configuration settings.
        """
        pass
 
    def setup(self):
        """
        Load segments and create time bins.
        """
        if "frame_segments_file" in self.source:
            xmldoc = ligolw_utils.load_filename(
                self.source.frame_segments_file,
                contenthandler=ligolw_segments.LIGOLWContentHandler
            )
            self.segments = ligolw_segments.segmenttable_get_by_name(xmldoc, "datasegments").coalesce()
        else:
            self.segments = segmentlistdict((ifo, segmentlist([self.span])) for ifo in self.ifos)
 
        if self.span != segment(0, 0):
            self.create_time_bins(start_pad=512, min_instruments=self.min_ifos)
 
    @staticmethod
    def to_ifo_list(ifos):
        """
        Given a string of IFO pairs (e.g. H1L1), return a list of IFOs.
        """
        return [ifos[2*n:2*n+2] for n in range(len(ifos) // 2)]
 
    @classmethod
    def load(cls, path):
        """
        Load configuration from a file given a file path.
        """
        with open(path, "r") as f:
            return cls(**yaml.safe_load(f))
 
 
class dotdict(dict):
    """
    A dictionary supporting dot notation.
 
    Implementation from https://gist.github.com/miku/dc6d06ed894bc23dfd5a364b7def5ed8.
 
    """
    __getattr__ = dict.get
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__
 
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        for k, v in self.items():
            if isinstance(v, dict):
                self[k] = dotdict(v)
 
 
def replace_keys(dict_, reverse=False):
    out = dict(dict_)
    for k, v in out.items():
        if isinstance(v, dict):
            out[k] = replace_keys(v)
    return {k.replace("_", "-") if reverse else k.replace("-", "_") : v for k, v in out.items()}