<commit_before>import ROOT
import copy
import dill
import pathos.multiprocessing as mp
from functools import partial
from PyAnalysisTools.base import _logger, InvalidInputError
from PyAnalysisTools.PlottingUtils.PlotConfig import parse_and_build_plot_config, parse_and_build_process_config, \
    get_histogram_definition, find_process_config
from PyAnalysisTools.base.YAMLHandle import YAMLLoader
from PyAnalysisTools.ROOTUtils.FileHandle import FileHandle
from PyAnalysisTools.PlottingUtils import set_batch_mode
from PyAnalysisTools.PlottingUtils import Formatting as FM
from PyAnalysisTools.PlottingUtils import HistTools as HT
from PyAnalysisTools.PlottingUtils import PlottingTools as PT
from PyAnalysisTools.PlottingUtils.PlotConfig import PlotConfig
from PyAnalysisTools.AnalysisTools.XSHandle import XSHandle
from PyAnalysisTools.ROOTUtils.ObjectHandle import merge_objects_by_process_type
from PyAnalysisTools.AnalysisTools import StatisticsTools as ST
from PyAnalysisTools.base.OutputHandle import OutputFileHandle
from PyAnalysisTools.ROOTUtils.ObjectHandle import get_objects_from_canvas_by_type

class BasePlotter(object):
    def __init__(self, **kwargs):
        if "input_files" not in kwargs:
            _logger.error("No input files provided")
            raise  InvalidInputError("No input files")
        if "plot_config_file" not in kwargs:
            _logger.error("No plot config file provided. Nothing to parse")
            raise InvalidInputError("No plot config file")
        if "process_config_file" not in kwargs:
            _logger.warning("No process config file provided. Unable to read process specific options.")
        if "xs_config_file" not in kwargs:
            _logger.error("No cross section file provided. No scaling will be applied.")
            kwargs.setdefault("xs_config_file", None)
        kwargs.setdefault("systematics", None)
        kwargs.setdefault("process_config_file", None)
        kwargs.setdefault("xs_config_file", None)
        kwargs.setdefault("batch", True)
        kwargs.setdefault("output_file_name", "plots.root")
        for k,v in kwargs.iteritems():
            setattr(self, k, v)
        set_batch_mode(kwargs["batch"])
        self.file_handles = [FileHandle(file_name=input_file, dataset_info=kwargs["xs_config_file"]) for input_file in self.input_files]
        self.xs_handle = XSHandle(kwargs["xs_config_file"])
        self.process_configs = self.parse_process_config()
        FM.load_atlas_style()
        self.statistical_uncertainty_hist = None
        self.histograms = {}
        self.initialise()
        self.event_yields = {}
        self.output_handle = OutputFileHandle(make_plotbook=self.common_config.make_plot_book, **kwargs)

    def parse_plot_config(self):
        _logger.debug("Try to parse plot config file")
        self.plot_configs, self.common_config = parse_and_build_plot_config(self.plot_config_file)
        if not hasattr(self, "lumi"):
            self.lumi = self.common_config.lumi

    def parse_process_config(self):
        if self.process_config_file is None:
            return None
        process_config = parse_and_build_process_config(self.process_config_file)
        return process_config

    def initialise(self):
        self.parse_plot_config()
        self.ncpu = min(self.ncpu, len(self.plot_configs))

    def read_cutflows(self):
        for file_handle in self.file_handles:
            if file_handle.process in self.event_yields:
                self.event_yields[file_handle.process] += file_handle.get_number_of_total_events()
            else:
                self.event_yields[file_handle.process] = file_handle.get_number_of_total_events()

    def retrieve_histogram(self, file_handle, plot_config):
        file_handle.open()
        hist = get_histogram_definition(plot_config)
        hist.SetName(hist.GetName() + file_handle.process)
        try:
            weight = None
            selection_cuts = ""
            if self.common_config.weight and plot_config.weight is not None:
                weight = self.common_config.weight
            if self.common_config.cuts:
                selection_cuts += "&&".join(self.common_config.cuts)
            if self.common_config.blind and self.process_configs[file_handle.process].type == "Data":
                if len(selection_cuts) != 0:
                    selection_cuts += " && "
                selection_cuts += " !({:s})".format(" && ".join(self.common_config.blind))
            file_handle.fetch_and_link_hist_to_tree(self.tree_name, hist, plot_config.dist, selection_cuts,
                                                    tdirectory=self.systematics, weight=weight)
            hist.SetName(hist.GetName() + "_" + file_handle.process)
            _logger.debug("try to access config for process %s" % file_handle.process)
            process_config = find_process_config(file_handle.process, self.process_configs)
            if process_config is None:
                _logger.error("Could not find process config for {:s}".format(file_handle.process))
                return None

        except Exception as e:
            raise e

        return hist

    def merge_histograms(self):
        def merge(histograms):
            for process, process_config in self.process_configs.iteritems():
                if not hasattr(process_config, "subprocesses"):
                    continue
                for sub_process in process_config.subprocesses:
                    if sub_process not in histograms.keys():
                        continue
                    if process not in histograms.keys():
                        new_hist_name = histograms[sub_process].GetName().replace(sub_process, process)
                        histograms[process] = histograms[sub_process].Clone(new_hist_name)
                    else:
                        histograms[process].Add(histograms[sub_process])
                    histograms.pop(sub_process)
        for plot_config, histograms in self.histograms.iteritems():
            merge(histograms)

    def calculate_ratios(self, hists, plot_config):
        ratio = hists["Data"].Clone("ratio_%s" % plot_config.dist)
        mc = None
        for key, hist in hists.iteritems():
            if key == "Data":
                continue
            if mc is None:
                mc = hist.Clone("mc_total_%s" % plot_config.dist)
                continue
            mc.Add(hist)
        ratio.Divide(mc)
        plot_config.name = "ratio_" + plot_config.name
        plot_config.ytitle = "data/MC"
        if "unit" in plot_config.__dict__.keys():
            plot_config.__dict__.pop("unit")
        plot_config.draw = "Marker"
        #ratios = [self.calculate_ratio(hist, reference) for hist in hists]
        if self.statistical_uncertainty_hist:
            plot_config_stat_unc_ratio = copy.copy(plot_config)
            plot_config_stat_unc_ratio.color = ROOT.kYellow
            plot_config_stat_unc_ratio.style = 1001
            plot_config_stat_unc_ratio.draw = "E2"
            plot_config_stat_unc_ratio.logy = False
            statistical_uncertainty_ratio = self.statistical_uncertainty_hist.Clone("stat.unc.ratio")
            statistical_uncertainty_ratio.Divide(mc)
            canvas = PT.plot_hist(statistical_uncertainty_ratio, plot_config_stat_unc_ratio)
            PT.add_histogram_to_canvas(canvas, ratio, plot_config)
        else:
            canvas = PT.plot_hist(ratio, plot_config)
        return canvas

    def fetch_histograms(self, file_handle, plot_config):
        if "data" in file_handle.process.lower() and (plot_config.no_data or self.common_config.no_data):
            return
        tmp = self.retrieve_histogram(file_handle, plot_config)
        return file_handle.process, tmp

    def read_histograms(self, plot_config):
        histograms = mp.ThreadPool(min(self.nfile_handles,
                                       len(self.file_handles))).map(partial(self.fetch_histograms,
                                                                            plot_config=plot_config), self.file_handles)
        return plot_config, histograms

    def apply_lumi_weights(self):
        for hist_set in self.histograms.values():
            for process, hist in hist_set.iteritems():
                if "data" in process:
                    continue
                cross_section_weight = self.xs_handle.get_lumi_scale_factor(process, self.lumi, self.event_yields[process])
                HT.scale(hist, cross_section_weight)

    def categorise_histograms(self, plot_config, histograms):
        _logger.debug("categorising {:d} histograms".format(len(histograms)))
        for process, hist in histograms:
            try:
                if process not in self.histograms[plot_config].keys():
                    self.histograms[plot_config][process] = hist
                else:
                    self.histograms[plot_config][process].Add(hist)
            except KeyError:
                self.histograms[plot_config] = {process: hist}

    def make_multidimensional_plot(self, plot_config, data):
        for process, histogram in data.iteritems():
            canvas = PT.plot_hist(histogram, plot_config)
            canvas.SetName("{:s}_{:s}".format(canvas.GetName(), process))
            canvas.SetRightMargin(0.15)
            FM.decorate_canvas(canvas, self.common_config, plot_config)
            self.output_handle.register_object(canvas)

    def make_plots(self):
        self.read_cutflows()
        fetched_histograms = mp.ThreadPool(min(self.ncpu, len(self.plot_configs))).map(self.read_histograms,
                                                                                self.plot_configs)

        for plot_config, histograms in fetched_histograms:
            histograms = filter(lambda hist: hist is not None, histograms)
            self.categorise_histograms(plot_config, histograms)
        self.apply_lumi_weights()
        #workaround due to missing worker node communication of regex process parsing
        for hist_set in self.histograms.values():
            for process_name in hist_set.keys():
                _ = find_process_config(process_name, self.process_configs)
        if self.common_config.merge:
            self.merge_histograms()
        for plot_config, data in self.histograms.iteritems():
            if self.common_config.normalise or plot_config.normalise:
                HT.normalise(data)
            if self.common_config.outline == "hist" and not plot_config.is_multidimensional:
                canvas = PT.plot_histograms(data, plot_config, self.common_config, self.process_configs)
            elif self.common_config.outline == "stack" and not plot_config.is_multidimensional:
                canvas = PT.plot_stack(data, plot_config, self.common_config, self.process_configs)
                stack = get_objects_from_canvas_by_type(canvas, "THStack")[0]
                self.statistical_uncertainty_hist = ST.get_statistical_uncertainty_from_stack(stack)
                #todo: temporary fix
                self.statistical_uncertainty_hist.SetMarkerStyle(1)
                plot_config_stat_unc = PlotConfig(name="stat.unc", dist=None, label="stat unc", draw="E2", style=3244,
                                                  color=ROOT.kBlack)
                PT.add_histogram_to_canvas(canvas, self.statistical_uncertainty_hist, plot_config_stat_unc)
                self.process_configs[plot_config_stat_unc.name] = plot_config_stat_unc
            elif plot_config.is_multidimensional:
                self.make_multidimensional_plot(plot_config, data)
                continue
            else:
                _logger.error("Unsupported outline option %s" % self.common_config.outline)
                raise InvalidInputError("Unsupported outline option")
            FM.decorate_canvas(canvas, self.common_config, plot_config)
            FM.add_legend_to_canvas(canvas, process_configs=self.process_configs)
            if hasattr(plot_config, "calcsig"):
                #todo: "Background" should be an actual type
                signal_hist = merge_objects_by_process_type(canvas, self.process_configs, "Signal")
                background_hist = merge_objects_by_process_type(canvas, self.process_configs, "Background")
                significance_hist = ST.get_significance(signal_hist, background_hist)
                canvas_significance_ratio = PT.add_ratio_to_canvas(canvas, significance_hist)
            self.output_handle.register_object(canvas)
            if hasattr(plot_config, "ratio") or hasattr(self.common_config, "ratio"):
                if plot_config.no_data or self.common_config.no_data or plot_config.is_multidimensional:
                    continue
                if self.common_config.ratio:
                    canvas_ratio = self.calculate_ratios(data, plot_config)
                    canvas_combined = PT.add_ratio_to_canvas(canvas, canvas_ratio)
                    self.output_handle.register_object(canvas_combined)
<commit_msg>implement less pedantic behaviour in case of missing metadata<commit_after>import ROOT
import copy
import dill
import pathos.multiprocessing as mp
from functools import partial
from PyAnalysisTools.base import _logger, InvalidInputError
from PyAnalysisTools.PlottingUtils.PlotConfig import parse_and_build_plot_config, parse_and_build_process_config, \
    get_histogram_definition, find_process_config
from PyAnalysisTools.base.YAMLHandle import YAMLLoader
from PyAnalysisTools.ROOTUtils.FileHandle import FileHandle
from PyAnalysisTools.PlottingUtils import set_batch_mode
from PyAnalysisTools.PlottingUtils import Formatting as FM
from PyAnalysisTools.PlottingUtils import HistTools as HT
from PyAnalysisTools.PlottingUtils import PlottingTools as PT
from PyAnalysisTools.PlottingUtils.PlotConfig import PlotConfig
from PyAnalysisTools.AnalysisTools.XSHandle import XSHandle
from PyAnalysisTools.ROOTUtils.ObjectHandle import merge_objects_by_process_type
from PyAnalysisTools.AnalysisTools import StatisticsTools as ST
from PyAnalysisTools.base.OutputHandle import OutputFileHandle
from PyAnalysisTools.ROOTUtils.ObjectHandle import get_objects_from_canvas_by_type

class BasePlotter(object):
    def __init__(self, **kwargs):
        if "input_files" not in kwargs:
            _logger.error("No input files provided")
            raise  InvalidInputError("No input files")
        if "plot_config_file" not in kwargs:
            _logger.error("No plot config file provided. Nothing to parse")
            raise InvalidInputError("No plot config file")
        if "process_config_file" not in kwargs:
            _logger.warning("No process config file provided. Unable to read process specific options.")
        if "xs_config_file" not in kwargs:
            _logger.error("No cross section file provided. No scaling will be applied.")
            kwargs.setdefault("xs_config_file", None)
        kwargs.setdefault("systematics", None)
        kwargs.setdefault("process_config_file", None)
        kwargs.setdefault("xs_config_file", None)
        kwargs.setdefault("batch", True)
        kwargs.setdefault("output_file_name", "plots.root")
        for k,v in kwargs.iteritems():
            setattr(self, k, v)
        set_batch_mode(kwargs["batch"])
        self.file_handles = [FileHandle(file_name=input_file, dataset_info=kwargs["xs_config_file"]) for input_file in self.input_files]
        self.xs_handle = XSHandle(kwargs["xs_config_file"])
        self.process_configs = self.parse_process_config()
        FM.load_atlas_style()
        self.statistical_uncertainty_hist = None
        self.histograms = {}
        self.initialise()
        self.event_yields = {}
        self.output_handle = OutputFileHandle(make_plotbook=self.common_config.make_plot_book, **kwargs)

    def parse_plot_config(self):
        _logger.debug("Try to parse plot config file")
        self.plot_configs, self.common_config = parse_and_build_plot_config(self.plot_config_file)
        if not hasattr(self, "lumi"):
            self.lumi = self.common_config.lumi

    def parse_process_config(self):
        if self.process_config_file is None:
            return None
        process_config = parse_and_build_process_config(self.process_config_file)
        return process_config

    def initialise(self):
        self.parse_plot_config()
        self.ncpu = min(self.ncpu, len(self.plot_configs))

    def read_cutflows(self):
        for file_handle in self.file_handles:
            if file_handle.process in self.event_yields:
                self.event_yields[file_handle.process] += file_handle.get_number_of_total_events()
            else:
                self.event_yields[file_handle.process] = file_handle.get_number_of_total_events()

    def retrieve_histogram(self, file_handle, plot_config):
        file_handle.open()
        hist = get_histogram_definition(plot_config)
        hist.SetName(hist.GetName() + file_handle.process)
        try:
            weight = None
            selection_cuts = ""
            if self.common_config.weight and plot_config.weight is not None:
                weight = self.common_config.weight
            if self.common_config.cuts:
                selection_cuts += "&&".join(self.common_config.cuts)
            if self.common_config.blind and self.process_configs[file_handle.process].type == "Data":
                if len(selection_cuts) != 0:
                    selection_cuts += " && "
                selection_cuts += " !({:s})".format(" && ".join(self.common_config.blind))
            file_handle.fetch_and_link_hist_to_tree(self.tree_name, hist, plot_config.dist, selection_cuts,
                                                    tdirectory=self.systematics, weight=weight)
            hist.SetName(hist.GetName() + "_" + file_handle.process)
            _logger.debug("try to access config for process %s" % file_handle.process)
            process_config = find_process_config(file_handle.process, self.process_configs)
            if process_config is None:
                _logger.error("Could not find process config for {:s}".format(file_handle.process))
                return None

        except Exception as e:
            raise e

        return hist

    def merge_histograms(self):
        def merge(histograms):
            for process, process_config in self.process_configs.iteritems():
                if not hasattr(process_config, "subprocesses"):
                    continue
                for sub_process in process_config.subprocesses:
                    if sub_process not in histograms.keys():
                        continue
                    if process not in histograms.keys():
                        new_hist_name = histograms[sub_process].GetName().replace(sub_process, process)
                        histograms[process] = histograms[sub_process].Clone(new_hist_name)
                    else:
                        histograms[process].Add(histograms[sub_process])
                    histograms.pop(sub_process)
        for plot_config, histograms in self.histograms.iteritems():
            merge(histograms)

    def calculate_ratios(self, hists, plot_config):
        if not "Data" in hists:
            _logger.error("Ratio requested but no data found")
            raise InvalidInputError("Missing data")
        ratio = hists["Data"].Clone("ratio_%s" % plot_config.dist)
        mc = None
        for key, hist in hists.iteritems():
            if key == "Data":
                continue
            if mc is None:
                mc = hist.Clone("mc_total_%s" % plot_config.dist)
                continue
            mc.Add(hist)
        ratio.Divide(mc)
        if hasattr(self.common_config, "ratio_config"):
            plot_config = self.common_config.ratio_config
        else:
            plot_config.name = "ratio_" + plot_config.name
            plot_config.ytitle = "data/MC"
        if "unit" in plot_config.__dict__.keys():
            plot_config.__dict__.pop("unit")
            plot_config.draw = "Marker"
        #ratios = [self.calculate_ratio(hist, reference) for hist in hists]
        if self.statistical_uncertainty_hist:
            plot_config_stat_unc_ratio = copy.copy(plot_config)
            plot_config_stat_unc_ratio.color = ROOT.kYellow
            plot_config_stat_unc_ratio.style = 1001
            plot_config_stat_unc_ratio.draw = "E2"
            plot_config_stat_unc_ratio.logy = False
            statistical_uncertainty_ratio = self.statistical_uncertainty_hist.Clone("stat.unc.ratio")
            statistical_uncertainty_ratio.Divide(mc)
            canvas = PT.plot_hist(statistical_uncertainty_ratio, plot_config_stat_unc_ratio)
            PT.add_histogram_to_canvas(canvas, ratio, plot_config)
        else:
            canvas = PT.plot_hist(ratio, plot_config)
        return canvas

    def fetch_histograms(self, file_handle, plot_config):
        if "data" in file_handle.process.lower() and (plot_config.no_data or self.common_config.no_data):
            return
        tmp = self.retrieve_histogram(file_handle, plot_config)
        return file_handle.process, tmp

    def read_histograms(self, plot_config):
        histograms = mp.ThreadPool(min(self.nfile_handles,
                                       len(self.file_handles))).map(partial(self.fetch_histograms,
                                                                            plot_config=plot_config), self.file_handles)
        return plot_config, histograms

    def apply_lumi_weights(self):
        for hist_set in self.histograms.values():
            for process, hist in hist_set.iteritems():
                if "data" in process:
                    continue
                cross_section_weight = self.xs_handle.get_lumi_scale_factor(process, self.lumi, self.event_yields[process])
                HT.scale(hist, cross_section_weight)

    def categorise_histograms(self, plot_config, histograms):
        _logger.debug("categorising {:d} histograms".format(len(histograms)))
        for process, hist in histograms:
            try:
                if process not in self.histograms[plot_config].keys():
                    self.histograms[plot_config][process] = hist
                else:
                    self.histograms[plot_config][process].Add(hist)
            except KeyError:
                self.histograms[plot_config] = {process: hist}

    def make_multidimensional_plot(self, plot_config, data):
        for process, histogram in data.iteritems():
            canvas = PT.plot_hist(histogram, plot_config)
            canvas.SetName("{:s}_{:s}".format(canvas.GetName(), process))
            canvas.SetRightMargin(0.15)
            FM.decorate_canvas(canvas, self.common_config, plot_config)
            self.output_handle.register_object(canvas)

    def make_plots(self):
        self.read_cutflows()
        fetched_histograms = mp.ThreadPool(min(self.ncpu, len(self.plot_configs))).map(self.read_histograms,
                                                                                self.plot_configs)

        for plot_config, histograms in fetched_histograms:
            histograms = filter(lambda hist: hist is not None, histograms)
            self.categorise_histograms(plot_config, histograms)
        self.apply_lumi_weights()
        #workaround due to missing worker node communication of regex process parsing
        for hist_set in self.histograms.values():
            for process_name in hist_set.keys():
                _ = find_process_config(process_name, self.process_configs)
        if self.common_config.merge:
            self.merge_histograms()
        for plot_config, data in self.histograms.iteritems():
            data = {k: v for k, v in data.iteritems() if v}
            if self.common_config.normalise or plot_config.normalise:
                HT.normalise(data)
            if self.common_config.outline == "hist" and not plot_config.is_multidimensional:
                canvas = PT.plot_histograms(data, plot_config, self.common_config, self.process_configs)
            elif self.common_config.outline == "stack" and not plot_config.is_multidimensional:
                canvas = PT.plot_stack(data, plot_config, self.common_config, self.process_configs)
                stack = get_objects_from_canvas_by_type(canvas, "THStack")[0]
                self.statistical_uncertainty_hist = ST.get_statistical_uncertainty_from_stack(stack)
                #todo: temporary fix
                self.statistical_uncertainty_hist.SetMarkerStyle(1)
                plot_config_stat_unc = PlotConfig(name="stat.unc", dist=None, label="stat unc", draw="E2", style=3244,
                                                  color=ROOT.kBlack)
                PT.add_histogram_to_canvas(canvas, self.statistical_uncertainty_hist, plot_config_stat_unc)
                self.process_configs[plot_config_stat_unc.name] = plot_config_stat_unc
            elif plot_config.is_multidimensional:
                self.make_multidimensional_plot(plot_config, data)
                continue
            else:
                _logger.error("Unsupported outline option %s" % self.common_config.outline)
                raise InvalidInputError("Unsupported outline option")
            FM.decorate_canvas(canvas, self.common_config, plot_config)
            FM.add_legend_to_canvas(canvas, process_configs=self.process_configs)
            if hasattr(plot_config, "calcsig"):
                #todo: "Background" should be an actual type
                signal_hist = merge_objects_by_process_type(canvas, self.process_configs, "Signal")
                background_hist = merge_objects_by_process_type(canvas, self.process_configs, "Background")
                significance_hist = ST.get_significance(signal_hist, background_hist)
                canvas_significance_ratio = PT.add_ratio_to_canvas(canvas, significance_hist)
            self.output_handle.register_object(canvas)
            if hasattr(plot_config, "ratio") or hasattr(self.common_config, "ratio"):
                if plot_config.no_data or self.common_config.no_data or plot_config.is_multidimensional:
                    continue
                if self.common_config.ratio:
                    try:
                        canvas_ratio = self.calculate_ratios(data, plot_config)
                        canvas_combined = PT.add_ratio_to_canvas(canvas, canvas_ratio)
                        self.output_handle.register_object(canvas_combined)
                    except InvalidInputError:
                        pass
<|endoftext|>


<commit_before>#!/usr/bin/python

# see http://docs.fedoraproject.org/en-US/Fedora_Draft_Documentation/0.1/html/RPM_Guide/ch16s04.html

import sys
import argparse
import os
import platform
import urlparse
import pkg

from scripts.lib import mappkgname

def build_type():
    debian_like = ["ubuntu", "debian"]
    rhel_like = ["fedora", "redhat", "centos"]

    dist = platform.linux_distribution(full_distribution_name=False)[0].lower()
    assert dist in debian_like + rhel_like

    if dist in debian_like:
        return "deb"
    elif dist in rhel_like:
        return "rpm"

def map_package_name_deb(name):
    """Map RPM package name to equivalent Deb names"""
    return mappkgname.map_package(name)

# Rules to build SRPM from SPEC
def build_srpm_from_spec(spec):
    srpmpath = spec.source_package_path()
    print '%s: %s %s' % (srpmpath, spec.specpath(),
                         " ".join(spec.source_paths()))

# Rules to download sources

# Assumes each RPM only needs one download - we have some multi-source
# packages but in all cases the additional sources are patches provided
# in the Git repository
def download_rpm_sources(spec):
    for (url, path) in zip(spec.source_urls(), spec.source_paths()):
        source = urlparse.urlparse(url)

        # Source comes from a remote HTTP server
        if source.scheme in ["http", "https"]:
            print '%s: %s' % (path, spec.specpath())
            print '\t@echo [CURL] $@'
            print '\t@curl --silent --show-error -L -o $@ %s' % url

        # Source comes from a local file or directory
        if source.scheme == "file":
            print '%s: %s $(shell find %s)' % (
                path, spec.specpath(), source.path)

            # Assume that the directory name is already what's expected by the
            # spec file, and prefix it with the version number in the tarball
            print '\t@echo [GIT] $@'
            dirname = "%s-%s" % (os.path.basename(source.path), spec.version())
            print '\t@git --git-dir=%s/.git '\
                'archive --prefix %s/ -o $@ HEAD' % (source.path, dirname)

# Rules to build RPMS from SRPMS (uses information from the SPECs to
# get packages)
def build_rpm_from_srpm(spec):
    # This doesn't generate the right Makefile fragment for a multi-target
    # rule - we may end up building too often, or not rebuilding correctly
    # on a partial build
    rpm_paths = spec.binary_package_paths()
    srpm_path = spec.source_package_path()
    for rpm_path in rpm_paths:
        print '%s: %s' % (rpm_path, srpm_path)

def package_to_rpm_map(specs):
    provides_to_rpm = {}
    for spec in specs:
        for provided in spec.provides():
            for rpmpath in spec.binary_package_paths():
                provides_to_rpm[provided] = rpmpath
    return provides_to_rpm

def buildrequires_for_rpm(spec, provides_to_rpm):
    for rpmpath in spec.binary_package_paths():
        for buildreq in spec.buildrequires():
            # Some buildrequires come from the system repository
            if provides_to_rpm.has_key(buildreq):
                buildreqrpm = provides_to_rpm[buildreq]
                print "%s: %s" % (rpmpath, buildreqrpm)

def parse_cmdline():
    """
    Parse command line options
    """
    parser = argparse.ArgumentParser(description=
        "Generate Makefile dependencies from RPM Spec files")
    parser.add_argument("specs", metavar="SPEC", nargs="+", help="spec file")
    parser.add_argument("-i", "--ignore", metavar="PKG", action="append",
                        default=[], help="package name to ignore")
    return parser.parse_args()

def main():
    args = parse_cmdline()
    spec_paths = args.specs
    ignore_list = args.ignore
    specs = {}

    for spec_path in spec_paths:
	if build_type() == "deb":
            spec = pkg.Spec(spec_path, target="deb",
                            map_name=map_package_name_deb)
	else:
            spec = pkg.Spec(spec_path, target="rpm")
        pkg_name = spec.name()
        if pkg_name in ignore_list:
            continue
        if os.path.splitext(os.path.basename(spec_path))[0] != pkg_name:
            sys.stderr.write(
                "error: spec file name '%s' does not match package name '%s'\n" %
                (spec_path, pkg_name))
            sys.exit(1)

        specs[os.path.basename(spec_path)] = spec

    provides_to_rpm = package_to_rpm_map(specs.values())

    for spec in specs.itervalues():
        build_srpm_from_spec(spec)
        download_rpm_sources(spec)
        build_rpm_from_srpm(spec)
        buildrequires_for_rpm(spec, provides_to_rpm)
        print ""

    # Generate targets to build all srpms and all rpms
    all_rpms = []
    all_srpms = []
    for spec in specs.itervalues():
        rpm_path = spec.binary_package_paths()[0]
        all_rpms.append(rpm_path)
        all_srpms.append(spec.source_package_path())
        print "%s: %s" % (spec.name(), rpm_path)
    print ""

    print "rpms: " + " \\\n\t".join(all_rpms)
    print ""
    print "srpms: " + " \\\n\t".join(all_srpms)
    print ""
    print "install: all"
    print "\t. scripts/%s/install.sh" % build_type()

if __name__ == "__main__":
    main()
<commit_msg>specdep: argparse is not available on CentOS 6.4<commit_after>#!/usr/bin/python

# see http://docs.fedoraproject.org/en-US/Fedora_Draft_Documentation/0.1/html/RPM_Guide/ch16s04.html

import sys
import os
import platform
import urlparse
import pkg
import getopt

from scripts.lib import mappkgname

def build_type():
    debian_like = ["ubuntu", "debian"]
    rhel_like = ["fedora", "redhat", "centos"]

    dist = platform.linux_distribution(full_distribution_name=False)[0].lower()
    assert dist in debian_like + rhel_like

    if dist in debian_like:
        return "deb"
    elif dist in rhel_like:
        return "rpm"

def map_package_name_deb(name):
    """Map RPM package name to equivalent Deb names"""
    return mappkgname.map_package(name)

# Rules to build SRPM from SPEC
def build_srpm_from_spec(spec):
    srpmpath = spec.source_package_path()
    print '%s: %s %s' % (srpmpath, spec.specpath(),
                         " ".join(spec.source_paths()))

# Rules to download sources

# Assumes each RPM only needs one download - we have some multi-source
# packages but in all cases the additional sources are patches provided
# in the Git repository
def download_rpm_sources(spec):
    for (url, path) in zip(spec.source_urls(), spec.source_paths()):
        source = urlparse.urlparse(url)

        # Source comes from a remote HTTP server
        if source.scheme in ["http", "https"]:
            print '%s: %s' % (path, spec.specpath())
            print '\t@echo [CURL] $@'
            print '\t@curl --silent --show-error -L -o $@ %s' % url

        # Source comes from a local file or directory
        if source.scheme == "file":
            print '%s: %s $(shell find %s)' % (
                path, spec.specpath(), source.path)

            # Assume that the directory name is already what's expected by the
            # spec file, and prefix it with the version number in the tarball
            print '\t@echo [GIT] $@'
            dirname = "%s-%s" % (os.path.basename(source.path), spec.version())
            print '\t@git --git-dir=%s/.git '\
                'archive --prefix %s/ -o $@ HEAD' % (source.path, dirname)

# Rules to build RPMS from SRPMS (uses information from the SPECs to
# get packages)
def build_rpm_from_srpm(spec):
    # This doesn't generate the right Makefile fragment for a multi-target
    # rule - we may end up building too often, or not rebuilding correctly
    # on a partial build
    rpm_paths = spec.binary_package_paths()
    srpm_path = spec.source_package_path()
    for rpm_path in rpm_paths:
        print '%s: %s' % (rpm_path, srpm_path)

def package_to_rpm_map(specs):
    provides_to_rpm = {}
    for spec in specs:
        for provided in spec.provides():
            for rpmpath in spec.binary_package_paths():
                provides_to_rpm[provided] = rpmpath
    return provides_to_rpm

def buildrequires_for_rpm(spec, provides_to_rpm):
    for rpmpath in spec.binary_package_paths():
        for buildreq in spec.buildrequires():
            # Some buildrequires come from the system repository
            if provides_to_rpm.has_key(buildreq):
                buildreqrpm = provides_to_rpm[buildreq]
                print "%s: %s" % (rpmpath, buildreqrpm)

def usage(name):
    """
    Print usage information
    """
    print "usage: %s [-h] [-i PKG] SPEC [SPEC ...]" % name

def parse_cmdline():
    """
    Parse command line options
    """

    try:
        opts, args = getopt.getopt(sys.argv[1:], "hi:", ["help", "ignore="])
    except getopt.GetoptError as err:
        usage(sys.argv[0])
        print str(err)
        sys.exit(1)

    ignore = []
    for opt, val in opts:
        if opt == "-i" or opt == "--ignore":
            ignore.append(val)
        else:
            usage(sys.argv[0])
            print "unknown option: %s" % opt
            sys.exit(1)

    if len(args) == 0:
        usage(sys.argv[0])
        print "%s: error: too few arguments" % sys.argv[0]
        sys.exit(1)

    return (ignore, args)

def main():
    ignore_list, spec_paths = parse_cmdline()
    specs = {}

    for spec_path in spec_paths:
	if build_type() == "deb":
            spec = pkg.Spec(spec_path, target="deb",
                            map_name=map_package_name_deb)
	else:
            spec = pkg.Spec(spec_path, target="rpm")
        pkg_name = spec.name()
        if pkg_name in ignore_list:
            continue
        if os.path.splitext(os.path.basename(spec_path))[0] != pkg_name:
            sys.stderr.write(
                "error: spec file name '%s' does not match package name '%s'\n" %
                (spec_path, pkg_name))
            sys.exit(1)

        specs[os.path.basename(spec_path)] = spec

    provides_to_rpm = package_to_rpm_map(specs.values())

    for spec in specs.itervalues():
        build_srpm_from_spec(spec)
        download_rpm_sources(spec)
        build_rpm_from_srpm(spec)
        buildrequires_for_rpm(spec, provides_to_rpm)
        print ""

    # Generate targets to build all srpms and all rpms
    all_rpms = []
    all_srpms = []
    for spec in specs.itervalues():
        rpm_path = spec.binary_package_paths()[0]
        all_rpms.append(rpm_path)
        all_srpms.append(spec.source_package_path())
        print "%s: %s" % (spec.name(), rpm_path)
    print ""

    print "rpms: " + " \\\n\t".join(all_rpms)
    print ""
    print "srpms: " + " \\\n\t".join(all_srpms)
    print ""
    print "install: all"
    print "\t. scripts/%s/install.sh" % build_type()

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before># Copyright 2016 Red Hat, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Error heirarchy for stratis cli.
"""

from ._stratisd_constants import (
    BLOCK_DEV_TIER_TO_NAME,
    STRATISD_ERROR_TO_NAME,
    BlockDevTiers,
)

class StratisCliError(Exception):
    """
    Top-level stratis cli error.
    """

class StratisCliRuntimeError(StratisCliError):
    """
    Exception raised while an action is being performed and as a result of
    the requested action.
    """

class StratisCliUserError(StratisCliRuntimeError):
    """
    Exception raised as a result of a user error.
    """

class StratisCliNoDeviceSizeChangeError(StratisCliUserError):
    """
    Exception raised when the user requests to extend a pool to include new
    space on data devices, but there isn't any.
    """

    def __init__(self):  # pylint: disable=super-init-not-called
        """
        Initializer.
        """

    def __str__(self):
        return (
            "None of the pool's devices seem to have increased in size since "
            "the pool was created; the pool's allocable space would not be "
            "changed."
        )

class StratisCliNoPropertyChangeError(StratisCliUserError):
    """
    Raised when the user requests that a property be changed to its existing
    value.
    """

    def __init__(self, value):
        """
        Initializer.

        :param bool value: overprovision value
        """
        # pylint: disable=super-init-not-called
        self.value = value

class StratisCliOverprovisionChangeError(StratisCliNoPropertyChangeError):
    """
    Raised when the user requests the same overprovision state that the pool
    already has.
    """

    def __str__(self):
        return f"Pool's overprovision mode is already set to {str(self.value).lower()}"

class StratisCliFsLimitChangeError(StratisCliNoPropertyChangeError):
    """
    Raised when the user requests the same FsLimti value that the pool already
    has.
    """

    def __str__(self):
        return f"Pool's filesystem limit is exactly {str(self.value).lower()}"

class StratisCliResourceNotFoundError(StratisCliUserError):
    """
    Raised if a request was made to stratisd to update a resource
    that was expected to exist but was not found.
    """

    def __init__(self, command, resource):
        """
        Initializer.

        :param str command: the executed command
        :param str resource: the target resource
        """
        # pylint: disable=super-init-not-called
        self.command = command
        self.resource = resource

    def __str__(self):
        return (
            f"Command '{self.command}' assumes that resource '{self.resource}' "
            f"exists but it was not found."
        )

class StratisCliPartialChangeError(StratisCliUserError):
    """
    Raised if a request made of stratisd must result in a partial or no change
    since some or all of the post-condition for the request already holds.

    Invariant: self.unchanged_resources != frozenset()
    """

    def __init__(self, command, changed_resources, unchanged_resources):
        """
        Initializer.

        :param str command: the command run that caused the error
        :param changed_resources: the target resources that would change
        :type changed_resources: frozenset of str
        :param unchanged_resources: the target resources that would not change
        :type unchanged_resources: frozenset of str

        Precondition: unchanged_resources != frozenset()
        """
        # pylint: disable=super-init-not-called
        self.command = command
        self.changed_resources = changed_resources
        self.unchanged_resources = unchanged_resources

    def __str__(self):
        if len(self.unchanged_resources) > 1:
            msg = (
                f"The '{self.command}' action has no effect for resources "
                f"{list(self.unchanged_resources)}"
            )
        else:
            msg = (
                f"The '{self.command}' action has no effect for resource "
                f"{list(self.unchanged_resources)[0]}"
            )

        if self.changed_resources != frozenset():
            if len(self.changed_resources) > 1:
                msg += f" but does for resources {list(self.changed_resources)}"
            else:
                msg += f" but does for resource {list(self.changed_resources)[0]}"

        return msg

class StratisCliNoChangeError(StratisCliPartialChangeError):
    """
    Raised if a request made of stratisd must result in no change since the
    post-condition of the request is already entirely satisfied. This is
    a special case of StratisCliPartialChangeError, where the change requested
    is so simple that it can only succeed or fail.
    """

    def __init__(self, command, resource):
        """
        Initializer.

        :param str command: the executed command
        :param str resource: the target resource
        """
        StratisCliPartialChangeError.__init__(
            self, command, frozenset(), frozenset([resource])
        )

class StratisCliNameConflictError(StratisCliUserError):
    """
    Raised if an item of the same name already exists.
    """

    def __init__(self, object_type, name):
        """
        Initializer.

        :param str object_type: the type of the object, pool, filesystem, etc.
        :param str name: the conflicting name
        """
        # pylint: disable=super-init-not-called
        self.object_type = object_type
        self.name = name

    def __str__(self):
        return f"A {self.object_type} named {self.name} already exists"

class StratisCliParserError(StratisCliRuntimeError):
    """
    Raised when user has passed arguments that the parser should reject
    but can not due to its configuration limitations.
    """

class StratisCliMissingClevisTangURLError(StratisCliParserError):
    """
    Raised during pool creation if --clevis=[tang|nbde] specified but no
    URL supplied.
    """

    def __str__(self):
        return (
            "Specified binding with Clevis Tang server, but URL was not "
            "specified. Use --tang-url option to specify tang URL."
        )

class StratisCliMissingClevisThumbprintError(StratisCliParserError):
    """
    Raised during pool creation if --clevis=[tang|nbde] specified but neither
    a thumbprint nor instructions to trust the URL were supplied.
    """

    def __str__(self):
        return (
            "Specified binding with Clevis Tang server, but neither "
            "--thumbprint nor --trust-url option was specified."
        )

class StratisCliIncoherenceError(StratisCliRuntimeError):
    """
    Raised if there was a disagreement about state between the CLI
    and the daemon. This can happen if the information derived from
    the GetManagedObjects() result does not correspond with the daemon's
    state. This is a very unlikely event, but could occur if two processes
    were issuing instructions to stratisd at the same time.
    """

class StratisCliSynthUeventError(StratisCliRuntimeError):
    """
    Raised if there was an error during a synthetic udev event generation.
    """

class StratisCliInUseError(StratisCliUserError):
    """
    Base class for if a request made of stratisd must result in a device being
    included in both data and cache tiers or in the same tier of two different
    pools.
    """

class StratisCliInUseOtherTierError(StratisCliInUseError):
    """
    Raised if a request made of stratisd must result in a device being
    included in both data and cache tiers.
    """

    def __init__(self, pools_to_blockdevs, added_as):
        """
        Initializer.

        :param pools_to_blockdevs: pools mapped to the blockdevs they own
        :type pools_to_blockdevs: dict of str * frozenset of str
        :param added_as: what tier the devices were to be added to
        :type added_as: _stratisd_constants.BlockDevTiers

        Precondition: pools_to_blockdevs != {}
        Precondition: all frozenset of str in pools_to_blockdevs have at least one item
        """
        # pylint: disable=super-init-not-called
        self.pools_to_blockdevs = pools_to_blockdevs
        self.added_as = added_as

    def __str__(self):
        (target_blockdev_tier, already_blockdev_tier) = (
            BLOCK_DEV_TIER_TO_NAME(self.added_as),
            BLOCK_DEV_TIER_TO_NAME(
                BlockDevTiers.DATA
                if self.added_as == BlockDevTiers.CACHE
                else BlockDevTiers.CACHE
            ),
        )

        msg = (
            f"At least one of the provided devices is already owned by an "
            f"existing pool's {already_blockdev_tier} tier"
        )

        for pool_name, blockdevs in self.pools_to_blockdevs.items():
            if len(blockdevs) > 1:
                msg += (
                    f"; devices {list(blockdevs)} would be added to the {target_blockdev_tier} "
                    f"tier but are already in use in the {already_blockdev_tier} tier "
                    f"of pool {pool_name}"
                )
            else:
                msg += (
                    "; device {list(blockdevs)[0]} would be added to the {target_blockdev_tier} "
                    f"tier but is already in use in the {already_blockdev_tier} tier "
                    f"of pool {pool_name}"
                )

        return msg

class StratisCliInUseSameTierError(StratisCliInUseError):
    """
    Raised if a request made of stratisd must result in a device being
    included in the same tier in two different pools.
    """

    def __init__(self, pools_to_blockdevs, added_as):
        """
        Initializer.

        :param pools_to_blockdevs: pools mapped to the blockdevs they own
        :type pools_to_blockdevs: dict of str * frozenset of str
        :param added_as: what tier the devices were to be added to
        :type added_as: _stratisd_constants.BlockDevTiers

        Precondition: pools_to_blockdevs != {}
        Precondition: all frozenset of str in pools_to_blockdevs have at least one item
        """
        # pylint: disable=super-init-not-called
        self.pools_to_blockdevs = pools_to_blockdevs
        self.added_as = added_as

    def __str__(self):
        blockdev_tier = BLOCK_DEV_TIER_TO_NAME(self.added_as)

        msg = (
            f"At least one of the provided devices is already owned by an "
            f"existing pool's {blockdev_tier} tier"
        )

        for pool_name, blockdevs in self.pools_to_blockdevs.items():
            if len(blockdevs) > 1:
                msg += (
                    f"; devices {list(blockdevs)} would be added to the {blockdev_tier} tier "
                    f"but are already in use in the {blockdev_tier} tier of pool {pool_name}"
                )
            else:
                msg += (
                    f"; device {list(blockdevs)[0]} would be added to the {blockdev_tier} tier "
                    f"but is already in use in the {blockdev_tier} tier of pool {pool_name}"
                )

        return msg

class StratisCliKeyfileNotFoundError(StratisCliUserError):
    """
    Raised if the user specified a keyfile which could not be found when
    setting or unsetting a key.
    """

    def __init__(self, keyfile_path):
        """
        Initializer.

        :param str keyfile_path: the unfound path
        """
        # pylint: disable=super-init-not-called
        self.keyfile_path = keyfile_path

    def __str__(self):
        return f'There is no keyfile at the path "{self.keyfile_path}"'

class StratisCliUnknownInterfaceError(StratisCliRuntimeError):
    """
    Error raised when code encounters an unexpected D-Bus interface name.
    """

    def __init__(self, interface_name):
        """
        Initializer.

        :param str interface_name: the unexpected interface name
        """
        # pylint: disable=super-init-not-called
        self._interface_name = interface_name

    def __str__(self):
        return f"unexpected interface name {self._interface_name}"

class StratisCliEngineError(StratisCliRuntimeError):
    """
    Raised if there was a failure due to an error in stratisd's engine.
    """

    def __init__(self, rc, message):
        """
        Initializer.

        :param rc int: the error code returned by the engine
        :param str message: whatever message accompanied the error code
        """
        # pylint: disable=super-init-not-called
        self.error_code = rc
        self.message = message

    def __str__(self):
        return f"{STRATISD_ERROR_TO_NAME(self.error_code)}: {self.message}"

class StratisCliActionError(StratisCliRuntimeError):
    """
    Raised if an action selected by the parser failed.
    """

    def __init__(self, command_line_args, namespace):
        """
        Initialize with parser-returned namespace.

        :param command_line_args: the arguments passed on the command line
        :type command_line_args: list of str
        :param Namespace namespace: the namespace constructed by the parser
        """
        # pylint: disable=super-init-not-called
        self.command_line_args = command_line_args
        self.namespace = namespace

    def __str__(self):
        fmt_str = (
            "Action selected by command-line arguments %s which were "
            "parsed to %s failed"
        )
        return fmt_str % (self.command_line_args, self.namespace)

class StratisCliGenerationError(StratisCliError):
    """
    Exception that occurs during generation of classes.
    """

class StratisCliEnvironmentError(StratisCliError):
    """
    Exception that occurs during processing of environment variables.
    """

class StratisCliStratisdVersionError(StratisCliRuntimeError):
    """
    Raised if stratisd version does not meet CLI version requirements.
    """

    def __init__(self, actual_version, minimum_version, maximum_version):
        """
        Initializer.
        :param tuple actual_version: stratisd's actual version
        :param tuple minimum_version: the minimum version required
        :param tuple maximum_version: the maximum version allowed
        """
        # pylint: disable=super-init-not-called
        self.actual_version = actual_version
        self.minimum_version = minimum_version
        self.maximum_version = maximum_version

    def __str__(self):
        fmt_str = (
            "stratisd version %s does not meet stratis version "
            "requirements; the version must be at least %s and no more "
            "than %s"
        )
        return fmt_str % (
            self.actual_version,
            self.minimum_version,
            self.maximum_version,
        )
<commit_msg>Improve error message<commit_after># Copyright 2016 Red Hat, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Error heirarchy for stratis cli.
"""

from ._stratisd_constants import (
    BLOCK_DEV_TIER_TO_NAME,
    STRATISD_ERROR_TO_NAME,
    BlockDevTiers,
)

class StratisCliError(Exception):
    """
    Top-level stratis cli error.
    """

class StratisCliRuntimeError(StratisCliError):
    """
    Exception raised while an action is being performed and as a result of
    the requested action.
    """

class StratisCliUserError(StratisCliRuntimeError):
    """
    Exception raised as a result of a user error.
    """

class StratisCliNoDeviceSizeChangeError(StratisCliUserError):
    """
    Exception raised when the user requests to extend a pool to include new
    space on data devices, but there isn't any.
    """

    def __init__(self):  # pylint: disable=super-init-not-called
        """
        Initializer.
        """

    def __str__(self):
        return (
            "None of the pool's devices seems to have increased in size; the "
            "pool's allocatable space would not be changed."
        )

class StratisCliNoPropertyChangeError(StratisCliUserError):
    """
    Raised when the user requests that a property be changed to its existing
    value.
    """

    def __init__(self, value):
        """
        Initializer.

        :param bool value: overprovision value
        """
        # pylint: disable=super-init-not-called
        self.value = value

class StratisCliOverprovisionChangeError(StratisCliNoPropertyChangeError):
    """
    Raised when the user requests the same overprovision state that the pool
    already has.
    """

    def __str__(self):
        return f"Pool's overprovision mode is already set to {str(self.value).lower()}"

class StratisCliFsLimitChangeError(StratisCliNoPropertyChangeError):
    """
    Raised when the user requests the same FsLimti value that the pool already
    has.
    """

    def __str__(self):
        return f"Pool's filesystem limit is exactly {str(self.value).lower()}"

class StratisCliResourceNotFoundError(StratisCliUserError):
    """
    Raised if a request was made to stratisd to update a resource
    that was expected to exist but was not found.
    """

    def __init__(self, command, resource):
        """
        Initializer.

        :param str command: the executed command
        :param str resource: the target resource
        """
        # pylint: disable=super-init-not-called
        self.command = command
        self.resource = resource

    def __str__(self):
        return (
            f"Command '{self.command}' assumes that resource '{self.resource}' "
            f"exists but it was not found."
        )

class StratisCliPartialChangeError(StratisCliUserError):
    """
    Raised if a request made of stratisd must result in a partial or no change
    since some or all of the post-condition for the request already holds.

    Invariant: self.unchanged_resources != frozenset()
    """

    def __init__(self, command, changed_resources, unchanged_resources):
        """
        Initializer.

        :param str command: the command run that caused the error
        :param changed_resources: the target resources that would change
        :type changed_resources: frozenset of str
        :param unchanged_resources: the target resources that would not change
        :type unchanged_resources: frozenset of str

        Precondition: unchanged_resources != frozenset()
        """
        # pylint: disable=super-init-not-called
        self.command = command
        self.changed_resources = changed_resources
        self.unchanged_resources = unchanged_resources

    def __str__(self):
        if len(self.unchanged_resources) > 1:
            msg = (
                f"The '{self.command}' action has no effect for resources "
                f"{list(self.unchanged_resources)}"
            )
        else:
            msg = (
                f"The '{self.command}' action has no effect for resource "
                f"{list(self.unchanged_resources)[0]}"
            )

        if self.changed_resources != frozenset():
            if len(self.changed_resources) > 1:
                msg += f" but does for resources {list(self.changed_resources)}"
            else:
                msg += f" but does for resource {list(self.changed_resources)[0]}"

        return msg

class StratisCliNoChangeError(StratisCliPartialChangeError):
    """
    Raised if a request made of stratisd must result in no change since the
    post-condition of the request is already entirely satisfied. This is
    a special case of StratisCliPartialChangeError, where the change requested
    is so simple that it can only succeed or fail.
    """

    def __init__(self, command, resource):
        """
        Initializer.

        :param str command: the executed command
        :param str resource: the target resource
        """
        StratisCliPartialChangeError.__init__(
            self, command, frozenset(), frozenset([resource])
        )

class StratisCliNameConflictError(StratisCliUserError):
    """
    Raised if an item of the same name already exists.
    """

    def __init__(self, object_type, name):
        """
        Initializer.

        :param str object_type: the type of the object, pool, filesystem, etc.
        :param str name: the conflicting name
        """
        # pylint: disable=super-init-not-called
        self.object_type = object_type
        self.name = name

    def __str__(self):
        return f"A {self.object_type} named {self.name} already exists"

class StratisCliParserError(StratisCliRuntimeError):
    """
    Raised when user has passed arguments that the parser should reject
    but can not due to its configuration limitations.
    """

class StratisCliMissingClevisTangURLError(StratisCliParserError):
    """
    Raised during pool creation if --clevis=[tang|nbde] specified but no
    URL supplied.
    """

    def __str__(self):
        return (
            "Specified binding with Clevis Tang server, but URL was not "
            "specified. Use --tang-url option to specify tang URL."
        )

class StratisCliMissingClevisThumbprintError(StratisCliParserError):
    """
    Raised during pool creation if --clevis=[tang|nbde] specified but neither
    a thumbprint nor instructions to trust the URL were supplied.
    """

    def __str__(self):
        return (
            "Specified binding with Clevis Tang server, but neither "
            "--thumbprint nor --trust-url option was specified."
        )

class StratisCliIncoherenceError(StratisCliRuntimeError):
    """
    Raised if there was a disagreement about state between the CLI
    and the daemon. This can happen if the information derived from
    the GetManagedObjects() result does not correspond with the daemon's
    state. This is a very unlikely event, but could occur if two processes
    were issuing instructions to stratisd at the same time.
    """

class StratisCliSynthUeventError(StratisCliRuntimeError):
    """
    Raised if there was an error during a synthetic udev event generation.
    """

class StratisCliInUseError(StratisCliUserError):
    """
    Base class for if a request made of stratisd must result in a device being
    included in both data and cache tiers or in the same tier of two different
    pools.
    """

class StratisCliInUseOtherTierError(StratisCliInUseError):
    """
    Raised if a request made of stratisd must result in a device being
    included in both data and cache tiers.
    """

    def __init__(self, pools_to_blockdevs, added_as):
        """
        Initializer.

        :param pools_to_blockdevs: pools mapped to the blockdevs they own
        :type pools_to_blockdevs: dict of str * frozenset of str
        :param added_as: what tier the devices were to be added to
        :type added_as: _stratisd_constants.BlockDevTiers

        Precondition: pools_to_blockdevs != {}
        Precondition: all frozenset of str in pools_to_blockdevs have at least one item
        """
        # pylint: disable=super-init-not-called
        self.pools_to_blockdevs = pools_to_blockdevs
        self.added_as = added_as

    def __str__(self):
        (target_blockdev_tier, already_blockdev_tier) = (
            BLOCK_DEV_TIER_TO_NAME(self.added_as),
            BLOCK_DEV_TIER_TO_NAME(
                BlockDevTiers.DATA
                if self.added_as == BlockDevTiers.CACHE
                else BlockDevTiers.CACHE
            ),
        )

        msg = (
            f"At least one of the provided devices is already owned by an "
            f"existing pool's {already_blockdev_tier} tier"
        )

        for pool_name, blockdevs in self.pools_to_blockdevs.items():
            if len(blockdevs) > 1:
                msg += (
                    f"; devices {list(blockdevs)} would be added to the {target_blockdev_tier} "
                    f"tier but are already in use in the {already_blockdev_tier} tier "
                    f"of pool {pool_name}"
                )
            else:
                msg += (
                    "; device {list(blockdevs)[0]} would be added to the {target_blockdev_tier} "
                    f"tier but is already in use in the {already_blockdev_tier} tier "
                    f"of pool {pool_name}"
                )

        return msg

class StratisCliInUseSameTierError(StratisCliInUseError):
    """
    Raised if a request made of stratisd must result in a device being
    included in the same tier in two different pools.
    """

    def __init__(self, pools_to_blockdevs, added_as):
        """
        Initializer.

        :param pools_to_blockdevs: pools mapped to the blockdevs they own
        :type pools_to_blockdevs: dict of str * frozenset of str
        :param added_as: what tier the devices were to be added to
        :type added_as: _stratisd_constants.BlockDevTiers

        Precondition: pools_to_blockdevs != {}
        Precondition: all frozenset of str in pools_to_blockdevs have at least one item
        """
        # pylint: disable=super-init-not-called
        self.pools_to_blockdevs = pools_to_blockdevs
        self.added_as = added_as

    def __str__(self):
        blockdev_tier = BLOCK_DEV_TIER_TO_NAME(self.added_as)

        msg = (
            f"At least one of the provided devices is already owned by an "
            f"existing pool's {blockdev_tier} tier"
        )

        for pool_name, blockdevs in self.pools_to_blockdevs.items():
            if len(blockdevs) > 1:
                msg += (
                    f"; devices {list(blockdevs)} would be added to the {blockdev_tier} tier "
                    f"but are already in use in the {blockdev_tier} tier of pool {pool_name}"
                )
            else:
                msg += (
                    f"; device {list(blockdevs)[0]} would be added to the {blockdev_tier} tier "
                    f"but is already in use in the {blockdev_tier} tier of pool {pool_name}"
                )

        return msg

class StratisCliKeyfileNotFoundError(StratisCliUserError):
    """
    Raised if the user specified a keyfile which could not be found when
    setting or unsetting a key.
    """

    def __init__(self, keyfile_path):
        """
        Initializer.

        :param str keyfile_path: the unfound path
        """
        # pylint: disable=super-init-not-called
        self.keyfile_path = keyfile_path

    def __str__(self):
        return f'There is no keyfile at the path "{self.keyfile_path}"'

class StratisCliUnknownInterfaceError(StratisCliRuntimeError):
    """
    Error raised when code encounters an unexpected D-Bus interface name.
    """

    def __init__(self, interface_name):
        """
        Initializer.

        :param str interface_name: the unexpected interface name
        """
        # pylint: disable=super-init-not-called
        self._interface_name = interface_name

    def __str__(self):
        return f"unexpected interface name {self._interface_name}"

class StratisCliEngineError(StratisCliRuntimeError):
    """
    Raised if there was a failure due to an error in stratisd's engine.
    """

    def __init__(self, rc, message):
        """
        Initializer.

        :param rc int: the error code returned by the engine
        :param str message: whatever message accompanied the error code
        """
        # pylint: disable=super-init-not-called
        self.error_code = rc
        self.message = message

    def __str__(self):
        return f"{STRATISD_ERROR_TO_NAME(self.error_code)}: {self.message}"

class StratisCliActionError(StratisCliRuntimeError):
    """
    Raised if an action selected by the parser failed.
    """

    def __init__(self, command_line_args, namespace):
        """
        Initialize with parser-returned namespace.

        :param command_line_args: the arguments passed on the command line
        :type command_line_args: list of str
        :param Namespace namespace: the namespace constructed by the parser
        """
        # pylint: disable=super-init-not-called
        self.command_line_args = command_line_args
        self.namespace = namespace

    def __str__(self):
        fmt_str = (
            "Action selected by command-line arguments %s which were "
            "parsed to %s failed"
        )
        return fmt_str % (self.command_line_args, self.namespace)

class StratisCliGenerationError(StratisCliError):
    """
    Exception that occurs during generation of classes.
    """

class StratisCliEnvironmentError(StratisCliError):
    """
    Exception that occurs during processing of environment variables.
    """

class StratisCliStratisdVersionError(StratisCliRuntimeError):
    """
    Raised if stratisd version does not meet CLI version requirements.
    """

    def __init__(self, actual_version, minimum_version, maximum_version):
        """
        Initializer.
        :param tuple actual_version: stratisd's actual version
        :param tuple minimum_version: the minimum version required
        :param tuple maximum_version: the maximum version allowed
        """
        # pylint: disable=super-init-not-called
        self.actual_version = actual_version
        self.minimum_version = minimum_version
        self.maximum_version = maximum_version

    def __str__(self):
        fmt_str = (
            "stratisd version %s does not meet stratis version "
            "requirements; the version must be at least %s and no more "
            "than %s"
        )
        return fmt_str % (
            self.actual_version,
            self.minimum_version,
            self.maximum_version,
        )
<|endoftext|>


<commit_before>"""Models.

Changes affecting results or their presentation should also update
constants.py `change_date`,
"""

from __future__ import annotations

from datetime import date, datetime, timedelta
from logging import INFO, basicConfig, getLogger
from sys import stdout
from typing import Dict, Generator, Tuple, Sequence, Optional

import numpy as np
import pandas as pd

from .constants import EPSILON, CHANGE_DATE
from .parameters import Parameters

basicConfig(
    level=INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=stdout,
)
logger = getLogger(__name__)

class SimSirModel:

    def __init__(self, p: Parameters):

        self.rates = {
            key: d.rate
            for key, d in p.dispositions.items()
        }

        self.days = {
            key: d.days
            for key, d in p.dispositions.items()
        }

        self.keys = ("susceptible", "infected", "recovered")

        # An estimate of the number of infected people on the day that
        # the first hospitalized case is seen
        #
        # Note: this should not be an integer.
        infected = (
            1.0 / p.market_share / p.hospitalized.rate
        )

        susceptible = p.population - infected

        gamma = 1.0 / p.infectious_days
        self.gamma = gamma

        self.susceptible = susceptible
        self.infected = infected
        self.recovered = p.recovered

        if p.date_first_hospitalized is None and p.doubling_time is not None:
            # Back-projecting to when the first hospitalized case would have been admitted
            logger.info('Using doubling_time: %s', p.doubling_time)

            intrinsic_growth_rate = get_growth_rate(p.doubling_time)

            self.beta = get_beta(intrinsic_growth_rate,  gamma, self.susceptible, 0.0)
            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)

            self.i_day = 0 # seed to the full length
            self.run_projection(p, [(self.beta, p.n_days)])
            self.i_day = i_day = int(get_argmin_ds(self.census_df, p.current_hospitalized))

            self.run_projection(p, self.gen_policy(p))

            logger.info('Set i_day = %s', i_day)
            p.date_first_hospitalized = p.current_date - timedelta(days=i_day)
            logger.info(
                'Estimated date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)

        elif p.date_first_hospitalized is not None and p.doubling_time is None:
            # Fitting spread parameter to observed hospital census (dates of 1 patient and today)
            self.i_day = (p.current_date - p.date_first_hospitalized).days
            logger.info(
                'Using date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)
            min_loss = 2.0**99
            dts = np.linspace(1, 15, 29)
            losses = np.zeros(dts.shape[0])
            self.current_hospitalized = p.current_hospitalized
            for i, i_dt in enumerate(dts):
                intrinsic_growth_rate = get_growth_rate(i_dt)
                self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
                self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)

                self.run_projection(p, self.gen_policy(p))
                loss = self.get_loss()
                losses[i] = loss

            p.doubling_time = dts[pd.Series(losses).argmin()]
            logger.info('Estimated doubling_time: %s', p.doubling_time)
            intrinsic_growth_rate = get_growth_rate(p.doubling_time)
            self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
            self.run_projection(p, self.gen_policy(p))

            self.population = p.population
        else:
            logger.info(
                'doubling_time: %s; date_first_hospitalized: %s',
                p.doubling_time,
                p.date_first_hospitalized,
            )
            raise AssertionError('doubling_time or date_first_hospitalized must be provided.')

        logger.info('len(np.arange(-i_day, n_days+1)): %s', len(np.arange(-self.i_day, p.n_days+1)))
        logger.info('len(raw_df): %s', len(self.raw_df))

        self.infected = self.raw_df['infected'].values[self.i_day]
        self.susceptible = self.raw_df['susceptible'].values[self.i_day]
        self.recovered = self.raw_df['recovered'].values[self.i_day]

        self.intrinsic_growth_rate = intrinsic_growth_rate

        # r_t is r_0 after distancing
        self.r_t = self.beta_t / gamma * susceptible
        self.r_naught = self.beta / gamma * susceptible

        doubling_time_t = 1.0 / np.log2(
            self.beta_t * susceptible - gamma + 1)
        self.doubling_time_t = doubling_time_t

        self.sim_sir_w_date_df = build_sim_sir_w_date_df(self.raw_df, p.current_date, self.keys)

        self.sim_sir_w_date_floor_df = build_floor_df(self.sim_sir_w_date_df, self.keys)
        self.admits_floor_df = build_floor_df(self.admits_df, p.dispositions.keys())
        self.census_floor_df = build_floor_df(self.census_df, p.dispositions.keys())

        self.daily_growth_rate = get_growth_rate(p.doubling_time)
        self.daily_growth_rate_t = get_growth_rate(self.doubling_time_t)

    def gen_policy(self, p: Parameters) -> Sequence[Tuple[float, int]]:
        if p.mitigation_date is not None:
            mitigation_day = -(p.current_date - p.mitigation_date).days
        else:
            mitigation_day = 0

        total_days = self.i_day + p.n_days

        if mitigation_day < -self.i_day:
            mitigation_day = -self.i_day

        pre_mitigation_days = self.i_day + mitigation_day
        post_mitigation_days = total_days - pre_mitigation_days

        return [
            (self.beta,   pre_mitigation_days),
            (self.beta_t, post_mitigation_days),
        ]

    def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
        self.raw_df = sim_sir_df(
            self.susceptible,
            self.infected,
            p.recovered,
            self.gamma,
            -self.i_day,
            policy
        )

        self.dispositions_df = build_dispositions_df(self.raw_df, self.rates, p.market_share, p.current_date)
        self.admits_df = build_admits_df(self.dispositions_df)
        self.census_df = build_census_df(self.admits_df, self.days)
        self.current_infected = self.raw_df.infected.loc[self.i_day]

    def get_loss(self) -> float:
        """Squared error: predicted vs. actual current hospitalized."""
        predicted = self.census_df.hospitalized.loc[self.i_day]
        return (self.current_hospitalized - predicted) ** 2.0

def get_argmin_ds(census_df: pd.DataFrame, current_hospitalized: float) -> float:
    # By design, this forbids choosing a day after the peak
    # If that's a problem, see #381
    peak_day = census_df.hospitalized.argmax()
    losses_df = (census_df.hospitalized[:peak_day] - current_hospitalized) ** 2.0
    return losses_df.argmin()

def get_beta(
    intrinsic_growth_rate: float,
    gamma: float,
    susceptible: float,
    relative_contact_rate: float
) -> float:
    return (
        (intrinsic_growth_rate + gamma)
        / susceptible
        * (1.0 - relative_contact_rate)
    )

def get_growth_rate(doubling_time: Optional[float]) -> float:
    """Calculates average daily growth rate from doubling time."""
    if doubling_time is None or doubling_time == 0.0:
        return 0.0
    return (2.0 ** (1.0 / doubling_time) - 1.0)

def sir(
    s: float, i: float, r: float, beta: float, gamma: float, n: float
) -> Tuple[float, float, float]:
    """The SIR model, one time step."""
    s_n = (-beta * s * i) + s
    i_n = (beta * s * i - gamma * i) + i
    r_n = gamma * i + r

    # TODO:
    #   Post check dfs for negative values and
    #   warn the user that their input data is bad.
    #   JL: I suspect that these adjustments covered bugs.

    #if s_n < 0.0:
    #    s_n = 0.0
    #if i_n < 0.0:
    #    i_n = 0.0
    #if r_n < 0.0:
    #    r_n = 0.0
    scale = n / (s_n + i_n + r_n)
    return s_n * scale, i_n * scale, r_n * scale

def gen_sir(
    s: float, i: float, r: float, gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
) -> Generator[Tuple[int, float, float, float], None, None]:
    """Simulate SIR model forward in time yielding tuples.
    Parameter order has changed to allow multiple (beta, n_days)
    to reflect multiple changing social distancing policies.
    """
    s, i, r = (float(v) for v in (s, i, r))
    n = s + i + r
    d = i_day
    for beta, n_days in policies:
        for _ in range(n_days):
            yield d, s, i, r
            s, i, r = sir(s, i, r, beta, gamma, n)
            d += 1
    yield d, s, i, r

def sim_sir_df(
    s: float, i: float, r: float,
    gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
) -> pd.DataFrame:
    """Simulate the SIR model forward in time."""
    return pd.DataFrame(
        data=gen_sir(s, i, r, gamma, i_day, policies),
        columns=("day", "susceptible", "infected", "recovered"),
    )

def build_sim_sir_w_date_df(
    raw_df: pd.DataFrame,
    current_date: datetime,
    keys: Sequence[str],
) -> pd.DataFrame:
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: raw_df[key]
            for key in keys
        }
    })

def build_floor_df(df, keys):
    """Build floor sim sir w date."""
    return pd.DataFrame({
        "day": df.day,
        "date": df.date,
        **{
            key: np.floor(df[key])
            for key in keys
        }
    })

def build_dispositions_df(
    raw_df: pd.DataFrame,
    rates: Dict[str, float],
    market_share: float,
    current_date: datetime,
) -> pd.DataFrame:
    """Build dispositions dataframe of patients adjusted by rate and market_share."""
    patients = raw_df.infected + raw_df.recovered
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: patients * rate * market_share
            for key, rate in rates.items()
        }
    })

def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
    """Build admits dataframe from dispositions."""
    admits_df = dispositions_df - dispositions_df.shift(1)
    admits_df.day = dispositions_df.day
    admits_df.date = dispositions_df.date
    return admits_df

def build_census_df(
    admits_df: pd.DataFrame,
    lengths_of_stay: Dict[str, int],
) -> pd.DataFrame:
    """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
    return pd.DataFrame({
        'day': admits_df.day,
        'date': admits_df.date,
        **{
            key: (
                admits_df[key].cumsum()
                - admits_df[key].cumsum().shift(los).fillna(0)
            )
            for key, los in lengths_of_stay.items()
        }
    })
<commit_msg>Don't allow fitting of doubling time to fit the forecast day against days past the admissions peak<commit_after>"""Models.

Changes affecting results or their presentation should also update
constants.py `change_date`,
"""

from __future__ import annotations

from datetime import date, datetime, timedelta
from logging import INFO, basicConfig, getLogger
from sys import stdout
from typing import Dict, Generator, Tuple, Sequence, Optional

import numpy as np
import pandas as pd

from .constants import EPSILON, CHANGE_DATE
from .parameters import Parameters

basicConfig(
    level=INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=stdout,
)
logger = getLogger(__name__)

class SimSirModel:

    def __init__(self, p: Parameters):

        self.rates = {
            key: d.rate
            for key, d in p.dispositions.items()
        }

        self.days = {
            key: d.days
            for key, d in p.dispositions.items()
        }

        self.keys = ("susceptible", "infected", "recovered")

        # An estimate of the number of infected people on the day that
        # the first hospitalized case is seen
        #
        # Note: this should not be an integer.
        infected = (
            1.0 / p.market_share / p.hospitalized.rate
        )

        susceptible = p.population - infected

        gamma = 1.0 / p.infectious_days
        self.gamma = gamma

        self.susceptible = susceptible
        self.infected = infected
        self.recovered = p.recovered

        if p.date_first_hospitalized is None and p.doubling_time is not None:
            # Back-projecting to when the first hospitalized case would have been admitted
            logger.info('Using doubling_time: %s', p.doubling_time)

            intrinsic_growth_rate = get_growth_rate(p.doubling_time)

            self.beta = get_beta(intrinsic_growth_rate,  gamma, self.susceptible, 0.0)
            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)

            self.i_day = 0 # seed to the full length
            self.run_projection(p, [(self.beta, p.n_days)])
            self.i_day = i_day = int(get_argmin_ds(self.census_df, p.current_hospitalized))

            self.run_projection(p, self.gen_policy(p))

            logger.info('Set i_day = %s', i_day)
            p.date_first_hospitalized = p.current_date - timedelta(days=i_day)
            logger.info(
                'Estimated date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)

        elif p.date_first_hospitalized is not None and p.doubling_time is None:
            # Fitting spread parameter to observed hospital census (dates of 1 patient and today)
            self.i_day = (p.current_date - p.date_first_hospitalized).days
            logger.info(
                'Using date_first_hospitalized: %s; current_date: %s; i_day: %s',
                p.date_first_hospitalized,
                p.current_date,
                self.i_day)
            min_loss = 2.0**99
            dts = np.linspace(1, 15, 29)
            losses = np.full(dts.shape[0], np.inf)
            self.current_hospitalized = p.current_hospitalized
            for i, i_dt in enumerate(dts):
                intrinsic_growth_rate = get_growth_rate(i_dt)
                self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
                self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)

                self.run_projection(p, self.gen_policy(p))

                # Skip values the would put the fit past peak
                peak_admits_day = self.admits_df.hospitalized.argmax()
                if peak_admits_day < 0:
                    continue

                loss = self.get_loss()
                losses[i] = loss

            p.doubling_time = dts[pd.Series(losses).argmin()]
            logger.info('Estimated doubling_time: %s', p.doubling_time)
            intrinsic_growth_rate = get_growth_rate(p.doubling_time)
            self.beta = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, 0.0)
            self.beta_t = get_beta(intrinsic_growth_rate, self.gamma, self.susceptible, p.relative_contact_rate)
            self.run_projection(p, self.gen_policy(p))

            self.population = p.population
        else:
            logger.info(
                'doubling_time: %s; date_first_hospitalized: %s',
                p.doubling_time,
                p.date_first_hospitalized,
            )
            raise AssertionError('doubling_time or date_first_hospitalized must be provided.')

        logger.info('len(np.arange(-i_day, n_days+1)): %s', len(np.arange(-self.i_day, p.n_days+1)))
        logger.info('len(raw_df): %s', len(self.raw_df))

        self.infected = self.raw_df['infected'].values[self.i_day]
        self.susceptible = self.raw_df['susceptible'].values[self.i_day]
        self.recovered = self.raw_df['recovered'].values[self.i_day]

        self.intrinsic_growth_rate = intrinsic_growth_rate

        # r_t is r_0 after distancing
        self.r_t = self.beta_t / gamma * susceptible
        self.r_naught = self.beta / gamma * susceptible

        doubling_time_t = 1.0 / np.log2(
            self.beta_t * susceptible - gamma + 1)
        self.doubling_time_t = doubling_time_t

        self.sim_sir_w_date_df = build_sim_sir_w_date_df(self.raw_df, p.current_date, self.keys)

        self.sim_sir_w_date_floor_df = build_floor_df(self.sim_sir_w_date_df, self.keys)
        self.admits_floor_df = build_floor_df(self.admits_df, p.dispositions.keys())
        self.census_floor_df = build_floor_df(self.census_df, p.dispositions.keys())

        self.daily_growth_rate = get_growth_rate(p.doubling_time)
        self.daily_growth_rate_t = get_growth_rate(self.doubling_time_t)

    def gen_policy(self, p: Parameters) -> Sequence[Tuple[float, int]]:
        if p.mitigation_date is not None:
            mitigation_day = -(p.current_date - p.mitigation_date).days
        else:
            mitigation_day = 0

        total_days = self.i_day + p.n_days

        if mitigation_day < -self.i_day:
            mitigation_day = -self.i_day

        pre_mitigation_days = self.i_day + mitigation_day
        post_mitigation_days = total_days - pre_mitigation_days

        return [
            (self.beta,   pre_mitigation_days),
            (self.beta_t, post_mitigation_days),
        ]

    def run_projection(self, p: Parameters, policy: Sequence[Tuple[float, int]]):
        self.raw_df = sim_sir_df(
            self.susceptible,
            self.infected,
            p.recovered,
            self.gamma,
            -self.i_day,
            policy
        )

        self.dispositions_df = build_dispositions_df(self.raw_df, self.rates, p.market_share, p.current_date)
        self.admits_df = build_admits_df(self.dispositions_df)
        self.census_df = build_census_df(self.admits_df, self.days)
        self.current_infected = self.raw_df.infected.loc[self.i_day]

    def get_loss(self) -> float:
        """Squared error: predicted vs. actual current hospitalized."""
        predicted = self.census_df.hospitalized.loc[self.i_day]
        return (self.current_hospitalized - predicted) ** 2.0

def get_argmin_ds(census_df: pd.DataFrame, current_hospitalized: float) -> float:
    # By design, this forbids choosing a day after the peak
    # If that's a problem, see #381
    peak_day = census_df.hospitalized.argmax()
    losses_df = (census_df.hospitalized[:peak_day] - current_hospitalized) ** 2.0
    return losses_df.argmin()

def get_beta(
    intrinsic_growth_rate: float,
    gamma: float,
    susceptible: float,
    relative_contact_rate: float
) -> float:
    return (
        (intrinsic_growth_rate + gamma)
        / susceptible
        * (1.0 - relative_contact_rate)
    )

def get_growth_rate(doubling_time: Optional[float]) -> float:
    """Calculates average daily growth rate from doubling time."""
    if doubling_time is None or doubling_time == 0.0:
        return 0.0
    return (2.0 ** (1.0 / doubling_time) - 1.0)

def sir(
    s: float, i: float, r: float, beta: float, gamma: float, n: float
) -> Tuple[float, float, float]:
    """The SIR model, one time step."""
    s_n = (-beta * s * i) + s
    i_n = (beta * s * i - gamma * i) + i
    r_n = gamma * i + r

    # TODO:
    #   Post check dfs for negative values and
    #   warn the user that their input data is bad.
    #   JL: I suspect that these adjustments covered bugs.

    #if s_n < 0.0:
    #    s_n = 0.0
    #if i_n < 0.0:
    #    i_n = 0.0
    #if r_n < 0.0:
    #    r_n = 0.0
    scale = n / (s_n + i_n + r_n)
    return s_n * scale, i_n * scale, r_n * scale

def gen_sir(
    s: float, i: float, r: float, gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
) -> Generator[Tuple[int, float, float, float], None, None]:
    """Simulate SIR model forward in time yielding tuples.
    Parameter order has changed to allow multiple (beta, n_days)
    to reflect multiple changing social distancing policies.
    """
    s, i, r = (float(v) for v in (s, i, r))
    n = s + i + r
    d = i_day
    for beta, n_days in policies:
        for _ in range(n_days):
            yield d, s, i, r
            s, i, r = sir(s, i, r, beta, gamma, n)
            d += 1
    yield d, s, i, r

def sim_sir_df(
    s: float, i: float, r: float,
    gamma: float, i_day: int, policies: Sequence[Tuple[float, int]]
) -> pd.DataFrame:
    """Simulate the SIR model forward in time."""
    return pd.DataFrame(
        data=gen_sir(s, i, r, gamma, i_day, policies),
        columns=("day", "susceptible", "infected", "recovered"),
    )

def build_sim_sir_w_date_df(
    raw_df: pd.DataFrame,
    current_date: datetime,
    keys: Sequence[str],
) -> pd.DataFrame:
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: raw_df[key]
            for key in keys
        }
    })

def build_floor_df(df, keys):
    """Build floor sim sir w date."""
    return pd.DataFrame({
        "day": df.day,
        "date": df.date,
        **{
            key: np.floor(df[key])
            for key in keys
        }
    })

def build_dispositions_df(
    raw_df: pd.DataFrame,
    rates: Dict[str, float],
    market_share: float,
    current_date: datetime,
) -> pd.DataFrame:
    """Build dispositions dataframe of patients adjusted by rate and market_share."""
    patients = raw_df.infected + raw_df.recovered
    day = raw_df.day
    return pd.DataFrame({
        "day": day,
        "date": day.astype('timedelta64[D]') + np.datetime64(current_date),
        **{
            key: patients * rate * market_share
            for key, rate in rates.items()
        }
    })

def build_admits_df(dispositions_df: pd.DataFrame) -> pd.DataFrame:
    """Build admits dataframe from dispositions."""
    admits_df = dispositions_df - dispositions_df.shift(1)
    admits_df.day = dispositions_df.day
    admits_df.date = dispositions_df.date
    return admits_df

def build_census_df(
    admits_df: pd.DataFrame,
    lengths_of_stay: Dict[str, int],
) -> pd.DataFrame:
    """Average Length of Stay for each disposition of COVID-19 case (total guesses)"""
    return pd.DataFrame({
        'day': admits_df.day,
        'date': admits_df.date,
        **{
            key: (
                admits_df[key].cumsum()
                - admits_df[key].cumsum().shift(los).fillna(0)
            )
            for key, los in lengths_of_stay.items()
        }
    })
<|endoftext|>


<commit_before># Copyright 2019 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Helper function to setup @npm_bazel_typescript dev dependencies.
"""

load("@bazel_gazelle//:deps.bzl", "go_repository")

def ts_setup_dev_workspace():
    """
    Setup the toolchain needed for local development, but not needed by users.

    These needs to be in a separate file from ts_setup_workspace() so as not
    to leak load statements.
    """

    ts_setup_workspace()

    go_repository(
        name = "com_github_kylelemons_godebug",
        commit = "9ff306d4fbead574800b66369df5b6144732d58e",  # v1.1.0
        importpath = "github.com/kylelemons/godebug",
    )

    go_repository(
        name = "com_github_mattn_go_isatty",
        commit = "1311e847b0cb909da63b5fecfb5370aa66236465",  # v0.0.8
        importpath = "github.com/mattn/go-isatty",
    )
<commit_msg>chore: update dependency com_github_mattn_go_isatty to v0.0.9<commit_after># Copyright 2019 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Helper function to setup @npm_bazel_typescript dev dependencies.
"""

load("@bazel_gazelle//:deps.bzl", "go_repository")

def ts_setup_dev_workspace():
    """
    Setup the toolchain needed for local development, but not needed by users.

    These needs to be in a separate file from ts_setup_workspace() so as not
    to leak load statements.
    """

    ts_setup_workspace()

    go_repository(
        name = "com_github_kylelemons_godebug",
        commit = "9ff306d4fbead574800b66369df5b6144732d58e",  # v1.1.0
        importpath = "github.com/kylelemons/godebug",
    )

    go_repository(
        name = "com_github_mattn_go_isatty",
        commit = "e1f7b56ace729e4a73a29a6b4fac6cd5fcda7ab3",  # v0.0.9
        importpath = "github.com/mattn/go-isatty",
    )
<|endoftext|>