Source code for idq.configparser

from collections import defaultdict
import os

import toml

from . import synthetic


class Config(object):
    """
    Store central configuration for iDQ.
    """

    def __init__(self, **kwargs):
        # load general parameters
        self.tag = kwargs["general"]["tag"]
        self.instrument = kwargs["general"]["instrument"]
        self.classifiers = set(kwargs["general"]["classifiers"])

        if "rootdir" in kwargs["general"]:
            self.rootdir = kwargs["general"]["rootdir"]
        else:
            self.rootdir = os.getcwd()

        # load sections
        self.defaults = kwargs["defaults"]
        self.samples = add_missing_kwargs(kwargs["samples"], **self.defaults)
        self.segments = add_missing_kwargs(kwargs["segments"], **self.defaults)
        self.condor = kwargs["condor"]

        self.train = add_missing_kwargs(kwargs["train"], **self.defaults)
        self.evaluate = add_missing_kwargs(kwargs["evaluate"], **self.defaults)
        self.calibrate = add_missing_kwargs(kwargs["calibrate"], **self.defaults)
        self.timeseries = add_missing_kwargs(kwargs["timeseries"], **self.defaults)
        self.report = add_missing_kwargs(kwargs["report"], **self.defaults)
        self.monitor = kwargs.get("monitor", {})

        # inherit instrument config from general section
        self.features = add_missing_kwargs(kwargs["features"], **self.defaults)
        self.features["instrument"] = self.instrument

        # load classifier config
        self.classifier_map = {
            c["name"]: add_missing_kwargs(c, **self.defaults)
            for c in kwargs["classifier"]
        }

        # add column / bound defaults
        for config in self.classifier_map.values():
            # column defaults
            for column in ("time", "significance", "frequency"):
                if column in self.features:
                    config.setdefault(column, self.features[column])
            # bound defaults
            config.setdefault("bounds", config2bounds(self.features.get("bounds", {})))

    @classmethod
    def load(cls, path):
        with open(path, "r") as f:
            return cls(**toml.load(f))

    def dump(self, path, rootdir=None):
        if not rootdir:
            rootdir = self.rootdir
        config = {
            "general": {
                "rootdir": rootdir,
                "tag": self.tag,
                "instrument": self.instrument,
                "classifiers": list(self.classifiers),
            },
            "samples": self.samples,
            "features": self.features,
            "condor": self.condor,
            "segments": self.segments,
            "train": self.train,
            "evaluate": self.evaluate,
            "calibrate": self.calibrate,
            "timeseries": self.timeseries,
            "report": self.report,
            "monitor": self.monitor,
            "classifier": list(self.classifier_map.values()),
            "defaults": self.defaults,
        }
        with open(path, "w") as f:
            f.write(toml.dumps(config))


def path2config(path, verbose=False):
    """
    read in the config from a path
    """
    if verbose:
        print("reading config from: " + path)
    if not os.path.exists(path):
        raise OSError("file does not exist!: " + path)
    return Config.load(path)


def config2bounds(bounds):
    """
    a special parser for setting up bounds for trigger sets
    """
    ans = {}
    for key, bound in bounds.items():
        min_val, max_val = bound
        try:
            min_val = int(min_val)
        except ValueError:
            min_val = float(min_val)
        try:
            max_val = int(max_val)
        except ValueError:
            max_val = float(max_val)

        ans[key] = (min_val, max_val)
    return ans


def add_missing_kwargs(kwargs, **new_kwargs):
    """
    add kwargs to existing ones if they are not present.
    """
    for key, val in new_kwargs.items():
        if key not in kwargs:
            kwargs[key] = val

    return kwargs


def path2streams(path, segs, verbose=False):
    """
    read in parameters from config and set up the appropriate data structures
    """
    with open(path, "r") as f:
        config = toml.load(f)

    # iterate through sections (streams) and add their configs to the object
    stream_names = set(stream["name"] for stream in config["stream"])

    streams = []
    channel2streams = defaultdict(list)
    for stream_name, options in zip(stream_names, config["stream"]):
        # iterate through required fields
        if "rate" not in options:
            raise KeyError('must specify "rate" for ' + stream_name)
        rate = options["rate"]

        if "frequency" not in options["dist"]:
            raise KeyError('must specify "dist.frequency" section for ' + stream_name)
        freq_opts = options["dist"]["frequency"]
        freq_distrib = synthetic.DistribWrapper(freq_opts["type"], *freq_opts["args"])

        if "snr" not in options["dist"]:
            raise KeyError('must specify "dist.snr" section for ' + stream_name)
        snr_opts = options["dist"]["snr"]
        snr_distrib = synthetic.DistribWrapper(snr_opts["type"], *snr_opts["args"])

        # instantiate the stream object
        stream = synthetic.SyntheticTriggerStream(
            stream_name, segs, rate, freq_distrib, snr_distrib
        )
        streams.append(stream)

        # process channels and jitters if available
        if "jitter" in config:
            for jitters in config["jitter"]:
                channel = jitters.pop("name")
                channel2streams[channel].append((stream, jitters))

    return streams, channel2streams