rs301_splot.C

Go to the documentation of this file.

/// \file
/// \ingroup tutorial_roostats
/// \notebook -js
/// SPlot tutorial
///
/// This tutorial shows an example of using SPlot to unfold two distributions.
/// The physics context for the example is that we want to know
/// the isolation distribution for real electrons from Z events
/// and fake electrons from QCD.  Isolation is our 'control' variable
/// To unfold them, we need a model for an uncorrelated variable that
/// discriminates between Z and QCD.  To do this, we use the invariant
/// mass of two electrons.  We model the Z with a Gaussian and the QCD
/// with a falling exponential.
///
/// Note, since we don't have real data in this tutorial, we need to generate
/// toy data.  To do that we need a model for the isolation variable for
/// both Z and QCD.  This is only used to generate the toy data, and would
/// not be needed if we had real data.
///
/// \macro_image
/// \macro_output
/// \macro_code
///
/// \author Kyle Cranmer

#include "RooRealVar.h"
#include "RooStats/SPlot.h"
#include "RooDataSet.h"
#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooChebychev.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooAddition.h"
#include "RooProduct.h"
#include "TCanvas.h"
#include "RooAbsPdf.h"
#include "RooFit.h"
#include "RooFitResult.h"
#include "RooWorkspace.h"
#include "RooConstVar.h"

// use this order for safety on library loading
using namespace RooFit;
using namespace RooStats;


// see below for implementation
void AddModel(RooWorkspace*);
void AddData(RooWorkspace*);
void DoSPlot(RooWorkspace*);
void MakePlots(RooWorkspace*);

void rs301_splot()
{

   // Create a new workspace to manage the project.
   RooWorkspace* wspace = new RooWorkspace("myWS");

   // add the signal and background models to the workspace.
   // Inside this function you will find a description our model.
   AddModel(wspace);

   // add some toy data to the workspace
   AddData(wspace);

   // inspect the workspace if you wish
   //  wspace->Print();

   // do sPlot.
   //This wil make a new dataset with sWeights added for every event.
   DoSPlot(wspace);

   // Make some plots showing the discriminating variable and
   // the control variable after unfolding.
   MakePlots(wspace);

   // cleanup
   delete wspace;

}


//____________________________________
void AddModel(RooWorkspace* ws){

   // Make models for signal (Higgs) and background (Z+jets and QCD)
   // In real life, this part requires an intelligent modeling
   // of signal and background -- this is only an example.

   // set range of observable
   Double_t lowRange = 00, highRange = 200;

   // make a RooRealVar for the observables
   RooRealVar invMass("invMass", "M_{inv}", lowRange, highRange,"GeV");
   RooRealVar isolation("isolation", "isolation", 0., 20., "GeV");


   // --------------------------------------
   // make 2-d model for Z including the invariant mass
   // distribution  and an isolation distribution which we want to
   // unfold from QCD.
   std::cout << "make z model" << std::endl;
   // mass model for Z
   RooRealVar mZ("mZ", "Z Mass", 91.2, lowRange, highRange);
   RooRealVar sigmaZ("sigmaZ", "Width of Gaussian", 2,0,10,"GeV");
   RooGaussian mZModel("mZModel", "Z+jets Model", invMass, mZ, sigmaZ);
   // we know Z mass
   mZ.setConstant();
   // we leave the width of the Z free during the fit in this example.

   // isolation model for Z.  Only used to generate toy MC.
   // the exponential is of the form exp(c*x).  If we want
   // the isolation to decay an e-fold every R GeV, we use
   // c = -1/R.
   RooConstVar zIsolDecayConst("zIsolDecayConst",
                              "z isolation decay  constant", -1);
   RooExponential zIsolationModel("zIsolationModel", "z isolation model",
                                 isolation, zIsolDecayConst);

   // make the combined Z model
   RooProdPdf zModel("zModel", "4-d model for Z",
                     RooArgSet(mZModel, zIsolationModel));

   // --------------------------------------
   // make QCD model

   std::cout << "make qcd model" << std::endl;
   // mass model for QCD.
   // the exponential is of the form exp(c*x).  If we want
   // the mass to decay an e-fold every R GeV, we use
   // c = -1/R.
   // We can leave this parameter free during the fit.
   RooRealVar qcdMassDecayConst("qcdMassDecayConst",
                              "Decay const for QCD mass spectrum",
                              -0.01, -100, 100,"1/GeV");
   RooExponential qcdMassModel("qcdMassModel", "qcd Mass Model",
                              invMass, qcdMassDecayConst);


   // isolation model for QCD.  Only used to generate toy MC
   // the exponential is of the form exp(c*x).  If we want
   // the isolation to decay an e-fold every R GeV, we use
   // c = -1/R.
   RooConstVar qcdIsolDecayConst("qcdIsolDecayConst",
                                 "Et resolution constant", -.1);
   RooExponential qcdIsolationModel("qcdIsolationModel", "QCD isolation model",
                                 isolation, qcdIsolDecayConst);

   // make the 2-d model
   RooProdPdf qcdModel("qcdModel", "2-d model for QCD",
      RooArgSet(qcdMassModel, qcdIsolationModel));

   // --------------------------------------
   // combined model

   // These variables represent the number of Z or QCD events
   // They will be fitted.
   RooRealVar zYield("zYield","fitted yield for Z",50 ,0.,1000) ;
   RooRealVar qcdYield("qcdYield","fitted yield for QCD", 100 ,0.,1000) ;

   // now make the combined model
   std::cout << "make full model" << std::endl;
   RooAddPdf model("model","z+qcd background models",
                  RooArgList(zModel, qcdModel),
                  RooArgList(zYield,qcdYield));


   // interesting for debugging and visualizing the model
   model.graphVizTree("fullModel.dot");

   std::cout << "import model" << std::endl;

   ws->import(model);
}

//____________________________________
void AddData(RooWorkspace* ws){
   // Add a toy dataset

   // how many events do we want?
   Int_t nEvents = 1000;

   // get what we need out of the workspace to make toy data
   RooAbsPdf* model = ws->pdf("model");
   RooRealVar* invMass = ws->var("invMass");
   RooRealVar* isolation = ws->var("isolation");

   // make the toy data
   std::cout << "make data set and import to workspace" << std::endl;
   RooDataSet* data = model->generate(RooArgSet(*invMass, *isolation),nEvents);

   // import data into workspace
   ws->import(*data, Rename("data"));

}

//____________________________________
void DoSPlot(RooWorkspace* ws){
   std::cout << "Calculate sWeights" << std::endl;

   // get what we need out of the workspace to do the fit
   RooAbsPdf* model = ws->pdf("model");
   RooRealVar* zYield = ws->var("zYield");
   RooRealVar* qcdYield = ws->var("qcdYield");
   RooDataSet* data = (RooDataSet*) ws->data("data");

   // fit the model to the data.
   model->fitTo(*data, Extended() );

   // The sPlot technique requires that we fix the parameters
   // of the model that are not yields after doing the fit.
   //
   // This *could* be done with the lines below, however this is taken care of
   // by the RooStats::SPlot constructor (or more precisely the AddSWeight
   // method).
   //
   //RooRealVar* sigmaZ = ws->var("sigmaZ");
   //RooRealVar* qcdMassDecayConst = ws->var("qcdMassDecayConst");
   //sigmaZ->setConstant();
   //qcdMassDecayConst->setConstant();


   RooMsgService::instance().setSilentMode(true);


   // Now we use the SPlot class to add SWeights to our data set
   // based on our model and our yield variables
   RooStats::SPlot* sData = new RooStats::SPlot("sData","An SPlot",
                                             *data, model, RooArgList(*zYield,*qcdYield) );


   // Check that our weights have the desired properties

   std::cout << "Check SWeights:" << std::endl;


   std::cout << std::endl <<  "Yield of Z is "
               << zYield->getVal() << ".  From sWeights it is "
               << sData->GetYieldFromSWeight("zYield") << std::endl;


   std::cout << "Yield of QCD is "
               << qcdYield->getVal() << ".  From sWeights it is "
               << sData->GetYieldFromSWeight("qcdYield") << std::endl
               << std::endl;

   for(Int_t i=0; i < 10; i++)
      {
      std::cout << "z Weight   " << sData->GetSWeight(i,"zYield")
                  << "   qcd Weight   " << sData->GetSWeight(i,"qcdYield")
                  << "  Total Weight   " << sData->GetSumOfEventSWeight(i)
                  << std::endl;
      }

   std::cout << std::endl;

   // import this new dataset with sWeights
   std::cout << "import new dataset with sWeights" << std::endl;
   ws->import(*data, Rename("dataWithSWeights"));


}

void MakePlots(RooWorkspace* ws){

   // Here we make plots of the discriminating variable (invMass) after the fit
   // and of the control variable (isolation) after unfolding with sPlot.
   std::cout << "make plots" << std::endl;

   // make our canvas
   TCanvas* cdata = new TCanvas("sPlot","sPlot demo", 400, 600);
   cdata->Divide(1,3);

   // get what we need out of the workspace
   RooAbsPdf* model = ws->pdf("model");
   RooAbsPdf* zModel = ws->pdf("zModel");
   RooAbsPdf* qcdModel = ws->pdf("qcdModel");

   RooRealVar* isolation = ws->var("isolation");
   RooRealVar* invMass = ws->var("invMass");

   // note, we get the dataset with sWeights
   RooDataSet* data = (RooDataSet*) ws->data("dataWithSWeights");

   // this shouldn't be necessary, need to fix something with workspace
   // do this to set parameters back to their fitted values.
   model->fitTo(*data, Extended() );

   //plot invMass for data with full model and individual components overlaid
   //  TCanvas* cdata = new TCanvas();
   cdata->cd(1);
   RooPlot* frame = invMass->frame() ;
   data->plotOn(frame ) ;
   model->plotOn(frame) ;
   model->plotOn(frame,Components(*zModel),LineStyle(kDashed), LineColor(kRed)) ;
   model->plotOn(frame,Components(*qcdModel),LineStyle(kDashed),LineColor(kGreen)) ;

   frame->SetTitle("Fit of model to discriminating variable");
   frame->Draw() ;


   // Now use the sWeights to show isolation distribution for Z and QCD.
   // The SPlot class can make this easier, but here we demonstrate in more
   // detail how the sWeights are used.  The SPlot class should make this
   // very easy and needs some more development.

   // Plot isolation for Z component.
   // Do this by plotting all events weighted by the sWeight for the Z component.
   // The SPlot class adds a new variable that has the name of the corresponding
   // yield + "_sw".
   cdata->cd(2);

   // create weighted data set
   RooDataSet * dataw_z = new RooDataSet(data->GetName(),data->GetTitle(),data,*data->get(),0,"zYield_sw") ;

   RooPlot* frame2 = isolation->frame() ;
   dataw_z->plotOn(frame2, DataError(RooAbsData::SumW2) ) ;

   frame2->SetTitle("isolation distribution for Z");
   frame2->Draw() ;

   // Plot isolation for QCD component.
   // Eg. plot all events weighted by the sWeight for the QCD component.
   // The SPlot class adds a new variable that has the name of the corresponding
   // yield + "_sw".
   cdata->cd(3);
   RooDataSet * dataw_qcd = new RooDataSet(data->GetName(),data->GetTitle(),data,*data->get(),0,"qcdYield_sw") ;
   RooPlot* frame3 = isolation->frame() ;
   dataw_qcd->plotOn(frame3,DataError(RooAbsData::SumW2) ) ;

   frame3->SetTitle("isolation distribution for QCD");
   frame3->Draw() ;

   //  cdata->SaveAs("SPlot.gif");

}