Strona domowa

Menu

Drzewa i plik część II

Zapis klasy

Na poprzedniej lekcji zapisywaliśmy dane w TNtuple które jest uproszczonym odpowiednikiem drzew ROOTowych. Teraz przejdziemy do drzew ROOTowych. Jak wspomniałem drzewa umożliwiaja zapis danych w formie bardziej zaawansowanych struktur - tzn. tutaj jedno wejście nie odpowiada strukturze jak w TNtuple ale klasie bądź grupie klas.

Zacznijmy od relatywnie prostego przykładu - chcemy zapisać sobie cząstki typu V0 w formie obiektu TParticle, cząstki takie rozpadają się na dwie córki, które również chcemy sobie zapisać. Przykładowe makro zapisujące je w drzewie może wyglądać tak:

void tree_w(){
    TFile *file  = new TFile("data.root","recreate");
    TTree *tree = new TTree("data","data");
    TParticle *pNeu = new TParticle();
    TParticle *pPos = new TParticle();
    TParticle *pNeg = new TParticle();
    tree->Branch("v0", &pNeu);
    tree->Branch("pos",&pPos);
    tree->Branch("neg",&pNeg);
    for(int i=0;i<100;i++){
        pNeu->SetPdgCode(3122);
        pNeu->SetMomentum(2, 2, 2, 2);
        pPos->SetPdgCode(2212);
        pPos->SetMomentum(1,1,1,1);
        pNeg->SetPdgCode(211);
        pNeg->SetMomentum(1,1,1,1);
        tree->Fill();
    }
    tree->Write();
    file->Close();
}

Natomiast po zapisie nasze drzewo wygląda tak:

Rys. 1. Plikt z drzewem i obiektami klasy TParticle. Pola z wykrzyknikiem nie są fizycznie zapisywane a są wyliczane.

Zauważyć tu można, że mimo iż zapisywaliśmy obiekty tej samej klasy (TParticle) mamy osobne nazwy branchy, co pozwala na łatwą identyfikację który TParticle to która cząstka (V0 czy któraś cząstka-córka). Odczytu takiego drzewa można dokonać w formie makra:

void tree_r(){
    TFile *file  = new TFile("data.root");
    TTree *tree = (TTree*)file->Get("data");
    TParticle *pNeu = nullptr;
    TParticle *pPos = nullptr;
    TParticle *pNeg = nullptr;
    tree->SetBranchAddress("v0", &pNeu);
    tree->SetBranchAddress("pos",&pPos);
    tree->SetBranchAddress("neg",&pNeg);
    for(int i=0;i<100;i++){
        tree->GetEntry(i);
        cout<GetPdgCode()<<" "<GetPdgCode()<<" "<<
                pNeg->GetPdgCode()<<endl;
    }
}

Tablice klonów

Poprzedni przykład zakładał że mamy zawsze 3 cząstki. Co jednak jeśli zapisujemy zderzenie i raz mamy 10,  a raz 20 cząstek? W tym przypadku przydałaby się jakaś tablica z TParticle. W tym przypadku możemy to zrobić na trzy sposoby:

  • zapisać wektor naszych obiektów
  • "zwinąć" nasze obiekty do jakieś klasy reprezentującej zderzenie
  • zapisać nasze obiekty w tablicy klonów

Dwie pierwsze opcje wymagają trochę bardziej złożonych operacji i zostaną omówione później. Zobaczmy jak można więc zapisywać tablice o różnych rozmiarach przy pomocy tablicy klonów:

void DoStuff(TLorentzVector *p){
    p->SetXYZT(gRandom->Gaus(0,1),gRandom->Gaus(0,1),gRandom->Gaus(0,1),0);
}
void tree_wtc(){
    TFile *file  = new TFile("data.root","recreate");
    TTree *tree = new TTree("tree","title");
    TClonesArray *tracks = new TClonesArray("TLorentzVector",100);
    tree->Branch("tracks", &tracks);
    for(int i=0;i<100;i++){
        tracks->Clear();
        Int_t n = (Int_t)(gRandom->Uniform(5,11));
        for(int j=0;j<n;j++){
            TLorentzVector *p = (TLorentzVector*)tracks->ConstructedAt(j);
            DoStuff(p);
        }
        tree->Fill();
    }
    tree->Write();
    file->Close();
}

Odczyt takiego pliku wyglądał będzie następująco:

void tree_wtcR(){
    TFile *file  = new TFile("data.root");
    TTree *tree = (TTree*)file->Get("tree");
    TClonesArray *tracks = nullptr;
    tree->SetBranchAddress("tracks", &tracks);
    for(int i=0;i<100;i++){
        tree->GetEntry(i);
        //liczba jest wypelnaniana przez TTree::GetEntry
        //nie trzeba wiec uzywac wolniejszej metody GetEntries
        Int_t n = tracks->GetEntriesFast();
        for(int j=0;j<n;j++){
            //tablica jest wypelniana przy TTree::GetEntry wiec nie trzeba alokowac
            //pamieci
            TLorentzVector *p = (TLorentzVector*)tracks->UncheckedAt(j);
            p->Print();
        }
    }
    tree->Write();
    file->Close();
}

Można tworzyć oczywiście kilka równoległych tablic klonów w jednym drzewie, przykładowo możemy stworzyć osobno branch na primary i global tracks:

void tree_wtcR2(){
    TFile *file  = new TFile("data.root","recreate");
    TTree *tree = new TTree("tree","title");
    TClonesArray *tracksG = new TClonesArray("TLorentzVector",100);
    TClonesArray *tracksP = new TClonesArray("TLorentzVector",100);
    tree->Branch("tracks_global", &tracksG);
    tree->Branch("tracks_primary", &tracksP);
    for(int i=0;i<100;i++){
        tracksP->Clear();
        tracksG->Clear();
        Int_t n = (Int_t)(gRandom->Uniform(5,11));
        for(int j=0;j<n;j++){
            TLorentzVector *p = (TLorentzVector*)tracksP->ConstructedAt(j);
            DoStuff(p);
        }
        n = (Int_t)(gRandom->Uniform(10,21));
        for(int j=0;j<n;j++){
            TLorentzVector *p = (TLorentzVector*)tracksG->ConstructedAt(j);
            DoStuff(p);
        }
        tree->Fill();
    }
    tree->Write();
    file->Close();
}

TChain

TChain to swojego rodzaju analogia do drzewa ale rozmieszczonego w różnych plikach. Dzięki temu zamiast analizować każdy plik oddzielnie, można je połączyć (bez fizycznego łączenia plików) w jeden wirtualny plik. Przykładowe użycie TChain wygląda tak:

void tree_wtcR3(){
    TChain *tree = new TChain("tree");

    for(int i=1;i<=3;i++){
        tree->AddFile(Form("data_%i.root",i));
    }
    TClonesArray *tracksG = nullptr;
    TClonesArray *tracksP = nullptr;

    tree->SetBranchAddress("tracks_primary", &tracksP);
    tree->SetBranchAddress("tracks_global", &tracksG);


    for(int i=0;iGetEntries();i++){
        tree->GetEntry(i);
        Int_t nPrim = tracksP->GetEntriesFast();
        Int_t nGlob = tracksG->GetEntriesFast();
        if(nPrim>=10)
            cout<<Form("Ev %i Nprim = %i    Nglob = %i",i,nPrim,nGlob)<<endl;
    }
}

 

SetBranchStatus

W przykładzie wyżej analizowaliśmy zderzenia zawierające global i primar tracki, interesowały nas tylko te zderzenia które miały więcej niż 9 primary tracków. Ponieważ alternatywna wersja kodu do podobnej analizy:

void tree_wtcR(){
    TChain *tree = new TChain("tree");

    for(int i=1;i<=3;i++){
        tree->AddFile(Form("data_%i.root",i));
    }
    TClonesArray *tracksG = new TClonesArray("TLorentzVector",10);
    TClonesArray *tracksP = nullptr;

    tree->SetBranchAddress("tracks_primary", &tracksP);
    tree->SetBranchAddress("tracks_global", &tracksG);

    tree->SetBranchStatus("tracks_global*",0);
    tree->SetBranchStatus("tracks_primary*",1);

    for(int i=0;iGetEntries();i++){
        tree->GetEntry(i,0);
        Int_t nPrim = tracksP->GetEntriesFast();
        Int_t nGlob = 0;;

        if(nPrim>=10){
            tree->GetEntry(i,1);
            nGlob = tracksG->GetEntriesFast();
            cout<<Form("Ev %i Nprim = %i    Nglob = %i",i,nPrim,nGlob)<<endl;
        }
    }
}

Kod jest bardzo podobny jednak występują pewne różnice:

  • pojawia się SetBranchStatus, metoda ta ustawia status brancha, branche o statusie 1 są zawsze wczytywane (tzn. gdy zrobimy TChain::GetEntry(i) lub TChain::GetEntry(i,0)), branche o statusie 0 są wczytywane tylko na żądanie (tj. poprzez TChain::GetEntry(i,1))
  • w SetBranchStatus używamy "tracks_global*" gwiazdka oznacza, że zablokowany zostanie odczyt jakichkolwiek danych które  się zaczynaja na "tracks_global*" (bez gwiazdki nie zablokujemy odczytu global tracków), gwiazdka działa trochę jak wyrażenia regularne np. SetBranchStatus("*",0) blokuje odczyt wszystkich branchy
  • alokujemy miejsce pod tracksG (ponieważ gałąź domyślnie nie jest ładowana musimy sami zarezerwować miejsce)
  • gdy mamy więcej niż 9 primary tracków chcemy mieć global tracki - dlatego tutaj ładujemy wszystkie branche metodą TChain::GetEntry(i,1)

Podobną "zabawę" z blokowaniem branchy można przeprowadzić na obiektach TTree.