Project

General

Profile

warp_gsimple.C

example script for "warping" a flux file - Robert Hatcher, 08/16/2012 12:39 PM

 
1
//========================================================================
2
//
3
//  This is a script for creating a modified GSimpleNtpFlux file from
4
//  an original file, but modifying entries.  Assumes it will be passed
5
//  the name of a single file.
6
//
7
//  User must supply two functions "warp_entry()" and "warp_meta()" 
8
//
9
//========================================================================
10
#include "Numerical/RandomGen.h"
11
#include "FluxDrivers/GSimpleNtpFlux.h"
12
#include "Utils/UnitUtils.h"
13

    
14
#include "TSystem.h"
15
#include "TStopwatch.h"
16
#include "TLorentzVector.h"
17
#include "TNtuple.h"
18
#include "TFile.h"
19
#include "TChain.h"
20

    
21
#include <iostream>
22
#include <iomanip>
23
#include <string>
24
#include <sstream>
25
#include <set>
26
#include <algorithm>
27

    
28
using namespace std;
29
using namespace genie;
30
using namespace genie::flux;
31

    
32
bool gDoWarp = true;
33

    
34
//========================================================================
35
// this is what the USER must supply
36

    
37
void warp_entry(GSimpleNtpEntry* entry, 
38
                GSimpleNtpAux*   aux,
39
                GSimpleNtpNuMI*  numi)
40
{
41
  if ( ! gDoWarp ) return;
42
  // this is a dumb weighting scheme
43
  if ( entry->E > 1.5 ) {
44
    // don't set absolute weight in case we started with weighted entries
45
    // scale whatever is there from GSimpleNtpFlux driver
46
    entry->wgt *= 0.5;
47
  }
48
  // change some flavors
49
  if ( gRandom->Rndm() < 0.25 ) {
50
    entry->pdg = ( entry->pdg > 0 ) ? 16 : -16;
51
  }
52

    
53
  // use all branches so there are no complaints from compiler
54
  static double trash = 0;
55
  if ( numi ) trash = numi->entryno;
56
  if ( aux  ) trash = aux->auxint.size();
57
  if ( trash == -1 ) cout << "trash was -1" << endl;
58
}
59

    
60
void warp_meta(GSimpleNtpMeta* meta, string fnamein)
61
{
62
  meta->infiles.push_back("THIS IS MDC WARPED FLUX FROM:");
63
  meta->infiles.push_back(fnamein);
64
  if ( ! gDoWarp ) {
65
    meta->infiles.push_back("NO ACTUAL WARPING APPLIED");
66
  } else {
67
    string msg ="USING WARP FUNCTION CODENAMED WHISKEY-TANGO-FOXTROT";
68
    meta->infiles.push_back(msg);
69
  }
70
} 
71

    
72
void update_flavors(GSimpleNtpMeta* meta, std::set<int>& gFlavors);
73

    
74
//========================================================================
75
// main routine
76

    
77
// assumes a single input file (meta-data handling)
78

    
79

    
80
void warp_gsimple(string fnameout="gsimple_output.root",
81
                  string fnamein="gsimple_input.root",
82
                  bool dowarp=true)
83
{
84
  
85
  gDoWarp = dowarp;
86

    
87
  string fnameinlong(gSystem->ExpandPathName(fnamein.c_str()));
88
  TFile* fin = TFile::Open(fnameinlong.c_str(),"READONLY");
89
  TTree* etreein = (TTree*)fin->Get("flux");
90
  TTree* mtreein = (TTree*)fin->Get("meta");
91
  genie::flux::GSimpleNtpEntry* entry_in = new genie::flux::GSimpleNtpEntry;
92
  genie::flux::GSimpleNtpNuMI*  numi_in  = new genie::flux::GSimpleNtpNuMI;
93
  genie::flux::GSimpleNtpAux*   aux_in   = new genie::flux::GSimpleNtpAux;
94
  genie::flux::GSimpleNtpMeta*  meta_in  = new genie::flux::GSimpleNtpMeta;
95

    
96
  long int nentries = etreein->GetEntries();
97

    
98
  int sba_status[4] = { -999, -999, -999, -999 };
99
  sba_status[0] = etreein->SetBranchAddress("entry",&entry_in);
100
  sba_status[1] = etreein->SetBranchAddress("numi",&numi_in);
101
  sba_status[2] = etreein->SetBranchAddress("aux",&aux_in);
102
  sba_status[3] = mtreein->SetBranchAddress("meta",&meta_in);
103
  cout << "SetBranchAddress results "
104
       << sba_status[0] << "," 
105
       << sba_status[1] << "," 
106
       << sba_status[2] << ","
107
       << sba_status[3] 
108
       << endl;
109
  bool donumi = ( sba_status[1] == 0 );
110
  bool doaux  = ( sba_status[2] == 0 );
111

    
112
  int nindices = mtreein->BuildIndex("metakey"); // tie metadata to entry
113
  cout << "saw " << nindices << " metakey indices" << endl;
114

    
115
  cout << "Creating:    " << fnameout << endl;
116
  cout << "Input file:  " << fnamein << endl;
117
  cout << "Branches:    entry " 
118
       << ( (doaux) ? "aux ":"" ) 
119
       << ( (donumi) ? "numi ":"" ) 
120
       << endl;
121
  if ( ! gDoWarp ) cout << "++++++++ NO ACTUAL WARP APPLIED" << endl;
122

    
123
  TFile* fout = TFile::Open(fnameout.c_str(),"RECREATE");
124
  TTree* fluxntp = new TTree("flux","a simple flux n-tuple");
125
  TTree* metantp = new TTree("meta","metadata for flux n-tuple");
126
  genie::flux::GSimpleNtpEntry* fentry = new genie::flux::GSimpleNtpEntry;
127
  genie::flux::GSimpleNtpAux*   faux   = new genie::flux::GSimpleNtpAux;
128
  genie::flux::GSimpleNtpNuMI*  fnumi  = new genie::flux::GSimpleNtpNuMI;
129
  genie::flux::GSimpleNtpMeta*  fmeta  = new genie::flux::GSimpleNtpMeta;
130

    
131
  fluxntp->Branch("entry",&fentry);
132
  if ( doaux  ) fluxntp->Branch("aux",&faux);
133
  if ( donumi ) fluxntp->Branch("numi",&fnumi);
134
  metantp->Branch("meta",&fmeta);
135

    
136
  cout << "=========================== Start " << endl;
137

    
138
  TStopwatch sw;
139
  sw.Start();
140

    
141
  const double large = 1.0e300;
142
  double minwgt = large, maxwgt = -large, maxenergy = -large;
143
  std::set<int> gFlavors;
144

    
145
  for ( long int ientry = 0; ientry < nentries; ++ientry ) {
146

    
147
    // reset what's been read in anticipation of a new entry
148
    entry_in->Reset();
149
    aux_in->Reset();
150
    numi_in->Reset();
151
    
152
    // read the next entry, get metadata if it's different
153
    //int nbytes = 
154
    etreein->GetEntry(ientry);
155
    UInt_t metakey = entry_in->metakey;
156
    if ( fmeta->metakey != metakey ) {
157
      // UInt_t oldkey = meta_in->metakey;
158
      int nbmeta = mtreein->GetEntryWithIndex(metakey);
159
      cout << "on entry " << ientry << " fetched metadata w/ key "
160
           << metakey << " read " << nbmeta << " bytes" << endl;
161
    }
162

    
163
    // reset what we're writing
164
    fentry->Reset();
165
    fnumi->Reset();
166
    faux->Reset();
167

    
168
    // copy read in objects to output objects
169
    *fentry = *entry_in;
170
    *fnumi  = *numi_in;
171
    *faux   = *aux_in;
172

    
173
    // mess with the values to your hearts content
174
    warp_entry(fentry,faux,fnumi);
175
    // keep track of any weights for the meta data
176
    if ( fentry->wgt < minwgt ) minwgt = fentry->wgt;
177
    if ( fentry->wgt > maxwgt ) maxwgt = fentry->wgt;
178
    // user might have changed an energy larger than any in original file
179
    if ( fentry->E > maxenergy ) maxenergy = fentry->E;
180
    // user might have changed a flavor keep track of all we see
181
    gFlavors.insert(fentry->pdg);
182

    
183
    // process currently held metadata after transition to metadata key
184
    if ( fmeta->metakey != meta_in->metakey ) {
185
      // new meta data found
186
      cout << "new meta found " << fmeta->metakey << "/" << meta_in->metakey << endl;
187
      if ( fmeta->metakey != 0 ) {
188
        // have metadata that needs adjustment and writing
189
        // before processing new metadata
190
        warp_meta(fmeta,fnamein);
191
        fmeta->minWgt = minwgt;
192
        fmeta->maxWgt = maxwgt;
193
        fmeta->maxEnergy = maxenergy;
194
        update_flavors(fmeta,gFlavors);
195
        gFlavors.clear();
196
        cout << "metantp->Fill() " << *fmeta << endl;
197
        metantp->Fill();
198
        // next meta-data restarts weight range (and energy max)
199
        minwgt    =  large;
200
        maxwgt    = -large;
201
        maxenergy = -large;
202
      }
203
      // get a copy of metadata for accumulating adjustments
204
      *fmeta = *meta_in;
205
      //cout << " copy meta_in " << *meta_in << endl << *fmeta << endl;
206
    }
207
 
208
    fluxntp->Fill();
209
  }
210

    
211
  // write last set of meta-data after all is done
212
  if ( fmeta->metakey != 0 ) {
213
    cout << "final meta flush " << fmeta->metakey << endl;
214
    // have metadata that needs adjustment and writing
215
    warp_meta(fmeta,fnamein);
216
    fmeta->minWgt = minwgt;
217
    fmeta->maxWgt = maxwgt;
218
    fmeta->maxEnergy = maxenergy;
219
    update_flavors(fmeta,gFlavors);
220
    gFlavors.clear();
221
    cout << "metantp->Fill() " << *fmeta << endl;
222
    metantp->Fill();
223
  }
224

    
225
  cout << "=========================== Complete " << endl;
226

    
227
  sw.Stop();
228
  cout << "Generated " << nentries 
229
       << endl
230
       << "Time to generate: " << endl;
231
  sw.Print();
232
  
233
  fout->cd();
234

    
235
  // write ntuples out
236
  fluxntp->Write();
237
  metantp->Write();
238
  fout->Close();
239

    
240
  cout << endl << endl;
241
}
242

    
243
void update_flavors(GSimpleNtpMeta* meta, std::set<int>& gFlavors)
244
{
245
  // add any new flavors here, if necessary
246
  std::set<int>::const_iterator flitr = gFlavors.begin();
247
  std::vector<Int_t>& flvlist = meta->pdglist;
248
  for ( ; flitr != gFlavors.end(); ++flitr ) {
249
    int  seen_pdg = *flitr;
250
    std::vector<Int_t>::iterator knwitr = 
251
      std::find(flvlist.begin(),flvlist.end(),seen_pdg);
252
    bool known_pdg = ( knwitr != flvlist.end() );
253
    if ( ! known_pdg ) meta->pdglist.push_back(seen_pdg);
254
  }
255
}