Remove hardcoded condor/ directory in processing scripts

StandardAnalysis/python/plotUtilities.py

@@ -2,12 +2,45 @@
 #       configuration causes a segfault when saving skims. The ROOT modules seem to not play nicely
 #       with CMSSW, so be sure to only use this for plot-making scripts.
 
+import warnings
+
 from OSUT3Analysis.Configuration.histogramUtilities import *
 from DisappTrks.StandardAnalysis.tdrstyle import *
 from DisappTrks.StandardAnalysis.utilities import *
 
 from ctypes import c_double # see https://root-forum.cern.ch/t/issue-with-using-integralanderror-with-pyroot/53182/2
 
+
+def get_hist(file_path, channel, hist):
+    """Retrieve a histogram from a ROOT file.
+
+    Args:
+        file_path: Full path to the ROOT file.
+        channel: Directory path within the ROOT file.
+        hist: Name of the histogram.
+
+    Returns:
+        A cloned histogram detached from the file.
+
+    Raises:
+        FileNotFoundError: If the ROOT file cannot be opened.
+        KeyError: If the histogram is not found in the file.
+    """
+    input_file = TFile(file_path)
+    if not input_file or input_file.IsZombie():
+        raise FileNotFoundError(f"Could not open ROOT file: {file_path}")
+
+    hist_path = f"{channel}/{hist}"
+    h0 = input_file.Get(hist_path)
+    if not h0:
+        input_file.Close()
+        raise KeyError(f"Histogram '{hist_path}' not found in file: {file_path}")
+
+    h = h0.Clone()
+    h.SetDirectory(0)
+    input_file.Close()
+    return h
+
 def setStyle(h, color = 1):
     h.SetLineColor(color)
     h.SetLineStyle(1)
@@ -54,72 +87,172 @@ def setAxisStyle(h, xTitle = "", yTitle = "", xRange = (0, 0), yRange = (0, 0)):
         h.GetYaxis().SetRangeUser(yRange[0], yRange[1])
 
 def getHist(sample, condor_dir, channel, hist, quietFailure = False):
-    dataset_file = "condor/%s/%s.root" % (condor_dir,sample)
-    inputFile = TFile(dataset_file)
-    h0 = inputFile.Get(channel + "/" + hist)
-    if not h0:
-        if not quietFailure: print("ERROR [getHist]: didn't find histogram ", channel+str("/")+hist, "in file", dataset_file)
+    """Deprecated: Use get_hist(file_path, channel, hist) instead."""
+    warnings.warn(
+        "getHist is deprecated and will be removed in a future release. "
+        "Use get_hist(file_path, channel, hist) instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
+    dataset_file = "condor/%s/%s.root" % (condor_dir, sample)
+    try:
+        return get_hist(dataset_file, channel, hist)
+    except (FileNotFoundError, KeyError) as e:
+        if not quietFailure:
+            print(f"ERROR [getHist]: {e}")
         return 0
-    h = h0.Clone()
-    h.SetDirectory(0)
-    inputFile.Close()
-    return h
 
-def getHistFromProjectionZ(sample, condor_dir, channel, hist, fiducialElectronSigmaCut, fiducialMuonSigmaCut, alternate1DHist = "", verbose = False):
-    countProjections = 0 if not hasattr (getHistFromProjectionZ, "countProjections") else getattr (getHistFromProjectionZ, "countProjections")
-    h3d = getHist (sample, condor_dir, channel, hist, quietFailure = True)
+def get_hist_from_projection_z(file_path, channel, hist, fiducial_electron_sigma_cut, fiducial_muon_sigma_cut, alternate_1d_hist="", verbose=False):
+    """Retrieve a histogram from a 3D projection with fiducial cuts.
+
+    Args:
+        file_path: Full path to the ROOT file.
+        channel: Directory path within the ROOT file.
+        hist: Name of the 3D histogram.
+        fiducial_electron_sigma_cut: Upper bound for electron sigma cut.
+        fiducial_muon_sigma_cut: Upper bound for muon sigma cut.
+        alternate_1d_hist: Fallback 1D histogram if 3D not found.
+        verbose: If True, print warning when using fallback.
+
+    Returns:
+        The projected histogram, or None if not found.
+
+    Raises:
+        KeyError: If neither the 3D histogram nor alternate is found.
+    """
+    count_projections = 0 if not hasattr(get_hist_from_projection_z, "count_projections") else getattr(get_hist_from_projection_z, "count_projections")
+
+    try:
+        h3d = get_hist(file_path, channel, hist)
+    except (FileNotFoundError, KeyError):
+        h3d = None
+
     if not h3d:
-        h = None
-        if alternate1DHist:
-            if verbose: print("WARNING: not applying fiducial cuts via projections.")
-            h = getHist (sample, condor_dir, channel, alternate1DHist)
-        return h
-
-    h = h3d.ProjectionZ (hist + "_pz" + str(countProjections),
-                         0, h3d.GetXaxis().FindBin(fiducialElectronSigmaCut) - 1,
-                         0, h3d.GetYaxis().FindBin(fiducialMuonSigmaCut) - 1, "e")
-
-    countProjections += 1
-    setattr (getHistFromProjectionZ, "countProjections", countProjections)
+        if alternate_1d_hist:
+            if verbose:
+                print("WARNING: not applying fiducial cuts via projections.")
+            return get_hist(file_path, channel, alternate_1d_hist)
+        return None
+
+    h = h3d.ProjectionZ(hist + "_pz" + str(count_projections),
+                        0, h3d.GetXaxis().FindBin(fiducial_electron_sigma_cut) - 1,
+                        0, h3d.GetYaxis().FindBin(fiducial_muon_sigma_cut) - 1, "e")
+
+    count_projections += 1
+    setattr(get_hist_from_projection_z, "count_projections", count_projections)
     return h
 
-def getHistIntegralFromProjectionZ(sample, condor_dir, channel, fiducialElectronSigmaCut, fiducialMuonSigmaCut):
+
+def getHistFromProjectionZ(sample, condor_dir, channel, hist, fiducialElectronSigmaCut, fiducialMuonSigmaCut, alternate1DHist="", verbose=False):
+    """Deprecated: Use get_hist_from_projection_z(file_path, ...) instead."""
+    warnings.warn(
+        "getHistFromProjectionZ is deprecated and will be removed in a future release. "
+        "Use get_hist_from_projection_z(file_path, channel, hist, ...) instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
+    dataset_file = "condor/%s/%s.root" % (condor_dir, sample)
+    try:
+        return get_hist_from_projection_z(dataset_file, channel, hist, fiducialElectronSigmaCut, fiducialMuonSigmaCut, alternate_1d_hist=alternate1DHist, verbose=verbose)
+    except (FileNotFoundError, KeyError) as e:
+        print(f"ERROR [getHistFromProjectionZ]: {e}")
+        return None
+
+
+def get_hist_integral_from_projection_z(file_path, channel, fiducial_electron_sigma_cut, fiducial_muon_sigma_cut):
+    """Get the yield from a 3D histogram projection with fiducial cuts.
+
+    Args:
+        file_path: Full path to the ROOT file.
+        channel: Directory path within the ROOT file.
+        fiducial_electron_sigma_cut: Upper bound for electron sigma cut.
+        fiducial_muon_sigma_cut: Upper bound for muon sigma cut.
+
+    Returns:
+        Tuple of (yield, statistical_error).
+
+    Raises:
+        KeyError: If the histogram is not found.
+    """
     hist = "Track-met Plots/metNoMuMinusOnePtVsMaxSigmaForFiducialTracks"
-    h = getHistFromProjectionZ (sample, condor_dir, channel, hist, fiducialElectronSigmaCut, fiducialMuonSigmaCut, alternate1DHist = "Met Plots/metNoMu")
-    statError_ = c_double(0.0) # changed from Double to c_double; see https://root-forum.cern.ch/t/issue-with-using-integralanderror-with-pyroot/53182/2
+    h = get_hist_from_projection_z(file_path, channel, hist, fiducial_electron_sigma_cut, fiducial_muon_sigma_cut, alternate_1d_hist="Met Plots/metNoMu")
+    statError_ = c_double(0.0)
     yield_ = h.IntegralAndError(0, -1, statError_)
     return (yield_, statError_.value)
 
-def getYield(sample,condor_dir,channel):
-    dataset_file = "condor/%s/%s.root" % (condor_dir,sample)
-    inputFile = TFile(dataset_file)
-    metPtHistogram = inputFile.Get(channel + "/Met Plots/metPt")
-    if not metPtHistogram:
-        print("ERROR: didn't find histogram ", channel+str("/Met Plots/metPt"), "in file ", dataset_file)
-        return 0
-    statError_ = c_double(0.0) # changed from Double to c_double; see https://root-forum.cern.ch/t/issue-with-using-integralanderror-with-pyroot/53182/2
-    yield_     = float(metPtHistogram.IntegralAndError (0, metPtHistogram.GetNbinsX () + 1, statError_))
 
-    inputFile.Close()
+def getHistIntegralFromProjectionZ(sample, condor_dir, channel, fiducialElectronSigmaCut, fiducialMuonSigmaCut):
+    """Deprecated: Use get_hist_integral_from_projection_z(file_path, ...) instead."""
+    warnings.warn(
+        "getHistIntegralFromProjectionZ is deprecated and will be removed in a future release. "
+        "Use get_hist_integral_from_projection_z(file_path, channel, ...) instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
+    dataset_file = "condor/%s/%s.root" % (condor_dir, sample)
+    try:
+        return get_hist_integral_from_projection_z(dataset_file, channel, fiducialElectronSigmaCut, fiducialMuonSigmaCut)
+    except (FileNotFoundError, KeyError) as e:
+        print(f"ERROR [getHistIntegralFromProjectionZ]: {e}")
+        return (0, 0)
+
+def get_yield(file_path, channel, hist):
+    """Get the yield (integral and error) of a histogram.
+
+    Args:
+        file_path: Full path to the ROOT file.
+        channel: Directory path within the ROOT file.
+        hist: Name of the histogram.
+
+    Returns:
+        Tuple of (yield, statistical_error).
+
+    Raises:
+        FileNotFoundError: If the ROOT file cannot be opened.
+        KeyError: If the histogram is not found in the file.
+    """
+    h = get_hist(file_path, channel, hist)
+    statError_ = c_double(0.0)
+    yield_ = float(h.IntegralAndError(0, h.GetNbinsX() + 1, statError_))
     return (yield_, statError_.value)
 
+
+def getYield(sample, condor_dir, channel):
+    """Deprecated: Use get_yield(file_path, channel, hist) instead."""
+    warnings.warn(
+        "getYield is deprecated and will be removed in a future release. "
+        "Use get_yield(file_path, channel, hist) instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
+    dataset_file = "condor/%s/%s.root" % (condor_dir, sample)
+    try:
+        return get_yield(dataset_file, channel, "Met Plots/metPt")
+    except (FileNotFoundError, KeyError) as e:
+        print(f"ERROR [getYield]: {e}")
+        return 0
+
 def getHistFromChannelDict(channel, hist, quietFailure = False):
+    """Deprecated: Use get_hist(file_path, channel, hist) instead."""
+    warnings.warn(
+        "getHistFromChannelDict is deprecated and will be removed in a future release. "
+        "Use get_hist(file_path, channel, hist) instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
     if "sample" not in channel or "condorDir" not in channel or "name" not in channel:
-        print("Bad channel given to getHistFromChannelDict: " + str(channel))
-        return
+        raise ValueError(f"Bad channel given to getHistFromChannelDict: {channel}")
     dataset_file = "condor/%s/%s.root" % (channel["condorDir"], channel["sample"])
-    inputFile = TFile(dataset_file)
-    h0 = inputFile.Get(channel["name"] + "Plotter/" + hist)
-    if not h0:
-        if not quietFailure: print("ERROR [quietFailure]: didn't find histogram ", channel["name"] + str("Plotter/") + hist, "in file", dataset_file)
+    channel_name = channel["name"] + "Plotter"
+    try:
+        return get_hist(dataset_file, channel_name, hist)
+    except (FileNotFoundError, KeyError) as e:
+        if not quietFailure:
+            print(f"ERROR [getHistFromChannelDict]: {e}")
         return 0
-    h = h0.Clone()
-    h.SetDirectory(0)
-    inputFile.Close()
-    return h
 
 def addChannelExtensions(histogram, channel, histName):
     if "extensions" in channel:
         for x in channel["extensions"]:
             histogram.Add (getHistFromChannelDict (x, histName))
-            
+