def print_evt_info_cjlst(tree):
    """A goofy way to print branch info for `evt` in TTree.
    
    TODO:
    - [ ] Print only up to 6 decimals for all floats.
    """
    # d_branch = {
    #     "LepPt" : "list",
    #     "LepEta" : "list",
    #     "LepLepId" : "list",
    #     "LepisID" : "array",
    #     "LepCombRelIsoPF" : "list",
    #     "CRflag" : "list",
    #     "Z1Mass" : "list",
    #     "Z2Mass" : "list",
    #     "ZZMass" : "list",
    # }
    # for branch, express_as in d_branch.items():
    #     if express_as == "":
    #         try:
    #             print(f"{branch}: {getattr(evt, branch)}")
    #         except AttributeError:
    #             # Branch doesn't exist.
    #             pass
    #     elif express_as == "list":
    #         try:
    #             print(f"{branch}: {list(getattr(evt, branch))}")
    #         except AttributeError:
    #             # Branch doesn't exist.
    #             pass

    announce("Analyzer: CJLST")

    raise RuntimeWarning(
        f"Modify this func to include proper handling of tight CJLST lep\n"
        f"Proper way:\n"
        f"  lep_pdgID = t.LepLepId[idx]\n"
        f"  lep_is_tightID = bool(t.LepisID[idx])\n"
        f"  lep_iso = t.LepCombRelIsoPF[idx]\n"
        f"  For electrons, only require: lep_is_tightID=True\n"
        f"  For muons, require: lep_is_tightID=True AND lep_iso < 0.35")
    print(
        f"tree.LepPt: {list(tree.LepPt)}\n"
        # f"tree.fsrPt: {list(tree.fsrPt)}\n"  # pT of FSR photons.
        f"tree.LepEta: {list(tree.LepEta)}\n"
        f"tree.LepLepId: {list(tree.LepLepId)}\n"
        f"tree.LepisLoose: {list(np.array(tree.LepisLoose, dtype=bool))}\n"
        f"tree.LepisID (tight ID): {list(np.array(tree.LepisID, dtype=bool))}\n"
        f"tree.LepCombRelIsoPF: {list(tree.LepCombRelIsoPF)}\n"
        f"tree.CRflag: {tree.CRflag} -> {CjlstFlag(tree.CRflag).name}\n"
        f"tree.Z1Mass: {tree.Z1Mass}\n"
        f"tree.Z2Mass: {tree.Z2Mass}\n"
        f"tree.ZZMass: {tree.ZZMass}\n")
Exemple #2
0
def print_totalredbkg_estimate(
    h1_data_3p1fpred_m4l,
    h1_zz_3p1fpred_m4l,
    d_data_3p1fpred_fs_hists,
    d_data_2p2fpred_fs_hists,
    d_zz_3p1fpred_fs_hists,
    neg_bin_removal=True
    ):
    """Prints RedBkg estimate before and after negative bin removal.
    
    estimate = (3P1F - 2P2F -ZZ)
    """
    msg = "(3P1F - 2P2F - ZZ) Data Pred: before neg bin removal"
    if neg_bin_removal:
        msg = msg.replace("before", "AFTER")

    announce(msg)
    h1_final_estimate = h1_data_3p1fpred_m4l.Clone()
    h1_final_estimate.Add(h1_data_2p2fpred_m4l, -1)
    print(f"total integral (3P1F - 2P2F) = {h1_final_estimate.Integral()}")
    h1_final_estimate.Add(h1_zz_3p1fpred_m4l, -1)
    if neg_bin_removal:
        h1_final_estimate = set_neg_bins_to_zero(h1_final_estimate)
        print(f"total integral (3P1F - 2P2F - ZZ) negbinremoval = {h1_final_estimate.Integral()}")
    else:
        print(f"total integral (3P1F - 2P2F - ZZ) = {h1_final_estimate.Integral()}")

    # Print each final state.
    for fs, h_3p1f, h_2p2f, h_3p1f_zz in zip(
            d_data_3p1fpred_fs_hists.keys(),
            d_data_3p1fpred_fs_hists.values(),
            d_data_2p2fpred_fs_hists.values(),
            d_zz_3p1fpred_fs_hists.values(),
            ):
        h1_final_estimate_fs = h_3p1f.Clone()
        h1_final_estimate_fs.Add(h_2p2f.Clone(), -1)
        h1_final_estimate_fs.Add(h_3p1f_zz.Clone(), -1)
        h1_final_estimate_fs_rmv = h1_final_estimate_fs.Clone()
        if neg_bin_removal:
            print("--- Negative bin removal ---")
            h1_final_estimate_fs_rmv = set_neg_bins_to_zero(h1_final_estimate_fs)
        print(
            f"finalstate = {fs}, "
            f"integral = {h1_final_estimate_fs_rmv.Integral():.2f}"
            )
Exemple #3
0
def print_evt_info_bbf(evt):
    """A goofy way to print branch info for `evt` in TTree.
    
    TODO:
    - [ ] Print only up to 6 decimals for all floats.
    """
    announce("Analyzer: BBF")
    d_branch = {
        "passedFullSelection": "",
        "passedZXCRSelection": "",
        "nZXCRFailedLeptons": "",
        "lep_Hindex": "list",
        "lep_RedBkgindex": "list",
        "lep_id": "list",
        "lep_pt": "list",
        "lep_eta": "list",
        "lep_phi": "list",
        "lep_mass": "list",
        "lepFSR_pt": "list",
        "lepFSR_eta": "list",
        "lepFSR_phi": "list",
        "lepFSR_mass": "list",
        "vtxLepFSR_BS_pt": "list",
        # "lep_RelIso" : "list",
        "lep_RelIsoNoFSR": "list",
        "lep_tightId": "list",
        "is2P2F": "",
        "is3P1F": "",
    }
    for branch, express_as in d_branch.items():
        if express_as == "":
            try:
                print(f"evt.{branch}: {getattr(evt, branch)}")
            except AttributeError:
                # Branch doesn't exist.
                pass
        elif express_as == "list":
            try:
                print(f"evt.{branch}: {list(getattr(evt, branch))}")
            except AttributeError:
                # Branch doesn't exist.
                pass
    print()
Exemple #4
0
def make_all_zcands(
    mylep_ls,
    explain_skipevent=False, verbose=False
    ):
    """Return list of valid Z candidates as MyZboson objects.
    
    This function DOES apply cuts to each Z!

    To form a valid Z candidate:
    - Leptons DO NOT HAVE TO be tight (loose+tightID+RelIso)!
    - Leptons must be OSSF.
    - 12 < m(ll, including FSR) < 120 GeV.

    NOTE: Also stores indices of Z as it appears in list of Z.
    """
    zcand_ls = []
    ndx_zvec = 0
    # Make a lep-by-lep comparison to find eligible Z candidates.
    for ndx_lep1, mylep1 in enumerate(mylep_ls[:-1]):
        start_ndx_lep2 = ndx_lep1 + 1
        for ndx_lep2, mylep2 in enumerate(
                mylep_ls[start_ndx_lep2:], start_ndx_lep2
                ):
            these_lep_idcs = [mylep1.ndx_lepvec, mylep2.ndx_lepvec]
            if verbose:
                announce(
                    f"  Considering leptons: {these_lep_idcs}",
                    n_center_pad_chars=1
                    )
            if not makes_valid_zcand(mylep1, mylep2, verbose=verbose):
                if verbose: print(f"    Failed to make valid Z cand.\n")
                continue
            # Found valid Z candidate.
            if verbose: print(f"    Made valid Z cand.\n")
            zcand = MyZboson(mylep1, mylep2, explain_skipevent=explain_skipevent)
            zcand.ndx_zcand_ls = ndx_zvec
            ndx_zvec += 1
            zcand_ls.extend((zcand,))
    return zcand_ls
Exemple #5
0
def analyze_single_evt(
    tree,
    run=None,
    lumi=None,
    event=None,
    entry=None,
    fw="bbf",
    which="all",
    evt_start=0,
    evt_end=-1,
    print_every=10000,
    infile_fakerates=None,
    genwgts_dct=None,
    dct_xs=None,
    LUMI_INT=59830,
    smartcut_ZapassesZ1sel=False,
    overwrite=False,
    skip_mass4l_lessthan0=False,
    match_lep_Hindex=False,
    recalc_masses=False,
    skip_passedFullSelection=True,
    stop_when_found_3p1f=True,
    keep_one_quartet=False,
    explain_skipevent=True,
    verbose=True,
):
    """Return list of event numbers that correspond to run, lumi, event.

    Event number in this case refers to its entry (row index) in `tree`.
    
    Also prints out event info (`run`:`lumi`:`event`) found in `tree`.
    
    Args:
        entry (int):
            Row in TTree.
        fw (str):
            Which framework to use: "bbf", "cjlst", "jake"
        which (str):
            Which instance of the event you want to select.
            Options: "first", anything else collects all such events.
        evt_start (int):
    """
    know_evtid = all(x is not None for x in (run, lumi, event))
    know_entry = entry is not None
    if (not know_entry) and know_evtid:
        # Specific entry (row in TTree) was not specified.
        # Use Run, Lumi, Event to find all, entries.
        assert all(x is not None for x in (run, lumi, event))
        if evt_end == -1:
            evt_end = tree.GetEntries()
        ls_entries = find_entries_using_runlumievent(tree,
                                                     run,
                                                     lumi,
                                                     event,
                                                     evt_start=evt_start,
                                                     evt_end=evt_end,
                                                     print_every=print_every,
                                                     which=which,
                                                     fw=fw)
    elif know_entry and (not know_evtid):
        # We know the entry, but not the Run, Lumi, Event.
        run, lumi, event = find_runlumievent_using_entry(tree=tree,
                                                         entry=entry,
                                                         fw=fw)
        ls_entries = [entry]
    elif know_entry and know_evtid:
        ls_entries = [entry]
    else:
        raise ValueError(f"Must specify either `entry` or `run, lumi, event`")

    for entry in ls_entries:
        print(f"Analyzing event #{entry} in {fw.upper()} framework"
              f" (event ID {run}:{lumi}:{event}).")

        tree.GetEntry(entry)
        if fw == "bbf":
            print_evt_info_bbf(tree)
        elif fw == "jake":
            announce("ANALYZER: Jake")
            select_evts_2P2F_3P1F_multiquartets(
                tree,
                infile_fakerates=infile_fakerates,
                genwgts_dct=genwgts_dct,
                dct_xs=dct_xs,
                outfile_root=None,
                outfile_json=None,
                name="Data",
                int_lumi=LUMI_INT,
                start_at_evt=entry,
                break_at_evt=entry + 1,
                fill_hists=False,
                explain_skipevent=explain_skipevent,
                verbose=verbose,
                print_every=1,
                smartcut_ZapassesZ1sel=False,
                overwrite=False,
                skip_mass4l_lessthan0=False,
                match_lep_Hindex=match_lep_Hindex,
                recalc_masses=False,
                skip_passedFullSelection=skip_passedFullSelection,
                stop_when_found_3p1f=stop_when_found_3p1f,
                keep_one_quartet=keep_one_quartet,
            )
            # store_evt = True
        elif fw == "cjlst":
            print_evt_info_cjlst(tree)

    # End loop over collected ls_entries.
    return ls_entries
Exemple #6
0
        (1, 2, 3, 4),
        (50, 50, 50, 50)
        ):
        prettyup_and_drawhist(d_data_2p2fpred_fs_hists[fs], printer, outpdf_path, y_max=y_max)

    # Need to draw 2P2F contribution to 3P1F.
    # prettyup_and_drawhist(h1_data_2p2fpred_m4l, printer, outpdf_path, y_max=120)
    # for fs, y_max in zip(
    #     (1, 2, 3, 4),
    #     (50, 50, 50, 50)
    #     ):
    #     prettyup_and_drawhist(d_data_2p2fpred_fs_hists[fs], printer, outpdf_path, y_max=y_max)

    printer.canv.Print(outpdf_path + "]")

    announce("3P1F Raw Data")
    # print(f"integral = {h1_data_3p1f_m4l.Integral():.2f}")
    print_integral_dict_hist(d_data_3p1f_fs_hists)

    announce("2P2F Raw Data")
    # print(f"integral = {h1_data_2p2f_m4l.Integral():.2f}")
    print_integral_dict_hist(d_data_2p2f_fs_hists)

    announce("3P1F Data Pred")
    # print(f"integral = {h1_data_3p1fpred_m4l.Integral():.2f}")
    print_integral_dict_hist(d_data_3p1fpred_fs_hists)

    announce("2P2F Data Pred")
    # print(f"integral = {h1_data_2p2fpred_m4l.Integral():.2f}")
    print_integral_dict_hist(d_data_2p2fpred_fs_hists)
Exemple #7
0
    def get_best_ZZcand_per_quart(self,
                                  mylep_ls,
                                  cr,
                                  verbose=False,
                                  explain_skipevent=False,
                                  smartcut_ZapassesZ1sel=False,
                                  run=None,
                                  lumi=None,
                                  event=None,
                                  entry=None):
        """Return a list of best ZZ cands that pass OS method selections.
        
        In the returned list:
            ls_valid_zz_cand_OS = [
                best_ZZ_quart1,
                best_ZZ_quart2,
                ...
                ]

        Leptons are sorted into quartets based on given control region `cr`.

        TODO: Add wrong charge/flavor functionality.

        Args:
            mylep_ls (list of MyLeptons):
                Leptons from only 1 event. Not limited to just 4 leptons.
                Will be sorted into different quartets based on methods
                (OS, WCF) and CRs (3P1F, 2P2F).
            cr (str): Control region ('3P1F' or '2P2F').

        Returns:
            list: Valid ZZ candidates that pass OS Method selections.
        """
        cr = cr.upper()
        if cr == '3P1F':
            ls_quartets = find_quartets_3pass1fail(mylep_ls)
        elif cr == '2P2F':
            ls_quartets = find_quartets_2pass2fail(mylep_ls)

        n_tot_combos = len(ls_quartets)
        if n_tot_combos == 0:
            # No chance to make quartets.
            return []

        if verbose:
            print(f"  Num of {cr} quartets to analyze: {n_tot_combos}")
        # List to hold valid ZZ cands that pass OS Method sel.
        ls_valid_zz_cand_OS = []
        for n_quartet, quart in enumerate(ls_quartets):
            if verbose:
                announce(f"Analyzing quartet #{n_quartet}")
                lep_ndcs = [mylep.ndx_lepvec for mylep in quart]
                print(f"  Lepton indices: {lep_ndcs}")
            # For this quartet, use OS Method logic to pick best ZZ cand.
            ls_zzcand = get_ZZcands_from_myleps_OSmethod(
                quart,
                verbose=verbose,
                explain_skipevent=explain_skipevent,
                smartcut_ZapassesZ1sel=smartcut_ZapassesZ1sel,
                run=run,
                lumi=lumi,
                event=event,
                entry=entry)

            if len(ls_zzcand) == 0:
                # No good ZZ candidate found.
                continue

            ls_valid_zz_cand_OS.append(ls_zzcand[0]  # Best ZZ cand.
                                       )
        return ls_valid_zz_cand_OS
Exemple #8
0
def select_evts_2P2F_3P1F_multiquartets(
    tree,
    infile_fakerates,
    genwgts_dct,
    dct_xs,
    outfile_root=None,
    outfile_json=None,
    name="",
    int_lumi=-1,
    start_at_evt=0,
    break_at_evt=-1,
    fill_hists=False,
    explain_skipevent=False,
    verbose=False,
    print_every=50000,
    smartcut_ZapassesZ1sel=False,
    overwrite=False,
    skip_mass4l_lessthan0=False,
    match_lep_Hindex=False,
    recalc_masses=False,
    skip_passedFullSelection=True,
    stop_when_found_3p1f=True,
    keep_one_quartet=False,
    use_multiquart_sel=True,
    sync_with_xBFAna=False,
):
    """Apply RedBkg multi-lepton quartet selection to all events in tree.

    NOTE:
        A "quartet" is a combination of 4-leptons.
        Each event may have multiple leptons quartets.
        If an event has both a 3P1F quartet and a 2P2F quartet,
        then the 3P1F one takes priority. The 2P2F quartet will be skipped.

    TODO Update below:
    Select events with:
        - Exactly 3 leptons passing tight selection and at least 1 failing.
        - Exactly 2 leptons passing tight selection and at least 2 failing.
        - When an event has >4 leptons, then multiple combinations
            of 2P2F/3P1F are possible.
            Suppose you have an event with 5 leptons, 3 of which pass tight
            selection and 2 fail tight selection.
            Then we have two different ways to make a 3P1F combo
            (two 3P1F subevents).
        - Does NOT select the BEST ZZ candidate per event.
            If there are >1 ZZ candidates that pass ZZ selections,
            then each ZZ cand is saved as a separate entry in the TTree.
            This means a single event can show up multiple times in a TTree,
            if there are multiple valid 4-lepton combinations ("quartets").
    
    Args:
        outfile_root (str):
            Path to store root file.
            TTree of selected events will be made.
            If `fill_hists` is True, then will store histograms.
        smartcut_ZapassesZ1sel (bool, optional):
            In the smart cut, the literature essentially says that if a Za
            looks like a more on-shell Z boson than the Z1 AND if the Zb is a
            low mass di-lep resonance, then veto the whole ZZ candidate.
            However, literature doesn't check for Za passing Z1 selections!
            Set this to True if you require Za to pass Z1 selections.
            Default is False.
        skip_mass4l_lessthan0 (bool, optional):
            If True, then skip events whose tree.mass4l <= 0.
            This is useful when you need to use the values that are already
            stored in the BBF NTuple.
            Default is False.
        match_lep_Hindex (bool, optional):
            Useful for reproducing results from the xBF Analyzer.
            If True, then only save events in which the selected quartet has:
                Z1 lepton indices that match:
                    lep_Hindex[0], lep_Hindex[1] (any order)
                Z2 lepton indices that match:
                    lep_Hindex[2], lep_Hindex[3] (any order)
            Since all quartets are automatically built (not necessarily saved)
            then this will find all the same events that xBF finds.
            It is "cheating" since we know the answer (lep_Hindex) already!
            Default is False.
        recalc_masses (bool, optional):
            If True, recalculate mass4l of selected lepton quartet and
            corresponding massZ1 and massZ2 values.
        stop_when_found_3p1f (bool, optional):
            If True, if at least one valid 3P1F ZZ candidate was found,
            do not build any 2P2F candidates. Defaults to True.
        use_multiquart_sel (bool, optional):
            If True, use the updated reducible background event selection
            logic ("multi-quartet").

    """
    assert not (use_multiquart_sel and sync_with_xBFAna), (
        f"Can't have use_multiquart_sel=True and sync_with_xBFAna=True.")

    if use_multiquart_sel:
        msg = f"use_multiquart_sel = True. Forcing these bools:"
        announce(msg, pad_char='#')
        print(f"  stop_when_found_3p1f = True\n"
              f"  match_lep_Hindex = False\n"
              f"  keep_one_quartet = False\n"
              f"  recalc_masses = True\n"
              f"  skip_mass4l_lessthan0 = False\n"
              f"  skip_passedFullSelection = True\n")
        stop_when_found_3p1f = True
        match_lep_Hindex = False
        keep_one_quartet = False
        recalc_masses = True
        skip_mass4l_lessthan0 = False
        skip_passedFullSelection = True

    if sync_with_xBFAna:
        msg = f"sync_with_xBFAna = True. Forcing these bools:"
        announce(msg, pad_char='#')
        print(f"  stop_when_found_3p1f = False\n"
              f"  match_lep_Hindex = True\n"
              f"  keep_one_quartet = True\n"
              f"  skip_mass4l_lessthan0 = True\n"
              f"  skip_passedFullSelection = True\n"
              f"  recalc_masses = False\n")
        stop_when_found_3p1f = False
        match_lep_Hindex = True
        keep_one_quartet = True
        skip_mass4l_lessthan0 = True
        skip_passedFullSelection = True
        recalc_masses = False

    if keep_one_quartet and not recalc_masses:
        warnings.warn(f"!!! You are keeping one quartet per event "
                      f"(keep_one_quartet=True) "
                      f"without recalculating masses(recalc_masses=False).\n"
                      f"!!! Therefore, mass info (mass4l, massZ1, massZ2) "
                      f"in tree may not correspond to saved quartet!")

    if stop_when_found_3p1f:
        print("A valid 3P1F quartet will ignore all 2P2F quartets.")

    if keep_one_quartet:
        announce("Saving only one quartet per event.")
        if match_lep_Hindex:
            print(f"  Since match_lep_Hindex=True, the one quartet saved\n"
                  f"  must have lepton indices from Z1 and Z2 agree with\n"
                  f"  those in lep_Hindex.")
            if not stop_when_found_3p1f:
                print("  A valid 2P2F quartet can be chosen over valid 3P1F.")
        else:
            print(f"  Since match_lep_Hindex=False, "
                  f"not sure which quartet to save!\n"
                  f"  Saving the first valid ZZ cand arbitrarily!.")

    if fill_hists:
        raise ValueError(f"fill_hists = True, but must review code first.")
        # Prep histograms.
        d_hists = {
            "Data": {
                "2p2f": {
                    "mass4l": h1_data_2p2f_m4l,
                    "n_quartets": h1_data_n2p2f_combos,
                },
                "3p1f": {
                    "mass4l": h1_data_3p1f_m4l,
                    "n_quartets": h1_data_n3p1f_combos,
                }
            },
            "ZZ": {
                "2p2f": {
                    "mass4l": h1_zz_2p2f_m4l,
                    "n_quartets": h1_zz_n2p2f_combos,
                },
                "3p1f": {
                    "mass4l": h1_zz_3p1f_m4l,
                    "n_quartets": h1_zz_n3p1f_combos,
                }
            },
        }

    evt_info_d = make_evt_info_d()  # Info for printing.
    evt_info_2p2f_3p1f_d = {}  # Info for json file.

    assert infile_fakerates is not None
    h_FRe_bar, h_FRe_end, h_FRmu_bar, h_FRmu_end = \
        retrieve_FR_hists(infile_fakerates)

    if outfile_json is not None:
        check_overwrite(outfile_json, overwrite=overwrite)

    # Make pointers to store new values.
    ptr_finalState = np.array([0], dtype=int)
    ptr_nZXCRFailedLeptons = np.array([0], dtype=int)
    ptr_is2P2F = np.array([0], dtype=int)  # Close enough to bool lol.
    ptr_is3P1F = np.array([0], dtype=int)
    ptr_isData = np.array([0], dtype=int)
    ptr_isMCzz = np.array([0], dtype=int)
    ptr_fr2_down = array('f', [0.])
    ptr_fr2 = array('f', [0.])
    ptr_fr2_up = array('f', [0.])
    ptr_fr3_down = array('f', [0.])
    ptr_fr3 = array('f', [0.])
    ptr_fr3_up = array('f', [0.])
    ptr_eventWeightFR_down = array('f', [0.])
    ptr_eventWeightFR = array('f', [0.])
    ptr_eventWeightFR_up = array('f', [0.])
    ptr_lep_RedBkgindex = array('i', [0, 0, 0, 0])
    ptr_mass4l = array('f', [0.])
    ptr_massZ1 = array('f', [0.])
    ptr_massZ2 = array('f', [0.])
    # ptr_mass4l_vtxFSR_BS = array('f', [0.])
    # ptr_eventWeight = array('f', [0.])

    if outfile_root is not None:
        check_overwrite(outfile_root, overwrite=overwrite)
        new_file = TFile.Open(outfile_root, "recreate")
        print("Cloning TTree.")
        new_tree = tree.CloneTree(0)  # Clone 0 events.

        # Make new corresponding branches in the TTree.
        new_tree.Branch("is2P2F", ptr_is2P2F, "is2P2F/I")
        new_tree.Branch("is3P1F", ptr_is3P1F, "is3P1F/I")
        new_tree.Branch("isData", ptr_isData, "isData/I")
        new_tree.Branch("isMCzz", ptr_isMCzz, "isMCzz/I")
        new_tree.Branch("fr2_down", ptr_fr2_down, "fr2_down/F")
        new_tree.Branch("fr2", ptr_fr2, "fr2/F")
        new_tree.Branch("fr2_up", ptr_fr2_up, "fr2_up/F")
        new_tree.Branch("fr3_down", ptr_fr3_down, "fr3_down/F")
        new_tree.Branch("fr3", ptr_fr3, "fr3/F")
        new_tree.Branch("fr3_up", ptr_fr3_up, "fr3_up/F")
        new_tree.Branch("eventWeightFR_down", ptr_eventWeightFR_down,
                        "eventWeightFR_down/F")
        new_tree.Branch("eventWeightFR", ptr_eventWeightFR, "eventWeightFR/F")
        new_tree.Branch("eventWeightFR_up", ptr_eventWeightFR_up,
                        "eventWeightFR_up/F")
        # Record the indices of the leptons in passing quartet.
        new_tree.Branch("lep_RedBkgindex", ptr_lep_RedBkgindex,
                        "lep_RedBkgindex[4]/I")

        # Modify existing values of branches.
        new_tree.SetBranchAddress("finalState", ptr_finalState)
        new_tree.SetBranchAddress("nZXCRFailedLeptons", ptr_nZXCRFailedLeptons)
        if recalc_masses:
            new_tree.SetBranchAddress("mass4l", ptr_mass4l)
            new_tree.SetBranchAddress("massZ1", ptr_massZ1)
            new_tree.SetBranchAddress("massZ2", ptr_massZ2)

    n_tot = tree.GetEntries()
    print(f"Total number of events: {n_tot}\n"
          f"Looking for >=4 leptons per event...")

    ####################
    #=== Event Loop ===#
    ####################
    isMCzz = 1 if name in "ZZ" else 0
    isData = 1 if name in "Data" else 0
    for evt_num in range(start_at_evt, n_tot):
        if evt_num == break_at_evt:
            break

        print_periodic_evtnum(evt_num, n_tot, print_every=print_every)

        tree.GetEntry(evt_num)
        run = tree.Run
        lumi = tree.LumiSect
        event = tree.Event
        evt_id = f"{run} : {lumi} : {event}"

        ###################################
        #=== Initial event selections. ===#
        ###################################
        try:
            if not tree.passedTrig:
                continue
        except AttributeError:
            # Branch 'passedTrig' doesn't exist.
            warnings.warn(f"Branch passedTrig probably doesn't exist!\n"
                          f"Ignoring passedTrig==1 criterion.")

        if skip_passedFullSelection and tree.passedFullSelection:
            if explain_skipevent:
                print_skipevent_msg("passedFullSelection == 1", evt_num, run,
                                    lumi, event)
            evt_info_d["n_evts_passedFullSelection"] += 1
            continue

        if skip_mass4l_lessthan0 and (tree.mass4l <= 0):
            evt_info_d["n_evts_skip_mass4l_le0"] += 1
            continue

        # Check the number of leptons in this event.
        n_tot_leps = len(tree.lep_pt)
        if verbose:
            print(f"  Total number of leptons found: {n_tot_leps}")

        # Ensure at least 4 leptons in event:
        if n_tot_leps < 4:
            if explain_skipevent:
                print_skipevent_msg("n_leps < 4", evt_num, run, lumi, event)
            evt_info_d["n_evts_lt4_leps"] += 1
            continue

        # Initialize ALL leptons (possibly >=4 leptons).
        mylep_ls = make_filled_mylep_ls(tree)
        n_leps_passing = get_n_myleps_passing(mylep_ls)
        n_leps_failing = get_n_myleps_failing(mylep_ls)
        if verbose:
            print(f"    Num leptons passing tight sel: {n_leps_passing}\n"
                  f"    Num leptons failing tight sel: {n_leps_failing}")
            for mylep in mylep_ls:
                mylep.print_info(oneline=True)

        if n_leps_passing < 2:
            evt_info_d["n_evts_lt2tightleps"] += 1
            if explain_skipevent:
                msg = f"  Contains {n_leps_passing} (< 2) tight leps."
                print_skipevent_msg(msg, evt_num, run, lumi, event)
            continue

        ####################################################
        #=== Find best ZZ cand for each lepton quartet. ===#
        ####################################################
        qtcat = QuartetCategorizer(
            mylep_ls,
            verbose=verbose,
            explain_skipevent=explain_skipevent,
            smartcut_ZapassesZ1sel=smartcut_ZapassesZ1sel,
            run=run,
            lumi=lumi,
            event=event,
            entry=evt_num,
            stop_when_found_3p1f=stop_when_found_3p1f)

        ls_valid_ZZcands_OS = qtcat.ls_valid_ZZcands_OS_3p1f + \
                                qtcat.ls_valid_ZZcands_OS_2p2f
        n_valid_ZZcands_OS = qtcat.n_valid_ZZcands_OS_3p1f + \
                                qtcat.n_valid_ZZcands_OS_2p2f

        if n_valid_ZZcands_OS == 0:
            if explain_skipevent:
                print_skipevent_msg("  No valid 2P2F or 3P1F ZZ candidates.",
                                    evt_num, run, lumi, event)
            evt_info_d["n_evts_novalid2P2For3P1F_ZZcands"] += 1
            continue

        if verbose:
            print(f"  Num valid OS ZZ cands 3P1F: "
                  f"{qtcat.n_valid_ZZcands_OS_3p1f}\n"
                  f"  Num valid OS ZZ cands 2P2F: "
                  f"{qtcat.n_valid_ZZcands_OS_2p2f}")

        ##############################################################
        #=== Record each valid ZZ cand per quartet in this event. ===#
        ##############################################################
        n_saved_quartets_3p1f = 0
        n_saved_quartets_2p2f = 0
        found_matching_lep_Hindex = False
        overall_evt_label = ''  # Either '3P1F' or '2P2F'.
        for ndx_zzcand, zzcand in enumerate(ls_valid_ZZcands_OS, 1):
            cr_str = zzcand.get_str_cr_os_method().upper()

            # Sanity checks.
            if cr_str == '3P1F':
                assert zzcand.check_valid_cand_os_3p1f()
            elif cr_str == '2P2F':
                assert zzcand.check_valid_cand_os_2p2f()
            else:
                raise ValueError(f"cr_str is not in ('3P1F', '2P2F').")

            assert zzcand.check_valid_cand_os_3p1f() \
                    ^ zzcand.check_valid_cand_os_2p2f()  # xor.

            if match_lep_Hindex:
                # Only save quartet if the lepton indices match lep_Hindex.
                # E.g., the three lists below are the same ZZ cand:
                # quartet_1 = [2, 3, 1, 4]
                # quartet_2 = [3, 2, 1, 4]
                # quartet_3 = [3, 2, 4, 1]
                lep_Hindex_ls = list(tree.lep_Hindex)
                try:
                    assert len(lep_Hindex_ls) == 4
                except AssertionError:
                    print(f'lep_Hindex_ls has weird vals: {lep_Hindex_ls}')
                    AssertionError

                idcs_lepHindex_z1 = lep_Hindex_ls[:2]
                idcs_lepHindex_z2 = lep_Hindex_ls[2:]
                idcs_myz1 = zzcand.z_fir.get_mylep_indices()
                idcs_myz2 = zzcand.z_sec.get_mylep_indices()
                # When checking if two sets are equal, order doesn't matter!
                same_z1 = (set(idcs_lepHindex_z1) == set(idcs_myz1))
                same_z2 = (set(idcs_lepHindex_z2) == set(idcs_myz2))
                if same_z1 and same_z2:
                    found_matching_lep_Hindex = True
                else:
                    evt_info_d["n_quartets_skip_lep_Hindex_mismatch"] += 1
                    continue

            # Event weight calculation.
            n_dataset_tot = float(genwgts_dct[name])
            evt_weight_calcd = get_evt_weight(dct_xs=dct_xs,
                                              Nickname=name,
                                              lumi=int_lumi,
                                              event=tree,
                                              n_dataset_tot=n_dataset_tot,
                                              orig_evt_weight=tree.eventWeight)

            # See which leptons from Z2 failed.
            n_fail_code = check_which_Z2_leps_failed(zzcand)

            mylep1_fromz1 = zzcand.z_fir.mylep1
            mylep2_fromz1 = zzcand.z_fir.mylep2
            mylep1_fromz2 = zzcand.z_sec.mylep1
            mylep2_fromz2 = zzcand.z_sec.mylep2
            #=== 3P1F quartet. ===#
            if n_fail_code == 2:
                # First lep from Z2 failed.
                # NOTE: Turns out this will never trigger for OS Method!
                # This is due to the way that the 3P1F quartets are built.
                # The 1 failing lepton is always placed as the 4th lepton.
                # Therefore, fr3 (the 4th lepton) will always be != 0.
                assert cr_str == '3P1F'
                assert mylep1_fromz2.is_loose
                fr2, fr2_err = get_fakerate_and_error_mylep(mylep1_fromz2,
                                                            h_FRe_bar,
                                                            h_FRe_end,
                                                            h_FRmu_bar,
                                                            h_FRmu_end,
                                                            verbose=verbose)
                fr2_down = calc_fakerate_down(fr2, fr2_err)  # Scale down.
                fr2_up = calc_fakerate_up(fr2, fr2_err)  # Scale up.
                fr3 = 0
                fr3_err = 0
                fr3_down = 0
                fr3_up = 0
                # Use fake rates to calculate new event weight.
                new_weight_down = (fr2_down /
                                   (1 - fr2_down)) * evt_weight_calcd
                new_weight = (fr2 / (1 - fr2)) * evt_weight_calcd
                new_weight_up = (fr2_up / (1 - fr2_up)) * evt_weight_calcd
            elif n_fail_code == 3:
                # Second lep from Z2 failed.
                assert cr_str == '3P1F'
                assert mylep2_fromz2.is_loose
                fr2 = 0
                fr2_err = 0
                fr2_down = 0
                fr2_up = 0
                fr3, fr3_err = get_fakerate_and_error_mylep(mylep2_fromz2,
                                                            h_FRe_bar,
                                                            h_FRe_end,
                                                            h_FRmu_bar,
                                                            h_FRmu_end,
                                                            verbose=verbose)
                fr3_down = calc_fakerate_down(fr3, fr3_err)  # Scale down.
                fr3_up = calc_fakerate_up(fr3, fr3_err)  # Scale up.
                # Use fake rates to calculate new event weight.
                new_weight_down = (fr3_down /
                                   (1 - fr3_down)) * evt_weight_calcd
                new_weight = (fr3 / (1 - fr3)) * evt_weight_calcd
                new_weight_up = (fr3_up / (1 - fr3_up)) * evt_weight_calcd
            #=== 2P2F quartet. ===#
            elif n_fail_code == 5:
                # Both leps failed.
                assert cr_str == '2P2F'
                fr2, fr2_err = get_fakerate_and_error_mylep(mylep1_fromz2,
                                                            h_FRe_bar,
                                                            h_FRe_end,
                                                            h_FRmu_bar,
                                                            h_FRmu_end,
                                                            verbose=verbose)
                fr3, fr3_err = get_fakerate_and_error_mylep(mylep2_fromz2,
                                                            h_FRe_bar,
                                                            h_FRe_end,
                                                            h_FRmu_bar,
                                                            h_FRmu_end,
                                                            verbose=verbose)
                fr2_down = calc_fakerate_down(fr2, fr2_err)  # Scale down.
                fr2_up = calc_fakerate_up(fr2, fr2_err)  # Scale up.
                fr3_down = calc_fakerate_down(fr3, fr3_err)  # Scale down.
                fr3_up = calc_fakerate_up(fr3, fr3_err)  # Scale up.
                # Use fake rates to calculate new event weight.
                new_weight_down = (fr2_down / (1 - fr2_down)) * (
                    fr3_down / (1 - fr3_down)) * evt_weight_calcd
                new_weight = (fr2 / (1 - fr2)) * (fr3 /
                                                  (1 - fr3)) * evt_weight_calcd
                new_weight_up = (fr2_up / (1 - fr2_up)) * (
                    fr3_up / (1 - fr3_up)) * evt_weight_calcd

            if verbose:
                announce(f"Valid ZZ cand is {cr_str}")
                print(f"  This fail code: {n_fail_code}\n"
                      f"  Z2 lep codes FOR DEBUGGING:\n"
                      f"  ===========================\n"
                      f"  0, if neither lep from Z2 failed\n"
                      f"  2, if first lep from Z2 failed\n"
                      f"  3, if second lep from Z2 failed\n"
                      f"  5, if both leps from Z2 failed\n"
                      f"  ===========================\n"
                      f"  fr2_down={fr2_down:.6f}, fr3_down={fr3_down:.6f}\n"
                      f"       fr2={fr2:.6f},      fr3={fr3:.6f}\n"
                      f"    fr2_up={fr2_up:.6f},   fr3_up={fr3_up:.6f}\n"
                      f"         tree.eventWeight = {tree.eventWeight:.6f}\n"
                      f"         evt_weight_calcd = {evt_weight_calcd:.6f}\n"
                      f"          new_weight_down = {new_weight_down:.6f}\n"
                      f"               new_weight = {new_weight:.6f}\n"
                      f"            new_weight_up = {new_weight_up:.6f}")
                zzname = f"SELECTED ZZ CAND ({ndx_zzcand}/{n_valid_ZZcands_OS})"
                zzcand.print_info(name=zzname)

            # Save this quartet in TTree. Fill branches.
            lep_idcs = [
                mylep1_fromz1.ndx_lepvec,
                mylep2_fromz1.ndx_lepvec,
                mylep1_fromz2.ndx_lepvec,
                mylep2_fromz2.ndx_lepvec,
            ]
            ptr_finalState[0] = dct_finalstates_str2int[
                zzcand.get_finalstate()]
            ptr_nZXCRFailedLeptons[0] = zzcand.get_num_failing_leps()
            ptr_is3P1F[0] = zzcand.check_valid_cand_os_3p1f()
            ptr_is2P2F[0] = zzcand.check_valid_cand_os_2p2f()
            ptr_isData[0] = isData
            ptr_isMCzz[0] = isMCzz
            ptr_fr2_down[0] = fr2_down
            ptr_fr2[0] = fr2
            ptr_fr2_up[0] = fr2_up
            ptr_fr3_down[0] = fr3_down
            ptr_fr3[0] = fr3
            ptr_fr3_up[0] = fr3_up
            ptr_eventWeightFR_down[0] = new_weight_down
            ptr_eventWeightFR[0] = new_weight
            ptr_eventWeightFR_up[0] = new_weight_up
            ptr_lep_RedBkgindex[0] = lep_idcs[0]
            ptr_lep_RedBkgindex[1] = lep_idcs[1]
            ptr_lep_RedBkgindex[2] = lep_idcs[2]
            ptr_lep_RedBkgindex[3] = lep_idcs[3]

            # Label entire event as either 3P1F or 2P2F.
            # The new RedBkg logic gives 3P1F priority over 2P2F,
            # i.e. if there are 3P1F and 2P2F quartets, label event 3P1F.
            #### However when finding the quartet whose lep indices match
            #### lep_Hindex (i.e., when syncing with the xBF Analyzer),
            #### then choose the quartet (ZZ) with the highest D_kin_bkg.
            #### In this case, 2P2F may override 3P1F.
            if stop_when_found_3p1f and (overall_evt_label == '3P1F') \
                                    and (cr_str == '2P2F'):
                # DON'T label as 2P2F!
                pass
            else:
                overall_evt_label = cr_str

            n_saved_quartets_3p1f += zzcand.check_valid_cand_os_3p1f()
            n_saved_quartets_2p2f += zzcand.check_valid_cand_os_2p2f()

            if recalc_masses:
                ptr_massZ1[0] = zzcand.z_fir.get_mass()
                ptr_massZ2[0] = zzcand.z_sec.get_mass()
                ptr_mass4l[0] = zzcand.get_m4l()
                # Getting an IndexError since there is a discrepancy in the length
                # of vectors, like `lep_pt` and `vtxLepFSR_BS_pt`.
                # ptr_mass4l[0] = calc_mass4l_from_idcs(
                #     tree, lep_idcs, kind="lepFSR"
                #     )
                # ptr_mass4l_vtxFSR_BS[0] = calc_mass4l_from_idcs(
                #     tree, lep_idcs, kind="vtxLepFSR_BS"
                #     )

            # Save each quartet (ZZcand) as a separate entry in TTree.
            if outfile_root is not None:
                new_tree.Fill()

            if verbose:
                print(f"** ZZ cand {ndx_zzcand} passed OS Method Sel: **\n"
                      f"   CR {cr_str}: {evt_id}, (row {evt_num})")

            if found_matching_lep_Hindex:
                # Matching quartet info has been saved. Done with event.
                # Current zzcand is retained for further analysis.
                break

            if keep_one_quartet and not match_lep_Hindex:
                # Save first valid ZZ cand.
                break
        # End loop over quartets.

        # Final counts.
        evt_info_d["n_sel_redbkg_evts"] += 1
        evt_info_d["n_saved_quartets_3p1f"] += n_saved_quartets_3p1f
        evt_info_d["n_saved_quartets_2p2f"] += n_saved_quartets_2p2f
        evt_info_d["n_saved_evts_3p1f"] += (overall_evt_label == '3P1F')
        evt_info_d["n_saved_evts_2p2f"] += (overall_evt_label == '2P2F')

        # if fill_hists:
        #     d_hists[name][cr]["mass4l"].Fill(tree.mass4l, new_weight)
        #     d_hists[name][cr]["n_quartets"].Fill(n_valid_ZZcands_OS, 1)

        # if evt_is_2p2plusf:
        # cr = "2p2f"
        # n_quart_3p1f = 0
        # n_quart_2p2f = n_valid_ZZcands_OS
        # evt_info_d["n_good_2p2f_evts"] += 1

        evt_info_2p2f_3p1f_d[evt_id] = {
            "num_combos_2p2f": n_saved_quartets_2p2f,
            "num_combos_3p1f": n_saved_quartets_3p1f,
        }

        # if fill_hists:
        #     h2_n3p1fcombos_n2p2fcombos.Fill(n_quart_2p2f, n_quart_3p1f, 1)
    print("End loop over events.")

    print("Event info:")
    pretty_print_dict(evt_info_d)

    if outfile_root is not None:
        print(f"Writing tree to root file:\n{outfile_root}")
        new_tree.Write()

        if fill_hists:
            print(f"Writing hists to root file:\n{outfile_root}")
            for d_name in d_hists.values():
                for d_cr in d_name.values():
                    for h in d_cr.values():
                        h.Write()
            h2_n3p1fcombos_n2p2fcombos.Write()

        new_file.Close()

    if outfile_json is not None:
        save_to_json(evt_info_2p2f_3p1f_d,
                     outfile_json,
                     overwrite=overwrite,
                     sort_keys=False)