eic-opticks/ana/evts.py

0001 #!/usr/bin/env python
0002 """
0003 evts.py
0004 ========
0005
0006 Try concatenation of event arrays::
0007
0008    ipython --pdb -i evts.py -- --pfx tds3gun --src natural
0009
0010    OR
0011
0012    ./evts.sh
0013
0014 """
0015 import json
0016 from opticks.ana.nload import A
0017 from opticks.ana.evt import Evt
0018
0019 class EvtConcatenate(object):
0020     mtems = ["fdom", "idom"]
0021     stems = ["ox", "rx", "bn", "ph"]
0022
0023     def __init__(self, ok):
0024
0025         self.ok = ok
0026
0027         self.aa = {}   # content arrays : ox, rx, bn, ph
0028         self.bb = {}
0029         self.am = {}   # metadata arrays : idom, fdom
0030         self.bm = {}
0031         self.aj = {}   # json metadata
0032         self.bj = {}
0033
0034         for stem in self.stems:
0035             self.aa[stem] = []
0036             self.bb[stem] = []
0037         pass
0038
0039     def load_all(self, n):
0040         for i in range(n):
0041             itag = 1+i
0042             self.load_one(itag)
0043         pass
0044
0045     def load_one(self, itag):
0046
0047         ok = self.ok
0048         aname = "a%d" % itag
0049         bname = "b%d" % itag
0050
0051         atag = "%s"% itag
0052         btag = "-%s"% itag
0053
0054         a = Evt(tag=atag, src=ok.src, det=ok.det, pfx=ok.pfx, args=ok)
0055         b = Evt(tag=btag, src=ok.src, det=ok.det, pfx=ok.pfx, args=ok)
0056
0057         ap = A.path_("ox", ok.src, atag, ok.det, ok.pfx)
0058         bp = A.path_("ox", ok.src, btag, ok.det, ok.pfx)
0059
0060         print(ap)
0061         print(a.ox.path)
0062
0063         print(bp)
0064         print(b.ox.path)
0065
0066         globals()[aname] = a
0067         globals()[bname] = b
0068
0069         if itag == 1:
0070             # TODO: should really assert-compare and add
0071             self.aj["parameters.json"] = a.metadata.parameters
0072             self.bj["parameters.json"] = b.metadata.parameters
0073         pass
0074
0075         ## necessary metadata arrays : fdom, idom
0076         for mtem in self.mtems:
0077             a_ = getattr(a, mtem)
0078             b_ = getattr(b, mtem)
0079
0080             if not mtem in self.am:
0081                 self.am[mtem] = a_
0082             else:
0083                 assert np.all( self.am[mtem] == a_ )
0084             pass
0085             if not mtem in self.bm:
0086                 self.bm[mtem] = b_
0087             else:
0088                 assert np.all( self.bm[mtem] == b_ )
0089             pass
0090         pass
0091
0092         ## content arrays : ox, rx, ph, ...
0093         for stem in self.stems:
0094
0095             a_ = getattr(a, stem)
0096             b_ = getattr(b, stem)
0097
0098             if not a_.missing and len(a_) > 0:
0099                 self.aa[stem].append(a_)
0100             pass
0101             if not b_.missing and len(b_) > 0:
0102                 self.bb[stem].append(b_)
0103             pass
0104         pass
0105     pass
0106
0107     def save_concat(self, jtag):
0108         log.info("save_concat jtag:%d" % jtag)
0109         ok = self.ok
0110         atag_c = "%s"%jtag
0111         btag_c = "-%s"%jtag
0112
0113         for mtem in self.mtems:
0114             ap = A.path_(mtem, ok.src, atag_c, ok.det, ok.pfx)
0115             bp = A.path_(mtem, ok.src, btag_c, ok.det, ok.pfx)
0116             ad = os.path.dirname(ap)
0117             bd = os.path.dirname(bp)
0118             if not os.path.isdir(ad):
0119                 os.makedirs(ad)
0120             pass
0121             if not os.path.isdir(bd):
0122                 os.makedirs(bd)
0123             pass
0124             log.info("save %s " % ap)
0125             np.save(ap, self.am[mtem] )
0126             log.info("save %s " % bp)
0127             np.save(bp, self.bm[mtem] )
0128
0129             last =  mtem == self.mtems[-1]
0130             if last:
0131                 for k in self.aj.keys():  ## keys are file names eg parameters.json
0132                     json.dump(self.aj[k], open(os.path.join(ad, k), "w"))
0133                     json.dump(self.bj[k], open(os.path.join(bd, k), "w"))
0134                 pass
0135             pass
0136         pass
0137
0138         for stem in self.stems:
0139             ap = A.path_(stem, ok.src, atag_c, ok.det, ok.pfx)
0140             bp = A.path_(stem, ok.src, btag_c, ok.det, ok.pfx)
0141
0142             aa_ = tuple(self.aa[stem])
0143             if len(aa_) > 0:
0144                 ac =  np.concatenate(aa_)
0145                 log.info("save %s " % ap)
0146                 np.save(ap, ac)
0147                 globals()["a_"+stem] = ac
0148             pass
0149
0150             bb_ = tuple(self.bb[stem])
0151             if len(bb_) > 0:
0152                 bc =  np.concatenate(bb_)
0153                 log.info("save %s " % bp)
0154                 np.save(bp, bc)
0155                 globals()["b_"+stem] = bc
0156             pass
0157         pass
0158
0159
0160 if __name__ == '__main__':
0161     from opticks.ana.main import opticks_main
0162     ok = opticks_main()
0163     ec = EvtConcatenate(ok)
0164     ec.load_all(7)           # load tags 1 to 7
0165     ec.save_concat(100)      # save into new tag 100
0166