eic-opticks/ana/gs.py

0001 #!/usr/bin/env python
0002 """
0003
0004 ::
0005
0006     gs.numPhotons[gs.numPhotons > 1]
0007
0008     In [13]: np.count_nonzero(gs.numPhotons == 1)   # reemission gensteps, need to be excluded
0009     Out[13]: 12275
0010
0011     In [14]: np.count_nonzero(gs.numPhotons > 1)
0012     Out[14]: 112
0013
0014
0015 """
0016 import os, numpy as np, argparse
0017 import logging
0018 log = logging.getLogger(__name__)
0019 from opticks.ana.OpticksGenstepEnum import OpticksGenstepEnum
0020 from opticks.ana.PDGCodeEnum import PDGCodeEnum
0021 from opticks.ana.nbase import count_unique_sorted
0022 X,Y,Z,T = 0,1,2,3
0023
0024
0025 def is_int(s):
0026     s = str(s)
0027     if s[0] in ('-', '+'):
0028         return s[1:].isdigit()
0029     return s.isdigit()
0030
0031 class GS(object):
0032     etyp = OpticksGenstepEnum()
0033     epdg = PDGCodeEnum()
0034
0035     @classmethod
0036     def parse_args(cls, doc, **kwa):
0037         np.set_printoptions(suppress=True, precision=3 )
0038         parser = argparse.ArgumentParser(doc)
0039         parser.add_argument(     "--pathtmpl",  help="Path template of genstep file", default=kwa.get("pathtmpl",None) )
0040         parser.add_argument(     "paths", nargs="*",  help="Paths of genstep files" )
0041         parser.add_argument(     "--level", default="info", help="logging level" )
0042         args = parser.parse_args()
0043         fmt = '[%(asctime)s] p%(process)s {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s'
0044         logging.basicConfig(level=getattr(logging,args.level.upper()), format=fmt)
0045         return args
0046
0047     def __init__(self, path, pathtmpl):
0048         if is_int(path):
0049             path = pathtmpl % int(path)
0050         pass
0051         f = np.load(os.path.expandvars(path))
0052         i = f.view(np.int32)
0053
0054         self.path = path
0055         self.f = f
0056         self.i = i
0057
0058         log.info(" path %s shape %s " % (self.path, str(self.f.shape)))
0059         self.check_counts()
0060         self.check_pdgcode()
0061         self.check_ranges()
0062
0063
0064     ID = property(lambda self:self.i[:,0,0])  # gencode: the OpticksGenstep enum value
0065     PID = property(lambda self:self.i[:,0,1]) # parentID
0066     numPhotons = property(lambda self:self.i[:,0,3])
0067
0068     xyzt = property(lambda self:self.f[:,1])
0069
0070     deltaPositionLength = property(lambda self:self.f[:,2])
0071
0072     pdgCode = property(lambda self:self.i[:,3,0])
0073     charge = property(lambda self:self.f[:,3,1])
0074     weight = property(lambda self:self.f[:,3,2])
0075     meanVelocity = property(lambda self:self.f[:,3,3])
0076
0077     def check_counts(self):
0078         log.info("check_counts")
0079         i = self.i
0080
0081         num_gensteps = len(i)
0082         num_photons = self.numPhotons.sum()
0083
0084         print("num_gensteps : %d " % num_gensteps )
0085         print("num_photons  : %d " % num_photons )
0086
0087         cu = count_unique_sorted(self.ID)
0088
0089         nph_tot = 0
0090         ngs_tot = 0
0091         fmt = " (%d)%-25s : ngs:%5d  npho:%5d "
0092         for typ,ngs in cu:
0093             sel = i[:,0,0] == typ
0094             nph = i[sel][:,0,3].sum()
0095             nph_tot += nph
0096             ngs_tot += ngs
0097             print(fmt % (typ,self.etyp(typ),ngs,nph ))
0098         pass
0099         print(fmt % (0, "TOTALS", ngs_tot, nph_tot))
0100
0101     def check_pdgcode(self):
0102         log.info("check_pdgcode")
0103         #cu = count_unique_sorted(self.pdgCode)  doesnt work with -ve coded
0104         #print(cu)
0105         pdgcodes, counts = np.unique(self.pdgCode, return_counts=True) # needs numpy > 1.9
0106         assert len(pdgcodes) == len(counts)
0107
0108         fmt = " %7d : %10s : %d "
0109         for i in range(len(pdgcodes)):
0110             pdgcode = pdgcodes[i]
0111             count = counts[i]
0112             print(fmt % (pdgcode,self.epdg(pdgcode),count ))
0113         pass
0114
0115     def check_ranges(self):
0116         log.info("check_ranges")
0117         f = self.f
0118
0119         xyzt = f[:,1]
0120         x,y,z,t = xyzt[:,X], xyzt[:,Y], xyzt[:,Z], xyzt[:,T]
0121
0122         tr = (t.min(), t.max())
0123         xr = (x.min(), x.max())
0124         yr = (y.min(), y.max())
0125         zr = (z.min(), z.max())
0126
0127         print(" tr %10.4f %10.4f " % tr )
0128         print(" xr %10.4f %10.4f " % xr )
0129         print(" yr %10.4f %10.4f " % yr )
0130         print(" zr %10.4f %10.4f " % zr )
0131
0132
0133
0134 if __name__ == '__main__':
0135     pathtmpl = "$TMP/source/evt/g4live/natural/%d/gs.npy"
0136     args = GS.parse_args(__doc__, pathtmpl=pathtmpl)
0137     log.info(args)
0138     for path in args.paths:
0139         gs = GS(path, args.pathtmpl)
0140     pass
0141
0142
0143