MakeCanvases.py

Go to the documentation of this file.

"""
  MakeCanvases.py:
   Build canvases used in validation web pages.
  
   Original author: J. Wolcott <jwolcott@fnal.gov>
              Date: September 2016
"""

import argparse
import copy
import itertools
import os.path
import sys

import ROOT

from controllers.ControllerBase import ControllerBase
from models.Organizational import  PlotCollectionKey
from models.Organizational import ComparisonRegistry
from models.Organizational import Configuration
from models.Organizational import PlotID
from models.Organizational import PlotLibrary
from models.Organizational import PlotRegistry
from models.Organizational import PlotSet
from models.PlotInfo import CanvasSummary
from models.PlotInfo import ComparisonSummary
from models.PlotInfo import Exposure
from tools.ImageMaker import ImageMaker
from tools import PathTools

OUTPUT_TYPES = [
#     "svg",
#     "json",
    "root",
]

if "json" in OUTPUT_TYPES:
    try:
        ROOT.TBufferJSON
    except AttributeError:
        OUTPUT_TYPES.remove("json")

GOOD_COLORS = [
        ROOT.kRed,
        ROOT.kBlue,
        ROOT.kGreen+3,
        ROOT.kBlack,

        ROOT.kOrange+2,  # brown
        ROOT.kMagenta+1,
        ROOT.kSpring-1,
        ROOT.kGray+2,

        ROOT.kRed-6,
        ROOT.kCyan+1,
        ROOT.kYellow+2,
        ROOT.kTeal-7,

#        ROOT.kViolet-4, # looks too much like kMagenta+1
#        ROOT.kRed+3,    # looks sorta black
#        ROOT.kViolet+7, # looks too much like kBlue
#        ROOT.kAzure+7,  # looks too much like kViolet-4
#        ROOT.kGreen,    # too light
]

class NoDirectory(object):
    """ Syntactic sugar object to make the setting of a TObject's directory to None easier.
        If NO_DIR is a NoDirectory instance, write
          NO_DIR >> obj
        to get back the object with its Directory set to None.
     """
    def __rshift__(self, obj):
        if not hasattr(obj, "SetDirectory"):
            return obj
        
        obj.SetDirectory(None)
        return obj
NO_DIR = NoDirectory()

class Warner(object):
    """ Syntactic sugar object to make writing a single instance of a warning to stderr
        (no matter how many times it's called) easier.
    """
    def __init__(self):
        self._warned = set()
    
    def __call__(self, warning):
        if warning in self._warned:
            return
        
        self._warned.add(warning)
        print >> sys.stderr, "\n\n\033[31m\033[1mWarning:\033[0m", warning, "\n"

SINGLE_WARN = Warner()

class HistogramExtractor(object):
    def __init__(self, accept_filters=[], reject_filters=[]):
        self.accept_filters = accept_filters
        self.reject_filters = reject_filters
        
    def KeepHist(self, name):
        if self.accept_filters:
            return any(name.startswith(flt) for flt in self.accept_filters)
        
        return not any(name.startswith(flt) for flt in self.reject_filters)
       
    
    def ExtractHistsFromDirectory(self, tdirectory, base_path=""):
        for tkey in tdirectory.GetListOfKeys():
            obj = tkey.ReadObj()
            name = obj.GetName()
            
            if isinstance(obj, ROOT.TH1) and self.KeepHist(name):
                splitted = name.rsplit("/", 1)
                if len(splitted) > 1:
                    path, name = splitted
                    obj.SetName(name)
                else:
                    path = ""
                obj.exposure = None
                yield base_path + ("/" if base_path and path else "") + path, (NO_DIR >> obj)
                
            # these were CAFAna Spectra, once
            elif isinstance(obj, ROOT.TDirectory) and \
                    obj.Get("type") and \
                    obj.Get("type").GetName() == "Spectrum" and \
                    self.KeepHist(obj.GetName()):
                h = obj.Get("hist")
                args = {
                    "POT": obj.Get("pot"),
                    "livetime": obj.Get("livetime"),
                }
                for exp_type, hist in args.items():
                    args[exp_type] = None if hist.GetBinContent(1) == 0 else hist.GetBinContent(1)
                h.exposure = Exposure(**args)
                h.SetName(obj.GetName())
                
                # labels for new-style CAFs are stored separately
                for label_idx, axis_name in enumerate("XYZ"):
                    title = obj.Get("label%d" % label_idx)
                    if title:
                        getattr(h, "Get%saxis" % axis_name)().SetTitle(str(title))
                
                yield base_path, (NO_DIR >> h)
                
            elif isinstance(obj, ROOT.TDirectory):
                path = name if not base_path else "/".join((base_path, name))
                for path, h in self.ExtractHistsFromDirectory(obj, base_path=path):
                    yield path, (NO_DIR >> h)

class CanvasMaker(ControllerBase):
    class _ReferenceHistComp(object):
        """ Helper type to do comparisons to a reference histogram.
            Used in the display variations within CanvasMaker. """
            
        def __init__(self, comparison_function, ignore_ref=False):
            self.comp_fn = comparison_function
            self.ignore_ref = ignore_ref  # basically a hack so that area normalization works (where you don't want the ref).  better solutions welcome
        
        def __call__(self, reference, subject):
            if not self.ignore_ref:
                ref_val = self.comp_fn(reference)
            sub_val = self.comp_fn(subject)
            
            try:
#                 print self.comp_fn
#                 print reference, ref_val
#                 print subject, sub_val
                return ref_val/sub_val if not self.ignore_ref else sub_val
            except:
                return None
         
    # variations in the arguments to the ImageMaker
    DISPLAY_VARIATIONS = {
        "Log y": {
            "enabled":  { "log_y": True },
            "disabled": { "log_y": False } 
        },
        
        "Normalization": {
            # use the _ReferenceHistComp for 'raw' as well
            # so that a normalization factor is generated for each histogram
            # (which is what the draw routine expects)
            "raw":      { "norm_factors": _ReferenceHistComp(lambda h: None) },
            "area":     {
                "norm_factors": _ReferenceHistComp(lambda h: 100./h.GetSumOfWeights() if h.GetSumOfWeights() > 0 else 1, ignore_ref=True),
                "axis_labels": ";Percent of entries"
            },
            "exposure": {
                "norm_factors": _ReferenceHistComp(lambda h: h.exposure.POT or h.exposure.livetime if hasattr(h, "exposure") and h.exposure else None)
            },
            
        },
    }
    DISPLAY_VARIATIONS_NOJSROOT = ["Log y",]

    def __init__(self, config):
        super(CanvasMaker, self).__init__(config)
        
        self.n_images_written = 0

        self.histogram_table = None
        self.histogram_extractor = HistogramExtractor(accept_filters=self.config.plot_accept_filters,
                                                      reject_filters=self.config.plot_reject_filters)
        
    def LoadHistograms(self):
        """ Sift through all the files in the configuration and discover histograms. """
        
        # only do once
        if self.histogram_table is not None:
            return
        
        self.histogram_table = PlotLibrary()

        # todo: would like to ensure that only input files that are newer than any pre-existing registry are remade.
        # however, want everything to be in the comparisons.  hmm. 
        
        for dataset_name, file_list in self.config.histogram_topdirs.iteritems():
            # maybe the "file list" is just one file
            if isinstance(file_list, basestring):
                file_list = [file_list,]
            
            for file_name in file_list:
                assert file_name.lower().endswith(".root"), "Can't open non-ROOT file input '%s'" % file_name
                
                f = ROOT.TFile(file_name)
                assert f, "Can't open ROOT file '%s'" % file_name
                
                for path, histogram in self.histogram_extractor.ExtractHistsFromDirectory(f):
                    if isinstance(histogram, (ROOT.TH3,)):
                        print >> sys.stderr, "Don't know how to handle 3D histograms.  Skipping histogram:", histogram.GetName()
                        continue
                    self.histogram_table[PlotCollectionKey(PlotID(histogram.GetName()), dataset_name, path)] = histogram
    
    def MakeCanvases(self):
        """ Master method for canvas making. """
        
        self.LoadHistograms()
        
        self.im = ImageMaker()

        self.plot_registry = PlotRegistry()
        self.comparison_registry = ComparisonRegistry()
        
        # supplemental collections that will be used for aggregations
        overlays =   PlotLibrary()
        sums     =   PlotLibrary()
        
        by_name_then_dataset = PlotLibrary()

        # make canvases for each histogram (or set of histograms, if there are categories)
        for key, histogram in self.histogram_table.iteritems():
            if histogram.GetTitle() == "":
                histogram.SetTitle(key.plot_id.name)
                
            # key with name and path but no category
            no_category_key = PlotCollectionKey(
                plot_id={
                    "name": key.plot_id.name,
                    "categories": [],
                },
                data_set=None,
                path=key.path
            )
            by_name_then_dataset.setdefault(no_category_key, {}).setdefault(key.data_set, PlotSet()).add(key)

            # couple extra things to do for category breakdowns...
            # (note that breakouts & sums are only made for the last category
            #  to keep the combinatorics under control)
            if len(key.plot_id.categories) > 0:
                # slice off the last 'category group' to make a collection of all the plots in the same category.
                # use the magic string in PlotID to identify it
                categories = key.plot_id.categories[:-1]
                categories.append(PlotID.OVERLAY_STRING)
                plotid_args = {
                    "name": key.plot_id.name,
                    "categories": categories,
                    "groups": key.plot_id.groups
                }
                plotid_args["categories"] = categories
                overlay_key = PlotCollectionKey( PlotID(plotid_args), key.data_set, key.path )
                overlays.setdefault(overlay_key, {})[key.plot_id.categories[-1]] = histogram

                # make a 'sum of all categories version'
                categories[-1] = PlotID.SUM_STRING
                plotid_args["categories"] = categories
                sum_key = PlotCollectionKey( PlotID(plotid_args), key.data_set, key.path )
                if sum_key not in sums:
                    sum_histogram = histogram.Clone( sum_key.plot_id.id_str )
                    sum_histogram.exposure = histogram.exposure
                    sum_histogram.SetTitle(histogram.GetTitle() or "Untitled plot")
                    sums[sum_key] = sum_histogram
                    overlays[overlay_key][PlotID.SUM_STRING] = sum_histogram
                else:
                    sums[sum_key].Add(histogram)
                
            self.WriteAndRegisterCanvas({histogram.GetTitle(): histogram}, self.PrepareOutputPath(key.data_set, key.path, key.plot_id.id_str), key)

        # also write out the sums separately
        for key, histogram in sums.iteritems():
            self.WriteAndRegisterCanvas(
                {histogram.GetTitle(): histogram},
                self.PrepareOutputPath(key.data_set, key.path, key.plot_id.id_str),
                key
            )
            self.histogram_table[key] = histogram
        

        # make overlay plots for plots with multiple categories
        for key, histogram_collection in overlays.iteritems():
            self.WriteAndRegisterCanvas(
                histogram_collection,
                self.PrepareOutputPath(key.data_set, key.path, key.plot_id.id_str),
                key,
                plot_order=self.config.plot_order,
            )
       
        # dataset comparisons.
        # a bit tricky:
        # we always assume that:
        #   * the 'deepest' category set fully partitions the available choices;
        #   * that categories 'higher' than the deepest one are interchangeable
        #     and should all be shown as options.
        # so we have to be pretty careful to do the comparisons correctly
        # if different datasets have different numbers of categories
        # (frequently happens when comparing data
        #  to MC interaction type breakdowns in the same distribution)...
        comparison_dirname = PathTools.ComparisonSubdirName(self.config.histogram_topdirs)
        for no_category_key, plots_by_dataset in by_name_then_dataset.iteritems():
            if len(plots_by_dataset) < 2:
                continue
            
            
            category_lists = {}
            max_num_gps = 0
            for dataset_name, plot_keys in plots_by_dataset.iteritems():
                category_lists[dataset_name] = []
                cl = category_lists[dataset_name]
                for plot_key in plot_keys:
                    for cat_idx, cat_name in enumerate(plot_key.plot_id.categories):
                        if len(cl) < cat_idx + 1:
                            cl.append(set())
                            max_num_gps = max(max_num_gps, cat_idx+1)
                        cl[cat_idx].add(cat_name)

            identical_keys = True
            for gp_idx in range(max_num_gps):
                # if the datasets don't even have the same *number* of categories, clearly their keys are different
                if len(set([len(category_lists[ds]) for ds in category_lists])) != 1:
                    identical_keys = False
                    break
                
                if len(set.union(*[category_lists[ds][gp_idx] for ds in category_lists])) != len(category_lists[category_lists.keys()[0]][gp_idx]):
                    identical_keys = False
                    break
            
            # now combine all the keys from the datasets together
            category_lists = [ set.union(*[cat_list[cat_idx] for ds, cat_list in category_lists.iteritems() if cat_idx < len(cat_list)]) for cat_idx in range(max_num_gps)]
            
            # remember, we assume that the last category group
            # fully partitions the space, so don't choose from
            # among the deepest level.
            # the exception here is if all datasets' plots have
            # exactly the same category partitioning; in that case
            # we should compare them category-by-category.
            if len(category_lists) > 0 and not identical_keys:
                del category_lists[-1]
            
            # go through the category options in sequence.
            # whenever we find a sum, that's the last level
            # for this dataset.
            for category_combination in itertools.product(*category_lists):
                to_plot = {}
                for dataset_name in plots_by_dataset:
                    categories = []
                    i = -2  # want to try [], [SUM_STRING], [category 1][SUM_STRING], ...
                    matched = False
                    no_sum = False
                    while not matched:
                        reduced_groups = {k: v for k,v in plot_key.plot_id.groups.iteritems() if v < len(categories)}
                        reduced_key = PlotCollectionKey(
                            plot_id={
                                "name": plot_key.plot_id.name,
                                "categories": categories,
                                "groups": reduced_groups
                            },
                            data_set = dataset_name,
                            path=no_category_key.path
                        )
                        
                        if reduced_key in self.histogram_table:
                            # we found a sum, but we haven't got one of each
                            # category group, so we know the longest key is longer 
                            # than this one.  we therefore don't want the sum.
                            # instead, we want to go around once more
                            # without the SUM_STRING to get the
                            # single category that matches the category_combination.
                            if len(categories) <= len(category_combination) and len(categories) > 0 and categories[-1] == PlotID.SUM_STRING:
                                if len(categories) > 0:
                                    del categories[-1]  # strip off SUM_STRING
                                no_sum = True
                            else:
                                matched = True
                                to_plot[dataset_name] = reduced_key
                                break
                        
                        i += 1
                        if i >= len(category_combination):
                            break
                        
                        if no_sum:
                            categories.append(category_combination[i])
                        else:
                            if i == -1:
                                categories.append(PlotID.SUM_STRING)
                            else:
                                categories.insert(-1, category_combination[i])
                        
                        
                    
                    if not matched:
                        print "Could not find plot for category sequence", category_combination, "in plot / dataset =", no_category_key.plot_id, "/", dataset_name
                        continue
                
                # if any of the plots are 2D, we (usually) can't do a useful comparison.
                # don't try unless forced to.
                if any(isinstance(self.histogram_table[key], ROOT.TH2) for key in to_plot.itervalues()) and not self.config.force_2d_comp:
                    SINGLE_WARN("comparisons will not be made for 2D plots.  Pass 'force_2d_comp' as true in your config to change.")
                    continue
            
                # find one key that contains all the others.
                # that'll be the one we use to identify this comparison.
                strictest_key = None
                for key in to_plot.values():
                    stricter = all(key.plot_id in other_key.plot_id for other_key in to_plot.values() if other_key != key)
                    if stricter:
                        strictest_key = key
                        break
                assert strictest_key, "Couldn't find one key that contains all the others in this collection.  Keys in collection: %s" % to_plot.values()
                
                # remove the last category (the one with the SUM_STRING in it)
                # (unless it doesn't have SUM_STRING in it because it's from the top level)
                # fixme: should have a way to edit the PlotID in-place, but would require
                # making 'PlotID.categories' and 'PlotID.groups' properties
                # too.  no time right now.
                ##### todo: want to ALSO do the sum_string one if the matching keys
                ####        across the different datasets are all from the same level
                categories = strictest_key.plot_id.categories
                if len(strictest_key.plot_id.categories) > 0 and categories[-1] == PlotID.SUM_STRING:
                    categories = strictest_key.plot_id.categories[:-1]
                comparison_key = PlotCollectionKey(
                    plot_id={
                        "name": strictest_key.plot_id.name,
                        "categories": categories,
                        "groups": {k: v for k,v in strictest_key.plot_id.groups.iteritems() if v < len(strictest_key.plot_id.categories)-1}
                    },
                    data_set=None,
                    path=no_category_key.path
                )
                
                # finally, draw & register
                plots = {ds: self.histogram_table[key] for ds, key in to_plot.iteritems()}
                ratio_denom = None
                if self.config.ratio_denom is not None and self.config.ratio_denom in plots:
                    ratio_denom = self.config.ratio_denom
                self.WriteAndRegisterCanvas(
                    histograms=plots,
                    outfile_stub=self.PrepareOutputPath(PathTools.DATASET_COMPARISON_DIRNAME, comparison_dirname, no_category_key.path, comparison_key.plot_id.id_str),
                    key=comparison_key,
                    include_ratio=True,
                    plot_order=self.config.plot_order,
                    ratio_denom=ratio_denom
                )
                
                if len(to_plot) == 2:
                    c = ComparisonSummary({key: self.histogram_table[key] for key in to_plot.itervalues()})
                    self.comparison_registry.comparisons[comparison_key] = c
               
        self.plot_registry.Serialize(self.config.validation_dir)
        self.comparison_registry.Serialize(self.config.validation_dir)
    
    def WriteAndRegisterCanvas(self, histograms, outfile_stub, key, include_ratio=False, plot_order=[], ratio_denom=None):
        outdir = os.path.dirname(outfile_stub)
        if not os.path.isdir(outdir):
            os.makedirs(outdir)
            
        # use the specified order (if relevant).
        # if any histograms not in the specified order, alphabetize them
        # (reverse alphabetized, actually, so that they wind up
        #  in descending alphabetical order in the legend).
        # also be sure to put all_sum at the top (drawn last)
        histogram_keys = [p for p in plot_order if p in histograms]
        histogram_keys += [p for p in reversed(sorted(histograms.keys())) if p not in plot_order]
        if PlotID.SUM_STRING in histogram_keys:
            histogram_keys.pop(histogram_keys.index(PlotID.SUM_STRING))
            histogram_keys.append(PlotID.SUM_STRING)
        histogram_objs = [histograms[p] for p in histogram_keys]
        ref_idx = histogram_keys.index(ratio_denom) if ratio_denom in histogram_keys else 0
        sample_hist = histogram_objs[ref_idx]
        
        # enumerate all the possible combinations
        option_groups = CanvasMaker.DISPLAY_VARIATIONS.keys()
        option_choices = copy.deepcopy(CanvasMaker.DISPLAY_VARIATIONS.values())  # don't want to run into aliasing issues
        
        # check the choices for _ReferenceHistComp objects,
        # which signify we have some more work to do.
        # for now, the 'reference' histogram is either the first one
        # or the one specified by ref_key
        for choice_list in option_choices:
            for choice_opts in choice_list.itervalues():
                for arg, val in choice_opts.items():
                    if isinstance(val, CanvasMaker._ReferenceHistComp):
                        choice_opts[arg] = [val(sample_hist, h) for h in histogram_objs]
        
        for option_choice_set in itertools.product(*option_choices):

            # special case: "exposure" normalization is dumb if we don't have exposure info
            if option_choice_set[option_groups.index("Normalization")] == "exposure" and \
                    not any(hasattr(h, "exposure") and h.exposure for h in histogram_objs):
                continue
            
            files_written = []
            for suffix in OUTPUT_TYPES:
                option_string = ""
                option_args = {}
                for option_idx, option_choice_name in enumerate(option_choice_set):
                    if suffix in ("json", "root") and option_groups[option_idx] in CanvasMaker.DISPLAY_VARIATIONS_NOJSROOT:
                        continue
                    
                    option_string += "{group=%s,cat=%s}" % (option_groups[option_idx], option_choice_name)
                    option_args.update(option_choices[option_idx][option_choice_name])
            
                args = {
                    "input_histograms": histogram_objs,
                    "ref_hist_idx":     ref_idx,
                    "labels":           histogram_keys,
                    "colours":          GOOD_COLORS,
                    "axis_labels":      "%s;%s" % (sample_hist.GetXaxis().GetTitle(), sample_hist.GetYaxis().GetTitle()), 
                    "save_as":          outfile_stub + option_string,
                    "suffixes":         ["." + suffix,],
                    "with_ratio":       include_ratio,
                }
                
                # sigh. special-case because sometimes the options only want to update ONE of the axis labels
                # (e.g. for normalization)
                if "axis_labels" in option_args:
                    xaxis_pairs, yaxis_pairs = zip(*(args["axis_labels"].split(";") for args in (option_args, args)))
                    option_args["axis_labels"] = ";".join( [update or default for update, default in (xaxis_pairs, yaxis_pairs) ]  )
                
                args.update(option_args)
    
                files_written += self.im.draw(**args)
                self.n_images_written += len(files_written)
                print "\rWrote %d canvases" % self.n_images_written,
                sys.stdout.flush()
#                 for f in files_written:
#                     print "  wrote canvas:", f
                
            
            plot_id = PlotID(key.plot_id.id_str + option_string)
            name = key.plot_id.id_str if len(histograms) > 1 else histogram_keys[0]
            info = CanvasSummary(name=name, plots=histogram_objs, labels=histogram_keys)
            registry_key = PlotCollectionKey(plot_id=plot_id, data_set=None if include_ratio else key.data_set, path=files_written[0])
            if include_ratio:
                self.plot_registry.AddComparisonPlot(registry_key, info)
            else:
                self.plot_registry.AddPlot(registry_key, info)

# bootstrap, load config, go!
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run validation.')
    parser.add_argument('yaml_config', metavar='yaml_config', 
                        help='JSON configuration as specified in README.txt')
    
    args = parser.parse_args()
    
    config = Configuration.FromYAML(args.yaml_config)
    
    maker = CanvasMaker(config)
    
    print "Building canvases..."
    print " (base output directory: %s)" % config.validation_dir
    maker.MakeCanvases()
    print "... done.  Bye."