std_candles.py

Go to the documentation of this file.

"""
  std_candles.py:
    Plots for energy standard candles studies.
    
    Requires ntuple(s) created by the EnergyStandardCandles_module.cc in this package.
    
  Original author: J. Wolcott <jwolcott@fnal.gov>
             Date: December 2015

"""

import array
import itertools
import numbers
import os, os.path
import random
import sys
import types

import ROOT
ROOT.gROOT.SetBatch(True)

# the ntuples contain custom objects.  can't load them without this.
# (cafe does basically the same thing for CAFs.)
import PyCintex
ROOT.gSystem.Load("libCalibration_dict")  # in the same package as this macro, so should be ok without modifying LD_LIBRARY_PATH further

sys.path.append("/nova/app/users/jwolcott/analysis")  # so that when I forget to add this to a package later somebody else can still find it
from plot_tools.NtupleUtils import TreeProcessor
from plot_tools import HistoTools

TESTING = False

TREE_NAME = "stdcandles/MichelE"

OUT_LOCATION = "/nova/ana/users/jwolcott/stdcandles"
OUTPUT_TYPES = ["eps", "png", "root"]

TRUE_ORIGIN_PDGS = [11, 13, 111, 211, 2112, 2212, 0]
#TRUE_ORIGIN_PDGS = [11, 2112, 0]
TRUE_ORIGIN_COLORS = {
        11: ROOT.kRed+1,
        13: ROOT.kBlue+1,
        211: ROOT.kGreen+2,
        111: ROOT.kViolet+2,
        2112: ROOT.kOrange-8,
        2212: ROOT.kBlack,
        0: ROOT.kGray,
}
LEGEND_ENTRIES = {
        11: "Michel e^{#pm}",
        13: "#mu^{#pm}",
        111: "#gamma from #pi^{0}",
        211: "#pi^{#pm}",
        2112: "#gamma from neutron capture",
        2212: "p^{+}",
        0: "Other",
}

NUCLEUS_LABELS = {
        0: "Other",
        1000170350: "^{35}Cl",
        1000220480: "^{48}Ti",
}

NUCLEUS_COLORS = {
        0: ROOT.kBlack,
        1000160320: ROOT.kMagenta+4,
        1000170350: ROOT.kBlue,
        1000170370: ROOT.kRed,
        1000220480: ROOT.kOrange+2,
        1000220490: ROOT.kGreen+2,
}

#PARENTFRONTZ_RANGE = (400, float("inf"))
PARENTFRONTZ_RANGE = None

NEUTRON_REWEIGHT = 1  # 0.58

MICHELID_RANGE = (-9, float("inf")) # keep everything except the junk
#MICHELID_RANGE = (0, float("inf"))   # good quality Michels only
#MICHELID_RANGE = (-6, -2)            # neutron capture rich sample


ENERGY_BINS = [0.5*i for i in range(120)] # + range(40, 100)
ZCOORD_BINS = [0,] + range(10, 1160, 50)

class StdCandlesPlotter(TreeProcessor):
        def __init__(self, testing=False):
                if "pc_" in NTUPLE_FILELIST:
                        suffix = "PC"
                elif "MC" in NTUPLE_FILELIST:
                        suffix = "MC"
                else:
                        suffix = "data"
                name = "StdCandles_" + suffix
                max_nevts = 1000 if testing else None
                super(StdCandlesPlotter, self).__init__(tree_name=TREE_NAME, file_list=NTUPLE_FILELIST, name=name, max_nevts=max_nevts)
                
                self.POT = 0

        def RegisterBranches(self):
                self._branches_to_enable = [
                        "run",
                        "subRun",
                        "spillNum",
                        "michel_candidate",
                ]
        
        def RegisterMethods(self):
                self._initialize_methods += [ self.Warnings, self.PrepareMichelEPlots ]
                self._process_methods.append( self.FillMichelEPlots )
                self._finalize_methods += [ self.FinalizeMichelEPlots, self.StorePOT ]
                
                self._tree_methods.append( self.AddPOT )

        # --------------------
        def IsMC(self):
                return "MC" in self.name or "PC" in self.name
                
        def SavePlots(self):
                c = ROOT.TCanvas()
                for name, plotobj in self.michel_histos.iteritems():
                        c.Clear()
                        plotobj.Draw()
                        obj_to_print = plotobj if isinstance(plotobj, ROOT.TCanvas) else c
                        for out_type in OUTPUT_TYPES:
                                obj_to_print.Print( os.path.join(OUT_LOCATION, self.name.replace("StdCandles_", ""), "MichelEPlots", "%s.%s" % (name, out_type)) )

        def Warnings(self, tree):
                if NEUTRON_REWEIGHT != 1:
                        print >> sys.stderr, "\nWARNING: neutron captures being reweighted by factor:", NEUTRON_REWEIGHT
                        
                if PARENTFRONTZ_RANGE is not None:
                        print >> sys.stderr, "\nWARNING: requiring parent tracks to have z coordinates in range", PARENTFRONTZ_RANGE
                        
                if MICHELID_RANGE is not None:
                        print "\nNote: Michel PID range:", MICHELID_RANGE
                
                print

        # --------------------

        # --------------------
        
        def AddPOT(self):
#               print dir(self.fChain)
                tfile = self.fChain.GetCurrentFile()
                POT_param = tfile.Get("stdcandles/POT")
                if not POT_param:
                        print >> sys.stderr, "WARNING: no POT information for file:", tfile.GetName()
                        return
                self.POT += POT_param.GetVal()
                
        def StorePOT(self):
                if not self.IsMC():
                        self.POT *= 1e12  # since the data ART files store it in units of 10^12
                print "Total POT:", self.POT
                objs_to_tag = []
                for plotname, plotobj in self.michel_histos.iteritems():
                        if hasattr(plotobj, "GetListOfPrimitives"):
                                for obj in plotobj.GetListOfPrimitives():
                                        if hasattr(obj, "GetListOfFunctions"):
                                                objs_to_tag.append(obj)
                                        elif hasattr(obj, "GetHists"):
                                                for o in obj.GetHists():
                                                        if hasattr(o, "GetListOfFunctions"):
                                                                objs_to_tag.append(o)
                                continue
                        elif hasattr(plotobj, "GetHists"):
                                for obj in obj.GetHists():
                                        if hasattr(obj, "GetListOfFunctions"):
                                                objs_to_tag.append(obj)
                                continue
                        
                        if not hasattr(plotobj, "GetListOfFunctions"):
                                print >> sys.stderr, "WARNING: can't store POT in plot '%s'" % plotname
                                continue
                        
                        objs_to_tag.append(plotobj)

#               print "Saving POT in objs:", objs_to_tag
                for obj in objs_to_tag:
                        param = ROOT.TParameter("double")("POT", self.POT)
                        ROOT.SetOwnership(param, False)   # the plotobj's destructor will want to delete this, so don't let it disappear when it goes out of scope
                        obj.GetListOfFunctions().Add(param)
                
        
        def PrepareMichelEPlots(self, chain):
                self.michel_histos = {}
                
                self.michel_histos["MichelEnergy"] = ROOT.TH1D("michel_energy", ";Candidate Michel electron E_{reco} (MeV);Events/MeV;", len(ENERGY_BINS)-1, array.array('d', ENERGY_BINS))
                self.michel_histos["MichelPID"] = ROOT.TH1D("michel_pid", ";Michel PID score;Events/MeV;", 120, -6, 6)
#               self.michel_histos["MichelPIDVsDeltaT"] = ROOT.TH2D("michel_pid_vs_deltaT", ";#Delta T (ns);Michel PID score;", 20, 0, 18000, 120, -6, 6)
#               self.michel_histos["MichelPIDVsDeltaT"].SetOption("colz")
                self.michel_histos["DeltaT"] = ROOT.TH1D("time_diff", ";Candidate Michel #Delta_{t} (ns);Events;", 20, 0, 18000)
                self.michel_histos["NHits"] = ROOT.TH1D("num_hits", ";N_{hit} in Michel candidates (MeV);Events;", 15, 0, 15)
                self.michel_histos["HitE"] = ROOT.TH1D("energy_per_hit", ";All E_{hit} in Michel candidates (MeV);Events;", 100, 0, 20)
                self.michel_histos["MeanHitE"] = ROOT.TH1D("mean_energy_per_hit", ";Candidate Michel electron E_{reco}/N_{hits} (MeV);Events;", 100, 0, 20)
                self.michel_histos["HitTimeRMS"] = ROOT.TH1D("hit_time_RMS", ";RMS of hit times within candidate Michel (ns);Events;", 50, 0, 25)
                self.michel_histos["ParentFrontZ"] = ROOT.TH1D("parent_front_Z", ";z-coord of parent track start (cm);Events/cm;", len(ZCOORD_BINS)-1, array.array('d', ZCOORD_BINS))
                
                if self.IsMC():
                        self.michel_histos["TrueMichelEnergy"] = ROOT.TH1D("true_michel_energy", ";Michel electron E_{true} (MeV);Events;", 100, 0, 100)
                        self.michel_histos["TrueMissingE_NonBrem"] = ROOT.TH1D("true_missing_E_nonbrem", ";Michel E_{true} - E_{'captured'} - Brem. E_{true} (MeV);Events;", 100, 0, 100)
                        self.michel_histos["TrueMissingE"] = ROOT.TH1D("true_missing_E", ";Michel E_{true} - E_{'captured'} - Brem. E_{true} - E_{inactive} (MeV);Events;", 100, -50, 50)
                        self.michel_histos["TrueDeadMaterialE"] = ROOT.TH1D("true_E_in_dead_materials", ";Estimated true E_{Michel} deposited in inactive materials (MeV);Events;", 100, 0, 100)
                        self.michel_histos["TrueMichelFracCaptured"] = ROOT.TH1D("michel_frac_captured", ";Fraction of true Michels' energy captured;Events;", 101, 0, 1.01)
                        self.michel_histos["TrueELostToBremVsRecoE"] = ROOT.TH2D("michel_E_brem", ";Michel E_{reco};True Michel energy lost to Brem. (MeV);", 100, 0, 100, 100, 0, 100)
                        
                        self.michel_histos_by_pdg = {}
                        for plotname, plot in self.michel_histos.iteritems():
                                if "TrueMichel" in plotname or isinstance(plot, ROOT.TH2D):
                                        continue
                                self.michel_histos_by_pdg[plotname] = {}
                                for pdg in TRUE_ORIGIN_PDGS:
                                        h = ROOT.TH1D(plot)
                                        h.SetName(h.GetName() + "_TruePID_%d" % pdg)
                                        h.SetLineColor(TRUE_ORIGIN_COLORS[pdg])
                                        self.michel_histos_by_pdg[plotname][pdg] = h
                
                self.michel_origins = []
                self.neutron_captures = {}

        def FillMichelEPlots(self):
                tree = self.fChain  # save some typing
        
                michel_candidate = tree.michel_candidate  # slightly less typing, and one less reference for Python to traverse
        
                if PARENTFRONTZ_RANGE is not None and not (PARENTFRONTZ_RANGE[0] <= michel_candidate.parentFrontZ <= PARENTFRONTZ_RANGE[1]):
                        return
                
                if MICHELID_RANGE is not None and not (MICHELID_RANGE[0] <= michel_candidate.MID < MICHELID_RANGE[1]):
                        return
                
                energy = sum(michel_candidate.hitEnergies)
                nHits = len(michel_candidate.hitEnergies)
                if energy <= 0:
                        return
                
                vals = {
                        "MichelEnergy": energy,
                        "MichelPID": michel_candidate.MID,
                        "NHits": nHits,
                        "HitE": michel_candidate.hitEnergies,
                        "DeltaT": michel_candidate.deltaT,
                        "MeanHitE": energy/nHits,
                        "HitTimeRMS": michel_candidate.t_RMS,
                        "ParentFrontZ": michel_candidate.parentFrontZ,
                }
                
                wgt = 1
                
#               print tree.energy
#               print len(michel_candidate.eContribs)                   
                if self.IsMC() and len(michel_candidate.eContribs) > 0:
                        origin_breakdown = {}
                        for contrib in michel_candidate.eContribs:
                                if hasattr(contrib, "ElostToBrem"):
                                        self.michel_histos["TrueMichelFracCaptured"].Fill(contrib.fracCaptured)
                                        self.michel_histos["TrueMichelEnergy"].Fill(contrib.trueFourP.E())
                                        
                                        if contrib.contribEFrac > 0.9:
                                                self.michel_histos["TrueMissingE_NonBrem"].Fill(contrib.trueFourP.E() - contrib.trueFourP.E()*contrib.fracCaptured -  contrib.ElostToBrem)
                                                self.michel_histos["TrueMissingE"].Fill(contrib.trueFourP.E() - contrib.trueFourP.E()*contrib.fracCaptured - contrib.ElostToBrem - contrib.ElostToDeadMaterials)
                                                self.michel_histos["TrueDeadMaterialE"].Fill(contrib.ElostToDeadMaterials)
                                                self.michel_histos["TrueELostToBremVsRecoE"].Fill(energy, contrib.ElostToBrem)
                                
                                pdg = abs(contrib.pdg)
                                if pdg not in TRUE_ORIGIN_PDGS:
                                        pdg = 0
#                               print pdg, contrib.contribEFrac
                                origin_breakdown[pdg] = origin_breakdown.setdefault(pdg, 0) + contrib.contribEFrac
                                
                                if hasattr(contrib, "nCaptureNucleus") and contrib.nCaptureNucleus != 0:
                                        self.neutron_captures[contrib.nCaptureNucleus] = self.neutron_captures.setdefault(contrib.nCaptureNucleus, 0) + 1
                        
                        total_frac = sum(origin_breakdown.values())
#                       print origin_breakdown, total_frac
                        assert abs(total_frac - 1) < 1e-5, "total fraction not 1!  sum: %f" % total_frac
                        
                        self.michel_origins.append(origin_breakdown)

                        for pdg, frac in origin_breakdown.iteritems():
                                if frac > 0.75:
                                        if pdg == 2112:
                                                wgt = NEUTRON_REWEIGHT
                                        for plotname, val in vals.iteritems():
                                                if plotname not in self.michel_histos:
                                                        continue
                                                try:
                                                        vs = list(val)
                                                except TypeError:
                                                        vs = [val,]
                                                for v in vs:
                                                        self.michel_histos_by_pdg[plotname][pdg].Fill(v, wgt)

#                       if set(origin_breakdown.keys()) != set([11,]):
#                               return

#               print tree.energy
                for plotname, val in vals.iteritems():
                        if plotname not in self.michel_histos:
                                continue
                        try:
                                vs = list(val)
                        except TypeError:
                                vs = [val,]
                        for v in vs:
                                self.michel_histos[plotname].Fill(v, wgt)
                
#               self.michel_histos["MichelPIDVsDeltaT"].Fill(michel_candidate.deltaT, michel_candidate.MID)
                

        def FinalizeMichelEPlots(self):
                hists = self.michel_histos.values()
                if self.IsMC():
#                       print "adding hists:",   self.michel_histos_by_pdg.values()
                        for coll in self.michel_histos_by_pdg.itervalues():
                                hists += coll.values()
#                               print coll, hists
                for p in hists:
#                       print p
                        axis_variable = [getattr(p, "Get%saxis" % axis)().IsVariableBinSize() for axis in ("X", "Y")]
                        if any(axis_variable):
                                p.Scale(1, "width")
#                               print "Scaled plot:", p.GetName()
                
                if self.IsMC():
                        self.MakeBremPlots()
                        self.MakeOverlayMCPlots()
                        self.MakeTrueOriginBreakdownPlot()
                        self.MakeNeutronCapturePlot()
                        
        # ---------------------
        
        def FillMichelBreakdownPlots(self, histo_stack):
                fracs = {}
                nCurrEvts = 0
                nBin = 1
                n_bins = min(200, len(self.michel_origins))
                n_evts_in_bin = float(len(self.michel_origins)) / n_bins
                nEvtsInBin_int = int(n_evts_in_bin)
                nEvtsInBin_frac = n_evts_in_bin - nEvtsInBin_int
                for breakdown in self.michel_origins:
                        # if this is all the events that should go into this bin,
                        # average and insert into the plot
                        if nCurrEvts >= nEvtsInBin_int + int(random.random() < nEvtsInBin_frac):
                                for pdg in TRUE_ORIGIN_PDGS:
                                        frac = 0
                                        if pdg in fracs:
                                                frac = fracs[pdg] / float(nCurrEvts)
                                        histo_stack[pdg].SetBinContent(nBin, frac)
#                                       print "  PDG %d: setting bin %d to have content %f" % (pdg, nBin, frac)
                                fracs = {}
                                nCurrEvts = 0
                                nBin += 1

#                       print breakdown
                        for pdg, frac in breakdown.iteritems():
#                                       LOG_DEBUG("StandardCandlePlots") << "  PDG " << pdgFracPair.first << " contributed fraction " << pdgFracPair.second << " to Michel";
                                fracs[pdg] = fracs.setdefault(pdg, 0) + frac
#                               LOG_DEBUG("StandardCandlePlots") << "  fracs is now: ";
#                               for (const auto & f : fracs)
#                                       LOG_DEBUG("StandardCandlePlots") << "   PDG " << f.first << ": " << f.second;
                        nCurrEvts += 1;

        def MakeBremPlots(self):
                twoD_plot = self.michel_histos["TrueELostToBremVsRecoE"]
                ROOT.SetOwnership(twoD_plot, False)
                canvas = ROOT.TCanvas()
                canvas.cd()
                twoD_plot.Draw("colz")
                prof = twoD_plot.ProfileX()
                prof.Draw("same")
                self.michel_histos["TrueELostToBremVsRecoE"] = canvas
                
                proj = twoD_plot.ProjectionY("TrueELostToBrem")
                self.michel_histos["TrueELostToBrem"] = proj
                
        def MakeNeutronCapturePlot(self):
#               print self.neutron_captures
                key_order = NUCLEUS_LABELS.keys()
                vals = dict( [(k, self.neutron_captures[k]) for k in key_order if k in self.neutron_captures])
                vals[0] = sum([v for k, v in self.neutron_captures.iteritems() if k not in vals])
                colors = array.array('i', [NUCLEUS_COLORS[pdg] for pdg in key_order])
                vals = array.array('d', [vals[k] for k in key_order if k in vals])
                n_nuclei = len(vals)
                pie = ROOT.TPie("NeutronCaptureNucleus", "", n_nuclei, vals, colors)
                for i, pdg in enumerate(key_order):
                        pie.SetEntryLabel(i, NUCLEUS_LABELS[pdg])
                pie.SetRadius(0.2)
                pie.SetLabelFormat("#splitline{%txt}{%perc}");
                self.michel_histos["NeutronCaptureNucleus"] = pie

        def MakeOverlayMCPlots(self):
                for plotname, plot in self.michel_histos.iteritems():
                        if "True" in plotname:
                                continue

                        ROOT.SetOwnership(plot, False)

                        legend = ROOT.TLegend(0.55, 0.6, 0.83, 0.9)
                        ROOT.SetOwnership(legend, False)

                        keys = self.michel_histos_by_pdg[plotname].keys()
                        total_exclusive = self.michel_histos_by_pdg[plotname][keys[0]].Clone("%s_sum_exclusive" % plotname)
                        ROOT.SetOwnership(total_exclusive, False)
                        for k in keys[1:]:
                                total_exclusive.Add(self.michel_histos_by_pdg[plotname][k])
                        total_exclusive.SetLineColor(ROOT.kBlack)
                        total_exclusive.SetLineStyle(3)

                        h_by_pdg = ROOT.THStack()
                        ROOT.SetOwnership(h_by_pdg, False)
                        hists_to_consider = {"Total": plot, "Total-exclusive": total_exclusive}
                        hists_to_consider.update(self.michel_histos_by_pdg[plotname])
                        for pdg in ["Total", "Total-exclusive"] + TRUE_ORIGIN_PDGS:
                                h = hists_to_consider[pdg]
#                               print "  adding histogram:", h
                                h_by_pdg.Add(h, "hist")
                                if pdg == "Total":
                                        legend_entry = "Total"
                                elif pdg == "Total-exclusive":
                                        legend_entry = "Total of exclusive categories"
                                else:
                                        legend_entry = LEGEND_ENTRIES[pdg]
                                legend.AddEntry(h, legend_entry, "l")
                        c = ROOT.TCanvas()
                        h_by_pdg.Draw("nostack")
                        HistoTools.SanitizeStackLabels(h_by_pdg)

                        # legend winds up overlapping with distributions in this one plot.
                        # not a wonderful place for special cases but will move if it gets unwieldy
                        if "MichelPID" in plotname:
                                print "fixing Michel PID plot legend"
                                legend.SetX1(0.2)
                                legend.SetX2(0.5)
                        legend.Draw()
                        
                        self.michel_histos[plotname] = c
                
        def MakeTrueOriginBreakdownPlot(self):
                legend = ROOT.TLegend(0.55, 0.5, 0.83, 0.8)
                legend.SetFillColor(ROOT.kWhite)
                legend.SetBorderSize(1)
                ROOT.SetOwnership(legend, False)

                n_bins = min(200, len(self.michel_origins))
                origin_by_pdg = {}
                for pdg in TRUE_ORIGIN_PDGS:
                        h = ROOT.TH1D("true_origin_%d" % pdg, ";Fraction of events;Fraction of true energy;", n_bins, 0, 1)
                        h.SetLineColor(TRUE_ORIGIN_COLORS[pdg])
                        h.SetFillColor(TRUE_ORIGIN_COLORS[pdg])
                        origin_by_pdg[pdg] = h
                        
                        legend.AddEntry(h, LEGEND_ENTRIES[pdg], "f")

                self.michel_origins.sort(
                        # yikes.
                        # basically: sort by amount of each PDG in priority order of TRUE_ORIGIN_PDGS
                        key=lambda bkdown:
                                sum(100**(len(TRUE_ORIGIN_PDGS)-i)*(bkdown[pdg] if pdg in bkdown else 0) for i, pdg in enumerate(TRUE_ORIGIN_PDGS))
                )
                
                self.FillMichelBreakdownPlots(origin_by_pdg)
                
                st = ROOT.THStack()
                ROOT.SetOwnership(st, False)
                for pdg in TRUE_ORIGIN_PDGS:
                        ROOT.SetOwnership(origin_by_pdg[pdg], False)
                        st.Add(origin_by_pdg[pdg], "hist")
                
                c = ROOT.TCanvas()
                st.Draw()
                legend.Draw()
                HistoTools.SanitizeStackLabels(st)
                self.michel_histos["TrueMichelOriginBreakdown"] = c

# -----------------

class PlotConfig(object):
        AREA_NORM = object()
        PEAK_NORM = object()
        POT_NORM = object()
        def __init__(self, name, path, mc_norm_factor=AREA_NORM):
                self.name = name
                self.path = path
                self.mc_norm_factor = mc_norm_factor

COMBINATION_PLOTS = [
        PlotConfig( name="MichelEnergy", path="MichelEPlots/MichelEnergy", mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig( name="MichelPID",    path="MichelEPlots/MichelPID",    mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig( name="DeltaT",       path="MichelEPlots/DeltaT",       mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig( name="NHits",        path="MichelEPlots/NHits",        mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig( name="HitE",         path="MichelEPlots/HitE",         mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig( name="HitTimeRMS",   path="MichelEPlots/HitTimeRMS",   mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig (name="MeanHitE",     path="MichelEPlots/MeanHitE",     mc_norm_factor=PlotConfig.POT_NORM ),
        PlotConfig (name="ParentFrontZ", path="MichelEPlots/ParentFrontZ", mc_norm_factor=PlotConfig.POT_NORM ),
]

class DataMCPlots(object):
        def __init__(self):
                self.data_plots_in = {}
                self.mc_plots_in = {}

                self.plots_out = {}
                
                self.LoadPlots()
                
        @staticmethod
        def GetMCDataStack(input):
                stack = None
                if hasattr(input, "GetListOfPrimitives"):
                        stacks = [o for o in input.GetListOfPrimitives() if hasattr(o, "GetHists")]
                        assert len(stacks) == 1, "Too many THStacks in collection '%s' for me to deduce which one you want!" % input.GetName()
                        stack = stacks[0]
                elif hasattr(input, "GetHists"):
                        stack = input
                assert stack, "Could not determine the stack to pull the sum out of!"
                
                return stack
        
        def LoadPlots(self):
                for plot_config in COMBINATION_PLOTS:
                        for source in ("data", "mc"):
                                collection = getattr(self, "%s_plots_in" % source)
                                
                                # get the capitalization right.  should've been more consistent
                                source_infix = source                           
                                if source == "mc":
                                        source_infix = "MC"
                                plotfile_name = os.path.join(OUT_LOCATION, source_infix, "%s.root" % plot_config.path)
                                if not os.path.isfile(plotfile_name) :
                                        raise Exception("Could not find necessary input file: '%s'" % plotfile_name)
                                
                                f = ROOT.TFile(plotfile_name)
                                keys = f.GetListOfKeys()
                                if len(keys) > 1:
                                        raise Exception("Don't know which of the %d objects in file you want: %s" % (len(objs), list(keys)))
                                elif len(keys) == 0:
                                        raise Exception("File has nothing in it: '%s'" % plotname_file)
                                
                                obj = keys.First().ReadObj()
                                assert obj, "Object read from file ('%s') is null!" % plotfile_name
                                if hasattr(obj, "SetDirectory"):
                                        obj.SetDirectory(None)  # otherwise they disappear when the file is closed
                                collection[plot_config.path] = obj
                                
        def MakePlots(self):
                for plot_config in COMBINATION_PLOTS:
                        method_name = "Make_%s_Plot" % plot_config.name
                        if hasattr(self, method_name) or self.Make_Default_Plot_Method(plot_config.name, plot_config.path, plot_config.mc_norm_factor):
                                getattr(self, method_name)()
                        else:
                                raise Exception("Don't know how to make '%s' plot (no method '%s' defined)" % (plot_config.name, method_name))
        
        def Make_DataMC_Overlay(self, in_varname):
                mc_canvas = self.mc_plots_in[in_varname]
                assert isinstance(mc_canvas, ROOT.TCanvas)
                mc_stack = [o for o in mc_canvas.GetListOfPrimitives() if isinstance(o, ROOT.THStack)][0]
                
                legend = ROOT.TLegend(0.65, 0.7, 0.85, 0.9)
                ROOT.SetOwnership(legend, False)
                
                data_canvas = self.data_plots_in[in_varname]
                data_plot = [o for o in data_canvas.GetListOfPrimitives() if isinstance(o, ROOT.TH1)][0]
                data_plot.SetMarkerSize(0)
                legend.AddEntry(data_plot, "Data", "le")
                
                # we're just going to do data-MC total
                # comparisons (to avoid over-complicating the figure).
                out_stack = ROOT.THStack(mc_stack)
                ROOT.SetOwnership(out_stack, False)
                for h in list(out_stack.GetHists()):
                        if "TruePID" in h.GetName() or "sum_exclusive" in h.GetName():
                                out_stack.GetHists().Remove(h)
                        else:
                                h.SetLineColor(ROOT.kRed)
                                legend.AddEntry(h, "MC", "l")
                
                out_stack.Add(data_plot, "pe")
                
                c = ROOT.TCanvas()
                c.cd()
                ROOT.SetOwnership(c, False)
                out_stack.Draw("nostack")
                HistoTools.SanitizeStackLabels(out_stack)
                legend.Draw()
                return c
        
        @staticmethod
        def NormalizeMCPlot(data_plot, mc_plot, mc_norm_factor=PlotConfig.AREA_NORM):
                
                assert data_plot and hasattr(mc_plot, "Scale")
                
                scale_factor = None
                if isinstance(mc_norm_factor, numbers.Number) and mc_norm_factor >= 0:
                        scale_factor = mc_norm_factor
                elif mc_norm_factor is PlotConfig.AREA_NORM:
                        scale_factor = data_plot.Integral() / mc_plot.Integral()
                elif mc_norm_factor is PlotConfig.PEAK_NORM:
                        scale_factor = data_plot.GetMaximum() / mc_plot.GetMaximum()
                elif mc_norm_factor is PlotConfig.POT_NORM:
                        data_pot = data_plot.GetListOfFunctions().FindObject("POT")
                        mc_pot = mc_plot.GetListOfFunctions().FindObject("POT")
                        assert hasattr(data_pot, "GetVal"), "Couldn't get POT from data plot: '%s'" % data_plot.GetName()
                        assert hasattr(mc_pot, "GetVal"), "Could't get POT values from MC plot: '%s'" % mc_plot.GetName() 
                        scale_factor = data_pot.GetVal() / mc_pot.GetVal()

                assert scale_factor is not None, "Could not understand your MC normalization factor: %s" % mc_norm_factor
                
                mc_plot.Scale(scale_factor)
                
        # this will cover most cases.
        # however, if you need especially fine-grained control,
        # you can define a Make_<your plot name>_Plot(self)
        # in which you can do whatever you want
        def Make_Default_Plot_Method(self, plot_name, plot_path, mc_norm_factor):
                def f(self):
                        canvas = self.Make_DataMC_Overlay(plot_path)
                        
                        # normalize...
                        stack = DataMCPlots.GetMCDataStack(canvas)
                        mc_plot = stack.GetHists()[0]
                        data_plot = stack.GetHists()[1]
                        
                        DataMCPlots.NormalizeMCPlot(data_plot, mc_plot, mc_norm_factor)
        
                        self.plots_out[plot_name] = canvas
                
                setattr(self, "Make_%s_Plot" % plot_name, types.MethodType(f, self))
                
                return True
        
        def Make_MichelPID_Plot(self):
                # don't use the stock method because the legend winds up overlapping with the distribution
                michel_config = [config for config in COMBINATION_PLOTS if config.name == "MichelPID"][0]
                canvas = self.Make_DataMC_Overlay(michel_config.path)
                leg = [o for o in canvas.GetListOfPrimitives() if isinstance(o, ROOT.TLegend)][0]
                leg.SetX1(0.2)
                leg.SetX2(0.5)
                
                stack = DataMCPlots.GetMCDataStack(canvas)
                mc_plot = stack.GetHists()[0]
                data_plot = stack.GetHists()[1]
                
                DataMCPlots.NormalizeMCPlot(data_plot, mc_plot, michel_config.mc_norm_factor)

                self.plots_out[michel_config.name] = canvas
                
                        
        def SavePlots(self):
                c = ROOT.TCanvas()
                for name, plotobj in self.plots_out.iteritems():
                        c.Clear()
                        plotobj.Draw()
                        obj_to_print = plotobj if isinstance(plotobj, ROOT.TCanvas) else c
                        for out_type in OUTPUT_TYPES:
                                obj_to_print.Print( os.path.join(OUT_LOCATION, "data+MC", "MichelEPlots", "%s.%s" % (name, out_type)) )
        
################################
NTUPLE_FILELISTS = {
        "test": "/nova/ana/users/jwolcott/ntuple_lists/test_list_MC.txt",
        "data": "/nova/ana/users/jwolcott/art/data/std_candles/playlist-jwolcott_data_ND_FA_small-20151208.txt",
        "MC":   "/nova/ana/users/jwolcott/art/mc/std_candles/playlist-jwolcott_MC_ND_FA_small-20151208.txt",
        "PC":   "/nova/ana/users/jwolcott/art/mc/pc_muplus_extraEM/playlist-jwolcott_pc_muplus_extraEM_small-20160108.txt",
}
        
if __name__ == "__main__":
        assert len(sys.argv) == 2 and (sys.argv[1] in NTUPLE_FILELISTS or sys.argv[1] == "comp"), \
                "This script must be given exactly one argument (the name of the ntuple list to use).\nAvailable options: %s" % NTUPLE_FILELISTS.keys()

        if sys.argv[1] == "comp":
                pl = DataMCPlots()
                pl.MakePlots()
                pl.SavePlots()
        else:
                NTUPLE_FILELIST = NTUPLE_FILELISTS[sys.argv[1]]
                pl = StdCandlesPlotter(testing=TESTING)
                pl.LoopTree()
                pl.SavePlots()