EWP-BplusToKstMuMu-AngAna/Code/FCNCFitter/sources/Run/toysfit.cc

//Renata Kopecna

#include <toysfit.hh>
#include <constants.hh>
#include <fitter.hh>
#include <bu2kstarmumu_pdf.hh>
#include <bu2kstarmumu_plotter.hh>
#include <folder.hh>
#include "generatetoys.hh"
#include <design.hh>
#include <helpers.hh>
#include <paths.hh>
#include <event.hh>

#include <spdlog.h>

//Fit the toys and save results somewhere
int toysfit(fcnc::options opts, unsigned int nBins, bool fitReference, basic_params params, bool onlySig, bool onlyBkg){

    opts.job_id = params.jobID; //Just to be sure
    spdlog::info("[TOYSFIT]\t Starting fit with jobID {0:d}", opts.job_id);

    //actually fit the toy events
    const bool do_fit = true;
    if(!do_fit) spdlog::warn("Fitting of the toy sample is disabled. Set hardcoded 'do_fit' to true in toysfit.cc");
    //It is useful to save the generated toys in order to debug
    const bool SaveTheToys = params.index==-1;
    const bool SaveProjections = params.index==-1;
    opts.plot_folder = get_ToysFitPlot_path(params);
    const bool save_init = (params.index==-1 || params.jobID ==0);
    if(!save_init)spdlog::warn("Saving the initialised parameters for the toy generation is disabled. Set hardcoded 'save_init' to true in toysfit.cc");

    const bool bkgFromHighMass = true;
    const bool bkgFromLowMass = false;

    const bool SimultaneousFit = true;

    //Fit ranges for angles
    double angleStepSize = 0.01;
    double PprimeRangeScale = 10.0;
    double angleRange = 1.0;
    if (params.usePprime) angleRange  *= PprimeRangeScale;

    double swaveStepSize = 0.1;

    //Set what parameters to fix
    fixConstr Bmass_FC(!fitReference || onlyBkg,false); //Fix/constrain B mass
    fixConstr f_sig_FC(false,false); //Fix/constrain f_sig?
    fixConstr ang_params_FC(false,false); //Fix/constrain angular parameters?
    fixConstr mass_params_FC(true,false); //Fix/constrain mass parameters?
    fixConstr bkg_ang_FC(false,false); //Fix/constrain angular background?
    fixConstr bkg_Kpi_FC(false,false); //Fix/constrain Kst mass background?
    fixConstr bkg_mass_FC(false, false); //Fix/constrain mass background?
    //This boolean decides whether to fit the bkg with order of five or of two //TODO: Should be less hardcody
    bool setCtkToTwo = !fitReference;
    opts.bkg_order_costhetak = 5;
    if(params.folding == 4 && !fitReference){
        setCtkToTwo = false;
        opts.bkg_order_costhetak = 4;
    }
    opts.flat_bkg = false;
    opts.fit_full_angular_bkg = true; //Do you want to fold also the bkg?

    //Init angular parameters according to SM (true) or previous measurements (false)?
    opts.initSM = true;

    //Weighted fit
    opts.weighted_fit = true;

    //Set s-wave parameters to fix
    fixConstr gammakstar_FC(true,false); //Fix/constrain gammakstar
    fixConstr gammakstarplus_FC(true,false); //Fix/constrain gammakstar
    fixConstr sWave_FC(onlyBkg || !fitReference ,false); //Fix/constrain sWave parameters?
    fixConstr FS_FC(onlyBkg || !fitReference,false);
    //FS is fixed no matter what
    //Fix to something else afterwards
    bool changeFS = false;

    //Do or do not fit s-wave
    opts.swave = !onlyBkg;//TODO
    //Fit kpi or not
    opts.fit_mkpi = opts.swave || onlyBkg;

    //Quick sanity checks
    if (onlyBkg && onlySig){
        spdlog::error("Cannot have onlyBkg && onlySig! Returning.");
        return 1;
    }

    if (bkgFromHighMass && bkgFromLowMass){
        spdlog::error("Cannot have background taken only from upper and only from lower mass sideband at the same time! Returning.");
        return 1;
    }

    //Save plots?
    opts.write_eps = SaveProjections;
    opts.write_pdf = false;
    //Plot less bins for signal channel as the stats are sad there
    if (!fitReference) opts.plots_m_bins = 30;
    //Plot less bins for signal channel as the stats are sad there
    if (!fitReference) opts.plots_mkpi_bins = 20;

    //Fit both mass and angles, use lambda for bkg fit
    opts.only_angles = false;
    opts.only_Bmass = false;
    opts.only_mkpi = false;

    //Use weights and squared hesse
    opts.shift_lh = false;
    opts.use_angular_acc = false;
    opts.hesse_postrun = true;
    opts.squared_hesse = true;//!(params.index == -1);// fitReference;
    opts.minos_errors = false;//(params.index == -1); //!fitReference;
    opts.multiply_eff = false; //Toy generated events have saved acceptance weights!

    //Fit angular BKG
    opts.individual_penalties = false;  //forces individual probabilities for sig and bkg to be positive

    //How would you like to fit mkpi?
    opts.use_mkpi = false;
    opts.simple_mkpi = false;
    opts.isobar = false;
    opts.generate_mkpi = opts.fit_mkpi || opts.use_mkpi;

    //Get names for the init parameters file and the fitted parameters file
    const std::string results_file = final_result_name_toys(params.jobID,fitReference, nBins, SimultaneousFit, params, opts.run, opts.only_angles, opts.only_Bmass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);
    std::string params_file =  init_params_name_toys(params.jobID,fitReference, nBins, SimultaneousFit, params, opts.run, opts.only_angles, opts.only_Bmass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);

    if (existsTest(results_file) && params.index != -1){
        spdlog::warn("[SKIP]\t\tJob already run as root-file '"+results_file+"' exists, continue to next job!");
        return 0;
    }

    //Set the available PDFs: this depends on the selected Run for now, so if only one is selected, run over one pdf, if both, use two
    std::vector<UInt_t> pdf_idx;
    if (params.Run == 1 || params.Run == 12) pdf_idx.push_back(1);
    if (params.Run == 2 || params.Run == 12) pdf_idx.push_back(2);

    //we are good to go, now; how many individual pdfs are used?
    const UInt_t nPDFs = pdf_idx.size();

    //current fitter, plotter, parameteres and pdfs:
    fcnc::fitter f(&opts);
    fcnc::options theOptions[nPDFs];
    fcnc::bu2kstarmumu_plotter * thePlotter[nPDFs];
    std::vector<fcnc::parameters*> theParams [nBins];
    std::vector<fcnc::pdf*>        theProbs  [nBins];
    std::vector< std::vector<fcnc::event>* >    selection[nBins];

    //Initialize common parameters
    std::vector<std::string> common_params = param_string(opts, false); //All the angular observables are shared
    if(SimultaneousFit) f.set_common_parameters(common_params);
    //--------------------------------
    //          Read events
    //--------------------------------

    //Init vector for the events to be saved, in case required
    std::vector<fcnc::event> eventsToSave;

    //Loop over PDFs and initialize parameters
    for(unsigned n = 0; n < nPDFs; n++){
        unsigned int idx = pdf_idx.at(n);
        opts.run = idx; //Set proper run to options
        spdlog::debug("Run {0:d}", opts.run);
        opts.name = get_ToysFit_label(params.jobID, fitReference, -1, nBins, SimultaneousFit, params, opts.run, opts.only_angles, opts.only_Bmass, onlySig, onlyBkg, bkgFromLowMass, bkgFromHighMass);

        opts.update_angle_ranges(); //Set angles in options back to defaults
        opts.update_efficiencies = true; //This ensures the acceptance weights to be taken into account
        opts.plot_label = "Toys";

        //Get options separate for all PDFs
        theOptions[n] = opts;

        for(UInt_t b = 0; b < nBins; b++){
            opts.q2_min = theOptions[n].TheQ2binsmin.front();
            opts.q2_max = theOptions[n].TheQ2binsmax.back();

            //-----------------------------------------
            //      Initialize 1D parameters
            //-----------------------------------------

            //create parameter sets:
            fcnc::bu2kstarmumu_parameters * leParameters = new fcnc::bu2kstarmumu_parameters(&theOptions[n]);

            //create PDFs
            fcnc::bu2kstarmumu_pdf * lePDF = new fcnc::bu2kstarmumu_pdf(&theOptions[n], leParameters);

            //cache the folding and set to -1 for loading mass parameters from non-folded fits to jpsi, MC, etc.
            int cached_folding = params.folding;
            params.folding = -1;

            //Inititalize q2 by hand
            leParameters->eff_q2.init_fixed( bin_center_q2(theOptions[n],b));

            //Init the fraction betweeen sig and bkg
            if (onlySig)      leParameters->f_sig.init_fixed(1.0);
            else if (onlyBkg) leParameters->f_sig.init_fixed(0.0);
            else{ //Or both, in that case load the f_sig from mass only fit
                initialize_parameters_from_MassFit(leParameters,fitReference,idx,b,{"f_sig"},params,f_sig_FC);
            }

            //**********************************//
            //                                  //
            //              Bmass               //
            //                                  //
            //**********************************//

            //Set the Bmass separatelly as it's range is important
            std::string fileName_dataMass = final_result_name_mass(true, 1, true, params, params.Run);
            leParameters->init_Bmass(fileName_dataMass,idx,0.5,Bmass_FC);


            //Init B mass fit from Jpsi MC            
            if (!onlyBkg){
                initialize_parameters_from_MCfit(leParameters,true,idx,b, {"alpha_1","alpha_2","n_1","n_2"},params, mass_params_FC);

                //Init the ratio of sigmas in signal MC/reference MC
                initialize_parameters_from_MassFit(leParameters,true,idx,b,{"m_sigma_1"},params,fixConstr(!fitReference,false));
                if (!fitReference){
                    leParameters->m_scale.init_fixed(get_sigmaRatio_fromMC(params,nBins,b,pdf_idx.at(n)));
                }
                else{
                    leParameters->m_scale.init_fixed(1.0);
                }
            }

            //Init the mass background from mass fit to reference for now
            //TODO: set as a proper vector depending on using 2 lambdas or anything
            if (!onlySig) initialize_parameters_from_MassFit(leParameters,fitReference,idx,b,{"m_lambda"},params,bkg_mass_FC);

            //**********************************//
            //                                  //
            //              Angles              //
            //                                  //
            //**********************************//

            //Init angular parameters from MC
            if (!onlyBkg){
                if (opts.initSM){
                    leParameters->init_angular_parameters(nBins,b,ang_params_FC.fix ? 0.0 : angleStepSize,1.0,false); //TODO: make it better
                }
                else{
                    if (fitReference) leParameters->init_ang_parameters_fromRefDavid(b,1.0,0.01);
                    else initialize_parameters_from_MCfit(leParameters,false,idx,b,param_string(opts, true),params, ang_params_FC);
                }
            }


            //Init the bkg parameters from the BKG fit file
            if (!onlySig){
                if (opts.flat_bkg) leParameters->use_default_bkg();
                else{
                    //First init the background for angles for all orders
                    initialize_parameters_from_BkgFit(leParameters, true,
                                                      bkgFromLowMass, bkgFromHighMass,false, 12,b,
                                                      PAR_BKG_STRING(-1,opts.bkg_order_costhetal,opts.bkg_order_costhetak),
                                                      params,fixConstr(true,false));
                    //Now let float only the ones actually used by the folding
                    initialize_parameters_from_BkgFit(leParameters, true,
                                                      bkgFromLowMass, bkgFromHighMass,false, 12,b,
                                                      PAR_BKG_STRING(cached_folding,opts.bkg_order_costhetal,opts.bkg_order_costhetak),
                                                      params,bkg_ang_FC);
                    if (opts.bkg_order_costhetak>=3 && cached_folding!=4){
                        leParameters->cbkgctk3.init(-2.25,-3.25,1.5,0.05);
                    }
                    //Then init the background for m_Kpi                    
                    if (opts.generate_mkpi){
                        initialize_parameters_from_BkgFit(leParameters, true,
                                                          bkgFromLowMass, bkgFromHighMass,
                                                          true, 12,
                                                          b,param_string_bkg_mkpi(),params,bkg_Kpi_FC);
                    }
                }
            }

            //**********************************//
            //                                  //
            //         S wave and Kstar         //
            //                                  //
            //**********************************//

            //Init the Kstar mass fit
            if(opts.generate_mkpi){
                //p-wave
                leParameters->init_mkpi_pWave_parameters(fitReference,gammakstar_FC.fix ? 0.0:0.01);
                //s-wave
                if (opts.swave) leParameters->init_mkpi_sWave_parameters(fitReference,gammakstarplus_FC.fix ? 0.0:0.001);
            }
            if (opts.swave){
                leParameters->init_sWave_parameters(sWave_FC.fix ? 0.0 : swaveStepSize);
                std::string RefMassFile = final_result_name_mass(true,1,true,params,params.Run);
                if (!fitReference) leParameters->FS.init(0.25,-0.25,1.0,FS_FC.fix ? 0.0 : 0.1);
                else leParameters->get_param_from_rootfile(RefMassFile, {"FS"}, idx, 0,FS_FC);
                if (changeFS) leParameters->FS.set_constant();
            }

            //set folding back to configured value using the cached value from above:
            params.folding = cached_folding;

            spdlog::info("[PDF{0:d}]\tSaved PDF and parameters!", n);

            //Save the events now
            std::vector<fcnc::event> * leEvents = new std::vector<fcnc::event>;
            if(theOptions[n].job_id >= 0){
                //Seed is opts->get_job_id()*opts->ncores+thread_id+12345 if opts->static_seed, else take from the internal clock backwards
                std::vector<fcnc::event> evts = generateToys(params, b, pdf_idx.at(n),
                                                             leParameters,theOptions[n]);
                //copy generated events to pointer vector object
                leEvents->clear();
                //count the number of dismissed events:
                unsigned int dismissed = 0;
                for(auto evt: evts){
                    theOptions[n].update_angle_ranges(); //Update angle ranges based on the folding
                    //filter events, that are outside the defined angular range:
                    if(!filterFldFour(&evt, &theOptions[n])){ //I don't think this should be here at all
                        dismissed++;
                        continue;
                    }
                    leEvents->push_back(evt);
                }
                spdlog::info("{0:d} generated events are outside the angular range and are dismissed from the sample", dismissed);
                assert(leEvents->size()+dismissed == evts.size());
            }
            else{ //Just load all events TODO
                std::vector<fcnc::event> evts = fcnc::load_events("/home/lhcb/kopecna/B2KstarMuMu/code/ewp-Bplus2Kstmumu-AngAna/FCNCfitter/Toys/Toy_toyfit__JpsiFit_1BIN_Run12_SimultaneousFit_6.root","Events", -1);                 //copy generated events to pointer vector object
                leEvents->clear();
                for(auto evt: evts)leEvents->push_back(evt);
                assert(leEvents->size() == evts.size());
            }

            spdlog::info("[PDF{0:d}]\tFinished generating the events: {1:d}", n, leEvents->size());
            if (!onlySig && setCtkToTwo){
                theOptions[n].bkg_order_costhetak = 2;
                leParameters->use_default_bkg();
                leParameters->init_angular_background_parameters(fitReference,0.01);
            }
            if (changeFS){
                leParameters->FS.init_fixed(0.0);
                leParameters->SS1.init_fixed(0.0);
                leParameters->SS2.init_fixed(0.0);
                leParameters->SS3.init_fixed(0.0);
                leParameters->SS4.init_fixed(0.0);
                leParameters->SS5.init_fixed(0.0);
            }

            leParameters->take_current_as_start();
            lePDF->load_coeffs_eff_phsp_4d();
            lePDF->update_cached_normalization(leParameters);
            theParams [b].push_back(leParameters);
            theProbs  [b].push_back(lePDF);
            lePDF->update_cached_efficiencies(leParameters, leEvents);

            if (SaveTheToys){
                //save all events into one vector, except for the fifth (extra wide bin) in the 5BIN scheme. this would be double counting events
                if(!(nBins == 5 && b == 4))eventsToSave.insert(eventsToSave.end(),leEvents->begin(),leEvents->end());
            }

            //save event vector in vector
            selection[b].push_back(leEvents);
        } //end loop over bins
        theOptions[n].update_efficiencies = false;//Prevent the weights to be applied several more times in fitter::fit
    } //end loop over PDFs

    //validate correct saving to of all events:
    for(unsigned n = 0; n < nPDFs; n++){
        for(unsigned int b = 0; b < nBins; b++){
            if(selection[b].at(n)->size() > 0){
                spdlog::info("[PDF{0:d}]\t[BIN{1:d}]\tDone!", n, b);
            }
            else{
                spdlog::critical("No events found for PDF={0:d} and q2-bin={1:d} Exit!", n, b);
                assert(0);
            }
        }
    }
    spdlog::info("Finished configurating all parameters, defining all pdfs and generating all toy event samples:");

    //If required to save the generated events, save them
    if (SaveTheToys){
        std::string eventsFile = get_finalToys_file(fitReference,nBins,SimultaneousFit,params,params.Run);
        replace(eventsFile,".root","_"+std::to_string(params.jobID)+".root"); //Add the number of the job to the file name
        fcnc::save_events(eventsFile,eventsToSave);
    }
    //Save the initial toy values in a root file

    std::vector<int>tmp[nBins]; //Create a vector array with zeros, so you can save init values before the fit
    for(unsigned int b = 0; b < nBins; b++)tmp[b] = std::vector<int>(nPDFs, 0);
    if(save_init) save_results(params_file,nBins,pdf_idx,tmp,theParams,SimultaneousFit,&opts);

    //Save the fit results
    std::vector<int>fit_results[nBins];

    //fit all bins:
    for(unsigned int b = 0; b < nBins; b++){
        spdlog::info("");
        spdlog::warn("[START]\tStart fit BIN{0:d}", b);

        //Int for fit status
        int fitresult = 0;

        if(!SimultaneousFit){
            for(UInt_t n = 0; n < nPDFs; n++){

                //Delete the texFile
                std::string tag = get_ToysFit_label(params.jobID, fitReference, b, nBins,
                                                    false, params, opts.run,
                                                    opts.only_angles, opts.only_Bmass,
                                                    onlySig, onlyBkg,
                                                    bkgFromLowMass, bkgFromHighMass);

                spdlog::info("[FIT{0:d}]\tRunning the fitter...", n);
                if(do_fit)fitresult = f.fit(theProbs[b].at(n), theParams[b].at(n), selection[b].at(n),tag);
                fit_results[b].push_back(fitresult);
            }
        }
        else{
            spdlog::info("[FIT]\tFitting simultaenously....");
            std::string tag = "";
            for(UInt_t n = 0; n < nPDFs; n++)for(UInt_t e = 0; e < selection[b].at(n)->size(); e++)spdlog::trace("Event for fitting: pdf={0:d}\t bin={1:d}\t evt={2:d}\t ctk={3:f}\t ctl={4:f}\t phi={5:f}\t m={6:f}\t mkpi={7:f}", n, b, e, selection[b].at(n)->at(e).costhetak, selection[b].at(n)->at(e).costhetal, selection[b].at(n)->at(e).phi, selection[b].at(n)->at(e).m, selection[b].at(n)->at(e).mkpi);

            if(do_fit)fitresult = f.fit(theProbs[b], theParams[b], selection[b], tag);
            fit_results[b].push_back(fitresult);
            spdlog::info("Q2BIN={0:d}: LLH={1:f}", b, f.likelihood());
        }

        //Print the fit results
        spdlog::info("[BIN{0:d}]: Fitresult: {1:d}", b, fitresult);
        spdlog::info("");
    }

    //Plot stuff if wanted
    if (SaveProjections){
        //If needed, also save the projections for the toys
        //plot each set of pdfs with the data points:
        bool plotPulls = true;
        for(UInt_t b = 0; b < nBins; b++){
            for(UInt_t n = 0; n < nPDFs; n++){
                thePlotter[n] = new fcnc::bu2kstarmumu_plotter(&theOptions[n]);

                //plot pdf and events:
                std::string eps_label = get_ToysFit_label(params.jobID, fitReference, b, nBins,
                                                          SimultaneousFit, params, pdf_idx.at(n),
                                                          opts.only_angles, opts.only_Bmass,
                                                          onlySig, onlyBkg,
                                                          bkgFromLowMass, bkgFromHighMass);

                theOptions[n].plot_label = "Toys";

                if (!fitReference) theOptions[n].plots_m_bins = 30;
                if (!fitReference) theOptions[n].plots_mkpi_bins = 20;
                theOptions[n].q2_label = q2_label(opts.TheQ2binsmin.at(b), opts.TheQ2binsmax.at(b));

                spdlog::info("[PLOT]\t"+eps_label);
                thePlotter[n]->SetPulls(plotPulls);
                thePlotter[n]->plot_data((fcnc::bu2kstarmumu_pdf*)theProbs[b].at(n), (fcnc::bu2kstarmumu_parameters*)theParams[b].at(n), selection[b].at(n), get_ToysFitPlot_path(params), eps_label, true);
                thePlotter[n]->plot_data((fcnc::bu2kstarmumu_pdf*)theProbs[b].at(n), (fcnc::bu2kstarmumu_parameters*)theParams[b].at(n), selection[b].at(n), get_ToysFitPlot_path(params), eps_label+"_incBkg", false);

            }

            std::vector<fcnc::bu2kstarmumu_pdf*> * prober = (std::vector<fcnc::bu2kstarmumu_pdf*> *)  & theProbs[b];

            std::string eps_label = get_ToysFit_label(params.jobID, fitReference, b, nBins,
                                                      SimultaneousFit, params, params.Run,
                                                      opts.only_angles, opts.only_Bmass,
                                                      onlySig, onlyBkg,
                                                      bkgFromLowMass, bkgFromHighMass);
            std::vector<fcnc::bu2kstarmumu_parameters*> * paramser = (std::vector<fcnc::bu2kstarmumu_parameters*> *) & theParams[b];
            thePlotter[0]->SetPulls(plotPulls);
            thePlotter[0]->plot_added_pdfs(prober, paramser, & selection[b], get_ToysFitPlot_path(params),  eps_label, true);
            thePlotter[0]->plot_added_pdfs(prober, paramser, & selection[b], get_ToysFitPlot_path(params),  eps_label+"_incBkg", false);
        }//end loop over bins
    }

    spdlog::info("[TOYFIT] Fit status results:");
    if(SimultaneousFit){
        for(UInt_t b = 0; b < nBins; b++)spdlog::info("[BIN{0:d}]: {1:d}", b, fit_results[b].at(0));
    }
    else{
        for(UInt_t b = 0; b < nBins; b++)for(UInt_t n = 0; n < nPDFs; n++)spdlog::info("[BIN{0:d}][PDF{1:d}]: {2:d}", b, n, fit_results[b].at(n));
    }
    if (params.index == -1) print_all_parameters(nBins, pdf_idx, theParams, spdlog::level::info);
    //saving results to root file
    save_results(results_file,nBins,pdf_idx,fit_results,theParams,SimultaneousFit,&opts);

    spdlog::info("[TOYFIT] Finished.");
    return 0;
}