generate_fd_fake_events.C

Go to the documentation of this file.

/// \file   generate_fd_fake_events.C
/// \brief  Produce fake FD files (but including the real 2018 data)
///         to aid in testing blessed plots before box opening
/// \author J. Wolcott <jwolcott@fnal.gov>
/// \date   Mar. 21, 2019

#include <numeric>
#include <deque>

#include "TRandom3.h"
#include "TTree.h"

#include "CAFAna/Analysis/Prod4Loaders.h"
#include "CAFAna/Core/FileReducer.h"
#include "CAFAna/Core/SAMQuerySource.h"
#include "CAFAna/Core/Utilities.h"
#include "CAFAna/Cuts/SpillCuts.h"
#include "3FlavorAna/Cuts/NueCuts2018.h"
#include "3FlavorAna/Cuts/NumuCuts2018.h"
#include "OscLib/OscCalcPMNSOpt.h"
#include "StandardRecord/SRSpill.h"
#include "StandardRecord/StandardRecord.h"
#include "StandardRecord/Proxy/SRProxy.h"
#include "Utilities/func/NuMITimeParams.h"

#include "fake_data.h"

using namespace ana;

namespace ana2019
{
  // amount of livetime (in seconds) assumed for each spill.
  // lifted from TimingCuts.cxx
  const double kFixedSpillLivetime = 1e-6 * util::kBeamWindowMicroSec;

  // lifted from  R18-02-25-art1.a/CAFAna/nue/reduce_nue_rhc.C (2018 decaf macro)
  const Var kCVN2017e([](const caf::SRProxy* sr){return sr->sel.cvn.nueid;});
  const Cut kNueFD2018DecafCut = kApplySecondAnalysisMask && kVeto && kNueDQ2017CVN && (kCVN2017e > 0.5 || kCVNSSe > 0.5);

  namespace fakedata
  {
    const std::string OUT_FILE_STUB = "fake_fd_events_2019_%s.root";

    /// base class: random number generator with easily settable seed...
    class EvtHashRNG
    {
      public:
        EvtHashRNG()
          : fRng(0)
        {}

        TRandom * RNG() { return &fRng; }
        void SetSeed(std::size_t seed) { fRng.SetSeed(seed); }

      private:
        mutable TRandom3 fRng;
    };

    /// Class for deciding whether to keep an event (either full StandardRecord or just SRSpill)
    /// based on a random number generator and a single (changable) probability
    class KeepEvtFunc : public EvtHashRNG
    {
      public:
        KeepEvtFunc(double keepProb)
          : fKeepProb(keepProb), fLastHash(0)
        {}

        virtual bool Keep(const caf::SRSpillProxy *sr)
        {
          return Keep(Hash(sr));
        }

        virtual bool Keep(const caf::SRProxy *sr)
        {
          auto hash = Hash({Hash(&sr->spill), (std::size_t)(sr->hdr.cycle), sr->hdr.subevt});
          return Keep(hash);
        }

        void SetKeepProb(double prob)
        {
          fKeepProb = prob;
        }

      protected:

        bool Keep(std::size_t hash)
        {
          if (hash != fLastHash)
          {
            SetSeed(hash);
            fLastHash = hash;
          }

         return RNG()->Uniform() < fKeepProb;
        }

      private:
        double fKeepProb;
        mutable std::size_t fLastHash;
    };

    /// Functional for deciding whether to keep an MC event
    /// based on a fixed probability (e.g., POT-based)
    /// and on oscillation probability.
    /// Can be passed to a CAFAna Cut object as its internal function
    class KeepMCEventFunc : public KeepEvtFunc
    {
      public:
        /// 'keepProb' is for all events (depends on POT only).
        /// 'NCwgt' is because we get NCs from all the swap files too
        ///   and have too many unless we weight them down...
        /// you can pass in 'kept' in case you want to know what kind of events were retained...
        KeepMCEventFunc(double keepProb, double NCwgt, std::map<std::pair<int, int>, unsigned int> * kept = nullptr)
            : KeepEvtFunc(keepProb), fNCwgt(NCwgt), fEventsKept(kept)
        {
          ResetOscCalcToDefault(&fCalc);
        }


        bool operator()(const caf::SRProxy * sr)
        {
          // short-circuit on the POT prob to avoid doing oscillation prob calculations unnecessarily
          if (!Keep(&sr->spill))
            return false;

          if (sr->mc.nnu != 1)
            return false;

          bool keep;
          if (sr->mc.nu[0].iscc)
            keep = RNG()->Uniform() < fCalc.P(sr->mc.nu[0].pdgorig, sr->mc.nu[0].pdg, sr->mc.nu[0].E);
          else
            keep = RNG()->Uniform() < fNCwgt;

          if (keep && fEventsKept)
          {
            auto flav = sr->mc.nu[0].iscc ? (std::make_pair(sr->mc.nu[0].pdgorig, sr->mc.nu[0].pdg))
                                          : (std::make_pair(0, 0));
            (*fEventsKept)[flav]++;
          }
          return keep;
        }

        auto EventsKept() const
        {
          return fEventsKept;
        }

      private:
        double fNCwgt;
        mutable osc::OscCalcPMNSOpt fCalc;

        mutable std::map<std::pair<int, int>, unsigned int> * fEventsKept;
    };
  } // namespace fakedata

  /// a very dumb implementation of a stream-wrapper that makes the text in one sequence of <<s some color and resets afterwards.
  /// helps output stick out since the SpectrumLoaders are fairly chatty in between useful output.
  /// resets the stream when it's destroyed so that subsequent text doesn't show up with color.
  /// use like:
  ///    ColorStream(std::cout)     << "This text, and this number " << 5736 << " show up as cyan!" << std::endl;
  ///    ColorStream(std::cout, 31) << "This text is red!" << std::endl;
  ///    std::cout << "This text has no special coloring" << std::endl;
  class ColorStream
  {
    public:
      ColorStream(std::ostream & os, unsigned int col=36)
        : fOS(os)
      {
        // printf so that it definitely all goes into the same stream with no flushes in between
        printf("\033[%d;1m", col);
      }
      ~ColorStream()
      {
        fOS << "\033[0;0m";
      }

      template <typename T>
      std::ostream & operator<<(const T & t)
      {
        return fOS << t;
      }

    private:
      std::ostream & fOS;
  };
} // namespace ana2019

// -------------------------------------------------

void generate_fd_fake_events(const std::string & sample)
{
  // ------------ step 0:
  // which sample are we working with?
  std::string fdDataDef;
  const Cut * decafCut;
  ECAFType decafType;
  if (sample == "nue")
  {
    decafType = kNueConcat;
    decafCut = &ana2019::kNueFD2018DecafCut;
    fdDataDef = "prod4_sumrestricteddecaf_fd_numi_rhc_full_goodruns_nue2018";
  }
  else if (sample == "numu")
  {
    decafType = kNumuConcat;
    decafCut = &kNumuDecafCutFD2018;
    fdDataDef = "prod4_sumrestricteddecaf_fd_numi_rhc_full_goodruns_numu2018";
  }
  else
  {
    std::cerr << "Unknown sample: '" << sample << "'" << std::endl;
    std::cerr << "I only know how to make fake data for 'numu' or 'nue' samples." << std::endl;
    return;
  }

  // ------------ step 1:
  // figure out how much POT and livetime we already have,
  // both in RHC data and the MC sample we'll be drawing from
  ana2019::ColorStream(ana2019::ColorStream(std::cout)) << "Counting POT in RHC MC and data datasets..." << std::endl << std::endl;

  const std::map<Loaders::SwappingConfig, std::string> kSwaps
  {
      {Loaders::kNonSwap,  "nonswap"},
      {Loaders::kFluxSwap, "fluxswap"},
      {Loaders::kTauSwap,  "tauswap"}
  };

  Prod4NomLoaders loaders(decafType, Loaders::kRHC);
  // the loaders don't normally know where the FD data is to prevent accidental unblinding.
  // but, since these fake events are meant to include the FD events
  // we've already selected in 2018, we need to make an exception.
  loaders.SetLoaderPath(fdDataDef, caf::kFARDET, Loaders::kData);

  loaders.SetSpillCut(kStandardSpillCuts);
  HistAxis dummyAxis("dummy", Binning::Simple(1, 0, 1), kUnweighted);
  std::map<Loaders::SwappingConfig, Spectrum> mcDummySpecs;
  std::map<Loaders::SwappingConfig, TH1D>     mcDummySpillHists;
  for (const auto swapPair : kSwaps)
  {
    auto & loader = loaders.GetLoader(caf::kFARDET, Loaders::kMC, kBeam, swapPair.first);
    mcDummySpecs.emplace(std::piecewise_construct,
                         std::forward_as_tuple(swapPair.first),
                         std::forward_as_tuple(loader, dummyAxis, *decafCut));
    mcDummySpillHists.emplace(std::piecewise_construct,
                              std::forward_as_tuple(swapPair.first),
                              std::forward_as_tuple(Form("spills_%d", swapPair.first), "", 1, 0, 2));
    loader.AddSpillHistogram(&mcDummySpillHists.at(swapPair.first), kSpillUnweighted, kStandardSpillCuts);
  }
  Spectrum dataDummySpec(loaders.GetLoader(caf::kFARDET, Loaders::kData), dummyAxis, *decafCut && kInBeamSpill);

  loaders.Go();

  // ------------ step 2:
  // draw events out of the MC according to the expected rates
  // (using the 2018 best fit oscillation parameters
  //  and the ratio of the desired "new data" and MC exposures).
  // simultaneously rewrite the MC spill records to match the "new data" POT & livetime
  const double fakeDataPOT = ana2019::fakedata::FULL_RHC_EXPOSURE - dataDummySpec.POT();
  const double fakeDataLivetime = fakeDataPOT / dataDummySpec.POT() * dataDummySpec.Livetime();
  double totalMCPOT = 0;
  for (const auto & swapPair : mcDummySpecs)
    totalMCPOT += swapPair.second.POT();

  ana2019::ColorStream(std::cout)   << "Found the following POT:" << std::endl;
  ana2019::ColorStream(std::cout)   << "    data --> " << dataDummySpec.POT() << std::endl;
  ana2019::ColorStream(std::cout)   << "         (which means to obtain " << ana2019::fakedata::FULL_RHC_EXPOSURE
                                      << " total RHC POT, we need to draw " << fakeDataPOT << " POT of fake data.)" << std::endl;
  for (const auto & swapPair : mcDummySpecs)
    ana2019::ColorStream(std::cout) << "      MC swap " << swapPair.first << " --> " << swapPair.second.POT()
                                      << " (" << mcDummySpillHists.at(swapPair.first).Integral() << " spills)" << std::endl;
  ana2019::ColorStream(std::cout)   << "    --> total MC =  " << totalMCPOT << std::endl;
  ana2019::ColorStream(std::cout)   << "Found data livetime = " << dataDummySpec.Livetime() << "s" << std::endl;
  ana2019::ColorStream(std::cout)   << "     so when scaled to " << ana2019::fakedata::FULL_RHC_EXPOSURE << " POT equivalent,"
                                      << " need " << fakeDataLivetime << "s of livetime in fake data." << std::endl;


  std::map<Loaders::SwappingConfig, std::string> intermediateFileNames;
  const double MCSwapScale = 1. / mcDummySpecs.size();   // so that we only count the NCs & the POT once over all the swaps...
  for (const auto & swapPair : mcDummySpecs)
  {
    auto swap = swapPair.first;

    // a further complication here.
    // we need to keep the full "fake data"'s worth of POT for the CC events,
    // since each possible oscillation pair (nue -> nue, nue -> numu, nue -> nutau, numu -> nue, ...)
    // is represented only once over the various swaps.
    // however, there are NCs in all of them,
    // and we also want to keep only enough spill records from all the swaps
    // to add to the total desired fake data POT,
    // so both NC events & spill records are reduced by the extra MCSwapScale from above.
    double keepProbCC = fakeDataPOT / swapPair.second.POT();
    const auto nSpills = mcDummySpillHists.at(swapPair.first).Integral();
    const double keepProbSpill = (fakeDataLivetime * MCSwapScale) / ana2019::kFixedSpillLivetime / nSpills;
    if (keepProbSpill < 0 || keepProbSpill > 1)
    {
      ana2019::ColorStream(std::cout, 31) << "Error: spill retention probability is invalid: " << keepProbSpill
                                          << std::endl;
      ana2019::ColorStream(std::cout, 31) << "Perhaps you have too few spills (" << nSpills << ") to add up to the necessary livetime"
                                          << " (note that they are " << ana2019::kFixedSpillLivetime << "s/spill)?" << std::endl;
      abort();
    }
    std::cout << std::endl;
    ana2019::ColorStream(std::cout) << "---------------" << std::endl;
    std::cout << std::endl;
    ana2019::ColorStream(std::cout) << "For swap " << swap << ":" << std::endl;
    ana2019::ColorStream(std::cout) << "       acceptance probability for CC events: " << keepProbCC << std::endl;
    ana2019::ColorStream(std::cout) << "       (for spill info & NC events, this is further modified by a factor of " << MCSwapScale
                                      << " since we will be sampling from multiple swaps.)" << std::endl;
    ana2019::ColorStream(std::cout) << "       N spills = " << nSpills << std::endl;
    ana2019::ColorStream(std::cout) << "       need to retain " <<  (fakeDataLivetime * MCSwapScale) / ana2019::kFixedSpillLivetime
                                      << " spills to get desired livetime" << std::endl;

    std::string intermediateFileName = "2019fakedata.intermediate-" + kSwaps.at(swap) + ".root";
    intermediateFileNames[swap] = intermediateFileName;

    std::map<std::pair<int, int>, unsigned int> eventsKept;
    // the second argument is the further reduction factor applied to NC events mentioned above
    ana2019::fakedata::KeepMCEventFunc keepMcEventFunc {keepProbCC, MCSwapScale, &eventsKept};
    Cut kKeepEventCut{keepMcEventFunc};
    // for spill records we need to scale down as mentioned above
    ana2019::fakedata::KeepEvtFunc keepSpillFuncObj{keepProbCC * MCSwapScale};
    FileReducer MCsampler(loaders.GetLoaderPath(caf::kFARDET, Loaders::kMC, kBeam, swap), intermediateFileName);
    MCsampler.AddEventCut(kKeepEventCut);
    MCsampler.AddEventCut(*decafCut);
    double totalPOT = 0;
    double totalLive = 0;
    // todo: it would be *better* for this to be a _cut_ on the spill, but FileReducer has no such concept.
    // instead we just reset those spills that were rejected and count on ROOT compression to shrink them down
    MCsampler.AddSpillReductionStep([&](caf::SRSpill *spill, const caf::SRSpillProxy *spillProxy)
                                    {
                                      bool potKeep = false;
                                      if (spill->isgoodspill && spill->spillpot > 0)
                                      {
                                        keepSpillFuncObj.SetKeepProb(keepProbCC * MCSwapScale);
                                        potKeep = keepSpillFuncObj.Keep(spillProxy);
                                      }

                                      // always keep a spill for livetime if we're keeping it for POT
                                      keepSpillFuncObj.SetKeepProb(keepProbSpill);
                                      bool livetimeKeep = potKeep || keepSpillFuncObj.Keep(spillProxy);

                                      if (potKeep)
                                        totalPOT += spillProxy->spillpot;
                                      else
                                        spill->spillpot = 0;

                                      if (livetimeKeep)
                                      {
                                        // MC livetimes are meaningless because we don't simulate all the cosmics-only spills
                                        // where most of the livetime is anyway.
                                        // so even though we override the livetime *amount* here for completeness,
                                        // the thing that actually makes the difference is that
                                        // we are keeping the right *number* of spill.isgoodspill==true
                                        // entries in the tree
                                        spill->livetime = ana2019::kFixedSpillLivetime;
                                        totalLive += ana2019::kFixedSpillLivetime;
                                      }

                                      if (!potKeep && !livetimeKeep)
                                      {
                                        // just reset the whole thing
                                        *spill = caf::SRSpill();

                                        // make sure no NaNs
                                        spill->spillpot = 0;
                                        spill->livetime = 0;
                                      }

                                      // no matter what, we need to treat this as a "data" spill
                                      // so that either its livetime is counted correctly
                                      // or so that it's disregarded
                                      spill->ismc = false;
                                    });

    MCsampler.Go();

    // update the POT, livetime, and event counts in the file, which are otherwise lying
    TFile fakeEventsFile(intermediateFileName.c_str(), "update");
    auto potHist = dynamic_cast<TH1 *>(fakeEventsFile.Get("TotalPOT"));
    potHist->SetBinContent(1, totalPOT);
    potHist->Write();
    auto liveHist = dynamic_cast<TH1 *>(fakeEventsFile.Get("TotalLivetime"));
    liveHist->SetBinContent(1, totalLive);
    liveHist->Write();

    auto evtTree = dynamic_cast<TTree *>(fakeEventsFile.Get("recTree"));
    auto evtHist = dynamic_cast<TH1 *>(fakeEventsFile.Get("TotalEvents"));
    std::size_t evts = evtTree->GetEntries();
    evtHist->SetBinContent(1, evts);
    evtHist->Write();
    fakeEventsFile.Close();

    std::cout << std::endl;
    ana2019::ColorStream(std::cout) << "    Drew " << totalPOT << " POT, which is " << (100 * totalPOT / (fakeDataPOT * MCSwapScale))
              << "% of the desired amount." << std::endl;
    ana2019::ColorStream(std::cout) << "    " << evts << " new events were sampled from the MC:" << std::endl;
    for (const auto & keptPair : eventsKept)
    {
      std::cout << "       ";
      std::pair<int, int> flav = keptPair.first;
      unsigned int kept = keptPair.second;
      if (flav.first == 0)
        ana2019::ColorStream(std::cout) << "NC";
      else
      {
        std::string out;
        for (const auto &v : {flav.first, flav.second})
        {
          if (std::abs(v) == 12)
            out += "nue";
          else if (std::abs(v) == 14)
            out += "numu";
          else if (std::abs(v) == 16)
            out += "nutau";
          else
            out += std::to_string(v);

          if (v < 0)
            out += "bar";

          out += " --> ";
        }
        out = out.substr(0, out.size() - 5);
        ana2019::ColorStream(std::cout) << out;
      }
      ana2019::ColorStream(std::cout) << " = " << kept << std::endl;
    }
  }


  // ------------ step 3:
  // add the 2018 data events back into the file.
  // while doing that, also sample cosmics from the timing sideband
  // and use those as "fake" cosmics by moving them into the beam window...
  std::vector<std::string> files;
  std::transform(intermediateFileNames.begin(), intermediateFileNames.end(), std::back_inserter(files),
                 [](const auto & swapFilePair){ return swapFilePair.second; }
  );
  SAMQuerySource samquerysource(loaders.GetLoaderPath(caf::kFARDET, Loaders::kData));
  files.reserve(files.size() + samquerysource.NFiles());
  files.insert(files.end(), samquerysource.GetFileNames().begin(), samquerysource.GetFileNames().end());
  FileReducer combiner(files, Form(ana2019::fakedata::OUT_FILE_STUB.c_str(), sample.c_str()));

  // this is a bit weird.
  // to get some 'simulated' cosmics, we're going to take the out-of-time cosmics
  // and, with probability corresponding to the (needed new exposure)/(actual exposure),
  // **move those events into the beam window**.
  // we're assuming that the livetime scales exactly with the POT, which won't be right
  // for the actual new data, but is fine for this fake data.
  double OOTKeepProb = util::kBeamWindowMicroSec / (util::kMaxTimingSidebandAfterMicroSec - util::kMinTimingSidebandAfterMicroSec);
  OOTKeepProb *= fakeDataPOT / dataDummySpec.POT();  // note that if the timing sideband were smaller than the beam window, this would result in probabilities > 1
  assert(OOTKeepProb < 1);

  ana2019::fakedata::KeepEvtFunc keepCosmicFunc{OOTKeepProb};
  combiner.AddEventCut(*decafCut);
  combiner.AddReductionStep([OOTKeepProb, &keepCosmicFunc, &decafCut](caf::StandardRecord * sr, const caf::SRProxy * srp){
    const auto t = sr->slc.meantime;
    if (t < 1000*util::kMinTimingSidebandAfterMicroSec || t > 1000*util::kMaxTimingSidebandAfterMicroSec)
      return;

    if (!keepCosmicFunc.Keep(srp))
      return;

    // are there other branches that need to be adjusted?
    double newTimeNS = util::kBeamWindowMinMicroSec + keepCosmicFunc.RNG()->Uniform() * (util::kBeamWindowMaxMicroSec - util::kBeamWindowMinMicroSec);
    newTimeNS *= 1000;
    std::deque<std::reference_wrapper<float>> toUpdate
    {
        sr->slc.meantime,
        sr->slc.starttime,
        sr->slc.endtime,
    };
    for (auto & vtx : sr->vtx.elastic)
      toUpdate.push_back(vtx.time);

    for (auto & val : toUpdate)
      val += newTimeNS - val;  // since aren't reassigning the reference, which is what operator=() would do
  });
  combiner.AddEventCut(*decafCut);
  combiner.Go();

  for (const auto & fileNamePair : intermediateFileNames)
    system( ("rm " + fileNamePair.second).c_str() );
}