LauFitDataTree.cc

Go to the documentation of this file.

// Copyright University of Warwick 2004 - 2013.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

// Authors:
// Thomas Latham
// John Back
// Paul Harrison

#include <cmath>
#include <cstdlib>
#include <iostream>

#include "TFile.h"
#include "TLeaf.h"
#include "TString.h"
#include "TSystem.h"

#include "LauFitDataTree.hh"

ClassImp(LauFitDataTree)


LauFitDataTree::LauFitDataTree(const TString& rootFileName, const TString& rootTreeName) :
        rootFileName_(rootFileName),
        rootTreeName_(rootTreeName),
        rootFile_(0),
        rootTree_(0),
        eventList_(0)
{
        if (rootFileName_ != "" && rootTreeName_ != "") {
                this->openFileAndTree();
        }
}

LauFitDataTree::~LauFitDataTree()
{
        if (rootFile_ && rootFile_->IsOpen()) {
                delete rootTree_; rootTree_ = 0;
                delete eventList_; eventList_ = 0;
        }
        delete rootFile_; rootFile_ = 0;
}

void LauFitDataTree::openFileAndTree()
{
        // first check whether we've already opened up the file or not
        if (!rootFile_) {
                // if not, first check the filename and if all ok create the file
                if (rootFileName_ == "") {
                        std::cerr << "ERROR in LauFitDataTree::openFileAndTree : Bad filename supplied, not creating file or tree." << std::endl;
                        return;
                }
                rootFile_ = TFile::Open(rootFileName_, "read");
                if (!rootFile_ || rootFile_->IsZombie() || !rootFile_->IsOpen()) {
                        std::cerr << "ERROR in LauFitDataTree::openFileAndTree : Problem opening file \"" << rootFileName_ << "\" for reading." << std::endl;
                        return;
                }
        }
        // check whether we've already created the tree
        if (!rootTree_) {
                // if not change to the file's directory and grab the tree
                rootFile_->cd();
                rootTree_ = dynamic_cast<TTree*>(rootFile_->Get(rootTreeName_));
                if (!rootTree_) {
                        std::cerr << "ERROR in LauFitDataTree::openFileAndTree : Problem reading tree \"" << rootTreeName_ << "\" from file." << std::endl;
                        return;
                }
                rootTree_->SetDirectory(rootFile_);
        }
}

Bool_t LauFitDataTree::findBranches()
{
        if (!rootTree_) {
                std::cerr << "ERROR in LauFitDataTree::findBranches : Invalid pointer to data tree." << std::endl;
                return kFALSE;
        }

        // this method should only be called once
        if ( ! leafNames_.empty() ) {
                std::cerr << "ERROR in LauFitDataTree::findBranches : Branches already found, not running again." << std::endl;
                return kFALSE;
        }

        leafNames_.clear();
        eventData_.clear();
        eventDataOut_.clear();
        leaves_.clear();
        treeEvents_.clear();
        fakeEvents_.clear();

        const UInt_t numBranches(this->nBranches());
        eventData_.reserve( numBranches );

        TObjArray* pLeaves = rootTree_->GetListOfLeaves();
        if (!pLeaves) {
                std::cerr << "ERROR in LauFitDataTree::findBranches : Problem retrieving leaves from the tree." << std::endl;
                return kFALSE;
        }
        TObjArray& leaves = *pLeaves;

        const UInt_t numLeaves = static_cast<UInt_t>( leaves.GetEntries() );
        if ( numBranches != numLeaves ) {
                std::cerr << "ERROR in LauFitDataTree::findBranches : List of leaves is not the same size as the number of branches - this is strange!" << std::endl;
                std::cerr << "                                      : There are " << numLeaves << " leaves and " << numBranches << " branches" << std::endl;
                return kFALSE;
        }

        for (UInt_t iLeaf = 0; iLeaf < numBranches; ++iLeaf) {

                TLeaf * leaf = dynamic_cast<TLeaf*>(leaves[iLeaf]);

                // we can't deal with arrays
                Int_t size = leaf->GetNdata();
                if (size != 1) {
                        std::cerr << "ERROR in LauFitDataTree::findBranches : Tree has array branches, can't deal with those." << std::endl;
                        return kFALSE;
                }

                // find the name and type of the leaf
                TString name = leaf->GetName();

                // initialise an entry in the maps to hold the value
                leafNames_[ name ] = iLeaf;
                leaves_.push_back( leaf );
                eventData_.push_back( 0.0 );
                eventDataOut_[ name ] = 0.0;
        }

        const UInt_t numEvents(this->nTotalEvents());

        std::cout << "INFO in LauFitDataTree::findBranches : Finished finding data tree branches." << std::endl;
        std::cout << "                                     : Tree contains " << numBranches << " branches and a total of " << numEvents << " events." << std::endl;

        return kTRUE;
}

void LauFitDataTree::readExperimentData(UInt_t iExpt)
{
        if (!rootTree_) {
                std::cerr << "ERROR in LauFitDataTree::readExperimentData : Invalid pointer to data tree." << std::endl;
                return;
        }

        rootFile_->cd();

        if (!eventList_) {
                eventList_ = new TEventList("elist","elist");
                eventList_->SetDirectory(rootFile_);
        }

        TString listName(eventList_->GetName());
        listName.Prepend(">>");
        TString selection("iExpt==");
        selection += iExpt;

        std::cout << "INFO in LauFitDataTree::readExperimentData : Setting tree to experiment number " << iExpt << "." << std::endl;
        rootTree_->Draw(listName,selection);
        std::cout << "                                           : Found " << this->nEvents() << " events." << std::endl;

        this->loadData();
}

UInt_t LauFitDataTree::nBranches() const
{
        if ( rootTree_ ) {
                return static_cast<UInt_t>(rootTree_->GetNbranches());
        } else if ( ! fakeEvents_.empty() ) {
                std::vector<LauEventData>::const_iterator fakeIter = fakeEvents_.begin();
                return (*fakeIter).size();
        } else {
                return 0;
        }
}

void LauFitDataTree::appendFakePoints( const std::vector<Double_t>& xCoords, const std::vector<Double_t>& yCoords)
{
        if ( xCoords.size() != yCoords.size() ) {
                std::cerr << "ERROR in LauFitDataTree::appendFakePoints : The co-ordinate vectors are not the same size!" << std::endl;
                gSystem->Exit(EXIT_FAILURE);
        }

        // clear the list of fake events and reserve enough space
        fakeEvents_.clear();
        fakeEvents_.reserve(xCoords.size());

        // loop over the event data zeroing everything
        for (LauEventData::iterator iter = eventData_.begin(); iter != eventData_.end(); ++iter) {
                *iter = 0.0;
        }

        // find the vector indices for the DP co-ordinates
        LauNameIndexMap::const_iterator iter = leafNames_.find( "m13Sq" );
        if ( iter == leafNames_.end() ) {
                std::cerr << "ERROR in LauFitDataTree::appendFakePoints : Can't find entry \"m13Sq\" in event data map." << std::endl;
                gSystem->Exit(EXIT_FAILURE);
        }
        UInt_t m13SqIdx = iter->second;

        iter = leafNames_.find( "m23Sq" );
        if ( iter == leafNames_.end() ) {
                std::cerr << "ERROR in LauFitDataTree::appendFakePoints : Can't find entry \"m23Sq\" in event data map." << std::endl;
                gSystem->Exit(EXIT_FAILURE);
        }
        UInt_t m23SqIdx = iter->second;

        // TODO - is this really needed?  Comment out for now.
        /*
        // and a fake tagCat which we don't care about, but the program
        // needs a number here for comptibility with conventional events.
        iter = leafNames_.find( "tagCat" );
        if ( iter == leafNames_.end() ) {
                std::cerr << "ERROR in LauFitDataTree::appendFakePoints : Can't find entry \"tagCat\" in event data map." << std::endl;
                gSystem->Exit(EXIT_FAILURE);
        }
        UInt_t m23SqIdx = iter->second;
         */

        for (UInt_t iEvt = 0; iEvt < xCoords.size() ; ++iEvt) {

                // set the values
                eventData_[ m13SqIdx ] = xCoords[iEvt];
                eventData_[ m23SqIdx ] = yCoords[iEvt];
                //eventData_[ tagCatIdx ] = 0.0;

                // append to the list of fake events
                fakeEvents_.push_back(eventData_);
        }
}

Bool_t LauFitDataTree::haveBranch(const TString& name) const
{
        LauNameIndexMap::const_iterator iter = leafNames_.find( name );
        return (iter != leafNames_.end());
}

void LauFitDataTree::disableAllBranches() const
{
         rootTree_->SetBranchStatus("*", 0);
}

void LauFitDataTree::enableAllBranches() const
{
         rootTree_->SetBranchStatus("*", 1);
}

void LauFitDataTree::enableBranch(const TString& name) const
{
         if (this->haveBranch(name)) {rootTree_->SetBranchStatus(name, 1);}
}

void LauFitDataTree::disableBranch(const TString& name) const
{
         if (this->haveBranch(name)) {rootTree_->SetBranchStatus(name, 0);}
}

void LauFitDataTree::loadData()
{
        // Clear any data from previous expts
        treeEvents_.clear();

        // Reserve enough space in the vector
        Long64_t nEvts = this->nEvents();
        treeEvents_.reserve( nEvts );

        const UInt_t nLeaves = this->nBranches();

        // Loop through the tree and store the data
        for ( Long64_t iEvt(0); iEvt < nEvts; ++iEvt ) {

                // Find which entry from the full tree contains the requested event
                Long64_t iEntry = eventList_ ? eventList_->GetEntry(iEvt) : iEvt;
                if (iEntry<0) { // this shouldn't happen, but just in case...
                        std::cerr << "ERROR in LauFitDataTree::loadData : Requested event not found." << std::endl;
                        gSystem->Exit(EXIT_FAILURE);
                }

                // Then retrieve that entry from the tree
                rootTree_->GetEntry(iEntry);

                // Copy the leaf values into eventData_
                for ( UInt_t iLeaf(0); iLeaf < nLeaves; ++iLeaf ) {

                        const TLeaf * leaf =  leaves_[ iLeaf ];

                        Double_t value = leaf->GetValue();
                        if ( std::isnan( value ) || std::isinf( value ) ) {
                                std::cerr << "ERROR in LauFitDataTree::loadData : Event " << iEvt << " has infinite or NaN entry for variable " << leaf->GetName() << std::endl;
                                gSystem->Exit(EXIT_FAILURE);
                        }

                        eventData_[ iLeaf ] = value;
                }

                // Store the event in the vector
                treeEvents_.push_back( eventData_ );
        }
}

void LauFitDataTree::readAllData()
{
        delete eventList_; eventList_ = 0;
        this->loadData();
}

const LauFitData& LauFitDataTree::getData(UInt_t iEvt) const
{
        // For a given event, within the current experiment,
        // retrieve the input fit variables

        const UInt_t numTreeEvents = this->nEvents();
        const UInt_t numFakeEvents = this->nFakeEvents();

        // Does the requested event come from the tree or from the fake events list?
        if ( iEvt < numTreeEvents ) {
                if ( iEvt > treeEvents_.size() ) { // this shouldn't happen, but just in case...
                        std::cerr << "ERROR in LauFitDataTree::getData : Requested event, " << iEvt << ", not found." << std::endl;
                        gSystem->Exit(EXIT_FAILURE);
                }
                // Retrieve that event from the vector
                eventData_ = treeEvents_[iEvt];
        } else if ( iEvt < (numTreeEvents + numFakeEvents) ) {
                // Find the entry from the fake events list contains the requested event
                UInt_t iEntry = iEvt - numTreeEvents;
                // Copy the entry into eventData_
                eventData_ = fakeEvents_[iEntry];
        } else {
                std::cerr << "ERROR in LauFitDataTree::getData : Requested event " << iEvt << " not found for " << rootTreeName_ << std::endl;
                gSystem->Exit(EXIT_FAILURE);
        }

        for ( LauNameIndexMap::const_iterator iter = leafNames_.begin(); iter != leafNames_.end(); ++iter ) {
                const TString& name = iter->first;
                UInt_t index = iter->second;

                LauFitData::iterator outIter = eventDataOut_.find( name );
                if ( outIter == eventDataOut_.end() ) {
                        std::cerr << "ERROR in LauFitDataTree::getData : Something very strange happening." << std::endl;
                        gSystem->Exit(EXIT_FAILURE);
                }

                outIter->second = eventData_[ index ];
        }
        return eventDataOut_;
}