Project

General

Profile

warp_gsimple.C

Shih-kai Lin, 02/28/2016 10:24 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
// function prototype
33
double pick_energy(double, double);
34

    
35
bool gDoWarp = true;
36

    
37
//========================================================================
38
// this is what the USER must supply
39

    
40
void warp_entry(GSimpleNtpEntry* entry, 
41
                GSimpleNtpAux*   aux,
42
                GSimpleNtpNuMI*  numi)
43
{
44
  if ( ! gDoWarp ) return;
45
  // this is a dumb weighting scheme
46
  //if ( entry->E > 1.5 ) {
47
    //// don't set absolute weight in case we started with weighted entries
48
    //// scale whatever is there from GSimpleNtpFlux driver
49
    //entry->wgt *= 0.5;
50
  //}
51
  //// change some flavors
52
  //if ( gRandom->Rndm() < 0.25 ) {
53
    //entry->pdg = ( entry->pdg > 0 ) ? 16 : -16;
54
  //}
55
  
56
  // record the 4 momentum before warp
57
  TLorentzVector fmb4 = TLorentzVector(entry->px, entry->py, entry->pz,
58
                                       entry->E);
59
  // sample an energy value from the 1/E distribution
60
  // min energy 0.01 GeV, where GENIE spline starts
61
  // max energy 100 GeV, by eye
62
  entry->E = pick_energy(0.01, 20);
63
  // scale momentum so that the new momentum is in the same direction as the
64
  // old one but obeys the energy-momentum relation
65
  double k = sqrt(1+(entry->E*entry->E-fmb4.E()*fmb4.E())/fmb4.Vect().Mag2());
66
  entry->px *= k;
67
  entry->py *= k;
68
  entry->pz *= k;
69

    
70
  // use all branches so there are no complaints from compiler
71
  static double trash = 0;
72
  if ( numi ) trash = numi->entryno;
73
  if ( aux  ) trash = aux->auxint.size();
74
  if ( trash == -1 ) cout << "trash was -1" << endl;
75
}
76

    
77
void warp_meta(GSimpleNtpMeta* meta, string fnamein)
78
{
79
  meta->infiles.push_back("THIS IS MDC WARPED FLUX FROM:");
80
  meta->infiles.push_back(fnamein);
81
  if ( ! gDoWarp ) {
82
    meta->infiles.push_back("NO ACTUAL WARPING APPLIED");
83
  } else {
84
    string msg ="USING WARP FUNCTION CODENAMED WHISKEY-TANGO-FOXTROT";
85
    meta->infiles.push_back(msg);
86
  }
87
} 
88

    
89
double pick_energy(double emin, double emax)
90
{
91
  // pick a distribution of shape 1/E
92
  // between the values [emin, emax]
93
  // must have low energy cut to avoid infrared divergence
94
  double c = log(emin);
95
  double A = 1. / (log(emax) - c);
96
  double r = gRandom->Rndm(); // flat variate between (0,1]
97
  return exp(r/A + c);
98
}
99

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

    
102
//========================================================================
103
// main routine
104

    
105
// assumes a single input file (meta-data handling)
106

    
107

    
108
void warp_gsimple(string fnameout="gsimple_output.root",
109
                  string fnamein="gsimple_input.root",
110
                  bool dowarp=true)
111
{
112
  
113
  gDoWarp = dowarp;
114

    
115
  string fnameinlong(gSystem->ExpandPathName(fnamein.c_str()));
116
  TFile* fin = TFile::Open(fnameinlong.c_str(),"READONLY");
117
  TTree* etreein = (TTree*)fin->Get("flux");
118
  TTree* mtreein = (TTree*)fin->Get("meta");
119
  genie::flux::GSimpleNtpEntry* entry_in = new genie::flux::GSimpleNtpEntry;
120
  genie::flux::GSimpleNtpNuMI*  numi_in  = new genie::flux::GSimpleNtpNuMI;
121
  genie::flux::GSimpleNtpAux*   aux_in   = new genie::flux::GSimpleNtpAux;
122
  genie::flux::GSimpleNtpMeta*  meta_in  = new genie::flux::GSimpleNtpMeta;
123

    
124
  long int nentries = etreein->GetEntries();
125

    
126
  int sba_status[4] = { -999, -999, -999, -999 };
127
  sba_status[0] = etreein->SetBranchAddress("entry",&entry_in);
128
  sba_status[1] = etreein->SetBranchAddress("numi",&numi_in);
129
  sba_status[2] = etreein->SetBranchAddress("aux",&aux_in);
130
  sba_status[3] = mtreein->SetBranchAddress("meta",&meta_in);
131
  cout << "SetBranchAddress results "
132
       << sba_status[0] << "," 
133
       << sba_status[1] << "," 
134
       << sba_status[2] << ","
135
       << sba_status[3] 
136
       << endl;
137
  bool donumi = ( sba_status[1] == 0 );
138
  bool doaux  = ( sba_status[2] == 0 );
139

    
140
  int nindices = mtreein->BuildIndex("metakey"); // tie metadata to entry
141
  cout << "saw " << nindices << " metakey indices" << endl;
142

    
143
  cout << "Creating:    " << fnameout << endl;
144
  cout << "Input file:  " << fnamein << endl;
145
  cout << "Branches:    entry " 
146
       << ( (doaux) ? "aux ":"" ) 
147
       << ( (donumi) ? "numi ":"" ) 
148
       << endl;
149
  if ( ! gDoWarp ) cout << "++++++++ NO ACTUAL WARP APPLIED" << endl;
150

    
151
  TFile* fout = TFile::Open(fnameout.c_str(),"RECREATE");
152
  TTree* fluxntp = new TTree("flux","a simple flux n-tuple");
153
  TTree* metantp = new TTree("meta","metadata for flux n-tuple");
154
  genie::flux::GSimpleNtpEntry* fentry = new genie::flux::GSimpleNtpEntry;
155
  genie::flux::GSimpleNtpAux*   faux   = new genie::flux::GSimpleNtpAux;
156
  genie::flux::GSimpleNtpNuMI*  fnumi  = new genie::flux::GSimpleNtpNuMI;
157
  genie::flux::GSimpleNtpMeta*  fmeta  = new genie::flux::GSimpleNtpMeta;
158

    
159
  fluxntp->Branch("entry",&fentry);
160
  if ( doaux  ) fluxntp->Branch("aux",&faux);
161
  if ( donumi ) fluxntp->Branch("numi",&fnumi);
162
  metantp->Branch("meta",&fmeta);
163

    
164
  cout << "=========================== Start " << endl;
165

    
166
  TStopwatch sw;
167
  sw.Start();
168

    
169
  const double large = 1.0e300;
170
  double minwgt = large, maxwgt = -large, maxenergy = -large;
171
  std::set<int> gFlavors;
172

    
173
  for ( long int ientry = 0; ientry < nentries; ++ientry ) {
174

    
175
    // reset what's been read in anticipation of a new entry
176
    entry_in->Reset();
177
    aux_in->Reset();
178
    numi_in->Reset();
179
    
180
    // read the next entry, get metadata if it's different
181
    //int nbytes = 
182
    etreein->GetEntry(ientry);
183
    UInt_t metakey = entry_in->metakey;
184
    if ( fmeta->metakey != metakey ) {
185
      // UInt_t oldkey = meta_in->metakey;
186
      int nbmeta = mtreein->GetEntryWithIndex(metakey);
187
      cout << "on entry " << ientry << " fetched metadata w/ key "
188
           << metakey << " read " << nbmeta << " bytes" << endl;
189
    }
190

    
191
    // reset what we're writing
192
    fentry->Reset();
193
    fnumi->Reset();
194
    faux->Reset();
195

    
196
    // copy read in objects to output objects
197
    *fentry = *entry_in;
198
    *fnumi  = *numi_in;
199
    *faux   = *aux_in;
200

    
201
    // mess with the values to your hearts content
202
    warp_entry(fentry,faux,fnumi);
203
    // keep track of any weights for the meta data
204
    if ( fentry->wgt < minwgt ) minwgt = fentry->wgt;
205
    if ( fentry->wgt > maxwgt ) maxwgt = fentry->wgt;
206
    // user might have changed an energy larger than any in original file
207
    if ( fentry->E > maxenergy ) maxenergy = fentry->E;
208
    // user might have changed a flavor keep track of all we see
209
    gFlavors.insert(fentry->pdg);
210

    
211
    // process currently held metadata after transition to metadata key
212
    if ( fmeta->metakey != meta_in->metakey ) {
213
      // new meta data found
214
      cout << "new meta found " << fmeta->metakey << "/" << meta_in->metakey << endl;
215
      if ( fmeta->metakey != 0 ) {
216
        // have metadata that needs adjustment and writing
217
        // before processing new metadata
218
        warp_meta(fmeta,fnamein);
219
        fmeta->minWgt = minwgt;
220
        fmeta->maxWgt = maxwgt;
221
        fmeta->maxEnergy = maxenergy;
222
        update_flavors(fmeta,gFlavors);
223
        gFlavors.clear();
224
        cout << "metantp->Fill() " << *fmeta << endl;
225
        metantp->Fill();
226
        // next meta-data restarts weight range (and energy max)
227
        minwgt    =  large;
228
        maxwgt    = -large;
229
        maxenergy = -large;
230
      }
231
      // get a copy of metadata for accumulating adjustments
232
      *fmeta = *meta_in;
233
      //cout << " copy meta_in " << *meta_in << endl << *fmeta << endl;
234
    }
235
 
236
    fluxntp->Fill();
237
  }
238

    
239
  // write last set of meta-data after all is done
240
  if ( fmeta->metakey != 0 ) {
241
    cout << "final meta flush " << fmeta->metakey << endl;
242
    // have metadata that needs adjustment and writing
243
    warp_meta(fmeta,fnamein);
244
    fmeta->minWgt = minwgt;
245
    fmeta->maxWgt = maxwgt;
246
    fmeta->maxEnergy = maxenergy;
247
    update_flavors(fmeta,gFlavors);
248
    gFlavors.clear();
249
    cout << "metantp->Fill() " << *fmeta << endl;
250
    metantp->Fill();
251
  }
252

    
253
  cout << "=========================== Complete " << endl;
254

    
255
  sw.Stop();
256
  cout << "Generated " << nentries 
257
       << endl
258
       << "Time to generate: " << endl;
259
  sw.Print();
260
  
261
  fout->cd();
262

    
263
  // write ntuples out
264
  fluxntp->Write();
265
  metantp->Write();
266
  fout->Close();
267

    
268
  cout << endl << endl;
269
}
270

    
271
void update_flavors(GSimpleNtpMeta* meta, std::set<int>& gFlavors)
272
{
273
  // add any new flavors here, if necessary
274
  std::set<int>::const_iterator flitr = gFlavors.begin();
275
  std::vector<Int_t>& flvlist = meta->pdglist;
276
  for ( ; flitr != gFlavors.end(); ++flitr ) {
277
    int  seen_pdg = *flitr;
278
    std::vector<Int_t>::iterator knwitr = 
279
      std::find(flvlist.begin(),flvlist.end(),seen_pdg);
280
    bool known_pdg = ( knwitr != flvlist.end() );
281
    if ( ! known_pdg ) meta->pdglist.push_back(seen_pdg);
282
  }
283
}