addEtaToTree.cxx

// Compile with ROOT 6 ACliC with
// .L addEtaToTree.cxx+
#include <iostream>
#include <map>
#include <vector>
#include <assert.h>

#include <TFile.h>
#include <TTree.h>
#include <TBranch.h>
#include <TString.h>
#include <TMath.h>
#include <TTreeReader.h>
#include <TSystem.h>


Double_t Eta(Double_t cosTheta) {

  Double_t theta = TMath::ACos(cosTheta);
  Double_t tanThetaDividedBy2 = TMath::Tan(theta/2.0);
  Double_t eta = -TMath::Log(tanThetaDividedBy2);

  return eta;
}

std::vector<std::pair<TBranch *, Double_t *> > fetchCosThetaBranches(TTree * tree) {

  std::vector<std::pair<TBranch *, Double_t *> > ret_vect;

  for(TObject * object: *(tree->GetListOfBranches())) {
    TBranch * branch = (TBranch *)object;
    TString branchName = TString(branch->GetName());

    if(branchName.Contains("CosTheta")) {
      Double_t * _tmp = new Double_t(-1.0);
      tree->SetBranchAddress(branchName.Data(), _tmp);

      ret_vect.push_back(std::pair<TBranch*, Double_t*>(branch, _tmp));
    }
  }

  return ret_vect;

}

std::vector<std::pair<TBranch *, Double_t *> > createEtaBranches(TTree * tree) {

  std::vector<std::pair<TBranch *, Double_t *> > ret_vect;

  for(TObject * object: *(tree->GetListOfBranches())) {
    TBranch * branch = (TBranch *)object;
    TString branchName = TString(branch->GetName());

    if(branchName.Contains("CosTheta")) {
      TString newBranchName = branchName.ReplaceAll("CosTheta", "ETA");
      Double_t * _tmp = new Double_t(-1.0);
      TBranch * newBranch = tree->Branch(newBranchName.Data(), _tmp);

      ret_vect.push_back(std::pair<TBranch*, Double_t*>(newBranch, _tmp));
    }
  }

  return ret_vect;
}


Int_t createCopyOfFile(TString source, TString destination) {
  return gSystem->CopyFile(source.Data(), destination.Data());
}

void fillCurrentEtaBranches(std::vector<std::pair<TBranch *, Double_t *> > branches) {

  for(auto& i: branches) {
    TBranch * branch = i.first;
    Double_t * branchAddress = i.second;

    TString branchName = TString(branch->GetName());
    if(branchName.Contains("CosTheta")) {
      for(auto& j: branches) {
        if(TString(j.first->GetName()) == branchName.ReplaceAll("CosTheta", "ETA")) {
          *j.second = Eta(*i.second);
          j.first->Fill();
        }
      }
    }
  }
}

Int_t addEtaToTree(TString initialPath, TString copyToPath, TString treeDirectory) {

  std::cout << "Copying file returned status code: "
            << createCopyOfFile(initialPath, copyToPath) << std::endl;

  TFile * newFile = TFile::Open(copyToPath.Data(), "update");
  assert(newFile);
  TTree * newTree = (TTree*)newFile->Get(treeDirectory.Data());
  assert(newTree);

  std::vector<std::pair<TBranch *, Double_t *> > branches = createEtaBranches(newTree);
  std::vector<std::pair<TBranch *, Double_t *> > cosThetaBranches = fetchCosThetaBranches(newTree);

  for(auto& i: cosThetaBranches) {
    branches.push_back(i);
  }

  // Not sure if this is a thing yet...
  for(Int_t i = 0; i < newTree->GetEntries(); ++i) {
    newTree->GetEntry(i);
    fillCurrentEtaBranches(branches);
  }

  newFile->Write();
  newFile->Close();

  return 0;

}