Organizational.py

Go to the documentation of this file.

"""
  models/Organizational.py:
   Models related to plot organization used in the validation.
  
   Original author: J. Wolcott <jwolcott@fnal.gov>
              Date: September 2016
"""

import collections
import cPickle
import os.path
import pprint
import re
import yaml

class Configuration(object):
    """ Central configuration for Validation run. """
    
    MANDATORY = [
        "histogram_topdirs",
        "output_dir",
        "validation_name",
    ]
    
    DEFAULTS = {
        "plot_accept_filters": [],  # accept only plots with names starting with any of these strings
        "plot_reject_filters": [],  # reject plots with names starting with any of these strings
        "static_files_url":    None,
        "make_dirs":           True,
        "ratio_denom":         None,  # in multi-dataset comparisons, which dataset is the denominator?
        "plot_order":          [],    # organize plots in this order (applies both to lowest level categories in single-sample plots and datasets in multi-dataset comparisons; any unrecognized name is ignored)
        "force_2d_comp":       False, # by default 2D plots won't be compared since the comparisons are usually useless.  Set to True to re-enable
    }

    @staticmethod
    def FromYAML(yaml_fname):
        return Configuration(**yaml.load(open(yaml_fname)))
        
    def __init__(self, **kwargs):
        # dispatch special args for treatment.  otherwise just set
        for kw, val in kwargs.iteritems():
            kw = kw.lower()
            if kw == "folders":
                self.folders = Configuration._ParseFolders(val)
            elif kw in Configuration.MANDATORY or kw in Configuration.DEFAULTS:
                setattr(self, kw, val)
            else:
                raise KeyError("Unknown configuration parameter '%s' with value: %s" % (kw, val))
        
        # set defaults...
        for kw, val in Configuration.DEFAULTS.iteritems():
            if not hasattr(self, kw):
                setattr(self, kw, val)
        
        self._Validate()
        
        self.validation_dir = os.path.join(self.output_dir, self.validation_name)
               
    @staticmethod 
    def _ParseFolders(folder_args):
        return [Folder.FromDict(folder_config) for folder_config in folder_args]
            
    
    def _Validate(self):
        """ Ensure mandatory attributes were specified, consistency between arguments, etc. """
        
        missing_attrs = [k for k in Configuration.MANDATORY if not hasattr(self, k)]
        if len(missing_attrs) > 0:
            raise ValueError("Configuration attributes %s are all mandatory (missing: %s)" % (Configuration.MANDATORY, missing_attrs))
        
        if len(self.plot_accept_filters) > 0 and len(self.plot_reject_filters) > 0:
            raise ValueError("Can't specify both 'plot_accept_filters' and 'plot_reject_filters' in Configuration")

class PlotID(object):
    """ Wrapper around plot identifier.  Simplifies access to 'group' interface. 
    
        ID string as input should have the following format:
        
           name_string<group1_string><group2_string>[possibly more <group_string>s]
        
        where the <group_string>s can either be formatted as
        
           {category_name}
        
        or
        
           {group=group name,cat=category name}
    
        Valid example:
           sample_plot{nocuts}{group=MC,cat=DIS}
    
        Note that if you specify a category using the category-only syntax ('{category_name}'),
        you must ensure it always occurs in the same sequence in any plots that are intended
        to be grouped together (i.e., have the same name_string).
    """
    
    CATEGORY_PATTERN = re.compile(r"(\{group=(?P<group>.+?),cat=(?P<group_cat>.+?)\}|\{(?P<cat>.+?)\})")
    OVERLAY_STRING = "all_overlay"
    SUM_STRING     = "all_sum"
    
    def __init__(self, init):
        # passed a path string 
        if isinstance(init, basestring):
            self.id_str = init
        # passed another PlotID instance -- just copy
        elif hasattr(init, "id_str"):
            self.id_str = init.id_str
        # passed a dictionary with the name, categories, and possibly groups specified
        elif hasattr(init, "__contains__") and all(prop in init for prop in ("name", "categories")):
            id_str = init["name"]
            groups_by_idx = {} if "groups" not in init else dict(zip(init["groups"].itervalues(), init["groups"].iterkeys()))
            for idx, cat in enumerate(init["categories"]):
                if idx in groups_by_idx:
                    id_str += "{group=%s,cat=%s}" % (groups_by_idx[idx], cat)
                else:
                    id_str += "{%s}" % cat
            self.id_str = id_str
            
            # categories, groups will now be re-deduced from the string
            # check that we got the same thing we began with...
            assert self.categories == init["categories"] and \
                ( ("groups" not in init and len(self.groups) == 0) or \
                  (init["groups"] == self.groups) ), \
                "Didn't correctly reconstruct groups and categories for PlotID.  Given: %s, %s; reconstructed: %s, %s" \
                   % (init["groups"], init["categories"], self.groups, self.categories)
                    
    @property
    def id_str(self):
        return self._id_str
    
    @id_str.setter
    def id_str(self, val):
        self._id_str = val
        self._Parse()

    def __contains__(self, other):
        """ This ID 'contains' another if:
              * they share the same 'name'
              * this ID has at least as long a category list
              * the categories match up to the end of the shorter list
              * the group lists contain no incompatible category links. """
        
        if not all( hasattr(other, prop) for prop in ("name", "categories", "groups") ):
            return False
    
        if self.name != other.name:
            return False
        
        if len(self.categories) > len(other.categories):
            return False
        
        min_length = min(len(self.categories), len(other.categories))
        if self.categories[:min_length] != other.categories[:min_length]:
            return False
        
        # make sure the category indices point to the same group names,
        # where they overlap... 
        for key in set(self.groups) & set(other.groups):
            if self.groups[key] != other.groups[key]:
                return False
        
        # also make sure that the groups always point to the same
        # category indices as well
        my_reverse_map = {v:k for k, v in self.groups.iteritems()}
        other_reverse_map = {v:k for k, v in other.groups.iteritems()}
        for key in set(my_reverse_map) & set(other_reverse_map):
            if my_reverse_map[key] != other_reverse_map[key]:
                return False
        
        return True
        

    def __eq__(self, other):
        return hasattr(other, "id_str") and self.id_str == other.id_str
    
    def __repr__(self):
        return "'%s'" % self.id_str
        
    def _Parse(self):
        """ Try to understand the ID string. """
        
        # sequential list of categories buried in the ID
        self.categories = []
        
        # the indices of named groups in the category list above
        # (so, e.g., the category within group "MC_breakdown" might
        #  correspond to index number 2 in the list:
        #   groups["MC_breakdown"] = 2)  
        self.groups = {}
        
        splitted = self.id_str.split("{", 1)
        self.name = splitted[0]
        if len(splitted) == 1:
            return
        
        remnants = "{" + splitted[1] # the initial "{" got eaten by the split... gotta put it back
        for match in PlotID.CATEGORY_PATTERN.finditer(remnants):
            if all([match.group(k) for k in ("group","group_cat")]):
                self.categories.append(match.group("group_cat"))
                self.groups[match.group("group")] = len(self.categories) - 1
            elif match.group("cat"):
                self.categories.append(match.group("cat"))
            else:
                # can't be a match any other way...
                raise ValueError("Something bizarre happened in parsing PlotID.  Matched string pieces: %s" % match.groupdict())
            
            # strip off the matched part
            remnants = remnants.replace(match.group(0), "", 1)
        
        if len(remnants) > 0:
            raise ValueError("Couldn't parse plot name.  Full string: '%s'; unprocessed leftover text: '%s'" % (self.id_str, remnants))

class PlotCollectionKey(object):
    def __init__(self, plot_id, data_set, path=""):
        # funky initialization syntax required
        # since __setattr__ overridden for immutability purposes below
        plot_id = plot_id if isinstance(plot_id, PlotID) else PlotID(plot_id)
        object.__setattr__(self, "plot_id", plot_id)
        object.__setattr__(self, "path", path)
        object.__setattr__(self, "data_set", data_set)
    
    def __hash__(self):
        return hash((self.plot_id.id_str, self.data_set, self.path))
    
    def __contains__(self, other):
        """ One Key 'contains' another if one of their PlotIDs contains the other
            and their other attributes are the same. """

        if not all(hasattr(other, attr) for attr in ("plot_id", "data_set", "path")):
            return False
       
        return self.plot_id in other.plot_id and \
                all(getattr(self, prop) == getattr(other, prop) for prop in ("data_set", "path"))
    
    def __eq__(self, other):
        if not all(hasattr(other, attr) for attr in ("plot_id", "data_set", "path")):
            return False
        return (self.plot_id, self.data_set, self.path) == (other.plot_id, other.data_set, other.path)

    # to cope with Python 2 bug that makes default __ne__() ~ 'not is' instead of '! __eq__()'
    def __ne__(self, other):
        return not (self == other)
    
    def __repr__(self):
        return "%s|%s|%s" % (self.plot_id, self.data_set, self.path)
    
    # keys for dictionaries aren't allowed to change.
    # this won't 100% guarantee it (some weird edge cases),
    # but is enough to prevent accidents
    def __delattr__(self, attr):
        raise TypeError("%s is intended to be immutable" % self.__class__)
    def __setattr__(self, attr, val):
        raise TypeError("%s is intended to be immutable" % self.__class__)

class PlotSet(collections.MutableSet):
    """ Collection of unique plot keys. """ 
    def __init__(self):
        self._cache = set()

    # emulate set interface.
    # defining these methods allows the collections.MutableSet mixin to deduce the rest
    # (see https://docs.python.org/2/library/collections.html#collections.MutableSet)
        
    def __contains__(self, obj):
        if not isinstance(obj, PlotCollectionKey):
            raise TypeError("%s items must be of type %s, not %s" % (self.__class__.__name__, PlotCollectionKey, type(obj)))
        return obj in self._cache
    
    def __iter__(self):
        return iter(self._cache)
    
    def __len__(self):
        return len(self._cache)
    
    def __repr__(self):
        return repr(self._cache)
    
    def add(self, obj):
        if not isinstance(obj, PlotCollectionKey):
            raise TypeError("%s items must be of type %s, not %s" % (self.__class__.__name__, PlotCollectionKey, type(obj)))
        
        self._cache.add(obj)
        
    def discard(self, obj):
        self._cache.discard(obj)
        
    ### addons
    
    def iter_plot_collections(self):
        """ Iterate over the collections of plots with unique names """
        
        work = list(self._cache)  # make a copy
        
        while len(work) > 0:
            collection = PlotSet()
            plot_name = work[0].plot_id.name
            collection.add([work.pop(0)])
            
            # iterate backwards because we're going to be removing stuff
            # and we don't want the indices of unchecked stuff jumping under us
            for plot_idx in range(len(work)-1, 0, -1):
                if work[plot_idx].plot_id.name == plot_name:
                    collection.add(work.pop(plot_idx))
            
            yield collection

class PlotLibrary(collections.MutableMapping):
    """ Lookup collection for plots, keyed by PlotID """
    
    def __init__(self):
        self._cache = {}
    
    def __repr__(self):
        return repr(self._cache)
    
    # emulate container interface.
    # defining these methods allows the collections.MutableMapping mixin to deduce the rest
    # (see https://docs.python.org/2/library/collections.html#collections.MutableMapping)
    def __contains__(self, key):
        return key in self._cache
    
    def __delitem__(self, key):
        del self._cache[key]
        
    def __getitem__(self, key):
        if not isinstance(key, PlotCollectionKey):
            raise TypeError("%s keys must be of type %s, not %s" % (self.__class__.__name__, PlotCollectionKey, type(key)))
        
        try:
            return self._cache[key]
        except KeyError:
            # customize the error message -- don't expose the internal dict
            raise KeyError("%s has no such plot: %s" % (self.__class__.__name__, key))

    def __iter__(self):
        return iter(self._cache)

    def __len__(self):
        return len(self._cache)
    
    def __setitem__(self, key, value):
        if not isinstance(key, PlotCollectionKey):
            raise TypeError("%s keys must be of type %s" % (self.__class__.__name__, PlotCollectionKey))

        self._cache[key] = value

class ComparisonRegistry(object):
    SERIALIZE_FILENAME = "comparison_registry.pkl"
    
    def __init__(self):
        self.comparisons = PlotLibrary()

    def Serialize(self, base_path):
        cPickle.dump(self, open(os.path.join(base_path, ComparisonRegistry.SERIALIZE_FILENAME), "w"))
    
    @staticmethod
    def Deserialize(file_name):
        obj = cPickle.load(open(file_name))
        if not isinstance(obj, ComparisonRegistry):
            raise TypeError("Object loaded from alleged ComparisonRegistry file is actually a %s" % type(obj))
        return obj
    
    

class PlotRegistry(object):
    """ Collect summaries of plots drawn and where they were written in the same place """
    
    SERIALIZE_FILENAME = "plot_registry.pkl"
    
    def __init__(self):
        self.plots_by_dataset = {}
        self.dataset_comparisons = PlotLibrary()
    
    def __repr__(self):
        return "PlotRegistry:\n by dataset: %s\n cross-dataset comparisons: %s" % (pprint.pformat(self.plots_by_dataset), pprint.pformat(self.dataset_comparisons))
    
    def AddPlot(self, plot_key, plot_summary):
        assert isinstance(plot_key, PlotCollectionKey)
        
        self.plots_by_dataset.setdefault(plot_key.data_set, PlotLibrary())[plot_key] = plot_summary
        
    def AddComparisonPlot(self, plot_key, plot_summary):
        self.dataset_comparisons[plot_key] = plot_summary
    
    def AddDatasetPlots(self, dataset, plot_summaries):
        collection = self.plots_by_dataset.setdefault(dataset, PlotLibrary())
        collection.update(plot_summaries)
        
    def Serialize(self, base_path):
        cPickle.dump(self, open(os.path.join(base_path, PlotRegistry.SERIALIZE_FILENAME), "w"))
    
    @staticmethod
    def Deserialize(file_name):
        obj = cPickle.load(open(file_name))
        if not isinstance(obj, PlotRegistry):
            raise TypeError("Object loaded from alleged PlotRegistry file is actually a %s" % type(obj))
        return obj

class Folder(object):
    """ Simple container for creating structure in plot collections. """
    
    ALLOWED_ATTRS = [
        "name",
        "plot_ids",
        "subfolders",
    ]
    
    @staticmethod
    def FromDict(config_dict):
        # JSON/YAML format (and thus Python dictionary) for a folder looks like:
        #     {
        #         "name": "folder1",
        #         "plot_ids": ["key1", "key2"],
        #     
        #         "subfolders": [
        #             {
        #                 "name": "subfolder1",
        #                 "plot_ids": [ "key1.1", "key1.2"    ]
        #             },
        #             
        #             {    
        #                 "name": "subfolder2",
        #                 "plot_ids": [    "key2.1", "key2.2" ]
        #             }
        #         ]
        #     }
        assert "name" in config_dict, "Folder configuration must specify 'name'"
        
        for attr in config_dict:
            assert attr in Folder.ALLOWED_ATTRS, "Unrecognized attribute in folder configuration: '%s'" % attr
        
        args = {attr: config_dict[attr] for attr in Folder.ALLOWED_ATTRS if attr in config_dict}
        
        return Folder(**args)     
    
    def __init__(self, name, plot_ids=[], subfolders=[]):
        self.name = name
        
        assert hasattr(plot_ids, "__iter__"), "Invalid configuration supplied for Folder plot entries: '%s'" % plot_ids
        self.plot_entries = [PlotID(p) if not isinstance(p, PlotID) else p for p in plot_ids]
        
        assert hasattr(subfolders, "__iter__"), "Invalid configuration supplied for Folder subfolders: '%s'" % subfolders
        self.subfolders = []
        for subfolder in subfolders:
            if isinstance(subfolder, Folder):
                self.subfolders.append(subfolder)
            elif hasattr(subfolder, "__contains__"):
                self.subfolders.append(Folder.FromDict(subfolder))