macros/wEEC_3D_unfolding_kinematics/wEEC_doUnfolding.C

0001 #include "analysisHelper.h"
0002 #include "RooUnfold.h"
0003 #include "RooUnfoldResponse.h"
0004 #include "RooUnfoldBayes.h"
0005
0006 // ─────────────────────────────────────────────────────────────────────────────
0007 //  wEEC_doUnfolding.C
0008 //
0009 //  Single-step 3D unfolding of the wEEC, one RooUnfoldBayes call per ΔΦ bin.
0010 //
0011 //  The flat-3D index defined in analysisHelper.h only covers kinematically
0012 //  valid (iL,iS) dijet pT pairs (trueSublPtBins[iS] < trueLeadPtBins[iL+1]),
0013 //  so every truth bin in the response matrix has a nonzero row sum and
0014 //  RooUnfold's Bayesian prior is fully nonzero. No runtime compression needed.
0015 //
0016 //  Dijet-level fakes and misses are encoded directly in the response matrix
0017 //  via RooUnfoldResponse::Fake() and ::Miss() in fillResponse.C. Tower-level
0018 //  fakes/misses within matched dijets are also in the response matrix.
0019 //  RooUnfold handles all of this internally.
0020 //
0021 //  Outputs per ΔΦ bin k:
0022 //    hWEEC3D_meas_{k}      — measured input clone (diagnostic)
0023 //    hWEEC3D_unfolded_{k}  — unfolded flat-3D histogram
0024 //
0025 //  Other outputs:
0026 //    hJetPt_meas, hJetPt_unfolded, hJetPt_truth  — dijet pT unfolding
0027 //
0028 //  Mode::kFull / kHalf: measured from MC response file (hWEEC3D_meas_{k})
0029 //  Mode::kData:         measured from measFile (hWEEC3D_data_{k})
0030 // ─────────────────────────────────────────────────────────────────────────────
0031 void wEEC_doUnfolding(const char* respFile,
0032                        const char* outFileName,
0033                        int         nIterWEEC   = 4,
0034                        Mode        mode        = Mode::kFull,
0035                        const char* measFile    = nullptr,
0036                        int         minCounts   = 0)   // bins with fewer raw entries are zeroed
0037 {
0038     const bool isData    = (mode == Mode::kData);
0039     const bool isClosure = !isData;
0040     const bool doTrim = false;
0041     if (isData && !measFile) {
0042         std::cerr << "kData mode requires measFile\n"; return;
0043     }
0044
0045     TFile* fResp = TFile::Open(respFile, "READ");
0046     if (!fResp || fResp->IsZombie()) {
0047         std::cerr << "Cannot open " << respFile << "\n"; return;
0048     }
0049
0050     TFile* fMeas = nullptr;
0051     if (isData) {
0052         fMeas = TFile::Open(measFile, "READ");
0053         if (!fMeas || fMeas->IsZombie()) {
0054             std::cerr << "Cannot open " << measFile << "\n"; return;
0055         }
0056     }
0057     TFile* fMeasSrc = fMeas ? fMeas : fResp;
0058
0059     TFile* fOut = new TFile(outFileName, "RECREATE");
0060
0061     auto meas3DName = [&](int k) -> std::string {
0062         return isData
0063             ? std::format("hWEEC3D_data_{}", k)
0064             : std::format("hWEEC3D_meas_{}", k);
0065     };
0066
0067     // ── dijet pT unfolding ────────────────────────────────────────────────
0068     {
0069         RooUnfoldResponse* respJetPt =
0070             (RooUnfoldResponse*)fResp->Get("response_jet_pT");
0071         if (!respJetPt) {
0072             std::cerr << "Warning: cannot find response_jet_pT — skipping jet pT unfolding\n";
0073         } else {
0074             TH1D* hJetPt_input = nullptr;
0075             if (mode == Mode::kFull) {
0076                 hJetPt_input = (TH1D*)fResp->Get("hJetPt_reco");
0077                 if (hJetPt_input) {
0078                     hJetPt_input = (TH1D*)hJetPt_input->Clone("hJetPt_meas");
0079                     hJetPt_input->SetDirectory(0);
0080                 }
0081             } else {
0082                 hJetPt_input = (TH1D*)fResp->Get("hJetPt_meas");
0083                 if (hJetPt_input) {
0084                     hJetPt_input = (TH1D*)hJetPt_input->Clone();
0085                     hJetPt_input->SetDirectory(0);
0086                 }
0087             }
0088
0089             if (hJetPt_input) {
0090                 RooUnfoldBayes unfoldJet(respJetPt, hJetPt_input, 4);
0091                 unfoldJet.HandleFakes(true);
0092                 TH1D* hJetPt_unfolded = (TH1D*)unfoldJet.Hunfold();
0093                 hJetPt_unfolded->SetDirectory(0);
0094                 hJetPt_unfolded->SetName("hJetPt_unfolded");
0095
0096                 fOut->cd();
0097                 hJetPt_input->Write("hJetPt_meas");
0098                 hJetPt_unfolded->Write();
0099
0100                 if (isClosure) {
0101                     TH1D* hJetPt_truth = (TH1D*)fResp->Get("hJetPt_truth");
0102                     if (hJetPt_truth) {
0103                         hJetPt_truth->SetDirectory(0);
0104                         hJetPt_truth->Write();
0105                     }
0106                 }
0107
0108                 delete hJetPt_input;
0109                 delete hJetPt_unfolded;
0110             }
0111         }
0112     }
0113
0114     // ── wEEC unfolding — one ΔΦ bin at a time ────────────────────────────
0115     for (int k = 0; k < nDphi; ++k)
0116     {
0117         TH1D* hMeas = (TH1D*)fMeasSrc->Get(meas3DName(k).c_str());
0118         if (!hMeas || hMeas->Integral() == 0) continue;
0119
0120         TH1D* hMeasClone = (TH1D*)hMeas->Clone(
0121             std::format("hWEEC3D_meas_{}", k).c_str());
0122         hMeasClone->SetDirectory(0);
0123
0124         RooUnfoldResponse* resp =
0125             (RooUnfoldResponse*)fResp->Get(
0126                 std::format("response_wEEC3D_{}", k).c_str());
0127         if (!resp) { delete hMeasClone; continue; }
0128
0129         // ── Count-based trimming ──────────────────────────────────────────
0130         // hWEEC3D_counts_{k} is a 2D histogram with the same (reco, truth)
0131         // binning as the response matrix, filled with unity weight for every
0132         // matched tower pair that entered a Fill call.  Any cell whose raw
0133         // count is below minCounts is individually zeroed in the cloned
0134         // response matrix; cells above threshold are left intact.
0135         // The measured histogram is not modified — only the response is
0136         // trimmed, so RooUnfold simply has no path to those truth bins from
0137         // those reco bins.
0138         TH2D* hCounts = (TH2D*)fResp->Get(
0139             std::format("hWEEC3D_counts_{}", k).c_str());
0140
0141         RooUnfoldResponse* respTrimmed = nullptr;
0142         // hResp2D bin indices: X = reco flat3D (1-based), Y = truth flat3D
0143         TH2D* hResp2D = (TH2D*)resp->Hresponse()->Clone();
0144         if (doTrim && hCounts) {
0145             for (int rBin = 1; rBin <= hResp2D->GetNbinsX(); ++rBin)
0146             for (int tBin = 1; tBin <= hResp2D->GetNbinsY(); ++tBin) {
0147                 double rawCount = hCounts->GetBinContent(rBin, tBin);
0148                 if (rawCount < minCounts)
0149                     hResp2D->SetBinContent(rBin, tBin, 0.0);
0150             }
0151             respTrimmed = new RooUnfoldResponse((TH1D*)resp->Hmeasured()->Clone(), (TH1D*)resp->Htruth()->Clone(), hResp2D);
0152         }
0153
0154
0155         RooUnfoldResponse *respToUse = (doTrim && respTrimmed != nullptr) ? respTrimmed : resp;
0156
0157
0158         std::cout << "Unfolding ΔΦ bin " << k << std::endl;
0159
0160         // All truth bins in the response matrix are kinematically valid by
0161         // construction (see analysisHelper.h), so the prior vector is fully
0162         // nonzero and no runtime compression is needed.
0163         RooUnfoldBayes unfold(respToUse, hMeasClone, nIterWEEC);
0164         unfold.HandleFakes(true);
0165         TH1D* hUnf = (TH1D*)unfold.Hunfold(RooUnfold::ErrorTreatment::kCovariance);
0166         hUnf->SetDirectory(0);
0167         hUnf->SetName(std::format("hWEEC3D_unfolded_{}", k).c_str());
0168         TMatrixD covM = (TMatrixD)unfold.Eunfold(RooUnfold::ErrorTreatment::kCovariance);
0169
0170         // Full covariance matrix C = M * V_meas * M^T where:
0171         //   M      = unfolding matrix (nTruth × nReco) from UnfoldingMatrix()
0172         //   V_meas = diagonal covariance of the measured histogram
0173         //
0174         // This is the analytic Bayesian error propagation formula, identical to
0175         // what RooUnfold computes internally for kCovariance mode.  It assumes
0176         // M is independent of the measurement — a good approximation for small
0177         // iteration counts (nIterWEEC = 4 here).
0178         //
0179         // RooUnfold appends one extra fake-handling bin to the truth axis, so
0180         // M has shape (nTrueFlat3D()+1) × nRecoFlat3D().  We trim the last row
0181         // to get a nTrueFlat3D() × nTrueFlat3D() covariance that matches hUnf.
0182         TMatrixD M = unfold.UnfoldingMatrix();   // (nT+1) × nR
0183         int nT = hUnf->GetNbinsX();              // nTrueFlat3D() — trim fake row
0184         int nR = hMeasClone->GetNbinsX();
0185
0186         // Build diagonal measured covariance from bin errors
0187         TMatrixD Vmeas(nR, nR);
0188         for (int i = 0; i < nR; ++i)
0189             Vmeas(i, i) = std::pow(hMeasClone->GetBinError(i + 1), 2);
0190
0191         // Trim M to nT rows (drop the fake-handling last row)
0192         TMatrixD Mtrim(nT, nR);
0193         for (int i = 0; i < nT; ++i)
0194         for (int j = 0; j < nR; ++j)
0195             Mtrim(i, j) = M(i, j);
0196
0197         TMatrixD MT(TMatrixD::kTransposed, Mtrim);
0198         TMatrixD cov = Mtrim * Vmeas * MT;   // nT × nT
0199
0200         // Reset hUnf bin errors to match the covariance diagonal.
0201         // Hunfold() without an error mode sets errors to sqrt(bin_content)
0202         // (Poisson approximation on weighted counts), which is not the correct
0203         // propagated uncertainty.  The diagonal of C = M*Vmeas*M^T is the
0204         // correct per-bin variance from measurement error propagation.
0205         hUnf->Sumw2();
0206         for (int i = 0; i < nT; ++i)
0207             //hUnf->SetBinError(i + 1, std::sqrt(std::max(cov(i, i), 0.0)));
0208             hUnf->SetBinError(i + 1, std::sqrt(std::max(covM(i, i), 0.0)));
0209
0210         // Sanity check: find first non-zero measured bin and compare
0211         if (k == 0) {
0212             for (int i = 0; i < nR; ++i) {
0213                 double v = hMeasClone->GetBinContent(i+1);
0214                 double e = hMeasClone->GetBinError(i+1);
0215                 if (v > 0 || e > 0) {
0216                     std::cout << std::format(
0217                         "  First non-zero meas bin {}: content={:.3e} error={:.3e} rel={:.3f}\n",
0218                         i, v, e, (v > 0 ? e/v : 0.0));
0219                     break;
0220                 }
0221             }
0222             for (int i = 0; i < nT; ++i) {
0223                 if (cov(i,i) > 0) {
0224                     std::cout << std::format(
0225                         "  First non-zero cov diag bin {}: cov={:.3e} hUnfErr^2={:.3e}\n",
0226                         i, cov(i,i), std::pow(hUnf->GetBinError(i+1), 2));
0227                     break;
0228                 }
0229             }
0230         }
0231
0232         fOut->cd();
0233         hMeasClone->Write();
0234         hUnf->Write();
0235         covM.Write(std::format("covDirectWEEC3D_{}",k).c_str());
0236         cov.Write(std::format("covWEEC3D_{}", k).c_str());
0237         if (doTrim && respTrimmed)
0238             respTrimmed->Write(std::format("response_wEEC3D_trimmed_{}", k).c_str());
0239
0240         delete hMeasClone;
0241         delete hUnf;
0242         if (doTrim && respTrimmed) delete respTrimmed;
0243     }
0244
0245     fResp->Close(); delete fResp;
0246     if (fMeas) { fMeas->Close(); delete fMeas; }
0247     fOut->Close();
0248     std::cout << "Written: " << outFileName << "\n";
0249 }